From c3cc138a54177afb14dabe78d9e4d9ab9b5f1832 Mon Sep 17 00:00:00 2001 From: amirejaz Date: Thu, 30 Jul 2026 21:57:59 +0100 Subject: [PATCH] Strip session ID from proxied initialize requests The transparent proxy forwarded a client-initiated initialize with an Mcp-Session-Id attached: either the client's own, or the backend session ID mapped from it by an earlier transparent re-initialization. MCP requires a re-initializing client to start a new session by sending a new InitializeRequest without a session ID, and a backend session that has already completed its handshake rejects a second one -- go-sdk v1.7+ answers `duplicate "initialize" received`. A client retrying its handshake on a live connection therefore hit a hard failure instead of getting a fresh session. Strip the header on initialize so the backend mints a new session, and exclude initialize from dial-error recovery: reinitializeAndReplay sends its own initialize and then replays the original request, which for an initialize handed the newly created backend session a second handshake. The 404 recovery path already made that exclusion. Unpin the session from the unreachable pod in its place, so the client's retry routes via the target service rather than looping against a dead address. Co-Authored-By: Claude Opus 5 (1M context) --- .../proxy/transparent/backend_routing_test.go | 122 ++++++++++++++++++ .../proxy/transparent/transparent_proxy.go | 72 +++++++++-- 2 files changed, 183 insertions(+), 11 deletions(-) diff --git a/pkg/transport/proxy/transparent/backend_routing_test.go b/pkg/transport/proxy/transparent/backend_routing_test.go index 55bd76d503..0b22f26496 100644 --- a/pkg/transport/proxy/transparent/backend_routing_test.go +++ b/pkg/transport/proxy/transparent/backend_routing_test.go @@ -712,3 +712,125 @@ func TestRoundTripReinitializesOnDialError(t *testing.T) { assert.Equal(t, http.StatusOK, resp.StatusCode, "client should see 200 after transparent re-init on dial error") assert.GreaterOrEqual(t, freshHit.Load(), int32(2), "fresh backend should receive initialize + replay") } + +// TestRoundTripStripsSessionIDFromInitialize verifies that a client-initiated +// initialize reaches the backend with no Mcp-Session-Id header, whether or not +// the session it carries maps to a backend session ID recorded by a prior +// transparent re-initialization. +// +// MCP requires a re-initializing client to "start a new session by sending a +// new InitializeRequest without a session ID attached". Forwarding the client's +// session ID — or the backend SID mapped from it — hands an already-initialized +// backend session a second handshake, which go-sdk v1.7+ rejects with +// `duplicate "initialize" received`, turning a client-side handshake retry into +// a hard failure. +func TestRoundTripStripsSessionIDFromInitialize(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + priorBackend bool // session carries a backend_sid from an earlier re-init + }{ + {name: "session mapped to a backend SID", priorBackend: true}, + {name: "session with no backend SID", priorBackend: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var receivedSIDs []string + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + receivedSIDs = append(receivedSIDs, r.Header.Get("Mcp-Session-Id")) + mu.Unlock() + w.Header().Set("Mcp-Session-Id", uuid.New().String()) + w.WriteHeader(http.StatusOK) + })) + defer backend.Close() + + proxy, addr := startProxy(t, backend.URL) + + clientSessionID := uuid.New().String() + sess := session.NewProxySession(clientSessionID) + if tt.priorBackend { + sess.SetMetadata(sessionMetadataBackendSID, uuid.New().String()) + } + require.NoError(t, proxy.sessionManager.AddSession(sess)) + + ctx := context.Background() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + "http://"+addr+"/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize"}`)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Mcp-Session-Id", clientSessionID) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + _ = resp.Body.Close() + + mu.Lock() + defer mu.Unlock() + require.Len(t, receivedSIDs, 1, "backend should receive exactly one request") + assert.Empty(t, receivedSIDs[0], "initialize must reach the backend with no Mcp-Session-Id") + }) + } +} + +// TestRoundTripDoesNotReplayInitializeOnDialError verifies that a dial error on +// an initialize does not trigger transparent re-initialization. +// +// reinitializeAndReplay sends its own initialize and then replays the original +// request; for an initialize that would give the freshly created backend session +// a second handshake — the exact failure the session-ID strip exists to avoid. +// The client is already starting a new session, so the proxy instead unpins the +// session from the unreachable pod and lets the error surface, leaving the +// client's retry to route via the target service. +func TestRoundTripDoesNotReplayInitializeOnDialError(t *testing.T) { + t.Parallel() + + // A server closed immediately after creation: its URL refuses connections. + dead := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})) + deadURL := dead.URL + dead.Close() + + var targetHits atomic.Int32 + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + targetHits.Add(1) + w.Header().Set("Mcp-Session-Id", uuid.New().String()) + w.WriteHeader(http.StatusOK) + })) + defer target.Close() + + proxy, addr := startProxy(t, target.URL) + + clientSessionID := uuid.New().String() + sess := session.NewProxySession(clientSessionID) + sess.SetMetadata(sessionMetadataBackendURL, deadURL) + sess.SetMetadata(sessionMetadataInitBody, `{"jsonrpc":"2.0","id":1,"method":"initialize"}`) + require.NoError(t, proxy.sessionManager.AddSession(sess)) + + ctx := context.Background() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + "http://"+addr+"/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize"}`)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Mcp-Session-Id", clientSessionID) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + _ = resp.Body.Close() + + assert.Zero(t, targetHits.Load(), + "the proxy must not send its own initialize to the target service for an initialize dial error") + + updated, ok := proxy.sessionManager.Get(normalizeSessionID(clientSessionID)) + require.True(t, ok, "session should survive the dial error") + backendURL, exists := updated.GetMetadataValue(sessionMetadataBackendURL) + require.True(t, exists, "backend_url should still be set") + assert.Equal(t, target.URL, backendURL, + "session must be unpinned from the unreachable pod so the client's retry routes via the target service") +} diff --git a/pkg/transport/proxy/transparent/transparent_proxy.go b/pkg/transport/proxy/transparent/transparent_proxy.go index e1c43b5d0c..8ca5e3bba6 100644 --- a/pkg/transport/proxy/transparent/transparent_proxy.go +++ b/pkg/transport/proxy/transparent/transparent_proxy.go @@ -642,16 +642,31 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) } } - // Capture the client-facing session ID before the backend SID rewrite below. - // Recovery and session cleanup paths must look up sessions by the client SID - // (the store key), not the backend SID that is written into the header. + // Capture the client-facing session ID before the header is rewritten or + // stripped below. Recovery and session cleanup paths must look up sessions by + // the client SID (the store key), not the backend SID that is written into the + // header. clientSID := req.Header.Get("Mcp-Session-Id") - // Rewrite the outbound Mcp-Session-Id to the backend's assigned session ID when - // the proxy transparently re-initialized the backend session. This is done here - // (after the guard check above) so the guard always sees the original client - // session ID and can look it up correctly in the session store. - if clientSID != "" { + switch { + case sawInitialize: + // An initialize request must never carry a session ID. MCP requires a + // client that is re-initializing to "start a new session by sending a new + // InitializeRequest without a session ID attached", and a backend session + // that has already completed its handshake rejects a second initialize on + // it -- go-sdk v1.7+ answers `duplicate "initialize" received`. Forwarding + // the client's stale SID -- or, worse, the backend SID mapped from it -- + // turns a client-side handshake retry into a hard failure. Strip it so the + // backend mints a fresh session instead; the response's Mcp-Session-Id is + // then adopted as a new session below. reinitializeAndReplay already does + // the same on the recovery path. + req.Header.Del("Mcp-Session-Id") + case clientSID != "": + // Rewrite the outbound Mcp-Session-Id to the backend's assigned session ID + // when the proxy transparently re-initialized the backend session. This is + // done here (after the guard check above) so the guard always sees the + // original client session ID and can look it up correctly in the session + // store. if sess, ok := t.p.sessionManager.Get(normalizeSessionID(clientSID)); ok { if backendSID, exists := sess.GetMetadataValue(sessionMetadataBackendSID); exists && backendSID != "" { req.Header.Set("Mcp-Session-Id", backendSID) @@ -681,10 +696,22 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) } // Dial error against a stored pod IP means the pod has been replaced. // Attempt transparent re-initialization so the client sees no error. + // + // An initialize request is excluded, as it already is on the 404 path + // below: reinitializeAndReplay sends its own initialize and then replays + // the original request, which for an initialize would hand the freshly + // created backend session a second handshake -- the very failure this + // function strips the session ID to avoid. The client is already starting + // a new session, so instead unpin it from the dead pod and let the error + // surface; the client's retry is then routed via the target service. if isDialError(err) { - req.Header.Set("Mcp-Session-Id", clientSID) - if reInitResp, reInitErr := t.recovery.reinitializeAndReplay(req, reqBody); reInitResp != nil || reInitErr != nil { - return reInitResp, reInitErr + if sawInitialize { + t.recovery.unpinSession(clientSID) + } else { + req.Header.Set("Mcp-Session-Id", clientSID) + if reInitResp, reInitErr := t.recovery.reinitializeAndReplay(req, reqBody); reInitResp != nil || reInitErr != nil { + return reInitResp, reInitErr + } } } slog.Error("failed to forward request", "error", err) @@ -992,6 +1019,29 @@ func isRedirectStatus(code int) bool { } } +// unpinSession clears the backend_url pin on the session identified by +// clientSID so the next request carrying it is routed to the target service +// rather than to a pod address that is no longer reachable. It is a no-op for +// an empty or unknown session ID. +// +// This is the same fallback reinitializeAndReplay applies when it has no stored +// initialize body to replay, reused for the initialize case where issuing the +// new handshake is the client's job, not the proxy's. +func (r *backendRecovery) unpinSession(clientSID string) { + if clientSID == "" { + return + } + sess, ok := r.sessions.Get(normalizeSessionID(clientSID)) + if !ok { + return + } + sess.SetMetadata(sessionMetadataBackendURL, r.targetURI) + if err := r.sessions.UpsertSession(sess); err != nil { + slog.Debug("failed to unpin session after dial error on initialize", + "session_id", clientSID, "error", err) + } +} + // reinitializeAndReplay is called when the proxy detects that the backend pod // that owned a session is no longer reachable (dial error) or has lost its // in-memory session state (backend returned 404). It transparently: