From 1fead0ba64525932f9ca1bd687ee622250106035 Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Tue, 17 Jun 2025 16:56:36 +0000 Subject: [PATCH] internal/mcp: rename Stream->Connection After implementing the Streamable transport, it is apparent to me that the name 'Stream' is confusing, as it can refer either to a stream in the streamable transport, or to a JSON-RPC stream. But the name 'Connection' is available, and a natural name for the result of the Transport.Connect interface method. Therefore, rename Stream to Connection, and update the streamable implementation to refer to its HTTP connections as 'streams', consistent with other SDKs. Change-Id: Ica55f84caabcbfa5a2fe029cc0c2db57da574875 Reviewed-on: https://go-review.googlesource.com/c/tools/+/682376 LUCI-TryBot-Result: Go LUCI Reviewed-by: Jonathan Amsterdam --- design/design.md | 32 +++++------ mcp/cmd.go | 4 +- mcp/mcp.go | 2 +- mcp/mcp_test.go | 2 +- mcp/sse.go | 38 ++++++------- mcp/streamable.go | 120 +++++++++++++++++++++--------------------- mcp/transport.go | 59 ++++++++++----------- mcp/transport_test.go | 12 +---- 8 files changed, 127 insertions(+), 142 deletions(-) diff --git a/design/design.md b/design/design.md index abe7d8d..af3ca6f 100644 --- a/design/design.md +++ b/design/design.md @@ -44,7 +44,7 @@ The JSON-RPC implementation is hidden, to avoid tight coupling. As described in ### JSON-RPC and Transports -The MCP is defined in terms of client-server communication over bidirectional JSON-RPC message streams. Specifically, version `2025-03-26` of the spec defines two transports: +The MCP is defined in terms of client-server communication over bidirectional JSON-RPC message connections. Specifically, version `2025-03-26` of the spec defines two transports: - **stdio**: communication with a subprocess over stdin/stdout. - **streamable http**: communication over a relatively complicated series of text/event-stream GET and HTTP POST requests. @@ -55,7 +55,7 @@ Additionally, version `2024-11-05` of the spec defined a simpler (yet stateful) Furthermore, the spec [states](https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#custom-transports) that it must be possible for users to define their own custom transports. -Given the diversity of the transport implementations, they can be challenging to abstract. However, since JSON-RPC requires a bidirectional stream, we can use this to model the MCP transport abstraction: +Given the diversity of the transport implementations, they can be challenging to abstract. However, since JSON-RPC requires a bidirectional connection, we can use this to model the MCP transport abstraction: ```go type ( @@ -68,11 +68,11 @@ type ( // A Transport is used to create a bidirectional connection between MCP client // and server. type Transport interface { - Connect(ctx context.Context) (Stream, error) + Connect(ctx context.Context) (Connection, error) } -// A Stream is a bidirectional jsonrpc2 Stream. -type Stream interface { +// A Connection is a bidirectional jsonrpc2 connection. +type Connection interface { Read(ctx context.Context) (JSONRPCMessage, error) Write(ctx context.Context, JSONRPCMessage) error Close() error @@ -81,9 +81,9 @@ type Stream interface { Methods accept a Go `Context` and return an `error`, as is idiomatic for APIs that do I/O. -A `Transport` is something that connects a logical JSON-RPC stream, and nothing more. Streams must be closeable in order to implement client and server shutdown, and therefore conform to the `io.Closer` interface. +A `Transport` is something that connects a logical JSON-RPC connection, and nothing more. Connections must be closeable in order to implement client and server shutdown, and therefore conform to the `io.Closer` interface. -Other SDKs define higher-level transports, with, for example, methods to send a notification or make a call. Those are jsonrpc2 operations on top of the logical stream, and the lower-level interface is easier to implement in most cases, which means it is easier to implement custom transports. +Other SDKs define higher-level transports, with, for example, methods to send a notification or make a call. Those are jsonrpc2 operations on top of the logical connection, and the lower-level interface is easier to implement in most cases, which means it is easier to implement custom transports. For our prototype, we've used an internal `jsonrpc2` package based on the Go language server `gopls`, which we propose to fork for the MCP SDK. It already handles concerns like client/server connection, request lifecycle, cancellation, and shutdown. @@ -95,7 +95,7 @@ The `Transport` interface here is lower-level than that of mcp-go, but serves a In the MCP Spec, the **stdio** transport uses newline-delimited JSON to communicate over stdin/stdout. It's possible to model both client side and server side of this communication with a shared type that communicates over an `io.ReadWriteCloser`. However, for the purposes of future-proofing, we should use a different types for client and server stdio transport. -The `CommandTransport` is the client side of the stdio transport, and connects by starting a command and binding its jsonrpc2 stream to its stdin/stdout. +The `CommandTransport` is the client side of the stdio transport, and connects by starting a command and streaming JSON-RPC messages over stdin/stdout. ```go // A CommandTransport is a [Transport] that runs a command and communicates @@ -107,7 +107,7 @@ type CommandTransport struct { /* unexported fields */ } func NewCommandTransport(cmd *exec.Command) *CommandTransport // Connect starts the command, and connects to it over stdin/stdout. -func (*CommandTransport) Connect(ctx context.Context) (Stream, error) { +func (*CommandTransport) Connect(ctx context.Context) (Connection, error) { ``` The `StdIOTransport` is the server side of the stdio transport, and connects by binding to `os.Stdin` and `os.Stdout`. @@ -119,7 +119,7 @@ type StdIOTransport struct { /* unexported fields */ } func NewStdIOTransport() *StdIOTransport -func (t *StdIOTransport) Connect(context.Context) (Stream, error) +func (t *StdIOTransport) Connect(context.Context) (Connection, error) ``` #### HTTP transports @@ -154,7 +154,7 @@ By default, the SSE handler creates messages endpoints with the `?sessionId=...` // A SSEServerTransport is a logical SSE session created through a hanging GET // request. // -// When connected, it returns the following [Stream] implementation: +// When connected, it returns the following [Connection] implementation: // - Writes are SSE 'message' events to the GET response. // - Reads are received from POSTs to the session endpoint, via // [SSEServerTransport.ServeHTTP]. @@ -178,8 +178,8 @@ func NewSSEServerTransport(endpoint string, w http.ResponseWriter) *SSEServerTra 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) +// See [SSEServerTransport] for more details on the [Connection] implementation. +func (*SSEServerTransport) Connect(context.Context) (Connection, error) ``` The SSE client transport is simpler, and hopefully self-explanatory. @@ -192,7 +192,7 @@ type SSEClientTransport struct { /* ... */ } func NewSSEClientTransport(url string) (*SSEClientTransport, error) { // Connect connects through the client endpoint. -func (*SSEClientTransport) Connect(ctx context.Context) (Stream, error) +func (*SSEClientTransport) Connect(ctx context.Context) (Connection, error) ``` The Streamable HTTP transports are similar to the SSE transport, albeit with a @@ -213,12 +213,12 @@ func (*StreamableHTTPHandler) Close() error type StreamableServerTransport struct { /* ... */ } func NewStreamableServerTransport(sessionID string) *StreamableServerTransport func (*StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request) -func (*StreamableServerTransport) Connect(context.Context) (Stream, error) +func (*StreamableServerTransport) Connect(context.Context) (Connection, 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) +func (*StreamableClientTransport) Connect(context.Context) (Connection, error) ``` **Differences from mcp-go**: In mcp-go, server authors create an `MCPServer`, populate it with tools, resources and so on, and then wrap it in an `SSEServer` or `StdioServer`. Users can manage their own sessions with `RegisterSession` and `UnregisterSession`. Rather than use a server constructor to get a distinct server for each connection, there is a concept of a "session tool" that overlays tools for a specific session. diff --git a/mcp/cmd.go b/mcp/cmd.go index 42af63b..163bb0c 100644 --- a/mcp/cmd.go +++ b/mcp/cmd.go @@ -29,7 +29,7 @@ func NewCommandTransport(cmd *exec.Cmd) *CommandTransport { } // Connect starts the command, and connects to it over stdin/stdout. -func (t *CommandTransport) Connect(ctx context.Context) (Stream, error) { +func (t *CommandTransport) Connect(ctx context.Context) (Connection, error) { stdout, err := t.cmd.StdoutPipe() if err != nil { return nil, err @@ -42,7 +42,7 @@ func (t *CommandTransport) Connect(ctx context.Context) (Stream, error) { if err := t.cmd.Start(); err != nil { return nil, err } - return newIOStream(&pipeRWC{t.cmd, stdout, stdin}), nil + return newIOConn(&pipeRWC{t.cmd, stdout, stdin}), nil } // A pipeRWC is an io.ReadWriteCloser that communicates with a subprocess over diff --git a/mcp/mcp.go b/mcp/mcp.go index e8ccd4d..91f1487 100644 --- a/mcp/mcp.go +++ b/mcp/mcp.go @@ -24,7 +24,7 @@ // instances using the [Server.Connect] method, which creates a // [ServerSession]. // -// A [Transport] connects a bidirectional [Stream] of jsonrpc2 messages. In +// A [Transport] connects a bidirectional [Connection] of jsonrpc2 messages. In // practice, transports in the MCP spec are are either client transports or // server transports. For example, the [StdIOTransport] is a server transport // that communicates over stdin/stdout, and its counterpart is a diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index fc2be19..a02ec22 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -597,8 +597,8 @@ func TestBatching(t *testing.T) { c := NewClient("testClient", "v1.0.0", nil) // TODO: this test is broken, because increasing the batch size here causes // 'initialize' to block. Therefore, we can only test with a size of 1. + // Since batching is being removed, we can probably just delete this. const batchSize = 1 - BatchSize(ct, batchSize) cs, err := c.Connect(ctx, ct) if err != nil { t.Fatal(err) diff --git a/mcp/sse.go b/mcp/sse.go index 578605f..187b14e 100644 --- a/mcp/sse.go +++ b/mcp/sse.go @@ -104,7 +104,7 @@ func NewSSEHandler(getServer func(request *http.Request) *Server) *SSEHandler { // A SSEServerTransport is a logical SSE session created through a hanging GET // request. // -// When connected, it returns the following [Stream] implementation: +// When connected, it returns the following [Connection] implementation: // - Writes are SSE 'message' events to the GET response. // - Reads are received from POSTs to the session endpoint, via // [SSEServerTransport.ServeHTTP]. @@ -120,7 +120,7 @@ type SSEServerTransport struct { mu sync.Mutex w http.ResponseWriter // the hanging response body closed bool // set when the stream is closed - done chan struct{} // closed when the stream is closed + done chan struct{} // closed when the connection is closed } // NewSSEServerTransport creates a new SSE transport for the given messages @@ -168,8 +168,8 @@ func (t *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 (t *SSEServerTransport) Connect(context.Context) (Stream, error) { +// See [SSEServerTransport] for more details on the [Connection] implementation. +func (t *SSEServerTransport) Connect(context.Context) (Connection, error) { t.mu.Lock() _, err := writeEvent(t.w, event{ name: "endpoint", @@ -179,7 +179,7 @@ func (t *SSEServerTransport) Connect(context.Context) (Stream, error) { if err != nil { return nil, err } - return sseServerStream{t}, nil + return sseServerConn{t}, nil } func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { @@ -257,14 +257,14 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) { } } -// sseServerStream implements the Stream interface for a single [SSEServerTransport]. -// It hides the Stream interface from the SSEServerTransport API. -type sseServerStream struct { +// sseServerConn implements the [Connection] interface for a single [SSEServerTransport]. +// It hides the Connection interface from the SSEServerTransport API. +type sseServerConn struct { t *SSEServerTransport } // Read implements jsonrpc2.Reader. -func (s sseServerStream) Read(ctx context.Context) (jsonrpc2.Message, error) { +func (s sseServerConn) Read(ctx context.Context) (jsonrpc2.Message, error) { select { case <-ctx.Done(): return nil, ctx.Err() @@ -276,7 +276,7 @@ func (s sseServerStream) Read(ctx context.Context) (jsonrpc2.Message, error) { } // Write implements jsonrpc2.Writer. -func (s sseServerStream) Write(ctx context.Context, msg jsonrpc2.Message) error { +func (s sseServerConn) Write(ctx context.Context, msg jsonrpc2.Message) error { if ctx.Err() != nil { return ctx.Err() } @@ -305,7 +305,7 @@ func (s sseServerStream) Write(ctx context.Context, msg jsonrpc2.Message) error // 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 sseServerStream) Close() error { +func (s sseServerConn) Close() error { s.t.mu.Lock() defer s.t.mu.Unlock() if !s.t.closed { @@ -339,7 +339,7 @@ func NewSSEClientTransport(baseURL string) *SSEClientTransport { } // Connect connects through the client endpoint. -func (c *SSEClientTransport) Connect(ctx context.Context) (Stream, error) { +func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) { req, err := http.NewRequestWithContext(ctx, "GET", c.sseEndpoint.String(), nil) if err != nil { return nil, err @@ -370,7 +370,7 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Stream, error) { } // From here on, the stream takes ownership of resp.Body. - s := &sseClientStream{ + s := &sseClientConn{ sseEndpoint: c.sseEndpoint, msgEndpoint: msgEndpoint, incoming: make(chan []byte, 100), @@ -485,12 +485,12 @@ func scanEvents(r io.Reader) iter.Seq2[event, error] { } } -// An sseClientStream is a logical jsonrpc2 stream that implements the client +// An sseClientConn is a logical jsonrpc2 connection that implements the client // half of the SSE protocol: // - Writes are POSTS to the session endpoint. // - Reads are SSE 'message' events, and pushes them onto a buffered channel. // - Close terminates the GET request. -type sseClientStream struct { +type sseClientConn struct { sseEndpoint *url.URL // SSE endpoint for the GET msgEndpoint *url.URL // session endpoint for POSTs incoming chan []byte // queue of incoming messages @@ -501,13 +501,13 @@ type sseClientStream struct { done chan struct{} // closed when the stream is closed } -func (c *sseClientStream) isDone() bool { +func (c *sseClientConn) isDone() bool { c.mu.Lock() defer c.mu.Unlock() return c.closed } -func (c *sseClientStream) Read(ctx context.Context) (jsonrpc2.Message, error) { +func (c *sseClientConn) Read(ctx context.Context) (jsonrpc2.Message, error) { select { case <-ctx.Done(): return nil, ctx.Err() @@ -528,7 +528,7 @@ func (c *sseClientStream) Read(ctx context.Context) (jsonrpc2.Message, error) { } } -func (c *sseClientStream) Write(ctx context.Context, msg jsonrpc2.Message) error { +func (c *sseClientConn) Write(ctx context.Context, msg jsonrpc2.Message) error { data, err := jsonrpc2.EncodeMessage(msg) if err != nil { return err @@ -552,7 +552,7 @@ func (c *sseClientStream) Write(ctx context.Context, msg jsonrpc2.Message) error return nil } -func (c *sseClientStream) Close() error { +func (c *sseClientConn) Close() error { c.mu.Lock() defer c.mu.Unlock() if !c.closed { diff --git a/mcp/streamable.go b/mcp/streamable.go index 4e6d386..c70d6d7 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -150,20 +150,20 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque // policy. func NewStreamableServerTransport(sessionID string) *StreamableServerTransport { return &StreamableServerTransport{ - id: sessionID, - incoming: make(chan jsonrpc2.Message, 10), - done: make(chan struct{}), - outgoingMessages: make(map[connID][]*streamableMsg), - signals: make(map[connID]chan struct{}), - requestConns: make(map[JSONRPCID]connID), - outstandingRequests: make(map[connID]map[JSONRPCID]struct{}), + id: sessionID, + incoming: make(chan jsonrpc2.Message, 10), + done: make(chan struct{}), + outgoingMessages: make(map[streamID][]*streamableMsg), + signals: make(map[streamID]chan struct{}), + requestStreams: make(map[JSONRPCID]streamID), + streamRequests: make(map[streamID]map[JSONRPCID]struct{}), } } // A StreamableServerTransport implements the [Transport] interface for a // single session. type StreamableServerTransport struct { - nextConnID atomic.Int64 // incrementing next connection ID + nextStreamID atomic.Int64 // incrementing next stream ID id string incoming chan jsonrpc2.Message // messages from the client to the server @@ -189,48 +189,48 @@ type StreamableServerTransport struct { // TODO(rfindley): simplify. // outgoingMessages is the collection of outgoingMessages messages, keyed by the logical - // connection ID where they should be delivered. + // stream ID where they should be delivered. // - // connID 0 is used for messages that don't correlate with an incoming + // streamID 0 is used for messages that don't correlate with an incoming // request. // // Lifecycle: outgoingMessages persists for the duration of the session. // // TODO(rfindley): garbage collect this data. For now, we save all outgoingMessages // messages for the lifespan of the transport. - outgoingMessages map[connID][]*streamableMsg + outgoingMessages map[streamID][]*streamableMsg - // signals maps a logical connection ID to a 1-buffered channel, owned by an + // signals maps a logical stream ID to a 1-buffered channel, owned by an // incoming HTTP request, that signals that there are messages available to // write into the HTTP response. Signals guarantees that at most one HTTP - // connection can receive messages for a logical connection. After claiming - // the connection, incoming requests should read from outgoing, to ensure + // response can receive messages for a logical stream. After claiming + // the stream, incoming requests should read from outgoing, to ensure // that no new messages are missed. // // Lifecycle: signals persists for the duration of an HTTP POST or GET - // request for the given connID. - signals map[connID]chan struct{} + // request for the given streamID. + signals map[streamID]chan struct{} - // requestConns maps incoming requests to their logical connection ID. + // requestStreams maps incoming requests to their logical stream ID. // - // Lifecycle: requestConns persists for the duration of the session. + // Lifecycle: requestStreams persists for the duration of the session. // // TODO(rfindley): clean up once requests are handled. - requestConns map[JSONRPCID]connID + requestStreams map[JSONRPCID]streamID // outstandingRequests tracks the set of unanswered incoming RPCs for each logical - // connection. + // stream. // - // When the server has responded to each request, the connection should be + // When the server has responded to each request, the stream should be // closed. // // Lifecycle: outstandingRequests values persist as until the requests have been // replied to by the server. Notably, NOT until they are sent to an HTTP // response, as delivery is not guaranteed. - outstandingRequests map[connID]map[JSONRPCID]struct{} + streamRequests map[streamID]map[JSONRPCID]struct{} } -type connID int64 +type streamID int64 // a streamableMsg is an SSE event with an index into its logical stream. type streamableMsg struct { @@ -241,7 +241,7 @@ type streamableMsg struct { // Connect implements the [Transport] interface. // // TODO(rfindley): Connect should return a new object. -func (s *StreamableServerTransport) Connect(context.Context) (Stream, error) { +func (s *StreamableServerTransport) Connect(context.Context) (Connection, error) { return s, nil } @@ -284,7 +284,7 @@ func (t *StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.R func (t *StreamableServerTransport) serveGET(w http.ResponseWriter, req *http.Request) { // connID 0 corresponds to the default GET request. - id, nextIdx := connID(0), 0 + id, nextIdx := streamID(0), 0 if len(req.Header.Values("Last-Event-ID")) > 0 { eid := req.Header.Get("Last-Event-ID") var ok bool @@ -338,15 +338,15 @@ func (t *StreamableServerTransport) servePOST(w http.ResponseWriter, req *http.R } // Update accounting for this request. - id := connID(t.nextConnID.Add(1)) + id := streamID(t.nextStreamID.Add(1)) signal := make(chan struct{}, 1) t.mu.Lock() if len(requests) > 0 { - t.outstandingRequests[id] = make(map[JSONRPCID]struct{}) + t.streamRequests[id] = make(map[JSONRPCID]struct{}) } for reqID := range requests { - t.requestConns[reqID] = id - t.outstandingRequests[id][reqID] = struct{}{} + t.requestStreams[reqID] = id + t.streamRequests[id][reqID] = struct{}{} } t.signals[id] = signal t.mu.Unlock() @@ -362,7 +362,7 @@ func (t *StreamableServerTransport) servePOST(w http.ResponseWriter, req *http.R t.streamResponse(w, req, id, 0, signal) } -func (t *StreamableServerTransport) streamResponse(w http.ResponseWriter, req *http.Request, id connID, nextIndex int, signal chan struct{}) { +func (t *StreamableServerTransport) streamResponse(w http.ResponseWriter, req *http.Request, id streamID, nextIndex int, signal chan struct{}) { defer func() { t.mu.Lock() delete(t.signals, id) @@ -404,7 +404,7 @@ stream: } t.mu.Lock() - nOutstanding := len(t.outstandingRequests[id]) + nOutstanding := len(t.streamRequests[id]) nOutgoing := len(t.outgoingMessages[id]) t.mu.Unlock() // If all requests have been handled and replied to, we can terminate this @@ -443,21 +443,21 @@ stream: } // Event IDs: encode both the logical connection ID and the index, as -// _, to be consistent with the typescript implementation. +// _, to be consistent with the typescript implementation. // formatEventID returns the event ID to use for the logical connection ID -// connID and message index idx. +// streamID and message index idx. // // See also [parseEventID]. -func formatEventID(id connID, idx int) string { +func formatEventID(id streamID, idx int) string { return fmt.Sprintf("%d_%d", id, idx) } -// parseEventID parses a Last-Event-ID value into a logical connection id and +// parseEventID parses a Last-Event-ID value into a logical stream id and // index. // // See also [formatEventID]. -func parseEventID(eventID string) (conn connID, idx int, ok bool) { +func parseEventID(eventID string) (conn streamID, idx int, ok bool) { _, err := fmt.Sscanf(eventID, "%d_%d", &conn, &idx) if err != nil || conn < 0 || idx < 0 { return 0, 0, false @@ -465,7 +465,7 @@ func parseEventID(eventID string) (conn connID, idx int, ok bool) { return conn, idx, true } -// Read implements the [Stream] interface. +// Read implements the [Connection] interface. func (t *StreamableServerTransport) Read(ctx context.Context) (jsonrpc2.Message, error) { select { case <-ctx.Done(): @@ -480,7 +480,7 @@ func (t *StreamableServerTransport) Read(ctx context.Context) (jsonrpc2.Message, } } -// Write implements the [Stream] interface. +// Write implements the [Connection] interface. func (t *StreamableServerTransport) Write(ctx context.Context, msg JSONRPCMessage) error { // Find the incoming request that this write relates to, if any. var forRequest, replyTo JSONRPCID @@ -501,10 +501,10 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg JSONRPCMessag // // For messages sent outside of a request context, this is the default // connection 0. - var forConn connID + var forConn streamID if forRequest.IsValid() { t.mu.Lock() - forConn = t.requestConns[forRequest] + forConn = t.requestStreams[forRequest] t.mu.Unlock() } @@ -519,7 +519,7 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg JSONRPCMessag return fmt.Errorf("session is closed") // TODO: should this be EOF? } - if _, ok := t.outstandingRequests[forConn]; !ok && forConn != 0 { + if _, ok := t.streamRequests[forConn]; !ok && forConn != 0 { // No outstanding requests for this connection, which means it is logically // done. This is a sequencing violation from the server, so we should report // a side-channel error here. Put the message on the general queue to avoid @@ -538,9 +538,9 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg JSONRPCMessag }) if replyTo.IsValid() { // Once we've put the reply on the queue, it's no longer outstanding. - delete(t.outstandingRequests[forConn], replyTo) - if len(t.outstandingRequests[forConn]) == 0 { - delete(t.outstandingRequests, forConn) + delete(t.streamRequests[forConn], replyTo) + if len(t.streamRequests[forConn]) == 0 { + delete(t.streamRequests, forConn) } } @@ -554,7 +554,7 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg JSONRPCMessag return nil } -// Close implements the [Stream] interface. +// Close implements the [Connection] interface. func (t *StreamableServerTransport) Close() error { t.mu.Lock() defer t.mu.Unlock() @@ -582,14 +582,14 @@ func NewStreamableClientTransport(url string) *StreamableClientTransport { // Connect implements the [Transport] interface. // -// The resulting Stream writes messages via POST requests to the transport URL -// with the Mcp-Session-Id header set, and reads messages from hanging -// requests. +// The resulting [Connection] writes messages via POST requests to the +// transport URL with the Mcp-Session-Id header set, and reads messages from +// hanging requests. // -// When closed, the stream issues a DELETE request to terminate the logical +// When closed, the connection issues a DELETE request to terminate the logical // session. -func (t *StreamableClientTransport) Connect(ctx context.Context) (Stream, error) { - return &streamableClientStream{ +func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, error) { + return &streamableClientConn{ url: t.url, client: http.DefaultClient, incoming: make(chan []byte, 100), @@ -597,7 +597,7 @@ func (t *StreamableClientTransport) Connect(ctx context.Context) (Stream, error) }, nil } -type streamableClientStream struct { +type streamableClientConn struct { url string sessionID string client *http.Client @@ -612,8 +612,8 @@ type streamableClientStream struct { err error } -// Read implements the [Stream] interface. -func (s *streamableClientStream) Read(ctx context.Context) (jsonrpc2.Message, error) { +// Read implements the [Connection] interface. +func (s *streamableClientConn) Read(ctx context.Context) (jsonrpc2.Message, error) { select { case <-ctx.Done(): return nil, ctx.Err() @@ -624,8 +624,8 @@ func (s *streamableClientStream) Read(ctx context.Context) (jsonrpc2.Message, er } } -// Write implements the [Stream] interface. -func (s *streamableClientStream) Write(ctx context.Context, msg jsonrpc2.Message) error { +// Write implements the [Connection] interface. +func (s *streamableClientConn) Write(ctx context.Context, msg jsonrpc2.Message) error { s.mu.Lock() if s.err != nil { s.mu.Unlock() @@ -661,7 +661,7 @@ func (s *streamableClientStream) Write(ctx context.Context, msg jsonrpc2.Message return nil } -func (s *streamableClientStream) postMessage(ctx context.Context, sessionID string, msg JSONRPCMessage) (string, error) { +func (s *streamableClientConn) postMessage(ctx context.Context, sessionID string, msg JSONRPCMessage) (string, error) { data, err := jsonrpc2.EncodeMessage(msg) if err != nil { return "", err @@ -697,7 +697,7 @@ func (s *streamableClientStream) postMessage(ctx context.Context, sessionID stri return sessionID, nil } -func (s *streamableClientStream) handleSSE(resp *http.Response) { +func (s *streamableClientConn) handleSSE(resp *http.Response) { defer resp.Body.Close() done := make(chan struct{}) @@ -718,8 +718,8 @@ func (s *streamableClientStream) handleSSE(resp *http.Response) { } } -// Close implements the [Stream] interface. -func (s *streamableClientStream) Close() error { +// Close implements the [Connection] interface. +func (s *streamableClientConn) Close() error { s.closeOnce.Do(func() { close(s.done) diff --git a/mcp/transport.go b/mcp/transport.go index 3e31c59..008fc3d 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -28,10 +28,10 @@ var ErrConnectionClosed = errors.New("connection closed") // Transports should be used for at most one call to [Server.Connect] or // [Client.Start]. type Transport interface { - // Connect returns the logical stream. + // Connect returns the logical JSON-RPC connection.. // // It is called exactly once by [Server.Connect] or [Client.Connect]. - Connect(ctx context.Context) (Stream, error) + Connect(ctx context.Context) (Connection, error) } type ( @@ -47,8 +47,8 @@ type ( JSONRPCResponse = jsonrpc2.Response ) -// A Stream is a bidirectional jsonrpc2 Stream. -type Stream interface { +// A Connection is a logical bidirectional JSON-RPC connection. +type Connection interface { Read(context.Context) (JSONRPCMessage, error) Write(context.Context, JSONRPCMessage) error Close() error // may be called concurrently by both peers @@ -66,8 +66,8 @@ type ioTransport struct { rwc io.ReadWriteCloser } -func (t *ioTransport) Connect(context.Context) (Stream, error) { - return newIOStream(t.rwc), nil +func (t *ioTransport) Connect(context.Context) (Connection, error) { + return newIOConn(t.rwc), nil } // NewStdIOTransport constructs a transport that communicates over @@ -100,12 +100,12 @@ type handler interface { func connect[H handler](ctx context.Context, t Transport, b binder[H]) (H, error) { var zero H - stream, err := t.Connect(ctx) + conn, err := t.Connect(ctx) if err != nil { return zero, err } // If logging is configured, write message logs. - reader, writer := jsonrpc2.Reader(stream), jsonrpc2.Writer(stream) + reader, writer := jsonrpc2.Reader(conn), jsonrpc2.Writer(conn) var ( h H preempter canceller @@ -118,7 +118,7 @@ func connect[H handler](ctx context.Context, t Transport, b binder[H]) (H, error _ = jsonrpc2.NewConnection(ctx, jsonrpc2.ConnectionConfig{ Reader: reader, Writer: writer, - Closer: stream, + Closer: conn, Bind: bind, Preempter: &preempter, OnDone: func() { @@ -187,23 +187,23 @@ func NewLoggingTransport(delegate Transport, w io.Writer) *LoggingTransport { return &LoggingTransport{delegate, w} } -// Connect connects the underlying transport, returning a [Stream] that writes +// Connect connects the underlying transport, returning a [Connection] that writes // logs to the configured destination. -func (t *LoggingTransport) Connect(ctx context.Context) (Stream, error) { +func (t *LoggingTransport) Connect(ctx context.Context) (Connection, error) { delegate, err := t.delegate.Connect(ctx) if err != nil { return nil, err } - return &loggingStream{delegate, t.w}, nil + return &loggingConn{delegate, t.w}, nil } -type loggingStream struct { - delegate Stream +type loggingConn struct { + delegate Connection w io.Writer } // loggingReader is a stream middleware that logs incoming messages. -func (s *loggingStream) Read(ctx context.Context) (jsonrpc2.Message, error) { +func (s *loggingConn) Read(ctx context.Context) (jsonrpc2.Message, error) { msg, err := s.delegate.Read(ctx) if err != nil { fmt.Fprintf(s.w, "read error: %v", err) @@ -218,7 +218,7 @@ func (s *loggingStream) Read(ctx context.Context) (jsonrpc2.Message, error) { } // loggingWriter is a stream middleware that logs outgoing messages. -func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) error { +func (s *loggingConn) Write(ctx context.Context, msg jsonrpc2.Message) error { err := s.delegate.Write(ctx, msg) if err != nil { fmt.Fprintf(s.w, "write error: %v", err) @@ -232,7 +232,7 @@ func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) error { return err } -func (s *loggingStream) Close() error { +func (s *loggingConn) Close() error { return s.delegate.Close() } @@ -255,14 +255,14 @@ func (r rwc) Close() error { return errors.Join(r.rc.Close(), r.wc.Close()) } -// An ioStream is a transport that delimits messages with newlines across +// An ioConn is a transport that delimits messages with newlines across // a bidirectional stream, and supports JSONRPC2 message batching. // // See https://github.com/ndjson/ndjson-spec for discussion of newline // delimited JSON. // // See [msgBatch] for more discussion of message batching. -type ioStream struct { +type ioConn struct { rwc io.ReadWriteCloser // the underlying stream in *json.Decoder // a decoder bound to rwc @@ -280,23 +280,18 @@ type ioStream struct { batches map[jsonrpc2.ID]*msgBatch // lazily allocated } -func newIOStream(rwc io.ReadWriteCloser) *ioStream { - return &ioStream{ +func newIOConn(rwc io.ReadWriteCloser) *ioConn { + return &ioConn{ rwc: rwc, in: json.NewDecoder(rwc), } } -// Connect returns the receiver, as a streamTransport is a logical stream. -func (t *ioStream) Connect(ctx context.Context) (Stream, error) { - return t, nil -} - // addBatch records a msgBatch for an incoming batch payload. // It returns an error if batch is malformed, containing previously seen IDs. // // See [msgBatch] for more. -func (t *ioStream) addBatch(batch *msgBatch) error { +func (t *ioConn) addBatch(batch *msgBatch) error { t.batchMu.Lock() defer t.batchMu.Unlock() for id := range batch.unresolved { @@ -319,7 +314,7 @@ func (t *ioStream) addBatch(batch *msgBatch) error { // The second result reports whether resp was part of a batch. If this is true, // the first result is nil if the batch is still incomplete, or the full set of // batch responses if resp completed the batch. -func (t *ioStream) updateBatch(resp *jsonrpc2.Response) ([]*jsonrpc2.Response, bool) { +func (t *ioConn) updateBatch(resp *jsonrpc2.Response) ([]*jsonrpc2.Response, bool) { t.batchMu.Lock() defer t.batchMu.Unlock() @@ -357,11 +352,11 @@ type msgBatch struct { responses []*jsonrpc2.Response } -func (t *ioStream) Read(ctx context.Context) (jsonrpc2.Message, error) { +func (t *ioConn) Read(ctx context.Context) (jsonrpc2.Message, error) { return t.read(ctx, t.in) } -func (t *ioStream) read(ctx context.Context, in *json.Decoder) (jsonrpc2.Message, error) { +func (t *ioConn) read(ctx context.Context, in *json.Decoder) (jsonrpc2.Message, error) { select { case <-ctx.Done(): return nil, ctx.Err() @@ -432,7 +427,7 @@ func readBatch(data []byte) (msgs []jsonrpc2.Message, isBatch bool, _ error) { return []jsonrpc2.Message{msg}, false, err } -func (t *ioStream) Write(ctx context.Context, msg jsonrpc2.Message) error { +func (t *ioConn) Write(ctx context.Context, msg jsonrpc2.Message) error { select { case <-ctx.Done(): return ctx.Err() @@ -479,7 +474,7 @@ func (t *ioStream) Write(ctx context.Context, msg jsonrpc2.Message) error { return err } -func (t *ioStream) Close() error { +func (t *ioConn) Close() error { return t.rwc.Close() } diff --git a/mcp/transport_test.go b/mcp/transport_test.go index 51ec584..83c6d7a 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -12,16 +12,6 @@ import ( "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" ) -// BatchSize causes a transport to collect n requests or notifications before -// sending a batch on the wire (responses are always sent in isolation). -// -// Exported for testing in the mcp_test package. -func BatchSize(t Transport, n int) { - if st, ok := t.(*ioStream); ok { - st.outgoingBatch = make([]jsonrpc2.Message, 0, n) - } -} - func TestBatchFraming(t *testing.T) { // This test checks that the ndjsonFramer can read and write JSON batches. // @@ -31,7 +21,7 @@ func TestBatchFraming(t *testing.T) { ctx := context.Background() r, w := io.Pipe() - tport := newIOStream(rwc{r, w}) + tport := newIOConn(rwc{r, w}) tport.outgoingBatch = make([]jsonrpc2.Message, 0, 2) // Read the two messages into a channel, for easy testing later.