Route MCP sessions to specific StatefulSet pods via headless DNS

cmd/thv-operator/controllers/mcpserver_runconfig.go

@@ -288,8 +288,8 @@ func (r *MCPServerReconciler) createRunConfigFromMCPServer(m *mcpv1alpha1.MCPSer
 	return runConfig, nil
 }
 
-// populateScalingConfig sets BackendReplicas and SessionRedis on the RunConfig from the MCPServer spec.
-// Fields are only set when present in the spec; nil means "not configured" and is left as-is.
+// populateScalingConfig sets BackendReplicas, SessionRedis, and HeadlessService on the RunConfig
+// from the MCPServer spec. Fields are only set when present in the spec; nil means "not configured".
 func populateScalingConfig(runConfig *runner.RunConfig, m *mcpv1alpha1.MCPServer) {
 	hasBackendReplicas := m.Spec.BackendReplicas != nil
 	hasRedis := m.Spec.SessionStorage != nil && m.Spec.SessionStorage.Provider == mcpv1alpha1.SessionStorageProviderRedis
@@ -305,6 +305,18 @@ func populateScalingConfig(runConfig *runner.RunConfig, m *mcpv1alpha1.MCPServer
 	if hasBackendReplicas {
 		val := *m.Spec.BackendReplicas
 		runConfig.ScalingConfig.BackendReplicas = &val
+
+		// Always populate headless service config when BackendReplicas is set.
+		// This enables the proxy runner to route each session to a specific pod via
+		// headless DNS (e.g. myserver-0.mcp-myserver-headless.default.svc.cluster.local)
+		// so sessions survive proxy-runner restarts. For single-replica StatefulSets,
+		// ordinal 0 is always selected deterministically.
+		runConfig.ScalingConfig.HeadlessService = &transporttypes.HeadlessServiceConfig{
+			StatefulSetName: m.Name,
+			ServiceName:     fmt.Sprintf("mcp-%s-headless", m.Name),
+			Namespace:       m.Namespace,
+			Replicas:        val,
+		}
 	}
 
 	if hasRedis {

pkg/container/kubernetes/client.go

@@ -447,7 +447,7 @@ func buildStatefulSetSpec(
 	spec := appsv1apply.StatefulSetSpec().
 		WithSelector(metav1apply.LabelSelector().
 			WithMatchLabels(map[string]string{"app": containerName})).
-		WithServiceName(containerName).
+		WithServiceName(fmt.Sprintf("mcp-%s-headless", containerName)).
 		WithTemplate(podTemplateSpec)
 	if options != nil && options.ScalingConfig != nil && options.ScalingConfig.BackendReplicas != nil {
 		spec = spec.WithReplicas(*options.ScalingConfig.BackendReplicas)

pkg/runner/config.go

@@ -247,6 +247,12 @@ type ScalingConfig struct {
 	// The Redis password is not included — it is injected as env var THV_SESSION_REDIS_PASSWORD.
 	// +optional
 	SessionRedis *SessionRedisConfig `json:"session_redis,omitempty" yaml:"session_redis,omitempty"`
+
+	// HeadlessService holds the information needed to construct pod-specific headless DNS URLs
+	// for session-affinity routing in Kubernetes StatefulSet deployments.
+	// Populated by the operator whenever BackendReplicas is set (including single-replica).
+	// +optional
+	HeadlessService *types.HeadlessServiceConfig `json:"headless_service,omitempty" yaml:"headless_service,omitempty"`
 }
 
 // SessionRedisConfig contains non-sensitive Redis connection parameters used for distributed

pkg/runner/runner.go

@@ -35,6 +35,7 @@ import (
 	"github.com/stacklok/toolhive/pkg/transport"
 	"github.com/stacklok/toolhive/pkg/transport/session"
 	"github.com/stacklok/toolhive/pkg/transport/types"
+	vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config"
 	"github.com/stacklok/toolhive/pkg/workloads/statuses"
 )
 
@@ -354,6 +355,11 @@ func (r *Runner) Run(ctx context.Context) error {
 		}
 	}
 
+	// Enable pod-specific session routing for Kubernetes StatefulSet backends.
+	if r.Config.ScalingConfig != nil && r.Config.ScalingConfig.HeadlessService != nil {
+		transportConfig.HeadlessService = r.Config.ScalingConfig.HeadlessService
+	}
+
 	// When Redis session storage is configured, create a Redis-backed session store
 	// so sessions are shared across proxy replicas instead of being pod-local.
 	if r.Config.ScalingConfig != nil && r.Config.ScalingConfig.SessionRedis != nil {
@@ -364,7 +370,7 @@ func (r *Runner) Run(ctx context.Context) error {
 		}
 		storage, err := session.NewRedisStorage(ctx, session.RedisConfig{
 			Addr:      redisCfg.Address,
-			Password:  os.Getenv(session.RedisPasswordEnvVar),
+			Password:  os.Getenv(vmcpconfig.RedisPasswordEnvVar),
 			DB:        int(redisCfg.DB),
 			KeyPrefix: keyPrefix,
 		}, session.DefaultSessionTTL)

pkg/transport/factory.go

@@ -73,6 +73,7 @@ func (*Factory) Create(config types.Config, opts ...Option) (types.Transport, er
 			config.Middlewares...,
 		)
 		httpTransport.sessionStorage = config.SessionStorage
+		httpTransport.headlessService = config.HeadlessService
 		tr = httpTransport
 	case types.TransportTypeStreamableHTTP:
 		httpTransport := NewHTTPTransport(
@@ -91,6 +92,7 @@ func (*Factory) Create(config types.Config, opts ...Option) (types.Transport, er
 			config.Middlewares...,
 		)
 		httpTransport.sessionStorage = config.SessionStorage
+		httpTransport.headlessService = config.HeadlessService
 		tr = httpTransport
 	case types.TransportTypeInspector:
 		// HTTP transport is not implemented yet

pkg/transport/http.go

@@ -74,6 +74,9 @@ type HTTPTransport struct {
 	// Mutex for protecting shared state
 	mutex sync.Mutex
 
+	// headlessService configures pod-specific routing for Kubernetes StatefulSet deployments.
+	headlessService *types.HeadlessServiceConfig
+
 	// sessionStorage overrides the default in-memory session store when set.
 	// Used for Redis-backed session sharing across replicas.
 	sessionStorage session.Storage
@@ -239,6 +242,33 @@ func (t *HTTPTransport) setTargetURI(targetURI string) {
 	t.targetURI = targetURI
 }
 
+// resolveTargetURI determines the proxy target URI, base path, and raw query from the
+// transport configuration. For remote MCP servers it parses the remote URL; for local
+// containers it returns the pre-configured targetURI.
+func (t *HTTPTransport) resolveTargetURI() (targetURI, remoteBasePath, remoteRawQuery string, err error) {
+	if t.remoteURL != "" {
+		remoteURL, err := url.Parse(t.remoteURL)
+		if err != nil {
+			return "", "", "", fmt.Errorf("failed to parse remote URL: %w", err)
+		}
+		targetURI = (&url.URL{Scheme: remoteURL.Scheme, Host: remoteURL.Host}).String()
+		remoteBasePath = remoteURL.Path
+		remoteRawQuery = remoteURL.RawQuery
+		slog.Debug("setting up transparent proxy to forward to remote URL",
+			"port", t.proxyPort, "target", targetURI, "base_path", remoteBasePath, "raw_query", remoteRawQuery)
+		return targetURI, remoteBasePath, remoteRawQuery, nil
+	}
+	if t.containerName == "" {
+		return "", "", "", transporterrors.ErrContainerNameNotSet
+	}
+	if t.targetURI == "" {
+		return "", "", "", fmt.Errorf("target URI not set for HTTP transport")
+	}
+	slog.Debug("setting up transparent proxy to forward to target",
+		"port", t.proxyPort, "target", t.targetURI)
+	return t.targetURI, "", "", nil
+}
+
 // Start initializes the transport and begins processing messages.
 // The transport is responsible for starting the container.
 //
@@ -251,52 +281,15 @@ func (t *HTTPTransport) Start(ctx context.Context) error {
 		return fmt.Errorf("container deployer not set")
 	}
 
-	// Determine target URI
-	var targetURI string
-
 	// remoteBasePath holds the path component from the remote URL (e.g., "/v2" from
 	// "https://mcp.asana.com/v2/mcp"). This must be prepended to incoming request
 	// paths so they reach the correct endpoint on the remote server.
-	var remoteBasePath string
-
 	// remoteRawQuery holds the raw query string from the remote URL (e.g.,
 	// "toolsets=core,alerting" from "https://mcp.example.com/mcp?toolsets=core,alerting").
 	// This must be forwarded on every outbound request or it is silently dropped.
-	var remoteRawQuery string
-
-	if t.remoteURL != "" {
-		// For remote MCP servers, construct target URI from remote URL
-		remoteURL, err := url.Parse(t.remoteURL)
-		if err != nil {
-			return fmt.Errorf("failed to parse remote URL: %w", err)
-		}
-		targetURI = (&url.URL{
-			Scheme: remoteURL.Scheme,
-			Host:   remoteURL.Host,
-		}).String()
-
-		// Extract the path prefix that needs to be prepended to incoming requests.
-		// The target URI only has scheme+host, so without this the remote path is lost.
-		remoteBasePath = remoteURL.Path
-
-		remoteRawQuery = remoteURL.RawQuery
-
-		//nolint:gosec // G706: logging proxy port and remote URL from config
-		slog.Debug("setting up transparent proxy to forward to remote URL",
-			"port", t.proxyPort, "target", targetURI, "base_path", remoteBasePath, "raw_query", remoteRawQuery)
-	} else {
-		if t.containerName == "" {
-			return transporterrors.ErrContainerNameNotSet
-		}
-
-		// For local containers, use the configured target URI
-		if t.targetURI == "" {
-			return fmt.Errorf("target URI not set for HTTP transport")
-		}
-		targetURI = t.targetURI
-		//nolint:gosec // G706: logging proxy port and target URI from config
-		slog.Debug("setting up transparent proxy to forward to target",
-			"port", t.proxyPort, "target", targetURI)
+	targetURI, remoteBasePath, remoteRawQuery, err := t.resolveTargetURI()
+	if err != nil {
+		return err
 	}
 
 	// Create middlewares slice
@@ -330,6 +323,21 @@ func (t *HTTPTransport) Start(ctx context.Context) error {
 		proxyOptions = append(proxyOptions, transparent.WithSessionStorage(t.sessionStorage))
 	}
 
+	// Inject Redis-backed session storage for cross-replica session sharing.
+	if t.sessionStorage != nil {
+		proxyOptions = append(proxyOptions, transparent.WithSessionStorage(t.sessionStorage))
+	}
+
+	// Enable pod-specific routing for Kubernetes StatefulSet backends.
+	// When configured, each new session is pinned to a specific pod via headless DNS
+	// so that session routing survives proxy-runner restarts.
+	if t.headlessService != nil {
+		proxyOptions = append(proxyOptions, transparent.WithPodHeadlessService(
+			t.headlessService.StatefulSetName, t.headlessService.ServiceName,
+			t.headlessService.Namespace, t.headlessService.Replicas,
+		))
+	}
+
 	// Create the transparent proxy
 	t.proxy = transparent.NewTransparentProxyWithOptions(
 		t.host,