用Go从零实现MCP协议SSE传输层：原理与实战

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/google/uuid"
)

type JSONRPCMessage struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      any             `json:"id,omitempty"`
	Method  string          `json:"method,omitempty"`
	Params  json.RawMessage `json:"params,omitempty"`
	Result  json.RawMessage `json:"result,omitempty"`
	Error   *JSONRPCError   `json:"error,omitempty"`
}

type JSONRPCError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

type Session struct {
	ID       string
	messages chan JSONRPCMessage
	done     chan struct{}
	ctx      context.Context
	cancel   context.CancelFunc
}

type SSETransport struct {
	sessions   map[string]*Session
	mu         sync.RWMutex
	basePath   string
	onMessage  func(session *Session, msg JSONRPCMessage)
	heartbeat  time.Duration
}

func NewSSETransport(basePath string, heartbeat time.Duration) *SSETransport {
	return &SSETransport{
		sessions:  make(map[string]*Session),
		basePath:  basePath,
		heartbeat: heartbeat,
	}
}

func (t *SSETransport) SetMessageHandler(handler func(session *Session, msg JSONRPCMessage)) {
	t.onMessage = handler
}

func (t *SSETransport) SendToSession(sessionID string, msg JSONRPCMessage) error {
	t.mu.RLock()
	session, exists := t.sessions[sessionID]
	t.mu.RUnlock()

	if !exists {
		return fmt.Errorf("session %s not found", sessionID)
	}

	select {
	case session.messages <- msg:
		return nil
	case <-session.ctx.Done():
		return fmt.Errorf("session %s closed", sessionID)
	case <-time.After(5 * time.Second):
		return fmt.Errorf("send to session %s timed out", sessionID)
	}
}

func (t *SSETransport) HandleSSE(w http.ResponseWriter, r *http.Request) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming not supported", http.StatusInternalServerError)
		return
	}

	ctx, cancel := context.WithCancel(r.Context())
	session := &Session{
		ID:       uuid.New().String(),
		messages: make(chan JSONRPCMessage, 64),
		done:     make(chan struct{}),
		ctx:      ctx,
		cancel:   cancel,
	}

	t.mu.Lock()
	t.sessions[session.ID] = session
	t.mu.Unlock()

	defer func() {
		t.mu.Lock()
		delete(t.sessions, session.ID)
		t.mu.Unlock()
		cancel()
		close(session.done)
		log.Printf("session %s disconnected", session.ID)
	}()

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")

	endpoint := fmt.Sprintf("%s/message?sessionId=%s", t.basePath, session.ID)
	fmt.Fprintf(w, "event: endpoint\ndata: %s\n\n", endpoint)
	flusher.Flush()
	log.Printf("session %s connected, endpoint: %s", session.ID, endpoint)

	heartbeatTicker := time.NewTicker(t.heartbeat)
	defer heartbeatTicker.Stop()

	for {
		select {
		case msg := <-session.messages:
			data, err := json.Marshal(msg)
			if err != nil {
				log.Printf("marshal error for session %s: %v", session.ID, err)
				continue
			}
			fmt.Fprintf(w, "event: message\ndata: %s\n\n", data)
			flusher.Flush()

		case <-heartbeatTicker.C:
			fmt.Fprintf(w, ": heartbeat\n\n")
			flusher.Flush()

		case <-ctx.Done():
			return
		}
	}
}

func (t *SSETransport) HandleMessage(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	sessionID := r.URL.Query().Get("sessionId")
	if sessionID == "" {
		http.Error(w, "missing sessionId", http.StatusBadRequest)
		return
	}

	t.mu.RLock()
	session, exists := t.sessions[sessionID]
	t.mu.RUnlock()

	if !exists {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}

	var msg JSONRPCMessage
	decoder := json.NewDecoder(r.Body)
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&msg); err != nil {
		resp := JSONRPCMessage{
			JSONRPC: "2.0",
			Error: &JSONRPCError{
				Code:    -32700,
				Message: "parse error: " + err.Error(),
			},
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadRequest)
		json.NewEncoder(w).Encode(resp)
		return
	}

	if t.onMessage != nil {
		t.onMessage(session, msg)
	}

	w.WriteHeader(http.StatusAccepted)
}

func main() {
	transport := NewSSETransport("", 30*time.Second)

	transport.SetMessageHandler(func(session *Session, msg JSONRPCMessage) {
		var response JSONRPCMessage
		response.JSONRPC = "2.0"
		response.ID = msg.ID

		switch msg.Method {
		case "initialize":
			result, _ := json.Marshal(map[string]any{
				"protocolVersion": "2024-11-05",
				"capabilities": map[string]any{
					"tools": map[string]any{},
				},
				"serverInfo": map[string]any{
					"name":    "go-mcp-server",
					"version": "1.0.0",
				},
			})
			response.Result = result

		case "notifications/initialized":
			return

		case "tools/list":
			result, _ := json.Marshal(map[string]any{
				"tools": []map[string]any{
					{
						"name":        "get_timestamp",
						"description": "Returns the current Unix timestamp",
						"inputSchema": map[string]any{
							"type":       "object",
							"properties": map[string]any{},
						},
					},
				},
			})
			response.Result = result

		case "tools/call":
			var params struct {
				Name string `json:"name"`
			}
			json.Unmarshal(msg.Params, &params)

			if params.Name == "get_timestamp" {
				content, _ := json.Marshal(map[string]any{
					"content": []map[string]any{
						{
							"type": "text",
							"text": fmt.Sprintf("Current timestamp: %d", time.Now().Unix()),
						},
					},
				})
				response.Result = content
			} else {
				response.Error = &JSONRPCError{
					Code:    -32601,
					Message: "tool not found: " + params.Name,
				}
			}

		default:
			response.Error = &JSONRPCError{
				Code:    -32601,
				Message: "method not found: " + msg.Method,
			}
		}

		if err := transport.SendToSession(session.ID, response); err != nil {
			log.Printf("failed to send response: %v", err)
		}
	})

	mux := http.NewServeMux()
	mux.HandleFunc("/sse", transport.HandleSSE)
	mux.HandleFunc("/message", transport.HandleMessage)

	server := &http.Server{
		Addr:         ":8080",
		Handler:      mux,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 0,
		IdleTimeout:  120 * time.Second,
	}

	log.Println("MCP SSE Server listening on :8080")
	log.Println("  SSE endpoint:     GET  /sse")
	log.Println("  Message endpoint: POST /message?sessionId=<id>")

	if err := server.ListenAndServe(); err != nil {
		log.Fatal(err)
	}
}

# 终端1：建立SSE连接
curl -N http://localhost:8080/sse

# 输出：
# event: endpoint
# data: /message?sessionId=550e8400-e29b-41d4-a716-446655440000

# 终端2：发送initialize请求
curl -X POST "http://localhost:8080/message?sessionId=550e8400-e29b-41d4-a716-446655440000" \
  -H "Content-Type: application/json" \
  -d '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{}}}'

# 终端1的SSE流会收到：
# event: message
# data: {"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2024-11-05","capabilities":{"tools":{}},"serverInfo":{"name":"go-mcp-server","version":"1.0.0"}}}

func (t *SSETransport) Shutdown(ctx context.Context) {
	t.mu.Lock()
	sessions := make([]*Session, 0, len(t.sessions))
	for _, s := range t.sessions {
		sessions = append(sessions, s)
	}
	t.mu.Unlock()

	for _, s := range sessions {
		s.cancel()
		select {
		case <-s.done:
		case <-ctx.Done():
			return
		}
	}
}

ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()

go func() {
	if err := server.ListenAndServe(); err != http.ErrServerClosed {
		log.Fatal(err)
	}
}()

<-ctx.Done()
log.Println("shutting down...")

shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()

transport.Shutdown(shutdownCtx)
server.Shutdown(shutdownCtx)

func (t *SSETransport) safeHandleMessage(session *Session, msg JSONRPCMessage) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("panic recovered in session %s: %v", session.ID, r)
			errResp := JSONRPCMessage{
				JSONRPC: "2.0",
				ID:      msg.ID,
				Error: &JSONRPCError{
					Code:    -32603,
					Message: "internal error",
				},
			}
			transport.SendToSession(session.ID, errResp)
		}
	}()
	t.onMessage(session, msg)
}