Strip session ID from proxied initialize requests

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")
+}

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: