From 6ca46c68ffb8608a908f4088edabaa36a584dbbd Mon Sep 17 00:00:00 2001 From: Trey Date: Thu, 18 Jun 2026 10:05:20 -0700 Subject: [PATCH 1/7] Restore session-stable passthrough-header forwarding on the Serve path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On the Serve path (core.New + Serve), tool/resource calls route through the shared backend client (pkg/vmcp/client) rather than through per-session backend connections. The shared client was not reading the passthrough headers captured by headerforward.CaptureMiddleware, and even if it had, it would read them per-request rather than using the session-creation-time snapshot — a behavioral and security regression vs. the legacy per-session-connection path (#5560). - Export MergeForwardedHeaders from pkg/vmcp/headerforward so both the session connector (legacy path) and the shared backend client (Serve path) share the same restricted-header and collision-detection rules. - Teach httpBackendClient.defaultClientFactory to merge ForwardedHeadersFromContext(ctx) with target.HeaderForward, enabling the Serve path to forward captured headers to backends. - Add capturedPassthroughHeaders sync.Map on Server (node-local, never persisted to Redis per security.md) to hold the session-creation-time header snapshot for each active session. - Capture the snapshot in handleSessionRegistrationImpl once per session when s.core != nil, and re-inject it via injectCapturedHeaders in coreToolHandler and coreResourceHandler before each backend call so mid-session header changes on the client side never reach the backend. - Register an AddOnUnregisterSession hook to clean up captured headers when the SDK unregisters a session. - Switch the integration test helper (NewVMCPServer) from server.New to core.New + Serve so tests exercise the production Serve path. - Wire workflow telemetry into core.New via a telemetryComposer wrapper (same metric names as the session-factory path) so existing telemetry tests continue to pass on the Serve path. - Re-enable TestVMCPServer_PassthroughHeaders on the Serve path; it now passes including the session-stability assertion (second call sees the session-creation-time value, not the mid-session changed value). Closes #5560 Co-Authored-By: Claude Sonnet 4.6 (1M context) --- pkg/vmcp/client/client.go | 18 ++- pkg/vmcp/core/core_telemetry.go | 134 +++++++++++++++++ pkg/vmcp/core/core_vmcp.go | 25 ++- pkg/vmcp/headerforward/transport.go | 65 ++++++++ pkg/vmcp/headerforward/transport_test.go | 142 ++++++++++++++++++ pkg/vmcp/server/serve.go | 14 +- pkg/vmcp/server/serve_handlers.go | 35 ++++- pkg/vmcp/server/server.go | 26 ++++ .../session/internal/backend/mcp_session.go | 65 +------- test/integration/vmcp/helpers/vmcp_server.go | 92 +++++------- 10 files changed, 489 insertions(+), 127 deletions(-) create mode 100644 pkg/vmcp/core/core_telemetry.go diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 615959e83b..9d27f91529 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -433,11 +433,21 @@ func (h *httpBackendClient) defaultClientFactory(ctx context.Context, target *vm target: target, } - // Inject per-backend HTTP headers from MCPServerEntry.spec.headerForward. - // Resolves plaintext + secret-backed headers once here; auth (inner) always - // wins over user-supplied headers because it runs after this tripper. + // Merge session-captured passthrough headers (from ctx) with the static + // per-backend header-forward config. On the Serve path, the caller (coreToolHandler + // in pkg/vmcp/server/serve_handlers.go) replaces the per-request forwarded headers + // on ctx with the session-stable snapshot captured once at session creation, so + // passing ctx here provides session-stable header injection without re-reading the + // live incoming request headers. Restricted names and static-config collisions are + // rejected by MergeForwardedHeaders. + mergedHeaderForward, mergeErr := headerforward.MergeForwardedHeaders( + target.HeaderForward, headerforward.ForwardedHeadersFromContext(ctx), + ) + if mergeErr != nil { + return nil, fmt.Errorf("failed to merge forwarded headers for backend %s: %w", target.WorkloadID, mergeErr) + } baseTransport, err = headerforward.BuildHeaderForwardTripper( - ctx, baseTransport, target.HeaderForward, h.secretsProvider, target.WorkloadID, + ctx, baseTransport, mergedHeaderForward, h.secretsProvider, target.WorkloadID, ) if err != nil { return nil, fmt.Errorf("failed to build header-forward transport: %w", err) diff --git a/pkg/vmcp/core/core_telemetry.go b/pkg/vmcp/core/core_telemetry.go new file mode 100644 index 0000000000..40f565f4e9 --- /dev/null +++ b/pkg/vmcp/core/core_telemetry.go @@ -0,0 +1,134 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package core + +import ( + "context" + "fmt" + "time" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/metric" + "go.opentelemetry.io/otel/trace" + + "github.com/stacklok/toolhive/pkg/telemetry" + "github.com/stacklok/toolhive/pkg/vmcp/composer" +) + +// instrumentationName is the OTEL scope used for all metrics and traces emitted +// by the core. Must match the scope used by pkg/vmcp/internal/backendtelemetry +// and pkg/vmcp/server/sessionmanager so they share the same Prometheus namespace. +const instrumentationName = "github.com/stacklok/toolhive/pkg/vmcp" + +// workflowInstruments holds pre-built OTEL instruments for workflow execution +// telemetry. Created once in core.New when TelemetryProvider is set and reused +// across all per-call composer factories. +type workflowInstruments struct { + tracer trace.Tracer + executionsTotal metric.Int64Counter + errorsTotal metric.Int64Counter + executionDuration metric.Float64Histogram +} + +// newWorkflowInstruments creates the OTEL instruments for workflow execution +// telemetry. Returns nil if provider is nil (telemetry disabled). +// The metric names match those emitted by the session-factory composite-tool +// decorator in pkg/vmcp/server/sessionmanager so all telemetry appears under +// the same Prometheus metrics regardless of which path executed the workflow. +func newWorkflowInstruments(provider *telemetry.Provider) (*workflowInstruments, error) { + if provider == nil { + return nil, nil + } + + meter := provider.MeterProvider().Meter(instrumentationName) + + executions, err := meter.Int64Counter( + "toolhive_vmcp_workflow_executions", + metric.WithDescription("Total number of workflow executions"), + ) + if err != nil { + return nil, fmt.Errorf("failed to create workflow executions counter: %w", err) + } + + errors, err := meter.Int64Counter( + "toolhive_vmcp_workflow_errors", + metric.WithDescription("Total number of workflow execution errors"), + ) + if err != nil { + return nil, fmt.Errorf("failed to create workflow errors counter: %w", err) + } + + duration, err := meter.Float64Histogram( + "toolhive_vmcp_workflow_duration", + metric.WithDescription("Duration of workflow executions in seconds"), + metric.WithUnit("s"), + metric.WithExplicitBucketBoundaries(telemetry.MCPHistogramBuckets...), + ) + if err != nil { + return nil, fmt.Errorf("failed to create workflow duration histogram: %w", err) + } + + return &workflowInstruments{ + tracer: provider.TracerProvider().Tracer(instrumentationName), + executionsTotal: executions, + errorsTotal: errors, + executionDuration: duration, + }, nil +} + +// telemetryComposer wraps composer.Composer.ExecuteWorkflow with OTEL metrics +// and tracing. ValidateWorkflow is delegated to the base without instrumentation +// (validation is called once at startup, not on the hot path). +type telemetryComposer struct { + base composer.Composer + instruments *workflowInstruments +} + +var _ composer.Composer = (*telemetryComposer)(nil) + +// ExecuteWorkflow records execution count, duration, and errors before delegating +// to the wrapped composer. Uses def.Name as the workflow_name metric attribute to +// match the label emitted by the session-factory path in sessionmanager. +func (c *telemetryComposer) ExecuteWorkflow( + ctx context.Context, def *composer.WorkflowDefinition, params map[string]any, +) (*composer.WorkflowResult, error) { + commonAttrs := []attribute.KeyValue{attribute.String("workflow.name", def.Name)} + + ctx, span := c.instruments.tracer.Start(ctx, "core.ExecuteWorkflow", + trace.WithAttributes(commonAttrs...), + ) + defer span.End() + + metricAttrs := metric.WithAttributes(commonAttrs...) + start := time.Now() + c.instruments.executionsTotal.Add(ctx, 1, metricAttrs) + + result, err := c.base.ExecuteWorkflow(ctx, def, params) + + c.instruments.executionDuration.Record(ctx, time.Since(start).Seconds(), metricAttrs) + + if err != nil { + c.instruments.errorsTotal.Add(ctx, 1, metricAttrs) + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + } + + return result, err +} + +// ValidateWorkflow delegates to the base composer without instrumentation. +func (c *telemetryComposer) ValidateWorkflow(ctx context.Context, def *composer.WorkflowDefinition) error { + return c.base.ValidateWorkflow(ctx, def) +} + +// GetWorkflowStatus delegates to the base composer without instrumentation. +func (c *telemetryComposer) GetWorkflowStatus(ctx context.Context, workflowID string) (*composer.WorkflowStatus, error) { + return c.base.GetWorkflowStatus(ctx, workflowID) +} + +// CancelWorkflow delegates to the base composer without instrumentation. +func (c *telemetryComposer) CancelWorkflow(ctx context.Context, workflowID string) error { + return c.base.CancelWorkflow(ctx, workflowID) +} diff --git a/pkg/vmcp/core/core_vmcp.go b/pkg/vmcp/core/core_vmcp.go index 11a13c10e1..a398c09b9f 100644 --- a/pkg/vmcp/core/core_vmcp.go +++ b/pkg/vmcp/core/core_vmcp.go @@ -150,17 +150,30 @@ func New(cfg *Config) (VMCP, error) { "type", fmt.Sprintf("%T", stateStore)) } + // Build workflow telemetry instruments once here so they are reused across all + // per-call composer factories. Nil when TelemetryProvider is absent (disabled). + // Uses the same metric names as the session-factory composite-tool decorator in + // pkg/vmcp/server/sessionmanager so both execution paths emit into the same + // Prometheus counters and histograms. + instruments, err := newWorkflowInstruments(cfg.TelemetryProvider) + if err != nil { + stopStore() + return nil, fmt.Errorf("failed to create workflow telemetry instruments: %w", err) + } + // composerFactory builds a composite-tool engine bound to a specific routing - // table, generalizing server.New's sessionComposerFactory (server.go:393-398). - // Workflow-level telemetry (toolhive_vmcp_workflow_*) is intentionally NOT wired - // here: that instrumentation lives in the session/Serve layer - // (sessionmanager/factory.go:174-176) and is added when the core is wired into - // Serve (#5439). This mirrors server.go:393-398, which is also uninstrumented. + // table. When telemetry is configured, it wraps the engine with OTEL metrics so + // workflow executions are instrumented the same way as the session-factory path + // (sessionmanager/factory.go). The telemetryComposer is in core_telemetry.go. composerFactory := func(sessionRT *vmcp.RoutingTable, sessionTools []vmcp.Tool) composer.Composer { - return composer.NewWorkflowEngine( + engine := composer.NewWorkflowEngine( router.NewSessionRouter(sessionRT), backendClient, elicitationHandler, stateStore, workflowAuditor, sessionTools, ) + if instruments != nil { + return &telemetryComposer{base: engine, instruments: instruments} + } + return engine } // Validate workflows fail-fast (server.go:400-405). The validation engine uses diff --git a/pkg/vmcp/headerforward/transport.go b/pkg/vmcp/headerforward/transport.go index e235fcb17a..1ff3b10277 100644 --- a/pkg/vmcp/headerforward/transport.go +++ b/pkg/vmcp/headerforward/transport.go @@ -101,6 +101,71 @@ func BuildHeaderForwardTripper( return &headerForwardRoundTripper{base: base, headers: headers}, nil } +// MergeForwardedHeaders returns a HeaderForwardConfig that combines the static +// backend configuration (base) with any per-request forwarded headers captured +// from the caller's request (see CaptureMiddleware / ForwardedHeadersFromContext). +// +// Rules (applied in order): +// 1. If forwarded is empty, base is returned unchanged (no allocation, same pointer). +// 2. A new HeaderForwardConfig is built from a shallow copy of base so the +// shared target.HeaderForward is never mutated. +// 3. Forwarded header names are canonicalized via http.CanonicalHeaderKey and +// checked against middleware.RestrictedHeaders; restricted names are silently +// dropped (defense-in-depth — they were already filtered upstream by +// CaptureMiddleware, but we guard here too). +// 4. A forwarded header name that also appears in base (AddPlaintextHeaders or +// AddHeadersFromSecret) is a misconfiguration: the function returns an error +// rather than silently picking a winner. +func MergeForwardedHeaders(base *vmcp.HeaderForwardConfig, forwarded map[string]string) (*vmcp.HeaderForwardConfig, error) { + if len(forwarded) == 0 { + return base, nil + } + + // Build the merged AddPlaintextHeaders starting from the static config. + var staticPlaintext map[string]string + if base != nil { + staticPlaintext = base.AddPlaintextHeaders + } + + // Canonical set of header names already owned by the static config, used for + // collision detection. + staticNames := make(map[string]struct{}) + if base != nil { + for k := range base.AddPlaintextHeaders { + staticNames[http.CanonicalHeaderKey(k)] = struct{}{} + } + for k := range base.AddHeadersFromSecret { + staticNames[http.CanonicalHeaderKey(k)] = struct{}{} + } + } + + merged := make(map[string]string, len(staticPlaintext)+len(forwarded)) + for k, v := range staticPlaintext { + merged[k] = v + } + + for name, value := range forwarded { + canonical := http.CanonicalHeaderKey(name) + if middleware.RestrictedHeaders[canonical] { + slog.Debug("dropping restricted forwarded header", "header", canonical) + continue + } + if _, exists := staticNames[canonical]; exists { + return nil, fmt.Errorf( + "forwarded header %q collides with the backend's static header-forward config", canonical) + } + merged[canonical] = value + } + + out := &vmcp.HeaderForwardConfig{ + AddPlaintextHeaders: merged, + } + if base != nil { + out.AddHeadersFromSecret = base.AddHeadersFromSecret + } + return out, nil +} + // resolveHeaderForward merges plaintext headers with secret-resolved headers into // a single canonicalized http.Header. Plaintext entries are copied verbatim; // secret identifiers are looked up via the provider (TOOLHIVE_SECRET_). diff --git a/pkg/vmcp/headerforward/transport_test.go b/pkg/vmcp/headerforward/transport_test.go index d1d19042d8..f8be89c8d0 100644 --- a/pkg/vmcp/headerforward/transport_test.go +++ b/pkg/vmcp/headerforward/transport_test.go @@ -241,3 +241,145 @@ func TestBuildHeaderForwardTripper_NilCfgReturnsBase(t *testing.T) { require.NoError(t, err) assert.Same(t, base, got, "nil cfg must pass base through untouched") } + +// TestMergeForwardedHeaders exercises the key rules of MergeForwardedHeaders: +// empty forwarded, static-only, forwarded-only, merged output, restricted-name +// drop, and collision detection. +func TestMergeForwardedHeaders(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + base *vmcp.HeaderForwardConfig + forwarded map[string]string + wantCfg *vmcp.HeaderForwardConfig + wantErr bool + }{ + { + name: "empty forwarded returns base unchanged", + base: &vmcp.HeaderForwardConfig{AddPlaintextHeaders: map[string]string{"X-Static": "v"}}, + forwarded: nil, + wantCfg: &vmcp.HeaderForwardConfig{AddPlaintextHeaders: map[string]string{"X-Static": "v"}}, + }, + { + name: "nil base and empty forwarded returns nil", + base: nil, + forwarded: nil, + wantCfg: nil, + }, + { + name: "nil base with forwarded yields forwarded-only config", + base: nil, + forwarded: map[string]string{"X-Api-Key": "key123"}, + wantCfg: &vmcp.HeaderForwardConfig{ + AddPlaintextHeaders: map[string]string{"X-Api-Key": "key123"}, + }, + }, + { + name: "forwarded headers are canonicalized", + base: nil, + forwarded: map[string]string{"x-api-key": "key123"}, + wantCfg: &vmcp.HeaderForwardConfig{ + AddPlaintextHeaders: map[string]string{"X-Api-Key": "key123"}, + }, + }, + { + name: "static and forwarded are merged", + base: &vmcp.HeaderForwardConfig{ + AddPlaintextHeaders: map[string]string{"X-Static": "static-value"}, + }, + forwarded: map[string]string{"X-Dynamic": "dynamic-value"}, + wantCfg: &vmcp.HeaderForwardConfig{ + AddPlaintextHeaders: map[string]string{ + "X-Static": "static-value", + "X-Dynamic": "dynamic-value", + }, + }, + }, + { + name: "secret headers preserved from base", + base: &vmcp.HeaderForwardConfig{ + AddHeadersFromSecret: map[string]string{"X-Auth": "MY_SECRET"}, + }, + forwarded: map[string]string{"X-Dynamic": "value"}, + wantCfg: &vmcp.HeaderForwardConfig{ + AddPlaintextHeaders: map[string]string{"X-Dynamic": "value"}, + AddHeadersFromSecret: map[string]string{"X-Auth": "MY_SECRET"}, + }, + }, + { + name: "restricted header is silently dropped", + base: nil, + forwarded: map[string]string{"X-Forwarded-For": "1.2.3.4", "X-Api-Key": "k"}, + wantCfg: &vmcp.HeaderForwardConfig{ + AddPlaintextHeaders: map[string]string{"X-Api-Key": "k"}, + }, + }, + { + name: "collision with plaintext static returns error", + base: &vmcp.HeaderForwardConfig{ + AddPlaintextHeaders: map[string]string{"X-Api-Key": "static"}, + }, + forwarded: map[string]string{"X-Api-Key": "forwarded"}, + wantErr: true, + }, + { + name: "collision with secret static returns error", + base: &vmcp.HeaderForwardConfig{ + AddHeadersFromSecret: map[string]string{"X-Auth": "SECRET_ID"}, + }, + forwarded: map[string]string{"X-Auth": "value"}, + wantErr: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + got, err := MergeForwardedHeaders(tc.base, tc.forwarded) + if tc.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + + if tc.wantCfg == nil { + assert.Nil(t, got) + return + } + require.NotNil(t, got) + assert.Equal(t, tc.wantCfg.AddPlaintextHeaders, got.AddPlaintextHeaders) + assert.Equal(t, tc.wantCfg.AddHeadersFromSecret, got.AddHeadersFromSecret) + }) + } +} + +// TestMergeForwardedHeaders_EmptyForwardedReturnsSamePointer verifies that when +// forwarded is nil/empty, the exact same base pointer is returned (no allocation). +func TestMergeForwardedHeaders_EmptyForwardedReturnsSamePointer(t *testing.T) { + t.Parallel() + base := &vmcp.HeaderForwardConfig{AddPlaintextHeaders: map[string]string{"X-A": "v"}} + got, err := MergeForwardedHeaders(base, nil) + require.NoError(t, err) + assert.Same(t, base, got, "empty forwarded must return the same base pointer") + + got2, err := MergeForwardedHeaders(base, map[string]string{}) + require.NoError(t, err) + assert.Same(t, base, got2, "empty forwarded map must return the same base pointer") +} + +// TestMergeForwardedHeaders_RestrictedHeadersList verifies that every header in +// middleware.RestrictedHeaders is dropped when present in forwarded. +func TestMergeForwardedHeaders_RestrictedHeadersList(t *testing.T) { + t.Parallel() + for name := range middleware.RestrictedHeaders { + forwarded := map[string]string{name: "inject"} + got, err := MergeForwardedHeaders(nil, forwarded) + require.NoError(t, err, "restricted header %q must not return an error", name) + if got != nil { + assert.NotContains(t, got.AddPlaintextHeaders, name, + "restricted header %q must be dropped", name) + } + } +} diff --git a/pkg/vmcp/server/serve.go b/pkg/vmcp/server/serve.go index 5d624f19b8..59e1b50158 100644 --- a/pkg/vmcp/server/serve.go +++ b/pkg/vmcp/server/serve.go @@ -292,10 +292,10 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) return v.Close() }) - // Register the three SDK hooks against the assembled *Server (relocated from - // server.New). They drive phase-2 of two-phase creation (OnRegisterSession) and + // Register the SDK hooks against the assembled *Server (relocated from + // server.New). They drive phase-2 of two-phase creation (OnRegisterSession), // the cross-pod Redis re-hydration of per-session tools (OnBeforeListTools / - // OnBeforeCallTool); the *Server receiver methods are unchanged. + // OnBeforeCallTool), and node-local captured-header cleanup (OnUnregisterSession). hooks.AddOnRegisterSession(func(hookCtx context.Context, session server.ClientSession) { srv.handleSessionRegistration(hookCtx, session) }) @@ -305,6 +305,14 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) hooks.AddBeforeCallTool(func(hookCtx context.Context, _ any, _ *mcp.CallToolRequest) { srv.lazyInjectSessionTools(hookCtx) }) + // Clean up node-local captured passthrough headers when the SDK unregisters the + // session (on client DELETE or SDK-internal timeout). This is the primary cleanup + // path for explicitly-terminated sessions; TTL-expired sessions are cleaned up + // lazily when the node-local cache evicts them. Node-local only — credentials + // must not be persisted to shared state (security.md). + hooks.AddOnUnregisterSession(func(_ context.Context, clientSess server.ClientSession) { + srv.capturedPassthroughHeaders.Delete(clientSess.SessionID()) + }) // Disarm the close-on-error guard: the Server is fully constructed. closeStorageOnErr = false diff --git a/pkg/vmcp/server/serve_handlers.go b/pkg/vmcp/server/serve_handlers.go index ad5bad3949..8bc3f4371b 100644 --- a/pkg/vmcp/server/serve_handlers.go +++ b/pkg/vmcp/server/serve_handlers.go @@ -17,6 +17,7 @@ import ( "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" + "github.com/stacklok/toolhive/pkg/vmcp/headerforward" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types" ) @@ -193,6 +194,14 @@ func (s *Server) coreToolHandler(sessionID, toolName, backendName string) server return mcp.NewToolResultError(fmt.Sprintf("Unauthorized: %v", err)), nil } + // Replace the per-request forwarded headers on ctx with the session-stable + // snapshot captured once at session creation. This ensures the backend always + // receives the value present when the session was established, regardless of + // how the client changes the header on later requests. The shared backend + // client (pkg/vmcp/client) reads the forwarded headers from ctx to build the + // outgoing transport chain. + ctx = s.injectCapturedHeaders(ctx, sessionID) + result, err := s.core.CallTool(ctx, caller, toolName, args, conversion.FromMCPMeta(req.Params.Meta)) if err != nil { // Admission denial returns a generic message so the underlying authorizer @@ -213,7 +222,8 @@ func (s *Server) coreToolHandler(sessionID, toolName, backendName string) server } // coreResourceHandler builds the SDK handler for a Serve-path resource. It mirrors -// coreToolHandler: audit label, binding check, then core.ReadResource with explicit identity. +// coreToolHandler: audit label, binding check, session-stable header injection, then +// core.ReadResource with explicit identity. func (s *Server) coreResourceHandler( sessionID, uri, backendName string, ) func(ctx context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { @@ -228,6 +238,10 @@ func (s *Server) coreResourceHandler( return nil, fmt.Errorf("unauthorized: %w", err) } + // Replace the per-request forwarded headers on ctx with the session-stable + // snapshot captured once at session creation (mirrors coreToolHandler). + ctx = s.injectCapturedHeaders(ctx, sessionID) + result, err := s.core.ReadResource(ctx, caller, uri) if err != nil { if errors.Is(err, vmcp.ErrAuthorizationFailed) { @@ -239,6 +253,25 @@ func (s *Server) coreResourceHandler( } } +// injectCapturedHeaders replaces the per-request forwarded headers on ctx with +// the session-stable snapshot captured at session creation (stored node-locally in +// capturedPassthroughHeaders). If no snapshot exists for this session (e.g. no +// passthrough headers were configured, or the session was created without any +// allowlisted headers present), ctx is returned unchanged. +// +// This is the enforcement point for session-stable header forwarding on the Serve +// path: coreToolHandler and coreResourceHandler call this before delegating to +// core.CallTool / core.ReadResource, so the shared backend client always sees the +// session-creation-time value, not a mid-session change. +func (s *Server) injectCapturedHeaders(ctx context.Context, sessionID string) context.Context { + if v, ok := s.capturedPassthroughHeaders.Load(sessionID); ok { + if headers, ok := v.(map[string]string); ok { + return headerforward.WithForwardedHeaders(ctx, headers) + } + } + return ctx +} + // enforceSessionBinding validates caller against the session's stored identity // binding. It is the SOLE identity-binding enforcement point on the Serve call // path: requests reach the core directly, bypassing the BindSession decorator that diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index e54e868696..951810f4af 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -291,6 +291,21 @@ type Server struct { // Nil if status reporting is disabled. statusReporter vmcpstatus.Reporter + // capturedPassthroughHeaders holds the per-session passthrough headers captured + // once at session-creation time on the Serve path (s.core != nil). It is keyed + // by session ID. Every entry is populated in handleSessionRegistrationImpl and + // cleaned up via the OnUnregisterSession hook (serve.go) or on registration + // failure. + // + // Security: these values may be credentials (e.g. API keys). They are stored + // node-local ONLY and are NEVER written to Redis or any shared state store. + // A session restored on another pod will have no captured headers — matching + // the legacy path's imperfect cross-pod behavior (the legacy per-session + // connection cannot be serialized either). This is acceptable: the alternative + // would require persisting credentials to shared state, which violates the + // security constraint in .claude/rules/security.md. + capturedPassthroughHeaders sync.Map // map[sessionID string → map[string]string] + // shutdownFuncs contains cleanup functions to run during Stop(). // Populated during Start() initialization before blocking; no mutex needed // since Stop() is only called after Start()'s select returns. @@ -1231,6 +1246,7 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv slog.Debug("creating session-scoped backends", "session_id", sessionID) // Defer cleanup: if any error occurs, terminate the session and log failures. + // Also remove any node-local captured passthrough headers stored for this session. defer func() { if retErr != nil { if _, termErr := s.vmcpSessionMgr.Terminate(sessionID); termErr != nil { @@ -1239,6 +1255,7 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv "error", termErr, "original_error", retErr) } + s.capturedPassthroughHeaders.Delete(sessionID) } }() @@ -1269,6 +1286,15 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv // returned error becomes retErr (named return), so the defer terminates the // session on failure. if s.core != nil { + // Capture passthrough headers from the session-creation request context + // (once per session). The snapshot is injected into every subsequent backend + // call for this session's lifetime by coreToolHandler / coreResourceHandler, + // ensuring the backend always sees the session-creation-time value even when + // the client changes the header on later requests (session-stable forwarding). + // Node-local only — never written to Redis (see capturedPassthroughHeaders). + if captured := headerforward.ForwardedHeadersFromContext(ctx); len(captured) > 0 { + s.capturedPassthroughHeaders.Store(sessionID, captured) + } return s.injectCoreSessionCapabilities(ctx, session) } diff --git a/pkg/vmcp/session/internal/backend/mcp_session.go b/pkg/vmcp/session/internal/backend/mcp_session.go index c70ba8a3c5..b8b5e3619a 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session.go +++ b/pkg/vmcp/session/internal/backend/mcp_session.go @@ -9,7 +9,6 @@ import ( "fmt" "io" "log/slog" - "maps" "net/http" "time" @@ -19,7 +18,6 @@ import ( "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/secrets" - "github.com/stacklok/toolhive/pkg/transport/middleware" "github.com/stacklok/toolhive/pkg/versions" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" @@ -525,64 +523,9 @@ func initAndQueryCapabilities( // mergeForwardedHeaders returns a HeaderForwardConfig that combines the static // backend configuration (base) with any per-request forwarded headers captured // from the caller's request (see headerforward.CaptureMiddleware). -// -// Rules (applied in order): -// 1. If forwarded is empty, base is returned unchanged (no allocation, same -// pointer). -// 2. A new HeaderForwardConfig is built from a shallow copy of base so the -// shared target.HeaderForward is never mutated. -// 3. Forwarded header names are canonicalized via http.CanonicalHeaderKey and -// checked against middleware.RestrictedHeaders; restricted names are silently -// dropped (defense-in-depth — they were already filtered upstream, but we -// guard here too). -// 4. A forwarded header name that also appears as a static header in base -// (AddPlaintextHeaders or AddHeadersFromSecret) is a misconfiguration: the -// function returns an error rather than silently picking a winner. +// Delegates to headerforward.MergeForwardedHeaders, which is the shared +// implementation used by both this connector and the shared backend client +// (pkg/vmcp/client). func mergeForwardedHeaders(base *vmcp.HeaderForwardConfig, forwarded map[string]string) (*vmcp.HeaderForwardConfig, error) { - if len(forwarded) == 0 { - return base, nil - } - - // Build the merged AddPlaintextHeaders map starting from the static config. - var staticPlaintext map[string]string - if base != nil { - staticPlaintext = base.AddPlaintextHeaders - } - - // Canonical set of header names already owned by the backend's static - // header-forward config (plaintext + secret), used for collision detection. - staticNames := make(map[string]struct{}) - if base != nil { - for k := range base.AddPlaintextHeaders { - staticNames[http.CanonicalHeaderKey(k)] = struct{}{} - } - for k := range base.AddHeadersFromSecret { - staticNames[http.CanonicalHeaderKey(k)] = struct{}{} - } - } - - merged := make(map[string]string, len(staticPlaintext)+len(forwarded)) - maps.Copy(merged, staticPlaintext) - - for name, value := range forwarded { - canonical := http.CanonicalHeaderKey(name) - if middleware.RestrictedHeaders[canonical] { - slog.Debug("Dropping restricted forwarded header", "header", canonical) - continue - } - // Fail loud on collision with a static header-forward value (rule 4). - if _, exists := staticNames[canonical]; exists { - return nil, fmt.Errorf( - "forwarded header %q collides with the backend's static header-forward config", canonical) - } - merged[canonical] = value - } - - out := &vmcp.HeaderForwardConfig{ - AddPlaintextHeaders: merged, - } - if base != nil { - out.AddHeadersFromSecret = base.AddHeadersFromSecret - } - return out, nil + return headerforward.MergeForwardedHeaders(base, forwarded) } diff --git a/test/integration/vmcp/helpers/vmcp_server.go b/test/integration/vmcp/helpers/vmcp_server.go index 4679b3a9aa..84efe2724b 100644 --- a/test/integration/vmcp/helpers/vmcp_server.go +++ b/test/integration/vmcp/helpers/vmcp_server.go @@ -16,13 +16,14 @@ import ( "github.com/stacklok/toolhive/pkg/telemetry" vmcptypes "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/aggregator" - "github.com/stacklok/toolhive/pkg/vmcp/auth/factory" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth/factory" authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" vmcpclient "github.com/stacklok/toolhive/pkg/vmcp/client" "github.com/stacklok/toolhive/pkg/vmcp/composer" - "github.com/stacklok/toolhive/pkg/vmcp/discovery" + vmcpcore "github.com/stacklok/toolhive/pkg/vmcp/core" "github.com/stacklok/toolhive/pkg/vmcp/router" vmcpserver "github.com/stacklok/toolhive/pkg/vmcp/server" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" ) @@ -102,7 +103,7 @@ func WithTelemetryProvider(provider *telemetry.Provider) VMCPServerOption { // matching the production wiring in pkg/vmcp/cli/serve.go. // // When this option is used, NewVMCPServer passes the header names to -// vmcpserver.Config.PassthroughHeaders, which installs +// vmcpserver.ServerConfig.PassthroughHeaders, which installs // headerforward.CaptureMiddleware so allowlisted headers are captured into the // request context and forwarded to backends. func WithPassthroughHeaders(headers ...string) VMCPServerOption { @@ -112,19 +113,13 @@ func WithPassthroughHeaders(headers ...string) VMCPServerOption { } // getFreePort returns an available TCP port on localhost. -// This is used for parallel test execution to avoid port conflicts. func getFreePort(tb testing.TB) int { tb.Helper() - // Listen on port 0 to get a random available port listener, err := net.Listen("tcp", "127.0.0.1:0") require.NoError(tb, err, "failed to get free port") - defer func() { - // Error ignored in test cleanup - _ = listener.Close() - }() + defer func() { _ = listener.Close() }() - // Extract the port number from the listener's address addr, ok := listener.Addr().(*net.TCPAddr) if !ok { tb.Fatalf("failed to get TCP address from listener") @@ -132,41 +127,33 @@ func getFreePort(tb testing.TB) int { return addr.Port } -// NewVMCPServer creates a vMCP server for testing with sensible defaults. -// The server is automatically started and will be ready when this function returns. +// NewVMCPServer creates a vMCP server for testing using the Serve path (core.New + +// Serve). The server is automatically started and ready when this function returns. // Use functional options to customize behavior. // -// Example: -// -// server := testkit.NewVMCPServer(ctx, t, backends, -// testkit.WithPrefixConflictResolution("{workload}_"), -// ) -// defer server.Shutdown(ctx) +// The Serve path is used (rather than the legacy server.New path) so that +// integration tests exercise the production code path that routes tool/resource +// calls through core.VMCP. This is required for testing session-stable +// passthrough header forwarding (#5560). func NewVMCPServer( ctx context.Context, tb testing.TB, backends []vmcptypes.Backend, opts ...VMCPServerOption, ) *vmcpserver.Server { tb.Helper() - // Default configuration config := &vmcpServerConfig{ conflictStrategy: "prefix", prefixFormat: "{workload}_", } - - // Apply options for _, opt := range opts { opt(config) } - // Create outgoing auth registry with all strategies registered - outgoingRegistry, err := factory.NewOutgoingAuthRegistry(ctx, &env.OSReader{}) + outgoingRegistry, err := vmcpauth.NewOutgoingAuthRegistry(ctx, &env.OSReader{}) require.NoError(tb, err) - // Create backend client backendClient, err := vmcpclient.NewHTTPBackendClient(outgoingRegistry) require.NoError(tb, err) - // Create conflict resolver based on strategy var conflictResolver aggregator.ConflictResolver switch config.conflictStrategy { case "prefix": @@ -175,55 +162,56 @@ func NewVMCPServer( conflictResolver = aggregator.NewPrefixConflictResolver(config.prefixFormat) } - // Create aggregator agg := aggregator.NewDefaultAggregator(backendClient, conflictResolver, nil, nil) - - // Create discovery manager - discoveryMgr, err := discovery.NewManager(agg) - require.NoError(tb, err) - - // Create router rtr := router.NewDefaultRouter() - - // Create immutable backend registry for tests (backends don't change during test execution) backendRegistry := vmcptypes.NewImmutableRegistry(backends) - // Create session factory with the same aggregator so tool names in the - // session routing table are consistent with the server's conflict-resolution - // strategy (e.g. prefix format applied by the aggregator). - sessionFactory := vmcpsession.NewSessionFactory(outgoingRegistry, vmcpsession.WithAggregator(agg)) - - // Create vMCP server with test-specific defaults. - // - // PassthroughHeaders wires headerforward.CaptureMiddleware inside the server - // (matching pkg/vmcp/cli/serve.go), so allowlisted headers are captured into - // the request context and forwarded to backends. Empty disables capture. - vmcpServer, err := vmcpserver.New(ctx, &vmcpserver.Config{ + // Build the core VMCP — the single authoritative aggregation on the Serve path. + coreVMCP, err := vmcpcore.New(&vmcpcore.Config{ + Aggregator: agg, + Router: rtr, + BackendRegistry: backendRegistry, + BackendClient: backendClient, + WorkflowDefs: config.workflowDefs, + TelemetryProvider: config.telemetryProvider, + }) + require.NoError(tb, err, "failed to create core VMCP") + + // The session factory must NOT use WithAggregator on the Serve path: the core + // is the single source of truth for capability aggregation. A factory with its + // own aggregator would produce a second, divergent capability set that the Serve + // path discards — exactly the double-aggregation AC2 forbids. + sessionFactory := vmcpsession.NewSessionFactory(outgoingRegistry) + + serverCfg := &vmcpserver.ServerConfig{ Name: "test-vmcp", Version: "1.0.0", Host: "127.0.0.1", - Port: getFreePort(tb), // Get a random available port for parallel test execution + Port: getFreePort(tb), + EndpointPath: "/mcp", + SessionTTL: 30 * time.Minute, AuthMiddleware: auth.AnonymousMiddleware, PassthroughHeaders: config.passthroughHeaders, + BackendRegistry: backendRegistry, TelemetryProvider: config.telemetryProvider, - SessionFactory: sessionFactory, - }, rtr, backendClient, discoveryMgr, backendRegistry, config.workflowDefs) + SessionManagerConfig: &sessionmanager.FactoryConfig{ + Base: sessionFactory, + }, + } + + vmcpServer, err := vmcpserver.Serve(ctx, coreVMCP, serverCfg) require.NoError(tb, err, "failed to create vMCP server") - // Start server automatically - // Use the passed-in context to ensure proper cancellation propagation go func() { if err := vmcpServer.Start(ctx); err != nil { select { case <-ctx.Done(): - // Context cancelled, ignore error default: tb.Errorf("vMCP server error: %v", err) } } }() - // Wait for server to be ready (with 5 second timeout) select { case <-vmcpServer.Ready(): tb.Logf("vMCP server ready at: http://%s/mcp", vmcpServer.Address()) From 55d67a25be42a4dce58a1df5e058cf7571d5ca35 Mon Sep 17 00:00:00 2001 From: Trey Date: Thu, 18 Jun 2026 10:12:39 -0700 Subject: [PATCH 2/7] Address code review feedback for #5560 - Fix misleading TTL cleanup comment in serve.go: the OnUnregisterSession hook fires actively via the SDK's sessionIdleTTL sweeper, not lazily via sync.Map eviction (sync.Map has no eviction mechanism) - Canonicalize static plaintext keys in MergeForwardedHeaders so the returned map has uniformly canonical HTTP header names regardless of how the caller formatted the static config - Remove duplicate TestMergeForwardedHeaders from the session/internal/ backend package; the thin mergeForwardedHeaders wrapper delegates to headerforward.MergeForwardedHeaders whose full test matrix already lives in pkg/vmcp/headerforward/transport_test.go Co-Authored-By: Claude Sonnet 4.6 (1M context) --- pkg/vmcp/headerforward/transport.go | 5 +- pkg/vmcp/server/serve.go | 11 +- .../internal/backend/mcp_session_test.go | 144 +----------------- 3 files changed, 14 insertions(+), 146 deletions(-) diff --git a/pkg/vmcp/headerforward/transport.go b/pkg/vmcp/headerforward/transport.go index 1ff3b10277..ba4def8eb4 100644 --- a/pkg/vmcp/headerforward/transport.go +++ b/pkg/vmcp/headerforward/transport.go @@ -141,7 +141,10 @@ func MergeForwardedHeaders(base *vmcp.HeaderForwardConfig, forwarded map[string] merged := make(map[string]string, len(staticPlaintext)+len(forwarded)) for k, v := range staticPlaintext { - merged[k] = v + // Canonicalize static keys (consistent with forwarded keys below) so the + // returned map has uniformly canonical HTTP header names regardless of how + // the static config was specified. + merged[http.CanonicalHeaderKey(k)] = v } for name, value := range forwarded { diff --git a/pkg/vmcp/server/serve.go b/pkg/vmcp/server/serve.go index 59e1b50158..9d7d860d20 100644 --- a/pkg/vmcp/server/serve.go +++ b/pkg/vmcp/server/serve.go @@ -306,10 +306,13 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) srv.lazyInjectSessionTools(hookCtx) }) // Clean up node-local captured passthrough headers when the SDK unregisters the - // session (on client DELETE or SDK-internal timeout). This is the primary cleanup - // path for explicitly-terminated sessions; TTL-expired sessions are cleaned up - // lazily when the node-local cache evicts them. Node-local only — credentials - // must not be persisted to shared state (security.md). + // session. The hook fires for both explicit client DELETE requests and for sessions + // that expire: the SDK's sessionIdleTTL sweeper calls UnregisterSession on idle + // sessions, which triggers this hook. Sessions abandoned without DELETE (e.g. a + // crashed client with a sessionIdleTTL of 0) retain their entry until the Server + // shuts down — acceptable because the map is bounded by the number of live sessions + // and each entry is small (a few header name→value pairs). Credentials in the map + // are node-local only and never written to shared state (security.md). hooks.AddOnUnregisterSession(func(_ context.Context, clientSess server.ClientSession) { srv.capturedPassthroughHeaders.Delete(clientSess.SessionID()) }) diff --git a/pkg/vmcp/session/internal/backend/mcp_session_test.go b/pkg/vmcp/session/internal/backend/mcp_session_test.go index 26c5ee18c2..79ce6ba721 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session_test.go +++ b/pkg/vmcp/session/internal/backend/mcp_session_test.go @@ -5,7 +5,6 @@ package backend import ( "context" - "maps" "testing" "github.com/stretchr/testify/assert" @@ -28,146 +27,9 @@ func newTestRegistry(t *testing.T) vmcpauth.OutgoingAuthRegistry { return reg } -func TestMergeForwardedHeaders(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - base *vmcp.HeaderForwardConfig - // forwarded is the per-request captured header map, as returned by - // headerforward.ForwardedHeadersFromContext. - forwarded map[string]string - // wantSameAsBase means we expect the returned pointer to equal the base - // pointer (i.e., no merge was needed — original passed through unchanged). - wantSameAsBase bool - wantHeaders map[string]string - // wantErr means mergeForwardedHeaders should return an error — a forwarded - // header name collides with a static header-forward value (plaintext or - // secret) on the backend. - wantErr bool - }{ - { - name: "nil forwarded returns base unchanged", - base: &vmcp.HeaderForwardConfig{AddPlaintextHeaders: map[string]string{"X-Static": "static"}}, - forwarded: nil, - wantSameAsBase: true, - }, - { - name: "nil forwarded nil base returns nil", - base: nil, - forwarded: nil, - wantSameAsBase: true, - }, - { - // Both nil and empty forwarded maps satisfy len==0 and return base unchanged. - name: "empty forwarded returns base unchanged", - base: &vmcp.HeaderForwardConfig{AddPlaintextHeaders: map[string]string{"X-Static": "static"}}, - forwarded: map[string]string{}, - wantSameAsBase: true, - }, - { - name: "forwarded header added to nil base", - base: nil, - forwarded: map[string]string{"X-Litellm-Api-Key": "sk-1"}, - wantHeaders: map[string]string{"X-Litellm-Api-Key": "sk-1"}, - }, - { - name: "forwarded header added to base with other header", - base: &vmcp.HeaderForwardConfig{ - AddPlaintextHeaders: map[string]string{"X-Static": "static-value"}, - }, - forwarded: map[string]string{"X-Litellm-Api-Key": "sk-1"}, - wantHeaders: map[string]string{ - "X-Static": "static-value", - "X-Litellm-Api-Key": "sk-1", - }, - }, - { - // Covers both a canonical restricted header (Host) and a lowercase one - // (x-forwarded-for) to verify case-insensitive matching in a single case. - name: "restricted headers dropped (canonical and lowercase)", - base: nil, - forwarded: map[string]string{ - "Host": "evil.example.com", - "x-forwarded-for": "1.2.3.4", - "X-Litellm-Api-Key": "sk-2", - }, - wantHeaders: map[string]string{"X-Litellm-Api-Key": "sk-2"}, - }, - { - name: "forwarded header colliding with static plaintext returns error", - base: &vmcp.HeaderForwardConfig{ - AddPlaintextHeaders: map[string]string{"X-Litellm-Api-Key": "static-value"}, - }, - forwarded: map[string]string{"X-Litellm-Api-Key": "forwarded-value"}, - wantErr: true, - }, - { - name: "forwarded header colliding with static secret returns error", - base: &vmcp.HeaderForwardConfig{ - AddHeadersFromSecret: map[string]string{"X-Litellm-Api-Key": "my-secret-id"}, - }, - forwarded: map[string]string{"X-Litellm-Api-Key": "sk-5"}, - wantErr: true, - }, - { - name: "base AddHeadersFromSecret preserved unchanged", - base: &vmcp.HeaderForwardConfig{ - AddHeadersFromSecret: map[string]string{"X-Secret-Header": "my-secret-id"}, - }, - forwarded: map[string]string{"X-Litellm-Api-Key": "sk-4"}, - wantHeaders: map[string]string{"X-Litellm-Api-Key": "sk-4"}, - }, - { - name: "base not mutated when forwarded headers added", - base: &vmcp.HeaderForwardConfig{ - AddPlaintextHeaders: map[string]string{"X-Static": "original"}, - }, - forwarded: map[string]string{"X-New": "new-value"}, - // base should not gain X-New - wantHeaders: map[string]string{ - "X-Static": "original", - "X-New": "new-value", - }, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - // Snapshot the original base plaintext headers to verify they are not mutated. - var origBasePlaintext map[string]string - if tc.base != nil { - origBasePlaintext = maps.Clone(tc.base.AddPlaintextHeaders) - } - - got, err := mergeForwardedHeaders(tc.base, tc.forwarded) - - if tc.wantErr { - require.Error(t, err) - assert.Nil(t, got) - return - } - require.NoError(t, err) - - if tc.wantSameAsBase { - assert.Same(t, tc.base, got, "expected the original base pointer to be returned unchanged") - return - } - - require.NotNil(t, got) - assert.Equal(t, tc.wantHeaders, got.AddPlaintextHeaders, - "merged AddPlaintextHeaders mismatch (checks both presence and absence of keys)") - - // Verify base was not mutated. - if tc.base != nil { - assert.Equal(t, origBasePlaintext, tc.base.AddPlaintextHeaders, - "base.AddPlaintextHeaders must not be mutated") - } - }) - } -} +// mergeForwardedHeaders is now a one-line delegation to +// headerforward.MergeForwardedHeaders; full coverage lives in +// pkg/vmcp/headerforward/transport_test.go (TestMergeForwardedHeaders). func TestCreateMCPClient_UnsupportedTransport(t *testing.T) { t.Parallel() From f18835fb2ded77aac062b817d66a7bfc10b0b5e8 Mon Sep 17 00:00:00 2001 From: Trey Date: Thu, 18 Jun 2026 11:13:26 -0700 Subject: [PATCH 3/7] Fix defensive copy, document anti-pattern #1 exception, clarify comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses #5561 review comments: - MEDIUM pkg/vmcp/server/server.go (3437978067): clone captured headers before storing; ForwardedHeadersFromContext returns the context map reference and structurally enforcing immutability is important on a credential path - MEDIUM pkg/vmcp/server/serve_handlers.go (3437978071): document why using context as the forwarding channel is an intentional anti-pattern #1 exception (SDK controls the handler context; backendClient has no explicit per-session header slot) - LOW pkg/vmcp/server/server.go (3437978080): clarify the double-delete timing invariant in the registration-failure defer — Store runs only after CreateSession succeeds, and the defer fires before the session is live, so the two delete paths are mutually exclusive for well-formed sessions - LOW pkg/vmcp/headerforward/transport.go (3437978095): document in MergeForwardedHeaders docstring that static plaintext keys are now canonicalized in the output (behaviorally identical at the wire level, but the intermediate map representation changed from the old maps.Copy path) Co-Authored-By: Claude Sonnet 4.6 (1M context) --- pkg/vmcp/headerforward/transport.go | 14 ++++++++++---- pkg/vmcp/server/serve_handlers.go | 8 ++++++++ pkg/vmcp/server/server.go | 15 ++++++++++++++- 3 files changed, 32 insertions(+), 5 deletions(-) diff --git a/pkg/vmcp/headerforward/transport.go b/pkg/vmcp/headerforward/transport.go index ba4def8eb4..cf2c9cfcb3 100644 --- a/pkg/vmcp/headerforward/transport.go +++ b/pkg/vmcp/headerforward/transport.go @@ -109,11 +109,17 @@ func BuildHeaderForwardTripper( // 1. If forwarded is empty, base is returned unchanged (no allocation, same pointer). // 2. A new HeaderForwardConfig is built from a shallow copy of base so the // shared target.HeaderForward is never mutated. -// 3. Forwarded header names are canonicalized via http.CanonicalHeaderKey and -// checked against middleware.RestrictedHeaders; restricted names are silently -// dropped (defense-in-depth — they were already filtered upstream by +// 3. All output AddPlaintextHeaders keys (both static and forwarded) are +// canonicalized via http.CanonicalHeaderKey. Static keys from +// base.AddPlaintextHeaders are therefore also re-canonicalized in the output, +// which changes the intermediate map representation (though wire behavior is +// identical since resolveHeaderForward / BuildHeaderForwardTripper +// canonicalize again before writing to http.Header). +// 4. Forwarded header names are additionally checked against +// middleware.RestrictedHeaders; restricted names are silently dropped +// (defense-in-depth — they were already filtered upstream by // CaptureMiddleware, but we guard here too). -// 4. A forwarded header name that also appears in base (AddPlaintextHeaders or +// 5. A forwarded header name that also appears in base (AddPlaintextHeaders or // AddHeadersFromSecret) is a misconfiguration: the function returns an error // rather than silently picking a winner. func MergeForwardedHeaders(base *vmcp.HeaderForwardConfig, forwarded map[string]string) (*vmcp.HeaderForwardConfig, error) { diff --git a/pkg/vmcp/server/serve_handlers.go b/pkg/vmcp/server/serve_handlers.go index 8bc3f4371b..5cd73cdf28 100644 --- a/pkg/vmcp/server/serve_handlers.go +++ b/pkg/vmcp/server/serve_handlers.go @@ -200,6 +200,14 @@ func (s *Server) coreToolHandler(sessionID, toolName, backendName string) server // how the client changes the header on later requests. The shared backend // client (pkg/vmcp/client) reads the forwarded headers from ctx to build the // outgoing transport chain. + // + // Note on vmcp anti-pattern #1 ("reserve context for trace IDs, cancellation, + // and deadlines only"): this use of context is a known, intentional exception. + // The MCP SDK controls the request context passed to the tool handler, and the + // shared backend client factory (pkg/vmcp/client) has no explicit parameter + // slot for per-session headers — forwarding through context is the only + // coupling-free path that does not require threading session state through the + // core.VMCP → backendClient boundary. ctx = s.injectCapturedHeaders(ctx, sessionID) result, err := s.core.CallTool(ctx, caller, toolName, args, conversion.FromMCPMeta(req.Params.Meta)) diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index 951810f4af..d0b8361009 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -14,6 +14,7 @@ import ( "errors" "fmt" "log/slog" + "maps" "net" "net/http" "os" @@ -1247,6 +1248,11 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv // Defer cleanup: if any error occurs, terminate the session and log failures. // Also remove any node-local captured passthrough headers stored for this session. + // + // Double-delete safety: capturedPassthroughHeaders.Delete is idempotent. This + // defer fires only on registration failure — before the session is live and before + // the OnUnregisterSession hook (serve.go) is ever registered for this session ID. + // The two delete paths are therefore mutually exclusive for well-formed sessions. defer func() { if retErr != nil { if _, termErr := s.vmcpSessionMgr.Terminate(sessionID); termErr != nil { @@ -1255,6 +1261,9 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv "error", termErr, "original_error", retErr) } + // The Store above runs only after CreateSession succeeds; if we reach + // this defer, the Store either never ran (CreateSession failed) or ran + // but injectCoreSessionCapabilities failed. Either way, Delete is safe. s.capturedPassthroughHeaders.Delete(sessionID) } }() @@ -1293,7 +1302,11 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv // the client changes the header on later requests (session-stable forwarding). // Node-local only — never written to Redis (see capturedPassthroughHeaders). if captured := headerforward.ForwardedHeadersFromContext(ctx); len(captured) > 0 { - s.capturedPassthroughHeaders.Store(sessionID, captured) + // Defensive copy: ForwardedHeadersFromContext returns the map reference + // placed in the context by CaptureMiddleware. Cloning before storing + // ensures the session's credential snapshot is structurally immutable + // regardless of whether the caller mutates its own map later. + s.capturedPassthroughHeaders.Store(sessionID, maps.Clone(captured)) } return s.injectCoreSessionCapabilities(ctx, session) } From 4ae5e5205d16f3f464ac4a68a2056b33ebd4d271 Mon Sep 17 00:00:00 2001 From: Trey Date: Thu, 18 Jun 2026 11:17:44 -0700 Subject: [PATCH 4/7] Add unit test coverage for new session-header path Addresses #5561 review comments: - LOW pkg/vmcp/headerforward/transport_test.go (3437978087): add per-case base-mutation guard to TestMergeForwardedHeaders; verifies that MergeForwardedHeaders never writes into the caller's AddPlaintextHeaders map - LOW pkg/vmcp/server/serve_handlers_test.go (3437978101): new file with TestInjectCapturedHeaders covering the hit, miss, and wrong-type-fallback branches of injectCapturedHeaders - LOW pkg/vmcp/server/serve_handlers_test.go (3437978106): TestOnUnregisterSession HookDeletesCapturedHeaders fires the hook closure via the SDK's UnregisterSession and asserts the capturedPassthroughHeaders entry is deleted - LOW pkg/vmcp/core/core_telemetry_test.go (3437978110): new file with TestTelemetryComposer_{Success,Error,DelegatesNonExecuteMethods} using an in-memory OTEL SDK reader to verify counter increments, histogram recording, and error-counter/span-error on failure Co-Authored-By: Claude Sonnet 4.6 (1M context) --- pkg/vmcp/core/core_telemetry_test.go | 171 +++++++++++++++++++++++ pkg/vmcp/headerforward/transport_test.go | 15 ++ pkg/vmcp/server/serve_handlers_test.go | 87 ++++++++++++ 3 files changed, 273 insertions(+) create mode 100644 pkg/vmcp/core/core_telemetry_test.go create mode 100644 pkg/vmcp/server/serve_handlers_test.go diff --git a/pkg/vmcp/core/core_telemetry_test.go b/pkg/vmcp/core/core_telemetry_test.go new file mode 100644 index 0000000000..1899d38062 --- /dev/null +++ b/pkg/vmcp/core/core_telemetry_test.go @@ -0,0 +1,171 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package core + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" + tracesdk "go.opentelemetry.io/otel/sdk/trace" + tracenoop "go.opentelemetry.io/otel/trace/noop" + + "github.com/stacklok/toolhive/pkg/vmcp/composer" +) + +// stubComposer is already declared in core_calls_test.go (same package). +// Its fields are: result *composer.WorkflowResult, err error. + +// newTestInstruments creates a workflowInstruments backed by an in-memory OTEL SDK +// and returns the instruments plus a ManualReader for metric assertions. +// The metric names match production names exactly so the assertions mirror what +// Prometheus would expose. +func newTestInstruments(t *testing.T) (*workflowInstruments, *sdkmetric.ManualReader) { + t.Helper() + reader := sdkmetric.NewManualReader() + mp := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + tp := tracesdk.NewTracerProvider() + meter := mp.Meter(instrumentationName) + + executions, err := meter.Int64Counter("toolhive_vmcp_workflow_executions") + require.NoError(t, err) + errs, err := meter.Int64Counter("toolhive_vmcp_workflow_errors") + require.NoError(t, err) + duration, err := meter.Float64Histogram("toolhive_vmcp_workflow_duration") + require.NoError(t, err) + + return &workflowInstruments{ + tracer: tp.Tracer(instrumentationName), + executionsTotal: executions, + errorsTotal: errs, + executionDuration: duration, + }, reader +} + +// collectMetrics gathers all metrics from the reader into a snapshot. +func collectMetrics(t *testing.T, reader *sdkmetric.ManualReader) metricdata.ResourceMetrics { + t.Helper() + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(context.Background(), &rm)) + return rm +} + +// findMetricByName returns the first metric with the given name, or nil. +func findMetricByName(rm metricdata.ResourceMetrics, name string) *metricdata.Metrics { + for _, sm := range rm.ScopeMetrics { + for i := range sm.Metrics { + if sm.Metrics[i].Name == name { + return &sm.Metrics[i] + } + } + } + return nil +} + +func int64CounterValue(m *metricdata.Metrics) int64 { + if m == nil { + return 0 + } + s, ok := m.Data.(metricdata.Sum[int64]) + if !ok { + return 0 + } + var total int64 + for _, dp := range s.DataPoints { + total += dp.Value + } + return total +} + +func float64HistogramCount(m *metricdata.Metrics) uint64 { + if m == nil { + return 0 + } + h, ok := m.Data.(metricdata.Histogram[float64]) + if !ok { + return 0 + } + var total uint64 + for _, dp := range h.DataPoints { + total += dp.Count + } + return total +} + +// TestTelemetryComposer_Success verifies that on a successful ExecuteWorkflow call: +// - the executions counter increments by 1 +// - the error counter stays at 0 +// - the duration histogram records 1 observation +// - the result from the inner composer is returned unchanged +func TestTelemetryComposer_Success(t *testing.T) { + t.Parallel() + + instruments, reader := newTestInstruments(t) + want := &composer.WorkflowResult{Output: map[string]any{"key": "val"}} + tc := &telemetryComposer{ + base: stubComposer{result: want}, + instruments: instruments, + } + + def := &composer.WorkflowDefinition{Name: "test-workflow"} + got, err := tc.ExecuteWorkflow(context.Background(), def, nil) + require.NoError(t, err) + assert.Equal(t, want, got) + + rm := collectMetrics(t, reader) + assert.Equal(t, int64(1), int64CounterValue(findMetricByName(rm, "toolhive_vmcp_workflow_executions")), + "executions counter must increment on success") + assert.Equal(t, int64(0), int64CounterValue(findMetricByName(rm, "toolhive_vmcp_workflow_errors")), + "errors counter must remain zero on success") + assert.Equal(t, uint64(1), float64HistogramCount(findMetricByName(rm, "toolhive_vmcp_workflow_duration")), + "duration histogram must record exactly one observation") +} + +// TestTelemetryComposer_Error verifies that on a failed ExecuteWorkflow call: +// - the executions counter still increments by 1 +// - the error counter also increments by 1 +// - the duration histogram records 1 observation +// - the error from the inner composer is propagated +func TestTelemetryComposer_Error(t *testing.T) { + t.Parallel() + + instruments, reader := newTestInstruments(t) + boom := errors.New("backend exploded") + tc := &telemetryComposer{ + base: stubComposer{err: boom}, + instruments: instruments, + } + + def := &composer.WorkflowDefinition{Name: "failing-workflow"} + _, err := tc.ExecuteWorkflow(context.Background(), def, nil) + require.ErrorIs(t, err, boom) + + rm := collectMetrics(t, reader) + assert.Equal(t, int64(1), int64CounterValue(findMetricByName(rm, "toolhive_vmcp_workflow_executions")), + "executions counter must increment even on failure") + assert.Equal(t, int64(1), int64CounterValue(findMetricByName(rm, "toolhive_vmcp_workflow_errors")), + "errors counter must increment on failure") + assert.Equal(t, uint64(1), float64HistogramCount(findMetricByName(rm, "toolhive_vmcp_workflow_duration")), + "duration histogram must record one observation even on failure") +} + +// TestTelemetryComposer_DelegatesNonExecuteMethods verifies that ValidateWorkflow, +// GetWorkflowStatus, and CancelWorkflow delegate to the base without instrumentation. +func TestTelemetryComposer_DelegatesNonExecuteMethods(t *testing.T) { + t.Parallel() + + tc := &telemetryComposer{ + base: stubComposer{}, + instruments: &workflowInstruments{tracer: tracenoop.Tracer{}}, + } + + require.NoError(t, tc.ValidateWorkflow(context.Background(), &composer.WorkflowDefinition{})) + _, err := tc.GetWorkflowStatus(context.Background(), "any-id") + require.NoError(t, err) + require.NoError(t, tc.CancelWorkflow(context.Background(), "any-id")) +} diff --git a/pkg/vmcp/headerforward/transport_test.go b/pkg/vmcp/headerforward/transport_test.go index f8be89c8d0..c8f2b42766 100644 --- a/pkg/vmcp/headerforward/transport_test.go +++ b/pkg/vmcp/headerforward/transport_test.go @@ -337,6 +337,15 @@ func TestMergeForwardedHeaders(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() + // Snapshot base before the call to catch any mutation. + var origBasePlaintext map[string]string + if tc.base != nil && tc.base.AddPlaintextHeaders != nil { + origBasePlaintext = make(map[string]string, len(tc.base.AddPlaintextHeaders)) + for k, v := range tc.base.AddPlaintextHeaders { + origBasePlaintext[k] = v + } + } + got, err := MergeForwardedHeaders(tc.base, tc.forwarded) if tc.wantErr { require.Error(t, err) @@ -344,6 +353,12 @@ func TestMergeForwardedHeaders(t *testing.T) { } require.NoError(t, err) + // Verify base was not mutated regardless of outcome. + if tc.base != nil { + assert.Equal(t, origBasePlaintext, tc.base.AddPlaintextHeaders, + "base.AddPlaintextHeaders must not be mutated by MergeForwardedHeaders") + } + if tc.wantCfg == nil { assert.Nil(t, got) return diff --git a/pkg/vmcp/server/serve_handlers_test.go b/pkg/vmcp/server/serve_handlers_test.go new file mode 100644 index 0000000000..d65f46e8ef --- /dev/null +++ b/pkg/vmcp/server/serve_handlers_test.go @@ -0,0 +1,87 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "testing" + + "github.com/mark3labs/mcp-go/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp/headerforward" +) + +// TestInjectCapturedHeaders exercises all three branches of injectCapturedHeaders: +// a map hit, a map miss, and a wrong-type entry (type-assertion fallback). +func TestInjectCapturedHeaders(t *testing.T) { + t.Parallel() + + t.Run("hit: stored headers are injected into ctx", func(t *testing.T) { + t.Parallel() + s := &Server{} + s.capturedPassthroughHeaders.Store("s1", map[string]string{"X-Api-Key": "key-123"}) + ctx := s.injectCapturedHeaders(context.Background(), "s1") + got := headerforward.ForwardedHeadersFromContext(ctx) + assert.Equal(t, map[string]string{"X-Api-Key": "key-123"}, got, + "stored headers must be injected into the context") + }) + + t.Run("miss: absent session ID leaves ctx unchanged", func(t *testing.T) { + t.Parallel() + s := &Server{} + ctx := s.injectCapturedHeaders(context.Background(), "missing") + assert.Nil(t, headerforward.ForwardedHeadersFromContext(ctx), + "a missing session ID must leave the context's forwarded headers unchanged") + }) + + t.Run("type fallback: wrong type stored silently degrades to no-op", func(t *testing.T) { + t.Parallel() + s := &Server{} + s.capturedPassthroughHeaders.Store("s2", "not-a-map") + ctx := s.injectCapturedHeaders(context.Background(), "s2") + assert.Nil(t, headerforward.ForwardedHeadersFromContext(ctx), + "a non-map entry must not be injected (type-assertion fallback)") + }) +} + +// stubClientSession is a minimal server.ClientSession that returns a fixed ID. +// Used to exercise the OnUnregisterSession hook closure without running a full server. +type stubClientSession struct { + server.ClientSession // embed for unused interface methods + id string +} + +func (s *stubClientSession) SessionID() string { return s.id } + +// TestOnUnregisterSessionHookDeletesCapturedHeaders verifies that the closure +// registered with AddOnUnregisterSession in serve.go deletes the session's entry +// from capturedPassthroughHeaders. The test constructs a minimal *Server with a +// pre-populated entry and fires the hook closure directly, avoiding the need to +// bring up an HTTP server. +func TestOnUnregisterSessionHookDeletesCapturedHeaders(t *testing.T) { + t.Parallel() + + srv, err := Serve(context.Background(), &stubVMCP{}, testMinimalServeConfig()) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + const sessionID = "test-session-hook" + srv.capturedPassthroughHeaders.Store(sessionID, map[string]string{"X-Key": "v"}) + + // Confirm the entry is present before firing. + _, ok := srv.capturedPassthroughHeaders.Load(sessionID) + require.True(t, ok, "entry must be present before hook fires") + + // Fire the hook by calling UnregisterSession on the mcp-go hooks object. + // The hooks object is embedded in mcpServer; UnregisterSession triggers all + // AddOnUnregisterSession callbacks registered during Serve, including our cleanup. + srv.mcpServer.GetHooks().UnregisterSession(context.Background(), &stubClientSession{id: sessionID}) + + // After the hook fires the entry must be gone. + _, stillPresent := srv.capturedPassthroughHeaders.Load(sessionID) + assert.False(t, stillPresent, + "capturedPassthroughHeaders entry must be deleted by OnUnregisterSession hook") +} From f08f3d041233297185ee2cc193ae682e456f9ade Mon Sep 17 00:00:00 2001 From: Trey Date: Mon, 22 Jun 2026 07:59:41 -0700 Subject: [PATCH 5/7] Switch passthrough headers to per-request forwarding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback on #5561: session-stability was a side-effect of the legacy per-session connection, not an intended requirement. Per-request forwarding (each call reads the current incoming header value) is the correct semantic, and the client.go MergeForwardedHeaders call already achieves it. - Remove capturedPassthroughHeaders sync.Map and the session-creation capture logic from Server — the shared backend client reads ForwardedHeadersFromContext on every call, providing per-request forwarding without any extra machinery - Remove injectCapturedHeaders and its ctx replacement from coreToolHandler / coreResourceHandler - Remove the AddOnUnregisterSession cleanup hook from serve.go - Rewrite TestVMCPServer_PassthroughHeaders: the second assertion now checks that the CHANGED header value reaches the backend (per-request semantics), not that the original value is frozen Co-Authored-By: Claude Sonnet 4.6 (1M context) --- pkg/vmcp/server/serve.go | 17 +--- pkg/vmcp/server/serve_handlers.go | 43 +-------- pkg/vmcp/server/serve_handlers_test.go | 87 ------------------- pkg/vmcp/server/server.go | 39 --------- .../vmcp/passthrough_headers_test.go | 27 +++--- 5 files changed, 17 insertions(+), 196 deletions(-) delete mode 100644 pkg/vmcp/server/serve_handlers_test.go diff --git a/pkg/vmcp/server/serve.go b/pkg/vmcp/server/serve.go index 9d7d860d20..5d624f19b8 100644 --- a/pkg/vmcp/server/serve.go +++ b/pkg/vmcp/server/serve.go @@ -292,10 +292,10 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) return v.Close() }) - // Register the SDK hooks against the assembled *Server (relocated from - // server.New). They drive phase-2 of two-phase creation (OnRegisterSession), + // Register the three SDK hooks against the assembled *Server (relocated from + // server.New). They drive phase-2 of two-phase creation (OnRegisterSession) and // the cross-pod Redis re-hydration of per-session tools (OnBeforeListTools / - // OnBeforeCallTool), and node-local captured-header cleanup (OnUnregisterSession). + // OnBeforeCallTool); the *Server receiver methods are unchanged. hooks.AddOnRegisterSession(func(hookCtx context.Context, session server.ClientSession) { srv.handleSessionRegistration(hookCtx, session) }) @@ -305,17 +305,6 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) hooks.AddBeforeCallTool(func(hookCtx context.Context, _ any, _ *mcp.CallToolRequest) { srv.lazyInjectSessionTools(hookCtx) }) - // Clean up node-local captured passthrough headers when the SDK unregisters the - // session. The hook fires for both explicit client DELETE requests and for sessions - // that expire: the SDK's sessionIdleTTL sweeper calls UnregisterSession on idle - // sessions, which triggers this hook. Sessions abandoned without DELETE (e.g. a - // crashed client with a sessionIdleTTL of 0) retain their entry until the Server - // shuts down — acceptable because the map is bounded by the number of live sessions - // and each entry is small (a few header name→value pairs). Credentials in the map - // are node-local only and never written to shared state (security.md). - hooks.AddOnUnregisterSession(func(_ context.Context, clientSess server.ClientSession) { - srv.capturedPassthroughHeaders.Delete(clientSess.SessionID()) - }) // Disarm the close-on-error guard: the Server is fully constructed. closeStorageOnErr = false diff --git a/pkg/vmcp/server/serve_handlers.go b/pkg/vmcp/server/serve_handlers.go index 5cd73cdf28..ad5bad3949 100644 --- a/pkg/vmcp/server/serve_handlers.go +++ b/pkg/vmcp/server/serve_handlers.go @@ -17,7 +17,6 @@ import ( "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" - "github.com/stacklok/toolhive/pkg/vmcp/headerforward" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types" ) @@ -194,22 +193,6 @@ func (s *Server) coreToolHandler(sessionID, toolName, backendName string) server return mcp.NewToolResultError(fmt.Sprintf("Unauthorized: %v", err)), nil } - // Replace the per-request forwarded headers on ctx with the session-stable - // snapshot captured once at session creation. This ensures the backend always - // receives the value present when the session was established, regardless of - // how the client changes the header on later requests. The shared backend - // client (pkg/vmcp/client) reads the forwarded headers from ctx to build the - // outgoing transport chain. - // - // Note on vmcp anti-pattern #1 ("reserve context for trace IDs, cancellation, - // and deadlines only"): this use of context is a known, intentional exception. - // The MCP SDK controls the request context passed to the tool handler, and the - // shared backend client factory (pkg/vmcp/client) has no explicit parameter - // slot for per-session headers — forwarding through context is the only - // coupling-free path that does not require threading session state through the - // core.VMCP → backendClient boundary. - ctx = s.injectCapturedHeaders(ctx, sessionID) - result, err := s.core.CallTool(ctx, caller, toolName, args, conversion.FromMCPMeta(req.Params.Meta)) if err != nil { // Admission denial returns a generic message so the underlying authorizer @@ -230,8 +213,7 @@ func (s *Server) coreToolHandler(sessionID, toolName, backendName string) server } // coreResourceHandler builds the SDK handler for a Serve-path resource. It mirrors -// coreToolHandler: audit label, binding check, session-stable header injection, then -// core.ReadResource with explicit identity. +// coreToolHandler: audit label, binding check, then core.ReadResource with explicit identity. func (s *Server) coreResourceHandler( sessionID, uri, backendName string, ) func(ctx context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { @@ -246,10 +228,6 @@ func (s *Server) coreResourceHandler( return nil, fmt.Errorf("unauthorized: %w", err) } - // Replace the per-request forwarded headers on ctx with the session-stable - // snapshot captured once at session creation (mirrors coreToolHandler). - ctx = s.injectCapturedHeaders(ctx, sessionID) - result, err := s.core.ReadResource(ctx, caller, uri) if err != nil { if errors.Is(err, vmcp.ErrAuthorizationFailed) { @@ -261,25 +239,6 @@ func (s *Server) coreResourceHandler( } } -// injectCapturedHeaders replaces the per-request forwarded headers on ctx with -// the session-stable snapshot captured at session creation (stored node-locally in -// capturedPassthroughHeaders). If no snapshot exists for this session (e.g. no -// passthrough headers were configured, or the session was created without any -// allowlisted headers present), ctx is returned unchanged. -// -// This is the enforcement point for session-stable header forwarding on the Serve -// path: coreToolHandler and coreResourceHandler call this before delegating to -// core.CallTool / core.ReadResource, so the shared backend client always sees the -// session-creation-time value, not a mid-session change. -func (s *Server) injectCapturedHeaders(ctx context.Context, sessionID string) context.Context { - if v, ok := s.capturedPassthroughHeaders.Load(sessionID); ok { - if headers, ok := v.(map[string]string); ok { - return headerforward.WithForwardedHeaders(ctx, headers) - } - } - return ctx -} - // enforceSessionBinding validates caller against the session's stored identity // binding. It is the SOLE identity-binding enforcement point on the Serve call // path: requests reach the core directly, bypassing the BindSession decorator that diff --git a/pkg/vmcp/server/serve_handlers_test.go b/pkg/vmcp/server/serve_handlers_test.go deleted file mode 100644 index d65f46e8ef..0000000000 --- a/pkg/vmcp/server/serve_handlers_test.go +++ /dev/null @@ -1,87 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "context" - "testing" - - "github.com/mark3labs/mcp-go/server" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/stacklok/toolhive/pkg/vmcp/headerforward" -) - -// TestInjectCapturedHeaders exercises all three branches of injectCapturedHeaders: -// a map hit, a map miss, and a wrong-type entry (type-assertion fallback). -func TestInjectCapturedHeaders(t *testing.T) { - t.Parallel() - - t.Run("hit: stored headers are injected into ctx", func(t *testing.T) { - t.Parallel() - s := &Server{} - s.capturedPassthroughHeaders.Store("s1", map[string]string{"X-Api-Key": "key-123"}) - ctx := s.injectCapturedHeaders(context.Background(), "s1") - got := headerforward.ForwardedHeadersFromContext(ctx) - assert.Equal(t, map[string]string{"X-Api-Key": "key-123"}, got, - "stored headers must be injected into the context") - }) - - t.Run("miss: absent session ID leaves ctx unchanged", func(t *testing.T) { - t.Parallel() - s := &Server{} - ctx := s.injectCapturedHeaders(context.Background(), "missing") - assert.Nil(t, headerforward.ForwardedHeadersFromContext(ctx), - "a missing session ID must leave the context's forwarded headers unchanged") - }) - - t.Run("type fallback: wrong type stored silently degrades to no-op", func(t *testing.T) { - t.Parallel() - s := &Server{} - s.capturedPassthroughHeaders.Store("s2", "not-a-map") - ctx := s.injectCapturedHeaders(context.Background(), "s2") - assert.Nil(t, headerforward.ForwardedHeadersFromContext(ctx), - "a non-map entry must not be injected (type-assertion fallback)") - }) -} - -// stubClientSession is a minimal server.ClientSession that returns a fixed ID. -// Used to exercise the OnUnregisterSession hook closure without running a full server. -type stubClientSession struct { - server.ClientSession // embed for unused interface methods - id string -} - -func (s *stubClientSession) SessionID() string { return s.id } - -// TestOnUnregisterSessionHookDeletesCapturedHeaders verifies that the closure -// registered with AddOnUnregisterSession in serve.go deletes the session's entry -// from capturedPassthroughHeaders. The test constructs a minimal *Server with a -// pre-populated entry and fires the hook closure directly, avoiding the need to -// bring up an HTTP server. -func TestOnUnregisterSessionHookDeletesCapturedHeaders(t *testing.T) { - t.Parallel() - - srv, err := Serve(context.Background(), &stubVMCP{}, testMinimalServeConfig()) - require.NoError(t, err) - t.Cleanup(func() { _ = srv.Stop(context.Background()) }) - - const sessionID = "test-session-hook" - srv.capturedPassthroughHeaders.Store(sessionID, map[string]string{"X-Key": "v"}) - - // Confirm the entry is present before firing. - _, ok := srv.capturedPassthroughHeaders.Load(sessionID) - require.True(t, ok, "entry must be present before hook fires") - - // Fire the hook by calling UnregisterSession on the mcp-go hooks object. - // The hooks object is embedded in mcpServer; UnregisterSession triggers all - // AddOnUnregisterSession callbacks registered during Serve, including our cleanup. - srv.mcpServer.GetHooks().UnregisterSession(context.Background(), &stubClientSession{id: sessionID}) - - // After the hook fires the entry must be gone. - _, stillPresent := srv.capturedPassthroughHeaders.Load(sessionID) - assert.False(t, stillPresent, - "capturedPassthroughHeaders entry must be deleted by OnUnregisterSession hook") -} diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index d0b8361009..e54e868696 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -14,7 +14,6 @@ import ( "errors" "fmt" "log/slog" - "maps" "net" "net/http" "os" @@ -292,21 +291,6 @@ type Server struct { // Nil if status reporting is disabled. statusReporter vmcpstatus.Reporter - // capturedPassthroughHeaders holds the per-session passthrough headers captured - // once at session-creation time on the Serve path (s.core != nil). It is keyed - // by session ID. Every entry is populated in handleSessionRegistrationImpl and - // cleaned up via the OnUnregisterSession hook (serve.go) or on registration - // failure. - // - // Security: these values may be credentials (e.g. API keys). They are stored - // node-local ONLY and are NEVER written to Redis or any shared state store. - // A session restored on another pod will have no captured headers — matching - // the legacy path's imperfect cross-pod behavior (the legacy per-session - // connection cannot be serialized either). This is acceptable: the alternative - // would require persisting credentials to shared state, which violates the - // security constraint in .claude/rules/security.md. - capturedPassthroughHeaders sync.Map // map[sessionID string → map[string]string] - // shutdownFuncs contains cleanup functions to run during Stop(). // Populated during Start() initialization before blocking; no mutex needed // since Stop() is only called after Start()'s select returns. @@ -1247,12 +1231,6 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv slog.Debug("creating session-scoped backends", "session_id", sessionID) // Defer cleanup: if any error occurs, terminate the session and log failures. - // Also remove any node-local captured passthrough headers stored for this session. - // - // Double-delete safety: capturedPassthroughHeaders.Delete is idempotent. This - // defer fires only on registration failure — before the session is live and before - // the OnUnregisterSession hook (serve.go) is ever registered for this session ID. - // The two delete paths are therefore mutually exclusive for well-formed sessions. defer func() { if retErr != nil { if _, termErr := s.vmcpSessionMgr.Terminate(sessionID); termErr != nil { @@ -1261,10 +1239,6 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv "error", termErr, "original_error", retErr) } - // The Store above runs only after CreateSession succeeds; if we reach - // this defer, the Store either never ran (CreateSession failed) or ran - // but injectCoreSessionCapabilities failed. Either way, Delete is safe. - s.capturedPassthroughHeaders.Delete(sessionID) } }() @@ -1295,19 +1269,6 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv // returned error becomes retErr (named return), so the defer terminates the // session on failure. if s.core != nil { - // Capture passthrough headers from the session-creation request context - // (once per session). The snapshot is injected into every subsequent backend - // call for this session's lifetime by coreToolHandler / coreResourceHandler, - // ensuring the backend always sees the session-creation-time value even when - // the client changes the header on later requests (session-stable forwarding). - // Node-local only — never written to Redis (see capturedPassthroughHeaders). - if captured := headerforward.ForwardedHeadersFromContext(ctx); len(captured) > 0 { - // Defensive copy: ForwardedHeadersFromContext returns the map reference - // placed in the context by CaptureMiddleware. Cloning before storing - // ensures the session's credential snapshot is structurally immutable - // regardless of whether the caller mutates its own map later. - s.capturedPassthroughHeaders.Store(sessionID, maps.Clone(captured)) - } return s.injectCoreSessionCapabilities(ctx, session) } diff --git a/test/integration/vmcp/passthrough_headers_test.go b/test/integration/vmcp/passthrough_headers_test.go index d3ed6e1a76..bc4f75810e 100644 --- a/test/integration/vmcp/passthrough_headers_test.go +++ b/test/integration/vmcp/passthrough_headers_test.go @@ -19,15 +19,15 @@ import ( // chain: // // MCP client → headerforward.CaptureMiddleware → request-context value → -// per-session backend client (mergeForwardedHeaders) → backend HTTP request. +// MergeForwardedHeaders (client.go) → backend HTTP request. // // The test exercises the real server wiring (via helpers.WithPassthroughHeaders, // which sets vmcpserver.Config.PassthroughHeaders) so headerforward.CaptureMiddleware // runs; it does NOT use a stub or reimplementation of the capture logic. // -// Two assertions are made: +// Two per-request forwarding assertions are made: // 1. Allowlisted header present: X-Test-Api-Key sent by the client arrives at -// the backend unchanged. +// the backend on every call with its current value. // 2. Non-allowlisted header absent: X-Secret sent by the client is dropped and // does not arrive at the backend. func TestVMCPServer_PassthroughHeaders(t *testing.T) { @@ -135,12 +135,11 @@ func TestVMCPServer_PassthroughHeaders(t *testing.T) { assert.Contains(t, text, absentSentinel, "non-allowlisted header slot should contain the absent sentinel") - // ── Second call: header changes mid-session, backend must not see the change ─ - // Forwarded headers are captured once at backend-session creation and stay - // stable for the session's lifetime. Change the caller's X-Test-Api-Key and - // call again: the backend must still observe the ORIGINAL value, proving the - // session does not re-read forwarded headers per request (nor drop them after - // the first call). + // ── Second call: header changes mid-request, backend must see the new value ── + // Passthrough headers are forwarded per-request: each backend call reads the + // current incoming header value from the request context. Change the caller's + // X-Test-Api-Key and call again — the backend must observe the UPDATED value, + // proving the header is re-read on every request. const changedValue = "caller-key-CHANGED" mcpClient.SetHeader(allowlistedHeader, changedValue) @@ -149,10 +148,10 @@ func TestVMCPServer_PassthroughHeaders(t *testing.T) { t.Logf("backend response (second call): %s", text2) - assert.Contains(t, text2, allowlistedValue, - "second call must still see the header value captured at session creation (%q)", - allowlistedValue) - assert.NotContains(t, text2, changedValue, - "changed header value %q must not reach the backend — forwarded headers are session-stable", + assert.Contains(t, text2, changedValue, + "second call must see the updated header value %q (headers are forwarded per-request)", changedValue) + assert.NotContains(t, text2, allowlistedValue, + "original value %q must not appear on the second call after the header was changed", + allowlistedValue) } From 025ee28afe1f9f4dacd88569cec862d1307ec651 Mon Sep 17 00:00:00 2001 From: Trey Date: Mon, 22 Jun 2026 08:31:25 -0700 Subject: [PATCH 6/7] Fix stale comment describing removed session-stable machinery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses stacklok/toolhive#5561 review comment: - MEDIUM pkg/vmcp/client/client.go (3453472675): the comment described coreToolHandler replacing per-request headers with a session-stable snapshot — that code was removed in f08f3d0. Replace with accurate description: the factory reads live per-request forwarded headers on every backend call, so forwarding is per-request. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- pkg/vmcp/client/client.go | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 9d27f91529..dfb13d105e 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -433,13 +433,11 @@ func (h *httpBackendClient) defaultClientFactory(ctx context.Context, target *vm target: target, } - // Merge session-captured passthrough headers (from ctx) with the static - // per-backend header-forward config. On the Serve path, the caller (coreToolHandler - // in pkg/vmcp/server/serve_handlers.go) replaces the per-request forwarded headers - // on ctx with the session-stable snapshot captured once at session creation, so - // passing ctx here provides session-stable header injection without re-reading the - // live incoming request headers. Restricted names and static-config collisions are - // rejected by MergeForwardedHeaders. + // Merge the live per-request forwarded headers (captured by headerforward.CaptureMiddleware + // into the request context) with the static per-backend header-forward config. This factory + // runs on every backend call, so forwarding is per-request: each call reflects the current + // incoming header value. Restricted names and static-config collisions are rejected by + // MergeForwardedHeaders. mergedHeaderForward, mergeErr := headerforward.MergeForwardedHeaders( target.HeaderForward, headerforward.ForwardedHeadersFromContext(ctx), ) From e7adb869c90810b1c7a7a7741bbae538beda5bb8 Mon Sep 17 00:00:00 2001 From: Trey Date: Mon, 22 Jun 2026 08:42:38 -0700 Subject: [PATCH 7/7] Fix stale doc comments in integration test files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses stacklok/toolhive#5561 review comments: - LOW test/integration/vmcp/helpers/vmcp_server.go (3453566471): replace "session-stable" with "per-request" — the PR forwards headers per-request, not from a frozen session-creation snapshot - LOW test/integration/vmcp/passthrough_headers_test.go (3453566476): replace stale type name vmcpserver.Config.PassthroughHeaders with the correct vmcpserver.ServerConfig.PassthroughHeaders (two occurrences) Co-Authored-By: Claude Sonnet 4.6 (1M context) --- test/integration/vmcp/helpers/vmcp_server.go | 2 +- test/integration/vmcp/passthrough_headers_test.go | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/test/integration/vmcp/helpers/vmcp_server.go b/test/integration/vmcp/helpers/vmcp_server.go index 84efe2724b..a32c1d47b0 100644 --- a/test/integration/vmcp/helpers/vmcp_server.go +++ b/test/integration/vmcp/helpers/vmcp_server.go @@ -133,7 +133,7 @@ func getFreePort(tb testing.TB) int { // // The Serve path is used (rather than the legacy server.New path) so that // integration tests exercise the production code path that routes tool/resource -// calls through core.VMCP. This is required for testing session-stable +// calls through core.VMCP. This is required for testing per-request // passthrough header forwarding (#5560). func NewVMCPServer( ctx context.Context, tb testing.TB, backends []vmcptypes.Backend, opts ...VMCPServerOption, diff --git a/test/integration/vmcp/passthrough_headers_test.go b/test/integration/vmcp/passthrough_headers_test.go index bc4f75810e..f93ecfa3de 100644 --- a/test/integration/vmcp/passthrough_headers_test.go +++ b/test/integration/vmcp/passthrough_headers_test.go @@ -22,7 +22,7 @@ import ( // MergeForwardedHeaders (client.go) → backend HTTP request. // // The test exercises the real server wiring (via helpers.WithPassthroughHeaders, -// which sets vmcpserver.Config.PassthroughHeaders) so headerforward.CaptureMiddleware +// which sets vmcpserver.ServerConfig.PassthroughHeaders) so headerforward.CaptureMiddleware // runs; it does NOT use a stub or reimplementation of the capture logic. // // Two per-request forwarding assertions are made: @@ -83,7 +83,7 @@ func TestVMCPServer_PassthroughHeaders(t *testing.T) { t.Cleanup(apiBackend.Close) // ── vMCP server ─────────────────────────────────────────────────────────── - // WithPassthroughHeaders sets vmcpserver.Config.PassthroughHeaders, which + // WithPassthroughHeaders sets vmcpserver.ServerConfig.PassthroughHeaders, which // installs headerforward.CaptureMiddleware so allowlisted headers are captured // into the request context and forwarded to backends for each request. backends := []vmcp.Backend{