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.