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

@@ -289,3 +289,146 @@ 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, normalizeSessionID(freshSessionID), backendSID)
+}
+
+// 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")
+}

pkg/transport/proxy/transparent/transparent_proxy.go

@@ -15,6 +15,7 @@ import (
 	"log/slog"
 	"net"
 	"net/http"
+	"net/http/httptrace"
 	"net/http/httputil"
 	"net/url"
 	"os"
@@ -153,6 +154,17 @@ const (
 	// It is written on initialize and read in the Rewrite closure to route follow-up requests
 	// to the same backend pod that handled the session's initialize request.
 	sessionMetadataBackendURL = "backend_url"
+
+	// sessionMetadataInitBody stores the raw JSON-RPC initialize request body.
+	// It is used to transparently re-initialize a backend session when the pod that
+	// originally handled initialize has been replaced (new IP or lost in-memory state).
+	sessionMetadataInitBody = "init_body"
+
+	// sessionMetadataBackendSID stores the backend's assigned Mcp-Session-Id when it
+	// diverges from the client-facing session ID after a transparent re-initialization.
+	// The Rewrite closure rewrites the outbound Mcp-Session-Id header to this value so
+	// the backend sees its own session ID while the client keeps its original one.
+	sessionMetadataBackendSID = "backend_sid"
 )
 
 // Option is a functional option for configuring TransparentProxy
@@ -436,12 +448,33 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error)
 		}
 	}
 
+	// Attach an httptrace to capture the actual backend pod IP after kube-proxy
+	// DNAT resolves the ClusterIP to a specific pod. The captured address is stored
+	// as backend_url so follow-up requests always reach the same pod, even after a
+	// proxy runner restart that would otherwise lose the in-memory routing state.
+	var capturedPodAddr string
+	if sawInitialize {
+		trace := &httptrace.ClientTrace{
+			GotConn: func(info httptrace.GotConnInfo) {
+				capturedPodAddr = info.Conn.RemoteAddr().String()
+			},
+		}
+		req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace))
+	}
+
 	resp, err := t.forward(req)
 	if err != nil {
 		if errors.Is(err, context.Canceled) {
 			// Expected during shutdown or client disconnect—silently ignore
 			return nil, err
 		}
+		// Dial error against a stored pod IP means the pod has been replaced.
+		// Attempt transparent re-initialization so the client sees no error.
+		if isDialError(err) {
+			if reInitResp, reInitErr := t.reinitializeAndReplay(req, reqBody); reInitResp != nil || reInitErr != nil {
+				return reInitResp, reInitErr
+			}
+		}
 		slog.Error("failed to forward request", "error", err)
 		return nil, err
 	}
@@ -471,6 +504,20 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error)
 		}
 	}
 
+	// Backend returned 404 for a non-initialize, non-DELETE request whose session IS
+	// known to the proxy. This means the backend pod lost its in-memory session state
+	// (e.g. it was restarted but got the same IP). Attempt transparent re-initialization
+	// so the client sees no error. DELETE is excluded because the session has already
+	// been cleaned up above and the 404 is the expected terminal response.
+	if resp.StatusCode == http.StatusNotFound && !sawInitialize && req.Method != http.MethodDelete {
+		if sid := req.Header.Get("Mcp-Session-Id"); sid != "" {
+			if reInitResp, reInitErr := t.reinitializeAndReplay(req, reqBody); reInitResp != nil || reInitErr != nil {
+				_ = resp.Body.Close()
+				return reInitResp, reInitErr
+			}
+		}
+	}
+
 	if resp.StatusCode == http.StatusOK {
 		// check if we saw a valid mcp header
 		ct := resp.Header.Get("Mcp-Session-Id")
@@ -480,14 +527,15 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error)
 			internalID := normalizeSessionID(ct)
 			if _, ok := t.p.sessionManager.Get(internalID); !ok {
 				sess := session.NewProxySession(internalID)
-				// Store targetURI as the default backend_url for this session.
-				// In single-replica deployments targetURI is already the pod address,
-				// so no override is needed. In multi-replica deployments the
-				// vMCP/operator layer is responsible for setting backend_url to the
-				// actual pod DNS name (e.g. http://mcp-server-0.mcp-server.default.svc:8080)
-				// before the request reaches this proxy; the Rewrite closure then reads
-				// that value and routes follow-up requests to the correct pod.
-				sess.SetMetadata(sessionMetadataBackendURL, t.p.targetURI)
+				// Store the actual pod IP (captured via GotConn) as backend_url so that
+				// after a proxy runner restart the session is routed to the same backend
+				// pod that handled initialize, not a random pod via ClusterIP.
+				sess.SetMetadata(sessionMetadataBackendURL, t.podBackendURL(capturedPodAddr))
+				// Store the initialize body so we can transparently re-initialize the
+				// backend session if the pod is later replaced or loses session state.
+				if len(reqBody) > 0 {
+					sess.SetMetadata(sessionMetadataInitBody, string(reqBody))
+				}
 				if err := t.p.sessionManager.AddSession(sess); err != nil {
 					//nolint:gosec // G706: session ID from HTTP response header
 					slog.Error("failed to create session from header",
@@ -553,6 +601,133 @@ func (t *tracingTransport) detectInitialize(body []byte) bool {
 	return false
 }
 
+// podBackendURL constructs a backend URL that targets the specific pod IP captured
+// via httptrace.GotConn, using the scheme from targetURI. Falls back to targetURI
+// when no address was captured (e.g. single-replica, connection reuse without a new conn).
+func (t *tracingTransport) podBackendURL(capturedAddr string) string {
+	if capturedAddr == "" {
+		return t.p.targetURI
+	}
+	parsed, err := url.Parse(t.p.targetURI)
+	if err != nil {
+		return t.p.targetURI
+	}
+	parsed.Host = capturedAddr
+	return parsed.String()
+}
+
+// isDialError reports whether err is a TCP dial failure, indicating that the
+// target host is unreachable (pod has been terminated or rescheduled).
+func isDialError(err error) bool {
+	var opErr *net.OpError
+	return errors.As(err, &opErr) && opErr.Op == "dial"
+}
+
+// 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:
+//  1. Re-sends the stored initialize body to the ClusterIP service so kube-proxy
+//     selects a healthy pod and the backend creates a new session.
+//  2. Captures the new pod IP via httptrace.GotConn and stores it as backend_url.
+//  3. Maps the client's original session ID to the new backend session ID.
+//  4. Replays the original client request so the client sees no error.
+//
+// Returns (nil, nil) when re-initialization is not applicable (no stored init
+// body, session unknown, or already routing via ClusterIP).
+func (t *tracingTransport) reinitializeAndReplay(req *http.Request, origBody []byte) (*http.Response, error) {
+	sid := req.Header.Get("Mcp-Session-Id")
+	if sid == "" {
+		return nil, nil
+	}
+	internalSID := normalizeSessionID(sid)
+	sess, ok := t.p.sessionManager.Get(internalSID)
+	if !ok {
+		return nil, nil
+	}
+
+	initBody, hasInit := sess.GetMetadataValue(sessionMetadataInitBody)
+	if !hasInit || initBody == "" {
+		// No stored init body — cannot re-initialize transparently.
+		// Reset backend_url to ClusterIP so the next request goes through
+		// kube-proxy and lets the client receive a clean 404 to re-initialize.
+		sess.SetMetadata(sessionMetadataBackendURL, t.p.targetURI)
+		_ = t.p.sessionManager.UpsertSession(sess)
+		return nil, nil
+	}
+
+	slog.Debug("backend session lost; transparently re-initializing",
+		"session_id", sid, "target", t.p.targetURI)
+
+	// Capture the new pod IP via GotConn on the re-initialize connection.
+	var capturedPodAddr string
+	trace := &httptrace.ClientTrace{
+		GotConn: func(info httptrace.GotConnInfo) {
+			capturedPodAddr = info.Conn.RemoteAddr().String()
+		},
+	}
+	initCtx := httptrace.WithClientTrace(req.Context(), trace)
+
+	// Build a fresh initialize request to the ClusterIP (no Mcp-Session-Id —
+	// the backend assigns a new session ID in the response).
+	parsedTarget, err := url.Parse(t.p.targetURI)
+	if err != nil {
+		return nil, nil
+	}
+	initURL := *req.URL
+	initURL.Scheme = parsedTarget.Scheme
+	initURL.Host = parsedTarget.Host
+
+	initReq, err := http.NewRequestWithContext(initCtx, http.MethodPost, initURL.String(), bytes.NewReader([]byte(initBody)))
+	if err != nil {
+		return nil, nil
+	}
+	initReq.Header.Set("Content-Type", "application/json")
+
+	initResp, err := t.forward(initReq)
+	if err != nil {
+		slog.Error("transparent re-initialize failed", "error", err)
+		return nil, err
+	}
+	_, _ = io.Copy(io.Discard, initResp.Body)
+	_ = initResp.Body.Close()
+
+	newBackendSID := initResp.Header.Get("Mcp-Session-Id")
+	if newBackendSID == "" {
+		slog.Debug("re-initialize response contained no Mcp-Session-Id; falling back to ClusterIP")
+		sess.SetMetadata(sessionMetadataBackendURL, t.p.targetURI)
+		_ = t.p.sessionManager.UpsertSession(sess)
+		return nil, nil
+	}
+
+	// Update session: point backend_url at the newly-discovered pod and record
+	// the backend session ID so Rewrite rewrites Mcp-Session-Id on outbound requests.
+	newPodURL := t.podBackendURL(capturedPodAddr)
+	sess.SetMetadata(sessionMetadataBackendURL, newPodURL)
+	sess.SetMetadata(sessionMetadataBackendSID, normalizeSessionID(newBackendSID))
+	if upsertErr := t.p.sessionManager.UpsertSession(sess); upsertErr != nil {
+		slog.Debug("failed to update session after re-initialize", "error", upsertErr)
+	}
+
+	// Replay the original client request to the new pod with the new backend SID.
+	// Use the captured pod address directly so we bypass the Rewrite closure
+	// (which still holds the old backend_url until the next session load).
+	replayHost := capturedPodAddr
+	if replayHost == "" {
+		replayHost = parsedTarget.Host
+	}
+	replayReq := req.Clone(req.Context())
+	replayReq.URL.Scheme = parsedTarget.Scheme
+	replayReq.URL.Host = replayHost
+	replayReq.Host = replayHost // keep Host header consistent with URL to avoid backend validation errors
+	replayReq.Header.Set("Mcp-Session-Id", newBackendSID)
+	replayReq.Body = io.NopCloser(bytes.NewReader(origBody))
+	replayReq.ContentLength = int64(len(origBody))
+
+	slog.Debug("replaying original request after transparent re-initialization",
+		"new_pod_url", newPodURL, "new_backend_sid", newBackendSID)
+	return t.forward(replayReq)
+}
+
 // modifyResponse modifies HTTP responses based on transport-specific requirements.
 // Delegates to the appropriate ResponseProcessor based on transport type.
 func (p *TransparentProxy) modifyResponse(resp *http.Response) error {
@@ -601,6 +776,13 @@ func (p *TransparentProxy) Start(ctx context.Context) error {
 							pr.Out.URL.Host = parsed.Host
 						}
 					}
+					// After a transparent re-initialization the proxy maps the client's
+					// session ID to the backend's newly-assigned session ID. Rewrite the
+					// outbound header so the backend sees its own ID while the client
+					// continues to use its original session ID unchanged.
+					if backendSID, exists := sess.GetMetadataValue(sessionMetadataBackendSID); exists && backendSID != "" {
+						pr.Out.Header.Set("Mcp-Session-Id", backendSID)
+					}
 				}
 			}