github--github-mcp-server/pkg/inventory/builder.go

package inventory

import (
	"context"
	"errors"
	"fmt"
	"maps"
	"slices"
	"strings"
)

var (
	// ErrUnknownTools is returned when tools specified via WithTools() are not recognized.
	ErrUnknownTools = errors.New("unknown tools specified in WithTools")
)

// mcpAppsFeatureFlag is the feature flag name that controls MCP Apps UI metadata.
// This is defined here to avoid importing pkg/github (which imports pkg/inventory).
// The value must match github.MCPAppsFeatureFlag.
const mcpAppsFeatureFlag = "remote_mcp_ui_apps"

// ToolFilter is a function that determines if a tool should be included.
// Returns true if the tool should be included, false to exclude it.
type ToolFilter func(ctx context.Context, tool *ServerTool) (bool, error)

// Builder builds a Registry with the specified configuration.
// Use NewBuilder to create a builder, chain configuration methods,
// then call Build() to create the final inventory.
//
// Example:
//
//	reg := NewBuilder().
//	    SetTools(tools).
//	    SetResources(resources).
//	    SetPrompts(prompts).
//	    WithDeprecatedAliases(aliases).
//	    WithReadOnly(true).
//	    WithToolsets([]string{"repos", "issues"}).
//	    WithFeatureChecker(checker).
//	    WithFilter(myFilter).
//	    Build()
type Builder struct {
	tools             []ServerTool
	resourceTemplates []ServerResourceTemplate
	prompts           []ServerPrompt
	deprecatedAliases map[string]string

	// Configuration options (processed at Build time)
	readOnly             bool
	toolsetIDs           []string // raw input, processed at Build()
	toolsetIDsIsNil      bool     // tracks if nil was passed (nil = defaults)
	additionalTools      []string // raw input, processed at Build()
	featureChecker       FeatureFlagChecker
	filters              []ToolFilter // filters to apply to all tools
	generateInstructions bool
}

// NewBuilder creates a new Builder.
func NewBuilder() *Builder {
	return &Builder{
		deprecatedAliases: make(map[string]string),
		toolsetIDsIsNil:   true, // default to nil (use defaults)
	}
}

// SetTools sets the tools for the inventory. Returns self for chaining.
func (b *Builder) SetTools(tools []ServerTool) *Builder {
	b.tools = tools
	return b
}

// SetResources sets the resource templates for the inventory. Returns self for chaining.
func (b *Builder) SetResources(resources []ServerResourceTemplate) *Builder {
	b.resourceTemplates = resources
	return b
}

// SetPrompts sets the prompts for the inventory. Returns self for chaining.
func (b *Builder) SetPrompts(prompts []ServerPrompt) *Builder {
	b.prompts = prompts
	return b
}

// WithDeprecatedAliases adds deprecated tool name aliases that map to canonical names.
// Returns self for chaining.
func (b *Builder) WithDeprecatedAliases(aliases map[string]string) *Builder {
	maps.Copy(b.deprecatedAliases, aliases)
	return b
}

// WithReadOnly sets whether only read-only tools should be available.
// When true, write tools are filtered out. Returns self for chaining.
func (b *Builder) WithReadOnly(readOnly bool) *Builder {
	b.readOnly = readOnly
	return b
}

func (b *Builder) WithServerInstructions() *Builder {
	b.generateInstructions = true
	return b
}

// WithToolsets specifies which toolsets should be enabled.
// Special keywords:
//   - "all": enables all toolsets
//   - "default": expands to toolsets marked with Default: true in their metadata
//
// Input strings are trimmed of whitespace and duplicates are removed.
// Pass nil to use default toolsets. Pass an empty slice to disable all toolsets.
// Returns self for chaining.
func (b *Builder) WithToolsets(toolsetIDs []string) *Builder {
	b.toolsetIDs = toolsetIDs
	b.toolsetIDsIsNil = toolsetIDs == nil
	return b
}

// WithTools specifies additional tools that bypass toolset filtering.
// These tools are additive - they will be included even if their toolset is not enabled.
// Read-only filtering still applies to these tools.
// Input is cleaned (trimmed, deduplicated) during Build().
// Deprecated tool aliases are automatically resolved to their canonical names during Build().
// Returns self for chaining.
func (b *Builder) WithTools(toolNames []string) *Builder {
	b.additionalTools = toolNames
	return b
}

// WithFeatureChecker sets the feature flag checker function.
// The checker receives a context (for actor extraction) and feature flag name,
// 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
	return b
}

// WithFilter adds a filter function that will be applied to all tools.
// Multiple filters can be added and are evaluated in order.
// If any filter returns false or an error, the tool is excluded.
// Returns self for chaining.
func (b *Builder) WithFilter(filter ToolFilter) *Builder {
	b.filters = append(b.filters, filter)
	return b
}

// WithExcludeTools specifies tools that should be disabled regardless of other settings.
// These tools will be excluded even if their toolset is enabled or they are in the
// additional tools list. This takes precedence over all other tool enablement settings.
// Input is cleaned (trimmed, deduplicated) before applying.
// Returns self for chaining.
func (b *Builder) WithExcludeTools(toolNames []string) *Builder {
	cleaned := cleanTools(toolNames)
	if len(cleaned) > 0 {
		b.filters = append(b.filters, CreateExcludeToolsFilter(cleaned))
	}
	return b
}

// CreateExcludeToolsFilter creates a ToolFilter that excludes tools by name.
// Any tool whose name appears in the excluded list will be filtered out.
// The input slice should already be cleaned (trimmed, deduplicated).
func CreateExcludeToolsFilter(excluded []string) ToolFilter {
	set := make(map[string]struct{}, len(excluded))
	for _, name := range excluded {
		set[name] = struct{}{}
	}
	return func(_ context.Context, tool *ServerTool) (bool, error) {
		_, blocked := set[tool.Tool.Name]
		return !blocked, nil
	}
}

// cleanTools trims whitespace and removes duplicates from tool names.
// Empty strings after trimming are excluded.
func cleanTools(tools []string) []string {
	seen := make(map[string]bool)
	var cleaned []string
	for _, name := range tools {
		trimmed := strings.TrimSpace(name)
		if trimmed == "" {
			continue
		}
		if !seen[trimmed] {
			seen[trimmed] = true
			cleaned = append(cleaned, trimmed)
		}
	}
	return cleaned
}

// Build creates the final Inventory with all configuration applied.
// This processes toolset filtering, tool name resolution, and sets up
// the inventory for use. The returned Inventory is ready for use with
// AvailableTools(), RegisterAll(), etc.
//
// Build returns an error if any tools specified via WithTools() are not recognized
// (i.e., they don't exist in the tool set and are not deprecated aliases).
// This ensures invalid tool configurations fail fast at build time.
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:           filters,
	}

	// Process toolsets and pre-compute metadata in a single pass
	r.enabledToolsets, r.unrecognizedToolsets, r.toolsetIDs, r.toolsetIDSet, r.defaultToolsetIDs, r.toolsetDescriptions = b.processToolsets()

	// Build set of valid tool names for validation
	validToolNames := make(map[string]bool, len(tools))
	for i := range tools {
		validToolNames[tools[i].Tool.Name] = true
	}

	// Process additional tools (clean, resolve aliases, and track unrecognized)
	if len(b.additionalTools) > 0 {
		cleanedTools := cleanTools(b.additionalTools)

		r.additionalTools = make(map[string]bool, len(cleanedTools))
		var unrecognizedTools []string
		for _, name := range cleanedTools {
			// Always include the original name - this handles the case where
			// the tool exists but is controlled by a feature flag that's OFF.
			r.additionalTools[name] = true
			// Also include the canonical name if this is a deprecated alias.
			// This handles the case where the feature flag is ON and only
			// the new consolidated tool is available.
			if canonical, isAlias := b.deprecatedAliases[name]; isAlias {
				r.additionalTools[canonical] = true
			} else if !validToolNames[name] {
				// Not a valid tool and not a deprecated alias - track as unrecognized
				unrecognizedTools = append(unrecognizedTools, name)
			}
		}

		// Error out if there are unrecognized tools
		if len(unrecognizedTools) > 0 {
			return nil, fmt.Errorf("%w: %s", ErrUnknownTools, strings.Join(unrecognizedTools, ", "))
		}
	}

	if b.generateInstructions {
		r.instructions = generateInstructions(r)
	}

	return r, nil
}

// processToolsets processes the toolsetIDs configuration and returns:
// - enabledToolsets map (nil means all enabled)
// - unrecognizedToolsets list for warnings
// - allToolsetIDs sorted list of all toolset IDs
// - toolsetIDSet map for O(1) HasToolset lookup
// - defaultToolsetIDs sorted list of default toolset IDs
// - toolsetDescriptions map of toolset ID to description
func (b *Builder) processToolsets() (map[ToolsetID]bool, []string, []ToolsetID, map[ToolsetID]bool, []ToolsetID, map[ToolsetID]string) {
	// Single pass: collect all toolset metadata together
	validIDs := make(map[ToolsetID]bool)
	defaultIDs := make(map[ToolsetID]bool)
	descriptions := make(map[ToolsetID]string)

	for i := range b.tools {
		t := &b.tools[i]
		validIDs[t.Toolset.ID] = true
		if t.Toolset.Default {
			defaultIDs[t.Toolset.ID] = true
		}
		if t.Toolset.Description != "" {
			descriptions[t.Toolset.ID] = t.Toolset.Description
		}
	}
	for i := range b.resourceTemplates {
		r := &b.resourceTemplates[i]
		validIDs[r.Toolset.ID] = true
		if r.Toolset.Default {
			defaultIDs[r.Toolset.ID] = true
		}
		if r.Toolset.Description != "" {
			descriptions[r.Toolset.ID] = r.Toolset.Description
		}
	}
	for i := range b.prompts {
		p := &b.prompts[i]
		validIDs[p.Toolset.ID] = true
		if p.Toolset.Default {
			defaultIDs[p.Toolset.ID] = true
		}
		if p.Toolset.Description != "" {
			descriptions[p.Toolset.ID] = p.Toolset.Description
		}
	}

	// Build sorted slices from the collected maps
	allToolsetIDs := make([]ToolsetID, 0, len(validIDs))
	for id := range validIDs {
		allToolsetIDs = append(allToolsetIDs, id)
	}
	slices.Sort(allToolsetIDs)

	defaultToolsetIDList := make([]ToolsetID, 0, len(defaultIDs))
	for id := range defaultIDs {
		defaultToolsetIDList = append(defaultToolsetIDList, id)
	}
	slices.Sort(defaultToolsetIDList)

	toolsetIDs := b.toolsetIDs

	// Check for "all" keyword - enables all toolsets
	for _, id := range toolsetIDs {
		if strings.TrimSpace(id) == "all" {
			return nil, nil, allToolsetIDs, validIDs, defaultToolsetIDList, descriptions // nil means all enabled
		}
	}

	// nil means use defaults, empty slice means no toolsets
	if b.toolsetIDsIsNil {
		toolsetIDs = []string{"default"}
	}

	// Expand "default" keyword, trim whitespace, collect other IDs, and track unrecognized
	seen := make(map[ToolsetID]bool)
	expanded := make([]ToolsetID, 0, len(toolsetIDs))
	var unrecognized []string

	for _, id := range toolsetIDs {
		trimmed := strings.TrimSpace(id)
		if trimmed == "" {
			continue
		}
		if trimmed == "default" {
			for _, defaultID := range defaultToolsetIDList {
				if !seen[defaultID] {
					seen[defaultID] = true
					expanded = append(expanded, defaultID)
				}
			}
		} else {
			tsID := ToolsetID(trimmed)
			if !seen[tsID] {
				seen[tsID] = true
				expanded = append(expanded, tsID)
				// Track if this toolset doesn't exist
				if !validIDs[tsID] {
					unrecognized = append(unrecognized, trimmed)
				}
			}
		}
	}

	if len(expanded) == 0 {
		return make(map[ToolsetID]bool), unrecognized, allToolsetIDs, validIDs, defaultToolsetIDList, descriptions
	}

	enabledToolsets := make(map[ToolsetID]bool, len(expanded))
	for _, id := range expanded {
		enabledToolsets[id] = true
	}
	return enabledToolsets, unrecognized, allToolsetIDs, validIDs, defaultToolsetIDList, descriptions
}

// mcpAppsMetaKeys lists the Meta keys controlled by the remote_mcp_ui_apps feature flag.
var mcpAppsMetaKeys = []string{
	"ui", // MCP Apps UI metadata
}

// stripMCPAppsMetadata removes MCP Apps UI metadata from tools when the
// remote_mcp_ui_apps feature flag is not enabled.
func stripMCPAppsMetadata(tools []ServerTool) []ServerTool {
	result := make([]ServerTool, 0, len(tools))
	for _, tool := range tools {
		if stripped := stripMetaKeys(tool, mcpAppsMetaKeys); stripped != nil {
			result = append(result, *stripped)
		} else {
			result = append(result, tool)
		}
	}
	return result
}

// stripMetaKeys removes the specified Meta keys from a single tool.
// Returns a modified copy if changes were made, nil otherwise.
func stripMetaKeys(tool ServerTool, keys []string) *ServerTool {
	if tool.Tool.Meta == nil || len(keys) == 0 {
		return nil
	}

	// Check if any of the specified keys exist
	hasKeys := false
	for _, key := range keys {
		if _, ok := tool.Tool.Meta[key]; ok {
			hasKeys = true
			break
		}
	}
	if !hasKeys {
		return nil
	}

	// Make a shallow copy and remove specified keys
	toolCopy := tool
	newMeta := make(map[string]any, len(tool.Tool.Meta))
	for k, v := range tool.Tool.Meta {
		if !slices.Contains(keys, k) {
			newMeta[k] = v
		}
	}

	if len(newMeta) == 0 {
		toolCopy.Tool.Meta = nil
	} else {
		toolCopy.Tool.Meta = newMeta
	}
	return &toolCopy
}