From 82a3d2c759185ee6c49cd6a7747b54442a3d24ca Mon Sep 17 00:00:00 2001 From: Mike Gabriel Date: Mon, 29 Dec 2025 13:03:49 -0500 Subject: [PATCH 1/6] feat(router): allow custom MCP tool names via @mcpTool directive --- router-tests/mcp_test.go | 58 +++++++--- .../mcp_operations/CustomToolName.graphql | 11 ++ router/pkg/mcpserver/server.go | 12 +- router/pkg/schemaloader/loader.go | 80 +++++++++++++- router/pkg/schemaloader/loader_test.go | 104 ++++++++++++++++++ 5 files changed, 242 insertions(+), 23 deletions(-) create mode 100644 router-tests/testdata/mcp_operations/CustomToolName.graphql diff --git a/router-tests/mcp_test.go b/router-tests/mcp_test.go index e89a2c0388..2fc0fa6fd2 100644 --- a/router-tests/mcp_test.go +++ b/router-tests/mcp_test.go @@ -38,19 +38,21 @@ func TestMCP(t *testing.T) { require.NoError(t, err) require.NotNil(t, resp) - require.Contains(t, resp.Tools, mcp.Tool{ - Name: "get_operation_info", - Description: "Provides instructions on how to execute the GraphQL operation via HTTP and how to integrate it into your application.", - InputSchema: mcp.ToolInputSchema{ - Type: "object", - Properties: map[string]interface{}{"operationName": map[string]interface{}{"description": "The exact name of the GraphQL operation to retrieve information for.", "enum": []interface{}{"UpdateMood", "MyEmployees"}, "type": "string"}}, - Required: []string{"operationName"}}, - RawInputSchema: json.RawMessage(nil), - Annotations: mcp.ToolAnnotation{ - Title: "Get GraphQL Operation Info", - ReadOnlyHint: mcp.ToBoolPtr(true), - }, - }) + var operationInfoTool *mcp.Tool + for i, tool := range resp.Tools { + if tool.Name == "get_operation_info" { + operationInfoTool = &resp.Tools[i] + break + } + } + require.NotNil(t, operationInfoTool, "get_operation_info tool should exist") + require.Equal(t, "Provides instructions on how to execute the GraphQL operation via HTTP and how to integrate it into your application.", operationInfoTool.Description) + operationNameProp := operationInfoTool.InputSchema.Properties["operationName"].(map[string]interface{}) + enum := operationNameProp["enum"].([]interface{}) + require.Len(t, enum, 3) + require.Contains(t, enum, "UpdateMood") + require.Contains(t, enum, "MyEmployees") + require.Contains(t, enum, "CustomNamedQuery") }) }) @@ -187,6 +189,36 @@ func TestMCP(t *testing.T) { }) }) + t.Run("List user Operations / Custom tool name via @mcpTool directive", func(t *testing.T) { + testenv.Run(t, &testenv.Config{ + MCP: config.MCPConfiguration{ + Enabled: true, + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + + toolsRequest := mcp.ListToolsRequest{} + resp, err := xEnv.MCPClient.ListTools(xEnv.Context, toolsRequest) + require.NoError(t, err) + require.NotNil(t, resp) + + var customTool *mcp.Tool + for i, tool := range resp.Tools { + if tool.Name == "get_employee_by_id" { + customTool = &resp.Tools[i] + break + } + } + + require.NotNil(t, customTool, "Tool get_employee_by_id should be found") + assert.Equal(t, "A query with a custom MCP tool name.", customTool.Description) + + for _, tool := range resp.Tools { + assert.NotEqual(t, "execute_operation_custom_named_query", tool.Name, + "Tool should not be registered with default generated name") + } + }) + }) + t.Run("Execute Operation Info", func(t *testing.T) { testenv.Run(t, &testenv.Config{ MCP: config.MCPConfiguration{ diff --git a/router-tests/testdata/mcp_operations/CustomToolName.graphql b/router-tests/testdata/mcp_operations/CustomToolName.graphql new file mode 100644 index 0000000000..c2d77d9838 --- /dev/null +++ b/router-tests/testdata/mcp_operations/CustomToolName.graphql @@ -0,0 +1,11 @@ +""" +A query with a custom MCP tool name. +""" +query CustomNamedQuery($id: Int!) @mcpTool(name: "get_employee_by_id") { + employee(id: $id) { + id + details { + forename + } + } +} diff --git a/router/pkg/mcpserver/server.go b/router/pkg/mcpserver/server.go index 173e2bf9d2..b970d503f3 100644 --- a/router/pkg/mcpserver/server.go +++ b/router/pkg/mcpserver/server.go @@ -536,8 +536,14 @@ func (s *GraphQLSchemaServer) registerTools() error { compiledSchema: compiledSchema, } - // Convert the operation name to snake_case for consistent tool naming - operationToolName := strcase.ToSnake(op.Name) + // Use custom tool name if provided via @mcpTool directive, otherwise generate default + var toolName string + if op.ToolName != "" { + toolName = op.ToolName + } else { + operationToolName := strcase.ToSnake(op.Name) + toolName = fmt.Sprintf("execute_operation_%s", operationToolName) + } // Use the operation description directly if provided, otherwise generate a default description var toolDescription string @@ -546,8 +552,6 @@ func (s *GraphQLSchemaServer) registerTools() error { } else { toolDescription = fmt.Sprintf("Executes the GraphQL operation '%s' of type %s.", op.Name, op.OperationType) } - - toolName := fmt.Sprintf("execute_operation_%s", operationToolName) tool := mcp.NewToolWithRawSchema( toolName, toolDescription, diff --git a/router/pkg/schemaloader/loader.go b/router/pkg/schemaloader/loader.go index cd3f53dad8..d997eae00e 100644 --- a/router/pkg/schemaloader/loader.go +++ b/router/pkg/schemaloader/loader.go @@ -18,6 +18,7 @@ import ( // Operation represents a GraphQL operation with its AST document and schema information type Operation struct { Name string + ToolName string // User-defined MCP tool name from @mcp(name: "...") directive (optional) FilePath string Document ast.Document OperationString string @@ -46,8 +47,7 @@ func NewOperationLoader(logger *zap.Logger, schemaDoc *ast.Document) *OperationL func (l *OperationLoader) LoadOperationsFromDirectory(dirPath string) ([]Operation, error) { var operations []Operation - // Create an operation validator - validator := astvalidation.DefaultOperationValidator() + validator := newMCPOperationValidator() // Walk through the directory and process GraphQL files err := filepath.WalkDir(dirPath, func(path string, d fs.DirEntry, err error) error { @@ -92,7 +92,10 @@ func (l *OperationLoader) LoadOperationsFromDirectory(dirPath string) ([]Operati return nil } - // Validate operation against schema + opDescription := extractOperationDescription(&opDoc) + toolName := extractMCPToolName(&opDoc) + stripMCPDirective(&opDoc) + validationReport := operationreport.Report{} validationState := validator.Validate(&opDoc, l.SchemaDocument, &validationReport) if validationState == astvalidation.Invalid { @@ -116,12 +119,10 @@ func (l *OperationLoader) LoadOperationsFromDirectory(dirPath string) ([]Operati } } - // Extract description from operation definition - opDescription := extractOperationDescription(&opDoc) - // Add to our list of operations operations = append(operations, Operation{ Name: opName, + ToolName: toolName, FilePath: path, Document: opDoc, OperationString: operationString, @@ -200,3 +201,70 @@ func extractOperationDescription(doc *ast.Document) string { } return "" } + +var mcpDirectiveName = []byte("mcpTool") +var mcpNameArgument = []byte("name") + +func newMCPOperationValidator() *astvalidation.OperationValidator { + return astvalidation.NewOperationValidator([]astvalidation.Rule{ + astvalidation.AllVariablesUsed(), + astvalidation.AllVariableUsesDefined(), + astvalidation.DocumentContainsExecutableOperation(), + astvalidation.OperationNameUniqueness(), + astvalidation.LoneAnonymousOperation(), + astvalidation.SubscriptionSingleRootField(), + astvalidation.FieldSelections(), + astvalidation.FieldSelectionMerging(), + astvalidation.KnownArguments(), + astvalidation.Values(), + astvalidation.ArgumentUniqueness(), + astvalidation.RequiredArguments(), + astvalidation.Fragments(), + astvalidation.DirectivesAreInValidLocations(), + astvalidation.DirectivesAreUniquePerLocation(), + astvalidation.VariableUniqueness(), + astvalidation.VariablesAreInputTypes(), + }) +} + +func extractMCPToolName(doc *ast.Document) string { + for _, ref := range doc.RootNodes { + if ref.Kind == ast.NodeKindOperationDefinition { + opDef := doc.OperationDefinitions[ref.Ref] + if !opDef.HasDirectives { + return "" + } + + directiveRef, exists := doc.DirectiveWithNameBytes(opDef.Directives.Refs, mcpDirectiveName) + if !exists { + return "" + } + + value, argExists := doc.DirectiveArgumentValueByName(directiveRef, mcpNameArgument) + if !argExists { + return "" + } + + if value.Kind == ast.ValueKindString { + return doc.StringValueContentString(value.Ref) + } + + return "" + } + } + return "" +} + +func stripMCPDirective(doc *ast.Document) { + for _, ref := range doc.RootNodes { + if ref.Kind == ast.NodeKindOperationDefinition { + opDef := &doc.OperationDefinitions[ref.Ref] + if !opDef.HasDirectives { + return + } + opDef.Directives.RemoveDirectiveByName(doc, "mcpTool") + opDef.HasDirectives = len(opDef.Directives.Refs) > 0 + return + } + } +} diff --git a/router/pkg/schemaloader/loader_test.go b/router/pkg/schemaloader/loader_test.go index b4573d89a5..729fc0167c 100644 --- a/router/pkg/schemaloader/loader_test.go +++ b/router/pkg/schemaloader/loader_test.go @@ -165,3 +165,107 @@ func TestLoadOperationsFromEmptyDirectory(t *testing.T) { require.NoError(t, err) assert.Len(t, operations, 0, "Empty directory should return no operations") } + +func TestExtractMCPToolName(t *testing.T) { + tests := []struct { + name string + query string + expected string + }{ + { + name: "with @mcpTool directive and name argument", + query: `query Foo @mcpTool(name: "custom_foo") { bar }`, + expected: "custom_foo", + }, + { + name: "without directive", + query: `query Foo { bar }`, + expected: "", + }, + { + name: "with @mcpTool but no name argument", + query: `query Foo @mcpTool { bar }`, + expected: "", + }, + { + name: "with different directive", + query: `query Foo @deprecated { bar }`, + expected: "", + }, + { + name: "with @mcpTool and other arguments but no name", + query: `query Foo @mcpTool(other: "value") { bar }`, + expected: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + doc, report := astparser.ParseGraphqlDocumentString(tt.query) + require.False(t, report.HasErrors()) + + result := extractMCPToolName(&doc) + assert.Equal(t, tt.expected, result) + }) + } +} + +func TestLoadOperationsWithMCPDirective(t *testing.T) { + tempDir := t.TempDir() + + testFiles := map[string]string{ + "WithCustomToolName.graphql": `"""Custom tool operation""" +query MyQuery @mcpTool(name: "custom_tool_name") { + employee(id: "1") { + id + name + } +}`, + "WithoutDirective.graphql": `query AnotherQuery { + employee(id: "1") { + id + name + } +}`, + } + + for filename, content := range testFiles { + err := os.WriteFile(filepath.Join(tempDir, filename), []byte(content), 0644) + require.NoError(t, err) + } + + schemaStr := ` +directive @mcpTool(name: String) on QUERY | MUTATION + +type Query { + employee(id: ID!): Employee +} + +type Employee { + id: ID! + name: String! +} +` + schemaDoc, report := astparser.ParseGraphqlDocumentString(schemaStr) + require.False(t, report.HasErrors()) + + err := asttransform.MergeDefinitionWithBaseSchema(&schemaDoc) + require.NoError(t, err) + + logger := zap.NewNop() + loader := NewOperationLoader(logger, &schemaDoc) + operations, err := loader.LoadOperationsFromDirectory(tempDir) + require.NoError(t, err) + require.Len(t, operations, 2) + + opMap := make(map[string]Operation) + for _, op := range operations { + opMap[op.Name] = op + } + + op1 := opMap["MyQuery"] + assert.Equal(t, "custom_tool_name", op1.ToolName) + + op2 := opMap["AnotherQuery"] + assert.Empty(t, op2.ToolName) +} From d5c4cda25f1d07941cb23a6aa7adf7b4e82b1e6f Mon Sep 17 00:00:00 2001 From: Mike Gabriel Date: Mon, 29 Dec 2025 13:09:16 -0500 Subject: [PATCH 2/6] chore(router): remove outdated comment from ToolName field --- router/pkg/schemaloader/loader.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/router/pkg/schemaloader/loader.go b/router/pkg/schemaloader/loader.go index d997eae00e..384b6b9efe 100644 --- a/router/pkg/schemaloader/loader.go +++ b/router/pkg/schemaloader/loader.go @@ -18,7 +18,7 @@ import ( // Operation represents a GraphQL operation with its AST document and schema information type Operation struct { Name string - ToolName string // User-defined MCP tool name from @mcp(name: "...") directive (optional) + ToolName string FilePath string Document ast.Document OperationString string From c2abf4db3e7106ee7b2ea45b03d9db1de70a88ee Mon Sep 17 00:00:00 2001 From: Mike Gabriel Date: Mon, 29 Dec 2025 13:15:17 -0500 Subject: [PATCH 3/6] docs(router): add docstrings for MCP directive helper functions --- router/pkg/schemaloader/loader.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/router/pkg/schemaloader/loader.go b/router/pkg/schemaloader/loader.go index 384b6b9efe..ad6750d026 100644 --- a/router/pkg/schemaloader/loader.go +++ b/router/pkg/schemaloader/loader.go @@ -205,6 +205,7 @@ func extractOperationDescription(doc *ast.Document) string { var mcpDirectiveName = []byte("mcpTool") var mcpNameArgument = []byte("name") +// newMCPOperationValidator creates an operation validator that allows the @mcpTool directive without schema definition func newMCPOperationValidator() *astvalidation.OperationValidator { return astvalidation.NewOperationValidator([]astvalidation.Rule{ astvalidation.AllVariablesUsed(), @@ -227,6 +228,7 @@ func newMCPOperationValidator() *astvalidation.OperationValidator { }) } +// extractMCPToolName extracts the custom tool name from the @mcpTool directive on an operation func extractMCPToolName(doc *ast.Document) string { for _, ref := range doc.RootNodes { if ref.Kind == ast.NodeKindOperationDefinition { @@ -255,6 +257,7 @@ func extractMCPToolName(doc *ast.Document) string { return "" } +// stripMCPDirective removes the @mcpTool directive from an operation before validation func stripMCPDirective(doc *ast.Document) { for _, ref := range doc.RootNodes { if ref.Kind == ast.NodeKindOperationDefinition { From bf4c87ea252fee2ea3941016b23a99ecf364204d Mon Sep 17 00:00:00 2001 From: Mike Gabriel Date: Mon, 29 Dec 2025 13:22:19 -0500 Subject: [PATCH 4/6] refactor(router): use mcpDirectiveName constant in stripMCPDirective --- router/pkg/schemaloader/loader.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/router/pkg/schemaloader/loader.go b/router/pkg/schemaloader/loader.go index ad6750d026..b0d570c445 100644 --- a/router/pkg/schemaloader/loader.go +++ b/router/pkg/schemaloader/loader.go @@ -265,7 +265,7 @@ func stripMCPDirective(doc *ast.Document) { if !opDef.HasDirectives { return } - opDef.Directives.RemoveDirectiveByName(doc, "mcpTool") + opDef.Directives.RemoveDirectiveByName(doc, string(mcpDirectiveName)) opDef.HasDirectives = len(opDef.Directives.Refs) > 0 return } From 57714a415003b53fae5e4062609e22d86118abcd Mon Sep 17 00:00:00 2001 From: Mike Gabriel Date: Mon, 29 Dec 2025 14:03:19 -0500 Subject: [PATCH 5/6] refactor(router): use DefaultOperationValidator for MCP operations The custom newMCPOperationValidator was unnecessary because stripMCPDirective removes the @mcpTool directive before validation runs. --- router/pkg/schemaloader/loader.go | 25 +------------------------ 1 file changed, 1 insertion(+), 24 deletions(-) diff --git a/router/pkg/schemaloader/loader.go b/router/pkg/schemaloader/loader.go index b0d570c445..ead08b4b6c 100644 --- a/router/pkg/schemaloader/loader.go +++ b/router/pkg/schemaloader/loader.go @@ -47,7 +47,7 @@ func NewOperationLoader(logger *zap.Logger, schemaDoc *ast.Document) *OperationL func (l *OperationLoader) LoadOperationsFromDirectory(dirPath string) ([]Operation, error) { var operations []Operation - validator := newMCPOperationValidator() + validator := astvalidation.DefaultOperationValidator() // Walk through the directory and process GraphQL files err := filepath.WalkDir(dirPath, func(path string, d fs.DirEntry, err error) error { @@ -205,29 +205,6 @@ func extractOperationDescription(doc *ast.Document) string { var mcpDirectiveName = []byte("mcpTool") var mcpNameArgument = []byte("name") -// newMCPOperationValidator creates an operation validator that allows the @mcpTool directive without schema definition -func newMCPOperationValidator() *astvalidation.OperationValidator { - return astvalidation.NewOperationValidator([]astvalidation.Rule{ - astvalidation.AllVariablesUsed(), - astvalidation.AllVariableUsesDefined(), - astvalidation.DocumentContainsExecutableOperation(), - astvalidation.OperationNameUniqueness(), - astvalidation.LoneAnonymousOperation(), - astvalidation.SubscriptionSingleRootField(), - astvalidation.FieldSelections(), - astvalidation.FieldSelectionMerging(), - astvalidation.KnownArguments(), - astvalidation.Values(), - astvalidation.ArgumentUniqueness(), - astvalidation.RequiredArguments(), - astvalidation.Fragments(), - astvalidation.DirectivesAreInValidLocations(), - astvalidation.DirectivesAreUniquePerLocation(), - astvalidation.VariableUniqueness(), - astvalidation.VariablesAreInputTypes(), - }) -} - // extractMCPToolName extracts the custom tool name from the @mcpTool directive on an operation func extractMCPToolName(doc *ast.Document) string { for _, ref := range doc.RootNodes { From 08ea5bd78537d17a90fc66c206f8225ebfcee60a Mon Sep 17 00:00:00 2001 From: Mike Gabriel Date: Mon, 29 Dec 2025 14:44:31 -0500 Subject: [PATCH 6/6] test(router): improve MCP test readability and structure --- router-tests/mcp_test.go | 35 ++++++++++++++++++++--------------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/router-tests/mcp_test.go b/router-tests/mcp_test.go index 2fc0fa6fd2..0506760655 100644 --- a/router-tests/mcp_test.go +++ b/router-tests/mcp_test.go @@ -38,21 +38,26 @@ func TestMCP(t *testing.T) { require.NoError(t, err) require.NotNil(t, resp) - var operationInfoTool *mcp.Tool - for i, tool := range resp.Tools { - if tool.Name == "get_operation_info" { - operationInfoTool = &resp.Tools[i] - break - } - } - require.NotNil(t, operationInfoTool, "get_operation_info tool should exist") - require.Equal(t, "Provides instructions on how to execute the GraphQL operation via HTTP and how to integrate it into your application.", operationInfoTool.Description) - operationNameProp := operationInfoTool.InputSchema.Properties["operationName"].(map[string]interface{}) - enum := operationNameProp["enum"].([]interface{}) - require.Len(t, enum, 3) - require.Contains(t, enum, "UpdateMood") - require.Contains(t, enum, "MyEmployees") - require.Contains(t, enum, "CustomNamedQuery") + require.Contains(t, resp.Tools, mcp.Tool{ + Name: "get_operation_info", + Description: "Provides instructions on how to execute the GraphQL operation via HTTP and how to integrate it into your application.", + InputSchema: mcp.ToolInputSchema{ + Type: "object", + Properties: map[string]interface{}{ + "operationName": map[string]interface{}{ + "description": "The exact name of the GraphQL operation to retrieve information for.", + "enum": []interface{}{"CustomNamedQuery", "UpdateMood", "MyEmployees"}, + "type": "string", + }, + }, + Required: []string{"operationName"}, + }, + RawInputSchema: json.RawMessage(nil), + Annotations: mcp.ToolAnnotation{ + Title: "Get GraphQL Operation Info", + ReadOnlyHint: mcp.ToBoolPtr(true), + }, + }) }) })