Add QueryWithFeatureChecker to minimize feature flag service calls

Adds two new methods to ToolIndex:

- UniqueFeatureFlags(): Returns all unique feature flags used by tools
- QueryWithFeatureChecker(): Queries with automatic flag checking

Key optimization: Instead of checking the same feature flag N times
for N tools that use it, each unique flag is checked exactly once.

For example, if 50 tools use 3 unique flags, this reduces feature flag
service calls from O(50) to O(3).

Test demonstrates: 50 tools using 3 flags results in exactly 3 checks.

Author: SamMorrowDrums

pkg/inventory/tool_index.go

@@ -280,3 +280,84 @@ func (idx *ToolIndex) ToolsetIDs() []ToolsetID {
 	sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
 	return ids
 }
+
+// UniqueFeatureFlags returns all unique feature flags referenced by tools in the index.
+// This allows callers to check only the flags that matter, once per request.
+func (idx *ToolIndex) UniqueFeatureFlags() []string {
+	seen := make(map[string]bool)
+	for flag := range idx.requiresFeature {
+		seen[flag] = true
+	}
+	for flag := range idx.disabledByFeature {
+		seen[flag] = true
+	}
+
+	flags := make([]string, 0, len(seen))
+	for flag := range seen {
+		flags = append(flags, flag)
+	}
+	sort.Strings(flags)
+	return flags
+}
+
+// QueryWithFeatureChecker performs a query, automatically checking feature flags.
+// Each unique flag is checked exactly once via the checker function.
+// This minimizes feature flag service calls from O(tools) to O(unique flags).
+func (idx *ToolIndex) QueryWithFeatureChecker(ctx context.Context, cfg QueryConfigWithChecker) QueryResult {
+	// Check each unique feature flag exactly once
+	var enabledFeatures, disabledFeatures []string
+
+	for flag := range idx.requiresFeature {
+		enabled, err := cfg.FeatureChecker(ctx, flag)
+		if err != nil || !enabled {
+			disabledFeatures = append(disabledFeatures, flag)
+		} else {
+			enabledFeatures = append(enabledFeatures, flag)
+		}
+	}
+
+	for flag := range idx.disabledByFeature {
+		// Only check if we haven't already
+		alreadyChecked := false
+		for _, f := range enabledFeatures {
+			if f == flag {
+				alreadyChecked = true
+				break
+			}
+		}
+		for _, f := range disabledFeatures {
+			if f == flag {
+				alreadyChecked = true
+				break
+			}
+		}
+		if !alreadyChecked {
+			enabled, err := cfg.FeatureChecker(ctx, flag)
+			if err != nil || !enabled {
+				disabledFeatures = append(disabledFeatures, flag)
+			} else {
+				enabledFeatures = append(enabledFeatures, flag)
+			}
+		}
+	}
+
+	// Delegate to the standard Query with resolved flags
+	return idx.Query(QueryConfig{
+		AllToolsets:      cfg.AllToolsets,
+		EnabledToolsets:  cfg.EnabledToolsets,
+		AdditionalTools:  cfg.AdditionalTools,
+		ReadOnly:         cfg.ReadOnly,
+		EnabledFeatures:  enabledFeatures,
+		DisabledFeatures: disabledFeatures,
+	})
+}
+
+// QueryConfigWithChecker is like QueryConfig but uses a checker function
+// instead of pre-resolved feature flag lists.
+type QueryConfigWithChecker struct {
+	AllToolsets     bool
+	EnabledToolsets []ToolsetID
+	AdditionalTools []string
+	ReadOnly        bool
+	FeatureChecker  FeatureFlagChecker // Reuses the existing FeatureFlagChecker type from filters.go
+}

pkg/inventory/tool_index_test.go

@@ -491,6 +491,122 @@ func TestToolIndex_ToolsetBitmap(t *testing.T) {
 	assert.True(t, bmX.IsEmpty())
 }
 
+func TestToolIndex_UniqueFeatureFlags(t *testing.T) {
+	t.Parallel()
+
+	testTools := []ServerTool{
+		mockServerToolInToolset("basic_tool", "tools", true),
+		mockServerToolWithFeatureFlag("advanced_tool", "tools", "feature_a", ""),
+		mockServerToolWithFeatureFlag("experimental_tool", "tools", "feature_b", ""),
+		mockServerToolWithFeatureFlag("legacy_tool", "tools", "", "feature_c"),
+		mockServerToolWithFeatureFlag("complex_tool", "tools", "feature_a", "feature_d"), // reuses feature_a
+	}
+
+	index := BuildToolIndex(testTools)
+
+	flags := index.UniqueFeatureFlags()
+
+	// Should have 4 unique flags: feature_a, feature_b, feature_c, feature_d
+	assert.Len(t, flags, 4)
+	assert.Contains(t, flags, "feature_a")
+	assert.Contains(t, flags, "feature_b")
+	assert.Contains(t, flags, "feature_c")
+	assert.Contains(t, flags, "feature_d")
+}
+
+func TestToolIndex_QueryWithFeatureChecker(t *testing.T) {
+	t.Parallel()
+
+	testTools := []ServerTool{
+		mockServerToolInToolset("basic_tool", "tools", true),
+		mockServerToolWithFeatureFlag("needs_feature_a", "tools", "feature_a", ""),
+		mockServerToolWithFeatureFlag("needs_feature_b", "tools", "feature_b", ""),
+		mockServerToolWithFeatureFlag("disabled_by_feature_c", "tools", "", "feature_c"),
+	}
+
+	index := BuildToolIndex(testTools)
+
+	// Track which flags were checked
+	checkedFlags := make(map[string]int)
+
+	checker := func(_ context.Context, flag string) (bool, error) {
+		checkedFlags[flag]++
+		// Enable feature_a, disable feature_b, enable feature_c
+		switch flag {
+		case "feature_a":
+			return true, nil
+		case "feature_b":
+			return false, nil
+		case "feature_c":
+			return true, nil // This will disable "disabled_by_feature_c"
+		default:
+			return false, nil
+		}
+	}
+
+	ctx := context.Background()
+	result := index.QueryWithFeatureChecker(ctx, QueryConfigWithChecker{
+		EnabledToolsets: []ToolsetID{"tools"},
+		ReadOnly:        false,
+		FeatureChecker:  checker,
+	})
+
+	// Each flag should be checked exactly once
+	assert.Equal(t, 1, checkedFlags["feature_a"], "feature_a should be checked once")
+	assert.Equal(t, 1, checkedFlags["feature_b"], "feature_b should be checked once")
+	assert.Equal(t, 1, checkedFlags["feature_c"], "feature_c should be checked once")
+
+	// Materialize and verify results
+	tools := index.Materialize(ctx, result)
+	names := make([]string, len(tools))
+	for i, tool := range tools {
+		names[i] = tool.Tool.Name
+	}
+
+	// basic_tool: no flags, included
+	// needs_feature_a: feature_a enabled, included
+	// needs_feature_b: feature_b disabled, excluded
+	// disabled_by_feature_c: feature_c enabled, excluded
+	assert.Len(t, tools, 2)
+	assert.Contains(t, names, "basic_tool")
+	assert.Contains(t, names, "needs_feature_a")
+	assert.NotContains(t, names, "needs_feature_b")
+	assert.NotContains(t, names, "disabled_by_feature_c")
+}
+
+func TestToolIndex_QueryWithFeatureChecker_MinimizesChecks(t *testing.T) {
+	t.Parallel()
+
+	// Create 50 tools that all use the same 3 feature flags
+	testTools := make([]ServerTool, 50)
+	for i := 0; i < 50; i++ {
+		flag := []string{"flag_a", "flag_b", "flag_c"}[i%3]
+		testTools[i] = mockServerToolWithFeatureFlag(
+			"tool_"+string(rune('a'+i%26)),
+			"tools",
+			flag, // All tools require one of 3 flags
+			"",
+		)
+	}
+
+	index := BuildToolIndex(testTools)
+
+	checkCount := 0
+	checker := func(_ context.Context, _ string) (bool, error) {
+		checkCount++
+		return true, nil
+	}
+
+	ctx := context.Background()
+	_ = index.QueryWithFeatureChecker(ctx, QueryConfigWithChecker{
+		EnabledToolsets: []ToolsetID{"tools"},
+		FeatureChecker:  checker,
+	})
+
+	// Should only check 3 unique flags, not 50 tools
+	assert.Equal(t, 3, checkCount, "Should check each unique flag exactly once")
+}
+
 func BenchmarkBuildToolIndex_130Tools(b *testing.B) {
 	// Create realistic toolset distribution
 	toolsets := []ToolsetID{"repos", "issues", "pull_requests", "users", "actions", "code_security", "projects", "notifications", "discussions", "experiments"}