Route MCP sessions to the originating backend pod using httptrace

When a proxy runner pod restarts it recovers sessions from Redis but
backend_url stored the ClusterIP, so kube-proxy could send follow-up
requests to a different backend pod that never handled initialize —
causing JSON-RPC -32001 "session not found" errors on the first request.

Use net/http/httptrace.GotConn to capture the actual backend pod IP
after kube-proxy DNAT on every initialize request, and store that as
backend_url instead of the ClusterIP URL. The existing Rewrite closure
already reads backend_url and pins routing to the correct pod; no
changes to that path are needed.

When the backend pod is later replaced (rescheduled to a new IP or
restarted in place and lost in-memory session state), the proxy now
re-initializes the backend session transparently rather than returning
404 to the client:

- Dial error (pod IP unreachable): re-init triggers on TCP failure
- Backend 404 (session lost, same IP): re-init triggers on response

In both cases the proxy replays the stored initialize body against the
ClusterIP, captures the new pod IP via GotConn, stores the new backend
session ID, rewrites outbound Mcp-Session-Id headers, and replays the
original client request — the client sees no error.

DELETE responses are excluded from the 404 re-init path since the
session is intentionally torn down in that case.

Closes #4575

Author: taskbot

pkg/transport/proxy/transparent/backend_routing_test.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -289,3 +290,230 @@ func TestRoundTripStoresBackendURLOnInitialize(t *testing.T) {
 	require.True(t, ok, "session should have backend_url metadata")
 	assert.Equal(t, backend.URL, backendURL)
 }
+
+// TestRoundTripStoresInitBodyOnInitialize verifies that the raw JSON-RPC initialize
+// request body is stored in session metadata so the proxy can transparently
+// re-initialize the backend session if the pod is later replaced.
+func TestRoundTripStoresInitBodyOnInitialize(t *testing.T) {
+	t.Parallel()
+
+	sessionID := uuid.New().String()
+	const initBody = `{"jsonrpc":"2.0","id":1,"method":"initialize"}`
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Mcp-Session-Id", sessionID)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer backend.Close()
+
+	proxy, addr := startProxy(t, backend.URL)
+
+	ctx := context.Background()
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
+		"http://"+addr+"/mcp",
+		strings.NewReader(initBody))
+	require.NoError(t, err)
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+	_ = resp.Body.Close()
+
+	sess, ok := proxy.sessionManager.Get(normalizeSessionID(sessionID))
+	require.True(t, ok, "session should have been created")
+	stored, exists := sess.GetMetadataValue(sessionMetadataInitBody)
+	require.True(t, exists, "init_body should be stored in session metadata")
+	assert.Equal(t, initBody, stored)
+}
+
+// TestRoundTripReinitializesOnBackend404 verifies that when the backend pod returns
+// 404 (session lost after restart on the same IP), the proxy transparently
+// re-initializes the backend session and replays the original request — client sees 200.
+func TestRoundTripReinitializesOnBackend404(t *testing.T) {
+	t.Parallel()
+
+	// staleBackend simulates a pod that has lost its in-memory session state.
+	var staleHit atomic.Int32
+	staleBackend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		staleHit.Add(1)
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	defer staleBackend.Close()
+
+	// freshBackend simulates a healthy pod: returns a session ID on initialize
+	// and 200 for all other requests.
+	freshSessionID := uuid.New().String()
+	var freshHit atomic.Int32
+	freshBackend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		freshHit.Add(1)
+		body, _ := io.ReadAll(r.Body)
+		if strings.Contains(string(body), `"initialize"`) {
+			w.Header().Set("Mcp-Session-Id", freshSessionID)
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer freshBackend.Close()
+
+	// targetURI (ClusterIP) points to freshBackend; the session has staleBackend as backend_url.
+	proxy, addr := startProxy(t, freshBackend.URL)
+
+	clientSessionID := uuid.New().String()
+	sess := session.NewProxySession(clientSessionID)
+	sess.SetMetadata(sessionMetadataBackendURL, staleBackend.URL)
+	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(`{"method":"tools/list"}`))
+	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.Equal(t, http.StatusOK, resp.StatusCode, "client should see 200 after transparent re-init")
+	assert.GreaterOrEqual(t, staleHit.Load(), int32(1), "stale backend should have been hit")
+	assert.GreaterOrEqual(t, freshHit.Load(), int32(2), "fresh backend should receive initialize + replay")
+
+	// Session should now have backend_sid mapping to the new backend session.
+	updated, ok := proxy.sessionManager.Get(normalizeSessionID(clientSessionID))
+	require.True(t, ok, "session should still exist after re-init")
+	backendSID, exists := updated.GetMetadataValue(sessionMetadataBackendSID)
+	require.True(t, exists, "backend_sid should be set after re-init")
+	assert.Equal(t, freshSessionID, backendSID, "backend_sid must be the raw value the backend issued, not normalized")
+}
+
+// TestRoundTripReinitializesPreservesNonUUIDBackendSessionID verifies that when the
+// backend issues a non-UUID Mcp-Session-Id on re-initialization, the proxy stores
+// and forwards the raw value — not a UUID v5 hash of it — on all subsequent requests.
+//
+// The normalization bug only manifests on the request AFTER the replay: the replay
+// sets Mcp-Session-Id directly from newBackendSID (bypassing Rewrite), but subsequent
+// requests go through the Rewrite closure which reads backend_sid from session metadata.
+// If backend_sid was stored as normalizeSessionID(newBackendSID), Rewrite would send
+// the wrong (hashed) value and the backend would reject every subsequent request.
+func TestRoundTripReinitializesPreservesNonUUIDBackendSessionID(t *testing.T) {
+	t.Parallel()
+
+	// Non-UUID opaque token, as some MCP servers issue.
+	const nonUUIDSessionID = "opaque-session-token-abc123"
+
+	staleBackend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	defer staleBackend.Close()
+
+	// receivedSIDs tracks Mcp-Session-Id values arriving on non-initialize requests,
+	// in order. Index 0 = replay (direct from reinitializeAndReplay), index 1 = second
+	// client request (routed through Rewrite reading backend_sid from session metadata).
+	var (
+		receivedMu   sync.Mutex
+		receivedSIDs []string
+	)
+	freshBackend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, _ := io.ReadAll(r.Body)
+		if strings.Contains(string(body), `"initialize"`) {
+			w.Header().Set("Mcp-Session-Id", nonUUIDSessionID)
+			w.WriteHeader(http.StatusOK)
+			return
+		}
+		receivedMu.Lock()
+		receivedSIDs = append(receivedSIDs, r.Header.Get("Mcp-Session-Id"))
+		receivedMu.Unlock()
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer freshBackend.Close()
+
+	proxy, addr := startProxy(t, freshBackend.URL)
+
+	clientSessionID := uuid.New().String()
+	sess := session.NewProxySession(clientSessionID)
+	sess.SetMetadata(sessionMetadataBackendURL, staleBackend.URL)
+	sess.SetMetadata(sessionMetadataInitBody, `{"jsonrpc":"2.0","id":1,"method":"initialize"}`)
+	require.NoError(t, proxy.sessionManager.AddSession(sess))
+
+	doRequest := func() *http.Response {
+		ctx := context.Background()
+		req, err := http.NewRequestWithContext(ctx, http.MethodPost,
+			"http://"+addr+"/mcp",
+			strings.NewReader(`{"method":"tools/list"}`))
+		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)
+		return resp
+	}
+
+	// First request: triggers re-init. The replay (inside reinitializeAndReplay) sets
+	// Mcp-Session-Id directly, so receivedSIDs[0] is always the raw value regardless
+	// of what is stored in session metadata.
+	resp1 := doRequest()
+	_ = resp1.Body.Close()
+	require.Equal(t, http.StatusOK, resp1.StatusCode)
+
+	// Second request: goes through the Rewrite closure, which reads backend_sid from
+	// session metadata. This is where the normalization bug manifests — if backend_sid
+	// was stored as normalizeSessionID(nonUUIDSessionID), Rewrite would forward the
+	// wrong hashed value and receivedSIDs[1] would not equal nonUUIDSessionID.
+	resp2 := doRequest()
+	_ = resp2.Body.Close()
+	require.Equal(t, http.StatusOK, resp2.StatusCode)
+
+	receivedMu.Lock()
+	defer receivedMu.Unlock()
+	require.Len(t, receivedSIDs, 2, "fresh backend should have received replay + second request")
+	assert.Equal(t, nonUUIDSessionID, receivedSIDs[0], "replay must forward raw non-UUID session ID")
+	assert.Equal(t, nonUUIDSessionID, receivedSIDs[1], "subsequent request via Rewrite must forward raw non-UUID session ID")
+}
+
+// TestRoundTripReinitializesOnDialError verifies that when the proxy cannot reach
+// the stored pod IP (dial error — pod rescheduled to a new IP), it transparently
+// re-initializes the backend session via the ClusterIP and replays the original
+// request — the client sees a 200.
+func TestRoundTripReinitializesOnDialError(t *testing.T) {
+	t.Parallel()
+
+	// Create a server and immediately close it so its URL refuses connections.
+	dead := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {}))
+	deadURL := dead.URL
+	dead.Close()
+
+	freshSessionID := uuid.New().String()
+	var freshHit atomic.Int32
+	freshBackend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		freshHit.Add(1)
+		body, _ := io.ReadAll(r.Body)
+		if strings.Contains(string(body), `"initialize"`) {
+			w.Header().Set("Mcp-Session-Id", freshSessionID)
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer freshBackend.Close()
+
+	proxy, addr := startProxy(t, freshBackend.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(`{"method":"tools/list"}`))
+	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.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")
+}