diff --git a/pkg/vmcp/aggregator/caching_aggregator.go b/pkg/vmcp/aggregator/caching_aggregator.go new file mode 100644 index 0000000000..2275c7209f --- /dev/null +++ b/pkg/vmcp/aggregator/caching_aggregator.go @@ -0,0 +1,140 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package aggregator + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "io" + "sort" + "time" + + lru "github.com/hashicorp/golang-lru/v2" + + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/headerforward" +) + +// capabilityCacheMaxEntries bounds the per-identity capability cache so it cannot grow +// without limit (one entry per distinct identity + forwarded-credential + backend-set key). +// Beyond it, the LRU evicts the least-recently-used entry. 1024 distinct active keys per +// node is generous for a vMCP instance; tune if real workloads exceed it. +const capabilityCacheMaxEntries = 1024 + +// cachingAggregator wraps an Aggregator and memoizes AggregateCapabilities for a bounded +// TTL, so the Serve path does not re-sweep every backend's tools/list on every tool call. +// +// On the New/Serve path, core.CallTool/ReadResource re-derive the advertised view on every +// call (the core holds no per-session cache); without this, a session with M calls performs +// ~M+1 full backend sweeps instead of the legacy ~1-per-session. This decorator restores +// once-per-(identity, TTL) freshness without coupling the core or Serve to a cache — it sits +// transparently below the core, wrapping the Aggregator the core already calls. +// +// It delegates the cache mechanics (LRU eviction, size bounding, thread-safety) to +// hashicorp/golang-lru and adds only a lazy TTL check on read. The base (non-expirable) LRU +// is used deliberately: the expirable variant runs a perpetual background cleanup goroutine +// with no Stop, which a per-server cache would leak (the repo's tests run under goleak). +// +// Security: backend enumeration is identity-dependent — what each backend returns depends on +// the credential presented to it — so the cache key MUST include the identity and the +// forwarded credentials, not just the backend set. Keying on (subject, forwarded headers, +// backend IDs) ensures one caller's capability view is never served to another, while a +// single caller's view is shared across their sessions. The key is a SHA-256 digest, so raw +// credential values are not retained as map keys. The cache is node-local and never persisted +// (it would be a credentialed view in shared state otherwise). +type cachingAggregator struct { + // Aggregator is the wrapped aggregator; embedding delegates every method except the + // AggregateCapabilities override below. + Aggregator + + ttl time.Duration + cache *lru.Cache[string, cacheEntry] +} + +type cacheEntry struct { + caps *AggregatedCapabilities + at time.Time +} + +// NewCachingAggregator wraps next so AggregateCapabilities results are memoized per identity +// for ttl, backed by a size-bounded LRU. A ttl <= 0 disables caching (next is returned +// unwrapped) so a misconfiguration cannot silently serve permanently-stale capabilities. A +// nil next is returned as-is so the downstream nil-aggregator validation (core.New) still +// fires rather than being masked by a non-nil wrapper. +func NewCachingAggregator(next Aggregator, ttl time.Duration) Aggregator { + if next == nil || ttl <= 0 { + return next + } + cache, err := lru.New[string, cacheEntry](capabilityCacheMaxEntries) + if err != nil { + // lru.New only errors on a non-positive size, which is a positive constant here, so + // this is unreachable; degrade to the uncached aggregator rather than panicking. + return next + } + return &cachingAggregator{Aggregator: next, ttl: ttl, cache: cache} +} + +// AggregateCapabilities returns a cached view when a fresh entry exists for the caller's +// identity + forwarded credentials + backend set, and otherwise sweeps the backends (via the +// wrapped aggregator) and caches the result. Errors are never cached. The returned value is +// treated as immutable by the core (it derives fresh per-call routers from it), so the cached +// pointer is shared rather than deep-copied. +func (c *cachingAggregator) AggregateCapabilities( + ctx context.Context, backends []vmcp.Backend, +) (*AggregatedCapabilities, error) { + key := cacheKey(ctx, backends) + if e, ok := c.cache.Get(key); ok && time.Since(e.at) < c.ttl { + return e.caps, nil + } + + // Miss/expiry: sweep with the lock released (Get/Add are individually locked) so callers + // with different keys are not serialized behind one backend sweep. Concurrent misses for + // the same key may each sweep once (last writer wins) — acceptable for a cold/expired key. + caps, err := c.Aggregator.AggregateCapabilities(ctx, backends) + if err != nil { + return nil, err + } + c.cache.Add(key, cacheEntry{caps: caps, at: time.Now()}) + return caps, nil +} + +// cacheKey derives a collision-resistant key from the inputs that drive backend enumeration: +// the caller's subject, the forwarded headers (passthrough credentials/scopes), and the +// backend ID set. Hashing keeps raw credential values out of the cache keys. +func cacheKey(ctx context.Context, backends []vmcp.Backend) string { + h := sha256.New() + + if id, ok := auth.IdentityFromContext(ctx); ok && id != nil { + _, _ = io.WriteString(h, id.Subject) + } + _, _ = h.Write([]byte{0}) + + fwd := headerforward.ForwardedHeadersFromContext(ctx) + fwdKeys := make([]string, 0, len(fwd)) + for k := range fwd { + fwdKeys = append(fwdKeys, k) + } + sort.Strings(fwdKeys) + for _, k := range fwdKeys { + _, _ = io.WriteString(h, k) + _, _ = h.Write([]byte{'='}) + _, _ = io.WriteString(h, fwd[k]) + _, _ = h.Write([]byte{0}) + } + _, _ = h.Write([]byte{0}) + + ids := make([]string, 0, len(backends)) + for _, b := range backends { + ids = append(ids, b.ID) + } + sort.Strings(ids) + for _, id := range ids { + _, _ = io.WriteString(h, id) + _, _ = h.Write([]byte{0}) + } + + return hex.EncodeToString(h.Sum(nil)) +} diff --git a/pkg/vmcp/aggregator/caching_aggregator_test.go b/pkg/vmcp/aggregator/caching_aggregator_test.go new file mode 100644 index 0000000000..e586b1844f --- /dev/null +++ b/pkg/vmcp/aggregator/caching_aggregator_test.go @@ -0,0 +1,126 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package aggregator_test + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/headerforward" +) + +func ctxWithSubject(subject string) context.Context { + return auth.WithIdentity(context.Background(), + &auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: subject}}) +} + +var testBackends = []vmcp.Backend{{ID: "b1"}, {ID: "b2"}} + +// TestCachingAggregator_CacheHitWithinTTL: two calls with the same identity and backend set +// within the TTL hit the cache, so the wrapped aggregator sweeps only once. +func TestCachingAggregator_CacheHitWithinTTL(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mock := mocks.NewMockAggregator(ctrl) + caps := &aggregator.AggregatedCapabilities{} + mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(caps, nil).Times(1) + + c := aggregator.NewCachingAggregator(mock, time.Hour) + ctx := ctxWithSubject("alice") + + first, err := c.AggregateCapabilities(ctx, testBackends) + require.NoError(t, err) + second, err := c.AggregateCapabilities(ctx, testBackends) + require.NoError(t, err) + assert.Same(t, first, second, "the cached view must be returned on a hit") + assert.Same(t, caps, second) +} + +// TestCachingAggregator_KeyedByIdentityAndHeaders: distinct subjects, and distinct forwarded +// credentials for the same subject, are distinct cache keys — each sweeps the backends. This +// is the security-critical property: one caller's view is never served to another. +func TestCachingAggregator_KeyedByIdentityAndHeaders(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mock := mocks.NewMockAggregator(ctrl) + // alice, bob, and alice-with-a-forwarded-token are three distinct keys → three sweeps. + mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil).Times(3) + + c := aggregator.NewCachingAggregator(mock, time.Hour) + + _, err := c.AggregateCapabilities(ctxWithSubject("alice"), testBackends) + require.NoError(t, err) + _, err = c.AggregateCapabilities(ctxWithSubject("bob"), testBackends) + require.NoError(t, err) + aliceWithToken := headerforward.WithForwardedHeaders(ctxWithSubject("alice"), + map[string]string{"Authorization": "Bearer xyz"}) + _, err = c.AggregateCapabilities(aliceWithToken, testBackends) + require.NoError(t, err) +} + +// TestCachingAggregator_TTLExpiry: once the TTL elapses, the next call re-sweeps. +func TestCachingAggregator_TTLExpiry(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mock := mocks.NewMockAggregator(ctrl) + mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil).Times(2) + + c := aggregator.NewCachingAggregator(mock, 20*time.Millisecond) + ctx := ctxWithSubject("alice") + + _, err := c.AggregateCapabilities(ctx, testBackends) + require.NoError(t, err) + time.Sleep(40 * time.Millisecond) + _, err = c.AggregateCapabilities(ctx, testBackends) + require.NoError(t, err) +} + +// TestCachingAggregator_DisabledWhenTTLNonPositive: a ttl <= 0 returns the wrapped aggregator +// unwrapped, so callers cannot silently get a permanently-stale cache. +func TestCachingAggregator_DisabledWhenTTLNonPositive(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mock := mocks.NewMockAggregator(ctrl) + + assert.Same(t, mock, aggregator.NewCachingAggregator(mock, 0)) + assert.Same(t, mock, aggregator.NewCachingAggregator(mock, -time.Second)) + + // A nil aggregator is returned as-is so the downstream nil check (core.New) still fires + // rather than being masked by a non-nil caching wrapper. + var nilAgg aggregator.Aggregator + assert.Nil(t, aggregator.NewCachingAggregator(nilAgg, time.Hour)) +} + +// TestCachingAggregator_ErrorNotCached: a failed sweep is not cached, so the next call retries +// the wrapped aggregator. +func TestCachingAggregator_ErrorNotCached(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mock := mocks.NewMockAggregator(ctrl) + gomock.InOrder( + mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(nil, errors.New("boom")), + mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil), + ) + + c := aggregator.NewCachingAggregator(mock, time.Hour) + ctx := ctxWithSubject("alice") + + _, err := c.AggregateCapabilities(ctx, testBackends) + require.Error(t, err) + _, err = c.AggregateCapabilities(ctx, testBackends) + require.NoError(t, err, "the error must not have been cached") +} diff --git a/pkg/vmcp/auth/factory/incoming.go b/pkg/vmcp/auth/factory/incoming.go index 313ff94c80..5618e619a2 100644 --- a/pkg/vmcp/auth/factory/incoming.go +++ b/pkg/vmcp/auth/factory/incoming.go @@ -123,6 +123,40 @@ func newCedarAuthzMiddleware( slog.Info("creating Cedar authorization middleware", "policies", len(authzCfg.Policies)) + authzConfig, err := buildCedarAuthzConfig(authzCfg) + if err != nil { + return nil, fmt.Errorf("failed to create authz config: %w", err) + } + + // Create the middleware using the existing factory + middlewareFn, err := authz.CreateMiddlewareFromConfig(authzConfig, serverName, passThroughTools) + if err != nil { + return nil, fmt.Errorf("failed to create Cedar middleware: %w", err) + } + + return middlewareFn, nil +} + +// BuildAuthzConfig builds the authorizer-agnostic *authz.Config that the vMCP core +// admission seam (core.Config.Authz) consumes from the incoming-auth config, or +// (nil, nil) when no Cedar policies are configured. +// +// It is the SAME config newCedarAuthzMiddleware builds the HTTP authz middleware from, +// surfaced so the composition root can feed it to core.New via server.Config.Authz once +// server.New routes through Serve. The nil return mirrors that middleware's nil result +// for the no-policies case, preserving allow-all parity (a nil core Authz is allow-all). +func BuildAuthzConfig(authzCfg *config.AuthzConfig) (*authz.Config, error) { + if authzCfg == nil || authzCfg.Type != "cedar" || len(authzCfg.Policies) == 0 { + return nil, nil + } + return buildCedarAuthzConfig(authzCfg) +} + +// buildCedarAuthzConfig converts the vMCP Cedar authz config into the authorizer-agnostic +// authorizers.Config (aliased as authz.Config) consumed by both the HTTP authz middleware +// (newCedarAuthzMiddleware) and the core admission seam (via BuildAuthzConfig). Callers +// guarantee authzCfg is non-nil with at least one policy. +func buildCedarAuthzConfig(authzCfg *config.AuthzConfig) (*authz.Config, error) { // Default EntitiesJSON to "[]" when the operator/CLI did not set it. Cedar // requires a valid JSON array; an empty string would fail to parse. entitiesJSON := authzCfg.EntitiesJSON @@ -149,19 +183,7 @@ func newCedarAuthzMiddleware( }, } - // Create the authz Config using the factory method - authzConfig, err := authorizers.NewConfig(cedarConfig) - if err != nil { - return nil, fmt.Errorf("failed to create authz config: %w", err) - } - - // Create the middleware using the existing factory - middlewareFn, err := authz.CreateMiddlewareFromConfig(authzConfig, serverName, passThroughTools) - if err != nil { - return nil, fmt.Errorf("failed to create Cedar middleware: %w", err) - } - - return middlewareFn, nil + return authorizers.NewConfig(cedarConfig) } // newOIDCAuthMiddleware creates OIDC authentication middleware. diff --git a/pkg/vmcp/cli/serve.go b/pkg/vmcp/cli/serve.go index 71cd595fd2..49290020f8 100644 --- a/pkg/vmcp/cli/serve.go +++ b/pkg/vmcp/cli/serve.go @@ -378,6 +378,16 @@ func Serve(ctx context.Context, cfg ServeConfig) error { return fmt.Errorf("failed to create authentication middleware: %w", err) } + // Build the authorizer-agnostic authz config the core admission seam consumes. It is + // the same config the HTTP authz middleware above is built from; server.New routes + // through Serve, which ignores the (vestigial) AuthzMiddleware and enforces authz in + // the core from this config instead. Nil when no Cedar policies are configured + // (allow-all parity). + authzConfig, err := authfactory.BuildAuthzConfig(vmcpCfg.IncomingAuth.Authz) + if err != nil { + return fmt.Errorf("failed to build authorization config: %w", err) + } + slog.Info(fmt.Sprintf("Incoming authentication configured: %s", vmcpCfg.IncomingAuth.Type)) namespace := vmcpNamespace() @@ -425,6 +435,11 @@ func Serve(ctx context.Context, cfg ServeConfig) error { OptimizerConfig: optCfg, SessionFactory: sessionFactory, SessionStorage: vmcpCfg.SessionStorage, + // Core collaborators: server.New routes through core.New + Serve, so the core + // is the single aggregator and authorizer. The aggregator is the same instance + // that backs discovery; Authz feeds the core admission seam (nil = allow-all). + Aggregator: agg, + Authz: authzConfig, }) // Assign Watcher only when backendWatcher is non-nil. A typed nil diff --git a/pkg/vmcp/core/core.go b/pkg/vmcp/core/core.go index 4e10fa9a09..6ba2cf31d0 100644 --- a/pkg/vmcp/core/core.go +++ b/pkg/vmcp/core/core.go @@ -104,7 +104,13 @@ type VMCP interface { // unadvertised name. Applies the same admission filter as ListPrompts. LookupPrompt(ctx context.Context, identity *auth.Identity, name string) (*vmcp.Prompt, error) - // Close releases core-held resources (backend connections, etc.). + // BackendHealth returns the backend health reporter the core owns, or nil when health + // monitoring is disabled. The core builds, starts, and (via Close) stops the monitor and + // filters capabilities with it; the transport layer uses this only to report on or sync + // backend health (e.g. the /health route and periodic status reporting). + BackendHealth() health.Reporter + + // Close releases core-held resources (backend connections, the health monitor, etc.). // Implementations must be idempotent: calling Close multiple times returns nil. Close() error } @@ -114,9 +120,9 @@ type VMCP interface { // *coreVMCP) lands in a later change. // // Cross-cutting TelemetryProvider/AuditConfig are consumed by both New and Serve -// (not a clean partition). HealthStatusProvider is the read-only health view -// built at the composition root and injected here; a nil provider means no health -// filtering (all backends included), matching today's no-monitor behavior. +// (not a clean partition). HealthMonitorConfig is the backend health monitoring +// configuration; the core builds, starts, and stops the monitor from it (a nil config +// disables monitoring, including all backends — matching today's no-monitor behavior). type Config struct { // Aggregator discovers and merges backend capabilities into the advertised set. Aggregator aggregator.Aggregator @@ -153,9 +159,10 @@ type Config struct { // AuditConfig is the cross-cutting audit configuration (also consumed by Serve). AuditConfig *audit.Config - // HealthStatusProvider is the read-only backend health view built at the - // composition root. Nil means no health filtering (all backends included). - HealthStatusProvider health.StatusProvider + // HealthMonitorConfig configures backend health monitoring. The core builds, starts, and + // (via Close) stops the monitor from it, filters capabilities with it, and exposes it via + // BackendHealth. Nil disables monitoring (no health filtering; all backends included). + HealthMonitorConfig *health.MonitorConfig // Elicitation sends MCP elicitation requests to the client and blocks for the // response. It is the domain-typed seam (vmcp anti-pattern #5: no mcp-go types) diff --git a/pkg/vmcp/core/core_vmcp.go b/pkg/vmcp/core/core_vmcp.go index a398c09b9f..fe1d381e69 100644 --- a/pkg/vmcp/core/core_vmcp.go +++ b/pkg/vmcp/core/core_vmcp.go @@ -48,10 +48,18 @@ type coreVMCP struct { backendRegistry vmcp.BackendRegistry backendClient vmcp.BackendClient - // health is the injected read-only backend health view. Nil means no health - // filtering (all backends included), matching today's no-monitor behavior. + // health is the read-only backend health view used for capability filtering. It is the + // healthMonitor below as a StatusProvider, or a true nil interface when monitoring is + // disabled/failed (so filterHealthyBackends includes all backends rather than calling + // through a typed-nil). health health.StatusProvider + // healthMonitor is the backend health monitor the core owns: built and started in New, + // stopped in Close, and exposed for reporting via BackendHealth. Nil when health + // monitoring is disabled or failed to start. (#5443 reversal: the monitor's lifecycle + // moved here from server.New/Serve so it has a single owner.) + healthMonitor *health.Monitor + // admission enforces the authorization decision for List* (filter) and // Call/Read/Get (deny) from one source. Never nil: New installs an allow-all // seam when authz is unconfigured. @@ -188,11 +196,23 @@ func New(cfg *Config) (VMCP, error) { return nil, fmt.Errorf("workflow validation failed: %w", err) } + // Build and start the backend health monitor (#5443 reversal: the core owns its + // lifecycle). The core filters capabilities with it and stops it in Close. It is started + // with context.Background() and torn down via Close — like the state store — since New has + // no request-scoped context. The monitor uses the undecorated backend client so health + // checks do not emit backend-call telemetry. Built last so few error paths must stop it. + healthMonitor, healthProvider, err := buildHealthMonitor(cfg) + if err != nil { + stopStore() + return nil, err + } + return &coreVMCP{ aggregator: cfg.Aggregator, backendRegistry: cfg.BackendRegistry, backendClient: backendClient, - health: cfg.HealthStatusProvider, + health: healthProvider, + healthMonitor: healthMonitor, admission: admission, workflowDefs: workflowDefs, composerFactory: composerFactory, @@ -200,6 +220,48 @@ func New(cfg *Config) (VMCP, error) { }, nil } +// buildHealthMonitor builds and starts the backend health monitor from cfg, returning the +// monitor (for Close/reporting) and its StatusProvider view (for capability filtering). +// Returns (nil, nil, nil) when health monitoring is not configured. A monitor that fails to +// start is logged and disabled — not fatal — preserving the legacy lenient behavior, so the +// returned monitor is nil in that case too. The StatusProvider is returned as a true nil +// interface when disabled so filterHealthyBackends does not call through a typed-nil. +func buildHealthMonitor(cfg *Config) (*health.Monitor, health.StatusProvider, error) { + if cfg.HealthMonitorConfig == nil { + slog.Info("health monitoring disabled") + return nil, nil, nil + } + + // Use the undecorated client so health checks do not emit backend-call telemetry. + monitor, err := health.NewMonitor( + cfg.BackendClient, cfg.BackendRegistry.List(context.Background()), *cfg.HealthMonitorConfig) + if err != nil { + return nil, nil, fmt.Errorf("failed to create health monitor: %w", err) + } + + if err := monitor.Start(context.Background()); err != nil { + slog.Warn("failed to start health monitor, disabling health monitoring", "error", err) + return nil, nil, nil + } + + slog.Info("health monitoring enabled", + "check_interval", cfg.HealthMonitorConfig.CheckInterval, + "unhealthy_threshold", cfg.HealthMonitorConfig.UnhealthyThreshold, + "timeout", cfg.HealthMonitorConfig.Timeout, + "degraded_threshold", cfg.HealthMonitorConfig.DegradedThreshold) + return monitor, monitor, nil +} + +// BackendHealth returns the core's backend health reporter, or nil when health monitoring is +// disabled. The core owns the monitor's lifecycle (build/start in New, stop in Close); the +// transport layer uses this only to report on or sync backend health. +func (c *coreVMCP) BackendHealth() health.Reporter { + if c.healthMonitor == nil { + return nil + } + return c.healthMonitor +} + // ListTools health-filters the registry, aggregates backend capabilities on // demand, appends the composite tools reachable in the current view, and returns // only the subset identity is admitted to call (the admission seam — the same @@ -293,7 +355,14 @@ func (c *coreVMCP) LookupPrompt(ctx context.Context, identity *auth.Identity, na // the underlying Stop closes a channel that cannot be closed twice, so the work // is guarded by sync.Once and subsequent calls return nil. func (c *coreVMCP) Close() error { - c.closeOnce.Do(c.stopStore) + c.closeOnce.Do(func() { + c.stopStore() + if c.healthMonitor != nil { + if err := c.healthMonitor.Stop(); err != nil { + slog.Warn("failed to stop health monitor", "error", err) + } + } + }) return nil } diff --git a/pkg/vmcp/core/core_vmcp_test.go b/pkg/vmcp/core/core_vmcp_test.go index 924143a064..69f6346d2e 100644 --- a/pkg/vmcp/core/core_vmcp_test.go +++ b/pkg/vmcp/core/core_vmcp_test.go @@ -9,6 +9,7 @@ import ( "errors" "log/slog" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -19,6 +20,7 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp/aggregator" aggmocks "github.com/stacklok/toolhive/pkg/vmcp/aggregator/mocks" "github.com/stacklok/toolhive/pkg/vmcp/composer" + "github.com/stacklok/toolhive/pkg/vmcp/health" vmcpmocks "github.com/stacklok/toolhive/pkg/vmcp/mocks" routermocks "github.com/stacklok/toolhive/pkg/vmcp/router/mocks" ) @@ -307,11 +309,6 @@ func TestListTools_HealthFiltersBeforeAggregating(t *testing.T) { {ID: "unhealthy", HealthStatus: vmcp.BackendHealthy}, // overridden by provider below {ID: "unknown", HealthStatus: vmcp.BackendHealthy}, // overridden by provider below } - cfg.HealthStatusProvider = fakeStatusProvider{ - "unhealthy": vmcp.BackendUnhealthy, - "unknown": vmcp.BackendUnknown, - } - m.reg.EXPECT().List(gomock.Any()).Return(backends) var got []vmcp.Backend m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).DoAndReturn( @@ -324,6 +321,13 @@ func TestListTools_HealthFiltersBeforeAggregating(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { _ = c.Close() }) + // The core builds its own monitor from HealthMonitorConfig in production; inject a stub + // health view directly (internal test) to exercise filtering without a running monitor. + c.(*coreVMCP).health = fakeStatusProvider{ + "unhealthy": vmcp.BackendUnhealthy, + "unknown": vmcp.BackendUnknown, + } + _, err = c.ListTools(context.Background(), nil) require.NoError(t, err) @@ -449,7 +453,6 @@ func TestIdentityNotLogged(t *testing.T) { // An excluded backend makes filterHealthyBackends emit a DEBUG log on every // aggregatedView, guaranteeing the buffer is non-empty so the assertion below // is not vacuous (it would otherwise pass if every path short-circuited). - cfg.HealthStatusProvider = fakeStatusProvider{"bad": vmcp.BackendUnhealthy} m.reg.EXPECT().List(gomock.Any()).Return([]vmcp.Backend{ {ID: "ok", HealthStatus: vmcp.BackendHealthy}, {ID: "bad", HealthStatus: vmcp.BackendHealthy}, @@ -461,6 +464,9 @@ func TestIdentityNotLogged(t *testing.T) { c, err := New(cfg) require.NoError(t, err) t.Cleanup(func() { _ = c.Close() }) + // Inject a stub health view (internal test) so filterHealthyBackends excludes "bad" and + // emits its DEBUG log (see comment above). + c.(*coreVMCP).health = fakeStatusProvider{"bad": vmcp.BackendUnhealthy} var buf bytes.Buffer prev := slog.Default() @@ -486,3 +492,53 @@ func TestIdentityNotLogged(t *testing.T) { require.NotEmpty(t, logs, "expected the driven methods to emit logs, else the assertion is vacuous") assert.NotContains(t, logs, secret, "identity token must never be logged") } + +// TestNew_HealthMonitorOwnedByCore verifies the core owns the backend health monitor +// (#5443 reversal): New builds and starts it from HealthMonitorConfig, exposes it via +// BackendHealth, drives initial checks, and Close stops it idempotently while leaving the +// reporter queryable. +func TestNew_HealthMonitorOwnedByCore(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + backends := []vmcp.Backend{{ID: testBackendID, Name: "be1", BaseURL: "http://be1:8080", TransportType: "sse"}} + // The monitor lists backends from the registry and health-checks them via the client. + m.reg.EXPECT().List(gomock.Any()).Return(backends).AnyTimes() + m.client.EXPECT().ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil).AnyTimes() + cfg.HealthMonitorConfig = &health.MonitorConfig{ + CheckInterval: 50 * time.Millisecond, + UnhealthyThreshold: 1, + Timeout: time.Second, + } + + c, err := New(cfg) + require.NoError(t, err) + + reporter := c.BackendHealth() + require.NotNil(t, reporter, "the core must expose the health monitor it built and started") + + require.Eventually(t, func() bool { + status, ok := reporter.QueryBackendStatus(testBackendID) + return ok && status == vmcp.BackendHealthy + }, 2*time.Second, 10*time.Millisecond, "the core-started monitor should report the backend healthy") + + // Close stops the monitor and is idempotent; the reporter stays valid (stopped, not nil). + require.NoError(t, c.Close()) + require.NoError(t, c.Close()) + assert.NotNil(t, c.BackendHealth()) +} + +// TestNew_HealthMonitorDisabledWhenNil verifies a nil HealthMonitorConfig leaves health +// monitoring disabled: the core builds no monitor and BackendHealth reports none. +func TestNew_HealthMonitorDisabledWhenNil(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + m.reg.EXPECT().List(gomock.Any()).Return(nil).AnyTimes() + + c, err := New(cfg) // HealthMonitorConfig is nil + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + assert.Nil(t, c.BackendHealth(), "nil HealthMonitorConfig must leave health monitoring disabled") +} diff --git a/pkg/vmcp/health/monitor.go b/pkg/vmcp/health/monitor.go index 82a6c72fac..7c5f3cead4 100644 --- a/pkg/vmcp/health/monitor.go +++ b/pkg/vmcp/health/monitor.go @@ -48,6 +48,32 @@ type StatusProvider interface { QueryBackendStatus(backendID string) (vmcp.BackendHealthStatus, bool) } +// Reporter is the read-and-sync surface a running Monitor exposes to the transport/status +// layer. The core owns the Monitor (builds it, starts it, stops it in Close, and filters +// capabilities with it); callers that only report on or sync backend health depend on this +// interface rather than the concrete *Monitor. *Monitor implements it. +type Reporter interface { + StatusProvider + + // GetBackendStatus returns the current health status for a backend, or an error if the + // backend is not monitored. + GetBackendStatus(backendID string) (vmcp.BackendHealthStatus, error) + // GetBackendState returns the full health state of a backend, or an error if unmonitored. + GetBackendState(backendID string) (*State, error) + // GetAllBackendStates returns the health states of all monitored backends. + GetAllBackendStates() map[string]*State + // GetHealthSummary returns an aggregate summary across all monitored backends. + GetHealthSummary() Summary + // WaitForInitialHealthChecks blocks until every backend has completed its first check. + WaitForInitialHealthChecks() + // UpdateBackends reconciles the monitored set with newBackends (dynamic registries). + UpdateBackends(newBackends []vmcp.Backend) + // BuildStatus assembles the aggregate vMCP status from current backend health. + BuildStatus() *vmcp.Status +} + +var _ Reporter = (*Monitor)(nil) + // backendCheck manages the health check goroutine lifecycle for a single backend. // It owns the backend snapshot and the cancel function for its goroutine, keeping // per-backend lifecycle mechanics out of the Monitor's coordination logic. diff --git a/pkg/vmcp/server/annotation_enrichment.go b/pkg/vmcp/server/annotation_enrichment.go deleted file mode 100644 index 6a1d2318e6..0000000000 --- a/pkg/vmcp/server/annotation_enrichment.go +++ /dev/null @@ -1,114 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "log/slog" - "net/http" - - "github.com/mark3labs/mcp-go/mcp" - - "github.com/stacklok/toolhive/pkg/authz/authorizers" - mcpparser "github.com/stacklok/toolhive/pkg/mcp" - "github.com/stacklok/toolhive/pkg/vmcp" - "github.com/stacklok/toolhive/pkg/vmcp/aggregator" - "github.com/stacklok/toolhive/pkg/vmcp/discovery" -) - -// AnnotationEnrichmentMiddleware creates middleware that reads tool annotations -// from the discovery context and injects them into the request context for -// the authz middleware to use. -// -// This middleware sits between discovery and authz in the middleware chain: -// -// ... -> discovery -> annotation-enrichment -> authz -> ... -// -// It only enriches context for tools/call requests. For all other request -// types, it passes through without modification. -func AnnotationEnrichmentMiddleware(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - ctx := r.Context() - - // Only enrich for tools/call requests where authz needs annotation data. - parsedReq := mcpparser.GetParsedMCPRequest(ctx) - if parsedReq == nil || parsedReq.Method != string(mcp.MethodToolsCall) { - next.ServeHTTP(w, r) - return - } - - toolName := parsedReq.ResourceID - if toolName == "" { - next.ServeHTTP(w, r) - return - } - - // Get discovered capabilities from context (set by discovery middleware). - caps, ok := discovery.DiscoveredCapabilitiesFromContext(ctx) - if !ok || caps == nil { - next.ServeHTTP(w, r) - return - } - - // Search all tool lists (backend tools and composite tools) for a match. - if ann := findToolAnnotations(toolName, caps); ann != nil { - ctx = authorizers.WithToolAnnotations(ctx, ann) - r = r.WithContext(ctx) - slog.Debug("enriched request context with tool annotations", - "tool", toolName, - "readOnlyHint", ann.ReadOnlyHint, - "destructiveHint", ann.DestructiveHint, - "idempotentHint", ann.IdempotentHint, - "openWorldHint", ann.OpenWorldHint, - ) - } - - next.ServeHTTP(w, r) - }) -} - -// findToolAnnotations searches for a tool by name in the aggregated capabilities -// and converts its vmcp.ToolAnnotations to the authorizers.ToolAnnotations format. -// Returns nil if the tool is not found or has no annotations. -func findToolAnnotations(toolName string, caps *aggregator.AggregatedCapabilities) *authorizers.ToolAnnotations { - // Search backend tools first, then composite tools. - for _, tool := range caps.Tools { - if tool.Name == toolName && tool.Annotations != nil { - return convertAnnotations(tool.Annotations) - } - } - for _, tool := range caps.CompositeTools { - if tool.Name == toolName && tool.Annotations != nil { - return convertAnnotations(tool.Annotations) - } - } - return nil -} - -// convertAnnotations converts vmcp.ToolAnnotations to authorizers.ToolAnnotations. -// Only authorization-relevant hint fields are mapped; informational fields like -// Title are intentionally omitted since they are not used in policy evaluation. -// Returns nil if the source annotations are nil or contain no hint fields. -// -// Kept identical to core.convertAnnotations (pkg/vmcp/core/admission.go), including -// the nil guard below — the two are intentional, temporary duplicates (the core -// cannot import this package: server imports core). #5441 makes this middleware inert -// on the Serve path via the AuthzMiddleware nil-guard; physical removal of this copy -// is deferred to Phase 3 (#5445). Do NOT delete it early on the import-cycle reasoning -// alone: the core admission seam (#5438) reuses the conversion logic, so this copy must -// stay in lockstep with core.convertAnnotations until the middleware itself is retired. -func convertAnnotations(ann *vmcp.ToolAnnotations) *authorizers.ToolAnnotations { - if ann == nil { - return nil - } - if ann.ReadOnlyHint == nil && ann.DestructiveHint == nil && - ann.IdempotentHint == nil && ann.OpenWorldHint == nil { - return nil - } - return &authorizers.ToolAnnotations{ - ReadOnlyHint: ann.ReadOnlyHint, - DestructiveHint: ann.DestructiveHint, - IdempotentHint: ann.IdempotentHint, - OpenWorldHint: ann.OpenWorldHint, - } -} diff --git a/pkg/vmcp/server/annotation_enrichment_test.go b/pkg/vmcp/server/annotation_enrichment_test.go deleted file mode 100644 index fb789f0ede..0000000000 --- a/pkg/vmcp/server/annotation_enrichment_test.go +++ /dev/null @@ -1,243 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "context" - "net/http" - "net/http/httptest" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/stacklok/toolhive/pkg/authz/authorizers" - mcpparser "github.com/stacklok/toolhive/pkg/mcp" - "github.com/stacklok/toolhive/pkg/vmcp" - "github.com/stacklok/toolhive/pkg/vmcp/aggregator" - "github.com/stacklok/toolhive/pkg/vmcp/discovery" -) - -func TestAnnotationEnrichmentMiddleware(t *testing.T) { - t.Parallel() - - boolPtr := func(b bool) *bool { return &b } - - tests := []struct { - name string - method string - resourceID string - capabilities *aggregator.AggregatedCapabilities - setDiscovery bool - setParsed bool - expectedAnnotation *authorizers.ToolAnnotations - }{ - { - name: "enriches_tools_call_with_annotations", - method: "tools/call", - resourceID: "my_tool", - setDiscovery: true, - setParsed: true, - capabilities: &aggregator.AggregatedCapabilities{ - Tools: []vmcp.Tool{ - { - Name: "my_tool", - Annotations: &vmcp.ToolAnnotations{ - ReadOnlyHint: boolPtr(true), - DestructiveHint: boolPtr(false), - }, - }, - }, - }, - expectedAnnotation: &authorizers.ToolAnnotations{ - ReadOnlyHint: boolPtr(true), - DestructiveHint: boolPtr(false), - }, - }, - { - name: "passes_through_for_non_tools_call", - method: "tools/list", - resourceID: "", - setDiscovery: true, - setParsed: true, - capabilities: &aggregator.AggregatedCapabilities{ - Tools: []vmcp.Tool{ - { - Name: "my_tool", - Annotations: &vmcp.ToolAnnotations{ - ReadOnlyHint: boolPtr(true), - }, - }, - }, - }, - expectedAnnotation: nil, - }, - { - name: "passes_through_when_no_parsed_request", - method: "", - resourceID: "", - setDiscovery: false, - setParsed: false, - capabilities: nil, - expectedAnnotation: nil, - }, - { - name: "passes_through_when_no_discovery_context", - method: "tools/call", - resourceID: "my_tool", - setDiscovery: false, - setParsed: true, - capabilities: nil, - expectedAnnotation: nil, - }, - { - name: "passes_through_when_tool_not_found", - method: "tools/call", - resourceID: "nonexistent_tool", - setDiscovery: true, - setParsed: true, - capabilities: &aggregator.AggregatedCapabilities{ - Tools: []vmcp.Tool{ - { - Name: "other_tool", - Annotations: &vmcp.ToolAnnotations{ - ReadOnlyHint: boolPtr(true), - }, - }, - }, - }, - expectedAnnotation: nil, - }, - { - name: "passes_through_when_tool_has_no_annotations", - method: "tools/call", - resourceID: "bare_tool", - setDiscovery: true, - setParsed: true, - capabilities: &aggregator.AggregatedCapabilities{ - Tools: []vmcp.Tool{ - { - Name: "bare_tool", - Annotations: nil, - }, - }, - }, - expectedAnnotation: nil, - }, - { - name: "passes_through_when_annotations_have_no_hints", - method: "tools/call", - resourceID: "empty_ann_tool", - setDiscovery: true, - setParsed: true, - capabilities: &aggregator.AggregatedCapabilities{ - Tools: []vmcp.Tool{ - { - Name: "empty_ann_tool", - Annotations: &vmcp.ToolAnnotations{ - Title: "Just a title, no hints", - }, - }, - }, - }, - expectedAnnotation: nil, - }, - { - name: "enriches_from_composite_tools", - method: "tools/call", - resourceID: "composite_tool", - setDiscovery: true, - setParsed: true, - capabilities: &aggregator.AggregatedCapabilities{ - Tools: []vmcp.Tool{}, - CompositeTools: []vmcp.Tool{ - { - Name: "composite_tool", - Annotations: &vmcp.ToolAnnotations{ - IdempotentHint: boolPtr(true), - OpenWorldHint: boolPtr(false), - }, - }, - }, - }, - expectedAnnotation: &authorizers.ToolAnnotations{ - IdempotentHint: boolPtr(true), - OpenWorldHint: boolPtr(false), - }, - }, - { - name: "empty_resource_id_passes_through", - method: "tools/call", - resourceID: "", - setDiscovery: true, - setParsed: true, - capabilities: &aggregator.AggregatedCapabilities{ - Tools: []vmcp.Tool{ - { - Name: "my_tool", - Annotations: &vmcp.ToolAnnotations{ - ReadOnlyHint: boolPtr(true), - }, - }, - }, - }, - expectedAnnotation: nil, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - var capturedAnnotation *authorizers.ToolAnnotations - handlerCalled := false - - inner := http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { - handlerCalled = true - capturedAnnotation = authorizers.ToolAnnotationsFromContext(r.Context()) - }) - - wrapped := AnnotationEnrichmentMiddleware(inner) - - // Build a minimal request. The MCP body is not needed because we inject - // the parsed request directly into context below. - req := httptest.NewRequest(http.MethodPost, "/mcp", nil) - - ctx := req.Context() - - // Set parsed MCP request in context if needed. - // Uses the exported MCPRequestContextKey to inject the parsed request, - // matching how ParsingMiddleware stores it in production. - if tt.setParsed && tt.method != "" { - parsedReq := &mcpparser.ParsedMCPRequest{ - Method: tt.method, - ResourceID: tt.resourceID, - } - ctx = context.WithValue(ctx, mcpparser.MCPRequestContextKey, parsedReq) - } - - // Set discovery context if needed - if tt.setDiscovery && tt.capabilities != nil { - ctx = discovery.WithDiscoveredCapabilities(ctx, tt.capabilities) - } - - req = req.WithContext(ctx) - recorder := httptest.NewRecorder() - - wrapped.ServeHTTP(recorder, req) - - require.True(t, handlerCalled, "inner handler should always be called") - - if tt.expectedAnnotation == nil { - assert.Nil(t, capturedAnnotation, "expected no annotations in context") - } else { - require.NotNil(t, capturedAnnotation, "expected annotations in context") - assert.Equal(t, tt.expectedAnnotation.ReadOnlyHint, capturedAnnotation.ReadOnlyHint) - assert.Equal(t, tt.expectedAnnotation.DestructiveHint, capturedAnnotation.DestructiveHint) - assert.Equal(t, tt.expectedAnnotation.IdempotentHint, capturedAnnotation.IdempotentHint) - assert.Equal(t, tt.expectedAnnotation.OpenWorldHint, capturedAnnotation.OpenWorldHint) - } - }) - } -} diff --git a/pkg/vmcp/server/authz_integration_test.go b/pkg/vmcp/server/authz_integration_test.go new file mode 100644 index 0000000000..762efcd84a --- /dev/null +++ b/pkg/vmcp/server/authz_integration_test.go @@ -0,0 +1,208 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/authz/authorizers" + "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" + "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" + authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" + vmcpclient "github.com/stacklok/toolhive/pkg/vmcp/client" + discoveryMocks "github.com/stacklok/toolhive/pkg/vmcp/discovery/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/router" + "github.com/stacklok/toolhive/pkg/vmcp/server" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// newCedarAuthzTestServer builds a vMCP server via the rerouted server.New (→ core.New + +// Serve) backed by the real "echo" MCP backend at backendURL. An AuthMiddleware injects a +// fixed authenticated principal (the Cedar authorizer must resolve one — a nil identity +// denies on missing-principal), and a Cedar authz.Config compiled from policies is supplied +// via Config.Authz. This exercises the full Config.Authz → deriveCoreConfig → core admission +// chain that replaced the legacy HTTP authz middleware on the Serve path. +func newCedarAuthzTestServer(t *testing.T, backendURL string, policies ...string) *httptest.Server { + t.Helper() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) + mockBackendRegistry := mocks.NewMockBackendRegistry(ctrl) + + backend := vmcp.Backend{ + ID: "real-backend", + Name: "real-backend", + BaseURL: backendURL, + TransportType: "streamable-http", + } + mockBackendRegistry.EXPECT().List(gomock.Any()).Return([]vmcp.Backend{backend}).AnyTimes() + mockBackendRegistry.EXPECT().Get(gomock.Any(), gomock.Any()).Return(&backend).AnyTimes() + mockDiscoveryMgr.EXPECT().Discover(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil).AnyTimes() + mockDiscoveryMgr.EXPECT().Stop().AnyTimes() + + authReg := vmcpauth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, authReg.RegisterStrategy( + authtypes.StrategyTypeUnauthenticated, + strategies.NewUnauthenticatedStrategy(), + )) + factory := vmcpsession.NewSessionFactory(authReg) + + backendClient, err := vmcpclient.NewHTTPBackendClient(authReg) + require.NoError(t, err) + // A priority resolver keeps raw tool names ("echo", not "real-backend_echo") so the + // Cedar policies below can name Tool::"echo". + resolver, err := aggregator.NewPriorityConflictResolver([]string{backend.Name}) + require.NoError(t, err) + agg := aggregator.NewDefaultAggregator(backendClient, resolver, nil, nil) + + // Inject a fixed authenticated identity on every request so the session binds to it at + // initialize and the Cedar authorizer can resolve the principal on subsequent calls. + identityMiddleware := func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + id := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: "user-123", + Name: "Test User", + Claims: map[string]any{"sub": "user-123", "name": "Test User"}, + }} + next.ServeHTTP(w, r.WithContext(auth.WithIdentity(r.Context(), id))) + }) + } + + authzCfg, err := authorizers.NewConfig(cedar.Config{ + Version: "1.0", + Type: cedar.ConfigType, + Options: &cedar.ConfigOptions{Policies: policies, EntitiesJSON: "[]"}, + }) + require.NoError(t, err) + + srv, err := server.New( + context.Background(), + &server.Config{ + // Name is non-empty: deriveCoreConfig forwards the raw server name to the core, + // and the Cedar admission seam requires it (resource entities are scoped to + // MCP::""). This is the raw-name-for-authz parity the reroute preserves. + Name: "test-vmcp", + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + SessionFactory: factory, + Aggregator: agg, + AuthMiddleware: identityMiddleware, + Authz: authzCfg, + }, + router.NewDefaultRouter(), + backendClient, + mockDiscoveryMgr, + mockBackendRegistry, + nil, + ) + require.NoError(t, err) + + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + return ts +} + +// TestIntegration_RealBackend_CedarAuthzGatesToolCall proves the core admission seam — the +// authorization boundary the New/Serve reroute moved off the HTTP authz middleware and onto +// Config.Authz — enforces a Cedar policy end-to-end over HTTP through the rerouted server.New. +// It covers both gates the core now owns: the list filter (FilterTools) and the call gate +// (AllowToolCall, whose denial is genericized so the authorizer detail never leaks). +func TestIntegration_RealBackend_CedarAuthzGatesToolCall(t *testing.T) { + t.Parallel() + + t.Run("list filter: an unpermitted tool is neither advertised nor callable", func(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + // Permit only an unrelated tool: the principal resolves, but "echo" is default-denied, + // so the core's FilterTools drops it and it is never registered with the session. + ts := newCedarAuthzTestServer(t, backendURL, + `permit(principal, action == Action::"call_tool", resource == Tool::"unrelated");`) + + client := NewMCPTestClient(t, ts.URL) + client.InitializeSession() + + // "echo" is filtered out, so there is no list signal to wait on; poll the call until + // the session is established and the SDK rejects the unregistered tool (transient + // pre-registration 404s are tolerated and retried). + var lastBody string + require.Eventually(t, func() bool { + resp := client.CallTool("echo", map[string]any{"input": "hello"}) + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + lastBody = string(b) + return resp.StatusCode == http.StatusOK && bytes.Contains(b, []byte("not found")) + }, 5*time.Second, 50*time.Millisecond, + "a tool filtered out by the Cedar policy must be rejected as an unknown tool") + assert.NotContains(t, lastBody, "hello", "a filtered tool must never reach the backend") + + // And it must be absent from tools/list. + listResp := client.ListTools() + defer listResp.Body.Close() + listBody, err := io.ReadAll(listResp.Body) + require.NoError(t, err) + assert.NotContains(t, string(listBody), `"echo"`, + "a tool denied by policy must not be advertised in tools/list") + }) + + t.Run("call gate: an advertised tool's call is denied by an arg-gated policy", func(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + // "echo" is permitted (so it is advertised and callable), but a forbid clause denies + // the specific call whose "input" argument is "blocked". The forbid does not fire at + // list time (no call args), so the tool is still advertised — this is the only way to + // reach the AllowToolCall gate, since the list filter and call gate otherwise agree. + ts := newCedarAuthzTestServer(t, backendURL, + `permit(principal, action == Action::"call_tool", resource == Tool::"echo");`, + `forbid(principal, action == Action::"call_tool", resource == Tool::"echo") `+ + `when { context has arg_input && context.arg_input == "blocked" };`) + + client := NewMCPTestClient(t, ts.URL) + client.InitializeSession() + waitForEchoTool(t, ts.URL, client.SessionID()) + + // Positive control: a call whose args satisfy the policy reaches the backend. + okResp := client.CallTool("echo", map[string]any{"input": "hello"}) + defer okResp.Body.Close() + okBody, err := io.ReadAll(okResp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, okResp.StatusCode, "body: %s", string(okBody)) + assert.Contains(t, string(okBody), "hello", "a permitted call must return the backend result") + + // Denied call: the core admission seam rejects it with a generic message that never + // leaks the underlying authorizer/policy detail. + denyResp := client.CallTool("echo", map[string]any{"input": "blocked"}) + defer denyResp.Body.Close() + denyBody, err := io.ReadAll(denyResp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, denyResp.StatusCode, "body: %s", string(denyBody)) + assert.Contains(t, string(denyBody), "call denied by authorization policy", + "an arg-gated policy denial must surface the genericized authorization message") + assert.NotContains(t, string(denyBody), "cedar", "the authorizer implementation detail must not leak") + assert.NotContains(t, string(denyBody), "forbid", "the policy text must not leak") + }) +} diff --git a/pkg/vmcp/server/body_limit_test.go b/pkg/vmcp/server/body_limit_test.go index a675922927..391825e850 100644 --- a/pkg/vmcp/server/body_limit_test.go +++ b/pkg/vmcp/server/body_limit_test.go @@ -40,7 +40,7 @@ func TestHandler_RejectsOversizedBody(t *testing.T) { srv, err := server.New( t.Context(), - &server.Config{Host: "127.0.0.1", Port: 0, SessionFactory: newNoopMockFactory(t)}, + &server.Config{Host: "127.0.0.1", Port: 0, SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil)}, mockRouter, mockBackendClient, mockDiscoveryMgr, diff --git a/pkg/vmcp/server/derive.go b/pkg/vmcp/server/derive.go index 6666cabd7c..6282c61098 100644 --- a/pkg/vmcp/server/derive.go +++ b/pkg/vmcp/server/derive.go @@ -11,7 +11,6 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp/aggregator" "github.com/stacklok/toolhive/pkg/vmcp/composer" "github.com/stacklok/toolhive/pkg/vmcp/core" - "github.com/stacklok/toolhive/pkg/vmcp/health" "github.com/stacklok/toolhive/pkg/vmcp/router" "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" ) @@ -50,11 +49,8 @@ func WithDefaults(cfg *Config) *Config { // values and deriveServerConfig applies no defaulting of its own. Port 0 still means // "OS-assigned". // -// Three collaborators the transport needs but server.Config does not carry directly are +// Two collaborators the transport needs but server.Config does not carry directly are // passed in by the composition root rather than reached out of cfg: -// - healthMon: the *health.Monitor built at the composition root (A2). Serve owns only -// its Start/Stop; the same instance's StatusProvider is injected into the core via -// deriveCoreConfig. Nil disables health monitoring. // - backendRegistry: the shared backend registry (also passed to deriveCoreConfig); // the Serve session layer enumerates it when opening a session's connections. // - sessionManagerConfig: the pre-assembled *sessionmanager.FactoryConfig. The legacy @@ -69,7 +65,6 @@ func WithDefaults(cfg *Config) *Config { // it are removed later (#5445). func deriveServerConfig( cfg *Config, - healthMon *health.Monitor, backendRegistry vmcp.BackendRegistry, sessionManagerConfig *sessionmanager.FactoryConfig, ) *ServerConfig { @@ -86,7 +81,6 @@ func deriveServerConfig( AuthInfoHandler: cfg.AuthInfoHandler, PassthroughHeaders: cfg.PassthroughHeaders, AuthServer: cfg.AuthServer, - HealthMonitor: healthMon, StatusReportingInterval: cfg.StatusReportingInterval, StatusReporter: cfg.StatusReporter, Watcher: cfg.Watcher, @@ -116,11 +110,10 @@ func deriveServerConfig( // aggregator, router, backend client, backend registry, and workflow definitions that // arrive as separate server.New parameters today; the authz config that feeds the // admission seam (server.Config carries only the pre-built AuthzMiddleware, never the -// authz.Config, so it cannot come from cfg); the elicitation requester; and the health -// StatusProvider — the read-only view of the same *health.Monitor deriveServerConfig -// places on ServerConfig.HealthMonitor (A2). A nil healthProvider means no health -// filtering; a nil authzCfg means allow-all (matching today's AuthzMiddleware != nil -// guard). +// authz.Config, so it cannot come from cfg); the elicitation requester; and the backend +// health monitor configuration (cfg.HealthMonitorConfig), from which the core builds, runs, +// and stops the monitor it owns (#5443 reversal). A nil HealthMonitorConfig means no health +// filtering; a nil authzCfg means allow-all (matching today's AuthzMiddleware != nil guard). func deriveCoreConfig( cfg *Config, agg aggregator.Aggregator, @@ -130,7 +123,6 @@ func deriveCoreConfig( workflowDefs map[string]*composer.WorkflowDefinition, authzCfg *authz.Config, elicitation vmcp.ElicitationRequester, - healthProvider health.StatusProvider, ) *core.Config { return &core.Config{ Aggregator: agg, @@ -141,9 +133,9 @@ func deriveCoreConfig( Authz: authzCfg, ServerName: cfg.Name, // Cross-cutting (also on ServerConfig) — R3, not a clean partition: - TelemetryProvider: cfg.TelemetryProvider, - AuditConfig: cfg.AuditConfig, - HealthStatusProvider: healthProvider, - Elicitation: elicitation, + TelemetryProvider: cfg.TelemetryProvider, + AuditConfig: cfg.AuditConfig, + HealthMonitorConfig: cfg.HealthMonitorConfig, + Elicitation: elicitation, } } diff --git a/pkg/vmcp/server/derive_test.go b/pkg/vmcp/server/derive_test.go index bba4c6198a..3e72dff60a 100644 --- a/pkg/vmcp/server/derive_test.go +++ b/pkg/vmcp/server/derive_test.go @@ -25,15 +25,6 @@ import ( routermocks "github.com/stacklok/toolhive/pkg/vmcp/router/mocks" ) -// stubStatusProvider is a non-nil health.StatusProvider for derivation tests; its -// behavior is not exercised — only its identity is asserted (it is passed through -// onto core.Config.HealthStatusProvider). -type stubStatusProvider struct{} - -func (*stubStatusProvider) QueryBackendStatus(string) (vmcp.BackendHealthStatus, bool) { - return vmcp.BackendHealthy, true -} - // populatedLegacyConfig returns a server.Config with every field deriveServerConfig // reads set to a distinctive non-zero value, so a dropped or wrong-source mapping // surfaces in the assertions below. AuthzMiddleware is set too, to prove derivation @@ -67,11 +58,10 @@ func TestDeriveServerConfigProjectsTransportFields(t *testing.T) { t.Parallel() cfg := populatedLegacyConfig() - healthMon := &health.Monitor{} registry := vmcp.NewImmutableRegistry([]vmcp.Backend{}) smCfg := testMinimalSessionManagerConfig() - got := deriveServerConfig(cfg, healthMon, registry, smCfg) + got := deriveServerConfig(cfg, registry, smCfg) // Scalars projected verbatim (cfg's values are all non-default, so cmp.Or returns them). assert.Equal(t, "vmcp-name", got.Name) @@ -98,7 +88,6 @@ func TestDeriveServerConfigProjectsTransportFields(t *testing.T) { assert.Same(t, cfg.AuditConfig, got.AuditConfig) // Collaborators passed in (not on server.Config) — threaded through, not from cfg. - assert.Same(t, healthMon, got.HealthMonitor) assert.Same(t, registry, got.BackendRegistry) assert.Same(t, smCfg, got.SessionManagerConfig) } @@ -110,13 +99,12 @@ func TestDeriveServerConfigProjectsTransportFields(t *testing.T) { func TestDeriveServerConfigPropagatesNilCrossCutting(t *testing.T) { t.Parallel() - // A deliberately-disabled subsystem (nil provider/config/monitor) must stay nil — + // A deliberately-disabled subsystem (nil provider/config) must stay nil — // derivation must not substitute a default or panic. - got := deriveServerConfig(&Config{}, nil, nil, nil) + got := deriveServerConfig(&Config{}, nil, nil) assert.Nil(t, got.TelemetryProvider) assert.Nil(t, got.AuditConfig) - assert.Nil(t, got.HealthMonitor) assert.Nil(t, got.BackendRegistry) assert.Nil(t, got.SessionManagerConfig) } @@ -133,7 +121,6 @@ func TestDeriveServerConfigMapsAllFields(t *testing.T) { got := deriveServerConfig( populatedLegacyConfig(), - &health.Monitor{}, vmcp.NewImmutableRegistry([]vmcp.Backend{}), testMinimalSessionManagerConfig(), ) @@ -146,9 +133,10 @@ func TestDeriveCoreConfigAssemblesCollaborators(t *testing.T) { ctrl := gomock.NewController(t) cfg := &Config{ - Name: "core-name", - TelemetryProvider: &telemetry.Provider{}, - AuditConfig: &audit.Config{}, + Name: "core-name", + TelemetryProvider: &telemetry.Provider{}, + AuditConfig: &audit.Config{}, + HealthMonitorConfig: &health.MonitorConfig{}, } agg := aggmocks.NewMockAggregator(ctrl) rt := routermocks.NewMockRouter(ctrl) @@ -157,9 +145,8 @@ func TestDeriveCoreConfigAssemblesCollaborators(t *testing.T) { workflowDefs := map[string]*composer.WorkflowDefinition{"wf": {}} authzCfg := &authz.Config{} elicitation := vmcpmocks.NewMockElicitationRequester(ctrl) - healthProvider := &stubStatusProvider{} - got := deriveCoreConfig(cfg, agg, rt, backendClient, registry, workflowDefs, authzCfg, elicitation, healthProvider) + got := deriveCoreConfig(cfg, agg, rt, backendClient, registry, workflowDefs, authzCfg, elicitation) // Collaborators passed through by identity. assert.Same(t, agg, got.Aggregator) @@ -168,7 +155,7 @@ func TestDeriveCoreConfigAssemblesCollaborators(t *testing.T) { assert.Same(t, registry, got.BackendRegistry) assert.Same(t, authzCfg, got.Authz) assert.Same(t, elicitation, got.Elicitation) - assert.Same(t, healthProvider, got.HealthStatusProvider) + assert.Same(t, cfg.HealthMonitorConfig, got.HealthMonitorConfig) assert.Equal(t, workflowDefs, got.WorkflowDefs) // ServerName uses the raw cfg.Name for authz parity (no transport default applied). @@ -194,13 +181,12 @@ func TestDeriveCoreConfigUsesRawServerNameAndNilDefaults(t *testing.T) { nil, // workflowDefs nil, // authzCfg nil, // elicitation - nil, // healthProvider ) assert.Empty(t, got.ServerName) assert.Nil(t, got.Authz) assert.Nil(t, got.Elicitation) - assert.Nil(t, got.HealthStatusProvider) + assert.Nil(t, got.HealthMonitorConfig) assert.Nil(t, got.WorkflowDefs) } @@ -214,9 +200,10 @@ func TestDeriveCoreConfigMapsAllFields(t *testing.T) { ctrl := gomock.NewController(t) cfg := &Config{ - Name: "core-name", - TelemetryProvider: &telemetry.Provider{}, - AuditConfig: &audit.Config{}, + Name: "core-name", + TelemetryProvider: &telemetry.Provider{}, + AuditConfig: &audit.Config{}, + HealthMonitorConfig: &health.MonitorConfig{}, } got := deriveCoreConfig( @@ -228,7 +215,6 @@ func TestDeriveCoreConfigMapsAllFields(t *testing.T) { map[string]*composer.WorkflowDefinition{"wf": {}}, &authz.Config{}, vmcpmocks.NewMockElicitationRequester(ctrl), - &stubStatusProvider{}, ) assertAllFieldsPopulated(t, *got, "core.Config", nil) @@ -244,14 +230,14 @@ func TestDeriveConfigsTreatInputAsReadOnly(t *testing.T) { authzFn := func(h http.Handler) http.Handler { return h } cfg := &Config{AuthzMiddleware: authzFn} // empty scalars: defaulting would mutate them. - _ = deriveServerConfig(cfg, nil, nil, nil) + _ = deriveServerConfig(cfg, nil, nil) _ = deriveCoreConfig( cfg, aggmocks.NewMockAggregator(ctrl), routermocks.NewMockRouter(ctrl), vmcpmocks.NewMockBackendClient(ctrl), vmcpmocks.NewMockBackendRegistry(ctrl), - nil, nil, nil, nil, + nil, nil, nil, ) // Transport defaults were NOT written back onto the caller's Config. diff --git a/pkg/vmcp/server/elicitation_latebound.go b/pkg/vmcp/server/elicitation_latebound.go new file mode 100644 index 0000000000..6428d95c9e --- /dev/null +++ b/pkg/vmcp/server/elicitation_latebound.go @@ -0,0 +1,60 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "fmt" + "sync" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// lateBoundElicitationRequester is a vmcp.ElicitationRequester whose backing +// requester is set once, after Serve builds the mcp-go server. +// +// server.New evaluates core.New(deriveCoreConfig(...)) before Serve, but the +// SDK-backed elicitation adapter (NewSDKElicitationAdapter) wraps the *server.MCPServer +// that Serve creates — a construction-order inversion. The core needs a non-nil +// ElicitationRequester at construction (core.New rejects a nil one when a configured +// workflow contains an elicitation step), but the requester is not invoked until a +// composite workflow runs an elicitation step at request time — always after New has +// bound the real adapter and before the server begins serving. So a nil target is +// unreachable in practice; RequestElicitation guards it anyway. +// +// Safe for concurrent use: bind happens once during construction (before serving), +// RequestElicitation reads under the same lock. +type lateBoundElicitationRequester struct { + mu sync.RWMutex + target vmcp.ElicitationRequester +} + +var _ vmcp.ElicitationRequester = (*lateBoundElicitationRequester)(nil) + +func newLateBoundElicitationRequester() *lateBoundElicitationRequester { + return &lateBoundElicitationRequester{} +} + +// bind sets the backing requester. New calls it exactly once, after Serve returns +// and before the server starts serving. +func (l *lateBoundElicitationRequester) bind(target vmcp.ElicitationRequester) { + l.mu.Lock() + defer l.mu.Unlock() + l.target = target +} + +// RequestElicitation forwards to the bound requester, returning an error if invoked +// before bind — which would mean an elicitation fired during construction rather than +// at request time. +func (l *lateBoundElicitationRequester) RequestElicitation( + ctx context.Context, req vmcp.ElicitationRequest, +) (*vmcp.ElicitationResult, error) { + l.mu.RLock() + target := l.target + l.mu.RUnlock() + if target == nil { + return nil, fmt.Errorf("elicitation requested before the SDK server was bound") + } + return target.RequestElicitation(ctx, req) +} diff --git a/pkg/vmcp/server/elicitation_latebound_test.go b/pkg/vmcp/server/elicitation_latebound_test.go new file mode 100644 index 0000000000..6152ae2040 --- /dev/null +++ b/pkg/vmcp/server/elicitation_latebound_test.go @@ -0,0 +1,130 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// fakeVMCPElicitationRequester is a vmcp.ElicitationRequester that records the +// forwarded request and returns a canned result/error, so the late-bound wrapper's +// delegation can be asserted. Not safe for concurrent use — single-goroutine tests +// only (the concurrency test below binds a stateless requester instead). +type fakeVMCPElicitationRequester struct { + calls int + gotRequest vmcp.ElicitationRequest + result *vmcp.ElicitationResult + err error +} + +func (f *fakeVMCPElicitationRequester) RequestElicitation( + _ context.Context, req vmcp.ElicitationRequest, +) (*vmcp.ElicitationResult, error) { + f.calls++ + f.gotRequest = req + return f.result, f.err +} + +// nopElicitationRequester is a stateless vmcp.ElicitationRequester for the +// concurrency test, where a recording fake would itself race. +type nopElicitationRequester struct{} + +func (nopElicitationRequester) RequestElicitation( + context.Context, vmcp.ElicitationRequest, +) (*vmcp.ElicitationResult, error) { + return &vmcp.ElicitationResult{Action: "accept"}, nil +} + +// TestLateBoundElicitationRequester_RequestBeforeBind covers the guard path: an +// elicitation fired before bind (i.e. during construction rather than at request +// time) returns an error and does not panic dereferencing the nil target. +func TestLateBoundElicitationRequester_RequestBeforeBind(t *testing.T) { + t.Parallel() + + l := newLateBoundElicitationRequester() + res, err := l.RequestElicitation(context.Background(), vmcp.ElicitationRequest{Message: "hi"}) + require.Error(t, err) + assert.Nil(t, res) + assert.Contains(t, err.Error(), "before the SDK server was bound") +} + +// TestLateBoundElicitationRequester_DelegatesAfterBind asserts that, once bound, the +// wrapper forwards the request verbatim to the bound requester and returns its result +// and error unchanged — the contract the core relies on at workflow-elicitation time. +func TestLateBoundElicitationRequester_DelegatesAfterBind(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + result *vmcp.ElicitationResult + err error + wantErr bool + }{ + { + "accept result forwarded", + &vmcp.ElicitationResult{Action: "accept", Content: map[string]any{"k": "v"}}, nil, false, + }, + {"delegated error forwarded", nil, errors.New("transport closed"), true}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + target := &fakeVMCPElicitationRequester{result: tc.result, err: tc.err} + l := newLateBoundElicitationRequester() + l.bind(target) + + req := vmcp.ElicitationRequest{Message: "approve?", RequestedSchema: map[string]any{"type": "object"}} + res, err := l.RequestElicitation(context.Background(), req) + + assert.Equal(t, 1, target.calls, "request must reach the bound requester exactly once") + assert.Equal(t, req, target.gotRequest, "request must be forwarded verbatim") + if tc.wantErr { + require.Error(t, err) + assert.ErrorIs(t, err, tc.err) + return + } + require.NoError(t, err) + assert.Same(t, tc.result, res, "the bound requester's result must be returned unchanged") + }) + } +} + +// TestLateBoundElicitationRequester_ConcurrentBindAndRequest exercises the RWMutex +// guarding the target field: readers calling RequestElicitation concurrently with the +// one-time bind must not race. Run under -race to detect unsynchronized access; the +// returned value is intentionally unasserted since it depends on bind ordering. +func TestLateBoundElicitationRequester_ConcurrentBindAndRequest(t *testing.T) { + t.Parallel() + + l := newLateBoundElicitationRequester() + + const readers = 16 + var wg sync.WaitGroup + wg.Add(readers) + for range readers { + go func() { + defer wg.Done() + _, _ = l.RequestElicitation(context.Background(), vmcp.ElicitationRequest{Message: "x"}) + }() + } + l.bind(nopElicitationRequester{}) // the one-time write, concurrent with the readers above + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("timeout waiting for concurrent elicitation readers") + } +} diff --git a/pkg/vmcp/server/health_monitoring_test.go b/pkg/vmcp/server/health_monitoring_test.go index b9a996466d..3b1253a4da 100644 --- a/pkg/vmcp/server/health_monitoring_test.go +++ b/pkg/vmcp/server/health_monitoring_test.go @@ -44,7 +44,7 @@ func TestServer_HealthMonitoring_Disabled(t *testing.T) { Host: "127.0.0.1", Port: 0, HealthMonitorConfig: nil, // Health monitoring disabled - SessionFactory: testMinimalFactory(), + SessionFactory: testMinimalFactory(), Aggregator: &stubAggregator{}, } backendRegistry := vmcp.NewImmutableRegistry(backends) @@ -52,8 +52,8 @@ func TestServer_HealthMonitoring_Disabled(t *testing.T) { require.NoError(t, err) require.NotNil(t, srv) - // Verify health monitor is nil - assert.Nil(t, srv.healthMonitor) + // Verify health monitoring is disabled (the core exposes no reporter). + assert.Nil(t, srv.backendHealth()) // Verify getter methods return appropriate responses when disabled status, err := srv.GetBackendHealthStatus("backend-1") @@ -113,7 +113,7 @@ func TestServer_HealthMonitoring_Enabled(t *testing.T) { Timeout: 5 * time.Second, DegradedThreshold: 2 * time.Second, }, - SessionFactory: testMinimalFactory(), + SessionFactory: testMinimalFactory(), Aggregator: &stubAggregator{}, } backendRegistry := vmcp.NewImmutableRegistry(backends) @@ -121,8 +121,8 @@ func TestServer_HealthMonitoring_Enabled(t *testing.T) { require.NoError(t, err) require.NotNil(t, srv) - // Verify health monitor is created - assert.NotNil(t, srv.healthMonitor) + // Verify the core built the health monitor (started in core.New) and exposes it. + assert.NotNil(t, srv.backendHealth()) // Start server in background mockDiscoveryMgr.EXPECT().Stop().AnyTimes() @@ -208,7 +208,7 @@ func TestServer_HealthMonitoring_StartupFailure(t *testing.T) { UnhealthyThreshold: 0, // Invalid config - will cause monitor creation to fail Timeout: 50 * time.Millisecond, }, - SessionFactory: testMinimalFactory(), + SessionFactory: testMinimalFactory(), Aggregator: &stubAggregator{}, } // This should fail during New() because of invalid health monitor config @@ -241,7 +241,7 @@ func TestServer_HandleBackendHealth_Disabled(t *testing.T) { Host: "127.0.0.1", Port: 0, HealthMonitorConfig: nil, - SessionFactory: testMinimalFactory(), + SessionFactory: testMinimalFactory(), Aggregator: &stubAggregator{}, } backendRegistry := vmcp.NewImmutableRegistry(backends) @@ -300,7 +300,7 @@ func TestServer_HandleBackendHealth_Enabled(t *testing.T) { UnhealthyThreshold: 3, Timeout: 5 * time.Second, }, - SessionFactory: testMinimalFactory(), + SessionFactory: testMinimalFactory(), Aggregator: &stubAggregator{}, } backendRegistry := vmcp.NewImmutableRegistry(backends) @@ -395,7 +395,7 @@ func TestServer_Stop_StopsHealthMonitor(t *testing.T) { UnhealthyThreshold: 3, Timeout: 5 * time.Second, }, - SessionFactory: testMinimalFactory(), + SessionFactory: testMinimalFactory(), Aggregator: &stubAggregator{}, } backendRegistry := vmcp.NewImmutableRegistry(backends) @@ -420,10 +420,8 @@ func TestServer_Stop_StopsHealthMonitor(t *testing.T) { return statusErr == nil && status == vmcp.BackendHealthy }, 2*time.Second, 10*time.Millisecond, "backend-1 should become healthy") - // Verify health monitor is running - srv.healthMonitorMu.RLock() - assert.NotNil(t, srv.healthMonitor) - srv.healthMonitorMu.RUnlock() + // Verify health monitor is running (owned by the core) + assert.NotNil(t, srv.backendHealth()) // Cancel context to trigger graceful shutdown cancel() @@ -436,11 +434,9 @@ func TestServer_Stop_StopsHealthMonitor(t *testing.T) { t.Fatal("timeout waiting for server to stop") } - // Verify health monitor still exists after stop (not set to nil) - // The monitor is stopped but the pointer remains valid - srv.healthMonitorMu.RLock() - assert.NotNil(t, srv.healthMonitor, "health monitor should still exist after stop") - srv.healthMonitorMu.RUnlock() + // Verify the reporter still exists after stop (core.Close stops the monitor but does not + // nil it), so the getter methods below keep working against the stopped monitor. + assert.NotNil(t, srv.backendHealth(), "health reporter should still exist after stop") // Verify getter methods still work (they query the stopped monitor) // This ensures no panics occur when accessing a stopped monitor diff --git a/pkg/vmcp/server/health_test.go b/pkg/vmcp/server/health_test.go index d89bae491f..9a3ac48270 100644 --- a/pkg/vmcp/server/health_test.go +++ b/pkg/vmcp/server/health_test.go @@ -71,7 +71,7 @@ func createTestServer(t *testing.T) *server.Server { Version: "1.0.0", Host: "127.0.0.1", Port: port, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, rt, mockBackendClient, mockDiscoveryMgr, backendRegistry, nil) require.NoError(t, err) @@ -186,7 +186,7 @@ func TestServer_SessionManager(t *testing.T) { Name: "test-vmcp", Version: "1.0.0", SessionTTL: 10 * time.Minute, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, rt, mockBackendClient, mockDiscoveryMgr, backendRegistry, nil) require.NoError(t, err) @@ -211,7 +211,7 @@ func TestServer_SessionManager(t *testing.T) { Name: "test-vmcp", Version: "1.0.0", SessionTTL: customTTL, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, rt, mockBackendClient, mockDiscoveryMgr, backendRegistry, nil) require.NoError(t, err) diff --git a/pkg/vmcp/server/integration_test.go b/pkg/vmcp/server/integration_test.go index d6c4325d45..e9753230bd 100644 --- a/pkg/vmcp/server/integration_test.go +++ b/pkg/vmcp/server/integration_test.go @@ -209,6 +209,7 @@ func TestIntegration_AggregatorToRouterToServer(t *testing.T) { Host: "127.0.0.1", Port: 4484, SessionFactory: newNoopMockFactory(t), + Aggregator: agg, }, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) require.NoError(t, err) @@ -515,7 +516,12 @@ func TestIntegration_AuditLogging(t *testing.T) { mock.EXPECT().Type().Return(transportsession.SessionType("")).AnyTimes() mock.EXPECT().GetData().Return(nil).AnyTimes() mock.EXPECT().SetData(gomock.Any()).AnyTimes() - mock.EXPECT().GetMetadata().Return(map[string]string{}).AnyTimes() + mock.EXPECT().GetMetadata().Return(map[string]string{ + vmcpsession.MetadataKeyIdentityBinding: "unauthenticated", + }).AnyTimes() + // Serve-path enforceSessionBinding reads the binding via GetMetadataValue. + mock.EXPECT().GetMetadataValue(vmcpsession.MetadataKeyIdentityBinding). + Return("unauthenticated", true).AnyTimes() mock.EXPECT().SetMetadata(gomock.Any(), gomock.Any()).AnyTimes() mock.EXPECT().Tools().Return(auditTools).AnyTimes() mock.EXPECT().Resources().Return(nil).AnyTimes() @@ -530,11 +536,18 @@ func TestIntegration_AuditLogging(t *testing.T) { return mock, nil }).AnyTimes() + // The core sources the advertised set by aggregating over mockBackendClient with the + // same prefix resolver the legacy discovery path used, so tools/call and resources/read + // route through the core and are audit-logged with the prefixed names. + auditAgg := aggregator.NewDefaultAggregator( + mockBackendClient, aggregator.NewPrefixConflictResolver("{workload}_"), nil, nil) + srv, err := server.New(ctx, &server.Config{ Host: "127.0.0.1", Port: 0, // Random port AuditConfig: auditConfig, SessionFactory: auditSessionFactory, + Aggregator: auditAgg, }, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) require.NoError(t, err) @@ -566,7 +579,9 @@ func TestIntegration_AuditLogging(t *testing.T) { // Test 1: Initialize request should be logged t.Run("initialize request is logged", func(t *testing.T) { initReq := map[string]any{ - "method": "initialize", + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", "params": map[string]any{ "protocolVersion": "2024-11-05", "capabilities": map[string]any{}, @@ -605,7 +620,9 @@ func TestIntegration_AuditLogging(t *testing.T) { require.NotEmpty(t, sessionID, "Session ID must be set from initialize test") toolsReq := map[string]any{ - "method": "tools/list", + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", } reqBody, err := json.Marshal(toolsReq) @@ -633,7 +650,9 @@ func TestIntegration_AuditLogging(t *testing.T) { require.NotEmpty(t, sessionID, "Session ID must be set from initialize test") toolCallReq := map[string]any{ - "method": "tools/call", + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", "params": map[string]any{ "name": "weather-service_get_weather", // Prefix added by aggregator "arguments": map[string]any{ @@ -675,7 +694,9 @@ func TestIntegration_AuditLogging(t *testing.T) { require.NotEmpty(t, sessionID, "Session ID must be set from initialize test") resourceReq := map[string]any{ - "method": "resources/read", + "jsonrpc": "2.0", + "id": 4, + "method": "resources/read", "params": map[string]any{ "uri": "weather://current", }, @@ -813,6 +834,7 @@ func TestIntegration_AuditLoggingWithAuth(t *testing.T) { AuditConfig: auditConfig, AuthMiddleware: identityMiddleware, SessionFactory: newNoopMockFactory(t), + Aggregator: newStubAggregator(nil), }, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) require.NoError(t, err) diff --git a/pkg/vmcp/server/readiness_test.go b/pkg/vmcp/server/readiness_test.go index bbd0592b20..badb4c6ed6 100644 --- a/pkg/vmcp/server/readiness_test.go +++ b/pkg/vmcp/server/readiness_test.go @@ -69,7 +69,7 @@ func TestReadinessEndpoint_StaticMode(t *testing.T) { Host: "127.0.0.1", Port: port, Watcher: nil, // Static mode - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry([]vmcp.Backend{}), nil) require.NoError(t, err) @@ -145,7 +145,7 @@ func TestReadinessEndpoint_DynamicMode_CacheSynced(t *testing.T) { Host: "127.0.0.1", Port: port, Watcher: mockWatcher, // Dynamic mode with synced cache - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewDynamicRegistry([]vmcp.Backend{}), nil) require.NoError(t, err) @@ -221,7 +221,7 @@ func TestReadinessEndpoint_DynamicMode_CacheNotSynced(t *testing.T) { Host: "127.0.0.1", Port: port, Watcher: mockWatcher, // Dynamic mode with unsynced cache - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewDynamicRegistry([]vmcp.Backend{}), nil) require.NoError(t, err) diff --git a/pkg/vmcp/server/serve.go b/pkg/vmcp/server/serve.go index 5d624f19b8..a4891547d1 100644 --- a/pkg/vmcp/server/serve.go +++ b/pkg/vmcp/server/serve.go @@ -19,7 +19,6 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" "github.com/stacklok/toolhive/pkg/vmcp/core" - "github.com/stacklok/toolhive/pkg/vmcp/health" "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" vmcpstatus "github.com/stacklok/toolhive/pkg/vmcp/status" ) @@ -85,11 +84,6 @@ type ServerConfig struct { // routes are registered on the mux alongside the protected resource metadata. AuthServer *asrunner.EmbeddedAuthServer - // HealthMonitor is the already-built backend health monitor (constructed at the - // composition root so the same instance's StatusProvider can be injected into the - // core). Serve owns only its Start/Stop lifecycle. Nil disables health monitoring. - HealthMonitor *health.Monitor - // StatusReportingInterval is the interval for reporting status updates. // If zero, the default reporting interval is used. StatusReportingInterval time.Duration @@ -176,15 +170,11 @@ type ServerConfig struct { // "/" MCP route serveable: the shared Handler guards the discovery middleware to the // legacy path (s.core == nil), and on the Serve path session registration and request // handlers call the core directly (#5442). The last transport subsystems — the embedded -// AS runner routes, the status reporter (and its periodic goroutine), the optimizer -// cleanup, and the health monitor's Start/Stop (#5443) — are driven from the -// carried-forward shared (*Server).Handler/Start/Stop using the fields Serve populates -// from ServerConfig below. Serve does NOT construct the health monitor: it receives the -// pre-built *health.Monitor via ServerConfig (nil ⇒ disabled) and owns only its -// lifecycle, while the composition root injects the same instance into the core as a -// health.StatusProvider. server.New is not yet routed through Serve and keeps its own -// copy of the session wiring until Phase 3, so this is purely additive and observable -// behavior is unchanged. +// AS runner routes, the status reporter (and its periodic goroutine), and the optimizer +// cleanup — are driven from the carried-forward shared (*Server).Handler/Start/Stop using +// the fields Serve populates from ServerConfig below. The backend health monitor is owned +// by the core (built/started in core.New, stopped in core.Close); Serve neither constructs +// nor starts it, and the status reporter reads it back through the core (#5443 reversal). // // Serve returns a vmcp.ErrInvalidConfig-wrapped error for a nil cfg, a nil core, or a // nil required collaborator (SessionManagerConfig or BackendRegistry). The session @@ -195,9 +185,9 @@ type ServerConfig struct { // but a nil discovery manager, router, and backend client — the request path goes // through the core, so those legacy collaborators are unused on the Serve path. The // shared Handler skips the discovery middleware when s.core != nil, so serving the "/" -// MCP route no longer nil-derefs. The embedded AS runner routes, the status reporter, -// the optimizer cleanup, and the health monitor's Start/Stop are driven from the shared -// Handler/Start/Stop via the fields Serve populates from ServerConfig (#5443). +// MCP route no longer nil-derefs. The embedded AS runner routes, the status reporter, and +// the optimizer cleanup are driven from the shared Handler/Start/Stop via the fields Serve +// populates from ServerConfig; the backend health monitor is owned by the core (#5443). func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) { if cfg == nil { return nil, fmt.Errorf("%w: nil server config", vmcp.ErrInvalidConfig) @@ -277,7 +267,6 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) // session manager (optimizerCleanup, appended to shutdownFuncs below). optimizerFactory: vmcpSessMgr.OptimizerFactory(), ready: make(chan struct{}), - healthMonitor: cfg.HealthMonitor, statusReporter: cfg.StatusReporter, } @@ -323,12 +312,12 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) // (*Server).Handler omits both the authz and annotation-enrichment blocks when // AuthzMiddleware is nil; authorization moves to the core admission seam (#5438). // The inert blocks stay in the shared Handler until Phase 3 (#5445) removes them. -// - HealthMonitorConfig: Serve receives the already-built *health.Monitor via -// ServerConfig.HealthMonitor (A2) and assigns it to the Server directly, so it -// never needs the monitor's construction config. +// - HealthMonitorConfig: the backend health monitor is owned by the core (built from the +// composition root's server.Config.HealthMonitorConfig in core.New, stopped in +// core.Close), so the Serve-path Server never needs the monitor or its config. // - StatusReporter: Serve assigns ServerConfig.StatusReporter to the Server field -// directly (like HealthMonitor); nothing reads Config.StatusReporter on the Serve -// path, so mapping it here would be dead. +// directly; nothing reads Config.StatusReporter on the Serve path, so mapping it here +// would be dead. // - SessionFactory, OptimizerFactory, OptimizerConfig: the vMCP session manager is // built directly in Serve from ServerConfig.SessionManagerConfig (a pre-built // *sessionmanager.FactoryConfig that carries the session factory and optimizer diff --git a/pkg/vmcp/server/serve_lifecycle_test.go b/pkg/vmcp/server/serve_lifecycle_test.go index 280ee7b2ad..609be45eee 100644 --- a/pkg/vmcp/server/serve_lifecycle_test.go +++ b/pkg/vmcp/server/serve_lifecycle_test.go @@ -12,12 +12,9 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "go.uber.org/mock/gomock" asrunner "github.com/stacklok/toolhive/pkg/authserver/runner" "github.com/stacklok/toolhive/pkg/vmcp" - "github.com/stacklok/toolhive/pkg/vmcp/health" - "github.com/stacklok/toolhive/pkg/vmcp/mocks" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" ) @@ -75,123 +72,6 @@ func startServeInBackground(t *testing.T, srv *Server) (stop func() error) { return stop } -// TestServeHealthMonitorDisabledWhenNil verifies that a nil ServerConfig.HealthMonitor -// leaves monitoring disabled on the Serve path: Serve stores no monitor and the getters -// report "disabled" without error, matching the no-monitor behavior of server.New. -func TestServeHealthMonitorDisabledWhenNil(t *testing.T) { - t.Parallel() - - srv, err := Serve(context.Background(), &stubVMCP{}, testMinimalServeConfig()) - require.NoError(t, err) - t.Cleanup(func() { _ = srv.Stop(context.Background()) }) - - assert.Nil(t, srv.healthMonitor, "nil ServerConfig.HealthMonitor must leave the monitor unset") - - status, err := srv.GetBackendHealthStatus("backend-1") - require.Error(t, err) - assert.Equal(t, vmcp.BackendUnknown, status) - assert.Contains(t, err.Error(), "health monitoring is disabled") - - assert.Equal(t, health.Summary{}, srv.GetHealthSummary()) -} - -// TestServeStartsAndStopsInjectedHealthMonitor verifies the health-monitor lifecycle on -// the Serve path: Serve stores the pre-built *health.Monitor it is given (it does NOT -// construct one), Start runs it (backends report healthy), and Stop stops it while -// retaining the struct so getters keep working — mirroring TestServer_Stop_StopsHealthMonitor -// for a monitor Serve was handed rather than one it built. -func TestServeStartsAndStopsInjectedHealthMonitor(t *testing.T) { - t.Parallel() - - ctrl := gomock.NewController(t) - mockBackendClient := mocks.NewMockBackendClient(ctrl) - mockBackendClient.EXPECT(). - ListCapabilities(gomock.Any(), gomock.Any()). - Return(&vmcp.CapabilityList{}, nil). - AnyTimes() - - backends := []vmcp.Backend{ - {ID: "backend-1", Name: "Backend 1", BaseURL: "http://localhost:8080", TransportType: "sse"}, - } - mon, err := health.NewMonitor(mockBackendClient, backends, health.MonitorConfig{ - CheckInterval: 50 * time.Millisecond, - UnhealthyThreshold: 1, - Timeout: 5 * time.Second, - }) - require.NoError(t, err) - - cfg := testMinimalServeConfig() - cfg.HealthMonitor = mon - cfg.BackendRegistry = vmcp.NewImmutableRegistry(backends) - - srv, err := Serve(context.Background(), &stubVMCP{}, cfg) - require.NoError(t, err) - - // Serve must reuse the injected instance, not build a new one (AC2: Serve does not - // call health.NewMonitor). - assert.Same(t, mon, srv.healthMonitor) - - stop := startServeInBackground(t, srv) - - require.Eventually(t, func() bool { - status, statusErr := srv.GetBackendHealthStatus("backend-1") - return statusErr == nil && status == vmcp.BackendHealthy - }, 2*time.Second, 10*time.Millisecond, "backend-1 should become healthy via the Serve-started monitor") - - require.NoError(t, stop()) - - // The monitor is stopped but the struct is retained (pointer stays valid). - srv.healthMonitorMu.RLock() - assert.Same(t, mon, srv.healthMonitor, "health monitor should still exist after Stop") - srv.healthMonitorMu.RUnlock() - - status, err := srv.GetBackendHealthStatus("backend-1") - assert.NoError(t, err, "getter should not error after Stop") - assert.NotEqual(t, vmcp.BackendUnknown, status, "should return last known status") -} - -// TestServeDisablesHealthMonitorOnStartFailure verifies graceful degradation when the -// injected monitor cannot start: Start logs a WARN and sets healthMonitor to nil so the -// getters report "disabled", and Serve.Start itself does not fail. The failure is forced -// with a monitor that has already been stopped — health.Monitor.Start refuses to restart. -func TestServeDisablesHealthMonitorOnStartFailure(t *testing.T) { - t.Parallel() - - ctrl := gomock.NewController(t) - mockBackendClient := mocks.NewMockBackendClient(ctrl) - - // No backends => Start spawns no health-check goroutines (ListCapabilities is never - // called) and Stop returns immediately, leaving the monitor in the "stopped" state. - mon, err := health.NewMonitor(mockBackendClient, nil, health.MonitorConfig{ - CheckInterval: time.Second, - UnhealthyThreshold: 1, - Timeout: time.Second, - }) - require.NoError(t, err) - require.NoError(t, mon.Start(context.Background())) - require.NoError(t, mon.Stop()) - - cfg := testMinimalServeConfig() - cfg.HealthMonitor = mon - - srv, err := Serve(context.Background(), &stubVMCP{}, cfg) - require.NoError(t, err) - - stop := startServeInBackground(t, srv) - - // Start must not fail; the un-restartable monitor is disabled (set to nil) instead. - require.Eventually(t, func() bool { - srv.healthMonitorMu.RLock() - defer srv.healthMonitorMu.RUnlock() - return srv.healthMonitor == nil - }, 2*time.Second, 10*time.Millisecond, "a monitor whose Start fails must be disabled") - - _, getErr := srv.GetBackendHealthStatus("backend-1") - assert.ErrorContains(t, getErr, "health monitoring is disabled") - - require.NoError(t, stop()) -} - // recordingReporter is a vmcpstatus.Reporter that records whether Start and its returned // shutdown func ran and counts ReportStatus calls, so the status-reporter lifecycle can // be asserted on the Serve path. diff --git a/pkg/vmcp/server/serve_session_test.go b/pkg/vmcp/server/serve_session_test.go index 0dc7e640f4..53653f1870 100644 --- a/pkg/vmcp/server/serve_session_test.go +++ b/pkg/vmcp/server/serve_session_test.go @@ -28,6 +28,7 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" "github.com/stacklok/toolhive/pkg/vmcp/core" + "github.com/stacklok/toolhive/pkg/vmcp/health" "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" sessionfactorymocks "github.com/stacklok/toolhive/pkg/vmcp/session/mocks" @@ -188,6 +189,8 @@ func (*fakeCore) LookupPrompt(context.Context, *auth.Identity, string) (*vmcp.Pr func (*fakeCore) Close() error { return nil } +func (*fakeCore) BackendHealth() health.Reporter { return nil } + // TestServeRegistersSessionHooks verifies that Serve registers the OnRegisterSession // hook and wires two-phase session creation: an MCP initialize triggers the hook, // which creates the session record (MakeSessionWithID) and injects the core's advertised diff --git a/pkg/vmcp/server/serve_test.go b/pkg/vmcp/server/serve_test.go index 0cfb344c54..e8a84113e3 100644 --- a/pkg/vmcp/server/serve_test.go +++ b/pkg/vmcp/server/serve_test.go @@ -63,7 +63,8 @@ func (*stubVMCP) LookupResource(context.Context, *auth.Identity, string) (*vmcp. func (*stubVMCP) LookupPrompt(context.Context, *auth.Identity, string) (*vmcp.Prompt, error) { return nil, nil } -func (s *stubVMCP) Close() error { s.closed = true; return nil } +func (s *stubVMCP) Close() error { s.closed = true; return nil } +func (*stubVMCP) BackendHealth() health.Reporter { return nil } // stubWatcher is a non-nil Watcher for the drift-guard test; its behavior is not // exercised there. @@ -320,6 +321,8 @@ func TestBuildServeConfigMapsSharedFields(t *testing.T) { "SessionFactory": {}, // session manager built in Serve from ServerConfig.SessionManagerConfig "OptimizerFactory": {}, // optimizer wiring carried on ServerConfig.SessionManagerConfig (FactoryConfig) "OptimizerConfig": {}, // optimizer wiring carried on ServerConfig.SessionManagerConfig (FactoryConfig) + "Aggregator": {}, // core collaborator: fed to core.New via deriveCoreConfig, not the transport + "Authz": {}, // core collaborator: fed to the core admission seam via deriveCoreConfig } // Every field set to a non-zero value so a dropped mapping surfaces as a zero @@ -339,7 +342,6 @@ func TestBuildServeConfigMapsSharedFields(t *testing.T) { AuthInfoHandler: http.NewServeMux(), PassthroughHeaders: []string{"x-test"}, AuthServer: &asrunner.EmbeddedAuthServer{}, - HealthMonitor: &health.Monitor{}, StatusReportingInterval: time.Second, StatusReporter: stubServeReporter{}, Watcher: stubWatcher{}, diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index e54e868696..75d324d188 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -21,13 +21,13 @@ import ( "sync" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" tcredis "github.com/stacklok/toolhive-core/redis" "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/auth" asrunner "github.com/stacklok/toolhive/pkg/authserver/runner" + "github.com/stacklok/toolhive/pkg/authz" "github.com/stacklok/toolhive/pkg/bodylimit" mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/recovery" @@ -35,16 +35,15 @@ import ( transportmiddleware "github.com/stacklok/toolhive/pkg/transport/middleware" transportsession "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" "github.com/stacklok/toolhive/pkg/vmcp/composer" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" "github.com/stacklok/toolhive/pkg/vmcp/core" "github.com/stacklok/toolhive/pkg/vmcp/discovery" "github.com/stacklok/toolhive/pkg/vmcp/headerforward" "github.com/stacklok/toolhive/pkg/vmcp/health" - "github.com/stacklok/toolhive/pkg/vmcp/internal/backendtelemetry" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" "github.com/stacklok/toolhive/pkg/vmcp/router" - "github.com/stacklok/toolhive/pkg/vmcp/server/adapter" "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" vmcpstatus "github.com/stacklok/toolhive/pkg/vmcp/status" @@ -86,6 +85,13 @@ const ( defaultServerVersion = "0.1.0" defaultHost = "127.0.0.1" defaultEndpointPath = "/mcp" + + // capabilityCacheTTL bounds how long the per-identity aggregated capability view is + // reused before re-sweeping backends. The core re-aggregates on every call (it holds no + // per-session cache), so without this the Serve path does a full backend tools/list sweep + // per tool call; a short TTL collapses bursts of calls into ~one sweep while keeping + // staleness tighter than the legacy once-per-session refresh. + capabilityCacheTTL = 30 * time.Second ) //go:generate mockgen -destination=mocks/mock_watcher.go -package=mocks -source=server.go Watcher @@ -204,6 +210,22 @@ type Config struct { // session persistence; the Redis password is read from the // THV_SESSION_REDIS_PASSWORD environment variable. SessionStorage *vmcpconfig.SessionStorageConfig + + // Aggregator merges backend capabilities into the advertised set. It is consumed + // by the domain core (core.Config.Aggregator) that New builds via deriveCoreConfig: + // the core is the single aggregator on the New/Serve path, so the session factory + // must NOT also aggregate. Required; New fails (via core.New) when nil. The + // composition root supplies it (the same aggregator that backs discovery). + Aggregator aggregator.Aggregator + + // Authz is the authorizer-agnostic authorization config fed to the core admission + // seam (core.Config.Authz). A nil value means authorization is unconfigured + // (allow-all), matching the legacy `AuthzMiddleware != nil` guard: the composition + // root populates it only when Cedar policies exist. It is distinct from the + // (vestigial) AuthzMiddleware field — the core enforces authz from this config, not + // from the HTTP middleware. When non-nil, Name must be non-empty (the Cedar resource + // entity name). + Authz *authz.Config } // Server is the Virtual MCP Server that aggregates multiple backends. @@ -239,16 +261,9 @@ type Server struct { listener net.Listener listenerMu sync.RWMutex - // Router for forwarding requests to backends - router router.Router - - // Backend client for making requests to backends - backendClient vmcp.BackendClient - - // Handler factory for creating MCP request handlers - handlerFactory *adapter.DefaultHandlerFactory - - // Discovery manager for lazy per-user capability discovery + // Discovery manager — retained only so Stop() drives its (now no-op) cleanup. The core + // replaced discovery's aggregation on the New/Serve path, so New no longer wires the + // router / backend client / handler factory / capability adapter the legacy body held. discoveryMgr discovery.Manager // Backend registry for capability discovery @@ -269,9 +284,6 @@ type Server struct { // Redis-backed storage for multi-pod deployments is not yet wired. sessionDataStorage transportsession.DataStorage - // Capability adapter for converting aggregator types to SDK types - capabilityAdapter *adapter.CapabilityAdapter - // vmcpSessionMgr manages session-scoped backend client lifecycle. vmcpSessionMgr SessionManager @@ -280,13 +292,6 @@ type Server struct { ready chan struct{} readyOnce sync.Once - // healthMonitor performs periodic health checks on backends. - // Nil if health monitoring is disabled. - // Protected by healthMonitorMu: RLock for reads (getter methods, HTTP handlers), - // Lock for writes (initialization, disabling on start failure). - healthMonitor *health.Monitor - healthMonitorMu sync.RWMutex - // statusReporter enables vMCP to report operational status to control plane. // Nil if status reporting is disabled. statusReporter vmcpstatus.Reporter @@ -332,13 +337,29 @@ func buildSessionDataStorage(ctx context.Context, cfg *Config) (transportsession return transportsession.NewRedisSessionDataStorage(ctx, redisCfg, keyPrefix, cfg.SessionTTL) } -// New creates a new Virtual MCP Server instance. +// New creates a Virtual MCP Server by composing the domain core (core.New) with the +// transport (Serve). +// +// It is the composition root for the in-memory Config form: it builds the backend +// health monitor (A2), assembles the core via core.New using the config projected by +// deriveCoreConfig, and hands that core to Serve with the transport config projected by +// deriveServerConfig. The 7-param signature is retained unchanged for existing callers +// (cli/serve.go and external embedders); the transport/core wiring it once performed +// inline now lives behind core.New + Serve. // // The backendRegistry parameter provides the list of available backends: // - For static mode (CLI), pass an immutable registry created from initial backends // - For dynamic mode (K8s), pass a DynamicRegistry that will be updated by the operator // -//nolint:gocyclo // Complexity from hook logic is acceptable +// Runtime-contract change for embedders (the signature is unchanged, but the body now +// routes through core.New + Serve): +// - Config.Aggregator is now REQUIRED (core.New rejects a nil aggregator); the core is +// the single source of the advertised capability set. +// - Config.AuthzMiddleware is vestigial: the Serve path never applies it. Authorization +// is enforced by the core admission seam from Config.Authz. An embedder that enforced +// Cedar authz only via AuthzMiddleware must set Config.Authz instead — New returns an +// error (ErrInvalidConfig) if AuthzMiddleware is set without Authz, so a silent +// allow-all fails fast at construction. func New( ctx context.Context, cfg *Config, @@ -348,202 +369,103 @@ func New( backendRegistry vmcp.BackendRegistry, workflowDefs map[string]*composer.WorkflowDefinition, ) (*Server, error) { - // Resolve transport defaults on a COPY. The composition root (cli) already resolves - // them at the edge via WithDefaults (a single defaulting list); New repeats - // it defensively so legacy direct callers and tests that build a Config by hand keep - // working — but without mutating the caller's value (go-style: copy before mutating - // caller input). That non-mutation is what lets #5445 hand the raw, un-defaulted - // cfg.Name to the core for Cedar authz parity. New's own defaulting goes away when - // #5445 reduces it to a Serve(core.New(...)) wrapper; until then WithDefaults is the - // single place the default list lives, shared with the edge. - cfg = WithDefaults(cfg) - - // Create hooks for SDK integration - hooks := &server.Hooks{} - - // Create mark3labs MCP server - mcpServer := server.NewMCPServer( - cfg.Name, - cfg.Version, - server.WithToolCapabilities(false), // We'll register tools dynamically - server.WithResourceCapabilities(false, false), // We'll register resources dynamically - server.WithLogging(), - server.WithHooks(hooks), - ) - - // Create SDK elicitation adapter for workflow engine - // This wraps the mark3labs SDK to provide elicitation functionality to the composer - sdkElicitationRequester := NewSDKElicitationAdapter(mcpServer) - - // Create elicitation handler for workflow engine - // This provides SDK-agnostic elicitation with security validation - elicitationHandler := composer.NewDefaultElicitationHandler(sdkElicitationRequester) - - // Decorate backend client with telemetry if provider is configured - // This must happen BEFORE creating the workflow engine so that workflow - // backend calls are instrumented when they occur during workflow execution. - if cfg.TelemetryProvider != nil { - var err error - // Get initial backends list from registry for telemetry setup - initialBackends := backendRegistry.List(ctx) - backendClient, err = backendtelemetry.MonitorBackends( - ctx, - cfg.TelemetryProvider.MeterProvider(), - cfg.TelemetryProvider.TracerProvider(), - initialBackends, - backendClient, - ) - if err != nil { - return nil, fmt.Errorf("failed to monitor backends: %w", err) - } - } - - // Create workflow auditor if audit config is provided - var workflowAuditor *audit.WorkflowAuditor - if cfg.AuditConfig != nil { - if err := cfg.AuditConfig.Validate(); err != nil { - return nil, fmt.Errorf("invalid audit configuration: %w", err) - } - var err error - workflowAuditor, err = audit.NewWorkflowAuditor(cfg.AuditConfig) - if err != nil { - return nil, fmt.Errorf("failed to create workflow auditor: %w", err) - } - slog.Info("workflow audit logging enabled") - } - - // Create workflow engine (composer) for executing composite tools - // The composer orchestrates multi-step workflows across backends - // Use in-memory state store with 5-minute cleanup interval and 1-hour max age for completed workflows - stateStore := composer.NewInMemoryStateStore(5*time.Minute, 1*time.Hour) - workflowComposer := composer.NewWorkflowEngine(rt, backendClient, elicitationHandler, stateStore, workflowAuditor, nil) - - // composerFactory builds a per-session workflow engine at session registration - // time, binding composite tool routing to the session's own routing table and - // tool list. This removes composite tools' dependency on the discovery middleware - // injecting DiscoveredCapabilities into the request context. - sessionComposerFactory := func(sessionRT *vmcp.RoutingTable, sessionTools []vmcp.Tool) composer.Composer { - return composer.NewWorkflowEngine( - router.NewSessionRouter(sessionRT), backendClient, elicitationHandler, stateStore, workflowAuditor, - sessionTools, - ) - } - - // Validate workflows (fail fast on invalid definitions) - var err error - workflowDefs, err = validateWorkflows(workflowComposer, workflowDefs) - if err != nil { - return nil, fmt.Errorf("workflow validation failed: %w", err) - } - - // Create session manager using StreamableSession as the transport-layer placeholder. - // StreamableSession is a lightweight implementation of transportsession.Session that - // handles disconnect tracking, TTL, and metadata for Streamable HTTP connections. - // It intentionally carries no vmcp-specific state — backend connections, routing - // tables, tool lists, and token binding all live in the separate sessionmanager.Manager, - // keyed by the same session ID. - sessionManager := transportsession.NewManager(cfg.SessionTTL, transportsession.NewStreamableSession) - - sessionDataStorage, err := buildSessionDataStorage(ctx, cfg) + // Resolve transport defaults for the transport side on a COPY (WithDefaults does not + // mutate cfg). The core receives the RAW cfg so its ServerName is the un-defaulted + // cfg.Name: Cedar authz keys on the real VirtualMCPServer name, not the synthetic + // "toolhive-vmcp" transport fallback. + resolved := WithDefaults(cfg) + + // AuthzMiddleware is vestigial on the New/Serve path (Serve never applies it; authz + // moved to the core admission seam). Setting it without a corresponding Config.Authz + // would silently degrade to allow-all — a security regression for an embedder that + // relied on it for Cedar authz — so fail fast instead of logging a warning the embedder + // is unlikely to notice. cli/serve.go sets Authz whenever it sets AuthzMiddleware, so + // this never fires in-tree. + if cfg.AuthzMiddleware != nil && cfg.Authz == nil { + return nil, fmt.Errorf("%w: Config.AuthzMiddleware is set but has no effect on the "+ + "New/Serve path; set Config.Authz to enforce authorization, or clear AuthzMiddleware", + vmcp.ErrInvalidConfig) + } + + // The core admission seam has no representation for the optimizer's meta-tools + // (find_tool/call_tool), so combining Authz with the optimizer would let those calls + // bypass Cedar evaluation entirely. Fail fast per the admission seam contract (see the + // core.Admission doc); the optimizer keeps its own HTTP-path authz until a focused PR. + if cfg.Authz != nil && cfg.OptimizerConfig != nil { + return nil, fmt.Errorf("%w: Config.Authz and Config.OptimizerConfig are mutually "+ + "exclusive; the optimizer meta-tools (find_tool, call_tool) are not represented "+ + "in the core admission seam", vmcp.ErrInvalidConfig) + } + + // Cedar resource entities are scoped to MCP::"" (see Config.Authz), so an empty + // Name with Authz set would key policies on MCP::"" and silently fail to match. core.New + // also rejects this (its admission seam), but fail at the construction root too for a + // clearer error consistent with the guards above. + if cfg.Authz != nil && cfg.Name == "" { + return nil, fmt.Errorf("%w: Config.Name is required when Config.Authz is set "+ + "(it is the Cedar resource entity name)", vmcp.ErrInvalidConfig) + } + + // The SDK elicitation adapter wraps the mcp-go server Serve builds below, so it + // cannot exist before core.New. Give the core a late-bound requester now and bind the + // real adapter to Serve's server before serving begins (RequestElicitation is only + // invoked at request time, after bind). + elicitation := newLateBoundElicitationRequester() + + // Wrap the aggregator in a per-identity caching decorator: the core re-derives the + // advertised view on every call, so without this the Serve path re-sweeps every backend's + // tools/list per tool call. The cache is keyed on identity + forwarded credentials, so it + // never serves one caller's capability view to another. + cachedAgg := aggregator.NewCachingAggregator(cfg.Aggregator, capabilityCacheTTL) + + coreVMCP, err := core.New(deriveCoreConfig( + cfg, cachedAgg, rt, backendClient, backendRegistry, workflowDefs, + cfg.Authz, elicitation, + )) if err != nil { - return nil, fmt.Errorf("failed to create session data storage: %w", err) + return nil, err } - // Close sessionDataStorage if New() returns an error after this point so the - // background cleanup goroutine does not leak. - closeStorageOnErr := true + // core.New started the workflow state store's cleanup goroutine and the backend health + // monitor (both owned by the core now). If Serve fails after this point, close the core so + // neither leaks (mirrors Serve's closeStorageOnErr guard); on success the core's lifecycle + // is owned by srv.Stop, so the guard is disarmed before returning. + closeCoreOnErr := true defer func() { - if closeStorageOnErr { - _ = sessionDataStorage.Close() + if closeCoreOnErr { + _ = coreVMCP.Close() } }() - // Create handler factory (used by adapter and for future dynamic registration) - handlerFactory := adapter.NewDefaultHandlerFactory(rt, backendClient) - - // Create capability adapter (single source of truth for converting aggregator types to SDK types) - capabilityAdapter := adapter.NewCapabilityAdapter(handlerFactory) - - // Create health monitor if configured - var healthMon *health.Monitor - if cfg.HealthMonitorConfig != nil { - // Get initial backends list from registry for health monitoring setup - initialBackends := backendRegistry.List(ctx) - healthMon, err = health.NewMonitor(backendClient, initialBackends, *cfg.HealthMonitorConfig) - if err != nil { - return nil, fmt.Errorf("failed to create health monitor: %w", err) - } - slog.Info("health monitoring enabled", - "check_interval", cfg.HealthMonitorConfig.CheckInterval, - "unhealthy_threshold", cfg.HealthMonitorConfig.UnhealthyThreshold, - "timeout", cfg.HealthMonitorConfig.Timeout, - "degraded_threshold", cfg.HealthMonitorConfig.DegradedThreshold) - } else { - slog.Info("health monitoring disabled") - } - - // Pass the whole factory config so the session manager constructs everything - // it needs (optimizer wiring, composite tool layers, telemetry instruments). + // On the New/Serve path the core is the single aggregator and the source of the + // advertised set; the session factory only opens per-session backend connections and + // binds identity. AdvertiseFromCore makes the session manager source tools from the + // core and (when the optimizer is enabled) build the Serve-layer optimizer over the + // core's tools instead of decorating the factory. Composite-tool workflows are owned + // by the core, so no WorkflowDefs/ComposerFactory are wired here. sessMgrCfg := &sessionmanager.FactoryConfig{ Base: cfg.SessionFactory, - WorkflowDefs: workflowDefs, - ComposerFactory: sessionComposerFactory, OptimizerConfig: cfg.OptimizerConfig, OptimizerFactory: cfg.OptimizerFactory, TelemetryProvider: cfg.TelemetryProvider, + AdvertiseFromCore: true, } - vmcpSessMgr, optimizerCleanup, err := sessionmanager.New(sessionDataStorage, sessMgrCfg, backendRegistry) + + srv, err := Serve(ctx, coreVMCP, deriveServerConfig(resolved, backendRegistry, sessMgrCfg)) if err != nil { return nil, err } - // Create Server instance - srv := &Server{ - config: cfg, - mcpServer: mcpServer, - router: rt, - backendClient: backendClient, - handlerFactory: handlerFactory, - discoveryMgr: discoveryMgr, - backendRegistry: backendRegistry, - sessionManager: sessionManager, - sessionDataStorage: sessionDataStorage, - capabilityAdapter: capabilityAdapter, - ready: make(chan struct{}), - healthMonitor: healthMon, - statusReporter: cfg.StatusReporter, - vmcpSessionMgr: vmcpSessMgr, - } + // Retain the discovery manager only so Stop() drives its (now no-op) cleanup for + // parity with the legacy path. The core replaced discovery's aggregation, so it is + // otherwise unused here: the shared Handler guards the discovery middleware to + // s.core == nil, which never holds for a Serve-built server. + srv.discoveryMgr = discoveryMgr - if optimizerCleanup != nil { - srv.shutdownFuncs = append(srv.shutdownFuncs, optimizerCleanup) - } - - // Register OnRegisterSession hook to inject capabilities after SDK registers session. - // See handleSessionRegistration for implementation details. - hooks.AddOnRegisterSession(func(ctx context.Context, session server.ClientSession) { - srv.handleSessionRegistration(ctx, session) - }) + // Bind the elicitation adapter to the SDK server Serve built so composite-workflow + // elicitation reaches the same mcp-go server that serves client traffic. + elicitation.bind(NewSDKElicitationAdapter(srv.MCPServer())) - // Register OnBeforeListTools hook for lazy session tool injection. - // - // When a session is reconstructed from Redis on a different pod (cross-pod sharing), - // the SDK's per-session tool store is empty because OnRegisterSession only fires - // during Initialize, which the client doesn't re-send to pod B. This hook lazily - // injects the tools from the VMCP session manager into the ephemeral SDK session - // before handleListTools reads from the per-session tool store. - hooks.AddBeforeListTools(func(ctx context.Context, _ any, _ *mcp.ListToolsRequest) { - srv.lazyInjectSessionTools(ctx) - }) - - // Register OnBeforeCallTool hook for the same reason as OnBeforeListTools. - // A client may call a tool directly without first calling tools/list, so we - // also need to ensure the tool handlers are registered before the call is routed. - hooks.AddBeforeCallTool(func(ctx context.Context, _ any, _ *mcp.CallToolRequest) { - srv.lazyInjectSessionTools(ctx) - }) - - // Disarm the close-on-error guard: Server is fully constructed. - closeStorageOnErr = false + closeCoreOnErr = false // Serve succeeded; srv.Stop now owns the core's lifecycle. return srv, nil } @@ -604,18 +526,16 @@ func (s *Server) Handler(_ context.Context) (http.Handler, error) { } // MCP endpoint - apply middleware chain (wrapping order, execution happens in reverse): - // Code wraps: auth+parser → rate-limit → audit → discovery → annotation-enrichment → - // authz → backend-enrichment → MCP-parsing → telemetry + // Code wraps: auth+parser → rate-limit → audit → discovery → backend-enrichment → + // MCP-parsing → telemetry // Execution order: recovery → header-val → auth+parser → rate-limit → audit → - // discovery → annotation-enrichment → authz → backend-enrichment → - // MCP-parsing → telemetry → handler + // discovery → backend-enrichment → MCP-parsing → telemetry → handler // - // The authz and annotation-enrichment layers are both guarded by - // s.config.AuthzMiddleware != nil: applied on the server.New path (authz on) and - // omitted on the Serve path, which leaves AuthzMiddleware nil so authorization - // moves to the core admission seam (#5438). Both blocks remain in this shared - // Handler — physical removal is deferred to Phase 3 (#5445), after server.New is - // routed through Serve and the legacy authz path is gone. + // The legacy HTTP authz and annotation-enrichment layers have been removed: every caller + // now routes through Serve, so authorization is enforced by the core admission seam + // (#5438) rather than HTTP middleware. The remaining legacy-only blocks (backend + // enrichment, discovery) are guarded by s.core == nil and are removed in the 302b + // follow-up. var mcpHandler http.Handler = streamableServer @@ -636,21 +556,6 @@ func (s *Server) Handler(_ context.Context) (http.Handler, error) { // Apply backend enrichment middleware (legacy path only — see withBackendEnrichment). mcpHandler = s.withBackendEnrichment(mcpHandler) - // Apply authorization middleware if configured (runs AFTER discovery in execution). - // Wrapping it here (before discovery wrap) means discovery runs first, then authz. - if s.config.AuthzMiddleware != nil { - mcpHandler = s.config.AuthzMiddleware(mcpHandler) - slog.Info("authorization middleware enabled for MCP endpoints (post-discovery)") - } - - // Apply annotation enrichment middleware (runs after discovery, before authz in execution). - // Reads tool annotations from discovered capabilities and injects them into the - // request context so the authz middleware can make annotation-aware decisions. - if s.config.AuthzMiddleware != nil { - mcpHandler = AnnotationEnrichmentMiddleware(mcpHandler) - slog.Info("annotation enrichment middleware enabled for MCP endpoints") - } - // Apply discovery middleware (runs after audit/auth middleware) — legacy path only. // Discovery middleware performs per-request capability aggregation with user context, // injecting the routing table into the request context (the discovery-into-context seam). @@ -666,16 +571,12 @@ func (s *Server) Handler(_ context.Context) (http.Handler, error) { // initialize-skip behavior is preserved here on the legacy path; physical removal of the // middleware and the context seam is deferred to Phase 3 (#5445). if s.core == nil { - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() - - var healthStatusProvider health.StatusProvider - if healthMon != nil { - healthStatusProvider = healthMon - } + // Dead on the Serve path: s.core is never nil after #5445 routes New through Serve, + // and the health monitor now lives in the core, so no StatusProvider is available + // here. Kept (with nil health) only until 302b removes the discovery middleware and + // its context seam (anti-pattern #1). mcpHandler = discovery.Middleware( - s.discoveryMgr, s.backendRegistry, s.vmcpSessionMgr, healthStatusProvider, + s.discoveryMgr, s.backendRegistry, s.vmcpSessionMgr, nil, discovery.WithSessionScopedRouting(), )(mcpHandler) slog.Info("discovery middleware enabled for lazy per-user capability discovery") @@ -756,8 +657,6 @@ func (s *Server) applyForwardedHeaderCapture(next http.Handler) http.Handler { } // Start starts the Virtual MCP Server and begins serving requests. -// -//nolint:gocyclo // Complexity from health monitoring and startup orchestration is acceptable func (s *Server) Start(ctx context.Context) error { // Build the HTTP handler (middleware chain, routes, mux) handler, err := s.Handler(ctx) @@ -808,23 +707,8 @@ func (s *Server) Start(ctx context.Context) error { close(s.ready) }) - // Start health monitor if configured - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() - - if healthMon != nil { - if err := healthMon.Start(ctx); err != nil { - // Log error and disable health monitoring - treat as if it wasn't configured - // This ensures getter methods correctly report monitoring as disabled - slog.Warn("failed to start health monitor, disabling health monitoring", "error", err) - s.healthMonitorMu.Lock() - s.healthMonitor = nil - s.healthMonitorMu.Unlock() - } else { - slog.Info("health monitor started") - } - } + // The backend health monitor is owned by the core (built and started in core.New, stopped + // in core.Close), so the server no longer starts or stops it here. // Start status reporter if configured if s.statusReporter != nil { @@ -894,16 +778,8 @@ func (s *Server) Stop(ctx context.Context) error { s.listener = nil s.listenerMu.Unlock() - // Stop health monitor to clean up health check goroutines - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() - - if healthMon != nil { - if err := healthMon.Stop(); err != nil { - errs = append(errs, fmt.Errorf("failed to stop health monitor: %w", err)) - } - } + // The backend health monitor is stopped by core.Close (the core owns it); Serve registered + // a shutdown function that closes the core, run in the loop below. // Run shutdown functions (e.g., status reporter cleanup, future components) for _, shutdown := range s.shutdownFuncs { @@ -1310,46 +1186,20 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv return nil } -// validateWorkflows validates workflow definitions, returning only the valid ones. -// -// This function: -// 1. Validates each workflow definition (cycle detection, tool references, etc.) -// 2. Returns error on first validation failure (fail-fast) -// -// Failing fast on invalid workflows provides immediate user feedback and prevents -// security issues (resource exhaustion from cycles, information disclosure from errors). -func validateWorkflows( - validator composer.Composer, - workflowDefs map[string]*composer.WorkflowDefinition, -) (map[string]*composer.WorkflowDefinition, error) { - if len(workflowDefs) == 0 { - return nil, nil - } - - validDefs := make(map[string]*composer.WorkflowDefinition, len(workflowDefs)) - - for name, def := range workflowDefs { - if err := validator.ValidateWorkflow(context.Background(), def); err != nil { - return nil, fmt.Errorf("invalid workflow definition '%s': %w", name, err) - } - - validDefs[name] = def - slog.Debug("validated workflow definition", "name", name) - } - - if len(validDefs) > 0 { - slog.Info("loaded valid composite tool workflows", "count", len(validDefs)) +// backendHealth returns the core-owned backend health reporter, or nil when health +// monitoring is disabled (the core owns the monitor's lifecycle; #5443 reversal). The +// s.core nil guard is defensive: a Serve-built server always has a non-nil core. +func (s *Server) backendHealth() health.Reporter { + if s.core == nil { + return nil } - - return validDefs, nil + return s.core.BackendHealth() } // GetBackendHealthStatus returns the health status of a specific backend. // Returns error if health monitoring is disabled or backend not found. func (s *Server) GetBackendHealthStatus(backendID string) (vmcp.BackendHealthStatus, error) { - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() + healthMon := s.backendHealth() if healthMon == nil { return vmcp.BackendUnknown, fmt.Errorf("health monitoring is disabled") @@ -1360,9 +1210,7 @@ func (s *Server) GetBackendHealthStatus(backendID string) (vmcp.BackendHealthSta // GetBackendHealthState returns the full health state of a specific backend. // Returns error if health monitoring is disabled or backend not found. func (s *Server) GetBackendHealthState(backendID string) (*health.State, error) { - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() + healthMon := s.backendHealth() if healthMon == nil { return nil, fmt.Errorf("health monitoring is disabled") @@ -1373,9 +1221,7 @@ func (s *Server) GetBackendHealthState(backendID string) (*health.State, error) // GetAllBackendHealthStates returns the health states of all backends. // Returns empty map if health monitoring is disabled. func (s *Server) GetAllBackendHealthStates() map[string]*health.State { - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() + healthMon := s.backendHealth() if healthMon == nil { return make(map[string]*health.State) @@ -1386,9 +1232,7 @@ func (s *Server) GetAllBackendHealthStates() map[string]*health.State { // GetHealthSummary returns a summary of backend health across all backends. // Returns zero-valued summary if health monitoring is disabled. func (s *Server) GetHealthSummary() health.Summary { - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() + healthMon := s.backendHealth() if healthMon == nil { return health.Summary{} @@ -1416,9 +1260,7 @@ type BackendHealthResponse struct { // Security Note: This endpoint is unauthenticated and may expose backend topology. // Consider applying authentication middleware if operating in multi-tenant mode. func (s *Server) handleBackendHealth(w http.ResponseWriter, _ *http.Request) { - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() + healthMon := s.backendHealth() response := BackendHealthResponse{ MonitoringEnabled: healthMon != nil, diff --git a/pkg/vmcp/server/server_test.go b/pkg/vmcp/server/server_test.go index 672217a712..e3497accfa 100644 --- a/pkg/vmcp/server/server_test.go +++ b/pkg/vmcp/server/server_test.go @@ -18,10 +18,9 @@ import ( "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/authz/authorizers" + "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/vmcp" - "github.com/stacklok/toolhive/pkg/vmcp/aggregator" - "github.com/stacklok/toolhive/pkg/vmcp/discovery" discoveryMocks "github.com/stacklok/toolhive/pkg/vmcp/discovery/mocks" "github.com/stacklok/toolhive/pkg/vmcp/mocks" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" @@ -71,7 +70,7 @@ func TestServerStartFailsWhenReporterStartFails(t *testing.T) { srv, err := server.New( context.Background(), - &server.Config{Host: "127.0.0.1", Port: 0, StatusReporter: sr, SessionFactory: newNoopMockFactory(t)}, + &server.Config{Host: "127.0.0.1", Port: 0, StatusReporter: sr, SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil)}, mockRouter, mockBackendClient, mockDiscoveryMgr, @@ -101,7 +100,7 @@ func TestServerStopRunsReporterShutdown(t *testing.T) { srv, err := server.New( context.Background(), - &server.Config{Host: "127.0.0.1", Port: 0, StatusReporter: sr, SessionFactory: newNoopMockFactory(t)}, + &server.Config{Host: "127.0.0.1", Port: 0, StatusReporter: sr, SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil)}, mockRouter, mockBackendClient, mockDiscoveryMgr, @@ -156,7 +155,7 @@ func TestNew(t *testing.T) { }{ { name: "applies all defaults", - config: &server.Config{SessionFactory: newNoopMockFactory(t)}, + config: &server.Config{SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil)}, expectedHost: "127.0.0.1", expectedPort: 4483, expectedPath: "/mcp", @@ -171,7 +170,7 @@ func TestNew(t *testing.T) { Host: "0.0.0.0", Port: 8080, EndpointPath: "/api/mcp", - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, expectedHost: "0.0.0.0", expectedPort: 8080, @@ -184,7 +183,7 @@ func TestNew(t *testing.T) { config: &server.Config{ Host: "192.168.1.1", Port: 9000, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, expectedHost: "192.168.1.1", expectedPort: 9000, @@ -299,7 +298,7 @@ func TestServer_Address(t *testing.T) { name: "default host with explicit port", config: &server.Config{ Port: 4483, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, expected: "127.0.0.1:4483", }, @@ -307,7 +306,7 @@ func TestServer_Address(t *testing.T) { name: "port 0 for dynamic allocation", config: &server.Config{ Port: 0, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, expected: "127.0.0.1:0", }, @@ -316,7 +315,7 @@ func TestServer_Address(t *testing.T) { config: &server.Config{ Host: "0.0.0.0", Port: 8080, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, expected: "0.0.0.0:8080", }, @@ -325,7 +324,7 @@ func TestServer_Address(t *testing.T) { config: &server.Config{ Host: "localhost", Port: 3000, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, expected: "localhost:3000", }, @@ -364,7 +363,7 @@ func TestServer_Stop(t *testing.T) { mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) mockDiscoveryMgr.EXPECT().Stop().Times(1) - s, err := server.New(context.Background(), &server.Config{SessionFactory: newNoopMockFactory(t)}, mockRouter, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry([]vmcp.Backend{}), nil) + s, err := server.New(context.Background(), &server.Config{SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil)}, mockRouter, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry([]vmcp.Backend{}), nil) require.NoError(t, err) err = s.Stop(context.Background()) require.NoError(t, err) @@ -385,6 +384,7 @@ func TestNew_NilSessionFactory_ReturnsError(t *testing.T) { context.Background(), &server.Config{ SessionFactory: nil, // deliberately omitted + Aggregator: newStubAggregator(nil), }, mockRouter, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry([]vmcp.Backend{}), nil, @@ -393,6 +393,32 @@ func TestNew_NilSessionFactory_ReturnsError(t *testing.T) { assert.Contains(t, err.Error(), "SessionFactory") } +// TestNew_NilAggregator_ReturnsError guards the now-required Config.Aggregator: the core +// is the single source of the advertised capability set, so server.New must fail (via +// core.New's validation) rather than silently construct a server with no aggregation. +func TestNew_NilAggregator_ReturnsError(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + mockRouter := routerMocks.NewMockRouter(ctrl) + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) + + _, err := server.New( + context.Background(), + &server.Config{ + SessionFactory: newNoopMockFactory(t), + Aggregator: nil, // deliberately omitted: now a required field + }, + mockRouter, mockBackendClient, mockDiscoveryMgr, + vmcp.NewImmutableRegistry([]vmcp.Backend{}), nil, + ) + require.Error(t, err) + assert.Contains(t, err.Error(), "Aggregator") +} + func TestNew_WithAuditConfig(t *testing.T) { t.Parallel() @@ -466,7 +492,7 @@ func TestNew_WithAuditConfig(t *testing.T) { config := &server.Config{ AuditConfig: tt.auditConfig, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), } s, err := server.New(context.Background(), config, mockRouter, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry([]vmcp.Backend{}), nil) @@ -499,7 +525,7 @@ func TestServerStopClosesOptimizerStore(t *testing.T) { srv, err := server.New( context.Background(), - &server.Config{Host: "127.0.0.1", Port: 0, OptimizerConfig: &optimizer.Config{}, SessionFactory: newNoopMockFactory(t)}, + &server.Config{Host: "127.0.0.1", Port: 0, OptimizerConfig: &optimizer.Config{}, SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil)}, mockRouter, mockBackendClient, mockDiscoveryMgr, @@ -552,7 +578,7 @@ func TestHandler_ReturnsNonNilHandler(t *testing.T) { srv, err := server.New( t.Context(), - &server.Config{Host: "127.0.0.1", Port: 0, SessionFactory: newNoopMockFactory(t)}, + &server.Config{Host: "127.0.0.1", Port: 0, SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil)}, mockRouter, mockBackendClient, mockDiscoveryMgr, @@ -594,7 +620,7 @@ func TestHandler_ReturnsErrorOnInvalidAuditConfig(t *testing.T) { Component: "vmcp-server", MaxDataSize: -1, }, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, mockRouter, mockBackendClient, @@ -630,7 +656,7 @@ func TestHandler_CanBeCalledMultipleTimes(t *testing.T) { srv, err := server.New( t.Context(), - &server.Config{Host: "127.0.0.1", Port: 0, SessionFactory: newNoopMockFactory(t)}, + &server.Config{Host: "127.0.0.1", Port: 0, SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil)}, mockRouter, mockBackendClient, mockDiscoveryMgr, @@ -683,7 +709,7 @@ func TestHandler_RegistersWellKnownRoutes(t *testing.T) { Host: "127.0.0.1", Port: 0, AuthInfoHandler: authInfoHandler, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), // AuthServer is not set here because the concrete type // *asrunner.EmbeddedAuthServer cannot be easily constructed in an // external test without a real auth server backing it. @@ -797,7 +823,7 @@ func TestAcceptHeaderValidation(t *testing.T) { srv, err := server.New( t.Context(), - &server.Config{Host: "127.0.0.1", Port: 0, SessionFactory: newNoopMockFactory(t)}, + &server.Config{Host: "127.0.0.1", Port: 0, SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil)}, mockRouter, mockBackendClient, mockDiscoveryMgr, @@ -862,118 +888,170 @@ func TestAcceptHeaderValidation(t *testing.T) { } } -// TestHandlerAppliesAuthzAndAnnotationOnlyWhenConfigured proves the shared -// (*Server).Handler applies BOTH the authz layer and the annotation-enrichment layer -// exactly when config.AuthzMiddleware != nil. This is the single guard that lets the -// same Handler serve both modes: the server.New path (AuthzMiddleware set) keeps -// enforcing authz with no regression, while the Serve path (AuthzMiddleware nil; see -// TestServeOmitsAuthzAndAnnotation) omits both layers and defers authorization to the -// core admission seam (#5438). The blocks are not deleted here — physical removal is -// Phase 3 (#5445). -func TestHandlerAppliesAuthzAndAnnotationOnlyWhenConfigured(t *testing.T) { +// newTestAuthzConfig builds a minimal, permissive Cedar authz config. server.New now +// requires Config.Authz to be set whenever Config.AuthzMiddleware is (the vestigial +// middleware must not silently degrade to allow-all), so tests that exercise AuthzMiddleware +// must supply it. The policy permits everything so the core admission seam never interferes +// with what these tests actually assert. +func newTestAuthzConfig(t *testing.T) *authorizers.Config { + t.Helper() + cfg, err := authorizers.NewConfig(cedar.Config{ + Version: "1.0", + Type: cedar.ConfigType, + Options: &cedar.ConfigOptions{Policies: []string{`permit(principal, action, resource);`}, EntitiesJSON: "[]"}, + }) + require.NoError(t, err) + return cfg +} + +// TestNew_AuthzMiddlewareWithoutAuthz_ReturnsError guards the fail-fast that replaced the +// former WARN: setting the vestigial Config.AuthzMiddleware without Config.Authz would +// silently degrade to allow-all on the New/Serve path, so server.New must reject it. +func TestNew_AuthzMiddlewareWithoutAuthz_ReturnsError(t *testing.T) { t.Parallel() - // Distinctive status only the observable authz layer writes, so its presence in the - // chain is unambiguous. - const sentinelStatus = http.StatusTeapot + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + mockRouter := routerMocks.NewMockRouter(ctrl) + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) - readOnly := true - caps := &aggregator.AggregatedCapabilities{ - Tools: []vmcp.Tool{{ - Name: "my_tool", - Annotations: &vmcp.ToolAnnotations{ReadOnlyHint: &readOnly}, - }}, - } + _, err := server.New(t.Context(), + &server.Config{ + SessionFactory: newNoopMockFactory(t), + Aggregator: newStubAggregator(nil), + AuthzMiddleware: func(h http.Handler) http.Handler { return h }, // set without Authz + }, + mockRouter, mockBackendClient, mockDiscoveryMgr, + vmcp.NewImmutableRegistry([]vmcp.Backend{}), nil, + ) + require.Error(t, err) + assert.ErrorIs(t, err, vmcp.ErrInvalidConfig) + assert.Contains(t, err.Error(), "AuthzMiddleware") +} - tests := []struct { - name string - withAuthz bool - wantAuthzApplied bool - }{ - {name: "applied when AuthzMiddleware set (server.New path)", withAuthz: true, wantAuthzApplied: true}, - {name: "omitted when AuthzMiddleware nil (Serve path)", withAuthz: false, wantAuthzApplied: false}, - } +// TestNew_AuthzWithOptimizer_ReturnsError guards the documented mutual exclusion: the core +// admission seam has no representation for the optimizer's meta-tools, so combining Authz +// with the optimizer would silently bypass Cedar. server.New must reject the combination. +func TestNew_AuthzWithOptimizer_ReturnsError(t *testing.T) { + t.Parallel() - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + mockRouter := routerMocks.NewMockRouter(ctrl) + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) - ctrl := gomock.NewController(t) - t.Cleanup(ctrl.Finish) - mockRouter := routerMocks.NewMockRouter(ctrl) - mockBackendClient := mocks.NewMockBackendClient(ctrl) - mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) - mockBackendRegistry := mocks.NewMockBackendRegistry(ctrl) - // List/Discover stay permissive (AnyTimes) but do not fire on this path: with - // WithSessionScopedRouting() and no Mcp-Session-Id, discovery.Middleware returns - // before aggregating. Stop fires via srv.Stop in the cleanup below. - mockBackendRegistry.EXPECT().List(gomock.Any()).Return(nil).AnyTimes() - mockDiscoveryMgr.EXPECT().Discover(gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() - mockDiscoveryMgr.EXPECT().Stop().AnyTimes() + _, err := server.New(t.Context(), + &server.Config{ + Name: "test-vmcp", + SessionFactory: newNoopMockFactory(t), + Aggregator: newStubAggregator(nil), + Authz: newTestAuthzConfig(t), + OptimizerConfig: &optimizer.Config{}, + }, + mockRouter, mockBackendClient, mockDiscoveryMgr, + vmcp.NewImmutableRegistry([]vmcp.Backend{}), nil, + ) + require.Error(t, err) + assert.ErrorIs(t, err, vmcp.ErrInvalidConfig) + assert.Contains(t, err.Error(), "OptimizerConfig") +} - var ( - authzApplied bool - annotationsSeen bool - ) - // Observable authz layer: records that it ran AND whether the - // annotation-enrichment layer (which executes immediately before it) injected - // the tool's annotations into ctx, then short-circuits with the sentinel. - authz := func(_ http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - authzApplied = true - annotationsSeen = authorizers.ToolAnnotationsFromContext(r.Context()) != nil - w.WriteHeader(sentinelStatus) - }) - } +// TestNew_AuthzWithoutName_ReturnsError guards the documented Config.Authz requirement: +// Cedar resource entities are scoped to MCP::"", so an empty Name with Authz set would +// silently key policies on MCP::"" and stop matching. server.New rejects it at the +// construction root (core.New also enforces it, with a deeper message). +func TestNew_AuthzWithoutName_ReturnsError(t *testing.T) { + t.Parallel() - cfg := &server.Config{Host: "127.0.0.1", Port: 0, SessionFactory: newNoopMockFactory(t)} - if tt.withAuthz { - cfg.AuthzMiddleware = authz - } + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + mockRouter := routerMocks.NewMockRouter(ctrl) + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) - srv, err := server.New(t.Context(), cfg, mockRouter, mockBackendClient, mockDiscoveryMgr, mockBackendRegistry, nil) - require.NoError(t, err) - t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + _, err := server.New(t.Context(), + &server.Config{ + // Name intentionally empty. + SessionFactory: newNoopMockFactory(t), + Aggregator: newStubAggregator(nil), + Authz: newTestAuthzConfig(t), + }, + mockRouter, mockBackendClient, mockDiscoveryMgr, + vmcp.NewImmutableRegistry([]vmcp.Backend{}), nil, + ) + require.Error(t, err) + assert.ErrorIs(t, err, vmcp.ErrInvalidConfig) + assert.Contains(t, err.Error(), "Name") +} - handler, err := srv.Handler(t.Context()) - require.NoError(t, err) +// TestNewIgnoresVestigialAuthzMiddleware proves that server.New — now routed through +// core.New + Serve — does NOT apply the HTTP authz or annotation-enrichment layers even +// when Config.AuthzMiddleware is set (alongside the now-required Config.Authz, mirroring +// cli/serve.go). deriveServerConfig/buildServeConfig drop AuthzMiddleware (it is vestigial +// on the New/Serve path), so the shared (*Server).Handler skips both blocks and +// authorization is enforced by the core admission seam (#5438) instead. The now-dead HTTP +// blocks remain in the shared Handler until the #5445 follow-up removes them; this guards +// that they never run via the public constructor. The Serve-path view of the same behavior +// is TestServeOmitsAuthzAndAnnotation (serve_test.go). +func TestNewIgnoresVestigialAuthzMiddleware(t *testing.T) { + t.Parallel() - // Craft a tools/call request. This chain has no auth-parser, so inject the - // parsed request and discovered capabilities directly into ctx (as - // ParsingMiddleware and the discovery middleware would). - // - // Load-bearing dependency: discovery.Middleware is built with - // WithSessionScopedRouting(), so a request with no Mcp-Session-Id passes straight - // through without touching the (mock) manager and WITHOUT rewriting ctx — - // preserving the injected capabilities for the annotation-enrichment layer. If - // that session-scoped pass-through ever changes to overwrite ctx, annotationsSeen - // silently goes false; the positive-case assert.True below is the guard. - ctx := context.WithValue(t.Context(), mcpparser.MCPRequestContextKey, - &mcpparser.ParsedMCPRequest{Method: "tools/call", ResourceID: "my_tool"}) - ctx = discovery.WithDiscoveredCapabilities(ctx, caps) - - req := httptest.NewRequest(http.MethodPost, "/mcp", nil).WithContext(ctx) - // Set Content-Type so the nil-authz request reaches a representative chain point - // (the inner SDK handler) rather than being rejected during content negotiation; - // both cases then exercise the chain identically. - req.Header.Set("Content-Type", "application/json") - rec := httptest.NewRecorder() + // Distinctive status the observable authz layer would write if it ever ran. + const sentinelStatus = http.StatusTeapot - // Both paths return synchronously: the configured case short-circuits at the - // observable authz layer (sentinel), and the nil case falls through to the inner - // SDK handler, which answers this POST without blocking. A direct call suffices — - // no goroutine/timeout scaffolding needed. - handler.ServeHTTP(rec, req) - - assert.Equal(t, tt.wantAuthzApplied, authzApplied) - if tt.wantAuthzApplied { - assert.Equal(t, sentinelStatus, rec.Code, "observable authz layer should have written the sentinel status") - assert.True(t, annotationsSeen, - "annotation-enrichment must run before authz and inject the tool annotations") - } else { - assert.NotEqual(t, sentinelStatus, rec.Code, - "no authz layer should run on the Serve-style (nil AuthzMiddleware) path") - } + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + mockRouter := routerMocks.NewMockRouter(ctrl) + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) + mockBackendRegistry := mocks.NewMockBackendRegistry(ctrl) + mockBackendRegistry.EXPECT().List(gomock.Any()).Return(nil).AnyTimes() + mockDiscoveryMgr.EXPECT().Discover(gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mockDiscoveryMgr.EXPECT().Stop().AnyTimes() + + // If this authz middleware is ever applied it records the fact and short-circuits + // with the sentinel status. + authzApplied := false + authz := func(_ http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + authzApplied = true + w.WriteHeader(sentinelStatus) }) } + + // AuthzMiddleware is set (alongside Authz, as cli/serve.go does), but server.New routes + // through Serve, which drops the HTTP AuthzMiddleware layer; authz is enforced by the + // core admission seam from Config.Authz instead. + cfg := &server.Config{ + Name: "test-vmcp", // required once Authz is set + Host: "127.0.0.1", + Port: 0, + SessionFactory: newNoopMockFactory(t), + Aggregator: newStubAggregator(nil), + AuthzMiddleware: authz, + Authz: newTestAuthzConfig(t), + } + srv, err := server.New(t.Context(), cfg, mockRouter, mockBackendClient, mockDiscoveryMgr, mockBackendRegistry, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + handler, err := srv.Handler(t.Context()) + require.NoError(t, err) + + // A tools/call request with a pre-parsed MCP request in ctx (this chain has no + // auth-parser). If the HTTP authz layer were applied it would short-circuit with the + // sentinel before the inner SDK handler. + ctx := context.WithValue(t.Context(), mcpparser.MCPRequestContextKey, + &mcpparser.ParsedMCPRequest{Method: "tools/call", ResourceID: "my_tool"}) + req := httptest.NewRequest(http.MethodPost, "/mcp", nil).WithContext(ctx) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + assert.False(t, authzApplied, + "server.New routes through Serve, which drops AuthzMiddleware; the HTTP authz layer must not run") + assert.NotEqual(t, sentinelStatus, rec.Code, + "no HTTP authz layer should short-circuit on the New/Serve path (authz is in the core)") } diff --git a/pkg/vmcp/server/session_management_integration_test.go b/pkg/vmcp/server/session_management_integration_test.go index 3b1eb6eba4..836294078c 100644 --- a/pkg/vmcp/server/session_management_integration_test.go +++ b/pkg/vmcp/server/session_management_integration_test.go @@ -106,6 +106,10 @@ func newMockFactory(t *testing.T, ctrl *gomock.Controller, tools []vmcp.Tool) (* mock.EXPECT().GetMetadata().Return(map[string]string{ vmcpsession.MetadataKeyIdentityBinding: "unauthenticated", }).AnyTimes() + // The Serve-path enforceSessionBinding reads the binding via the single-key + // GetMetadataValue (not the full-map GetMetadata); return the same sentinel. + mock.EXPECT().GetMetadataValue(vmcpsession.MetadataKeyIdentityBinding). + Return("unauthenticated", true).AnyTimes() mock.EXPECT().SetMetadata(gomock.Any(), gomock.Any()).AnyTimes() toolsCopy := make([]vmcp.Tool, len(tools)) copy(toolsCopy, tools) @@ -149,10 +153,29 @@ func newMockFactory(t *testing.T, ctrl *gomock.Controller, tools []vmcp.Tool) (* // serverOptions holds optional configuration extensions for buildTestServerWithOptions. type serverOptions struct { + // tools is the advertised backend tool set. Now that server.New routes through + // core.New + Serve, the core sources the advertised set from the Aggregator (not the + // session factory), so tests provide it here; buildTestServerWithOptions builds a stub + // aggregator from it (and a routing table so composite-tool reachability and call + // routing work). + tools []vmcp.Tool workflowDefs map[string]*composer.WorkflowDefinition optimizerFactory func(context.Context, []mcpsdk.ServerTool) (optimizer.Optimizer, error) } +// capsFromTools builds the AggregatedCapabilities the stub aggregator returns: the tools +// plus a routing table mapping each tool name to a backend target keyed by the tool name. +// This mirrors the routing table the legacy mock session built (newMockFactory), so the +// core advertises the tools, resolves composite-tool reachability, and routes tools/call +// through BackendClient.CallTool. +func capsFromTools(tools []vmcp.Tool) *aggregator.AggregatedCapabilities { + rt := &vmcp.RoutingTable{Tools: make(map[string]*vmcp.BackendTarget, len(tools))} + for i := range tools { + rt.Tools[tools[i].Name] = &vmcp.BackendTarget{WorkloadID: tools[i].Name} + } + return &aggregator.AggregatedCapabilities{Tools: tools, RoutingTable: rt} +} + // buildTestServer constructs a vMCP server with session management enabled, // backed by mock discovery infrastructure, and returns the httptest.Server // and the session factory so tests can inspect state. @@ -161,9 +184,10 @@ type serverOptions struct { func buildTestServer( t *testing.T, factory vmcpsession.MultiSessionFactory, + tools []vmcp.Tool, ) *httptest.Server { t.Helper() - return buildTestServerWithOptions(t, factory, serverOptions{}) + return buildTestServerWithOptions(t, factory, serverOptions{tools: tools}) } // buildTestServerWithOptions is like buildTestServer but accepts optional workflow @@ -183,14 +207,21 @@ func buildTestServerWithOptions( mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) mockBackendRegistry := mocks.NewMockBackendRegistry(ctrl) - // The discovery middleware calls List() + Discover() for every initialize request. - // Return an empty (non-nil) AggregatedCapabilities so the middleware does not - // dereference a nil pointer when logging tool/resource counts. + // List() is consumed by the core's on-demand aggregation; the discovery middleware + // is guarded off on the Serve path (s.core != nil), so Discover() is no longer called + // (the AnyTimes expectations tolerate zero calls). Return an empty (non-nil) result. emptyAggCaps := &aggregator.AggregatedCapabilities{} mockBackendRegistry.EXPECT().List(gomock.Any()).Return(nil).AnyTimes() mockDiscoveryMgr.EXPECT().Discover(gomock.Any(), gomock.Any()).Return(emptyAggCaps, nil).AnyTimes() // Stop is called when the server is stopped (not via httptest but via session manager cleanup). mockDiscoveryMgr.EXPECT().Stop().AnyTimes() + // tools/call routes through core.CallTool → BackendClient.CallTool (the session factory's + // own CallTool is bypassed on the Serve path). Return a deterministic result so call + // tests can assert on it. + mockBackendClient.EXPECT(). + CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(&vmcp.ToolCallResult{Content: []vmcp.Content{{Type: "text", Text: "fake result"}}}, nil). + AnyTimes() rt := router.NewDefaultRouter() @@ -202,6 +233,7 @@ func buildTestServerWithOptions( SessionTTL: 5 * time.Minute, SessionFactory: factory, OptimizerFactory: opts.optimizerFactory, + Aggregator: newStubAggregator(capsFromTools(opts.tools)), }, rt, mockBackendClient, @@ -255,7 +287,7 @@ func TestIntegration_SessionManagement_Initialize(t *testing.T) { testTool := vmcp.Tool{Name: "test-tool", Description: "a test tool"} factory, state := newMockFactory(t, ctrl, []vmcp.Tool{testTool}) - ts := buildTestServer(t, factory) + ts := buildTestServer(t, factory, []vmcp.Tool{testTool}) // Step 1: Send initialize request. initReq := map[string]any{ @@ -306,13 +338,15 @@ func TestIntegration_SessionManagement_Initialize(t *testing.T) { require.Equal(t, http.StatusOK, toolResp.StatusCode, "tool call should succeed; body: %s", string(body)) - // The mock session's CallTool should have been invoked. + // On the New/Serve path the tool call routes through core.CallTool → + // BackendClient.CallTool; the factory session's own CallTool is bypassed. Assert the + // backend result surfaced rather than the (now-unused) mock-session CallTool flag. state.mu.Lock() lastSession := state.lastSession state.mu.Unlock() require.NotNil(t, lastSession, "factory should have created a session") - assert.True(t, state.callToolCalled.Load(), - "CallTool on the mock session should have been invoked by the tool call request") + assert.Contains(t, string(body), "fake result", + "tool call should route through the core to the backend client and return its result") } // deleteMCP sends a DELETE request to /mcp with the given session ID and @@ -346,7 +380,7 @@ func TestIntegration_SessionManagement_Termination(t *testing.T) { testTool := vmcp.Tool{Name: "test-tool", Description: "a test tool"} factory, state := newMockFactory(t, ctrl, []vmcp.Tool{testTool}) - ts := buildTestServer(t, factory) + ts := buildTestServer(t, factory, []vmcp.Tool{testTool}) // Step 1: Initialize and obtain a session ID. initReq := map[string]any{ @@ -382,12 +416,9 @@ func TestIntegration_SessionManagement_Termination(t *testing.T) { state.mu.Unlock() require.NotNil(t, lastSession, "factory should have created a session") - // Subsequent requests with the terminated session ID are rejected. - // After Terminate() deletes the session from storage, the discovery middleware passes - // through (no session found → skip capability injection), and the SDK's Validate() - // returns HTTP 404 for the unknown session ID. - // This request also triggers the lazy eviction: GetMultiSession → checkSession → - // ErrExpired → onEvict → Close(). + // Subsequent requests with the terminated session ID are rejected: Terminate() + // deleted the session from storage, so the SDK's Validate() returns HTTP 404 for the + // unknown session ID before any handler runs. toolCallReq := map[string]any{ "jsonrpc": "2.0", "id": 2, @@ -402,13 +433,14 @@ func TestIntegration_SessionManagement_Termination(t *testing.T) { assert.Equal(t, http.StatusNotFound, postResp.StatusCode, "request with terminated session ID should be rejected") - // Close() is called lazily by onEvict when the stale cache entry is - // evicted on the first GetMultiSession call after Terminate deleted the - // session from storage (triggered by the POST above). - assert.Eventually(t, func() bool { - return state.closed.Load() - }, 2*time.Second, 10*time.Millisecond, - "Close() should have been called on the MultiSession after termination") + // Note: unlike the legacy server.New path, a rejected request no longer drives lazy + // cache eviction. The legacy discovery middleware called GetMultiSession on every + // request (which evicted the stale entry and closed the MultiSession); on the + // New/Serve path that middleware is gone and the 404 is returned by the SDK's + // Validate() before any handler, so the node-local MultiSession is closed when the + // cache entry is later evicted by TTL or capacity. MultiSession.Close-on-eviction is + // covered by the sessionmanager cache tests; the parity target here — DELETE + // terminates the session and subsequent requests are rejected — is asserted above. } // TestIntegration_SessionManagement_TokenBinding verifies end-to-end token binding security: @@ -436,7 +468,7 @@ func TestIntegration_SessionManagement_TokenBinding(t *testing.T) { ctrl := gomock.NewController(t) testTool := vmcp.Tool{Name: "echo", Description: "echoes input"} factory, state := newMockFactory(t, ctrl, []vmcp.Tool{testTool}) - ts := buildTestServer(t, factory) + ts := buildTestServer(t, factory, []vmcp.Tool{testTool}) tokenA := "bearer-token-A" tokenB := "bearer-token-B" @@ -655,6 +687,7 @@ func TestIntegration_SessionManagement_CompositeTools(t *testing.T) { } ts := buildTestServerWithOptions(t, factory, serverOptions{ + tools: []vmcp.Tool{backendTool}, workflowDefs: map[string]*composer.WorkflowDefinition{ "composite-tool": workflowDef, }, @@ -717,6 +750,7 @@ func TestIntegration_SessionManagement_CompositeToolConflict(t *testing.T) { } ts := buildTestServerWithOptions(t, factory, serverOptions{ + tools: []vmcp.Tool{{Name: sharedName, Description: "backend version"}}, workflowDefs: map[string]*composer.WorkflowDefinition{ sharedName: workflowDef, }, @@ -813,6 +847,7 @@ func TestIntegration_SessionManagement_CompositeToolsFilteredForSession(t *testi } ts := buildTestServerWithOptions(t, factory, serverOptions{ + tools: []vmcp.Tool{allowedTool}, workflowDefs: map[string]*composer.WorkflowDefinition{ "accessible-workflow": accessibleDef, "restricted-workflow": restrictedDef, @@ -864,6 +899,7 @@ func TestIntegration_SessionManagement_OptimizerMode(t *testing.T) { factory, _ := newMockFactory(t, ctrl, []vmcp.Tool{testTool}) ts := buildTestServerWithOptions(t, factory, serverOptions{ + tools: []vmcp.Tool{testTool}, optimizerFactory: func(_ context.Context, _ []mcpsdk.ServerTool) (optimizer.Optimizer, error) { return &fakeOptimizer{}, nil }, diff --git a/pkg/vmcp/server/session_management_realbackend_integration_test.go b/pkg/vmcp/server/session_management_realbackend_integration_test.go index 79ff3f5351..c5ec5b9afd 100644 --- a/pkg/vmcp/server/session_management_realbackend_integration_test.go +++ b/pkg/vmcp/server/session_management_realbackend_integration_test.go @@ -22,6 +22,7 @@ import ( vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" + vmcpclient "github.com/stacklok/toolhive/pkg/vmcp/client" discoveryMocks "github.com/stacklok/toolhive/pkg/vmcp/discovery/mocks" "github.com/stacklok/toolhive/pkg/vmcp/mocks" "github.com/stacklok/toolhive/pkg/vmcp/router" @@ -44,7 +45,6 @@ func newRealTestHandler(t *testing.T, backendURL string) http.Handler { ctrl := gomock.NewController(t) t.Cleanup(ctrl.Finish) - mockBackendClient := mocks.NewMockBackendClient(ctrl) mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) mockBackendRegistry := mocks.NewMockBackendRegistry(ctrl) @@ -55,9 +55,12 @@ func newRealTestHandler(t *testing.T, backendURL string) http.Handler { TransportType: "streamable-http", } - // BackendRegistry.List() is called by CreateSession() to build the backend list. - // Discover() is not called in the session management path (WithSessionScopedRouting skips discovery). + // BackendRegistry.List() is consumed by the core's on-demand aggregation and by + // CreateSession(). Get() resolves the backend display name during session registration + // (audit labelling). Discover() is no longer called (the discovery middleware is guarded + // off on the Serve path); the AnyTimes expectation tolerates zero calls. mockBackendRegistry.EXPECT().List(gomock.Any()).Return([]vmcp.Backend{backend}).AnyTimes() + mockBackendRegistry.EXPECT().Get(gomock.Any(), gomock.Any()).Return(&backend).AnyTimes() mockDiscoveryMgr.EXPECT().Discover(gomock.Any(), gomock.Any()). Return(&aggregator.AggregatedCapabilities{}, nil).AnyTimes() mockDiscoveryMgr.EXPECT().Stop().AnyTimes() @@ -69,6 +72,16 @@ func newRealTestHandler(t *testing.T, backendURL string) http.Handler { )) factory := vmcpsession.NewSessionFactory(authReg) + // Real backend client + aggregator: server.New routes through core.New + Serve, so the + // core sources the advertised set by querying the real backend through this client and + // routes tools/call back through it. A priority resolver keeps raw tool names (the + // prefix resolver would rename "echo" → "real-backend_echo"). + backendClient, err := vmcpclient.NewHTTPBackendClient(authReg) + require.NoError(t, err) + resolver, err := aggregator.NewPriorityConflictResolver([]string{backend.Name}) + require.NoError(t, err) + agg := aggregator.NewDefaultAggregator(backendClient, resolver, nil, nil) + rt := router.NewDefaultRouter() srv, err := server.New( context.Background(), @@ -77,9 +90,10 @@ func newRealTestHandler(t *testing.T, backendURL string) http.Handler { Port: 0, SessionTTL: 5 * time.Minute, SessionFactory: factory, + Aggregator: agg, }, rt, - mockBackendClient, + backendClient, mockDiscoveryMgr, mockBackendRegistry, nil, diff --git a/pkg/vmcp/server/status.go b/pkg/vmcp/server/status.go index b29f8335e8..84e830e022 100644 --- a/pkg/vmcp/server/status.go +++ b/pkg/vmcp/server/status.go @@ -79,12 +79,8 @@ func (s *Server) buildStatusResponse(ctx context.Context) StatusResponse { // /status reflects the same runtime health as /api/backends/health. // Skip the call — and the map allocation — entirely when monitoring is // disabled. Reading from a nil map is safe in Go and returns zero values. - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() - var liveHealthStates map[string]*health.State - if healthMon != nil { + if healthMon := s.backendHealth(); healthMon != nil { liveHealthStates = healthMon.GetAllBackendStates() } diff --git a/pkg/vmcp/server/status_reporting.go b/pkg/vmcp/server/status_reporting.go index fc384cc10b..21a8bc18b8 100644 --- a/pkg/vmcp/server/status_reporting.go +++ b/pkg/vmcp/server/status_reporting.go @@ -59,10 +59,7 @@ func (s *Server) periodicStatusReporting(ctx context.Context, config StatusRepor // Wait for initial health checks to complete before first status report // This ensures that the first status report has accurate health information // rather than reporting with backendCount=0 before checks complete - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() - if healthMon != nil { + if healthMon := s.backendHealth(); healthMon != nil { slog.Debug("waiting for initial health checks to complete before first status report") healthMon.WaitForInitialHealthChecks() slog.Debug("initial health checks complete, proceeding with status reporting") @@ -120,20 +117,15 @@ func (s *Server) reportStatus(ctx context.Context, reporter vmcpstatus.Reporter) if dynamicReg, ok := s.backendRegistry.(vmcp.DynamicRegistry); ok { currentBackends := dynamicReg.List(ctx) slog.Debug("refreshing backends from registry", "backends", len(currentBackends)) - s.healthMonitorMu.RLock() - healthMon := s.healthMonitor - s.healthMonitorMu.RUnlock() - if healthMon != nil { + if healthMon := s.backendHealth(); healthMon != nil { healthMon.UpdateBackends(currentBackends) } } - // Build status from health monitor if available + // Build status from the core-owned health monitor if available var status *vmcp.Status - - s.healthMonitorMu.RLock() - if s.healthMonitor != nil { - status = s.healthMonitor.BuildStatus() + if healthMon := s.backendHealth(); healthMon != nil { + status = healthMon.BuildStatus() } else { // No health monitor - create minimal status status = &vmcp.Status{ @@ -142,7 +134,6 @@ func (s *Server) reportStatus(ctx context.Context, reporter vmcpstatus.Reporter) Timestamp: time.Now(), } } - s.healthMonitorMu.RUnlock() // Log status at debug level slog.Debug("reporting status", diff --git a/pkg/vmcp/server/status_test.go b/pkg/vmcp/server/status_test.go index b5a1645bad..92db210ba3 100644 --- a/pkg/vmcp/server/status_test.go +++ b/pkg/vmcp/server/status_test.go @@ -251,7 +251,7 @@ func createTestServerWithHealthMonitor( Port: port, GroupRef: groupRef, HealthMonitorConfig: healthMonCfg, - SessionFactory: newNoopMockFactory(t), + SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), }, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) require.NoError(t, err) diff --git a/pkg/vmcp/server/stub_aggregator_test.go b/pkg/vmcp/server/stub_aggregator_test.go new file mode 100644 index 0000000000..15b6423aa4 --- /dev/null +++ b/pkg/vmcp/server/stub_aggregator_test.go @@ -0,0 +1,68 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "context" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" +) + +// stubAggregator is a drop-in aggregator.Aggregator for server tests that route through +// core.New (which requires a non-nil Aggregator). Now that server.New delegates to +// core.New + Serve, the core is the single source of the advertised capability set, so +// tests provide the tools the core should advertise here rather than through the session +// factory. +// +// AggregateCapabilities — the only method the core invokes — returns a fixed result; the +// other interface methods are unreachable on the New/Serve path and fail loudly if called. +type stubAggregator struct { + caps *aggregator.AggregatedCapabilities +} + +var _ aggregator.Aggregator = (*stubAggregator)(nil) + +// newStubAggregator returns a stub whose AggregateCapabilities yields caps. A nil caps +// yields an empty (no-capability) result, which is all transport-only tests need. +func newStubAggregator(caps *aggregator.AggregatedCapabilities) *stubAggregator { + if caps == nil { + caps = &aggregator.AggregatedCapabilities{} + } + return &stubAggregator{caps: caps} +} + +func (s *stubAggregator) AggregateCapabilities( + context.Context, []vmcp.Backend, +) (*aggregator.AggregatedCapabilities, error) { + return s.caps, nil +} + +func (*stubAggregator) QueryCapabilities(context.Context, vmcp.Backend) (*aggregator.BackendCapabilities, error) { + panic("stubAggregator.QueryCapabilities: unexpected call on the New/Serve path") +} + +func (*stubAggregator) QueryAllCapabilities( + context.Context, []vmcp.Backend, +) (map[string]*aggregator.BackendCapabilities, error) { + panic("stubAggregator.QueryAllCapabilities: unexpected call on the New/Serve path") +} + +func (*stubAggregator) ResolveConflicts( + context.Context, map[string]*aggregator.BackendCapabilities, +) (*aggregator.ResolvedCapabilities, error) { + panic("stubAggregator.ResolveConflicts: unexpected call on the New/Serve path") +} + +func (*stubAggregator) MergeCapabilities( + context.Context, *aggregator.ResolvedCapabilities, vmcp.BackendRegistry, +) (*aggregator.AggregatedCapabilities, error) { + panic("stubAggregator.MergeCapabilities: unexpected call on the New/Serve path") +} + +func (*stubAggregator) ProcessPreQueriedCapabilities( + context.Context, map[string][]vmcp.Tool, map[string]*vmcp.BackendTarget, +) ([]vmcp.Tool, []vmcp.Tool, map[string]*vmcp.BackendTarget, error) { + panic("stubAggregator.ProcessPreQueriedCapabilities: unexpected call on the New/Serve path") +} diff --git a/pkg/vmcp/server/telemetry_integration_test.go b/pkg/vmcp/server/telemetry_integration_test.go index cd8e245bc9..30b00263ef 100644 --- a/pkg/vmcp/server/telemetry_integration_test.go +++ b/pkg/vmcp/server/telemetry_integration_test.go @@ -10,7 +10,6 @@ import ( "errors" "io" "net/http" - "sync" "testing" "time" @@ -30,37 +29,18 @@ import ( ) // --------------------------------------------------------------------------- -// backendAwareTestSession / backendAwareTestFactory +// backendAwareTestFactory // --------------------------------------------------------------------------- -// Used by TestIntegration_TelemetryMiddleware to verify that tool calls reach -// the monitorBackends-wrapped backend client so backend-level metrics are recorded. - -// backendClientRef holds a vmcp.BackendClient that can be set after server -// creation, once the monitorBackends-wrapped client is available. -type backendClientRef struct { - mu sync.Mutex - client vmcp.BackendClient -} - -func (r *backendClientRef) set(c vmcp.BackendClient) { - r.mu.Lock() - defer r.mu.Unlock() - r.client = c -} +// Provides a bound MultiSession for TestIntegration_TelemetryMiddleware. On the New/Serve +// path the core sources the advertised set and routes tool calls through its own +// monitorBackends-wrapped backend client, so this factory only needs to open a bound +// session — its capability/call methods are not exercised by the Serve path. -func (r *backendClientRef) get() vmcp.BackendClient { - r.mu.Lock() - defer r.mu.Unlock() - return r.client -} - -// backendAwareTestSession delegates CallTool to the wrapped backend client so -// that monitorBackends instrumentation is exercised during tool calls. +// backendAwareTestSession is a minimal bound MultiSession. type backendAwareTestSession struct { transportsession.Session tools []vmcp.Tool routingTable *vmcp.RoutingTable - clientRef *backendClientRef } func (s *backendAwareTestSession) Tools() []vmcp.Tool { return s.tools } @@ -70,66 +50,63 @@ func (*backendAwareTestSession) Prompts() []vmcp.Prompt { return func (*backendAwareTestSession) BackendSessions() map[string]string { return nil } func (s *backendAwareTestSession) GetRoutingTable() *vmcp.RoutingTable { return s.routingTable } func (*backendAwareTestSession) Close() error { return nil } -func (s *backendAwareTestSession) CallTool( - ctx context.Context, _ *auth.Identity, toolName string, args map[string]any, meta map[string]any, + +// CallTool is not exercised on the Serve path (the core routes tool calls); it returns an +// empty result to satisfy the MultiSession interface. +func (*backendAwareTestSession) CallTool( + context.Context, *auth.Identity, string, map[string]any, map[string]any, ) (*vmcp.ToolCallResult, error) { - client := s.clientRef.get() - if s.routingTable == nil || client == nil { - return &vmcp.ToolCallResult{Content: []vmcp.Content{}}, nil - } - target, ok := s.routingTable.Tools[toolName] - if !ok { - return &vmcp.ToolCallResult{Content: []vmcp.Content{}}, nil - } - return client.CallTool(ctx, target, toolName, args, meta) + return &vmcp.ToolCallResult{Content: []vmcp.Content{}}, nil } func (*backendAwareTestSession) ReadResource( - _ context.Context, _ *auth.Identity, _ string, + context.Context, *auth.Identity, string, ) (*vmcp.ResourceReadResult, error) { return nil, errors.New("not implemented") } func (*backendAwareTestSession) GetPrompt( - _ context.Context, _ *auth.Identity, _ string, _ map[string]any, + context.Context, *auth.Identity, string, map[string]any, ) (*vmcp.PromptGetResult, error) { return nil, errors.New("not implemented") } -// backendAwareTestFactory creates backendAwareTestSessions. +// backendAwareTestFactory creates bound backendAwareTestSessions. type backendAwareTestFactory struct { tools []vmcp.Tool routingTable *vmcp.RoutingTable - clientRef *backendClientRef } var _ vmcpsession.MultiSessionFactory = (*backendAwareTestFactory)(nil) -func newBackendAwareTestFactory(tools []vmcp.Tool, rt *vmcp.RoutingTable) (*backendAwareTestFactory, *backendClientRef) { - ref := &backendClientRef{} - return &backendAwareTestFactory{tools: tools, routingTable: rt, clientRef: ref}, ref +func newBackendAwareTestFactory(tools []vmcp.Tool, rt *vmcp.RoutingTable) *backendAwareTestFactory { + return &backendAwareTestFactory{tools: tools, routingTable: rt} } func (f *backendAwareTestFactory) MakeSessionWithID( _ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, ) (vmcpsession.MultiSession, error) { - return &backendAwareTestSession{ - Session: transportsession.NewStreamableSession(id), - tools: f.tools, - routingTable: f.routingTable, - clientRef: f.clientRef, - }, nil + return f.newSession(id), nil } func (f *backendAwareTestFactory) RestoreSession( _ context.Context, id string, _ map[string]string, _ []*vmcp.Backend, ) (vmcpsession.MultiSession, error) { - return &backendAwareTestSession{ + return f.newSession(id), nil +} + +// newSession builds a session carrying the unauthenticated identity-binding sentinel. +// The Serve-path enforceSessionBinding reads it via GetMetadataValue before each +// core.CallTool; this test has no auth middleware, so the caller is anonymous and the +// sentinel admits it. +func (f *backendAwareTestFactory) newSession(id string) *backendAwareTestSession { + sess := &backendAwareTestSession{ Session: transportsession.NewStreamableSession(id), tools: f.tools, routingTable: f.routingTable, - clientRef: f.clientRef, - }, nil + } + sess.SetMetadata(vmcpsession.MetadataKeyIdentityBinding, "unauthenticated") + return sess } // TestIntegration_TelemetryMiddleware tests that the vMCP server records telemetry @@ -253,7 +230,14 @@ func TestIntegration_TelemetryMiddleware(t *testing.T) { // client with monitorBackends() which instruments outgoing backend calls. // Use backendAwareTestFactory so that CallTool delegates to the monitorBackends-wrapped // backendClient, ensuring toolhive_vmcp_backend_requests metrics are recorded. - telemetryFactory, clientRef := newBackendAwareTestFactory(telemetryTools, telemetryRoutingTable) + // The core sources the advertised set by aggregating over mockBackendClient (prefix + // resolver → "search-svc_search"). core.New wraps this same client with monitorBackends + // for telemetry, and core.CallTool routes tool calls through that wrapped client — so + // the backend instrumentation (toolhive_vmcp_backend_requests) is exercised without the + // session factory needing to hold the wrapped client. + telemetryAgg := aggregator.NewDefaultAggregator( + mockBackendClient, aggregator.NewPrefixConflictResolver("{workload}_"), nil, nil) + telemetryFactory := newBackendAwareTestFactory(telemetryTools, telemetryRoutingTable) srv, err := New(ctx, &Config{ Name: "telemetry-vmcp", Version: "1.0.0", @@ -261,11 +245,9 @@ func TestIntegration_TelemetryMiddleware(t *testing.T) { Port: 0, // Random available port TelemetryProvider: telemetryProvider, SessionFactory: telemetryFactory, + Aggregator: telemetryAgg, }, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) require.NoError(t, err) - // Wire the monitorBackends-wrapped client into the session factory so that - // tool calls go through the telemetry instrumentation layer. - clientRef.set(srv.backendClient) // Start server serverCtx, cancelServer := context.WithCancel(ctx) diff --git a/pkg/vmcp/server/testfactory_test.go b/pkg/vmcp/server/testfactory_test.go index e9b7e62145..6ad6aff45d 100644 --- a/pkg/vmcp/server/testfactory_test.go +++ b/pkg/vmcp/server/testfactory_test.go @@ -9,6 +9,7 @@ import ( "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" ) @@ -37,3 +38,47 @@ func (*minimalTestFactory) RestoreSession( ) (vmcpsession.MultiSession, error) { return nil, fmt.Errorf("minimalTestFactory: RestoreSession not implemented in test helper") } + +// stubAggregator is a drop-in aggregator.Aggregator for internal package server tests +// that route through core.New (which requires a non-nil Aggregator). AggregateCapabilities +// — the only method the core invokes — returns an empty result; the rest are unreachable on +// the New/Serve path and fail loudly if called. Internal tests only need the empty +// (no-capability) result; the external server_test stub (stub_aggregator_test.go) takes +// configurable caps. +type stubAggregator struct{} + +var _ aggregator.Aggregator = (*stubAggregator)(nil) + +func (*stubAggregator) AggregateCapabilities( + context.Context, []vmcp.Backend, +) (*aggregator.AggregatedCapabilities, error) { + return &aggregator.AggregatedCapabilities{}, nil +} + +func (*stubAggregator) QueryCapabilities(context.Context, vmcp.Backend) (*aggregator.BackendCapabilities, error) { + panic("stubAggregator.QueryCapabilities: unexpected call on the New/Serve path") +} + +func (*stubAggregator) QueryAllCapabilities( + context.Context, []vmcp.Backend, +) (map[string]*aggregator.BackendCapabilities, error) { + panic("stubAggregator.QueryAllCapabilities: unexpected call on the New/Serve path") +} + +func (*stubAggregator) ResolveConflicts( + context.Context, map[string]*aggregator.BackendCapabilities, +) (*aggregator.ResolvedCapabilities, error) { + panic("stubAggregator.ResolveConflicts: unexpected call on the New/Serve path") +} + +func (*stubAggregator) MergeCapabilities( + context.Context, *aggregator.ResolvedCapabilities, vmcp.BackendRegistry, +) (*aggregator.AggregatedCapabilities, error) { + panic("stubAggregator.MergeCapabilities: unexpected call on the New/Serve path") +} + +func (*stubAggregator) ProcessPreQueriedCapabilities( + context.Context, map[string][]vmcp.Tool, map[string]*vmcp.BackendTarget, +) ([]vmcp.Tool, []vmcp.Tool, map[string]*vmcp.BackendTarget, error) { + panic("stubAggregator.ProcessPreQueriedCapabilities: unexpected call on the New/Serve path") +}