diff --git a/cmd/github-mcp-server/generate_docs.go b/cmd/github-mcp-server/generate_docs.go index 791a40b886..a9a6c20f83 100644 --- a/cmd/github-mcp-server/generate_docs.go +++ b/cmd/github-mcp-server/generate_docs.go @@ -61,9 +61,8 @@ func generateReadmeDocs(readmePath string) error { // The README documents the default user experience: tools that are // enabled with no special flags set. Installing a checker that reports - // every flag as disabled excludes tools gated by FeatureFlagEnable and - // keeps the legacy variants of tools gated by FeatureFlagDisable, so - // flag-gated duplicates don't appear twice. + // every flag as disabled keeps the default variants selected by functional + // feature rules, so flag-gated duplicates don't appear twice. // Build() can only fail if WithTools specifies invalid tools - not used here r, _ := github.NewInventory(t). WithToolsets([]string{"all"}). diff --git a/docs/feature-flags.md b/docs/feature-flags.md index fe72955b08..33f1f9ea82 100644 --- a/docs/feature-flags.md +++ b/docs/feature-flags.md @@ -35,15 +35,54 @@ Only flags listed in [`AllowedFeatureFlags`](../pkg/github/feature_flags.go) can be enabled by end users. Insiders-only flags are not user-toggleable. +## Declaring tool availability + +Tools, resources, and prompts use `inventory.NewFeatureRule` when feature flags +change whether they are available. Each rule declares the flags it references +and evaluates them with a fail-closed `FeatureResolver`, so normal Go boolean +expressions can represent AND, OR, NOT, and mixed conditions: + +```go +tool.FeatureRule = inventory.NewFeatureRule( + []inventory.FeatureFlag{x, y}, + func(featureAsBool inventory.FeatureResolver) bool { + return !(featureAsBool(x) && featureAsBool(y)) + }, +) +``` + +Library consumers migrating existing inventory declarations should replace +`FeatureFlagEnable`, `FeatureFlagEnableAll`, and `FeatureFlagDisable` on +`ServerTool`, `ServerResourceTemplate`, and `ServerPrompt` with `FeatureRule`. +`FeatureFlagChecker` and `ToolDependencies.IsFeatureEnabled` continue to accept +string flag names. + +Rules are evaluated lazily after request narrowing and static availability +filters. Normal Go short-circuiting avoids checks that cannot affect the result, +while one request-owned memo ensures each flag actually reached is resolved at +most once across tools, resources, prompts, and `deps.IsFeatureEnabled`. +Unavailable named tool calls skip rule evaluation but remain registered so the +handler can return the specific client-availability error. + +Feature predicates are pure and may depend only on their resolver. Construction +validates every combination of up to 16 declared flags, so an undeclared lookup +fails immediately even when ordinary evaluation would short-circuit that +branch. + +The inventory's string-based checker owns request feature state. Once installed, +that state is authoritative; a checker stored on tool dependencies is used only +as a fallback when handlers are invoked directly without request state. +Feature checkers must not call `ResolveFeature`; nested resolution fails the +owning check closed. + --- ## Tools affected by each flag -The list below is regenerated from the Go source. For each user-controllable -feature flag, it lists every tool whose **inventory or input schema** differs -from the default — either because the flag introduces a new tool, or because -it selects a flag-aware variant of an existing tool. Flags that only affect -runtime behavior (such as output formatting) won't appear here. +The list below is regenerated by comparing the default tool surface with each +user-controllable flag enabled individually. Complex multi-flag rules may +require separate documentation. Flags that only affect runtime behavior (such +as output formatting) won't appear here. diff --git a/docs/insiders-features.md b/docs/insiders-features.md index 6191857ac8..4255e5bdd9 100644 --- a/docs/insiders-features.md +++ b/docs/insiders-features.md @@ -207,6 +207,11 @@ Insiders is a **meta feature flag** — the same shape as `default` or `all` for 3. **Insiders expansion.** If insiders mode is on (`--insiders`, `/insiders` route, or `X-MCP-Insiders: true`), every flag in [`InsidersFeatureFlags`](../pkg/github/feature_flags.go) is unioned in. The insiders expansion is **not** re-validated against the allowlist — insiders is a server-controlled switch that can reach internal-only flags. 4. **Server-side fallback (remote server only).** Any flag not yet decided falls back to the remote server's feature manager, which can roll a feature out independently of user input or insiders membership. +For tool availability, functional rules declare the flags they may read and are +evaluated lazily after request narrowing. Short-circuiting skips unnecessary +checks, and request-owned state memoizes each flag that is reached. The same +state backs `deps.IsFeatureEnabled`. + `AllowedFeatureFlags` and `InsidersFeatureFlags` are deliberately independent sets: - A flag in **`AllowedFeatureFlags` only** is a regular opt-in: users can turn it on, but insiders does not auto-enable it. Granular issues/PRs flags work this way. @@ -219,5 +224,6 @@ Insiders is a **meta feature flag** — the same shape as `default` or `all` for 2. Add it to `AllowedFeatureFlags` if end users should be able to opt in via `--features`, `X-MCP-Features`, or the `features` URL query parameter. 3. Add it to `InsidersFeatureFlags` if insiders mode should turn it on automatically. -4. Gate the behavior on the concrete flag (`deps.IsFeatureEnabled(ctx, FeatureFlagX)`), never on `cfg.InsidersMode`. There is a `TestGitHubPackageDoesNotReadInsidersMode` guard test that fails if `pkg/github` reads `InsidersMode` directly. -5. The MCP-diff CI workflow picks up new entries in `AllowedFeatureFlags` automatically — see `.github/workflows/mcp-diff.yml`. +4. For tool availability, attach an `inventory.NewFeatureRule` that declares every flag used by its predicate. For behavior inside a handler, use `deps.IsFeatureEnabled(ctx, FeatureFlagX)`. +5. Gate on concrete flags, never on `cfg.InsidersMode`. There is a `TestGitHubPackageDoesNotReadInsidersMode` guard test that fails if `pkg/github` reads `InsidersMode` directly. +6. The MCP-diff CI workflow picks up new entries in `AllowedFeatureFlags` automatically — see `.github/workflows/mcp-diff.yml`. diff --git a/internal/ghmcp/server.go b/internal/ghmcp/server.go index dadc05744b..f713a44026 100644 --- a/internal/ghmcp/server.go +++ b/internal/ghmcp/server.go @@ -256,7 +256,7 @@ type StdioServerConfig struct { EnabledTools []string // EnabledFeatures is a list of feature flags that are enabled - // Items with FeatureFlagEnable matching an entry in this list will be available + // Tool feature rules evaluate entries in this list. EnabledFeatures []string // ReadOnly indicates if we should only register read-only tools diff --git a/pkg/context/mcp_info.go b/pkg/context/mcp_info.go index af474b13a0..fc09e6b139 100644 --- a/pkg/context/mcp_info.go +++ b/pkg/context/mcp_info.go @@ -3,6 +3,8 @@ package context import ( "context" "encoding/json" + + "github.com/modelcontextprotocol/go-sdk/mcp" ) type mcpMethodInfoCtx string @@ -22,6 +24,10 @@ type MCPMethodInfo struct { ItemName string // RawArguments contains the unmaterialized tool arguments for tools/call requests. RawArguments json.RawMessage + // ProtocolVersion and ClientCapabilities describe the requesting MCP client + // when stateless HTTP parsing makes them available before registration. + ProtocolVersion string + ClientCapabilities *mcp.ClientCapabilities } // DecodeArguments materializes tool arguments when request middleware needs diff --git a/pkg/github/actions_test.go b/pkg/github/actions_test.go index 964bc95a6b..f390c31771 100644 --- a/pkg/github/actions_test.go +++ b/pkg/github/actions_test.go @@ -574,7 +574,7 @@ func Test_ActionsGetJobLogs(t *testing.T) { // Note: consolidated ActionsGetJobLogs has same tool name "get_job_logs" as the individual tool // but with different descriptions. We skip toolsnap validation here since the individual // tool's toolsnap already exists and is tested in Test_GetJobLogs. - // The consolidated tool has FeatureFlagEnable set, so only one will be active at a time. + // The functional feature rules ensure only one variant is active at a time. assert.Equal(t, "get_job_logs", toolDef.Tool.Name) assert.NotEmpty(t, toolDef.Tool.Description) inputSchema := toolDef.Tool.InputSchema.(*jsonschema.Schema) diff --git a/pkg/github/csv_output_test.go b/pkg/github/csv_output_test.go index 5cc6fe7e54..44dbff6280 100644 --- a/pkg/github/csv_output_test.go +++ b/pkg/github/csv_output_test.go @@ -29,20 +29,18 @@ func TestCSVOutputAppliedToDefaultListTools(t *testing.T) { require.Len(t, available, 2) listing := requireToolByName(t, available, "list_things") - assert.Empty(t, listing.FeatureFlagEnable) - assert.Empty(t, listing.FeatureFlagDisable) + assert.True(t, listing.FeatureRule.IsZero()) getting := requireToolByName(t, available, "get_thing") - assert.Empty(t, getting.FeatureFlagEnable) - assert.Empty(t, getting.FeatureFlagDisable) + assert.True(t, getting.FeatureRule.IsZero()) } } func TestCSVOutputAppliesToFlagGatedListTools(t *testing.T) { enabledOnly := testCSVOutputTool("list_things", `[{"number":1}]`) - enabledOnly.FeatureFlagEnable = FeatureFlagFileBlame + enabledOnly.FeatureRule = featureEnabledRule(FeatureFlagFileBlame) disabledOnly := testCSVOutputTool("list_legacy_things", `[{"number":2}]`) - disabledOnly.FeatureFlagDisable = []string{FeatureFlagFileBlame} + disabledOnly.FeatureRule = featureDisabledRule(FeatureFlagFileBlame) tools := withCSVOutput([]inventory.ServerTool{enabledOnly, disabledOnly}) require.Len(t, tools, 2) diff --git a/pkg/github/dependencies.go b/pkg/github/dependencies.go index 0de3e3b279..333b1d9c07 100644 --- a/pkg/github/dependencies.go +++ b/pkg/github/dependencies.go @@ -6,7 +6,6 @@ import ( "fmt" "log/slog" "net/http" - "os" ghcontext "github.com/github/github-mcp-server/pkg/context" "github.com/github/github-mcp-server/pkg/http/transport" @@ -95,7 +94,7 @@ type ToolDependencies interface { GetContentWindowSize() int // IsFeatureEnabled checks if a feature flag is enabled. - IsFeatureEnabled(ctx context.Context, flagName string) bool + IsFeatureEnabled(ctx context.Context, flag string) bool // Logger returns the structured logger, optionally enriched with // request-scoped data from ctx. Integrators provide their own slog.Handler @@ -204,22 +203,11 @@ func (d BaseDeps) Metrics(ctx context.Context) metrics.Metrics { // GetRequestStateSealer implements RequestStateSealerProvider. func (d BaseDeps) GetRequestStateSealer() RequestStateSealer { return d.StateSealer } -// IsFeatureEnabled checks if a feature flag is enabled. -// Returns false if the feature checker is nil, flag name is empty, or an error occurs. -// This allows tools to conditionally change behavior based on feature flags. -func (d BaseDeps) IsFeatureEnabled(ctx context.Context, flagName string) bool { - if d.featureChecker == nil || flagName == "" { - return false - } - - enabled, err := d.featureChecker(ctx, flagName) - if err != nil { - // Log error but don't fail the tool - treat as disabled - fmt.Fprintf(os.Stderr, "Feature flag check error for %q: %v\n", flagName, err) - return false - } - - return enabled +// IsFeatureEnabled checks if a feature flag is enabled. Request feature state +// is authoritative when present; the dependency checker is a fallback for +// direct handler invocation. Empty names and checker errors resolve false. +func (d BaseDeps) IsFeatureEnabled(ctx context.Context, flag string) bool { + return inventory.ResolveFeature(ctx, d.featureChecker, inventory.FeatureFlag(flag)) } // NewTool creates a ServerTool that retrieves ToolDependencies from context at call time. @@ -495,18 +483,9 @@ func (d *RequestDeps) Metrics(ctx context.Context) metrics.Metrics { return d.obsv.Metrics(ctx) } -// IsFeatureEnabled checks if a feature flag is enabled. -func (d *RequestDeps) IsFeatureEnabled(ctx context.Context, flagName string) bool { - if d.featureChecker == nil || flagName == "" { - return false - } - - enabled, err := d.featureChecker(ctx, flagName) - if err != nil { - // Log error but don't fail the tool - treat as disabled - fmt.Fprintf(os.Stderr, "Feature flag check error for %q: %v\n", flagName, err) - return false - } - - return enabled +// IsFeatureEnabled checks if a feature flag is enabled. Request feature state +// is authoritative when present; the dependency checker is a fallback for +// direct handler invocation. +func (d *RequestDeps) IsFeatureEnabled(ctx context.Context, flag string) bool { + return inventory.ResolveFeature(ctx, d.featureChecker, inventory.FeatureFlag(flag)) } diff --git a/pkg/github/feature_flags.go b/pkg/github/feature_flags.go index a388f30d6b..94f335ebe5 100644 --- a/pkg/github/feature_flags.go +++ b/pkg/github/feature_flags.go @@ -1,6 +1,10 @@ package github -import "slices" +import ( + "slices" + + "github.com/github/github-mcp-server/pkg/inventory" +) // MCPAppsFeatureFlag is the feature flag name for MCP Apps (interactive UI forms). const MCPAppsFeatureFlag = "remote_mcp_ui_apps" @@ -71,6 +75,33 @@ type FeatureFlags struct { LockdownMode bool } +func featureEnabledRule(feature string) inventory.FeatureRule { + flag := inventory.FeatureFlag(feature) + return inventory.NewFeatureRule( + []inventory.FeatureFlag{flag}, + func(featureAsBool inventory.FeatureResolver) bool { + return featureAsBool(flag) + }, + ) +} + +func featureDisabledRule(feature string) inventory.FeatureRule { + flag := inventory.FeatureFlag(feature) + return inventory.NewFeatureRule( + []inventory.FeatureFlag{flag}, + func(featureAsBool inventory.FeatureResolver) bool { + return !featureAsBool(flag) + }, + ) +} + +var ( + issuesGranularFeatureRule = featureEnabledRule(FeatureFlagIssuesGranular) + issuesConsolidatedFeatureRule = featureDisabledRule(FeatureFlagIssuesGranular) + pullRequestsGranularFeatureRule = featureEnabledRule(FeatureFlagPullRequestsGranular) + pullRequestsConsolidatedRule = featureDisabledRule(FeatureFlagPullRequestsGranular) +) + // ResolveFeatureFlags computes the effective set of enabled feature flags by: // 1. Taking the user-supplied flags (from --features or HTTP request // configuration) and @@ -89,9 +120,9 @@ type FeatureFlags struct { // Returns a set (map) for O(1) lookup by the feature checker. func ResolveFeatureFlags(enabledFeatures []string, insidersMode bool) map[string]bool { effective := make(map[string]bool) - for _, f := range enabledFeatures { - if slices.Contains(AllowedFeatureFlags, f) { - effective[f] = true + for _, feature := range enabledFeatures { + if slices.Contains(AllowedFeatureFlags, feature) { + effective[feature] = true } } if insidersMode { diff --git a/pkg/github/feature_flags_benchmark_test.go b/pkg/github/feature_flags_benchmark_test.go new file mode 100644 index 0000000000..79aa17d55d --- /dev/null +++ b/pkg/github/feature_flags_benchmark_test.go @@ -0,0 +1,188 @@ +package github + +import ( + "context" + "fmt" + "sync/atomic" + "testing" + + "github.com/github/github-mcp-server/pkg/inventory" + "github.com/github/github-mcp-server/pkg/translations" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +func BenchmarkFeatureInventory(b *testing.B) { + for _, distribution := range featureBenchmarkDistributions() { + b.Run(distribution.name, func(b *testing.B) { + b.Run("build", func(b *testing.B) { + var calls atomic.Int64 + b.ReportAllocs() + for b.Loop() { + _, err := featureBenchmarkBuilder(distribution, &calls).Build() + if err != nil { + b.Fatal(err) + } + } + b.ReportMetric(float64(calls.Load())/float64(b.N), "checks/op") + }) + + builder := featureBenchmarkBuilder(distribution, nil) + b.Run("preconstructed-builder", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + if _, err := builder.Build(); err != nil { + b.Fatal(err) + } + } + }) + + b.Run("tools-list", func(b *testing.B) { + inv, calls := featureBenchmarkInventory(b, distribution) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = inv.ForMCPRequest(inventory.MCPMethodToolsList, "").ToolsForRegistration(context.Background()) + } + b.ReportMetric(float64(calls.Load())/float64(b.N), "checks/op") + }) + + b.Run("read-only-tools-list", func(b *testing.B) { + var calls atomic.Int64 + checker := func(_ context.Context, flag string) (bool, error) { + calls.Add(1) + return distribution.enabled["*"] || distribution.enabled[flag], nil + } + tools := AllTools(translations.NullTranslationHelper) + inv, err := inventory.NewBuilder(). + SetTools(tools). + SetResources(AllResources(translations.NullTranslationHelper)). + SetPrompts(AllPrompts(translations.NullTranslationHelper)). + WithToolsets([]string{"all"}). + WithReadOnly(true). + WithFeatureChecker(checker). + Build() + if err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = inv.ForMCPRequest(inventory.MCPMethodToolsList, "").ToolsForRegistration(context.Background()) + } + b.ReportMetric(float64(calls.Load())/float64(b.N), "checks/op") + }) + + b.Run("unflagged-tool-call", func(b *testing.B) { + inv, calls := featureBenchmarkInventory(b, distribution) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = inv.ForMCPRequest(inventory.MCPMethodToolsCall, "get_commit").ToolsForRegistration(context.Background()) + } + b.ReportMetric(float64(calls.Load())/float64(b.N), "checks/op") + }) + + b.Run("gated-tool-call", func(b *testing.B) { + inv, calls := featureBenchmarkInventory(b, distribution) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = inv.ForMCPRequest(inventory.MCPMethodToolsCall, "get_file_blame").ToolsForRegistration(context.Background()) + } + b.ReportMetric(float64(calls.Load())/float64(b.N), "checks/op") + }) + + b.Run("ui-tool-call", func(b *testing.B) { + inv, calls := featureBenchmarkInventory(b, distribution) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = inv.ForMCPRequest(inventory.MCPMethodToolsCall, "ui_get").ToolsForRegistration(context.Background()) + } + b.ReportMetric(float64(calls.Load())/float64(b.N), "checks/op") + }) + + b.Run("direct-handler-checks", func(b *testing.B) { + var calls atomic.Int64 + checker := func(_ context.Context, flag string) (bool, error) { + calls.Add(1) + return distribution.enabled["*"] || distribution.enabled[flag], nil + } + b.ReportAllocs() + for b.Loop() { + ctx := inventory.WithFeatureState(context.Background(), checker) + _ = inventory.ResolveFeature(ctx, checker, inventory.FeatureFlag(FeatureFlagCSVOutput)) + _ = inventory.ResolveFeature(ctx, checker, inventory.FeatureFlag(FeatureFlagCSVOutput)) + } + b.ReportMetric(float64(calls.Load())/float64(b.N), "checks/op") + }) + + b.Run("build-list-register", func(b *testing.B) { + var calls atomic.Int64 + b.ReportAllocs() + for b.Loop() { + inv, err := featureBenchmarkBuilder(distribution, &calls).Build() + if err != nil { + b.Fatal(err) + } + inv = inv.ForMCPRequest(inventory.MCPMethodToolsList, "") + server := mcp.NewServer(&mcp.Implementation{Name: "benchmark", Version: "v0"}, nil) + inv.RegisterAll(context.Background(), server, nil) + } + b.ReportMetric(float64(calls.Load())/float64(b.N), "checks/op") + }) + }) + } +} + +type featureBenchmarkDistribution struct { + name string + enabled map[string]bool +} + +func featureBenchmarkDistributions() []featureBenchmarkDistribution { + return []featureBenchmarkDistribution{ + {name: "all-false"}, + { + name: "mixed", + enabled: map[string]bool{ + MCPAppsFeatureFlag: true, + FeatureFlagFileBlame: true, + FeatureFlagIssuesGranular: true, + FeatureFlagIssueDependencies: true, + }, + }, + {name: "all-true", enabled: map[string]bool{"*": true}}, + } +} + +func featureBenchmarkBuilder(distribution featureBenchmarkDistribution, calls *atomic.Int64) *inventory.Builder { + checker := func(_ context.Context, flag string) (bool, error) { + if calls != nil { + calls.Add(1) + } + return distribution.enabled["*"] || distribution.enabled[flag], nil + } + tools := AllTools(translations.NullTranslationHelper) + for i, baseCount := 0, len(tools); len(tools) < 139; i++ { + tool := tools[i%baseCount] + tool.Tool.Name = fmt.Sprintf("%s_remote_%d", tool.Tool.Name, i) + tools = append(tools, tool) + } + return inventory.NewBuilder(). + SetTools(tools). + SetResources(AllResources(translations.NullTranslationHelper)). + SetPrompts(AllPrompts(translations.NullTranslationHelper)). + WithToolsets([]string{"all"}). + WithFeatureChecker(checker) +} + +func featureBenchmarkInventory(b *testing.B, distribution featureBenchmarkDistribution) (*inventory.Inventory, *atomic.Int64) { + b.Helper() + var calls atomic.Int64 + inv, err := featureBenchmarkBuilder(distribution, &calls).Build() + if err != nil { + b.Fatal(err) + } + return inv, &calls +} diff --git a/pkg/github/feature_flags_test.go b/pkg/github/feature_flags_test.go index dafacfa79a..cc3fbf0837 100644 --- a/pkg/github/feature_flags_test.go +++ b/pkg/github/feature_flags_test.go @@ -17,12 +17,12 @@ import ( ) // RemoteMCPEnthusiasticGreeting is a dummy test feature flag . -const RemoteMCPEnthusiasticGreeting = "remote_mcp_enthusiastic_greeting" +const RemoteMCPEnthusiasticGreeting inventory.FeatureFlag = "remote_mcp_enthusiastic_greeting" -func featureCheckerFor(enabledFlags ...string) func(context.Context, string) (bool, error) { +func featureCheckerFor(enabledFlags ...inventory.FeatureFlag) inventory.FeatureFlagChecker { enabled := make(map[string]bool, len(enabledFlags)) for _, flag := range enabledFlags { - enabled[flag] = true + enabled[string(flag)] = true } return func(_ context.Context, flagName string) (bool, error) { return enabled[flagName], nil @@ -47,7 +47,7 @@ func HelloWorldTool(t translations.TranslationHelperFunc) inventory.ServerTool { // Check feature flag to determine greeting style greeting := "Hello, world!" - if deps.IsFeatureEnabled(ctx, RemoteMCPEnthusiasticGreeting) { + if deps.IsFeatureEnabled(ctx, string(RemoteMCPEnthusiasticGreeting)) { greeting += " Welcome to the future of MCP! 🎉" } @@ -91,7 +91,7 @@ func TestHelloWorld_ConditionalBehavior_Featureflag(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - var enabledFlags []string + var enabledFlags []inventory.FeatureFlag if tt.featureFlagEnabled { enabledFlags = append(enabledFlags, RemoteMCPEnthusiasticGreeting) } @@ -237,7 +237,7 @@ func TestResolveFeatureFlags(t *testing.T) { func TestThreadResolutionReasonToolVariants(t *testing.T) { tests := []struct { name string - flags []string + flags []inventory.FeatureFlag host utils.HostType toolName string hasReason bool @@ -248,30 +248,30 @@ func TestThreadResolutionReasonToolVariants(t *testing.T) { }, { name: "consolidated flag on", - flags: []string{FeatureFlagThreadResolutionReason}, + flags: []inventory.FeatureFlag{FeatureFlagThreadResolutionReason}, toolName: "pull_request_review_write", hasReason: true, }, { name: "granular flag off", - flags: []string{FeatureFlagPullRequestsGranular}, + flags: []inventory.FeatureFlag{inventory.FeatureFlag(FeatureFlagPullRequestsGranular)}, toolName: "resolve_review_thread", }, { name: "granular flag on", - flags: []string{FeatureFlagPullRequestsGranular, FeatureFlagThreadResolutionReason}, + flags: []inventory.FeatureFlag{inventory.FeatureFlag(FeatureFlagPullRequestsGranular), FeatureFlagThreadResolutionReason}, toolName: "resolve_review_thread", hasReason: true, }, { name: "consolidated flag on GHES", - flags: []string{FeatureFlagThreadResolutionReason}, + flags: []inventory.FeatureFlag{FeatureFlagThreadResolutionReason}, host: utils.HostTypeGHES, toolName: "pull_request_review_write", }, { name: "granular flag on GHES", - flags: []string{FeatureFlagPullRequestsGranular, FeatureFlagThreadResolutionReason}, + flags: []inventory.FeatureFlag{inventory.FeatureFlag(FeatureFlagPullRequestsGranular), FeatureFlagThreadResolutionReason}, host: utils.HostTypeGHES, toolName: "resolve_review_thread", }, diff --git a/pkg/github/find_duplicate.go b/pkg/github/find_duplicate.go index 65cb6f2002..236f53a489 100644 --- a/pkg/github/find_duplicate.go +++ b/pkg/github/find_duplicate.go @@ -178,6 +178,6 @@ func FindDuplicate(t translations.TranslationHelperFunc) inventory.ServerTool { result = attachRepoVisibilityIFCLabel(ctx, deps, client, owner, repo, result, ifc.LabelRepoUserContent) return result, nil, nil }) - st.FeatureFlagEnable = FeatureFlagDuplicateDetection + st.FeatureRule = featureEnabledRule(FeatureFlagDuplicateDetection) return st } diff --git a/pkg/github/find_duplicate_test.go b/pkg/github/find_duplicate_test.go index 9e20d958b4..b2c032989b 100644 --- a/pkg/github/find_duplicate_test.go +++ b/pkg/github/find_duplicate_test.go @@ -8,6 +8,7 @@ import ( "testing" "github.com/github/github-mcp-server/internal/toolsnaps" + "github.com/github/github-mcp-server/pkg/inventory" "github.com/github/github-mcp-server/pkg/translations" "github.com/google/jsonschema-go/jsonschema" "github.com/stretchr/testify/assert" @@ -20,8 +21,8 @@ func Test_FindDuplicate(t *testing.T) { // Verify tool definition once (flag-gated variant snap). serverTool := FindDuplicate(translations.NullTranslationHelper) tool := serverTool.Tool - require.NoError(t, toolsnaps.Test(tool.Name+"_ff_"+FeatureFlagDuplicateDetection, tool)) - require.Equal(t, FeatureFlagDuplicateDetection, serverTool.FeatureFlagEnable) + require.NoError(t, toolsnaps.Test(tool.Name+"_ff_"+string(FeatureFlagDuplicateDetection), tool)) + require.Equal(t, []inventory.FeatureFlag{FeatureFlagDuplicateDetection}, serverTool.FeatureRule.Features()) assert.Equal(t, "find_duplicate", tool.Name) assert.NotEmpty(t, tool.Description) diff --git a/pkg/github/granular_tools_test.go b/pkg/github/granular_tools_test.go index 5ef0c0a662..456129a0fe 100644 --- a/pkg/github/granular_tools_test.go +++ b/pkg/github/granular_tools_test.go @@ -21,9 +21,16 @@ import ( ) func granularToolsForToolset(toolsetID inventory.ToolsetID, featureFlag string) []inventory.ServerTool { + flag := inventory.FeatureFlag(featureFlag) var result []inventory.ServerTool for _, tool := range AllTools(translations.NullTranslationHelper) { - if tool.Toolset.ID == toolsetID && tool.FeatureFlagEnable == featureFlag && len(tool.FeatureFlagEnableAll) == 0 { + features := tool.FeatureRule.Features() + usesFeature := false + for _, feature := range features { + usesFeature = usesFeature || feature == flag + } + if tool.Toolset.ID == toolsetID && usesFeature && + tool.FeatureRule.Enabled(func(feature inventory.FeatureFlag) bool { return feature == flag }) { result = append(result, tool) } } @@ -102,7 +109,7 @@ func TestIssuesGranularToolset(t *testing.T) { t.Run("all granular tools have correct feature flag", func(t *testing.T) { for _, tool := range granularToolsForToolset(ToolsetMetadataIssues.ID, FeatureFlagIssuesGranular) { - assert.Equal(t, FeatureFlagIssuesGranular, tool.FeatureFlagEnable, "tool %s", tool.Tool.Name) + assert.Equal(t, []inventory.FeatureFlag{inventory.FeatureFlag(FeatureFlagIssuesGranular)}, tool.FeatureRule.Features(), "tool %s", tool.Tool.Name) } }) } @@ -138,7 +145,7 @@ func TestPullRequestsGranularToolset(t *testing.T) { t.Run("all granular tools have correct feature flag", func(t *testing.T) { for _, tool := range granularToolsForToolset(ToolsetMetadataPullRequests.ID, FeatureFlagPullRequestsGranular) { - assert.Equal(t, FeatureFlagPullRequestsGranular, tool.FeatureFlagEnable, "tool %s", tool.Tool.Name) + assert.Contains(t, tool.FeatureRule.Features(), inventory.FeatureFlag(FeatureFlagPullRequestsGranular), "tool %s", tool.Tool.Name) } }) } diff --git a/pkg/github/issue_dependencies.go b/pkg/github/issue_dependencies.go index 246192736e..dc78c9eccc 100644 --- a/pkg/github/issue_dependencies.go +++ b/pkg/github/issue_dependencies.go @@ -103,7 +103,7 @@ Options are: return utils.NewToolResultError(fmt.Sprintf("unknown method: %s", method)), nil, nil } }) - st.FeatureFlagEnable = FeatureFlagIssueDependencies + st.FeatureRule = featureEnabledRule(FeatureFlagIssueDependencies) return st } @@ -321,7 +321,7 @@ Options are: result, err := writeIssueDependency(ctx, client, method, blocked, blocking) return result, nil, err }) - st.FeatureFlagEnable = FeatureFlagIssueDependencies + st.FeatureRule = featureEnabledRule(FeatureFlagIssueDependencies) return st } diff --git a/pkg/github/issue_dependencies_test.go b/pkg/github/issue_dependencies_test.go index 6af9c504ed..35baa815ee 100644 --- a/pkg/github/issue_dependencies_test.go +++ b/pkg/github/issue_dependencies_test.go @@ -8,6 +8,7 @@ import ( "testing" "github.com/github/github-mcp-server/internal/toolsnaps" + "github.com/github/github-mcp-server/pkg/inventory" "github.com/github/github-mcp-server/pkg/translations" "github.com/google/jsonschema-go/jsonschema" "github.com/stretchr/testify/assert" @@ -34,8 +35,8 @@ func Test_IssueDependencyRead(t *testing.T) { // Verify tool definition once (flag-gated variant snap) serverTool := IssueDependencyRead(translations.NullTranslationHelper) tool := serverTool.Tool - require.NoError(t, toolsnaps.Test(tool.Name+"_ff_"+FeatureFlagIssueDependencies, tool)) - require.Equal(t, FeatureFlagIssueDependencies, serverTool.FeatureFlagEnable) + require.NoError(t, toolsnaps.Test(tool.Name+"_ff_"+string(FeatureFlagIssueDependencies), tool)) + require.Equal(t, []inventory.FeatureFlag{FeatureFlagIssueDependencies}, serverTool.FeatureRule.Features()) assert.Equal(t, "issue_dependency_read", tool.Name) assert.NotEmpty(t, tool.Description) @@ -189,8 +190,8 @@ func Test_IssueDependencyWrite(t *testing.T) { // Verify tool definition once (flag-gated variant snap) serverTool := IssueDependencyWrite(translations.NullTranslationHelper) tool := serverTool.Tool - require.NoError(t, toolsnaps.Test(tool.Name+"_ff_"+FeatureFlagIssueDependencies, tool)) - require.Equal(t, FeatureFlagIssueDependencies, serverTool.FeatureFlagEnable) + require.NoError(t, toolsnaps.Test(tool.Name+"_ff_"+string(FeatureFlagIssueDependencies), tool)) + require.Equal(t, []inventory.FeatureFlag{FeatureFlagIssueDependencies}, serverTool.FeatureRule.Features()) assert.Equal(t, "issue_dependency_write", tool.Name) assert.NotEmpty(t, tool.Description) diff --git a/pkg/github/issues.go b/pkg/github/issues.go index fd7ea36873..9450476b19 100644 --- a/pkg/github/issues.go +++ b/pkg/github/issues.go @@ -1679,7 +1679,7 @@ func SubIssueWrite(t translations.TranslationHelperFunc) inventory.ServerTool { return utils.NewToolResultError(fmt.Sprintf("unknown method: %s", method)), nil, nil } }) - st.FeatureFlagDisable = []string{FeatureFlagIssuesGranular} + st.FeatureRule = issuesConsolidatedFeatureRule return st } @@ -2700,7 +2700,7 @@ Options are: return utils.NewToolResultError("invalid method, must be either 'create' or 'update'"), nil, nil } }) - st.FeatureFlagDisable = []string{FeatureFlagIssuesGranular} + st.FeatureRule = issuesConsolidatedFeatureRule return st } diff --git a/pkg/github/issues_granular.go b/pkg/github/issues_granular.go index 0b1cdd0e7c..c6544495c6 100644 --- a/pkg/github/issues_granular.go +++ b/pkg/github/issues_granular.go @@ -108,7 +108,7 @@ func issueUpdateTool( return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -234,7 +234,7 @@ func GranularCreateIssue(t translations.TranslationHelperFunc) inventory.ServerT return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -460,7 +460,7 @@ func GranularUpdateIssueAssignees(t translations.TranslationHelperFunc) inventor return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -678,7 +678,7 @@ func GranularUpdateIssueLabels(t translations.TranslationHelperFunc) inventory.S return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -866,7 +866,7 @@ func GranularUpdateIssueType(t translations.TranslationHelperFunc) inventory.Ser return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -1064,7 +1064,7 @@ func GranularUpdateIssueState(t translations.TranslationHelperFunc) inventory.Se return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -1138,7 +1138,7 @@ func GranularAddSubIssue(t translations.TranslationHelperFunc) inventory.ServerT return result, nil, err }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -1207,7 +1207,7 @@ func GranularRemoveSubIssue(t translations.TranslationHelperFunc) inventory.Serv return result, nil, err }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -1292,7 +1292,7 @@ func GranularReprioritizeSubIssue(t translations.TranslationHelperFunc) inventor return result, nil, err }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -1580,7 +1580,7 @@ func GranularSetIssueFields(t translations.TranslationHelperFunc) inventory.Serv return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -1662,7 +1662,7 @@ func GranularAddIssueReaction(t translations.TranslationHelperFunc) inventory.Se return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } @@ -1744,6 +1744,6 @@ func GranularAddIssueCommentReaction(t translations.TranslationHelperFunc) inven return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagIssuesGranular + st.FeatureRule = issuesGranularFeatureRule return st } diff --git a/pkg/github/issues_test.go b/pkg/github/issues_test.go index e8b4cd2c13..a894641ddf 100644 --- a/pkg/github/issues_test.go +++ b/pkg/github/issues_test.go @@ -1833,7 +1833,7 @@ func Test_CreateIssue(t *testing.T) { serverTool := IssueWrite(translations.NullTranslationHelper) tool := serverTool.Tool require.NoError(t, toolsnaps.Test(tool.Name, tool)) - require.Empty(t, serverTool.FeatureFlagEnable) + require.Equal(t, []inventory.FeatureFlag{inventory.FeatureFlag(FeatureFlagIssuesGranular)}, serverTool.FeatureRule.Features()) assert.Equal(t, "issue_write", tool.Name) assert.NotEmpty(t, tool.Description) diff --git a/pkg/github/pullrequests.go b/pkg/github/pullrequests.go index 925e7d55a0..5cee8b3231 100644 --- a/pkg/github/pullrequests.go +++ b/pkg/github/pullrequests.go @@ -1169,7 +1169,7 @@ func UpdatePullRequest(t translations.TranslationHelperFunc) inventory.ServerToo return utils.NewToolResultText(string(r)), nil, nil }) - st.FeatureFlagDisable = []string{FeatureFlagPullRequestsGranular} + st.FeatureRule = pullRequestsConsolidatedRule return st } @@ -1907,12 +1907,24 @@ Available methods: } }) if withResolutionReason { - st.FeatureFlagEnable = FeatureFlagThreadResolutionReason - st.FeatureFlagDisable = []string{FeatureFlagPullRequestsGranular} + st.FeatureRule = inventory.NewFeatureRule( + []inventory.FeatureFlag{FeatureFlagThreadResolutionReason, inventory.FeatureFlag(FeatureFlagPullRequestsGranular)}, + func(featureAsBool inventory.FeatureResolver) bool { + return featureAsBool(FeatureFlagThreadResolutionReason) && + !featureAsBool(inventory.FeatureFlag(FeatureFlagPullRequestsGranular)) + }, + ) } else { - st.FeatureFlagDisable = []string{FeatureFlagPullRequestsGranular} - if cfg.hostType != utils.HostTypeGHES { - st.FeatureFlagDisable = append(st.FeatureFlagDisable, FeatureFlagThreadResolutionReason) + if cfg.hostType == utils.HostTypeGHES { + st.FeatureRule = pullRequestsConsolidatedRule + } else { + st.FeatureRule = inventory.NewFeatureRule( + []inventory.FeatureFlag{FeatureFlagThreadResolutionReason, inventory.FeatureFlag(FeatureFlagPullRequestsGranular)}, + func(featureAsBool inventory.FeatureResolver) bool { + return !featureAsBool(FeatureFlagThreadResolutionReason) && + !featureAsBool(inventory.FeatureFlag(FeatureFlagPullRequestsGranular)) + }, + ) } } return st @@ -2459,7 +2471,7 @@ func AddCommentToPendingReview(t translations.TranslationHelperFunc) inventory.S }) return result, nil, err }) - st.FeatureFlagDisable = []string{FeatureFlagPullRequestsGranular} + st.FeatureRule = pullRequestsConsolidatedRule return st } diff --git a/pkg/github/pullrequests_granular.go b/pkg/github/pullrequests_granular.go index 3a4b9fc810..a32723b3c7 100644 --- a/pkg/github/pullrequests_granular.go +++ b/pkg/github/pullrequests_granular.go @@ -103,7 +103,7 @@ func prUpdateTool( return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagPullRequestsGranular + st.FeatureRule = pullRequestsGranularFeatureRule return st } @@ -272,7 +272,7 @@ func GranularUpdatePullRequestDraftState(t translations.TranslationHelperFunc) i return utils.NewToolResultText("pull request marked as ready for review"), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagPullRequestsGranular + st.FeatureRule = pullRequestsGranularFeatureRule return st } @@ -351,7 +351,7 @@ func GranularRequestPullRequestReviewers(t translations.TranslationHelperFunc) i return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagPullRequestsGranular + st.FeatureRule = pullRequestsGranularFeatureRule return st } @@ -436,7 +436,7 @@ func GranularCreatePullRequestReview(t translations.TranslationHelperFunc) inven return result, nil, err }, ) - st.FeatureFlagEnable = FeatureFlagPullRequestsGranular + st.FeatureRule = pullRequestsGranularFeatureRule return st } @@ -500,7 +500,7 @@ func GranularSubmitPendingPullRequestReview(t translations.TranslationHelperFunc return result, nil, err }, ) - st.FeatureFlagEnable = FeatureFlagPullRequestsGranular + st.FeatureRule = pullRequestsGranularFeatureRule return st } @@ -555,7 +555,7 @@ func GranularDeletePendingPullRequestReview(t translations.TranslationHelperFunc return result, nil, err }, ) - st.FeatureFlagEnable = FeatureFlagPullRequestsGranular + st.FeatureRule = pullRequestsGranularFeatureRule return st } @@ -666,7 +666,7 @@ func GranularAddPullRequestReviewComment(t translations.TranslationHelperFunc) i return result, nil, err }, ) - st.FeatureFlagEnable = FeatureFlagPullRequestsGranular + st.FeatureRule = pullRequestsGranularFeatureRule return st } @@ -748,11 +748,25 @@ func granularResolveReviewThread(t translations.TranslationHelperFunc, withResol return result, nil, err }, ) - st.FeatureFlagEnable = FeatureFlagPullRequestsGranular - if withResolutionReason { - st.FeatureFlagEnableAll = []string{FeatureFlagThreadResolutionReason} - } else if cfg.hostType != utils.HostTypeGHES { - st.FeatureFlagDisable = []string{FeatureFlagThreadResolutionReason} + switch { + case withResolutionReason: + st.FeatureRule = inventory.NewFeatureRule( + []inventory.FeatureFlag{inventory.FeatureFlag(FeatureFlagPullRequestsGranular), FeatureFlagThreadResolutionReason}, + func(featureAsBool inventory.FeatureResolver) bool { + return featureAsBool(inventory.FeatureFlag(FeatureFlagPullRequestsGranular)) && + featureAsBool(FeatureFlagThreadResolutionReason) + }, + ) + case cfg.hostType == utils.HostTypeGHES: + st.FeatureRule = pullRequestsGranularFeatureRule + default: + st.FeatureRule = inventory.NewFeatureRule( + []inventory.FeatureFlag{inventory.FeatureFlag(FeatureFlagPullRequestsGranular), FeatureFlagThreadResolutionReason}, + func(featureAsBool inventory.FeatureResolver) bool { + return featureAsBool(inventory.FeatureFlag(FeatureFlagPullRequestsGranular)) && + !featureAsBool(FeatureFlagThreadResolutionReason) + }, + ) } return st } @@ -797,7 +811,7 @@ func GranularUnresolveReviewThread(t translations.TranslationHelperFunc) invento return result, nil, err }, ) - st.FeatureFlagEnable = FeatureFlagPullRequestsGranular + st.FeatureRule = pullRequestsGranularFeatureRule return st } @@ -879,6 +893,6 @@ func GranularAddPullRequestReviewCommentReaction(t translations.TranslationHelpe return utils.NewToolResultText(string(r)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagPullRequestsGranular + st.FeatureRule = pullRequestsGranularFeatureRule return st } diff --git a/pkg/github/repositories.go b/pkg/github/repositories.go index 8dfa19b4a2..326bea51c6 100644 --- a/pkg/github/repositories.go +++ b/pkg/github/repositories.go @@ -3030,7 +3030,7 @@ func GetFileBlame(t translations.TranslationHelperFunc) inventory.ServerTool { return utils.NewToolResultText(string(payload)), nil, nil }, ) - st.FeatureFlagEnable = FeatureFlagFileBlame + st.FeatureRule = featureEnabledRule(FeatureFlagFileBlame) return st } diff --git a/pkg/github/repositories_test.go b/pkg/github/repositories_test.go index 71b04faa3e..d88c13e7a5 100644 --- a/pkg/github/repositories_test.go +++ b/pkg/github/repositories_test.go @@ -5971,7 +5971,7 @@ func Test_GetFileBlame(t *testing.T) { // get_file_blame is gated so it is not advertised unless the feature flag // (or insiders mode) opts it in. - assert.Equal(t, FeatureFlagFileBlame, serverTool.FeatureFlagEnable, "get_file_blame must be gated behind the file_blame feature flag") + assert.Equal(t, []inventory.FeatureFlag{FeatureFlagFileBlame}, serverTool.FeatureRule.Features()) schema, ok := tool.InputSchema.(*jsonschema.Schema) require.True(t, ok, "InputSchema should be *jsonschema.Schema") diff --git a/pkg/github/server.go b/pkg/github/server.go index b8f0197889..6e9b5b7566 100644 --- a/pkg/github/server.go +++ b/pkg/github/server.go @@ -35,7 +35,7 @@ type MCPServerConfig struct { EnabledTools []string // EnabledFeatures is a list of feature flags that are enabled - // Items with FeatureFlagEnable matching an entry in this list will be available + // Tool feature rules evaluate entries in this list. EnabledFeatures []string // ReadOnly indicates if we should only offer read-only tools @@ -113,6 +113,7 @@ func NewMCPServer(ctx context.Context, cfg *MCPServerConfig, deps ToolDependenci // Add middlewares. Order matters - for example, the error context middleware should be applied last so that it runs FIRST (closest to the handler) to ensure all errors are captured, // and any middleware that needs to read or modify the context should be before it. ghServer.AddReceivingMiddleware(middleware...) + ghServer.AddReceivingMiddleware(injectFeatureStateMiddleware(inv)) ghServer.AddReceivingMiddleware(InjectDepsMiddleware(deps)) ghServer.AddReceivingMiddleware(addGitHubAPIErrorToContext) @@ -138,6 +139,14 @@ func NewMCPServer(ctx context.Context, cfg *MCPServerConfig, deps ToolDependenci return ghServer, nil } +func injectFeatureStateMiddleware(inv *inventory.Inventory) mcp.Middleware { + return func(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { + return next(inv.WithFeatureState(ctx), method, req) + } + } +} + // ResolvedEnabledToolsets determines which toolsets should be enabled based on config. // Returns nil for "use defaults", empty slice for "none", or explicit list. func ResolvedEnabledToolsets(enabledToolsets []string, enabledTools []string) []string { diff --git a/pkg/github/server_test.go b/pkg/github/server_test.go index 07cb63c85f..bd8b9abd6b 100644 --- a/pkg/github/server_test.go +++ b/pkg/github/server_test.go @@ -62,10 +62,12 @@ func (s stubDeps) GetRawClient(ctx context.Context) (*raw.Client, error) { func (s stubDeps) GetRepoAccessCache(_ context.Context) (*lockdown.RepoAccessCache, error) { return s.repoAccessCache, nil } -func (s stubDeps) GetT() translations.TranslationHelperFunc { return s.t } -func (s stubDeps) GetFlags(_ context.Context) FeatureFlags { return s.flags } -func (s stubDeps) GetContentWindowSize() int { return s.contentWindowSize } -func (s stubDeps) IsFeatureEnabled(_ context.Context, _ string) bool { return false } +func (s stubDeps) GetT() translations.TranslationHelperFunc { return s.t } +func (s stubDeps) GetFlags(_ context.Context) FeatureFlags { return s.flags } +func (s stubDeps) GetContentWindowSize() int { return s.contentWindowSize } +func (s stubDeps) IsFeatureEnabled(_ context.Context, _ string) bool { + return false +} func (s stubDeps) Logger(_ context.Context) *slog.Logger { return s.obsv.Logger() } @@ -192,6 +194,27 @@ func TestNewMCPServer_CreatesSuccessfully(t *testing.T) { // is already tested in pkg/github/*_test.go. } +func TestFeatureStateMiddlewareCachesHandlerChecks(t *testing.T) { + var calls int + checker := func(_ context.Context, flag string) (bool, error) { + calls++ + return flag == "enabled", nil + } + inv, err := NewInventory(translations.NullTranslationHelper). + WithFeatureChecker(checker). + Build() + require.NoError(t, err) + + next := func(ctx context.Context, _ string, _ mcp.Request) (mcp.Result, error) { + assert.True(t, inventory.ResolveFeature(ctx, nil, "enabled")) + assert.True(t, inventory.ResolveFeature(ctx, nil, "enabled")) + return nil, nil + } + _, err = injectFeatureStateMiddleware(inv)(next)(context.Background(), "tools/call", nil) + require.NoError(t, err) + assert.Equal(t, 1, calls) +} + // advertisedServerCapabilities connects an in-memory client to the given server // and returns the capabilities the server advertised during initialization. func advertisedServerCapabilities(t *testing.T, server *mcp.Server) *mcp.ServerCapabilities { diff --git a/pkg/github/tools_validation_test.go b/pkg/github/tools_validation_test.go index cdc12348ae..3eacc8e1c2 100644 --- a/pkg/github/tools_validation_test.go +++ b/pkg/github/tools_validation_test.go @@ -166,34 +166,78 @@ func TestToolReadOnlyHintConsistency(t *testing.T) { } } -// TestNoDuplicateToolNames ensures all tools have unique names +// TestNoDuplicateToolNames ensures duplicate names cannot be enabled together. func TestNoDuplicateToolNames(t *testing.T) { tools := AllTools(stubTranslation) - seen := make(map[string]bool) - featureFlagged := make(map[string]bool) + toolsByName := make(map[string][]inventory.ServerTool) + for _, tool := range tools { + toolsByName[tool.Tool.Name] = append(toolsByName[tool.Tool.Name], tool) + } // get_label is intentionally in both issues and labels toolsets for conformance // with original behavior where it was registered in both - allowedDuplicates := map[string]bool{ - "get_label": true, + for name, variants := range toolsByName { + if name == "get_label" || len(variants) < 2 { + continue + } + assert.False(t, featureDeclarationsOverlap(variants), "tool variants for %q can be enabled together", name) } +} - // First pass: identify tools that have feature flags (mutually exclusive at runtime) - for _, tool := range tools { - if tool.FeatureFlagEnable != "" || len(tool.FeatureFlagDisable) > 0 { - featureFlagged[tool.Tool.Name] = true +func featureDeclarationsOverlap(variants []inventory.ServerTool) bool { + positions := make(map[inventory.FeatureFlag]uint) + for _, variant := range variants { + for _, feature := range variant.FeatureRule.Features() { + if _, ok := positions[feature]; !ok { + positions[feature] = uint(len(positions)) + } } } + if len(positions) > 16 { + return true + } - for _, tool := range tools { - name := tool.Tool.Name - // Allow duplicates for explicitly allowed tools and feature-flagged tools - if !allowedDuplicates[name] && !featureFlagged[name] { - assert.False(t, seen[name], - "Duplicate tool name found: %q", name) + for assignment := range 1 << len(positions) { + enabled := 0 + featureAsBool := func(feature inventory.FeatureFlag) bool { + return assignment&(1< 1 { + return true } - seen[name] = true } + return false +} + +func TestFeatureRulesOverlap(t *testing.T) { + flag := inventory.FeatureFlag("flag") + enabled := inventory.ServerTool{FeatureRule: inventory.NewFeatureRule([]inventory.FeatureFlag{flag}, func(featureAsBool inventory.FeatureResolver) bool { + return featureAsBool(flag) + })} + disabled := inventory.ServerTool{FeatureRule: inventory.NewFeatureRule([]inventory.FeatureFlag{flag}, func(featureAsBool inventory.FeatureResolver) bool { + return !featureAsBool(flag) + })} + otherFlag := inventory.FeatureFlag("other") + otherEnabled := inventory.ServerTool{FeatureRule: inventory.NewFeatureRule([]inventory.FeatureFlag{otherFlag}, func(featureAsBool inventory.FeatureResolver) bool { + return featureAsBool(otherFlag) + })} + ungated := inventory.ServerTool{} + + assert.True(t, featureDeclarationsOverlap([]inventory.ServerTool{enabled, enabled})) + assert.True(t, featureDeclarationsOverlap([]inventory.ServerTool{enabled, otherEnabled})) + assert.True(t, featureDeclarationsOverlap([]inventory.ServerTool{ungated, enabled})) + assert.False(t, featureDeclarationsOverlap([]inventory.ServerTool{enabled, disabled})) +} + +func TestMCPAppsFeatureFlagMatchesInventory(t *testing.T) { + inv, err := NewInventory(stubTranslation).Build() + require.NoError(t, err) + assert.Contains(t, inv.RequiredFeatures(), inventory.FeatureFlag(MCPAppsFeatureFlag)) } // TestNoDuplicateResourceNames ensures all resources have unique names diff --git a/pkg/github/ui_capability_test.go b/pkg/github/ui_capability_test.go index 1c49ee15be..812d6fde89 100644 --- a/pkg/github/ui_capability_test.go +++ b/pkg/github/ui_capability_test.go @@ -5,6 +5,7 @@ import ( "testing" ghcontext "github.com/github/github-mcp-server/pkg/context" + "github.com/github/github-mcp-server/pkg/inventory" "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -95,17 +96,17 @@ func Test_shouldDeferToForm_featureFlags(t *testing.T) { tests := []struct { name string - enabledFlags []string + enabledFlags []inventory.FeatureFlag want bool }{ { name: "MCP Apps enabled defers to form", - enabledFlags: []string{MCPAppsFeatureFlag}, + enabledFlags: []inventory.FeatureFlag{MCPAppsFeatureFlag}, want: true, }, { name: "form deferral disabled executes directly", - enabledFlags: []string{ + enabledFlags: []inventory.FeatureFlag{ MCPAppsFeatureFlag, MCPAppsDisableFormDeferralFeatureFlag, }, @@ -113,7 +114,7 @@ func Test_shouldDeferToForm_featureFlags(t *testing.T) { }, { name: "form deferral opt-out does not enable MCP Apps", - enabledFlags: []string{MCPAppsDisableFormDeferralFeatureFlag}, + enabledFlags: []inventory.FeatureFlag{MCPAppsDisableFormDeferralFeatureFlag}, want: false, }, { diff --git a/pkg/github/ui_tools.go b/pkg/github/ui_tools.go index 62bba06ef6..b8d6cd31a5 100644 --- a/pkg/github/ui_tools.go +++ b/pkg/github/ui_tools.go @@ -98,7 +98,7 @@ func UIGet(t translations.TranslationHelperFunc) inventory.ServerTool { return utils.NewToolResultError(fmt.Sprintf("unknown method: %s", method)), nil, nil } }) - st.FeatureFlagEnable = MCPAppsFeatureFlag + st.FeatureRule = featureEnabledRule(MCPAppsFeatureFlag) return st } diff --git a/pkg/github/ui_tools_test.go b/pkg/github/ui_tools_test.go index 4a4981875b..d400752bdf 100644 --- a/pkg/github/ui_tools_test.go +++ b/pkg/github/ui_tools_test.go @@ -12,6 +12,7 @@ import ( "github.com/github/github-mcp-server/internal/githubv4mock" "github.com/github/github-mcp-server/internal/toolsnaps" + "github.com/github/github-mcp-server/pkg/inventory" "github.com/github/github-mcp-server/pkg/translations" "github.com/google/go-github/v89/github" "github.com/google/jsonschema-go/jsonschema" @@ -105,7 +106,7 @@ func Test_UIGet(t *testing.T) { assert.Contains(t, tool.InputSchema.(*jsonschema.Schema).Properties, "repo") assert.ElementsMatch(t, tool.InputSchema.(*jsonschema.Schema).Required, []string{"method", "owner"}) assert.True(t, tool.Annotations.ReadOnlyHint, "ui_get should be read-only") - assert.Equal(t, MCPAppsFeatureFlag, serverTool.FeatureFlagEnable, "ui_get should be gated on the MCP Apps feature flag") + assert.Equal(t, []inventory.FeatureFlag{MCPAppsFeatureFlag}, serverTool.FeatureRule.Features()) // ui_get must be app-only so the host hides it from the agent's tool list // while keeping it callable by the views (MCP Apps 2026-01-26 spec). diff --git a/pkg/http/handler.go b/pkg/http/handler.go index e4a9d198ec..1aad3ed01d 100644 --- a/pkg/http/handler.go +++ b/pkg/http/handler.go @@ -23,6 +23,11 @@ import ( const subscriptionsListenMethod = "subscriptions/listen" +// InventoryFactoryFunc builds the inventory for one HTTP request. All context +// values required by its feature checker must be installed before this runs: +// feature availability is resolved immediately afterward, before MCP receiving +// middleware can run. Handler-only lazy checks still see receiving-middleware +// context. type InventoryFactoryFunc func(r *http.Request) (*inventory.Inventory, error) // GitHubMCPServerFactoryFunc is a function type for creating a new MCP Server instance. @@ -214,6 +219,10 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { if methodInfo, ok := ghcontext.MCPMethod(r.Context()); ok && methodInfo != nil { invToUse = inv.ForMCPRequest(methodInfo.Method, methodInfo.ItemName) } + // Tool registration must know availability before the MCP server exists. + // Remote consumers install user identity in HTTP middleware before the + // inventory factory, so their per-user checker has its full context here. + r = r.WithContext(invToUse.WithFeatureState(r.Context())) ghServer, err := h.githubMcpServerFactory(r, h.deps, invToUse, &github.MCPServerConfig{ Version: h.config.Version, diff --git a/pkg/http/handler_test.go b/pkg/http/handler_test.go index 406f845897..ea37ec2a6a 100644 --- a/pkg/http/handler_test.go +++ b/pkg/http/handler_test.go @@ -60,12 +60,19 @@ func (f allScopesFetcher) FetchTokenScopes(_ context.Context, _ string) ([]strin var _ scopes.FetcherInterface = allScopesFetcher{} -func mockToolWithFeatureFlag(name, toolsetID string, readOnly bool, enableFlag, disableFlag string) inventory.ServerTool { +func mockToolWithFeatureFlag(name, toolsetID string, readOnly bool, enableFlag, disableFlag inventory.FeatureFlag) inventory.ServerTool { tool := mockTool(name, toolsetID, readOnly) - tool.FeatureFlagEnable = enableFlag + features := make([]inventory.FeatureFlag, 0, 2) + if enableFlag != "" { + features = append(features, enableFlag) + } if disableFlag != "" { - tool.FeatureFlagDisable = []string{disableFlag} + features = append(features, disableFlag) } + tool.FeatureRule = inventory.NewFeatureRule(features, func(featureAsBool inventory.FeatureResolver) bool { + return (enableFlag == "" || featureAsBool(enableFlag)) && + (disableFlag == "" || !featureAsBool(disableFlag)) + }) return tool } @@ -763,7 +770,9 @@ func TestStaticInventoryPreservesPerRequestFeatureVariants(t *testing.T) { available := inv.AvailableTools(ctx) require.Len(t, available, 1) assert.Equal(t, "list_issues", available[0].Tool.Name) - assert.Equal(t, github.FeatureFlagCSVOutput, available[0].FeatureFlagEnable) + assert.True(t, available[0].FeatureRule.Enabled(func(flag inventory.FeatureFlag) bool { + return flag == github.FeatureFlagCSVOutput + })) } func TestStaticInventoryDisablesOnlyDeleteRepository(t *testing.T) { @@ -1037,6 +1046,74 @@ func TestCrossOriginProtection(t *testing.T) { } } +func TestFeatureResolutionUsesOuterHTTPContext(t *testing.T) { + type userContextKey struct{} + const ( + userValue = "remote-user" + featureFlag = inventory.FeatureFlag("remote-feature") + ) + + var checkerCalls int + tool := mockTool("feature_tool", "test", true) + tool.FeatureRule = inventory.NewFeatureRule( + []inventory.FeatureFlag{featureFlag}, + func(featureAsBool inventory.FeatureResolver) bool { + return featureAsBool(featureFlag) + }, + ) + inventoryFactory := func(_ *http.Request) (*inventory.Inventory, error) { + checker := func(ctx context.Context, flag string) (bool, error) { + checkerCalls++ + return flag == string(featureFlag) && ctx.Value(userContextKey{}) == userValue, nil + } + return inventory.NewBuilder(). + SetTools([]inventory.ServerTool{tool}). + WithToolsets([]string{"all"}). + WithFeatureChecker(checker). + Build() + } + + apiHost, err := utils.NewAPIHost("https://api.github.com") + require.NoError(t, err) + handler := NewHTTPMcpHandler( + context.Background(), + &ServerConfig{Version: "test"}, + nil, + translations.NullTranslationHelper, + slog.Default(), + apiHost, + WithInventoryFactory(inventoryFactory), + WithGitHubMCPServerFactory(func(r *http.Request, _ github.ToolDependencies, inv *inventory.Inventory, _ *github.MCPServerConfig) (*mcp.Server, error) { + assert.True(t, inventory.ResolveFeature(r.Context(), nil, featureFlag)) + require.Len(t, inv.AvailableTools(r.Context()), 1) + return mcp.NewServer(&mcp.Implementation{Name: "test", Version: "0.0.1"}, nil), nil + }), + WithScopeFetcher(allScopesFetcher{}), + ) + + router := chi.NewRouter() + router.Use(func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), userContextKey{}, userValue))) + }) + }) + handler.RegisterMiddleware(router) + handler.RegisterRoutes(router) + + body := `{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"_meta":{"io.modelcontextprotocol/protocolVersion":"2026-07-28","io.modelcontextprotocol/clientInfo":{"name":"test","version":"1.0.0"},"io.modelcontextprotocol/clientCapabilities":{}}}}` + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(headers.ContentTypeHeader, headers.ContentTypeJSON) + req.Header.Set(headers.AcceptHeader, strings.Join([]string{headers.ContentTypeJSON, headers.ContentTypeEventStream}, ", ")) + req.Header.Set("Mcp-Protocol-Version", "2026-07-28") + req.Header.Set("Mcp-Method", "tools/list") + req.Header.Set(headers.AuthorizationHeader, "ghs_test-token") + + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, req) + require.Equal(t, http.StatusOK, recorder.Code, "response body: %s", recorder.Body.String()) + assert.Equal(t, 1, checkerCalls) +} + func TestHTTPToolMinimumProtocolVersion(t *testing.T) { apiHost, err := utils.NewAPIHost("https://api.github.com") require.NoError(t, err) diff --git a/pkg/http/middleware/mcp_parse.go b/pkg/http/middleware/mcp_parse.go index 0b56902261..bcc44f01f5 100644 --- a/pkg/http/middleware/mcp_parse.go +++ b/pkg/http/middleware/mcp_parse.go @@ -7,6 +7,7 @@ import ( "net/http" ghcontext "github.com/github/github-mcp-server/pkg/context" + "github.com/modelcontextprotocol/go-sdk/mcp" ) // mcpJSONRPCRequest represents the structure of an MCP JSON-RPC request. @@ -21,7 +22,11 @@ type mcpJSONRPCRequest struct { // For prompts/get // Name is shared with tools/call // For resources/read - URI string `json:"uri,omitempty"` + URI string `json:"uri,omitempty"` + Meta struct { + ProtocolVersion string `json:"io.modelcontextprotocol/protocolVersion,omitempty"` + ClientCapabilities *mcp.ClientCapabilities `json:"io.modelcontextprotocol/clientCapabilities,omitempty"` + } `json:"_meta"` } `json:"params"` } @@ -101,7 +106,11 @@ func parseMCPMethodInfo(body []byte) (*ghcontext.MCPMethodInfo, error) { return nil, nil } - methodInfo := &ghcontext.MCPMethodInfo{Method: mcpReq.Method} + methodInfo := &ghcontext.MCPMethodInfo{ + Method: mcpReq.Method, + ProtocolVersion: mcpReq.Params.Meta.ProtocolVersion, + ClientCapabilities: mcpReq.Params.Meta.ClientCapabilities, + } switch mcpReq.Method { case "tools/call": methodInfo.ItemName = mcpReq.Params.Name diff --git a/pkg/http/middleware/mcp_parse_test.go b/pkg/http/middleware/mcp_parse_test.go index e067f7808a..8034dc29a6 100644 --- a/pkg/http/middleware/mcp_parse_test.go +++ b/pkg/http/middleware/mcp_parse_test.go @@ -9,22 +9,25 @@ import ( "testing" ghcontext "github.com/github/github-mcp-server/pkg/context" + "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestWithMCPParse(t *testing.T) { tests := []struct { - name string - method string - path string - body string - expectInfo bool - expectedMethod string - expectedItem string - expectedRaw string - expectedArgs map[string]any - expectArgsError bool + name string + method string + path string + body string + expectInfo bool + expectedMethod string + expectedItem string + expectedRaw string + expectedArgs map[string]any + expectedProtocol string + expectedForm bool + expectArgsError bool }{ { name: "health check path is skipped", @@ -76,6 +79,19 @@ func TestWithMCPParse(t *testing.T) { expectInfo: true, expectedMethod: "tools/list", }, + { + name: "tools/list parses client availability", + method: http.MethodPost, + path: "/mcp", + body: `{"jsonrpc":"2.0","method":"tools/list","params":{"_meta":{ + "io.modelcontextprotocol/protocolVersion":"2026-07-28", + "io.modelcontextprotocol/clientCapabilities":{"elicitation":{"form":{}}} + }}}`, + expectInfo: true, + expectedMethod: "tools/list", + expectedProtocol: "2026-07-28", + expectedForm: true, + }, { name: "tools/call parses name", method: http.MethodPost, @@ -158,6 +174,12 @@ func TestWithMCPParse(t *testing.T) { require.NotNil(t, capturedInfo) assert.Equal(t, tt.expectedMethod, capturedInfo.Method) assert.Equal(t, tt.expectedItem, capturedInfo.ItemName) + assert.Equal(t, tt.expectedProtocol, capturedInfo.ProtocolVersion) + if tt.expectedForm { + require.NotNil(t, capturedInfo.ClientCapabilities) + require.NotNil(t, capturedInfo.ClientCapabilities.Elicitation) + assert.Equal(t, &mcp.FormElicitationCapabilities{}, capturedInfo.ClientCapabilities.Elicitation.Form) + } if tt.expectedRaw != "" { assert.JSONEq(t, tt.expectedRaw, string(capturedInfo.RawArguments)) } diff --git a/pkg/inventory/builder.go b/pkg/inventory/builder.go index 9ecaca1f57..60bda764b6 100644 --- a/pkg/inventory/builder.go +++ b/pkg/inventory/builder.go @@ -17,7 +17,7 @@ var ( // 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" +const mcpAppsFeatureFlag FeatureFlag = "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. @@ -65,19 +65,19 @@ func NewBuilder() *Builder { // SetTools sets the tools for the inventory. Returns self for chaining. func (b *Builder) SetTools(tools []ServerTool) *Builder { - b.tools = tools + b.tools = slices.Clone(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 + b.resourceTemplates = slices.Clone(resources) return b } // SetPrompts sets the prompts for the inventory. Returns self for chaining. func (b *Builder) SetPrompts(prompts []ServerPrompt) *Builder { - b.prompts = prompts + b.prompts = slices.Clone(prompts) return b } @@ -125,15 +125,10 @@ func (b *Builder) WithTools(toolNames []string) *Builder { 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. +// WithFeatureChecker sets the feature flag checker function. Inventory items +// declare their feature dependencies and functional availability rules through +// FeatureRule. Checks are deduplicated into request-owned resolution state; +// errors are logged and treated as disabled. // // When the checker is nil, no feature-flag filter is installed; tools, // resources, and prompts pass through feature-flag gating unchanged. The @@ -212,15 +207,7 @@ 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, diff --git a/pkg/inventory/features.go b/pkg/inventory/features.go new file mode 100644 index 0000000000..25665fde82 --- /dev/null +++ b/pkg/inventory/features.go @@ -0,0 +1,269 @@ +package inventory + +import ( + "context" + "fmt" + "os" + "slices" + "sync" +) + +const maxFeatureRuleFlags = 16 + +// FeatureFlag identifies a feature consistently across inventory consumers. +type FeatureFlag string + +// FeatureFlagChecker resolves one feature flag for the current request. Checkers +// must not call ResolveFeature; nested resolution fails the owning check closed. +type FeatureFlagChecker func(ctx context.Context, flag string) (bool, error) + +// FeatureResolver returns the resolved value of a feature flag. +// Implementations absorb resolution errors and fail closed. +type FeatureResolver func(flag FeatureFlag) bool + +// FeaturePredicate determines whether an inventory item is available. Predicates +// must be pure: their result may depend only on calls to the supplied resolver. +type FeaturePredicate func(featureAsBool FeatureResolver) bool + +// FeatureRule declares the feature flags used by an availability predicate. +// The predicate resolves reached flags lazily with normal Go boolean semantics, +// while request state deduplicates repeated checks. +type FeatureRule struct { + features []FeatureFlag + predicate FeaturePredicate +} + +// NewFeatureRule creates an availability rule over the supplied feature flags. +func NewFeatureRule(features []FeatureFlag, predicate FeaturePredicate) FeatureRule { + declared := make([]FeatureFlag, 0, len(features)) + for _, feature := range features { + if feature == "" { + continue + } + if slices.Contains(declared, feature) { + continue + } + declared = append(declared, feature) + } + rule := FeatureRule{ + features: declared, + predicate: predicate, + } + if len(declared) > 0 && predicate == nil { + panic("feature rule declares flags without a predicate") + } + rule.validate() + return rule +} + +func (r FeatureRule) validate() { + if r.predicate == nil { + return + } + if len(r.features) > maxFeatureRuleFlags { + panic(fmt.Sprintf("feature rule declares %d flags; maximum is %d", len(r.features), maxFeatureRuleFlags)) + } + + for assignment := range 1 << len(r.features) { + r.evaluate(func(feature FeatureFlag) bool { + for i, declared := range r.features { + if feature == declared { + return assignment&(1< "new_tool" // - User specifies --tools=old_tool // Expected behavior: @@ -1850,7 +2116,7 @@ func TestWithMCPApps_EnabledPreservesUIMetadata(t *testing.T) { // Feature checker enables MCP Apps - UI meta should be preserved mcpAppsChecker := func(_ context.Context, flag string) (bool, error) { - return flag == mcpAppsFeatureFlag, nil + return flag == string(mcpAppsFeatureFlag), nil } reg := mustBuild(t, NewBuilder(). SetTools([]ServerTool{toolWithUI}). diff --git a/pkg/inventory/resources.go b/pkg/inventory/resources.go index 2dd07ae0fe..62a846120f 100644 --- a/pkg/inventory/resources.go +++ b/pkg/inventory/resources.go @@ -16,12 +16,8 @@ type ServerResourceTemplate struct { HandlerFunc ResourceHandlerFunc // Toolset identifies which toolset this resource belongs to Toolset ToolsetMetadata - // FeatureFlagEnable specifies a feature flag that must be enabled for this resource - // to be available. If set and the flag is not enabled, the resource is omitted. - FeatureFlagEnable string - // FeatureFlagDisable specifies feature flags that, when any is enabled, cause this - // resource to be omitted. Used to disable resources when a feature flag is on. - FeatureFlagDisable []string + // FeatureRule controls whether this resource is available. + FeatureRule FeatureRule } // HasHandler returns true if this resource has a handler function. diff --git a/pkg/inventory/server_tool.go b/pkg/inventory/server_tool.go index 9c458c3d12..2bc2769593 100644 --- a/pkg/inventory/server_tool.go +++ b/pkg/inventory/server_tool.go @@ -94,17 +94,9 @@ type ServerTool struct { // and handlers are only created when needed. HandlerFunc HandlerFunc - // FeatureFlagEnable specifies a feature flag that must be enabled for this tool - // to be available. If set and the flag is not enabled, the tool is omitted. - FeatureFlagEnable string - - // FeatureFlagEnableAll specifies additional feature flags that must all be enabled - // for this tool to be available. - FeatureFlagEnableAll []string - - // FeatureFlagDisable specifies feature flags that, when any is enabled, cause this - // tool to be omitted. Used to disable tools when a feature flag is on. - FeatureFlagDisable []string + // FeatureRule declares and evaluates the feature flags that control whether + // this tool is available. Its zero value leaves the tool available. + FeatureRule FeatureRule // Enabled is an optional function called at build/filter time to determine // if this tool should be available. If nil, the tool is considered enabled diff --git a/pkg/inventory/tool_availability.go b/pkg/inventory/tool_availability.go index 937726a8c7..5f7edcee88 100644 --- a/pkg/inventory/tool_availability.go +++ b/pkg/inventory/tool_availability.go @@ -4,6 +4,7 @@ import ( "context" "fmt" + ghcontext "github.com/github/github-mcp-server/pkg/context" "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -26,6 +27,14 @@ type toolAvailability struct { requiredElicitationMode ElicitationMode } +type toolFeatureDecision uint8 + +const ( + evaluateToolFeatureRule toolFeatureDecision = iota + excludeToolBeforeFeatureRule + includeToolWithoutFeatureRule +) + func (st *ServerTool) availability() toolAvailability { return toolAvailability{ minimumProtocolVersion: st.MinimumProtocolVersion, @@ -37,6 +46,40 @@ func (a toolAvailability) unrestricted() bool { return a.minimumProtocolVersion == "" && a.requiredElicitationMode == "" } +func featureDecisionForToolAvailability(ctx context.Context, availability toolAvailability) toolFeatureDecision { + if availability.unrestricted() { + return evaluateToolFeatureRule + } + info, ok := ghcontext.MCPMethod(ctx) + if !ok || info == nil || (info.Method != MCPMethodToolsList && info.Method != MCPMethodToolsCall) { + return evaluateToolFeatureRule + } + + available, known := knownToolAvailability(info.ProtocolVersion, info.ClientCapabilities, availability) + if !known || available { + return evaluateToolFeatureRule + } + if info.Method == MCPMethodToolsCall { + return includeToolWithoutFeatureRule + } + return excludeToolBeforeFeatureRule +} + +func knownToolAvailability(protocolVersion string, capabilities *mcp.ClientCapabilities, availability toolAvailability) (bool, bool) { + protocolKnown := availability.minimumProtocolVersion == "" || protocolVersion != "" + if protocolKnown && !protocolVersionAllowed(protocolVersion, availability.minimumProtocolVersion) { + return false, true + } + capabilitiesKnown := availability.requiredElicitationMode == "" || capabilities != nil + if capabilitiesKnown && !elicitationModeSupported(capabilities, availability.requiredElicitationMode) { + return false, true + } + if !protocolKnown || !capabilitiesKnown { + return false, false + } + return true, true +} + func addToolAvailabilityMiddleware(server *mcp.Server, tools []ServerTool) { availabilityByName := make(map[string]toolAvailability) for _, tool := range tools { diff --git a/pkg/inventory/tool_availability_test.go b/pkg/inventory/tool_availability_test.go index 1f76e3cb26..d52e4ff607 100644 --- a/pkg/inventory/tool_availability_test.go +++ b/pkg/inventory/tool_availability_test.go @@ -5,7 +5,9 @@ import ( "encoding/json" "errors" "testing" + "time" + ghcontext "github.com/github/github-mcp-server/pkg/context" "github.com/google/jsonschema-go/jsonschema" "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/stretchr/testify/assert" @@ -140,6 +142,146 @@ func TestToolAvailability(t *testing.T) { } } +func TestToolCallAvailabilitySkipsFeatureChecksOnlyWhenUnavailable(t *testing.T) { + tests := []struct { + name string + minimumProtocolVersion string + requiredElicitationMode ElicitationMode + methodInfo *ghcontext.MCPMethodInfo + clientCapabilities *mcp.ClientCapabilities + legacyProtocol bool + wantCheckerCalls int + wantHandlerCalls int + wantError string + }{ + { + name: "protocol unavailable", + minimumProtocolVersion: ProtocolVersionMultiRoundTrip, + methodInfo: &ghcontext.MCPMethodInfo{ + Method: MCPMethodToolsCall, + ProtocolVersion: "2025-11-25", + ClientCapabilities: &mcp.ClientCapabilities{}, + }, + legacyProtocol: true, + wantError: `Tool "restricted" requires MCP protocol version 2026-07-28 or later.`, + wantCheckerCalls: 0, + }, + { + name: "elicitation unavailable", + requiredElicitationMode: ElicitationModeForm, + methodInfo: &ghcontext.MCPMethodInfo{ + Method: MCPMethodToolsCall, + ProtocolVersion: ProtocolVersionMultiRoundTrip, + ClientCapabilities: &mcp.ClientCapabilities{ + Elicitation: &mcp.ElicitationCapabilities{URL: &mcp.URLElicitationCapabilities{}}, + }, + }, + clientCapabilities: &mcp.ClientCapabilities{ + Elicitation: &mcp.ElicitationCapabilities{URL: &mcp.URLElicitationCapabilities{}}, + }, + wantError: `Tool "restricted" requires client support for form elicitation.`, + wantCheckerCalls: 0, + }, + { + name: "eligible", + minimumProtocolVersion: ProtocolVersionMultiRoundTrip, + requiredElicitationMode: ElicitationModeForm, + methodInfo: &ghcontext.MCPMethodInfo{ + Method: MCPMethodToolsCall, + ProtocolVersion: ProtocolVersionMultiRoundTrip, + ClientCapabilities: &mcp.ClientCapabilities{ + Elicitation: &mcp.ElicitationCapabilities{Form: &mcp.FormElicitationCapabilities{}}, + }, + }, + clientCapabilities: &mcp.ClientCapabilities{ + Elicitation: &mcp.ElicitationCapabilities{Form: &mcp.FormElicitationCapabilities{}}, + }, + wantCheckerCalls: 1, + wantHandlerCalls: 1, + }, + { + name: "unknown metadata defers", + minimumProtocolVersion: ProtocolVersionMultiRoundTrip, + requiredElicitationMode: ElicitationModeForm, + methodInfo: &ghcontext.MCPMethodInfo{Method: MCPMethodToolsCall}, + clientCapabilities: &mcp.ClientCapabilities{ + Elicitation: &mcp.ElicitationCapabilities{Form: &mcp.FormElicitationCapabilities{}}, + }, + wantCheckerCalls: 1, + wantHandlerCalls: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + checkerCalls := 0 + handlerCalls := 0 + tool := availabilityTestTool( + "restricted", + tt.minimumProtocolVersion, + tt.requiredElicitationMode, + func() { handlerCalls++ }, + ) + if tt.wantError != "" { + tool.Tool.Meta = map[string]any{"ui": map[string]any{"resourceUri": "ui://restricted"}} + } + tool.FeatureRule = NewFeatureRule([]FeatureFlag{"feature"}, func(featureAsBool FeatureResolver) bool { + return featureAsBool("feature") + }) + inv, err := NewBuilder(). + SetTools([]ServerTool{tool}). + WithToolsets([]string{"all"}). + WithFeatureChecker(func(context.Context, string) (bool, error) { + time.Sleep(time.Millisecond) + checkerCalls++ + return true, nil + }). + Build() + require.NoError(t, err) + + ctx := ghcontext.WithMCPMethodInfo(context.Background(), tt.methodInfo) + server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "v0.0.1"}, nil) + inv.ForMCPRequest(MCPMethodToolsCall, "restricted").RegisterTools(ctx, server, nil) + require.Equal(t, tt.wantCheckerCalls, checkerCalls) + if tt.legacyProtocol { + server.AddReceivingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, request mcp.Request) (mcp.Result, error) { + if method == "server/discover" { + return nil, errors.New("legacy server does not support discovery") + } + return next(ctx, method, request) + } + }) + } + + serverTransport, clientTransport := mcp.NewInMemoryTransports() + serverSession, err := server.Connect(context.Background(), serverTransport, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = serverSession.Close() }) + client := mcp.NewClient(&mcp.Implementation{Name: "test-client", Version: "v0.0.1"}, &mcp.ClientOptions{ + Capabilities: tt.clientCapabilities, + }) + clientSession, err := client.Connect(context.Background(), clientTransport, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = clientSession.Close() }) + + result, err := clientSession.CallTool(context.Background(), &mcp.CallToolParams{Name: "restricted"}) + require.NoError(t, err) + require.Equal(t, tt.wantCheckerCalls, checkerCalls) + require.Equal(t, tt.wantHandlerCalls, handlerCalls) + if tt.wantError == "" { + assert.False(t, result.IsError) + return + } + require.True(t, result.IsError) + require.Len(t, result.Content, 1) + content, ok := result.Content[0].(*mcp.TextContent) + require.True(t, ok) + assert.Equal(t, tt.wantError, content.Text) + }) + } +} + func availabilityTestTool(name, minimumProtocolVersion string, requiredElicitationMode ElicitationMode, onCall func()) ServerTool { return ServerTool{ Tool: mcp.Tool{ diff --git a/script/print-mcp-diff-configs/main.go b/script/print-mcp-diff-configs/main.go index 421c9fce41..ada18c4063 100644 --- a/script/print-mcp-diff-configs/main.go +++ b/script/print-mcp-diff-configs/main.go @@ -141,7 +141,7 @@ func baseEntries() []baseEntry { }}, } - flags := append([]string(nil), github.AllowedFeatureFlags...) + flags := github.HeaderAllowedFeatureFlags() sort.Strings(flags) for _, f := range flags { entries = append(entries, baseEntry{ @@ -208,7 +208,7 @@ func (s settings) toHeaders() map[string]string { } func firstFeatureFlag() string { - flags := append([]string(nil), github.AllowedFeatureFlags...) + flags := github.HeaderAllowedFeatureFlags() if len(flags) == 0 { return "" }