diff --git a/pkg/context/mcp_info.go b/pkg/context/mcp_info.go index fc09e6b139..178906efa4 100644 --- a/pkg/context/mcp_info.go +++ b/pkg/context/mcp_info.go @@ -3,6 +3,7 @@ package context import ( "context" "encoding/json" + "sync" "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -22,8 +23,22 @@ type MCPMethodInfo struct { // ItemName is the name of the specific item being accessed (tool name, resource URI, prompt name) // Only populated for call/get methods (tools/call, prompts/get, resources/read) ItemName string + // Owner and Repo are best-effort exact-key string values from tools/call arguments. + // Invalid or missing values are represented by an empty string. + Owner string + Repo string + // Arguments contains the decoded tools/call arguments after DecodeArguments is + // called. The returned map is shared and must not be mutated. + // + // Prefer RawArguments and DecodeArguments in new code. + Arguments map[string]any // RawArguments contains the unmaterialized tool arguments for tools/call requests. RawArguments json.RawMessage + + decodeMu sync.Mutex + decodeDone bool + decodeErr error + decodedArguments map[string]any // ProtocolVersion and ClientCapabilities describe the requesting MCP client // when stateless HTTP parsing makes them available before registration. ProtocolVersion string @@ -31,18 +46,39 @@ type MCPMethodInfo struct { } // DecodeArguments materializes tool arguments when request middleware needs -// call-specific values. Invalid argument shapes are returned to the caller so -// the request can continue to the tool handler's normal validation path. +// call-specific values. Concurrent calls on the same MCPMethodInfo are safe. +// The returned map is cached and shared; callers must treat it as read-only and +// must not mutate MCPMethodInfo argument fields after the first call. +// Invalid argument shapes are returned to the caller so the request can +// continue to the tool handler's normal validation path. func (info *MCPMethodInfo) DecodeArguments() (map[string]any, error) { + info.decodeMu.Lock() + defer info.decodeMu.Unlock() + + if info.decodeDone { + return info.decodedArguments, info.decodeErr + } + if info.Arguments != nil { + info.decodedArguments = info.Arguments + info.decodeDone = true + return info.decodedArguments, nil + } if len(info.RawArguments) == 0 { + info.decodeDone = true return nil, nil } var arguments map[string]any if err := json.Unmarshal(info.RawArguments, &arguments); err != nil { + info.decodeErr = err + info.decodeDone = true return nil, err } - return arguments, nil + info.decodedArguments = arguments + info.Arguments = arguments + info.decodeDone = true + + return info.decodedArguments, nil } // WithMCPMethodInfo stores the MCPMethodInfo in the context. diff --git a/pkg/context/mcp_info_test.go b/pkg/context/mcp_info_test.go new file mode 100644 index 0000000000..0f969914c9 --- /dev/null +++ b/pkg/context/mcp_info_test.go @@ -0,0 +1,70 @@ +package context + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestMCPMethodInfoDecodeArgumentsCachesResult(t *testing.T) { + info := &MCPMethodInfo{RawArguments: []byte(`{"owner":"github","repo":"github-mcp-server"}`)} + + decoded, err := info.DecodeArguments() + require.NoError(t, err) + require.NotNil(t, decoded) + assert.Equal(t, "github", decoded["owner"]) + assert.Equal(t, decoded, info.Arguments) + + cached, err := info.DecodeArguments() + require.NoError(t, err) + assert.Equal(t, decoded, cached) +} + +func TestMCPMethodInfoDecodeArgumentsUsesExistingArguments(t *testing.T) { + arguments := map[string]any{"owner": "github"} + info := &MCPMethodInfo{Arguments: arguments} + + decoded, err := info.DecodeArguments() + require.NoError(t, err) + assert.Equal(t, arguments, decoded) +} + +func TestMCPMethodInfoDecodeArgumentsCachesErrors(t *testing.T) { + info := &MCPMethodInfo{RawArguments: []byte(`["not","an","object"]`)} + + _, err := info.DecodeArguments() + require.Error(t, err) + _, cachedErr := info.DecodeArguments() + assert.Same(t, err, cachedErr) +} + +func TestMCPMethodInfoDecodeArgumentsConcurrentCalls(t *testing.T) { + info := &MCPMethodInfo{ + RawArguments: []byte(`{"owner":"github","repo":"github-mcp-server","nested":{"path":"README.md"}}`), + } + + var wg sync.WaitGroup + errors := make(chan error, 32) + for range 32 { + wg.Add(1) + go func() { + defer wg.Done() + decoded, err := info.DecodeArguments() + if err != nil { + errors <- err + return + } + if decoded["owner"] != "github" { + errors <- assert.AnError + } + }() + } + wg.Wait() + close(errors) + + for err := range errors { + require.NoError(t, err) + } +} diff --git a/pkg/http/middleware/mcp_parse.go b/pkg/http/middleware/mcp_parse.go index bcc44f01f5..f5276cae03 100644 --- a/pkg/http/middleware/mcp_parse.go +++ b/pkg/http/middleware/mcp_parse.go @@ -115,6 +115,7 @@ func parseMCPMethodInfo(body []byte) (*ghcontext.MCPMethodInfo, error) { case "tools/call": methodInfo.ItemName = mcpReq.Params.Name methodInfo.RawArguments = mcpReq.Params.Arguments + populateToolArgumentCompatibilityFields(methodInfo) case "prompts/get": methodInfo.ItemName = mcpReq.Params.Name case "resources/read": @@ -122,3 +123,21 @@ func parseMCPMethodInfo(body []byte) (*ghcontext.MCPMethodInfo, error) { } return methodInfo, nil } + +func populateToolArgumentCompatibilityFields(methodInfo *ghcontext.MCPMethodInfo) { + if len(methodInfo.RawArguments) == 0 { + return + } + + var arguments map[string]json.RawMessage + if err := json.Unmarshal(methodInfo.RawArguments, &arguments); err != nil { + return + } + + if owner, ok := arguments["owner"]; ok { + _ = json.Unmarshal(owner, &methodInfo.Owner) + } + if repo, ok := arguments["repo"]; ok { + _ = json.Unmarshal(repo, &methodInfo.Repo) + } +} diff --git a/pkg/http/middleware/mcp_parse_test.go b/pkg/http/middleware/mcp_parse_test.go index 8034dc29a6..e87a24fe5d 100644 --- a/pkg/http/middleware/mcp_parse_test.go +++ b/pkg/http/middleware/mcp_parse_test.go @@ -24,6 +24,8 @@ func TestWithMCPParse(t *testing.T) { expectedMethod string expectedItem string expectedRaw string + expectedOwner string + expectedRepo string expectedArgs map[string]any expectedProtocol string expectedForm bool @@ -110,6 +112,8 @@ func TestWithMCPParse(t *testing.T) { expectedMethod: "tools/call", expectedItem: "get_file_contents", expectedRaw: `{"owner":"github","repo":"github-mcp-server","path":"README.md"}`, + expectedOwner: "github", + expectedRepo: "github-mcp-server", expectedArgs: map[string]any{"owner": "github", "repo": "github-mcp-server", "path": "README.md"}, }, { @@ -123,6 +127,28 @@ func TestWithMCPParse(t *testing.T) { expectedRaw: `"not an object"`, expectArgsError: true, }, + { + name: "invalid owner type does not block repo extraction", + method: http.MethodPost, + path: "/mcp", + body: `{"jsonrpc":"2.0","method":"tools/call","params":{"name":"get_file_contents","arguments":{"owner":123,"repo":"github-mcp-server"}}}`, + expectInfo: true, + expectedMethod: "tools/call", + expectedItem: "get_file_contents", + expectedRepo: "github-mcp-server", + }, + { + name: "compatibility fields use exact keys and last duplicates", + method: http.MethodPost, + path: "/mcp", + body: `{"jsonrpc":"2.0","method":"tools/call","params":{"name":"get_file_contents","arguments":{"Owner":"ignored","owner":"first","owner":"second","repo":"one","repo":"two"}}}`, + expectInfo: true, + expectedMethod: "tools/call", + expectedItem: "get_file_contents", + expectedOwner: "second", + expectedRepo: "two", + expectedArgs: map[string]any{"Owner": "ignored", "owner": "second", "repo": "two"}, + }, { name: "prompts/get parses name", method: http.MethodPost, @@ -183,6 +209,8 @@ func TestWithMCPParse(t *testing.T) { if tt.expectedRaw != "" { assert.JSONEq(t, tt.expectedRaw, string(capturedInfo.RawArguments)) } + assert.Equal(t, tt.expectedOwner, capturedInfo.Owner) + assert.Equal(t, tt.expectedRepo, capturedInfo.Repo) decodedArgs, err := capturedInfo.DecodeArguments() if tt.expectArgsError { assert.Error(t, err) @@ -191,6 +219,7 @@ func TestWithMCPParse(t *testing.T) { } if tt.expectedArgs != nil { assert.Equal(t, tt.expectedArgs, decodedArgs) + assert.Equal(t, decodedArgs, capturedInfo.Arguments) } } else { assert.False(t, infoCaptured, "MCPMethodInfo should not be present in context")