Skip to content

Wire SessionStorage into MCPRemoteProxy for HA support #5237

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aron-muon wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Wire SessionStorage into MCPRemoteProxy for HA support #5237

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e181090
Wire SessionStorage into MCPRemoteProxy for HA support
aron-muon May 9, 2026
c52a373
chore: manifests
aron-muon May 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions cmd/thv-operator/api/v1beta1/mcpremoteproxy_types.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,22 @@ type MCPRemoteProxySpec struct {
// +kubebuilder:default=ClientIP
// +optional
SessionAffinity string `json:"sessionAffinity,omitempty"`

// SessionStorage configures session storage for stateful horizontal scaling.
// When nil, no session storage is configured and the proxy falls back to
// pod-local in-memory session state — incompatible with multi-replica
// deployments behind load balancers that don't preserve client-IP affinity
// (e.g. AWS ALB across multiple AZs).
//
// The transparent proxy validates `Mcp-Session-Id` against this store on
// every non-initialize request (see pkg/transport/proxy/transparent/
// transparent_proxy.go) and rewrites client-facing session IDs to backend
// session IDs using session metadata. Both lookups require shared state
// across replicas.
//
// Mirrors MCPServer.spec.sessionStorage and VirtualMCPServer.spec.sessionStorage.
// +optional
Copy link
Copy Markdown
Contributor

@yrobla yrobla May 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider documenting the sessionAffinity + sessionStorage interaction.

These two fields now coexist without any cross-reference, which creates a footgun. When sessionStorage.provider == redis the correct HA setup also requires sessionAffinity: None — otherwise the k8s Service's ClientIP stickiness makes the Redis store redundant. Conversely, sessionAffinity: None without Redis is the broken configuration this PR was written to fix.

At minimum, the SessionStorage field comment should say: "When using the Redis provider, also set sessionAffinity: None so the Service routes requests round-robin and all replicas rely on the shared session store rather than pod-local state."

A stronger option is a kubebuilder XValidation that requires sessionAffinity == "None" when sessionStorage.provider == "redis", catching the misconfiguration at admission time.

SessionStorage *SessionStorageConfig `json:"sessionStorage,omitempty"`
}

// MCPRemoteProxyStatus defines the observed state of MCPRemoteProxy
Expand Down
5 changes: 5 additions & 0 deletions cmd/thv-operator/api/v1beta1/zz_generated.deepcopy.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

30 changes: 30 additions & 0 deletions cmd/thv-operator/controllers/mcpremoteproxy_deployment.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
ctrlutil "github.com/stacklok/toolhive/cmd/thv-operator/pkg/controllerutil"
"github.com/stacklok/toolhive/cmd/thv-operator/pkg/runconfig/configmap/checksum"
"github.com/stacklok/toolhive/pkg/container/kubernetes"
"github.com/stacklok/toolhive/pkg/transport/session"
)

// deploymentForMCPRemoteProxy returns a MCPRemoteProxy Deployment object
Expand Down Expand Up @@ -232,9 +233,38 @@ func (r *MCPRemoteProxyReconciler) buildEnvVarsForProxy(
}
}

// Add THV_SESSION_REDIS_PASSWORD when sessionStorage uses a passwordRef.
// The non-sensitive parts (address/db/keyPrefix) are populated into the
// runconfig by populateScalingConfigForRemoteProxy; the password is
// injected separately so it never lands in the ConfigMap.
env = append(env, buildRedisPasswordEnvVarForRemoteProxy(proxy)...)

return ctrlutil.EnsureRequiredEnvVars(ctx, env)
}

// buildRedisPasswordEnvVarForRemoteProxy returns the THV_SESSION_REDIS_PASSWORD
// env var sourced from spec.sessionStorage.passwordRef when sessionStorage uses
// the redis provider; returns nil otherwise. Mirrors VirtualMCPServer's
// buildRedisPasswordEnvVar in virtualmcpserver_deployment.go.
func buildRedisPasswordEnvVarForRemoteProxy(proxy *mcpv1beta1.MCPRemoteProxy) []corev1.EnvVar {
if proxy.Spec.SessionStorage == nil ||
proxy.Spec.SessionStorage.Provider != mcpv1beta1.SessionStorageProviderRedis ||
proxy.Spec.SessionStorage.PasswordRef == nil {
return nil
}
return []corev1.EnvVar{{
Name: session.RedisPasswordEnvVar,
ValueFrom: &corev1.EnvVarSource{
SecretKeyRef: &corev1.SecretKeySelector{
LocalObjectReference: corev1.LocalObjectReference{
Name: proxy.Spec.SessionStorage.PasswordRef.Name,
},
Key: proxy.Spec.SessionStorage.PasswordRef.Key,
},
},
}}
}

// buildOIDCClientSecretEnvVars returns OIDC client secret env vars when the proxy
// references an MCPOIDCConfig with an inline client secret. Returns nil otherwise.
func (r *MCPRemoteProxyReconciler) buildOIDCClientSecretEnvVars(
Expand Down
64 changes: 64 additions & 0 deletions cmd/thv-operator/controllers/mcpremoteproxy_deployment_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import (

mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1"
ctrlutil "github.com/stacklok/toolhive/cmd/thv-operator/pkg/controllerutil"
"github.com/stacklok/toolhive/pkg/transport/session"
)

// TestDeploymentForMCPRemoteProxy tests deployment generation
Expand Down Expand Up @@ -1305,3 +1306,66 @@ func TestMCPRemoteProxyServiceNeedsUpdate(t *testing.T) {
})
}
}

// TestBuildRedisPasswordEnvVarForRemoteProxy mirrors VirtualMCPServer's
// TestBuildRedisPasswordEnvVar — the env var must be injected only when
// sessionStorage uses the redis provider AND a passwordRef is set, and it
// must always be a SecretKeyRef (never a plaintext value).
func TestBuildRedisPasswordEnvVarForRemoteProxy(t *testing.T) {
t.Parallel()

passwordRef := &mcpv1beta1.SecretKeyRef{Name: "redis-secret", Key: "password"}

tests := []struct {
name string
storage *mcpv1beta1.SessionStorageConfig
expectEnVar bool
}{
{
name: "nil sessionStorage produces no env var",
storage: nil,
expectEnVar: false,
},
{
name: "memory provider produces no env var",
storage: &mcpv1beta1.SessionStorageConfig{Provider: "memory"},
expectEnVar: false,
},
{
name: "redis without passwordRef produces no env var",
storage: &mcpv1beta1.SessionStorageConfig{Provider: mcpv1beta1.SessionStorageProviderRedis, Address: "redis:6379"},
expectEnVar: false,
},
{
name: "redis with passwordRef produces THV_SESSION_REDIS_PASSWORD",
storage: &mcpv1beta1.SessionStorageConfig{
Provider: mcpv1beta1.SessionStorageProviderRedis,
Address: "redis:6379",
PasswordRef: passwordRef,
},
expectEnVar: true,
},
}

for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
proxy := &mcpv1beta1.MCPRemoteProxy{
ObjectMeta: metav1.ObjectMeta{Name: "test-proxy", Namespace: "default"},
Spec: mcpv1beta1.MCPRemoteProxySpec{SessionStorage: tc.storage},
}
env := buildRedisPasswordEnvVarForRemoteProxy(proxy)
if tc.expectEnVar {
require.Len(t, env, 1)
assert.Equal(t, session.RedisPasswordEnvVar, env[0].Name)
assert.Empty(t, env[0].Value, "must not use plaintext Value")
require.NotNil(t, env[0].ValueFrom)
require.NotNil(t, env[0].ValueFrom.SecretKeyRef)
assert.Equal(t, passwordRef.Name, env[0].ValueFrom.SecretKeyRef.Name)
assert.Equal(t, passwordRef.Key, env[0].ValueFrom.SecretKeyRef.Key)
} else {
assert.Empty(t, env)
}
})
}
}
28 changes: 28 additions & 0 deletions cmd/thv-operator/controllers/mcpremoteproxy_runconfig.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,9 +173,37 @@ func (r *MCPRemoteProxyReconciler) createRunConfigFromMCPRemoteProxy(
return nil, fmt.Errorf("failed to populate middleware configs: %w", err)
}

// Populate ScalingConfig.SessionRedis from spec.sessionStorage so the
// proxy runner has the address/db/keyPrefix needed to construct a
// shared Redis-backed session store. The Redis password is intentionally
// excluded here — it is injected as the THV_SESSION_REDIS_PASSWORD env
// var by buildRedisPasswordEnvVar in mcpremoteproxy_deployment.go.
populateScalingConfigForRemoteProxy(runConfig, proxy)
Copy link
Copy Markdown
Contributor

@yrobla yrobla May 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ordering bug (latent): populateScalingConfigForRemoteProxy must be called before PopulateMiddlewareConfigs, not after.

PopulateMiddlewareConfigs → addRateLimitMiddleware reads config.ScalingConfig.SessionRedis and returns an error when it is nil and a rate-limit config is present. The equivalent call in mcpserver_runconfig.go explicitly documents this:

// Must run before PopulateMiddlewareConfigs because rate limiting reads SessionRedis.
populateScalingConfig(runConfig, m)
if err := runner.PopulateMiddlewareConfigs(runConfig); err != nil { ... }

MCPRemoteProxy has no rateLimiting field today so this is harmless right now, but the moment rate limiting is added (a natural follow-on given MCPServer already has it) this ordering will cause a confusing runtime error. Please swap the two blocks:

Suggested change
populateScalingConfigForRemoteProxy(runConfig, proxy)
// Populate ScalingConfig.SessionRedis from spec.sessionStorage so the
// proxy runner has the address/db/keyPrefix needed to construct a
// shared Redis-backed session store. The Redis password is intentionally
// excluded here — it is injected as the THV_SESSION_REDIS_PASSWORD env
// var by buildRedisPasswordEnvVar in mcpremoteproxy_deployment.go.
// Must run before PopulateMiddlewareConfigs because rate limiting reads SessionRedis.
populateScalingConfigForRemoteProxy(runConfig, proxy)
// Populate middleware configs from the configuration fields
// This ensures that middleware_configs is properly set for serialization
if err := runner.PopulateMiddlewareConfigs(runConfig); err != nil {
return nil, fmt.Errorf("failed to populate middleware configs: %w", err)
}
return runConfig, nil


return runConfig, nil
}

// populateScalingConfigForRemoteProxy mirrors populateScalingConfig from
// mcpserver_runconfig.go but for MCPRemoteProxy (which has no
// BackendReplicas concept). When MCPRemoteProxy.spec.sessionStorage uses
// the redis provider, this populates runner.ScalingConfig.SessionRedis with
// the non-sensitive connection parameters.
func populateScalingConfigForRemoteProxy(runConfig *runner.RunConfig, proxy *mcpv1beta1.MCPRemoteProxy) {
if proxy.Spec.SessionStorage == nil ||
proxy.Spec.SessionStorage.Provider != mcpv1beta1.SessionStorageProviderRedis {
return
}

if runConfig.ScalingConfig == nil {
runConfig.ScalingConfig = &runner.ScalingConfig{}
}
runConfig.ScalingConfig.SessionRedis = &runner.SessionRedisConfig{
Address: proxy.Spec.SessionStorage.Address,
DB: proxy.Spec.SessionStorage.DB,
KeyPrefix: proxy.Spec.SessionStorage.KeyPrefix,
}
}

// resolveAndAddOIDCConfig resolves OIDC configuration from the shared MCPOIDCConfigRef,
// adds the appropriate runner options, and returns the resolved config.
func (r *MCPRemoteProxyReconciler) resolveAndAddOIDCConfig(
Expand Down
86 changes: 86 additions & 0 deletions cmd/thv-operator/controllers/mcpremoteproxy_runconfig_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -856,3 +856,89 @@ func TestLabelsForRunConfigRemoteProxy(t *testing.T) {
result := labelsForRunConfigRemoteProxy("test-proxy")
assert.Equal(t, expected, result)
}

// TestPopulateScalingConfigForRemoteProxy tests SessionRedis injection into RunConfig
// from MCPRemoteProxy.spec.sessionStorage. Mirrors TestPopulateScalingConfig in
// mcpserver_runconfig_test.go. MCPRemoteProxy has no BackendReplicas concept (no
// backend StatefulSet), so we only cover the SessionRedis path.
func TestPopulateScalingConfigForRemoteProxy(t *testing.T) {
t.Parallel()

tests := []struct {
name string
storage *mcpv1beta1.SessionStorageConfig
expected func(t *testing.T, sc *runner.ScalingConfig)
}{
{
name: "nil sessionStorage — ScalingConfig stays nil",
storage: nil,
expected: func(t *testing.T, sc *runner.ScalingConfig) {
t.Helper()
assert.Nil(t, sc)
},
},
{
name: "memory provider — ScalingConfig stays nil",
storage: &mcpv1beta1.SessionStorageConfig{Provider: "memory"},
expected: func(t *testing.T, sc *runner.ScalingConfig) {
t.Helper()
assert.Nil(t, sc)
},
},
{
name: "redis — address/db/keyPrefix written to SessionRedis",
storage: &mcpv1beta1.SessionStorageConfig{
Provider: mcpv1beta1.SessionStorageProviderRedis,
Address: "redis.default.svc:6379",
DB: 2,
KeyPrefix: "thv:",
},
expected: func(t *testing.T, sc *runner.ScalingConfig) {
t.Helper()
require.NotNil(t, sc)
require.NotNil(t, sc.SessionRedis)
assert.Equal(t, "redis.default.svc:6379", sc.SessionRedis.Address)
assert.Equal(t, int32(2), sc.SessionRedis.DB)
assert.Equal(t, "thv:", sc.SessionRedis.KeyPrefix)
},
},
{
name: "redis with passwordRef — password NOT in SessionRedis",
storage: &mcpv1beta1.SessionStorageConfig{
Provider: mcpv1beta1.SessionStorageProviderRedis,
Address: "redis:6379",
PasswordRef: &mcpv1beta1.SecretKeyRef{
Name: "redis-secret",
Key: "password",
},
},
expected: func(t *testing.T, sc *runner.ScalingConfig) {
t.Helper()
require.NotNil(t, sc)
require.NotNil(t, sc.SessionRedis)
assert.Equal(t, "redis:6379", sc.SessionRedis.Address)
// Password must NOT leak into the RunConfig — it's injected
// separately as the THV_SESSION_REDIS_PASSWORD pod env var
// by buildRedisPasswordEnvVarForRemoteProxy.
data, err := json.Marshal(sc)
require.NoError(t, err)
assert.NotContains(t, string(data), "redis-secret")
assert.NotContains(t, string(data), "password")
},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
runConfig := &runner.RunConfig{}
proxy := &mcpv1beta1.MCPRemoteProxy{
Spec: mcpv1beta1.MCPRemoteProxySpec{
SessionStorage: tt.storage,
},
}
populateScalingConfigForRemoteProxy(runConfig, proxy)
tt.expected(t, runConfig.ScalingConfig)
})
}
}
Loading
Loading