stacklok / toolhive / 24212213832

66%

Build Type

push

github

Committed by web-flow

Commit Message

Route MCP sessions to the originating backend pod using httptrace (#4673)

* 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

* Reimplement tracingTransport logic + add e2e tests

Extract re-initialization logic from tracingTransport into a dedicated
backendRecovery type backed by a narrow recoverySessionStore interface
(Get + UpsertSession) and a forward func. tracingTransport now owns only
the request lifecycle (session guard, initialize detection, httptrace,
session creation) and delegates all forwarding and recovery to
backendRecovery.

This makes reinitializeAndReplay and podBackendURL testable without
standin... (continued)

Coverage Stats

152 of 164 new or added lines in 1 file covered. (92.68%)

22 existing lines in 5 files now uncovered.

56838 of 87202 relevant lines covered (65.18%)

62.13 hits per line