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
standing up a full TransparentProxy: backend_recovery_test.go covers
all recovery paths (no session, no init body, happy path, forward
error, missing new session ID) using a stubSessionStore and inline
httptest servers.

tracingTransport.forward() and the base field are removed; all network
I/O goes through recovery.forward — a single source of truth for the
underlying transport.

Also integrates the E2E acceptance test from #4574 that exercises
backendReplicas=2 + proxy runner restart, verifying that sessions are
routed to the correct backend pod after re-initialization.

Author: taskbot

pkg/transport/proxy/transparent/backend_recovery_test.go

@@ -0,0 +1,277 @@
+// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package transparent
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/stacklok/toolhive/pkg/transport/session"
+)
+
+// stubSessionStore is a minimal in-memory recoverySessionStore for unit tests.
+type stubSessionStore struct {
+	sessions map[string]session.Session
+}
+
+func newStubStore(sessions ...session.Session) *stubSessionStore {
+	m := make(map[string]session.Session)
+	for _, s := range sessions {
+		m[s.ID()] = s
+	}
+	return &stubSessionStore{sessions: m}
+}
+
+func (s *stubSessionStore) Get(id string) (session.Session, bool) {
+	sess, ok := s.sessions[id]
+	return sess, ok
+}
+
+func (s *stubSessionStore) UpsertSession(sess session.Session) error {
+	s.sessions[sess.ID()] = sess
+	return nil
+}
+
+// newRecovery builds a backendRecovery backed by the given store and forward func.
+func newRecovery(targetURL string, store recoverySessionStore, fwd func(*http.Request) (*http.Response, error)) *backendRecovery {
+	return &backendRecovery{
+		targetURI: targetURL,
+		forward:   fwd,
+		sessions:  store,
+	}
+}
+
+// TestBackendRecoveryNoSession verifies that reinitializeAndReplay returns
+// (nil, nil) when the request carries no Mcp-Session-Id.
+func TestBackendRecoveryNoSession(t *testing.T) {
+	t.Parallel()
+
+	r := newRecovery("http://cluster-ip:8080", newStubStore(), nil)
+	req, err := http.NewRequest(http.MethodPost, "http://cluster-ip:8080/mcp",
+		strings.NewReader(`{"method":"tools/list"}`))
+	require.NoError(t, err)
+
+	resp, err := r.reinitializeAndReplay(req, nil)
+	assert.Nil(t, resp)
+	assert.NoError(t, err)
+}
+
+// TestBackendRecoveryUnknownSession verifies that reinitializeAndReplay returns
+// (nil, nil) when the session ID is not in the store.
+func TestBackendRecoveryUnknownSession(t *testing.T) {
+	t.Parallel()
+
+	r := newRecovery("http://cluster-ip:8080", newStubStore(), nil)
+	req, err := http.NewRequest(http.MethodPost, "http://cluster-ip:8080/mcp",
+		strings.NewReader(`{"method":"tools/list"}`))
+	require.NoError(t, err)
+	req.Header.Set("Mcp-Session-Id", uuid.New().String())
+
+	resp, err := r.reinitializeAndReplay(req, nil)
+	assert.Nil(t, resp)
+	assert.NoError(t, err)
+}
+
+// TestBackendRecoveryNoInitBody verifies that when the session has no stored
+// init body, reinitializeAndReplay resets backend_url to the ClusterIP and
+// returns (nil, nil) so the caller falls through to a 404 the client can handle.
+func TestBackendRecoveryNoInitBody(t *testing.T) {
+	t.Parallel()
+
+	const clusterIP = "http://cluster-ip:8080"
+	clientSID := uuid.New().String()
+	sess := session.NewProxySession(clientSID)
+	sess.SetMetadata(sessionMetadataBackendURL, "http://10.0.0.5:8080") // stale pod IP
+	store := newStubStore(sess)
+
+	r := newRecovery(clusterIP, store, nil)
+	req, err := http.NewRequest(http.MethodPost, clusterIP+"/mcp",
+		strings.NewReader(`{"method":"tools/list"}`))
+	require.NoError(t, err)
+	req.Header.Set("Mcp-Session-Id", clientSID)
+
+	resp, err := r.reinitializeAndReplay(req, nil)
+	assert.Nil(t, resp)
+	assert.NoError(t, err)
+
+	// backend_url should be reset to ClusterIP so the next request routes correctly.
+	updated, ok := store.Get(clientSID)
+	require.True(t, ok)
+	backendURL, exists := updated.GetMetadataValue(sessionMetadataBackendURL)
+	require.True(t, exists)
+	assert.Equal(t, clusterIP, backendURL, "backend_url should be reset to ClusterIP when no init body")
+}
+
+// TestBackendRecoveryHappyPath verifies the full re-init flow: the stored
+// initialize body is replayed to the ClusterIP, the new backend session ID is
+// captured, the session is updated, and the original request is replayed — all
+// without standing up a full TransparentProxy.
+func TestBackendRecoveryHappyPath(t *testing.T) {
+	t.Parallel()
+
+	const initBody = `{"jsonrpc":"2.0","id":1,"method":"initialize"}`
+	newBackendSID := uuid.New().String()
+	var (
+		forwardMu    sync.Mutex
+		forwardCalls []string
+	)
+
+	// Backend: returns a session ID on initialize, 200 otherwise.
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, _ := io.ReadAll(r.Body)
+		forwardMu.Lock()
+		forwardCalls = append(forwardCalls, r.Header.Get("Mcp-Session-Id"))
+		forwardMu.Unlock()
+		if strings.Contains(string(body), `"initialize"`) {
+			w.Header().Set("Mcp-Session-Id", newBackendSID)
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer backend.Close()
+
+	clientSID := uuid.New().String()
+	sess := session.NewProxySession(clientSID)
+	sess.SetMetadata(sessionMetadataInitBody, initBody)
+	store := newStubStore(sess)
+
+	r := newRecovery(backend.URL, store, http.DefaultTransport.RoundTrip)
+
+	origBody := []byte(`{"method":"tools/list"}`)
+	req, err := http.NewRequest(http.MethodPost, backend.URL+"/mcp",
+		bytes.NewReader(origBody))
+	require.NoError(t, err)
+	req.Header.Set("Mcp-Session-Id", clientSID)
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := r.reinitializeAndReplay(req, origBody)
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+	_ = resp.Body.Close()
+
+	// Verify session was updated with new backend SID and a pod URL.
+	updated, ok := store.Get(clientSID)
+	require.True(t, ok)
+	backendSID, exists := updated.GetMetadataValue(sessionMetadataBackendSID)
+	require.True(t, exists)
+	assert.Equal(t, newBackendSID, backendSID)
+
+	backendURL, exists := updated.GetMetadataValue(sessionMetadataBackendURL)
+	require.True(t, exists)
+	assert.NotEmpty(t, backendURL)
+
+	// Two forward calls: initialize + replay. The initialize must not carry
+	// a session ID; the replay must carry the new backend SID.
+	forwardMu.Lock()
+	defer forwardMu.Unlock()
+	require.Len(t, forwardCalls, 2, "forward should be called for initialize and replay")
+	assert.Empty(t, forwardCalls[0], "initialize request must not carry Mcp-Session-Id")
+	assert.Equal(t, newBackendSID, forwardCalls[1], "replay must carry the new backend SID")
+}
+
+// TestBackendRecoveryReinitForwardError verifies that a forward error during
+// re-initialization is returned to the caller.
+func TestBackendRecoveryReinitForwardError(t *testing.T) {
+	t.Parallel()
+
+	// Server that is immediately closed — all connections will be refused.
+	dead := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {}))
+	deadURL := dead.URL
+	dead.Close()
+
+	clientSID := uuid.New().String()
+	sess := session.NewProxySession(clientSID)
+	sess.SetMetadata(sessionMetadataInitBody, `{"jsonrpc":"2.0","id":1,"method":"initialize"}`)
+	store := newStubStore(sess)
+
+	r := newRecovery(deadURL, store, http.DefaultTransport.RoundTrip)
+
+	req, err := http.NewRequest(http.MethodPost, deadURL+"/mcp",
+		strings.NewReader(`{"method":"tools/list"}`))
+	require.NoError(t, err)
+	req.Header.Set("Mcp-Session-Id", clientSID)
+
+	resp, err := r.reinitializeAndReplay(req, []byte(`{"method":"tools/list"}`))
+	assert.Nil(t, resp)
+	assert.Error(t, err, "forward error during re-init should be returned")
+}
+
+// TestBackendRecoveryNoNewSessionID verifies that when the re-initialize
+// response carries no Mcp-Session-Id, reinitializeAndReplay resets backend_url
+// to ClusterIP and returns (nil, nil).
+func TestBackendRecoveryNoNewSessionID(t *testing.T) {
+	t.Parallel()
+
+	// Backend that returns no Mcp-Session-Id on initialize.
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK) // no Mcp-Session-Id header
+	}))
+	defer backend.Close()
+
+	clientSID := uuid.New().String()
+	sess := session.NewProxySession(clientSID)
+	sess.SetMetadata(sessionMetadataInitBody, `{"jsonrpc":"2.0","id":1,"method":"initialize"}`)
+	sess.SetMetadata(sessionMetadataBackendURL, "http://10.0.0.5:8080")
+	store := newStubStore(sess)
+
+	// targetURI points to backend (so the init request succeeds), but we verify
+	// that backend_url is reset to targetURI when no session ID comes back.
+	r := newRecovery(backend.URL, store, http.DefaultTransport.RoundTrip)
+
+	req, err := http.NewRequest(http.MethodPost, backend.URL+"/mcp",
+		strings.NewReader(`{"method":"tools/list"}`))
+	require.NoError(t, err)
+	req.Header.Set("Mcp-Session-Id", clientSID)
+
+	resp, err := r.reinitializeAndReplay(req, []byte(`{"method":"tools/list"}`))
+	assert.Nil(t, resp)
+	assert.NoError(t, err)
+
+	updated, ok := store.Get(clientSID)
+	require.True(t, ok)
+	backendURL, exists := updated.GetMetadataValue(sessionMetadataBackendURL)
+	require.True(t, exists)
+	assert.Equal(t, backend.URL, backendURL, "backend_url should fall back to targetURI when no new session ID")
+}
+
+// TestPodBackendURLWithCapturedAddr verifies that a captured pod IP replaces the
+// host in targetURI while preserving the scheme.
+func TestPodBackendURLWithCapturedAddr(t *testing.T) {
+	t.Parallel()
+
+	r := &backendRecovery{targetURI: "http://cluster-ip:8080"}
+	got := r.podBackendURL("10.0.0.5:8080")
+	assert.Equal(t, "http://10.0.0.5:8080", got)
+}
+
+// TestPodBackendURLFallback verifies that an empty captured address falls back
+// to targetURI unchanged.
+func TestPodBackendURLFallback(t *testing.T) {
+	t.Parallel()
+
+	r := &backendRecovery{targetURI: "http://cluster-ip:8080"}
+	got := r.podBackendURL("")
+	assert.Equal(t, "http://cluster-ip:8080", got)
+}
+
+// TestPodBackendURLHTTPSFallback verifies that an HTTPS targetURI is never
+// rewritten to a pod IP. IP-literal HTTPS URLs fail TLS verification because
+// server certificates are issued for hostnames, not pod IPs.
+func TestPodBackendURLHTTPSFallback(t *testing.T) {
+	t.Parallel()
+
+	r := &backendRecovery{targetURI: "https://mcp.example.com/mcp"}
+	got := r.podBackendURL("1.2.3.4:443")
+	assert.Equal(t, "https://mcp.example.com/mcp", got,
+		"HTTPS target must not be rewritten to a pod IP")
+}

pkg/transport/proxy/transparent/backend_routing_test.go

@@ -172,13 +172,13 @@ func TestRoundTripReturns404ForUnknownSession(t *testing.T) {
 	}))
 	defer backend.Close()
 
-	tt := &tracingTransport{base: http.DefaultTransport, p: NewTransparentProxyWithOptions(
+	tt := newTracingTransport(http.DefaultTransport, NewTransparentProxyWithOptions(
 		"localhost", 0, backend.URL,
 		nil, nil, nil,
 		false, false, "sse",
 		nil, nil, "", false,
 		nil,
-	)}
+	))
 
 	req, err := http.NewRequest(http.MethodPost, backend.URL+"/mcp",
 		strings.NewReader(`{"method":"tools/list"}`))
@@ -207,13 +207,13 @@ func TestRoundTripAllowsInitializeWithUnknownSession(t *testing.T) {
 	}))
 	defer backend.Close()
 
-	tt := &tracingTransport{base: http.DefaultTransport, p: NewTransparentProxyWithOptions(
+	tt := newTracingTransport(http.DefaultTransport, NewTransparentProxyWithOptions(
 		"localhost", 0, backend.URL,
 		nil, nil, nil,
 		false, false, "sse",
 		nil, nil, "", false,
 		nil,
-	)}
+	))
 
 	req, err := http.NewRequest(http.MethodPost, backend.URL+"/mcp",
 		strings.NewReader(`{"method":"initialize"}`))
@@ -239,13 +239,13 @@ func TestRoundTripAllowsBatchInitializeWithUnknownSession(t *testing.T) {
 	}))
 	defer backend.Close()
 
-	tt := &tracingTransport{base: http.DefaultTransport, p: NewTransparentProxyWithOptions(
+	tt := newTracingTransport(http.DefaultTransport, NewTransparentProxyWithOptions(
 		"localhost", 0, backend.URL,
 		nil, nil, nil,
 		false, false, "sse",
 		nil, nil, "", false,
 		nil,
-	)}
+	))
 
 	req, err := http.NewRequest(http.MethodPost, backend.URL+"/mcp",
 		strings.NewReader(`[{"method":"initialize"},{"method":"tools/list"}]`))
@@ -470,6 +470,63 @@ func TestRoundTripReinitializesPreservesNonUUIDBackendSessionID(t *testing.T) {
 	assert.Equal(t, nonUUIDSessionID, receivedSIDs[1], "subsequent request via Rewrite must forward raw non-UUID session ID")
 }
 
+// TestRoundTripReinitializesAfterPriorReinit verifies that re-initialization
+// triggers correctly on a second failure when the session already has a
+// backend_sid from a prior re-init. Without the clientSID capture fix,
+// RoundTrip rewrites the header to backend_sid before calling reinitializeAndReplay,
+// which then looks up the session by the (wrong) backend SID and finds nothing.
+func TestRoundTripReinitializesAfterPriorReinit(t *testing.T) {
+	t.Parallel()
+
+	firstBackendSID := uuid.New().String()
+	secondBackendSID := uuid.New().String()
+
+	// staleBackend: returns 404 to trigger re-init.
+	staleBackend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	defer staleBackend.Close()
+
+	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", secondBackendSID)
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer freshBackend.Close()
+
+	proxy, addr := startProxy(t, freshBackend.URL)
+
+	// Session pre-populated as if a prior re-init already happened:
+	// backend_url points to staleBackend, backend_sid is set to firstBackendSID.
+	clientSessionID := uuid.New().String()
+	sess := session.NewProxySession(clientSessionID)
+	sess.SetMetadata(sessionMetadataBackendURL, staleBackend.URL)
+	sess.SetMetadata(sessionMetadataInitBody, `{"jsonrpc":"2.0","id":1,"method":"initialize"}`)
+	sess.SetMetadata(sessionMetadataBackendSID, firstBackendSID)
+	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: re-init must use client SID for session lookup, not backend SID")
+	assert.GreaterOrEqual(t, freshHit.Load(), int32(2),
+		"fresh backend should receive re-initialize + replay")
+}
+
 // 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