Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
97 changes: 0 additions & 97 deletions pkg/vmcp/aggregator/advertised_backendid_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -311,100 +311,3 @@ func TestDefaultAggregator_AdvertisingFilterPreservesBackendID(t *testing.T) {
})
}
}

func TestDefaultAggregator_ProcessPreQueriedCapabilitiesCarryBackendID(t *testing.T) {
t.Parallel()

newTarget := func(id string) *vmcp.BackendTarget {
return &vmcp.BackendTarget{
WorkloadID: id,
WorkloadName: id + "-name",
BaseURL: "http://" + id + ":8080",
TransportType: "streamable-http",
HealthStatus: vmcp.BackendHealthy,
}
}

manualConfigs := []*config.WorkloadToolConfig{
{Workload: "backend1", Overrides: map[string]*config.ToolOverride{"fetch": {Name: "custom_fetch_b1"}}},
{Workload: "backend2", Overrides: map[string]*config.ToolOverride{"fetch": {Name: "custom_fetch_b2"}}},
}

tests := []struct {
name string
strategy vmcp.ConflictResolutionStrategy
priorityOrder []string
wlConfigs []*config.WorkloadToolConfig
aggCfg *config.AggregationConfig
}{
{
name: "prefix strategy",
strategy: vmcp.ConflictStrategyPrefix,
},
{
name: "priority strategy",
strategy: vmcp.ConflictStrategyPriority,
priorityOrder: []string{"backend1", "backend2"},
},
{
name: "manual strategy",
strategy: vmcp.ConflictStrategyManual,
wlConfigs: manualConfigs,
aggCfg: &config.AggregationConfig{Tools: manualConfigs},
},
{
name: "prefix strategy with global ExcludeAllTools",
strategy: vmcp.ConflictStrategyPrefix,
aggCfg: &config.AggregationConfig{ExcludeAllTools: true},
},
{
name: "prefix strategy with per-workload Filter",
strategy: vmcp.ConflictStrategyPrefix,
aggCfg: &config.AggregationConfig{
Tools: []*config.WorkloadToolConfig{{Workload: "backend1", Filter: []string{"tool_a"}}},
},
},
}

for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()

toolsByBackend := map[string][]vmcp.Tool{
"backend1": {newTestTool("fetch", "backend1"), newTestTool("tool_a", "backend1")},
"backend2": {newTestTool("fetch", "backend2"), newTestTool("tool_b", "backend2")},
}
targets := map[string]*vmcp.BackendTarget{
"backend1": newTarget("backend1"),
"backend2": newTarget("backend2"),
}

resolver := resolverForStrategy(t, tt.strategy, tt.priorityOrder, tt.wlConfigs)
agg := NewDefaultAggregator(nil, resolver, tt.aggCfg, nil)

advertised, allResolved, routingTable, err := agg.ProcessPreQueriedCapabilities(
context.Background(), toolsByBackend, targets,
)
require.NoError(t, err)

// allResolvedTools is the full set regardless of advertising filter;
// every entry must carry a BackendID for composite-tool schema lookup.
require.NotEmpty(t, allResolved)
for _, tool := range allResolved {
assert.NotEmptyf(t, tool.BackendID,
"resolved tool %q must carry a non-empty BackendID", tool.Name)
}
// Every advertised (post-filter) tool must carry a BackendID, and its
// routing target must agree on the backend identity.
for _, tool := range advertised {
assert.NotEmptyf(t, tool.BackendID,
"advertised tool %q must carry a non-empty BackendID", tool.Name)
target := routingTable[tool.Name]
require.NotNilf(t, target, "routing target for %q", tool.Name)
assert.Equal(t, tool.BackendID, target.WorkloadID,
"advertised tool %q BackendID must match its routing target", tool.Name)
}
})
}
}
19 changes: 0 additions & 19 deletions pkg/vmcp/aggregator/aggregator.go
Original file line number Diff line number Diff line change
Expand Up @@ -61,25 +61,6 @@ type Aggregator interface {
// 2. Resolve conflicts
// 3. Merge into final view
AggregateCapabilities(ctx context.Context, backends []vmcp.Backend) (*AggregatedCapabilities, error)

// ProcessPreQueriedCapabilities applies the same aggregation pipeline (overrides,
// conflict resolution, advertising filter) to tools that have already been fetched
// from live backends. Used by the session management path to reuse aggregator
// logic without re-querying backends over HTTP.
//
// toolsByBackend maps backend WorkloadID → raw tools as returned by the backend.
// targets maps backend WorkloadID → the pre-built BackendTarget for that backend.
//
// Returns:
// - advertisedTools: resolved tools that pass the advertising filter (for MCP clients)
// - allResolvedTools: all resolved tools including non-advertised ones (for schema lookup)
// - toolsRouting: routing table keyed by resolved name; each entry has OriginalCapabilityName
// set so that GetBackendCapabilityName() translates back to the raw backend name.
ProcessPreQueriedCapabilities(
ctx context.Context,
toolsByBackend map[string][]vmcp.Tool,
targets map[string]*vmcp.BackendTarget,
) (advertisedTools []vmcp.Tool, allResolvedTools []vmcp.Tool, toolsRouting map[string]*vmcp.BackendTarget, err error)
}

// BackendCapabilities contains the raw capabilities from a single backend.
Expand Down
63 changes: 0 additions & 63 deletions pkg/vmcp/aggregator/default_aggregator.go
Original file line number Diff line number Diff line change
Expand Up @@ -495,69 +495,6 @@ func (a *defaultAggregator) AggregateCapabilities(
return aggregated, nil
}

// ProcessPreQueriedCapabilities implements Aggregator.ProcessPreQueriedCapabilities.
// It reuses processBackendTools, ResolveConflicts, and shouldAdvertiseTool so that
// the session path applies identical transforms to the aggregation path.
func (a *defaultAggregator) ProcessPreQueriedCapabilities(
ctx context.Context,
toolsByBackend map[string][]vmcp.Tool,
targets map[string]*vmcp.BackendTarget,
) ([]vmcp.Tool, []vmcp.Tool, map[string]*vmcp.BackendTarget, error) {
// Step 1: Apply per-backend overrides (renames, description changes).
processed := make(map[string]*BackendCapabilities, len(toolsByBackend))
for backendID, rawTools := range toolsByBackend {
processed[backendID] = &BackendCapabilities{
BackendID: backendID,
Tools: processBackendTools(ctx, backendID, rawTools, a.toolConfigMap[backendID]),
}
}

// Step 2: Resolve naming conflicts across backends.
resolved, err := a.ResolveConflicts(ctx, processed)
if err != nil {
return nil, nil, nil, err
}

// Step 3: Build advertised list, all-resolved list, and routing table.
// advertisedTools is the subset shown to MCP clients (post-filter).
// allResolvedTools includes every resolved tool regardless of advertising filter,
// so that workflow engines can look up InputSchema for type coercion even when
// a backend tool is hidden from clients via excludeAll or filter configuration.
var advertisedTools []vmcp.Tool
var allResolvedTools []vmcp.Tool
routingTable := make(map[string]*vmcp.BackendTarget, len(resolved.Tools))

for _, rt := range resolved.Tools {
target, ok := targets[rt.BackendID]
if !ok {
slog.Warn("ProcessPreQueriedCapabilities: no target for backend, skipping tool",
"backend", rt.BackendID, "tool", rt.ResolvedName)
continue
}
// Clone the target and record the actual backend capability name for call routing.
// rt.OriginalName is the post-override name; reverse the override map to get the
// actual name the backend itself uses.
t := *target
t.OriginalCapabilityName = actualBackendCapabilityName(a.toolConfigMap, rt.BackendID, rt.OriginalName)
routingTable[rt.ResolvedName] = &t

resolved := vmcp.Tool{
Name: rt.ResolvedName,
Description: rt.Description,
InputSchema: rt.InputSchema,
OutputSchema: rt.OutputSchema,
Annotations: rt.Annotations,
BackendID: rt.BackendID,
}
allResolvedTools = append(allResolvedTools, resolved)
if a.shouldAdvertiseTool(rt.BackendID, rt.OriginalName) {
advertisedTools = append(advertisedTools, resolved)
}
}

return advertisedTools, allResolvedTools, routingTable, nil
}

// actualBackendCapabilityName returns the real capability name the backend uses,
// reversing any per-backend override rename that processBackendTools may have applied.
//
Expand Down
Loading
Loading