refactor(inventory): make feature-flag gating a regular ToolFilter

Move tool feature-flag evaluation out of isToolEnabled and into a
ToolFilter installed at the head of the pipeline by Build() when
WithFeatureChecker received a non-nil checker. The 'no checker = no
filtering' contract is now expressed structurally (the filter isn't
installed) instead of by a runtime nil check inside the helper.

Resources and prompts have no filter pipeline, so they call the now-pure
featureFlagAllowed helper behind an explicit r.featureChecker != nil
guard at the iteration site.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: SamMorrowDrums

pkg/inventory/builder.go

@@ -127,11 +127,20 @@ func (b *Builder) WithTools(toolNames []string) *Builder {
 
 // WithFeatureChecker sets the feature flag checker function.
 // The checker receives a context (for actor extraction) and feature flag name,
-// returns (enabled, error). If error occurs, it will be logged and treated as false.
-// If checker is nil, feature flag filtering is skipped entirely — the inventory
-// then returns the full upper bound, including dual-name variants. The
-// per-request inventory must always install a checker so MCP registration
-// (which only serves one tool per name) sees a deduplicated set.
+// and returns (enabled, error). Errors are logged and treated as "not enabled".
+//
+// When the checker is non-nil, Build() installs a feature-flag ToolFilter
+// at the head of the filter pipeline so that tools annotated with
+// FeatureFlagEnable / FeatureFlagDisable are gated accordingly. Resources
+// and prompts use the same checker via an explicit guard at their iteration
+// site.
+//
+// When the checker is nil, no feature-flag filter is installed; tools,
+// resources, and prompts pass through feature-flag gating unchanged. The
+// per-request inventory in HTTP mode must always install a checker so that
+// MCP registration (which can only serve a given tool name once) sees a
+// deduplicated set of dual-name variants.
+//
 // Returns self for chaining.
 func (b *Builder) WithFeatureChecker(checker FeatureFlagChecker) *Builder {
 	b.featureChecker = checker
@@ -203,14 +212,24 @@ func cleanTools(tools []string) []string {
 func (b *Builder) Build() (*Inventory, error) {
 	tools := b.tools
 
+	// Install the feature-flag filter at the head of the pipeline so that
+	// flag-gated tools are excluded before any user-supplied WithFilter sees
+	// them. Doing this in Build() (rather than inside WithFeatureChecker)
+	// keeps the install idempotent — repeated WithFeatureChecker calls
+	// replace the checker without stacking duplicate filters.
+	filters := b.filters
+	if b.featureChecker != nil {
+		filters = append([]ToolFilter{createFeatureFlagFilter(b.featureChecker)}, filters...)
+	}
+
 	r := &Inventory{
 		tools:             tools,
 		resourceTemplates: b.resourceTemplates,
 		prompts:           b.prompts,
 		deprecatedAliases: b.deprecatedAliases,
 		readOnly:          b.readOnly,
 		featureChecker:    b.featureChecker,
-		filters:           b.filters,
+		filters:           filters,
 	}
 
 	// Process toolsets and pre-compute metadata in a single pass

pkg/inventory/filters.go

@@ -35,37 +35,57 @@ func (r *Inventory) checkFeatureFlag(ctx context.Context, flagName string) bool
 	return enabled
 }
 
-// isFeatureFlagAllowed checks if an item passes feature flag filtering.
-//   - If no feature checker is configured, feature flags are not evaluated and
-//     the item passes through. This gives the unfiltered upper bound used by
-//     HTTP mode at startup, before per-request feature checkers exist. The
-//     per-request inventory always installs a checker, so any dual-name
-//     variants are resolved before MCP registration (which can only serve a
-//     given tool name once).
+// featureFlagAllowed reports whether an item with the given enable/disable
+// flag pair is permitted under the supplied checker. The checker must be
+// non-nil — callers that don't want feature filtering should not call this at
+// all (this is also the contract for createFeatureFlagFilter, which is only
+// installed when WithFeatureChecker received a non-nil checker).
+//
 //   - If FeatureFlagEnable is set, the item is only allowed if the flag is enabled.
 //   - If FeatureFlagDisable is set, the item is excluded if the flag is enabled.
-func (r *Inventory) isFeatureFlagAllowed(ctx context.Context, enableFlag, disableFlag string) bool {
-	if r.featureChecker == nil {
-		return true
-	}
-	// Check enable flag - item requires this flag to be on
-	if enableFlag != "" && !r.checkFeatureFlag(ctx, enableFlag) {
-		return false
+func featureFlagAllowed(ctx context.Context, checker FeatureFlagChecker, enableFlag, disableFlag string) bool {
+	if enableFlag != "" {
+		enabled, err := checker(ctx, enableFlag)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Feature flag check error for %q: %v\n", enableFlag, err)
+			return false
+		}
+		if !enabled {
+			return false
+		}
 	}
-	// Check disable flag - item is excluded if this flag is on
-	if disableFlag != "" && r.checkFeatureFlag(ctx, disableFlag) {
-		return false
+	if disableFlag != "" {
+		enabled, err := checker(ctx, disableFlag)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Feature flag check error for %q: %v\n", disableFlag, err)
+			return false
+		}
+		if enabled {
+			return false
+		}
 	}
 	return true
 }
 
+// createFeatureFlagFilter returns a ToolFilter that gates tools on their
+// FeatureFlagEnable / FeatureFlagDisable annotations using the given checker.
+// Builder.Build() installs this filter exactly once when WithFeatureChecker
+// has been called with a non-nil checker, so "no feature filtering" is
+// expressed structurally — by the absence of the filter — rather than by a
+// runtime nil check inside the filter itself.
+func createFeatureFlagFilter(checker FeatureFlagChecker) ToolFilter {
+	return func(ctx context.Context, tool *ServerTool) (bool, error) {
+		return featureFlagAllowed(ctx, checker, tool.FeatureFlagEnable, tool.FeatureFlagDisable), nil
+	}
+}
+
 // isToolEnabled checks if a specific tool is enabled based on current filters.
 // Filter evaluation order:
 //  1. Tool.Enabled (tool self-filtering)
-//  2. FeatureFlagEnable/FeatureFlagDisable
-//  3. Read-only filter
-//  4. Builder filters (via WithFilter)
-//  5. Toolset/additional tools
+//  2. Read-only filter
+//  3. Builder filters (via WithFilter; the feature-flag filter, when
+//     installed via WithFeatureChecker, runs as part of this step)
+//  4. Toolset/additional tools
 func (r *Inventory) isToolEnabled(ctx context.Context, tool *ServerTool) bool {
 	// 1. Check tool's own Enabled function first
 	if tool.Enabled != nil {
@@ -78,15 +98,11 @@ func (r *Inventory) isToolEnabled(ctx context.Context, tool *ServerTool) bool {
 			return false
 		}
 	}
-	// 2. Check feature flags
-	if !r.isFeatureFlagAllowed(ctx, tool.FeatureFlagEnable, tool.FeatureFlagDisable) {
-		return false
-	}
-	// 3. Check read-only filter (applies to all tools)
+	// 2. Check read-only filter (applies to all tools)
 	if r.readOnly && !tool.IsReadOnly() {
 		return false
 	}
-	// 4. Apply builder filters
+	// 3. Apply builder filters (includes the feature-flag filter when set)
 	for _, filter := range r.filters {
 		allowed, err := filter(ctx, tool)
 		if err != nil {
@@ -97,11 +113,11 @@ func (r *Inventory) isToolEnabled(ctx context.Context, tool *ServerTool) bool {
 			return false
 		}
 	}
-	// 5. Check if tool is in additionalTools (bypasses toolset filter)
+	// 4. Check if tool is in additionalTools (bypasses toolset filter)
 	if r.additionalTools != nil && r.additionalTools[tool.Tool.Name] {
 		return true
 	}
-	// 5. Check toolset filter
+	// 4. Check toolset filter
 	if !r.isToolsetEnabled(tool.Toolset.ID) {
 		return false
 	}
@@ -160,8 +176,11 @@ func (r *Inventory) AvailableResourceTemplates(ctx context.Context) []ServerReso
 	var result []ServerResourceTemplate
 	for i := range r.resourceTemplates {
 		res := &r.resourceTemplates[i]
-		// Check feature flags
-		if !r.isFeatureFlagAllowed(ctx, res.FeatureFlagEnable, res.FeatureFlagDisable) {
+		// Resources have no filter pipeline, so feature gating runs inline.
+		// The featureChecker != nil guard mirrors the structural "no checker
+		// = no filtering" contract used for tools (where the absence of a
+		// pipeline step expresses the same thing).
+		if r.featureChecker != nil && !featureFlagAllowed(ctx, r.featureChecker, res.FeatureFlagEnable, res.FeatureFlagDisable) {
 			continue
 		}
 		if r.isToolsetEnabled(res.Toolset.ID) {
@@ -189,8 +208,9 @@ func (r *Inventory) AvailablePrompts(ctx context.Context) []ServerPrompt {
 	var result []ServerPrompt
 	for i := range r.prompts {
 		prompt := &r.prompts[i]
-		// Check feature flags
-		if !r.isFeatureFlagAllowed(ctx, prompt.FeatureFlagEnable, prompt.FeatureFlagDisable) {
+		// Prompts have no filter pipeline; see AvailableResourceTemplates for
+		// the rationale behind the explicit nil guard.
+		if r.featureChecker != nil && !featureFlagAllowed(ctx, r.featureChecker, prompt.FeatureFlagEnable, prompt.FeatureFlagDisable) {
 			continue
 		}
 		if r.isToolsetEnabled(prompt.Toolset.ID) {

pkg/inventory/registry_test.go

@@ -1652,10 +1652,10 @@ func TestFilteredToolsMatchesAvailableTools(t *testing.T) {
 func TestFilteringOrder(t *testing.T) {
 	// Test that filters are applied in the correct order:
 	// 1. Tool.Enabled
-	// 2. Feature flags
-	// 3. Read-only
-	// 4. Builder filters
-	// 5. Toolset/additional tools
+	// 2. Read-only
+	// 3. Builder filters (feature-flag filter is at the head of this list
+	//    when WithFeatureChecker is set)
+	// 4. Toolset/additional tools
 
 	callOrder := []string{}
 
@@ -1688,8 +1688,9 @@ func TestFilteringOrder(t *testing.T) {
 
 	_ = reg.AvailableTools(context.Background())
 
-	// Expected order: Enabled, FeatureFlag, ReadOnly (stops here because it's write tool)
-	expectedOrder := []string{"Enabled", "FeatureFlag"}
+	// Expected order: Enabled, then Read-only stops (write tool, read-only mode);
+	// neither the feature-flag filter nor the user filter is reached.
+	expectedOrder := []string{"Enabled"}
 	if len(callOrder) != len(expectedOrder) {
 		t.Errorf("Expected %d checks, got %d: %v", len(expectedOrder), len(callOrder), callOrder)
 	}