feat: add metadata support to MCP tools #475
Conversation
Summary of ChangesHello @philippslang, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the MCP toolset by integrating robust metadata support. It provides a flexible mechanism for passing contextual metadata, such as tracing identifiers or authentication tokens, from the ADK environment directly to underlying MCP tool invocations. This change is crucial for maintaining context and enabling advanced features like distributed tracing across tool boundaries, thereby improving the observability and manageability of tool interactions within the system. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for forwarding metadata in MCP tool calls by adding a MetadataProvider. The changes are well-structured, with updates to the configuration, toolset creation, and tool execution logic to handle the new metadata provider. The core logic correctly handles cases where the provider is nil or returns nil metadata.
The feature is accompanied by a comprehensive set of tests covering the success case, the absence of a provider, and a provider that returns nil. I've left one comment regarding test code structure to improve maintainability by reducing duplication. Overall, this is a solid contribution.
| func TestMetadataProvider(t *testing.T) { | ||
| var receivedMeta map[string]any | ||
|
|
||
| echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { | ||
| receivedMeta = req.Params.Meta | ||
| return nil, struct{ Message string }{Message: "ok"}, nil | ||
| } | ||
|
|
||
| clientTransport, serverTransport := mcp.NewInMemoryTransports() | ||
|
|
||
| server := mcp.NewServer(&mcp.Implementation{Name: "test_server", Version: "v1.0.0"}, nil) | ||
| mcp.AddTool(server, &mcp.Tool{Name: "echo_tool", Description: "echoes input"}, echoToolFunc) | ||
| _, err := server.Connect(t.Context(), serverTransport, nil) | ||
| if err != nil { | ||
| t.Fatal(err) | ||
| } | ||
|
|
||
| testMetadata := map[string]any{ | ||
| "request_id": "test-123", | ||
| "user_id": "user-456", | ||
| "nested_data": map[string]any{"key": "value"}, | ||
| } | ||
| metadataProvider := func(ctx tool.Context) map[string]any { | ||
| return testMetadata | ||
| } | ||
|
|
||
| ts, err := mcptoolset.New(mcptoolset.Config{ | ||
| Transport: clientTransport, | ||
| MetadataProvider: metadataProvider, | ||
| }) | ||
| if err != nil { | ||
| t.Fatalf("Failed to create MCP tool set: %v", err) | ||
| } | ||
|
|
||
| invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) | ||
| readonlyCtx := icontext.NewReadonlyContext(invCtx) | ||
| tools, err := ts.Tools(readonlyCtx) | ||
| if err != nil { | ||
| t.Fatalf("Failed to get tools: %v", err) | ||
| } | ||
|
|
||
| if len(tools) != 1 { | ||
| t.Fatalf("Expected 1 tool, got %d", len(tools)) | ||
| } | ||
|
|
||
| fnTool, ok := tools[0].(toolinternal.FunctionTool) | ||
| if !ok { | ||
| t.Fatal("Tool does not implement FunctionTool interface") | ||
| } | ||
|
|
||
| toolCtx := toolinternal.NewToolContext(invCtx, "", nil) | ||
| result, err := fnTool.Run(toolCtx, map[string]any{}) | ||
| if err != nil { | ||
| t.Fatalf("Failed to run tool: %v", err) | ||
| } | ||
|
|
||
| if result == nil { | ||
| t.Fatal("Expected non-nil result") | ||
| } | ||
|
|
||
| if diff := cmp.Diff(testMetadata, receivedMeta); diff != "" { | ||
| t.Errorf("metadata mismatch (-want +got):\n%s", diff) | ||
| } | ||
| } | ||
|
|
||
| func TestMetadataProviderNil(t *testing.T) { | ||
| echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { | ||
| if req.Params.Meta != nil { | ||
| t.Errorf("Expected nil metadata, got %v", req.Params.Meta) | ||
| } | ||
| return nil, struct{ Message string }{Message: "ok"}, nil | ||
| } | ||
|
|
||
| clientTransport, serverTransport := mcp.NewInMemoryTransports() | ||
|
|
||
| server := mcp.NewServer(&mcp.Implementation{Name: "test_server", Version: "v1.0.0"}, nil) | ||
| mcp.AddTool(server, &mcp.Tool{Name: "echo_tool", Description: "echoes input"}, echoToolFunc) | ||
| _, err := server.Connect(t.Context(), serverTransport, nil) | ||
| if err != nil { | ||
| t.Fatal(err) | ||
| } | ||
|
|
||
| ts, err := mcptoolset.New(mcptoolset.Config{ | ||
| Transport: clientTransport, | ||
| // MetadataProvider is nil | ||
| }) | ||
| if err != nil { | ||
| t.Fatalf("Failed to create MCP tool set: %v", err) | ||
| } | ||
|
|
||
| invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) | ||
| readonlyCtx := icontext.NewReadonlyContext(invCtx) | ||
| tools, err := ts.Tools(readonlyCtx) | ||
| if err != nil { | ||
| t.Fatalf("Failed to get tools: %v", err) | ||
| } | ||
|
|
||
| fnTool, ok := tools[0].(toolinternal.FunctionTool) | ||
| if !ok { | ||
| t.Fatal("Tool does not implement FunctionTool interface") | ||
| } | ||
|
|
||
| toolCtx := toolinternal.NewToolContext(invCtx, "", nil) | ||
| _, err = fnTool.Run(toolCtx, map[string]any{}) | ||
| if err != nil { | ||
| t.Fatalf("Failed to run tool: %v", err) | ||
| } | ||
| } | ||
|
|
||
| func TestMetadataProviderReturnsNil(t *testing.T) { | ||
| var receivedMeta map[string]any | ||
| var metaCalled bool | ||
|
|
||
| echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { | ||
| metaCalled = true | ||
| receivedMeta = req.Params.Meta | ||
| return nil, struct{ Message string }{Message: "ok"}, nil | ||
| } | ||
|
|
||
| clientTransport, serverTransport := mcp.NewInMemoryTransports() | ||
|
|
||
| server := mcp.NewServer(&mcp.Implementation{Name: "test_server", Version: "v1.0.0"}, nil) | ||
| mcp.AddTool(server, &mcp.Tool{Name: "echo_tool", Description: "echoes input"}, echoToolFunc) | ||
| _, err := server.Connect(t.Context(), serverTransport, nil) | ||
| if err != nil { | ||
| t.Fatal(err) | ||
| } | ||
|
|
||
| metadataProvider := func(ctx tool.Context) map[string]any { | ||
| return nil | ||
| } | ||
|
|
||
| ts, err := mcptoolset.New(mcptoolset.Config{ | ||
| Transport: clientTransport, | ||
| MetadataProvider: metadataProvider, | ||
| }) | ||
| if err != nil { | ||
| t.Fatalf("Failed to create MCP tool set: %v", err) | ||
| } | ||
|
|
||
| invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) | ||
| readonlyCtx := icontext.NewReadonlyContext(invCtx) | ||
| tools, err := ts.Tools(readonlyCtx) | ||
| if err != nil { | ||
| t.Fatalf("Failed to get tools: %v", err) | ||
| } | ||
|
|
||
| fnTool, ok := tools[0].(toolinternal.FunctionTool) | ||
| if !ok { | ||
| t.Fatal("Tool does not implement FunctionTool interface") | ||
| } | ||
|
|
||
| toolCtx := toolinternal.NewToolContext(invCtx, "", nil) | ||
| _, err = fnTool.Run(toolCtx, map[string]any{}) | ||
| if err != nil { | ||
| t.Fatalf("Failed to run tool: %v", err) | ||
| } | ||
|
|
||
| if !metaCalled { | ||
| t.Fatal("Tool was not called") | ||
| } | ||
|
|
||
| if receivedMeta != nil { | ||
| t.Errorf("Expected nil metadata when provider returns nil, got %v", receivedMeta) | ||
| } | ||
| } |
There was a problem hiding this comment.
The three new tests TestMetadataProvider, TestMetadataProviderNil, and TestMetadataProviderReturnsNil contain a significant amount of duplicated setup and execution logic. This can make the tests harder to maintain.
Consider refactoring the common code into a test helper function. This helper could handle setting up the MCP server, creating the toolset, and running the tool. The individual test cases would then just provide the specific metadataProvider and echoToolFunc.
For example, you could have a helper like this:
func runMetadataTest(t *testing.T, provider mcptoolset.MetadataProvider, toolFunc any) {
t.Helper()
clientTransport, serverTransport := mcp.NewInMemoryTransports()
server := mcp.NewServer(&mcp.Implementation{Name: "test_server", Version: "v1.0.0"}, nil)
mcp.AddTool(server, &mcp.Tool{Name: "echo_tool", Description: "echoes input"}, toolFunc)
_, err := server.Connect(t.Context(), serverTransport, nil)
if err != nil {
t.Fatal(err)
}
ts, err := mcptoolset.New(mcptoolset.Config{
Transport: clientTransport,
MetadataProvider: provider,
})
if err != nil {
t.Fatalf("Failed to create MCP tool set: %v", err)
}
invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{})
readonlyCtx := icontext.NewReadonlyContext(invCtx)
tools, err := ts.Tools(readonlyCtx)
if err != nil {
t.Fatalf("Failed to get tools: %v", err)
}
if len(tools) != 1 {
t.Fatalf("Expected 1 tool, got %d", len(tools))
}
fnTool, ok := tools[0].(toolinternal.FunctionTool)
if !ok {
t.Fatal("Tool does not implement FunctionTool interface")
}
toolCtx := toolinternal.NewToolContext(invCtx, "", nil)
_, err = fnTool.Run(toolCtx, map[string]any{})
if err != nil {
t.Fatalf("Failed to run tool: %v", err)
}
}Then your tests would become much simpler. For example, TestMetadataProviderNil would look like:
func TestMetadataProviderNil(t *testing.T) {
echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) {
if req.Params.Meta != nil {
t.Errorf("Expected nil metadata, got %v", req.Params.Meta)
}
return nil, struct{ Message string }{Message: "ok"}, nil
}
runMetadataTest(t, nil, echoToolFunc)
}This would improve readability and maintainability.
There was a problem hiding this comment.
mcp.AddTool(server, &mcp.Tool{Name: "echo_tool", Description: "echoes input"}, toolFunc) is not valid code since In cannot be inferred. Also this way the actual asserts are outside of the tests. I think in this case some duplication helps readability and cognitive load when looking at the various tests.
|
This function is urgently needed. |
2 participants
Add a metadata provider to MCP toolset to support a flexible way setting of metadata in tool calls. Addresses #469 and extends metadata support in ADK to MCP tool calls.