Skip to content

Add partial failure mode support to vMCP aggregator #3533

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Closed
yrobla wants to merge 3 commits into from
Closed

Add partial failure mode support to vMCP aggregator #3533

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3556e04
Add circuit breaker pattern to vMCP health monitoring
taskbot Jan 30, 2026
d5c6c36
Integrate circuit breaker with vMCP routing and aggregation
taskbot Jan 30, 2026
1602ebe
Add partial failure mode support to vMCP aggregator
taskbot Jan 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 21 additions & 1 deletion cmd/vmcp/app/commands.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -342,7 +342,14 @@ func runServe(cmd *cobra.Command, _ []string) error {
if telemetryProvider != nil {
tracerProvider = telemetryProvider.TracerProvider()
}
agg := aggregator.NewDefaultAggregator(backendClient, conflictResolver, cfg.Aggregation.Tools, tracerProvider)

// Get failure mode from config (default to "fail" if not specified)
failureMode := config.PartialFailureModeFail
if cfg.Operational != nil && cfg.Operational.FailureHandling != nil && cfg.Operational.FailureHandling.PartialFailureMode != "" {
failureMode = cfg.Operational.FailureHandling.PartialFailureMode
}

agg := aggregator.NewDefaultAggregator(backendClient, conflictResolver, cfg.Aggregation.Tools, tracerProvider, failureMode)

// Use DynamicRegistry for version-based cache invalidation
// Works in both standalone (CLI with YAML config) and Kubernetes (operator-deployed) modes
Expand Down Expand Up @@ -433,6 +440,19 @@ func runServe(cmd *cobra.Command, _ []string) error {
Timeout: defaults.Timeout,
DegradedThreshold: defaults.DegradedThreshold,
}

// Configure circuit breaker if enabled
if cfg.Operational.FailureHandling.CircuitBreaker != nil && cfg.Operational.FailureHandling.CircuitBreaker.Enabled {
healthMonitorConfig.CircuitBreaker = &health.CircuitBreakerConfig{
Enabled: true,
FailureThreshold: cfg.Operational.FailureHandling.CircuitBreaker.FailureThreshold,
Timeout: time.Duration(cfg.Operational.FailureHandling.CircuitBreaker.Timeout),
}
logger.Infof("Circuit breaker enabled (failure threshold: %d, timeout: %v)",
cfg.Operational.FailureHandling.CircuitBreaker.FailureThreshold,
time.Duration(cfg.Operational.FailureHandling.CircuitBreaker.Timeout))
}

logger.Info("Health monitoring configured from operational settings")
}

Expand Down
160 changes: 117 additions & 43 deletions pkg/vmcp/aggregator/default_aggregator.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,17 +26,20 @@ type defaultAggregator struct {
conflictResolver ConflictResolver
toolConfigMap map[string]*config.WorkloadToolConfig // Maps backend ID to tool config
tracer trace.Tracer
failureMode string // "fail" or "best_effort"
}

// NewDefaultAggregator creates a new default aggregator implementation.
// conflictResolver handles tool name conflicts across backends.
// workloadConfigs specifies per-backend tool filtering and overrides.
// tracerProvider is used to create a tracer for distributed tracing (pass nil for no tracing).
// failureMode determines behavior when backends are unavailable ("fail" or "best_effort").
func NewDefaultAggregator(
backendClient vmcp.BackendClient,
conflictResolver ConflictResolver,
workloadConfigs []*config.WorkloadToolConfig,
tracerProvider trace.TracerProvider,
failureMode string,
) Aggregator {
// Build tool config map for quick lookup by backend ID
toolConfigMap := make(map[string]*config.WorkloadToolConfig)
Expand All @@ -54,11 +57,17 @@ func NewDefaultAggregator(
tracer = noop.NewTracerProvider().Tracer("github.com/stacklok/toolhive/pkg/vmcp/aggregator")
}

// Default to fail mode if not specified or invalid
if failureMode != config.PartialFailureModeBestEffort {
failureMode = config.PartialFailureModeFail
}

return &defaultAggregator{
backendClient: backendClient,
conflictResolver: conflictResolver,
toolConfigMap: toolConfigMap,
tracer: tracer,
failureMode: failureMode,
}
}

Expand Down Expand Up @@ -290,7 +299,7 @@ func (a *defaultAggregator) MergeCapabilities(
span.End()
}()

logger.Debugf("Merging capabilities into final view")
logger.Debugf("Merging capabilities into final view (failure mode: %s)", a.failureMode)

// Create routing table
routingTable := &vmcp.RoutingTable{
Expand All @@ -299,68 +308,97 @@ func (a *defaultAggregator) MergeCapabilities(
Prompts: make(map[string]*vmcp.BackendTarget),
}

// Track backend health for failure mode enforcement
healthyBackends := make(map[string]bool)
unhealthyBackends := make(map[string]string) // Maps backend ID to health status

// Convert resolved tools to final vmcp.Tool format
tools := make([]vmcp.Tool, 0, len(resolved.Tools))
for _, resolvedTool := range resolved.Tools {
// Look up full backend information from registry
backend := registry.Get(ctx, resolvedTool.BackendID)
if backend == nil {
logger.Warnf("Backend %s not found in registry for tool %s, skipping",
resolvedTool.BackendID, resolvedTool.ResolvedName)
continue
}

// Track backend health
if backend.HealthStatus.IsHealthyForRouting() {
healthyBackends[backend.ID] = true
} else {
unhealthyBackends[backend.ID] = string(backend.HealthStatus)
}
Comment on lines +326 to +331

@jerm-dro jerm-dro Jan 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Suggestion: The factoring of the defaultAggregator leaves a bit to be desired and I think this change motivates improving it.

Background

The defaultAggregator implements the Aggregator interface, which has 4 unused public methods: QueryCapabilities, QueryAllCapabilities, ResolveConflicts, and MergeCapabilities.

Suggestion

Rather than inlining more logic into defaultAggregator, lets:

  1. Remove the unused methods from the interface.
  2. Implement a circuitBreakerAggregator for your new logic. This implements the new and improved Aggregator interface and decorates any arbitrary aggregator:
type circuitBreakerAggregator struct {
   // this would really just be the defaultAggregator
    inner Aggregator
}


func (c circuitBreakerAggregator) AggregateCapabilities(...) {
    // check all the backends for health
    err := c.enforceFailureMode(...)
    if err != nil {
        return nil, err
    }
    return inner.AggregateCapabilities(...)
}

I think this makes your change cleaner, because it doesn't add any complexity to the existing defaultAggregator's construction or runtime behavior. It would also enable pretty straightforward unit testing of the circuitBreakerAggregator, because you don't need the whole defaultAggregator to test it.

What do you think?


// Filter out tools from unhealthy backends
if !backend.HealthStatus.IsHealthyForRouting() {
logger.Debugf("Skipping tool %s from unhealthy backend %s (status: %s)",
resolvedTool.ResolvedName, backend.Name, backend.HealthStatus)
continue
}

tools = append(tools, vmcp.Tool{
Name: resolvedTool.ResolvedName,
Description: resolvedTool.Description,
InputSchema: resolvedTool.InputSchema,
BackendID: resolvedTool.BackendID,
})

// Look up full backend information from registry
backend := registry.Get(ctx, resolvedTool.BackendID)
if backend == nil {
logger.Warnf("Backend %s not found in registry for tool %s, creating minimal target",
resolvedTool.BackendID, resolvedTool.ResolvedName)
routingTable.Tools[resolvedTool.ResolvedName] = &vmcp.BackendTarget{
WorkloadID: resolvedTool.BackendID,
OriginalCapabilityName: resolvedTool.OriginalName,
}
} else {
// Use the backendToTarget helper from registry package
target := vmcp.BackendToTarget(backend)
// Store the original tool name for forwarding to backend
target.OriginalCapabilityName = resolvedTool.OriginalName
routingTable.Tools[resolvedTool.ResolvedName] = target
}
// Use the backendToTarget helper from registry package
target := vmcp.BackendToTarget(backend)
// Store the original tool name for forwarding to backend
target.OriginalCapabilityName = resolvedTool.OriginalName
routingTable.Tools[resolvedTool.ResolvedName] = target
}

// Add resources to routing table
// Add resources to routing table (with health filtering)
resources := make([]vmcp.Resource, 0, len(resolved.Resources))
for _, resource := range resolved.Resources {
backend := registry.Get(ctx, resource.BackendID)
if backend == nil {
logger.Warnf("Backend %s not found in registry for resource %s, creating minimal target",
logger.Warnf("Backend %s not found in registry for resource %s, skipping",
resource.BackendID, resource.URI)
routingTable.Resources[resource.URI] = &vmcp.BackendTarget{
WorkloadID: resource.BackendID,
OriginalCapabilityName: resource.URI,
}
} else {
target := vmcp.BackendToTarget(backend)
// Store the original resource URI for forwarding to backend
target.OriginalCapabilityName = resource.URI
routingTable.Resources[resource.URI] = target
continue
}

// Filter out resources from unhealthy backends
if !backend.HealthStatus.IsHealthyForRouting() {
logger.Debugf("Skipping resource %s from unhealthy backend %s (status: %s)",
resource.URI, backend.Name, backend.HealthStatus)
continue
}

resources = append(resources, resource)

target := vmcp.BackendToTarget(backend)
// Store the original resource URI for forwarding to backend
target.OriginalCapabilityName = resource.URI
routingTable.Resources[resource.URI] = target
}

// Add prompts to routing table
// Add prompts to routing table (with health filtering)
prompts := make([]vmcp.Prompt, 0, len(resolved.Prompts))
for _, prompt := range resolved.Prompts {
backend := registry.Get(ctx, prompt.BackendID)
if backend == nil {
logger.Warnf("Backend %s not found in registry for prompt %s, creating minimal target",
logger.Warnf("Backend %s not found in registry for prompt %s, skipping",
prompt.BackendID, prompt.Name)
routingTable.Prompts[prompt.Name] = &vmcp.BackendTarget{
WorkloadID: prompt.BackendID,
OriginalCapabilityName: prompt.Name,
}
} else {
target := vmcp.BackendToTarget(backend)
// Store the original prompt name for forwarding to backend
target.OriginalCapabilityName = prompt.Name
routingTable.Prompts[prompt.Name] = target
continue
}

// Filter out prompts from unhealthy backends
if !backend.HealthStatus.IsHealthyForRouting() {
logger.Debugf("Skipping prompt %s from unhealthy backend %s (status: %s)",
prompt.Name, backend.Name, backend.HealthStatus)
continue
}

prompts = append(prompts, prompt)

target := vmcp.BackendToTarget(backend)
// Store the original prompt name for forwarding to backend
target.OriginalCapabilityName = prompt.Name
routingTable.Prompts[prompt.Name] = target
}

// Determine conflict strategy used
Expand All @@ -376,16 +414,16 @@ func (a *defaultAggregator) MergeCapabilities(
// Create final aggregated view
aggregated := &AggregatedCapabilities{
Tools: tools,
Resources: resolved.Resources,
Prompts: resolved.Prompts,
Resources: resources,
Prompts: prompts,
SupportsLogging: resolved.SupportsLogging,
SupportsSampling: resolved.SupportsSampling,
RoutingTable: routingTable,
Metadata: &AggregationMetadata{
BackendCount: 0, // Will be set by caller
ToolCount: len(tools),
ResourceCount: len(resolved.Resources),
PromptCount: len(resolved.Prompts),
ResourceCount: len(resources),
PromptCount: len(prompts),
ConflictStrategy: conflictStrategy,
},
}
Expand All @@ -400,9 +438,45 @@ func (a *defaultAggregator) MergeCapabilities(
logger.Infof("Merged capabilities: %d tools, %d resources, %d prompts",
aggregated.Metadata.ToolCount, aggregated.Metadata.ResourceCount, aggregated.Metadata.PromptCount)

// Enforce partial failure mode
if err := a.enforceFailureMode(healthyBackends, unhealthyBackends); err != nil {
return nil, err
}

return aggregated, nil
}

// enforceFailureMode checks backend health and enforces the configured failure mode.
// In fail mode, returns error if all backends are unavailable.
// In best_effort mode, logs info about unavailable backends.
func (a *defaultAggregator) enforceFailureMode(healthyBackends map[string]bool, unhealthyBackends map[string]string) error {
if a.failureMode == config.PartialFailureModeFail {
// In fail mode, check if we have any healthy backends
if len(healthyBackends) == 0 && len(unhealthyBackends) > 0 {
// All backends are unhealthy - fail the aggregation
backendStatuses := make([]string, 0, len(unhealthyBackends))
for backendID, status := range unhealthyBackends {
backendStatuses = append(backendStatuses, fmt.Sprintf("%s(%s)", backendID, status))
}
return fmt.Errorf("all backends unavailable in fail mode: %v", backendStatuses)
}

// Log warning if some backends are unhealthy
if len(unhealthyBackends) > 0 {
logger.Warnf("Fail mode: %d healthy backends, %d unhealthy backends (continuing with healthy)",
len(healthyBackends), len(unhealthyBackends))
}
} else {
// In best_effort mode, just log if we're operating with reduced capacity
if len(unhealthyBackends) > 0 {
logger.Infof("Best effort mode: %d healthy backends, %d unhealthy backends (continuing with available)",
len(healthyBackends), len(unhealthyBackends))
}
}

return nil
}

// AggregateCapabilities is a convenience method that performs the full aggregation pipeline:
// 1. Create backend registry
// 2. Query all backends
Expand Down
Loading
Loading