From 7bfb777c34bba40a045b7a7af940759ec1fd64b9 Mon Sep 17 00:00:00 2001 From: Edson Date: Thu, 1 Oct 2026 01:58:31 -0400 Subject: [PATCH 1/3] mcp: forward session state through LoggingTransport loggingConn did not implement clientConnection. Client.Connect looks the interface up on the connection and skips the call when the assertion fails, so wrapping a client transport in a LoggingTransport silently dropped the session state update: streamableClientConn never learned the initialize result, so for sessions below 2026-07-28 it sent no Mcp-Protocol-Version header after initialize and never opened the standalone SSE GET stream. #1267 added the server-side forwarder, but both interfaces named their method sessionUpdated with different parameter types, so one wrapper could not implement both. Rename them to clientSessionUpdated and serverSessionUpdated, and have loggingConn forward both, and propagateCancellation, to the delegate when it implements them. TestStreamableClientLoggingTransport fails without the change. Fixes #1328 --- mcp/client.go | 4 +- mcp/server.go | 2 +- mcp/streamable.go | 2 +- mcp/streamable_client.go | 2 +- mcp/streamable_client_test.go | 79 +++++++++++++++++++++++- mcp/transport.go | 44 ++++++++++---- mcp/transport_test.go | 110 +++++++++++++++++++++++++++++++++- 7 files changed, 223 insertions(+), 20 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 3b24b2e9d..dcb2f0c37 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 819bc08f4..f9d436c55 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -2365,7 +2365,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 928b65edc..a7ba942c3 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 @@ -2250,3 +2250,80 @@ 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) { + getKey := streamableRequestKey{"GET", "123", "", ""} + 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, + }, + getKey: { + 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) + } + }) + } +} diff --git a/mcp/transport.go b/mcp/transport.go index 2b04ae9c6..022b5da51 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 @@ -440,17 +440,37 @@ 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) +) + +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 @@ -623,7 +643,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 90f7cb628..11846c205 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -141,7 +141,7 @@ func TestIOConnRead(t *testing.T) { if requested == "" { requested = tt.protocolVersion } - tr.sessionUpdated(ServerSessionState{ + tr.serverSessionUpdated(ServerSessionState{ InitializeParams: &InitializeParams{ ProtocolVersion: requested, }, @@ -389,7 +389,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)) @@ -418,6 +418,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. @@ -555,3 +622,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") + } +} From ba5ade31615920a543a2641c69f917451c3aa3b6 Mon Sep 17 00:00:00 2001 From: Edson Date: Thu, 22 Oct 2026 14:22:52 -0400 Subject: [PATCH 2/3] mcp: address review nits in LoggingTransport forwarding Inline the GET request key in the logging transport test, and assert at compile time that loggingConn implements jsonrpc2.ResponseDropper. --- mcp/streamable_client_test.go | 3 +-- mcp/transport.go | 7 ++++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index a7ba942c3..28fceb616 100644 --- a/mcp/streamable_client_test.go +++ b/mcp/streamable_client_test.go @@ -2268,7 +2268,6 @@ func TestStreamableClientLoggingTransport(t *testing.T) { name = "logging" } t.Run(name, func(t *testing.T) { - getKey := streamableRequestKey{"GET", "123", "", ""} fake := &fakeStreamableServer{ t: t, responses: fakeResponses{ @@ -2288,7 +2287,7 @@ func TestStreamableClientLoggingTransport(t *testing.T) { body: jsonBody(t, resp(2, &emptyResult{}, nil)), wantProtocolVersion: protocolVersion20251125, }, - getKey: { + {"GET", "123", "", ""}: { header: header{"Content-Type": "text/event-stream"}, wantProtocolVersion: protocolVersion20251125, }, diff --git a/mcp/transport.go b/mcp/transport.go index 022b5da51..3a925b127 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -449,9 +449,10 @@ func (s *loggingConn) Close() error { // 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) + _ clientConnection = (*loggingConn)(nil) + _ serverConnection = (*loggingConn)(nil) + _ cancellationPropagator = (*loggingConn)(nil) + _ jsonrpc2.ResponseDropper = (*loggingConn)(nil) ) func (s *loggingConn) clientSessionUpdated(state clientSessionState) { From 7c0bdfc81aa2ddb32c738dc72487a768440b6d7f Mon Sep 17 00:00:00 2001 From: Edson Date: Mon, 26 Oct 2026 21:17:17 -0400 Subject: [PATCH 3/3] mcp: restore closing braces lost in the merge of main The merge in 7b02956 dropped the end of TestStreamableClientLoggingTransport, so mcp/streamable_client_test.go no longer parsed. --- mcp/streamable_client_test.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index 86b2ae373..b65e8a20b 100644 --- a/mcp/streamable_client_test.go +++ b/mcp/streamable_client_test.go @@ -2324,6 +2324,8 @@ func TestStreamableClientLoggingTransport(t *testing.T) { 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