diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 615959e83b..dfb13d105e 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -433,11 +433,19 @@ 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 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), + ) + 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_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/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..cf2c9cfcb3 100644 --- a/pkg/vmcp/headerforward/transport.go +++ b/pkg/vmcp/headerforward/transport.go @@ -101,6 +101,80 @@ 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. 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). +// 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) { + 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 { + // 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 { + 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..c8f2b42766 100644 --- a/pkg/vmcp/headerforward/transport_test.go +++ b/pkg/vmcp/headerforward/transport_test.go @@ -241,3 +241,160 @@ 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() + + // 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) + return + } + 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 + } + 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/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/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() diff --git a/test/integration/vmcp/helpers/vmcp_server.go b/test/integration/vmcp/helpers/vmcp_server.go index 4679b3a9aa..a32c1d47b0 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 per-request +// 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()) diff --git a/test/integration/vmcp/passthrough_headers_test.go b/test/integration/vmcp/passthrough_headers_test.go index d3ed6e1a76..f93ecfa3de 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 +// which sets vmcpserver.ServerConfig.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) { @@ -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{ @@ -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) }