diff --git a/mcp/client.go b/mcp/client.go index 891fab5c3..3d01c643c 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -330,7 +330,7 @@ func (c *Client) Connect(ctx context.Context, t Transport, opts *ClientSessionOp if err == nil { cs.state.InitializeResult = discRes if hc, ok := cs.mcpConn.(clientConnection); ok { - hc.sessionUpdated(cs.state) + hc.clientSessionUpdated(cs.state) } subscribeParams := &SubscriptionsListenParams{ Notifications: &NotificationSubscriptions{}, @@ -396,7 +396,7 @@ func (c *Client) Connect(ctx context.Context, t Transport, opts *ClientSessionOp } cs.state.InitializeResult = res if hc, ok := cs.mcpConn.(clientConnection); ok { - hc.sessionUpdated(cs.state) + hc.clientSessionUpdated(cs.state) } req2 := &initializedClientRequest{Session: cs, Params: &InitializedParams{}} if err := handleNotify(ctx, notificationInitialized, req2); err != nil { diff --git a/mcp/server.go b/mcp/server.go index 899c67b3d..47d71cf83 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -1643,7 +1643,7 @@ func (ss *ServerSession) updateState(mut func(*ServerSessionState)) { copy := ss.state ss.mu.Unlock() if c, ok := ss.mcpConn.(serverConnection); ok { - c.sessionUpdated(copy) + c.serverSessionUpdated(copy) } } diff --git a/mcp/streamable.go b/mcp/streamable.go index 01ab31052..d11526e7a 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -2373,7 +2373,7 @@ type streamableClientConn struct { var _ clientConnection = (*streamableClientConn)(nil) -func (c *streamableClientConn) sessionUpdated(state clientSessionState) { +func (c *streamableClientConn) clientSessionUpdated(state clientSessionState) { c.mu.Lock() c.initializedResult = state.InitializeResult c.mu.Unlock() diff --git a/mcp/streamable_client.go b/mcp/streamable_client.go index 74940b3ad..e2f9b2640 100644 --- a/mcp/streamable_client.go +++ b/mcp/streamable_client.go @@ -107,7 +107,7 @@ The client handles both response formats: # Standalone SSE Stream -After initialization, [streamableClientConn.sessionUpdated] triggers +After initialization, [streamableClientConn.clientSessionUpdated] triggers [streamableClientConn.connectStandaloneSSE] to open a GET request for server-initiated messages: diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index b93a8f5db..b65e8a20b 100644 --- a/mcp/streamable_client_test.go +++ b/mcp/streamable_client_test.go @@ -1504,7 +1504,7 @@ func TestStreamableClientConnect_DiscoverSuccess(t *testing.T) { // streamableClientConn.setMCPHeaders' selection of the Mcp-Protocol-Version // header value. // -// Ordinarily initializedResult is populated by sessionUpdated, called +// Ordinarily initializedResult is populated by clientSessionUpdated, called // through a type assertion to the unexported clientConnection interface // (see Client.Connect). That assertion silently fails, leaving // initializedResult nil for the life of the session, whenever the @@ -2251,6 +2251,82 @@ func TestStreamableClientConnect_ColonlessSSELine(t *testing.T) { } } +// TestStreamableClientLoggingTransport verifies that wrapping a +// StreamableClientTransport in a LoggingTransport (documented as a decorator +// for any Transport) does not change the session's behavior: +// once initialize completes, the client must still send the negotiated +// Mcp-Protocol-Version header on every subsequent request (2025-06-18 ยง +// Protocol Version Header: "the client MUST include the MCP-Protocol-Version +// header on all subsequent requests") and still open the standalone SSE GET +// stream, exactly as it does when the transport is not wrapped. +func TestStreamableClientLoggingTransport(t *testing.T) { + ctx := context.Background() + + for _, wrap := range []bool{false, true} { + name := "plain" + if wrap { + name = "logging" + } + t.Run(name, func(t *testing.T) { + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + body: jsonBody(t, initResp), + }, + {"POST", "123", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + {"POST", "123", methodPing, ""}: { + header: header{"Content-Type": "application/json"}, + body: jsonBody(t, resp(2, &emptyResult{}, nil)), + wantProtocolVersion: protocolVersion20251125, + }, + {"GET", "123", "", ""}: { + header: header{"Content-Type": "text/event-stream"}, + wantProtocolVersion: protocolVersion20251125, + }, + {"DELETE", "123", "", ""}: { + optional: true, + }, + }, + } + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + var transport Transport = &StreamableClientTransport{Endpoint: httpServer.URL} + if wrap { + transport = &LoggingTransport{Transport: transport, Writer: io.Discard} + } + client := NewClient(testImpl, nil) + // Pin to 2025-11-25: the canned initialize response has id=1, and + // the standalone GET stream only exists before 2026-07-28. + session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatalf("client.Connect() failed: %v", err) + } + if err := session.Ping(ctx, nil); err != nil { + t.Fatalf("Ping() failed: %v", err) + } + + // Give the standalone SSE GET a moment to be issued. + time.Sleep(100 * time.Millisecond) + + if err := session.Close(); err != nil { + t.Errorf("closing session: %v", err) + } + if missing := fake.missingRequests(); len(missing) > 0 { + t.Errorf("did not receive expected requests: %v", missing) + } + }) + } +} + // TestStreamableClientRedirects verifies that the client follows a redirect // within the origin of the endpoint, and that a redirect to another origin is // not followed and fails only the redirected request. See #1429. diff --git a/mcp/transport.go b/mcp/transport.go index 0329626a7..4247bd807 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -102,8 +102,8 @@ type Connection interface { type clientConnection interface { Connection - // sessionUpdated is called whenever the client session state changes. - sessionUpdated(clientSessionState) + // clientSessionUpdated is called whenever the client session state changes. + clientSessionUpdated(clientSessionState) } // A serverConnection is a Connection that is specific to the MCP server. @@ -115,8 +115,8 @@ type clientConnection interface { type serverConnection interface { Connection - // sessionUpdated is called whenever the server session state changes. - sessionUpdated(ServerSessionState) + // serverSessionUpdated is called whenever the server session state changes. + serverSessionUpdated(ServerSessionState) } // DefaultMaxLineLength is the default maximum number of bytes buffered while @@ -466,17 +466,38 @@ func (s *loggingConn) DropResponse(id jsonrpc.ID) { } } -// sessionUpdated implements [serverConnection] by forwarding to the delegate, -// whose behavior may depend on the negotiated protocol version (see -// [ioConn.DropResponse]). -func (s *loggingConn) sessionUpdated(state ServerSessionState) { +func (s *loggingConn) Close() error { + return s.delegate.Close() +} + +// The methods below forward the optional connection interfaces to the +// delegate. Client and server sessions look these up with a type assertion on +// the connection, so without them wrapping a transport in a LoggingTransport +// would silently change its behavior. +var ( + _ clientConnection = (*loggingConn)(nil) + _ serverConnection = (*loggingConn)(nil) + _ cancellationPropagator = (*loggingConn)(nil) + _ jsonrpc2.ResponseDropper = (*loggingConn)(nil) +) + +func (s *loggingConn) clientSessionUpdated(state clientSessionState) { + if c, ok := s.delegate.(clientConnection); ok { + c.clientSessionUpdated(state) + } +} + +// serverSessionUpdated forwards to the delegate, whose behavior may depend on +// the negotiated protocol version (see [ioConn.DropResponse]). +func (s *loggingConn) serverSessionUpdated(state ServerSessionState) { if c, ok := s.delegate.(serverConnection); ok { - c.sessionUpdated(state) + c.serverSessionUpdated(state) } } -func (s *loggingConn) Close() error { - return s.delegate.Close() +func (s *loggingConn) propagateCancellation() bool { + cp, ok := s.delegate.(cancellationPropagator) + return ok && cp.propagateCancellation() } // A rwc binds an io.ReadCloser and io.WriteCloser together to create an @@ -649,7 +670,7 @@ func (r *frameLimitReader) resetFrame() { r.count = 0 } func (c *ioConn) SessionID() string { return "" } -func (c *ioConn) sessionUpdated(state ServerSessionState) { +func (c *ioConn) serverSessionUpdated(state ServerSessionState) { protocolVersion := state.NegotiatedProtocolVersion if protocolVersion == "" { // 2025-03-26 is used, because it's the last spec version diff --git a/mcp/transport_test.go b/mcp/transport_test.go index cd844de3c..c2fb985de 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -325,7 +325,7 @@ func TestIOConnRead(t *testing.T) { if requested == "" { requested = tt.protocolVersion } - tr.sessionUpdated(ServerSessionState{ + tr.serverSessionUpdated(ServerSessionState{ InitializeParams: &InitializeParams{ ProtocolVersion: requested, }, @@ -573,7 +573,7 @@ func TestIOConnDropResponse(t *testing.T) { wc: nopCloserWriter{&out}, }) t.Cleanup(func() { tr.Close() }) - tr.sessionUpdated(ServerSessionState{NegotiatedProtocolVersion: test.version}) + tr.serverSessionUpdated(ServerSessionState{NegotiatedProtocolVersion: test.version}) if test.drop { tr.DropResponse(jsonrpc2.Int64ID(1)) @@ -602,6 +602,73 @@ func TestIOConnDropResponse(t *testing.T) { } } +// TestLoggingTransportServerBatchRejection verifies that wrapping a stdio-style +// transport in a LoggingTransport (as examples/server/{everything,memory, +// sequentialthinking} do) does not change what the server accepts: once a +// session has negotiated 2025-06-18 or later, a JSON-RPC batch must be rejected +// (see TestIOConnRead "batching new protocol") whether or not the connection +// is wrapped. +func TestLoggingTransportServerBatchRejection(t *testing.T) { + for _, wrap := range []bool{false, true} { + name := "plain" + if wrap { + name = "logging" + } + t.Run(name, func(t *testing.T) { + ctx := context.Background() + inR, inW := io.Pipe() // test -> server stdin + outR, outW := io.Pipe() // server stdout -> test + var tr Transport = &IOTransport{Reader: inR, Writer: outW} + if wrap { + tr = &LoggingTransport{Transport: tr, Writer: io.Discard} + } + s := NewServer(testImpl, nil) + ss, err := s.Connect(ctx, tr, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + lines := bufio.NewScanner(outR) + send := func(msg string) { + t.Helper() + if _, err := io.WriteString(inW, msg+"\n"); err != nil { + t.Fatalf("write %q: %v", msg, err) + } + } + recv := func() (string, bool) { + ch := make(chan string, 1) + go func() { + if lines.Scan() { + ch <- lines.Text() + } + close(ch) + }() + select { + case line, ok := <-ch: + return line, ok + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for the server to respond") + return "", false + } + } + + send(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"c","version":"1"}}}`) + if line, ok := recv(); !ok || !strings.Contains(line, `"protocolVersion":"2025-11-25"`) { + t.Fatalf("initialize: got %q, want a 2025-11-25 InitializeResult", line) + } + send(`{"jsonrpc":"2.0","method":"notifications/initialized"}`) + + // A batch is not supported at 2025-06-18 or later: the server must + // not answer it. + send(`[{"jsonrpc":"2.0","id":2,"method":"ping"},{"jsonrpc":"2.0","id":3,"method":"ping"}]`) + if line, ok := recv(); ok { + t.Errorf("server answered a JSON-RPC batch on a 2025-11-25 session: %q", line) + } + }) + } +} + // TestIOConnCancelledCallGetsNoResponse checks end to end, over an in-memory // transport wrapped for logging, that on a 2026-07-28 session a call the // client cancels with notifications/cancelled gets no response. @@ -739,3 +806,42 @@ func TestLoggingConnDropResponse(t *testing.T) { conn.DropResponse(jsonrpc.ID{}) }) } + +// optionalConn is a Connection implementing every optional connection +// interface, recording the calls it receives. +type optionalConn struct { + Connection + clientState *clientSessionState + serverState *ServerSessionState +} + +func (c *optionalConn) clientSessionUpdated(s clientSessionState) { c.clientState = &s } +func (c *optionalConn) serverSessionUpdated(s ServerSessionState) { c.serverState = &s } +func (c *optionalConn) propagateCancellation() bool { return true } + +// TestLoggingConnForwardsOptionalInterfaces verifies that loggingConn forwards +// the optional connection interfaces to its delegate, and is a no-op for a +// delegate that implements none of them. +func TestLoggingConnForwardsOptionalInterfaces(t *testing.T) { + delegate := &optionalConn{} + lc := &loggingConn{delegate: delegate, w: io.Discard} + + lc.clientSessionUpdated(clientSessionState{InitializeResult: &InitializeResult{ProtocolVersion: protocolVersion20251125}}) + if delegate.clientState == nil || delegate.clientState.InitializeResult.ProtocolVersion != protocolVersion20251125 { + t.Errorf("clientSessionUpdated was not forwarded: got %+v", delegate.clientState) + } + lc.serverSessionUpdated(ServerSessionState{NegotiatedProtocolVersion: protocolVersion20251125}) + if delegate.serverState == nil || delegate.serverState.NegotiatedProtocolVersion != protocolVersion20251125 { + t.Errorf("serverSessionUpdated was not forwarded: got %+v", delegate.serverState) + } + if !lc.propagateCancellation() { + t.Error("propagateCancellation() = false, want the delegate's true") + } + + plain := &loggingConn{delegate: struct{ Connection }{}, w: io.Discard} + plain.clientSessionUpdated(clientSessionState{}) + plain.serverSessionUpdated(ServerSessionState{}) + if plain.propagateCancellation() { + t.Error("propagateCancellation() = true for a delegate without the interface, want false") + } +}