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: