Add simple ToolOverrides for tools with variant definitions

SamMorrowDrums · SamMorrowDrums · commit fe5b222a17d6 · 2025-12-18T00:11:05.000+01:00
Provides a lightweight mechanism for tools that have different definitions
based on runtime conditions (e.g., feature flags, capabilities).

Designed for the small number of tools (~2) that need variant handling:
- ToolOverride: condition + replacement definition
- ToolOverrides: map[string]ToolOverride with Apply() and ApplyToTools()
- Lazy allocation: no allocs when no overrides match

This is intentionally simple - only 60 lines - because most tools don't
need variants and the full VariantIndex would be overkill.
diff --git a/pkg/inventory/tool_variants.go b/pkg/inventory/tool_variants.go
@@ -0,0 +1,86 @@
+package inventory
+
+import "context"
+
+// ToolOverride allows replacing a tool's definition based on runtime conditions.
+// Use this for the small number of tools that have different schemas/handlers
+// depending on features, capabilities, or environment.
+type ToolOverride struct {
+	// ToolName is the canonical tool name to override
+	ToolName string
+
+	// Condition returns true if this override should apply
+	Condition func(ctx context.Context) (bool, error)
+
+	// Override is the replacement tool definition
+	Override ServerTool
+}
+
+// ToolOverrides is a simple map for the few tools that need variant handling.
+// Key is the tool name, value is the override to check.
+type ToolOverrides map[string]ToolOverride
+
+// Apply checks if an override should be used for the given tool.
+// Returns the override if condition matches, nil otherwise.
+func (o ToolOverrides) Apply(ctx context.Context, toolName string) *ServerTool {
+	override, ok := o[toolName]
+	if !ok {
+		return nil
+	}
+
+	if override.Condition == nil {
+		return &override.Override
+	}
+
+	matches, err := override.Condition(ctx)
+	if err != nil || !matches {
+		return nil
+	}
+
+	return &override.Override
+}
+
+// ApplyToTools applies overrides to a list of tools, returning a new list
+// with overridden tools replaced. Tools without overrides are unchanged.
+// If no overrides match, returns the original slice (no allocation).
+func (o ToolOverrides) ApplyToTools(ctx context.Context, tools []*ServerTool) []*ServerTool {
+	if len(o) == 0 {
+		return tools
+	}
+
+	// First pass: check if any overrides apply (avoid allocation if not)
+	var result []*ServerTool
+	for i, tool := range tools {
+		override, hasOverride := o[tool.Tool.Name]
+		if !hasOverride {
+			if result != nil {
+				result[i] = tool
+			}
+			continue
+		}
+
+		// Check condition
+		var applies bool
+		if override.Condition == nil {
+			applies = true
+		} else if matches, err := override.Condition(ctx); err == nil && matches {
+			applies = true
+		}
+
+		if applies {
+			// Lazy allocation only when we find a match
+			if result == nil {
+				result = make([]*ServerTool, len(tools))
+				copy(result[:i], tools[:i])
+			}
+			result[i] = &override.Override
+		} else if result != nil {
+			result[i] = tool
+		}
+	}
+
+	if result == nil {
+		return tools // No overrides matched, return original
+	}
+	return result
+}
diff --git a/pkg/inventory/tool_variants_test.go b/pkg/inventory/tool_variants_test.go
@@ -0,0 +1,185 @@
+package inventory
+
+import (
+	"context"
+	"testing"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/stretchr/testify/assert"
+)
+
+func makeTool(name string) ServerTool {
+	return ServerTool{
+		Tool: mcp.Tool{
+			Name:        name,
+			Description: "Tool " + name,
+		},
+	}
+}
+
+func makeOverride(name, desc string) ServerTool {
+	return ServerTool{
+		Tool: mcp.Tool{
+			Name:        name,
+			Description: desc,
+		},
+	}
+}
+
+func TestToolOverrides_Apply(t *testing.T) {
+	t.Parallel()
+
+	overrides := ToolOverrides{
+		"create_issue": {
+			ToolName:  "create_issue",
+			Condition: func(_ context.Context) (bool, error) { return true, nil },
+			Override:  makeOverride("create_issue", "Enterprise variant"),
+		},
+	}
+
+	ctx := context.Background()
+
+	// Tool with override
+	result := overrides.Apply(ctx, "create_issue")
+	assert.NotNil(t, result)
+	assert.Equal(t, "Enterprise variant", result.Tool.Description)
+
+	// Tool without override
+	result = overrides.Apply(ctx, "list_repos")
+	assert.Nil(t, result)
+}
+
+func TestToolOverrides_Apply_ConditionFalse(t *testing.T) {
+	t.Parallel()
+
+	overrides := ToolOverrides{
+		"create_issue": {
+			ToolName:  "create_issue",
+			Condition: func(_ context.Context) (bool, error) { return false, nil },
+			Override:  makeOverride("create_issue", "Enterprise variant"),
+		},
+	}
+
+	ctx := context.Background()
+
+	// Condition doesn't match - no override
+	result := overrides.Apply(ctx, "create_issue")
+	assert.Nil(t, result)
+}
+
+func TestToolOverrides_Apply_NilCondition(t *testing.T) {
+	t.Parallel()
+
+	overrides := ToolOverrides{
+		"create_issue": {
+			ToolName: "create_issue",
+			// nil Condition - always applies
+			Override: makeOverride("create_issue", "Always applied"),
+		},
+	}
+
+	ctx := context.Background()
+
+	result := overrides.Apply(ctx, "create_issue")
+	assert.NotNil(t, result)
+	assert.Equal(t, "Always applied", result.Tool.Description)
+}
+
+func TestToolOverrides_ApplyToTools(t *testing.T) {
+	t.Parallel()
+
+	tools := []*ServerTool{
+		ptr(makeTool("create_issue")),
+		ptr(makeTool("list_repos")),
+		ptr(makeTool("get_me")),
+	}
+
+	overrides := ToolOverrides{
+		"create_issue": {
+			ToolName:  "create_issue",
+			Condition: func(_ context.Context) (bool, error) { return true, nil },
+			Override:  makeOverride("create_issue", "Enterprise create_issue"),
+		},
+	}
+
+	ctx := context.Background()
+	result := overrides.ApplyToTools(ctx, tools)
+
+	assert.Len(t, result, 3)
+	assert.Equal(t, "Enterprise create_issue", result[0].Tool.Description)
+	assert.Equal(t, "Tool list_repos", result[1].Tool.Description)
+	assert.Equal(t, "Tool get_me", result[2].Tool.Description)
+}
+
+func TestToolOverrides_ApplyToTools_Empty(t *testing.T) {
+	t.Parallel()
+
+	tools := []*ServerTool{
+		ptr(makeTool("create_issue")),
+	}
+
+	overrides := ToolOverrides{}
+
+	ctx := context.Background()
+	result := overrides.ApplyToTools(ctx, tools)
+
+	// Empty overrides returns original slice
+	assert.Equal(t, tools, result)
+}
+
+func ptr(t ServerTool) *ServerTool {
+	return &t
+}
+
+func BenchmarkToolOverrides_ApplyToTools(b *testing.B) {
+	// 130 tools, 2 overrides (realistic)
+	tools := make([]*ServerTool, 130)
+	for i := range tools {
+		tools[i] = ptr(makeTool("tool_" + string(rune('a'+i%26))))
+	}
+
+	overrides := ToolOverrides{
+		"tool_a": {
+			ToolName:  "tool_a",
+			Condition: func(_ context.Context) (bool, error) { return true, nil },
+			Override:  makeOverride("tool_a", "Override A"),
+		},
+		"tool_b": {
+			ToolName:  "tool_b",
+			Condition: func(_ context.Context) (bool, error) { return true, nil },
+			Override:  makeOverride("tool_b", "Override B"),
+		},
+	}
+
+	ctx := context.Background()
+
+	b.ReportAllocs() // Only count allocs in the hot loop
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_ = overrides.ApplyToTools(ctx, tools)
+	}
+}
+
+func BenchmarkToolOverrides_ApplyToTools_NoMatch(b *testing.B) {
+	// 130 tools, overrides don't match any - should be zero alloc
+	tools := make([]*ServerTool, 130)
+	for i := range tools {
+		tools[i] = ptr(makeTool("tool_" + string(rune('a'+i%26))))
+	}
+
+	// Override exists but for a tool not in list
+	overrides := ToolOverrides{
+		"nonexistent_tool": {
+			ToolName: "nonexistent_tool",
+			Override: makeOverride("nonexistent_tool", "Override"),
+		},
+	}
+
+	ctx := context.Background()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_ = overrides.ApplyToTools(ctx, tools)
+	}
+}
