From 37ee399c3faa442ee6da871d84a9af4eb8497a60 Mon Sep 17 00:00:00 2001 From: lorenzozanee Date: Mon, 10 Aug 2026 06:59:16 +0800 Subject: [PATCH] Honor vMCP partial failure mode and timeouts Signed-off-by: lorenzozanee --- pkg/vmcp/aggregator/aggregator.go | 4 +- pkg/vmcp/aggregator/default_aggregator.go | 80 ++++- .../default_aggregator_operational_test.go | 273 ++++++++++++++++++ pkg/vmcp/cli/serve.go | 17 +- pkg/vmcp/cli/serve_test.go | 47 +++ pkg/vmcp/core/core_vmcp.go | 5 +- 6 files changed, 418 insertions(+), 8 deletions(-) create mode 100644 pkg/vmcp/aggregator/default_aggregator_operational_test.go diff --git a/pkg/vmcp/aggregator/aggregator.go b/pkg/vmcp/aggregator/aggregator.go index 12a976044e..b077f705ff 100644 --- a/pkg/vmcp/aggregator/aggregator.go +++ b/pkg/vmcp/aggregator/aggregator.go @@ -41,7 +41,9 @@ type Aggregator interface { QueryCapabilities(ctx context.Context, backend vmcp.Backend) (*BackendCapabilities, error) // QueryAllCapabilities queries all backends for their capabilities in parallel. - // Handles backend failures gracefully (logs and continues with remaining backends). + // How backend failures are treated follows the configured partial failure mode: + // "best_effort" (or unset) logs and continues with remaining backends, while + // "fail" makes the first failing backend fail the whole query. QueryAllCapabilities(ctx context.Context, backends []vmcp.Backend) (map[string]*BackendCapabilities, error) // ResolveConflicts applies conflict resolution strategy to handle diff --git a/pkg/vmcp/aggregator/default_aggregator.go b/pkg/vmcp/aggregator/default_aggregator.go index c8d07d8ae2..3d732573a7 100644 --- a/pkg/vmcp/aggregator/default_aggregator.go +++ b/pkg/vmcp/aggregator/default_aggregator.go @@ -11,6 +11,7 @@ import ( "slices" "sort" "sync" + "time" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -38,6 +39,44 @@ type defaultAggregator struct { // construction. promptNaming promptNaming tracer trace.Tracer + + // partialFailureMode wires operational.failureHandling.partialFailureMode + // ("fail" | "best_effort"). "" (no WithOperationalConfig option, or a nil + // config) preserves the pre-wiring behavior: best-effort. See + // QueryAllCapabilities. + partialFailureMode string + // timeoutDefault is the default per-backend query timeout. Zero (unset) + // leaves no per-query deadline. + timeoutDefault time.Duration + // timeoutPerWorkload overrides timeoutDefault for named backends. + timeoutPerWorkload map[string]time.Duration +} + +// Option configures the default aggregator. +type Option func(*defaultAggregator) + +// WithOperationalConfig wires operational settings (partial failure mode and +// backend request timeouts) into the aggregator so the configured behavior +// actually takes effect. Nil and zero-valued fields leave the aggregator's +// default behavior (best-effort, no per-query deadline) unchanged. +func WithOperationalConfig(cfg *config.OperationalConfig) Option { + return func(a *defaultAggregator) { + if cfg == nil { + return + } + if cfg.FailureHandling != nil { + a.partialFailureMode = cfg.FailureHandling.PartialFailureMode + } + if cfg.Timeouts != nil { + a.timeoutDefault = time.Duration(cfg.Timeouts.Default) + if len(cfg.Timeouts.PerWorkload) > 0 { + a.timeoutPerWorkload = make(map[string]time.Duration, len(cfg.Timeouts.PerWorkload)) + for backendID, d := range cfg.Timeouts.PerWorkload { + a.timeoutPerWorkload[backendID] = time.Duration(d) + } + } + } + } } // NewDefaultAggregator creates a new default aggregator implementation. @@ -45,11 +84,13 @@ type defaultAggregator struct { // aggregationConfig specifies aggregation settings including tool filtering/overrides, // excludeAllTools, and defaultToolVisibility. // tracerProvider is used to create a tracer for distributed tracing (pass nil for no tracing). +// opts apply operational settings such as partial failure mode and backend request timeouts. func NewDefaultAggregator( backendClient vmcp.BackendClient, conflictResolver ConflictResolver, aggregationConfig *config.AggregationConfig, tracerProvider trace.TracerProvider, + opts ...Option, ) Aggregator { // Build tool config map for quick lookup by backend ID toolConfigMap := make(map[string]*config.WorkloadToolConfig) @@ -76,7 +117,7 @@ func NewDefaultAggregator( tracer = noop.NewTracerProvider().Tracer("github.com/stacklok/toolhive/pkg/vmcp/aggregator") } - return &defaultAggregator{ + a := &defaultAggregator{ backendClient: backendClient, conflictResolver: conflictResolver, toolConfigMap: toolConfigMap, @@ -85,6 +126,10 @@ func NewDefaultAggregator( promptNaming: promptNamingFromConfig(aggregationConfig), tracer: tracer, } + for _, opt := range opts { + opt(a) + } + return a } // QueryCapabilities queries a single backend for its MCP capabilities. @@ -146,7 +191,17 @@ func (a *defaultAggregator) QueryCapabilities(ctx context.Context, backend vmcp. } // QueryAllCapabilities queries all backends for their capabilities in parallel. -// Handles backend failures gracefully (logs and continues with remaining backends). +// Each backend query runs under the timeout configured via +// WithOperationalConfig (the default, or a per-workload override). +// +// How backend failures are treated is controlled by +// operational.failureHandling.partialFailureMode: +// - "best_effort" (or unset): a failing backend is logged and skipped, and +// the capabilities of the healthy backends are returned. +// - "fail": the first failing backend fails the whole query; the errgroup +// cancels the in-flight queries to the remaining backends. +// +// The function always fails when no backend returned capabilities. func (a *defaultAggregator) QueryAllCapabilities( ctx context.Context, backends []vmcp.Backend, @@ -178,10 +233,27 @@ func (a *defaultAggregator) QueryAllCapabilities( for _, backend := range backends { backend := backend // Capture loop variable g.Go(func() error { - caps, err := a.QueryCapabilities(ctx, backend) + // Apply the configured per-backend timeout, if any. + queryCtx := ctx + timeout := a.timeoutDefault + if t, ok := a.timeoutPerWorkload[backend.ID]; ok { + timeout = t + } + if timeout > 0 { + var cancel context.CancelFunc + queryCtx, cancel = context.WithTimeout(ctx, timeout) + defer cancel() + } + + caps, err := a.QueryCapabilities(queryCtx, backend) if err != nil { - // Log the error but continue with other backends + // Log the error but continue with other backends. slog.Warn("failed to query backend", "backend", backend.ID, "error", err) + if a.partialFailureMode == "fail" { + // Fail fast: returning an error cancels the derived context, + // stopping the remaining in-flight backend queries. + return err + } return nil // Don't fail the entire operation } diff --git a/pkg/vmcp/aggregator/default_aggregator_operational_test.go b/pkg/vmcp/aggregator/default_aggregator_operational_test.go new file mode 100644 index 0000000000..d25af803be --- /dev/null +++ b/pkg/vmcp/aggregator/default_aggregator_operational_test.go @@ -0,0 +1,273 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package aggregator + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/config" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" +) + +// TestDefaultAggregator_QueryAllCapabilities_FailMode pins the contract of +// operational.failureHandling.partialFailureMode=fail: when any backend query +// fails, QueryAllCapabilities must surface an error instead of silently +// continuing with the remaining backends (the pre-wiring best-effort behavior). +func TestDefaultAggregator_QueryAllCapabilities_FailMode(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + backends := []vmcp.Backend{ + newTestBackend("backend1"), + newTestBackend("backend2"), + } + + caps1 := newTestCapabilityList(withTools(newTestTool("tool1", "backend1"))) + + mockClient.EXPECT().ListCapabilities(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + if target.WorkloadID == "backend1" { + return caps1, nil + } + return nil, errors.New("connection timeout") + }).Times(2) + + agg := NewDefaultAggregator(mockClient, nil, nil, nil, + WithOperationalConfig(&config.OperationalConfig{ + FailureHandling: &config.FailureHandlingConfig{ + PartialFailureMode: "fail", + }, + })) + + result, err := agg.QueryAllCapabilities(context.Background(), backends) + require.Error(t, err) + assert.Nil(t, result) + assert.ErrorIs(t, err, ErrBackendQueryFailed) +} + +// TestDefaultAggregator_QueryAllCapabilities_FailModeCancelsInflight pins the +// fail-fast contract of partialFailureMode=fail at the goroutine level: when one +// backend query fails, the errgroup must cancel the derived context so the +// remaining in-flight backend queries stop immediately instead of running to the +// outer deadline. backend2 blocks on ctx.Done() with no per-query timeout +// configured, so only errgroup cancellation can unblock it. Under a regression +// that collects the error instead of returning it, backend2 stays blocked until +// the outer deadline, blowing the elapsed upper bound below. +func TestDefaultAggregator_QueryAllCapabilities_FailModeCancelsInflight(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + backends := []vmcp.Backend{ + newTestBackend("backend1"), + newTestBackend("backend2"), + } + + // backend1 fails immediately; backend2 blocks until the derived context is + // cancelled. No timeouts are configured, so nothing but fail-fast + // cancellation can unblock backend2. + mockClient.EXPECT().ListCapabilities(gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + if target.WorkloadID == "backend1" { + return nil, errors.New("backend1 unavailable") + } + <-ctx.Done() + return nil, ctx.Err() + }).Times(2) + + agg := NewDefaultAggregator(mockClient, nil, nil, nil, + WithOperationalConfig(&config.OperationalConfig{ + FailureHandling: &config.FailureHandlingConfig{ + PartialFailureMode: "fail", + }, + })) + + // The outer deadline is only a safety net so the test cannot hang forever + // under a regression: with correct fail-fast wiring backend2 unblocks in + // milliseconds, far below the 2s upper bound (and well under the 3s net). + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + start := time.Now() + result, err := agg.QueryAllCapabilities(ctx, backends) + elapsed := time.Since(start) + + require.Error(t, err) + assert.Nil(t, result) + assert.Less(t, elapsed, 2*time.Second, + "fail-fast must cancel in-flight queries; backend2 blocked until the outer deadline") +} + +// TestDefaultAggregator_QueryAllCapabilities_BestEffortMode pins the contract of +// operational.failureHandling.partialFailureMode=best_effort: a failing backend +// is logged and skipped, and the capabilities of the healthy backends are +// returned. +func TestDefaultAggregator_QueryAllCapabilities_BestEffortMode(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + backends := []vmcp.Backend{ + newTestBackend("backend1"), + newTestBackend("backend2"), + } + + caps1 := newTestCapabilityList(withTools(newTestTool("tool1", "backend1"))) + + mockClient.EXPECT().ListCapabilities(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + if target.WorkloadID == "backend1" { + return caps1, nil + } + return nil, errors.New("connection timeout") + }).Times(2) + + agg := NewDefaultAggregator(mockClient, nil, nil, nil, + WithOperationalConfig(&config.OperationalConfig{ + FailureHandling: &config.FailureHandlingConfig{ + PartialFailureMode: "best_effort", + }, + })) + + result, err := agg.QueryAllCapabilities(context.Background(), backends) + require.NoError(t, err) + require.Len(t, result, 1) + assert.Contains(t, result, "backend1") + assert.NotContains(t, result, "backend2") +} + +// TestDefaultAggregator_QueryAllCapabilities_TimeoutPropagation pins the contract +// of operational.timeouts: each backend query runs under a context whose deadline +// reflects the per-backend timeout (the default, or the perWorkload override). +func TestDefaultAggregator_QueryAllCapabilities_TimeoutPropagation(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + backends := []vmcp.Backend{ + newTestBackend("backend1"), + newTestBackend("backend2"), + } + + caps1 := newTestCapabilityList(withTools(newTestTool("tool1", "backend1"))) + caps2 := newTestCapabilityList(withTools(newTestTool("tool2", "backend2"))) + + var mu sync.Mutex + observed := make(map[string]time.Duration) + recordDeadline := func(backendID string, ctx context.Context) { + deadline, ok := ctx.Deadline() + if !ok { + return + } + mu.Lock() + observed[backendID] = time.Until(deadline) + mu.Unlock() + } + + mockClient.EXPECT().ListCapabilities(gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + recordDeadline(target.WorkloadID, ctx) + if target.WorkloadID == "backend1" { + return caps1, nil + } + return caps2, nil + }).Times(2) + + agg := NewDefaultAggregator(mockClient, nil, nil, nil, + WithOperationalConfig(&config.OperationalConfig{ + Timeouts: &config.TimeoutConfig{ + Default: config.Duration(100 * time.Millisecond), + PerWorkload: map[string]config.Duration{"backend2": config.Duration(200 * time.Millisecond)}, + }, + })) + + result, err := agg.QueryAllCapabilities(context.Background(), backends) + require.NoError(t, err) + require.Len(t, result, 2) + + want := map[string]time.Duration{ + "backend1": 100 * time.Millisecond, + "backend2": 200 * time.Millisecond, + } + mu.Lock() + defer mu.Unlock() + for backendID, wantTimeout := range want { + got, ok := observed[backendID] + require.Truef(t, ok, "no query deadline was recorded for backend %s", backendID) + assert.GreaterOrEqual(t, got, wantTimeout-50*time.Millisecond, + "deadline for %s must not be shorter than the configured timeout", backendID) + assert.LessOrEqual(t, got, wantTimeout+50*time.Millisecond, + "deadline for %s must not exceed the configured timeout", backendID) + } +} + +// TestDefaultAggregator_QueryAllCapabilities_FailModeWithTimeout pins the combined +// contract of partialFailureMode=fail and a per-backend timeout: a backend that +// hangs past its deadline is failed fast and the error surfaces (wrapping +// context.DeadlineExceeded) rather than blocking the whole aggregation. +func TestDefaultAggregator_QueryAllCapabilities_FailModeWithTimeout(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + backends := []vmcp.Backend{ + newTestBackend("backend1"), + newTestBackend("backend2"), + } + + caps1 := newTestCapabilityList(withTools(newTestTool("tool1", "backend1"))) + + // backend1 answers immediately; backend2 blocks until its per-query + // deadline fires and then reports the deadline error. Waiting on + // ctx.Done() keeps the test deterministic (no wall-clock sleeps). + mockClient.EXPECT().ListCapabilities(gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + if target.WorkloadID == "backend1" { + return caps1, nil + } + <-ctx.Done() + return nil, ctx.Err() + }).Times(2) + + agg := NewDefaultAggregator(mockClient, nil, nil, nil, + WithOperationalConfig(&config.OperationalConfig{ + FailureHandling: &config.FailureHandlingConfig{ + PartialFailureMode: "fail", + }, + Timeouts: &config.TimeoutConfig{ + Default: config.Duration(10 * time.Millisecond), + }, + })) + + // The outer deadline is only a safety net so the test cannot hang forever + // under a regression: with the timeout wired, backend2 unblocks after ~10ms + // and the aggregation fails fast. The elapsed upper bound below + // distinguishes that from a regression where the per-backend timeout is not + // wired and only the outer deadline unblocks backend2 (~3s). + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + start := time.Now() + result, err := agg.QueryAllCapabilities(ctx, backends) + require.Error(t, err) + assert.Nil(t, result) + assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, time.Since(start), 1*time.Second, + "per-backend timeout must fire in ~10ms, not the outer deadline") +} diff --git a/pkg/vmcp/cli/serve.go b/pkg/vmcp/cli/serve.go index 0f7efa9bcb..b151aa9b58 100644 --- a/pkg/vmcp/cli/serve.go +++ b/pkg/vmcp/cli/serve.go @@ -206,7 +206,7 @@ func Serve(ctx context.Context, cfg ServeConfig) error { if telemetryProvider != nil { tracerProvider = telemetryProvider.TracerProvider() } - agg := aggregator.NewDefaultAggregator(backendClient, conflictResolver, vmcpCfg.Aggregation, tracerProvider) + agg := newAggregator(backendClient, conflictResolver, vmcpCfg, tracerProvider) // DynamicRegistry tracks backends for dynamic discovery in Kubernetes mode. dynamicRegistry := vmcp.NewDynamicRegistry(backends) @@ -473,6 +473,21 @@ func Serve(ctx context.Context, cfg ServeConfig) error { return srv.Start(ctx) } +// newAggregator constructs the default aggregator and wires the vMCP +// configuration's operational settings (partial failure mode, backend request +// timeouts) into it. Extracted from Serve so the production wiring is directly +// testable: dropping the WithOperationalConfig option below would silently +// revert production to best-effort behavior. +func newAggregator( + backendClient vmcp.BackendClient, + conflictResolver aggregator.ConflictResolver, + vmcpCfg *config.Config, + tracerProvider trace.TracerProvider, +) aggregator.Aggregator { + return aggregator.NewDefaultAggregator(backendClient, conflictResolver, vmcpCfg.Aggregation, tracerProvider, + aggregator.WithOperationalConfig(vmcpCfg.Operational)) +} + // embeddingManager is the minimal interface over *EmbeddingServiceManager needed // by the Serve lifecycle. Defined here to allow stub injection in unit tests; // production code passes a *EmbeddingServiceManager. diff --git a/pkg/vmcp/cli/serve_test.go b/pkg/vmcp/cli/serve_test.go index 88e05312cd..070db86756 100644 --- a/pkg/vmcp/cli/serve_test.go +++ b/pkg/vmcp/cli/serve_test.go @@ -4,6 +4,8 @@ package cli import ( + "context" + "errors" "fmt" "os" "path/filepath" @@ -237,6 +239,51 @@ func TestGenerateQuickModeConfig(t *testing.T) { } } +// TestNewAggregator_WiresOperationalFailureMode pins the production aggregator +// wiring behind Serve (via newAggregator): the +// operational.failureHandling.partialFailureMode=fail setting from the loaded +// vMCP config must reach the aggregator the CLI constructs. If the +// WithOperationalConfig option is dropped from newAggregator, the aggregator +// silently reverts to best-effort behavior and this test fails. +func TestNewAggregator_WiresOperationalFailureMode(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := vmcpmocks.NewMockBackendClient(ctrl) + conflictResolver := aggregatormocks.NewMockConflictResolver(ctrl) + + backends := []vmcp.Backend{ + {ID: "backend1", Name: "backend1", BaseURL: "http://127.0.0.1:9001/sse", TransportType: "sse"}, + {ID: "backend2", Name: "backend2", BaseURL: "http://127.0.0.1:9002/sse", TransportType: "sse"}, + } + + caps2 := &vmcp.CapabilityList{Tools: []vmcp.Tool{}} + mockClient.EXPECT().ListCapabilities(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + if target.WorkloadID == "backend1" { + return nil, errors.New("backend1 unavailable") + } + return caps2, nil + }).Times(2) + + vmcpCfg := &config.Config{ + Operational: &config.OperationalConfig{ + FailureHandling: &config.FailureHandlingConfig{ + PartialFailureMode: "fail", + }, + }, + } + + agg := newAggregator(mockClient, conflictResolver, vmcpCfg, nil) + result, err := agg.QueryAllCapabilities(context.Background(), backends) + + // With fail mode wired, backend1's error must surface and no partial result + // may be returned. + require.Error(t, err) + assert.Nil(t, result) +} + // TestServe_NeitherConfigNorGroup verifies that Serve returns an error when // both --config and --group are absent. func TestServe_NeitherConfigNorGroup(t *testing.T) { diff --git a/pkg/vmcp/core/core_vmcp.go b/pkg/vmcp/core/core_vmcp.go index e5730a4b21..47f8dbcf50 100644 --- a/pkg/vmcp/core/core_vmcp.go +++ b/pkg/vmcp/core/core_vmcp.go @@ -565,8 +565,9 @@ func (c *coreVMCP) aggregatedView(ctx context.Context) (*aggregator.AggregatedCa // decides the backend set: aggregatedView passes the health-filtered subset (the // data path), while authorizedBackends passes the full registry (backend // visibility, per #5741: health is a status, not a visibility filter). It exists so -// both share one aggregation error-wrap. Unreachable backends fail their live -// capability query and simply contribute nothing. +// both share one aggregation error-wrap. Whether a backend that fails its live +// capability query fails the whole aggregation or merely contributes nothing +// follows the configured partial failure mode (fail vs best_effort). func (c *coreVMCP) aggregateBackends( ctx context.Context, backends []vmcp.Backend, ) (*aggregator.AggregatedCapabilities, error) {