From e0d4191c70facf41475b75f7dd72b32df4312f21 Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Thu, 8 May 2025 15:12:40 +0000 Subject: [PATCH] internal/mcp: a new extensibility point for HTTP transports Expose the SSEServerTransport type, so that clients can customize the handling of sessions and endpoints by writing their own HTTP handler. This is similar to the behavior of the typescript SDK. Also, finish documenting transports in the design doc. Change-Id: Ief5aa8424ba2946a9615e366969a4750f40224a8 Reviewed-on: https://go-review.googlesource.com/c/tools/+/671015 LUCI-TryBot-Result: Go LUCI Reviewed-by: Jonathan Amsterdam Auto-Submit: Robert Findley --- design/design.md | 83 ++++++++++++++++++-- mcp/sse.go | 199 ++++++++++++++++++++++++++++------------------- 2 files changed, 193 insertions(+), 89 deletions(-) diff --git a/design/design.md b/design/design.md index 6b6a29e..6b76e5d 100644 --- a/design/design.md +++ b/design/design.md @@ -128,7 +128,7 @@ type CommandTransport struct { /* unexported fields */ } func NewCommandTransport(cmd *exec.Command) *CommandTransport // Connect starts the command, and connects to it over stdin/stdout. -func (t *CommandTransport) Connect(ctx context.Context) (Stream, error) { +func (*CommandTransport) Connect(ctx context.Context) (Stream, error) { ``` The `StdIOTransport` is the server side of the stdio transport, and connects by @@ -139,7 +139,7 @@ binding to `os.Stdin` and `os.Stdout`. // JSON over stdin/stdout. type StdIOTransport struct { /* unexported fields */ } -func NewStdIOTransport() *StdIOTransport { +func NewStdIOTransport() *StdIOTransport func (t *StdIOTransport) Connect(context.Context) (Stream, error) ``` @@ -174,13 +174,80 @@ Notably absent are options to hook into the request handling for the purposes of authentication or context injection. These concerns are better handled using standard HTTP middleware patterns. - +By default, the SSE handler creates messages endpoints with the +`?sessionId=...` query parameter. Users that want more control over the +management of sessions and session endpoints may write their own handler, and +create `SSEServerTransport` instances themselves, for incoming GET requests. - +```go +// A SSEServerTransport is a logical SSE session created through a hanging GET +// request. +// +// When connected, it it returns the following [Stream] implementation: +// - Writes are SSE 'message' events to the GET response. +// - Reads are received from POSTs to the session endpoint, via +// [SSEServerTransport.ServeHTTP]. +// - Close terminates the hanging GET. +type SSEServerTransport struct { /* ... */ } + +// NewSSEServerTransport creates a new SSE transport for the given messages +// endpoint, and hanging GET response. +// +// Use [SSEServerTransport.Connect] to initiate the flow of messages. +// +// The transport is itself an [http.Handler]. It is the caller's responsibility +// to ensure that the resulting transport serves HTTP requests on the given +// session endpoint. +func NewSSEServerTransport(endpoint string, w http.ResponseWriter) *SSEServerTransport + +// ServeHTTP handles POST requests to the transport endpoint. +func (*SSEServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request) + +// Connect sends the 'endpoint' event to the client. +// See [SSEServerTransport] for more details on the [Stream] implementation. +func (*SSEServerTransport) Connect(context.Context) (Stream, error) +``` + +The SSE client transport is simpler, and hopefully self-explanatory. + +```go +type SSEClientTransport struct { /* ... */ } + +// NewSSEClientTransport returns a new client transport that connects to the +// SSE server at the provided URL. +// +// NewSSEClientTransport panics if the given URL is invalid. +func NewSSEClientTransport(url string) *SSEClientTransport { + +// Connect connects through the client endpoint. +func (*SSEClientTransport) Connect(ctx context.Context) (Stream, error) +``` + +The Streamable HTTP transports are similar to the SSE transport, albeit with a +more complicated implementation. For brevity, we summarize only the differences +from the equivalent SSE types: + +```go +// The StreamableHandler interface is symmetrical to the SSEHandler. +type StreamableHandler struct { /* unexported fields */ } +func NewStreamableHandler(getServer func(request *http.Request) *Server) *StreamableHandler +func (*StreamableHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) +func (*StreamableHandler) Close() error + +// Unlike the SSE transport, the streamable transport constructor accepts a +// session ID, not an endpoint, along with the http response for the request +// that created the session. It is the caller's responsibility to delegate +// requests to this session. +type StreamableServerTransport struct { /* ... */ } +func NewStreamableServerTransport(sessionID string, w http.ResponseWriter) *StreamableServerTransport +func (*StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request) +func (*StreamableServerTransport) Connect(context.Context) (Stream, error) + +// The streamable client handles reconnection transparently to the user. +type StreamableClientTransport struct { /* ... */ } +func NewStreamableClientTransport(url string) *StreamableClientTransport { +func (*StreamableClientTransport) Connect(context.Context) (Stream, error) +``` ### Protocol types diff --git a/mcp/sse.go b/mcp/sse.go index 2104755..ef8330d 100644 --- a/mcp/sse.go +++ b/mcp/sse.go @@ -67,27 +67,38 @@ type SSEHandler struct { onConnection func(*ServerConnection) // for testing; must not block mu sync.Mutex - sessions map[string]*sseSession + sessions map[string]*SSEServerTransport } -// NewSSEHandler returns a new [SSEHandler] that is ready to serve HTTP. +// NewSSEHandler returns a new [SSEHandler] that creates and manages MCP +// sessions created via incoming HTTP requests. // -// The getServer function is used to bind created servers for new sessions. It -// is OK for getServer to return the same server multiple times. +// Sessions are created when the client issues a GET request to the server, +// which must accept text/event-stream responses (server-sent events). +// For each such request, a new [SSEServerTransport] is created with a distinct +// messages endpoint, and connected to the server returned by getServer. It is +// up to the user whether getServer returns a distinct [Server] for each new +// request, or reuses an existing server. +// +// The SSEHandler also handles requests to the message endpoints, by +// delegating them to the relevant server transport. func NewSSEHandler(getServer func(request *http.Request) *Server) *SSEHandler { return &SSEHandler{ getServer: getServer, - sessions: make(map[string]*sseSession), + sessions: make(map[string]*SSEServerTransport), } } -// A sseSession is a logical jsonrpc2 stream implementing the server side of -// MCP SSE transport, initiated through the hanging GET -// - Writes are SSE 'message' events to the GET response body. -// - Reads are received from POSTs to the session endpoing, mediated through a -// buffered channel. +// A SSEServerTransport is a logical SSE session created through a hanging GET +// request. +// +// When connected, it returns the following [Stream] implementation: +// - Writes are SSE 'message' events to the GET response. +// - Reads are received from POSTs to the session endpoint, via +// [SSEServerTransport.ServeHTTP]. // - Close terminates the hanging GET. -type sseSession struct { +type SSEServerTransport struct { + endpoint string incoming chan jsonrpc2.Message // queue of incoming messages; never closed // We must guard both pushes to the incoming queue and writes to the response @@ -100,9 +111,63 @@ type sseSession struct { done chan struct{} // closed when the stream is closed } -// Connect returns the receiver, as an sseSession is a logical stream. -func (s *sseSession) Connect(context.Context) (Stream, error) { - return s, nil +// NewSSEServerTransport creates a new SSE transport for the given messages +// endpoint, and hanging GET response. +// +// Use [SSEServerTransport.Connect] to initiate the flow of messages. +// +// The transport is itself an [http.Handler]. It is the caller's responsibility +// to ensure that the resulting transport serves HTTP requests on the given +// session endpoint. +// +// Most callers should instead use an [SSEHandler], which transparently handles +// the delegation to SSEServerTransports. +func NewSSEServerTransport(endpoint string, w http.ResponseWriter) *SSEServerTransport { + return &SSEServerTransport{ + endpoint: endpoint, + w: w, + incoming: make(chan jsonrpc2.Message, 100), + done: make(chan struct{}), + } +} + +// ServeHTTP handles POST requests to the transport endpoint. +func (t *SSEServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request) { + // Read and parse the message. + data, err := io.ReadAll(req.Body) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + // Optionally, we could just push the data onto a channel, and let the + // message fail to parse when it is read. This failure seems a bit more + // useful + msg, err := jsonrpc2.DecodeMessage(data) + if err != nil { + http.Error(w, "failed to parse body", http.StatusBadRequest) + return + } + select { + case t.incoming <- msg: + w.WriteHeader(http.StatusAccepted) + case <-t.done: + http.Error(w, "session closed", http.StatusBadRequest) + } +} + +// Connect sends the 'endpoint' event to the client. +// See [SSEServerTransport] for more details on the [Stream] implementation. +func (t *SSEServerTransport) Connect(context.Context) (Stream, error) { + t.mu.Lock() + _, err := writeEvent(t.w, event{ + name: "endpoint", + data: []byte(t.endpoint), + }) + t.mu.Unlock() + if err != nil { + return nil, err + } + return sseServerStream{t}, nil } func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { @@ -125,26 +190,7 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { return } - // Read and parse the message. - data, err := io.ReadAll(req.Body) - if err != nil { - http.Error(w, "failed to read body", http.StatusBadRequest) - return - } - // Optionally, we could just push the data onto a channel, and let the - // message fail to parse when it is read. This failure seems a bit more - // useful - msg, err := jsonrpc2.DecodeMessage(data) - if err != nil { - http.Error(w, "failed to parse body", http.StatusBadRequest) - return - } - select { - case session.incoming <- msg: - w.WriteHeader(http.StatusAccepted) - case <-session.done: - http.Error(w, "session closed", http.StatusBadRequest) - } + session.ServeHTTP(w, req) return } @@ -163,16 +209,27 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { w.Header().Set("Connection", "keep-alive") sessionID = randText() - session := &sseSession{ - w: w, - incoming: make(chan jsonrpc2.Message, 100), - done: make(chan struct{}), + endpoint, err := req.URL.Parse("?sessionid=" + sessionID) + if err != nil { + http.Error(w, "internal error: failed to create endpoint", http.StatusInternalServerError) + return } - defer session.Close() + + transport := NewSSEServerTransport(endpoint.RequestURI(), w) + + // The session is terminated when the request exits. + h.mu.Lock() + h.sessions[sessionID] = transport + h.mu.Unlock() + defer func() { + h.mu.Lock() + delete(h.sessions, sessionID) + h.mu.Unlock() + }() // TODO(hxjiang): getServer returns nil will panic. server := h.getServer(req) - cc, err := server.Connect(req.Context(), session, nil) + cc, err := server.Connect(req.Context(), transport, nil) if err != nil { http.Error(w, "connection failed", http.StatusInternalServerError) return @@ -182,52 +239,32 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { } defer cc.Close() // close the transport when the GET exits - // The session is terminated when the request exits. - h.mu.Lock() - h.sessions[sessionID] = session - h.mu.Unlock() - defer func() { - h.mu.Lock() - delete(h.sessions, sessionID) - h.mu.Unlock() - }() - - endpoint, err := req.URL.Parse("?sessionid=" + sessionID) - if err != nil { - http.Error(w, "internal error: failed to create endpoint", http.StatusInternalServerError) - return - } - - session.mu.Lock() - _, err = writeEvent(w, event{ - name: "endpoint", - data: []byte(endpoint.RequestURI()), - }) - session.mu.Unlock() - if err != nil { - return // too late to write the status header - } - select { case <-req.Context().Done(): - case <-session.done: + case <-transport.done: } } +// sseServerStream implements the Stream interface for a single [SSEServerTransport]. +// It hides the Stream interface from the SSEServerTransport API. +type sseServerStream struct { + t *SSEServerTransport +} + // Read implements jsonrpc2.Reader. -func (s *sseSession) Read(ctx context.Context) (jsonrpc2.Message, int64, error) { +func (s sseServerStream) Read(ctx context.Context) (jsonrpc2.Message, int64, error) { select { case <-ctx.Done(): return nil, 0, ctx.Err() - case msg := <-s.incoming: + case msg := <-s.t.incoming: return msg, 0, nil - case <-s.done: + case <-s.t.done: return nil, 0, io.EOF } } // Write implements jsonrpc2.Writer. -func (s *sseSession) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) { +func (s sseServerStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) { if ctx.Err() != nil { return 0, ctx.Err() } @@ -237,17 +274,17 @@ func (s *sseSession) Write(ctx context.Context, msg jsonrpc2.Message) (int64, er return 0, err } - s.mu.Lock() - defer s.mu.Unlock() + s.t.mu.Lock() + defer s.t.mu.Unlock() // Note that it is invalid to write to a ResponseWriter after ServeHTTP has // exited, and so we must lock around this write and check isDone, which is // set before the hanging GET exits. - if s.closed { + if s.t.closed { return 0, io.EOF } - n, err := writeEvent(s.w, event{name: "message", data: data}) + n, err := writeEvent(s.t.w, event{name: "message", data: data}) return int64(n), err } @@ -256,12 +293,12 @@ func (s *sseSession) Write(ctx context.Context, msg jsonrpc2.Message) (int64, er // It must be safe to call Close more than once, as the close may // asynchronously be initiated by either the server closing its connection, or // by the hanging GET exiting. -func (s *sseSession) Close() error { - s.mu.Lock() - defer s.mu.Unlock() - if !s.closed { - s.closed = true - close(s.done) +func (s sseServerStream) Close() error { + s.t.mu.Lock() + defer s.t.mu.Unlock() + if !s.t.closed { + s.t.closed = true + close(s.t.done) } return nil }