From 92fcbdc041cabf585a804e4513e2435ac0a744ec Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 11:06:37 +0000 Subject: [PATCH 01/20] resolve conflicts --- tool/mcptoolset/set.go | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tool/mcptoolset/set.go b/tool/mcptoolset/set.go index 53e424f55..8b51be8bf 100644 --- a/tool/mcptoolset/set.go +++ b/tool/mcptoolset/set.go @@ -24,6 +24,16 @@ import ( "google.golang.org/adk/tool" ) +// MetadataProvider is a callback function that extracts metadata from the tool context +// to be forwarded to MCP tool calls. The returned map[string]any will be set as the +// Meta field on mcp.CallToolParams. +// +// This allows forwarding request-scoped metadata (e.g., from A2A requests) to downstream +// MCP servers for tracing, authentication, or other purposes. +// +// If the provider returns nil, no metadata is attached to the MCP call. +type MetadataProvider func(ctx tool.Context) map[string]any + // New returns MCP ToolSet. // MCP ToolSet connects to a MCP Server, retrieves MCP Tools into ADK Tools and // passes them to the LLM. @@ -66,6 +76,9 @@ type Config struct { // If ToolFilter is nil, then all tools are returned. // tool.StringPredicate can be convenient if there's a known fixed list of tool names. ToolFilter tool.Predicate + // MetadataProvider is an optional callback that provides metadata to forward + // to MCP tool calls. If nil, no metadata is forwarded. + MetadataProvider MetadataProvider // RequireConfirmation flags whether the tools from this toolset must always ask for user confirmation // before execution. If set to true, the ADK framework will automatically initiate From d8a94b7e03d9fbdb3296bf98a8b3e63cdf101ac8 Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 11:14:03 +0000 Subject: [PATCH 02/20] Server and toolset extension --- server/adka2a/metadata_context.go | 51 +++++ server/adka2a/metadata_context_test.go | 77 +++++++ tool/mcptoolset/providers.go | 159 +++++++++++++ tool/mcptoolset/providers_test.go | 297 +++++++++++++++++++++++++ 4 files changed, 584 insertions(+) create mode 100644 server/adka2a/metadata_context.go create mode 100644 server/adka2a/metadata_context_test.go create mode 100644 tool/mcptoolset/providers.go create mode 100644 tool/mcptoolset/providers_test.go diff --git a/server/adka2a/metadata_context.go b/server/adka2a/metadata_context.go new file mode 100644 index 000000000..4d630364e --- /dev/null +++ b/server/adka2a/metadata_context.go @@ -0,0 +1,51 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package adka2a + +import ( + "context" +) + +type a2aMetadataCtxKey struct{} + +// A2AMetadata contains metadata from A2A requests that can be propagated +// through the context to downstream services like MCP tool servers. +type A2AMetadata struct { + // TaskID is the A2A task identifier. + TaskID string + // ContextID is the A2A context identifier. + ContextID string + // RequestMetadata contains arbitrary metadata from the A2A request. + RequestMetadata map[string]any + // MessageMetadata contains metadata from the A2A message. + MessageMetadata map[string]any +} + +// ContextWithA2AMetadata returns a new context with A2A metadata attached. +// This can be used in BeforeExecuteCallback to attach A2A request metadata +// to the context for downstream propagation to MCP tool calls. +func ContextWithA2AMetadata(ctx context.Context, meta *A2AMetadata) context.Context { + return context.WithValue(ctx, a2aMetadataCtxKey{}, meta) +} + +// A2AMetadataFromContext retrieves A2A metadata from the context. +// Returns nil if no metadata is present. +func A2AMetadataFromContext(ctx context.Context) *A2AMetadata { + meta, ok := ctx.Value(a2aMetadataCtxKey{}).(*A2AMetadata) + if !ok { + return nil + } + return meta +} diff --git a/server/adka2a/metadata_context_test.go b/server/adka2a/metadata_context_test.go new file mode 100644 index 000000000..634479745 --- /dev/null +++ b/server/adka2a/metadata_context_test.go @@ -0,0 +1,77 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package adka2a + +import ( + "context" + "testing" + + "github.com/google/go-cmp/cmp" +) + +func TestA2AMetadataContext(t *testing.T) { + testCases := []struct { + name string + meta *A2AMetadata + }{ + { + name: "full metadata", + meta: &A2AMetadata{ + TaskID: "task-123", + ContextID: "ctx-456", + RequestMetadata: map[string]any{"trace_id": "trace-789"}, + MessageMetadata: map[string]any{"correlation_id": "corr-abc"}, + }, + }, + { + name: "task id only", + meta: &A2AMetadata{ + TaskID: "task-123", + }, + }, + { + name: "empty metadata", + meta: &A2AMetadata{}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + + // Initially no metadata + if got := A2AMetadataFromContext(ctx); got != nil { + t.Errorf("A2AMetadataFromContext() = %v, want nil", got) + } + + // Add metadata + ctx = ContextWithA2AMetadata(ctx, tc.meta) + + // Should retrieve the same metadata + got := A2AMetadataFromContext(ctx) + if diff := cmp.Diff(tc.meta, got); diff != "" { + t.Errorf("A2AMetadataFromContext() mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestA2AMetadataFromContext_NoMetadata(t *testing.T) { + ctx := context.Background() + got := A2AMetadataFromContext(ctx) + if got != nil { + t.Errorf("A2AMetadataFromContext() = %v, want nil", got) + } +} diff --git a/tool/mcptoolset/providers.go b/tool/mcptoolset/providers.go new file mode 100644 index 000000000..dcbd29ae1 --- /dev/null +++ b/tool/mcptoolset/providers.go @@ -0,0 +1,159 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcptoolset + +import ( + "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/tool" +) + +// A2AMetadataProvider creates a MetadataProvider that forwards A2A request metadata +// to MCP tool calls. This is useful when an ADK agent is exposed via A2A and calls +// downstream MCP tools. +// +// The forwarded metadata includes: +// - "a2a:task_id": The A2A task ID (if present) +// - "a2a:context_id": The A2A context ID (if present) +// - Any keys specified in forwardKeys from request/message metadata +// +// If forwardKeys is nil or empty, all request and message metadata keys are forwarded. +// If forwardKeys is non-empty, only the specified keys are forwarded. +// +// Example usage: +// +// // Forward all A2A metadata +// mcptoolset.New(mcptoolset.Config{ +// Transport: transport, +// MetadataProvider: mcptoolset.A2AMetadataProvider(nil), +// }) +// +// // Forward only specific keys +// mcptoolset.New(mcptoolset.Config{ +// Transport: transport, +// MetadataProvider: mcptoolset.A2AMetadataProvider([]string{"trace_id", "correlation_id"}), +// }) +func A2AMetadataProvider(forwardKeys []string) MetadataProvider { + keySet := make(map[string]bool) + for _, k := range forwardKeys { + keySet[k] = true + } + + return func(ctx tool.Context) map[string]any { + a2aMeta := adka2a.A2AMetadataFromContext(ctx) + if a2aMeta == nil { + return nil + } + + result := make(map[string]any) + + // Always include task and context IDs if present + if a2aMeta.TaskID != "" { + result["a2a:task_id"] = a2aMeta.TaskID + } + if a2aMeta.ContextID != "" { + result["a2a:context_id"] = a2aMeta.ContextID + } + + // Forward selected or all metadata keys + forwardMetadata := func(source map[string]any) { + for k, v := range source { + if len(keySet) == 0 || keySet[k] { + result[k] = v + } + } + } + + forwardMetadata(a2aMeta.RequestMetadata) + forwardMetadata(a2aMeta.MessageMetadata) + + if len(result) == 0 { + return nil + } + return result + } +} + +// SessionStateMetadataProvider creates a MetadataProvider that reads metadata +// from session state keys. This is useful for non-A2A scenarios where metadata +// is stored in session state. +// +// The stateKeys map specifies which session state keys to read and how to name +// them in the MCP metadata. For example: +// +// mcptoolset.New(mcptoolset.Config{ +// Transport: transport, +// MetadataProvider: mcptoolset.SessionStateMetadataProvider(map[string]string{ +// "temp:trace_id": "x-trace-id", +// "temp:request_id": "x-request-id", +// }), +// }) +// +// This would read "temp:trace_id" from state and forward it as "x-trace-id" in MCP metadata. +func SessionStateMetadataProvider(stateKeys map[string]string) MetadataProvider { + return func(ctx tool.Context) map[string]any { + if len(stateKeys) == 0 { + return nil + } + + result := make(map[string]any) + state := ctx.ReadonlyState() + + for stateKey, metaKey := range stateKeys { + if val, err := state.Get(stateKey); err == nil { + result[metaKey] = val + } + } + + if len(result) == 0 { + return nil + } + return result + } +} + +// ChainMetadataProviders combines multiple MetadataProviders into one. +// Each provider is called in order, and later providers can override +// keys set by earlier providers. +// +// Example usage: +// +// mcptoolset.New(mcptoolset.Config{ +// Transport: transport, +// MetadataProvider: mcptoolset.ChainMetadataProviders( +// mcptoolset.A2AMetadataProvider(nil), +// mcptoolset.SessionStateMetadataProvider(map[string]string{ +// "temp:custom_field": "custom-field", +// }), +// ), +// }) +func ChainMetadataProviders(providers ...MetadataProvider) MetadataProvider { + return func(ctx tool.Context) map[string]any { + result := make(map[string]any) + for _, p := range providers { + if p == nil { + continue + } + if meta := p(ctx); meta != nil { + for k, v := range meta { + result[k] = v + } + } + } + if len(result) == 0 { + return nil + } + return result + } +} diff --git a/tool/mcptoolset/providers_test.go b/tool/mcptoolset/providers_test.go new file mode 100644 index 000000000..c9df9ba66 --- /dev/null +++ b/tool/mcptoolset/providers_test.go @@ -0,0 +1,297 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcptoolset_test + +import ( + "context" + "iter" + "testing" + + "github.com/google/go-cmp/cmp" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/artifact" + "google.golang.org/adk/memory" + "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/session" + "google.golang.org/adk/tool" + "google.golang.org/adk/tool/mcptoolset" +) + +// mockToolContext implements tool.Context for testing +type mockToolContext struct { + context.Context + state *mockState +} + +var _ tool.Context = (*mockToolContext)(nil) + +func (m *mockToolContext) FunctionCallID() string { return "test-function-call-id" } +func (m *mockToolContext) Actions() *session.EventActions { return nil } +func (m *mockToolContext) SearchMemory(context.Context, string) (*memory.SearchResponse, error) { return nil, nil } +func (m *mockToolContext) UserContent() *genai.Content { return nil } +func (m *mockToolContext) InvocationID() string { return "test-invocation-id" } +func (m *mockToolContext) AgentName() string { return "test-agent" } +func (m *mockToolContext) ReadonlyState() session.ReadonlyState { return m.state } +func (m *mockToolContext) UserID() string { return "test-user" } +func (m *mockToolContext) AppName() string { return "test-app" } +func (m *mockToolContext) SessionID() string { return "test-session" } +func (m *mockToolContext) Branch() string { return "" } +func (m *mockToolContext) Artifacts() agent.Artifacts { return &mockArtifacts{} } +func (m *mockToolContext) State() session.State { return m.state } + +type mockArtifacts struct{} + +func (m *mockArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { + return nil, nil +} +func (m *mockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } +func (m *mockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { return nil, nil } +func (m *mockArtifacts) LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error) { + return nil, nil +} + +type mockState struct { + data map[string]any +} + +func (m *mockState) Get(key string) (any, error) { + if v, ok := m.data[key]; ok { + return v, nil + } + return nil, session.ErrStateKeyNotExist +} + +func (m *mockState) Set(key string, val any) error { + if m.data == nil { + m.data = make(map[string]any) + } + m.data[key] = val + return nil +} + +func (m *mockState) All() iter.Seq2[string, any] { + return func(yield func(string, any) bool) { + for k, v := range m.data { + if !yield(k, v) { + return + } + } + } +} + +func TestA2AMetadataProvider(t *testing.T) { + testCases := []struct { + name string + a2aMeta *adka2a.A2AMetadata + forwardKeys []string + want map[string]any + }{ + { + name: "no a2a metadata in context", + a2aMeta: nil, + want: nil, + }, + { + name: "forward all metadata", + a2aMeta: &adka2a.A2AMetadata{ + TaskID: "task-123", + ContextID: "ctx-456", + RequestMetadata: map[string]any{"trace_id": "trace-789"}, + MessageMetadata: map[string]any{"correlation_id": "corr-abc"}, + }, + forwardKeys: nil, + want: map[string]any{ + "a2a:task_id": "task-123", + "a2a:context_id": "ctx-456", + "trace_id": "trace-789", + "correlation_id": "corr-abc", + }, + }, + { + name: "forward only specific keys", + a2aMeta: &adka2a.A2AMetadata{ + TaskID: "task-123", + ContextID: "ctx-456", + RequestMetadata: map[string]any{"trace_id": "trace-789", "ignored": "value"}, + MessageMetadata: map[string]any{"correlation_id": "corr-abc"}, + }, + forwardKeys: []string{"trace_id"}, + want: map[string]any{ + "a2a:task_id": "task-123", + "a2a:context_id": "ctx-456", + "trace_id": "trace-789", + }, + }, + { + name: "task id only", + a2aMeta: &adka2a.A2AMetadata{ + TaskID: "task-123", + }, + forwardKeys: nil, + want: map[string]any{ + "a2a:task_id": "task-123", + }, + }, + { + name: "empty metadata returns nil", + a2aMeta: &adka2a.A2AMetadata{}, + forwardKeys: nil, + want: nil, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + if tc.a2aMeta != nil { + ctx = adka2a.ContextWithA2AMetadata(ctx, tc.a2aMeta) + } + + mockCtx := &mockToolContext{Context: ctx, state: &mockState{}} + provider := mcptoolset.A2AMetadataProvider(tc.forwardKeys) + got := provider(mockCtx) + + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Errorf("A2AMetadataProvider() mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestSessionStateMetadataProvider(t *testing.T) { + testCases := []struct { + name string + stateData map[string]any + stateKeys map[string]string + want map[string]any + }{ + { + name: "empty state keys", + stateKeys: nil, + want: nil, + }, + { + name: "read from state", + stateData: map[string]any{ + "temp:trace_id": "trace-123", + "temp:request_id": "req-456", + }, + stateKeys: map[string]string{ + "temp:trace_id": "x-trace-id", + "temp:request_id": "x-request-id", + }, + want: map[string]any{ + "x-trace-id": "trace-123", + "x-request-id": "req-456", + }, + }, + { + name: "missing state key ignored", + stateData: map[string]any{ + "temp:trace_id": "trace-123", + }, + stateKeys: map[string]string{ + "temp:trace_id": "x-trace-id", + "temp:missing": "x-missing", + }, + want: map[string]any{ + "x-trace-id": "trace-123", + }, + }, + { + name: "no matching keys returns nil", + stateData: map[string]any{}, + stateKeys: map[string]string{ + "temp:missing": "x-missing", + }, + want: nil, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + mockCtx := &mockToolContext{ + Context: context.Background(), + state: &mockState{data: tc.stateData}, + } + + provider := mcptoolset.SessionStateMetadataProvider(tc.stateKeys) + got := provider(mockCtx) + + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Errorf("SessionStateMetadataProvider() mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestChainMetadataProviders(t *testing.T) { + testCases := []struct { + name string + providers []mcptoolset.MetadataProvider + want map[string]any + }{ + { + name: "no providers", + providers: nil, + want: nil, + }, + { + name: "nil provider in chain", + providers: []mcptoolset.MetadataProvider{ + nil, + func(ctx tool.Context) map[string]any { + return map[string]any{"key": "value"} + }, + }, + want: map[string]any{"key": "value"}, + }, + { + name: "later provider overrides earlier", + providers: []mcptoolset.MetadataProvider{ + func(ctx tool.Context) map[string]any { + return map[string]any{"key1": "first", "key2": "first"} + }, + func(ctx tool.Context) map[string]any { + return map[string]any{"key2": "second", "key3": "second"} + }, + }, + want: map[string]any{"key1": "first", "key2": "second", "key3": "second"}, + }, + { + name: "all nil returns nil", + providers: []mcptoolset.MetadataProvider{ + func(ctx tool.Context) map[string]any { return nil }, + func(ctx tool.Context) map[string]any { return nil }, + }, + want: nil, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + mockCtx := &mockToolContext{Context: context.Background(), state: &mockState{}} + + provider := mcptoolset.ChainMetadataProviders(tc.providers...) + got := provider(mockCtx) + + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Errorf("ChainMetadataProviders() mismatch (-want +got):\n%s", diff) + } + }) + } +} From bf51f59c7145592345f2e633e08cc1e88f3d2cca Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 11:40:38 +0000 Subject: [PATCH 03/20] Example --- examples/a2a_mcp_metadata/main.go | 258 ++++++++++++++++++++++++++++++ 1 file changed, 258 insertions(+) create mode 100644 examples/a2a_mcp_metadata/main.go diff --git a/examples/a2a_mcp_metadata/main.go b/examples/a2a_mcp_metadata/main.go new file mode 100644 index 000000000..0b6215f62 --- /dev/null +++ b/examples/a2a_mcp_metadata/main.go @@ -0,0 +1,258 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// This example demonstrates forwarding A2A request metadata to MCP tool calls. +// +// This example shows how to: +// 1. Create an A2A server that exposes an ADK agent with MCP tools +// 2. Use BeforeExecuteCallback to attach A2A metadata to the context +// 3. Configure the MCP toolset to forward metadata to MCP tool calls +// 4. Access the forwarded metadata in MCP tool handlers +// + +// ┌─────────────┐ A2A Request ┌─────────────────┐ MCP CallTool ┌────────────────┐ +// │ A2A Client │ ──────────────────▶ │ ADK Agent │ ─────────────────▶ │ MCP Server │ +// │ │ (with metadata) │ (A2A + MCP) │ (with metadata) │ (echo tool) │ +// └─────────────┘ └─────────────────┘ └────────────────┘ +// │ +// │ BeforeExecuteCallback +// │ attaches A2A metadata +// │ to context +// ▼ +// ┌─────────────────┐ +// │ A2AMetadata in │ +// │ context.Context │ +// └─────────────────┘ +// │ +// │ MetadataProvider +// │ extracts metadata +// │ from context +// ▼ +// ┌─────────────────┐ +// │ mcp.CallTool- │ +// │ Params.Meta │ +// └─────────────────┘ + +package main + +import ( + "context" + "fmt" + "log" + "net" + "net/http" + "net/url" + "os" + "os/signal" + + "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/a2asrv" + "github.com/modelcontextprotocol/go-sdk/mcp" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/agent/llmagent" + "google.golang.org/adk/agent/remoteagent" + "google.golang.org/adk/cmd/launcher" + "google.golang.org/adk/cmd/launcher/full" + "google.golang.org/adk/model/gemini" + "google.golang.org/adk/runner" + "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/session" + "google.golang.org/adk/tool" + "google.golang.org/adk/tool/mcptoolset" +) + +// EchoInput defines the input schema for the echo tool. +type EchoInput struct { + Message string `json:"message" jsonschema:"The message to echo back"` +} + +// EchoOutput defines the output schema for the echo tool. +type EchoOutput struct { + Echo string `json:"echo" jsonschema:"The echoed message"` + Metadata map[string]any `json:"metadata" jsonschema:"Metadata received from the request"` +} + +// EchoWithMetadata is an MCP tool that echoes back the input message along with +// any metadata that was forwarded from the A2A request. +func EchoWithMetadata(ctx context.Context, req *mcp.CallToolRequest, input EchoInput) (*mcp.CallToolResult, EchoOutput, error) { + // The metadata forwarded from A2A is available in req.Params.Meta + metadata := make(map[string]any) + if req.Params.Meta != nil { + metadata = req.Params.Meta + } + + // Log the received metadata for demonstration + log.Printf("[MCP Tool] Received metadata: %v", metadata) + + return nil, EchoOutput{ + Echo: fmt.Sprintf("You said: %s", input.Message), + Metadata: metadata, + }, nil +} + +// createMCPTransport creates an in-memory MCP server with the echo tool. +func createMCPTransport(ctx context.Context) mcp.Transport { + clientTransport, serverTransport := mcp.NewInMemoryTransports() + + server := mcp.NewServer(&mcp.Implementation{ + Name: "metadata_demo_server", + Version: "v1.0.0", + }, nil) + + mcp.AddTool(server, &mcp.Tool{ + Name: "echo_with_metadata", + Description: "Echoes back the message along with any metadata from the request. Use this to verify metadata forwarding.", + }, EchoWithMetadata) + + _, err := server.Connect(ctx, serverTransport, nil) + if err != nil { + log.Fatalf("Failed to connect MCP server: %v", err) + } + + return clientTransport +} + +// createAgent creates an LLM agent with MCP tools configured to forward A2A metadata. +func createAgent(ctx context.Context) agent.Agent { + model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + APIKey: os.Getenv("GOOGLE_API_KEY"), + }) + if err != nil { + log.Fatalf("Failed to create model: %v", err) + } + + // Create MCP toolset with metadata forwarding enabled + mcpToolSet, err := mcptoolset.New(mcptoolset.Config{ + Transport: createMCPTransport(ctx), + // MetadataProvider extracts A2A metadata from the context and forwards it to MCP tools. + // A2AMetadataProvider(nil) forwards all metadata fields. + // You can also specify specific keys to forward: A2AMetadataProvider([]string{"trace_id"}) + MetadataProvider: mcptoolset.A2AMetadataProvider(nil), + }) + if err != nil { + log.Fatalf("Failed to create MCP tool set: %v", err) + } + + a, err := llmagent.New(llmagent.Config{ + Name: "metadata_demo_agent", + Model: model, + Description: "An agent that demonstrates A2A to MCP metadata forwarding.", + Instruction: `You are a helpful assistant that can echo messages back to users. +When the user asks you to echo something or test metadata, use the echo_with_metadata tool. +The tool will show any metadata that was forwarded from the A2A request.`, + Toolsets: []tool.Toolset{mcpToolSet}, + }) + if err != nil { + log.Fatalf("Failed to create agent: %v", err) + } + + return a +} + +// startA2AServer starts an HTTP server exposing the agent via A2A protocol. +// It uses BeforeExecuteCallback to attach A2A metadata to the context. +func startA2AServer(ctx context.Context) string { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + log.Fatalf("Failed to bind to a port: %v", err) + } + + baseURL := &url.URL{Scheme: "http", Host: listener.Addr().String()} + log.Printf("Starting A2A server on %s", baseURL.String()) + + go func() { + a := createAgent(ctx) + agentPath := "/invoke" + + agentCard := &a2a.AgentCard{ + Name: a.Name(), + Skills: adka2a.BuildAgentSkills(a), + PreferredTransport: a2a.TransportProtocolJSONRPC, + URL: baseURL.JoinPath(agentPath).String(), + Capabilities: a2a.AgentCapabilities{Streaming: true}, + } + + executor := adka2a.NewExecutor(adka2a.ExecutorConfig{ + RunnerConfig: runner.Config{ + AppName: a.Name(), + Agent: a, + SessionService: session.InMemoryService(), + }, + // BeforeExecuteCallback is called before each agent execution. + // Here we extract A2A request metadata and attach it to the context + // so it can be forwarded to MCP tool calls. + BeforeExecuteCallback: func(ctx context.Context, reqCtx *a2asrv.RequestContext) (context.Context, error) { + log.Printf("[A2A Server] Received request with TaskID: %s, ContextID: %s", reqCtx.TaskID, reqCtx.ContextID) + + // Extract metadata from the A2A request + meta := &adka2a.A2AMetadata{ + TaskID: string(reqCtx.TaskID), + ContextID: reqCtx.ContextID, + } + + // Include reqest-level metadata + if reqCtx.Metadata != nil { + meta.RequestMetadata = reqCtx.Metadata + log.Printf("[A2A Server] Request metadata: %v", reqCtx.Metadata) + } + + // Include message-level metadata + if reqCtx.Message != nil && reqCtx.Message.Metadata != nil { + meta.MessageMetadata = reqCtx.Message.Metadata + log.Printf("[A2A Server] Message metadata: %v", reqCtx.Message.Metadata) + } + + return adka2a.ContextWithA2AMetadata(ctx, meta), nil + }, + }) + + mux := http.NewServeMux() + mux.Handle(a2asrv.WellKnownAgentCardPath, a2asrv.NewStaticAgentCardHandler(agentCard)) + mux.Handle(agentPath, a2asrv.NewJSONRPCHandler(a2asrv.NewHandler(executor))) + + if err := http.Serve(listener, mux); err != nil { + log.Printf("A2A server stopped: %v", err) + } + }() + + return baseURL.String() +} + +func main() { + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + + // Start A2A server with the agent + a2aServerAddress := startA2AServer(ctx) + + // Create a remote agent that connects to the A2A server + remoteAgent, err := remoteagent.NewA2A(remoteagent.A2AConfig{ + Name: "Remote Metadata Demo Agent", + AgentCardSource: a2aServerAddress, + }) + if err != nil { + log.Fatalf("Failed to create remote agent: %v", err) + } + + config := &launcher.Config{ + AgentLoader: agent.NewSingleLoader(remoteAgent), + } + + l := full.NewLauncher() + if err = l.Execute(ctx, config, os.Args[1:]); err != nil { + log.Fatalf("Run failed: %v\n\n%s", err, l.CommandLineSyntax()) + } +} From 9c14fb2e04f6cd5c3d366c2f7aa50d2c5a4517a0 Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 12:41:47 +0000 Subject: [PATCH 04/20] Remove protocol interdependencies --- examples/a2a_mcp_metadata/main.go | 2 +- tool/mcptoolset/providers.go | 69 +------------------ tool/mcptoolset/providers_test.go | 108 ++++-------------------------- 3 files changed, 16 insertions(+), 163 deletions(-) diff --git a/examples/a2a_mcp_metadata/main.go b/examples/a2a_mcp_metadata/main.go index 0b6215f62..9e2f53c78 100644 --- a/examples/a2a_mcp_metadata/main.go +++ b/examples/a2a_mcp_metadata/main.go @@ -140,7 +140,7 @@ func createAgent(ctx context.Context) agent.Agent { // MetadataProvider extracts A2A metadata from the context and forwards it to MCP tools. // A2AMetadataProvider(nil) forwards all metadata fields. // You can also specify specific keys to forward: A2AMetadataProvider([]string{"trace_id"}) - MetadataProvider: mcptoolset.A2AMetadataProvider(nil), + MetadataProvider: adka2a.A2AMetadataProvider(nil), }) if err != nil { log.Fatalf("Failed to create MCP tool set: %v", err) diff --git a/tool/mcptoolset/providers.go b/tool/mcptoolset/providers.go index dcbd29ae1..ce4c03ebf 100644 --- a/tool/mcptoolset/providers.go +++ b/tool/mcptoolset/providers.go @@ -15,76 +15,9 @@ package mcptoolset import ( - "google.golang.org/adk/server/adka2a" "google.golang.org/adk/tool" ) -// A2AMetadataProvider creates a MetadataProvider that forwards A2A request metadata -// to MCP tool calls. This is useful when an ADK agent is exposed via A2A and calls -// downstream MCP tools. -// -// The forwarded metadata includes: -// - "a2a:task_id": The A2A task ID (if present) -// - "a2a:context_id": The A2A context ID (if present) -// - Any keys specified in forwardKeys from request/message metadata -// -// If forwardKeys is nil or empty, all request and message metadata keys are forwarded. -// If forwardKeys is non-empty, only the specified keys are forwarded. -// -// Example usage: -// -// // Forward all A2A metadata -// mcptoolset.New(mcptoolset.Config{ -// Transport: transport, -// MetadataProvider: mcptoolset.A2AMetadataProvider(nil), -// }) -// -// // Forward only specific keys -// mcptoolset.New(mcptoolset.Config{ -// Transport: transport, -// MetadataProvider: mcptoolset.A2AMetadataProvider([]string{"trace_id", "correlation_id"}), -// }) -func A2AMetadataProvider(forwardKeys []string) MetadataProvider { - keySet := make(map[string]bool) - for _, k := range forwardKeys { - keySet[k] = true - } - - return func(ctx tool.Context) map[string]any { - a2aMeta := adka2a.A2AMetadataFromContext(ctx) - if a2aMeta == nil { - return nil - } - - result := make(map[string]any) - - // Always include task and context IDs if present - if a2aMeta.TaskID != "" { - result["a2a:task_id"] = a2aMeta.TaskID - } - if a2aMeta.ContextID != "" { - result["a2a:context_id"] = a2aMeta.ContextID - } - - // Forward selected or all metadata keys - forwardMetadata := func(source map[string]any) { - for k, v := range source { - if len(keySet) == 0 || keySet[k] { - result[k] = v - } - } - } - - forwardMetadata(a2aMeta.RequestMetadata) - forwardMetadata(a2aMeta.MessageMetadata) - - if len(result) == 0 { - return nil - } - return result - } -} - // SessionStateMetadataProvider creates a MetadataProvider that reads metadata // from session state keys. This is useful for non-A2A scenarios where metadata // is stored in session state. @@ -132,7 +65,7 @@ func SessionStateMetadataProvider(stateKeys map[string]string) MetadataProvider // mcptoolset.New(mcptoolset.Config{ // Transport: transport, // MetadataProvider: mcptoolset.ChainMetadataProviders( -// mcptoolset.A2AMetadataProvider(nil), +// adka2a.A2AMetadataProvider(nil), // mcptoolset.SessionStateMetadataProvider(map[string]string{ // "temp:custom_field": "custom-field", // }), diff --git a/tool/mcptoolset/providers_test.go b/tool/mcptoolset/providers_test.go index c9df9ba66..671080464 100644 --- a/tool/mcptoolset/providers_test.go +++ b/tool/mcptoolset/providers_test.go @@ -25,7 +25,6 @@ import ( "google.golang.org/adk/agent" "google.golang.org/adk/artifact" "google.golang.org/adk/memory" - "google.golang.org/adk/server/adka2a" "google.golang.org/adk/session" "google.golang.org/adk/tool" "google.golang.org/adk/tool/mcptoolset" @@ -39,27 +38,27 @@ type mockToolContext struct { var _ tool.Context = (*mockToolContext)(nil) -func (m *mockToolContext) FunctionCallID() string { return "test-function-call-id" } -func (m *mockToolContext) Actions() *session.EventActions { return nil } +func (m *mockToolContext) FunctionCallID() string { return "test-function-call-id" } +func (m *mockToolContext) Actions() *session.EventActions { return nil } func (m *mockToolContext) SearchMemory(context.Context, string) (*memory.SearchResponse, error) { return nil, nil } -func (m *mockToolContext) UserContent() *genai.Content { return nil } -func (m *mockToolContext) InvocationID() string { return "test-invocation-id" } -func (m *mockToolContext) AgentName() string { return "test-agent" } -func (m *mockToolContext) ReadonlyState() session.ReadonlyState { return m.state } -func (m *mockToolContext) UserID() string { return "test-user" } -func (m *mockToolContext) AppName() string { return "test-app" } -func (m *mockToolContext) SessionID() string { return "test-session" } -func (m *mockToolContext) Branch() string { return "" } -func (m *mockToolContext) Artifacts() agent.Artifacts { return &mockArtifacts{} } -func (m *mockToolContext) State() session.State { return m.state } +func (m *mockToolContext) UserContent() *genai.Content { return nil } +func (m *mockToolContext) InvocationID() string { return "test-invocation-id" } +func (m *mockToolContext) AgentName() string { return "test-agent" } +func (m *mockToolContext) ReadonlyState() session.ReadonlyState { return m.state } +func (m *mockToolContext) UserID() string { return "test-user" } +func (m *mockToolContext) AppName() string { return "test-app" } +func (m *mockToolContext) SessionID() string { return "test-session" } +func (m *mockToolContext) Branch() string { return "" } +func (m *mockToolContext) Artifacts() agent.Artifacts { return &mockArtifacts{} } +func (m *mockToolContext) State() session.State { return m.state } type mockArtifacts struct{} func (m *mockArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { return nil, nil } -func (m *mockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } -func (m *mockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { return nil, nil } +func (m *mockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } +func (m *mockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { return nil, nil } func (m *mockArtifacts) LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error) { return nil, nil } @@ -93,85 +92,6 @@ func (m *mockState) All() iter.Seq2[string, any] { } } -func TestA2AMetadataProvider(t *testing.T) { - testCases := []struct { - name string - a2aMeta *adka2a.A2AMetadata - forwardKeys []string - want map[string]any - }{ - { - name: "no a2a metadata in context", - a2aMeta: nil, - want: nil, - }, - { - name: "forward all metadata", - a2aMeta: &adka2a.A2AMetadata{ - TaskID: "task-123", - ContextID: "ctx-456", - RequestMetadata: map[string]any{"trace_id": "trace-789"}, - MessageMetadata: map[string]any{"correlation_id": "corr-abc"}, - }, - forwardKeys: nil, - want: map[string]any{ - "a2a:task_id": "task-123", - "a2a:context_id": "ctx-456", - "trace_id": "trace-789", - "correlation_id": "corr-abc", - }, - }, - { - name: "forward only specific keys", - a2aMeta: &adka2a.A2AMetadata{ - TaskID: "task-123", - ContextID: "ctx-456", - RequestMetadata: map[string]any{"trace_id": "trace-789", "ignored": "value"}, - MessageMetadata: map[string]any{"correlation_id": "corr-abc"}, - }, - forwardKeys: []string{"trace_id"}, - want: map[string]any{ - "a2a:task_id": "task-123", - "a2a:context_id": "ctx-456", - "trace_id": "trace-789", - }, - }, - { - name: "task id only", - a2aMeta: &adka2a.A2AMetadata{ - TaskID: "task-123", - }, - forwardKeys: nil, - want: map[string]any{ - "a2a:task_id": "task-123", - }, - }, - { - name: "empty metadata returns nil", - a2aMeta: &adka2a.A2AMetadata{}, - forwardKeys: nil, - want: nil, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - ctx := context.Background() - if tc.a2aMeta != nil { - ctx = adka2a.ContextWithA2AMetadata(ctx, tc.a2aMeta) - } - - mockCtx := &mockToolContext{Context: ctx, state: &mockState{}} - provider := mcptoolset.A2AMetadataProvider(tc.forwardKeys) - got := provider(mockCtx) - - if diff := cmp.Diff(tc.want, got); diff != "" { - t.Errorf("A2AMetadataProvider() mismatch (-want +got):\n%s", diff) - } - }) - } -} - func TestSessionStateMetadataProvider(t *testing.T) { testCases := []struct { name string From 9acadf46412cfee61654688c09e1e7eff56dc1de Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 12:44:02 +0000 Subject: [PATCH 05/20] Metadata provider --- server/adka2a/metadata_provider.go | 83 ++++++++++++ server/adka2a/metadata_provider_test.go | 172 ++++++++++++++++++++++++ 2 files changed, 255 insertions(+) create mode 100644 server/adka2a/metadata_provider.go create mode 100644 server/adka2a/metadata_provider_test.go diff --git a/server/adka2a/metadata_provider.go b/server/adka2a/metadata_provider.go new file mode 100644 index 000000000..73dd255f9 --- /dev/null +++ b/server/adka2a/metadata_provider.go @@ -0,0 +1,83 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package adka2a + +import ( + "google.golang.org/adk/tool" +) + +// A2AMetadataProvider creates a function that extracts A2A request metadata +// from the context. The returned function can be used as mcptoolset.MetadataProvider +// to forward A2A metadata to MCP tool calls. +// +// The forwarded metadata includes: +// - "a2a:task_id": The A2A task ID (if present) +// - "a2a:context_id": The A2A context ID (if present) +// - Any keys specified in forwardKeys from request/message metadata +// +// If forwardKeys is nil or empty, all request and message metadata keys are forwarded. +// If forwardKeys is non-empty, only the specified keys are forwarded. +// +// Example usage: +// +// mcptoolset.New(mcptoolset.Config{ +// Transport: transport, +// MetadataProvider: adka2a.A2AMetadataProvider(nil), // forward all +// }) +// +// mcptoolset.New(mcptoolset.Config{ +// Transport: transport, +// MetadataProvider: adka2a.A2AMetadataProvider([]string{"trace_id"}), // forward specific keys +// }) +func A2AMetadataProvider(forwardKeys []string) func(tool.Context) map[string]any { + keySet := make(map[string]bool) + for _, k := range forwardKeys { + keySet[k] = true + } + + return func(ctx tool.Context) map[string]any { + a2aMeta := A2AMetadataFromContext(ctx) + if a2aMeta == nil { + return nil + } + + result := make(map[string]any) + + // Always include task and context IDs if present + if a2aMeta.TaskID != "" { + result["a2a:task_id"] = a2aMeta.TaskID + } + if a2aMeta.ContextID != "" { + result["a2a:context_id"] = a2aMeta.ContextID + } + + // Forward selected or all metadata keys + forwardMetadata := func(source map[string]any) { + for k, v := range source { + if len(keySet) == 0 || keySet[k] { + result[k] = v + } + } + } + + forwardMetadata(a2aMeta.RequestMetadata) + forwardMetadata(a2aMeta.MessageMetadata) + + if len(result) == 0 { + return nil + } + return result + } +} diff --git a/server/adka2a/metadata_provider_test.go b/server/adka2a/metadata_provider_test.go new file mode 100644 index 000000000..480b548e4 --- /dev/null +++ b/server/adka2a/metadata_provider_test.go @@ -0,0 +1,172 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package adka2a_test + +import ( + "context" + "iter" + "testing" + + "github.com/google/go-cmp/cmp" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/artifact" + "google.golang.org/adk/memory" + "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/session" + "google.golang.org/adk/tool" +) + +// mockToolContext implements tool.Context for testing +type mockToolContext struct { + context.Context + state *mockState +} + +var _ tool.Context = (*mockToolContext)(nil) + +func (m *mockToolContext) FunctionCallID() string { return "test-function-call-id" } +func (m *mockToolContext) Actions() *session.EventActions { return nil } +func (m *mockToolContext) SearchMemory(context.Context, string) (*memory.SearchResponse, error) { return nil, nil } +func (m *mockToolContext) UserContent() *genai.Content { return nil } +func (m *mockToolContext) InvocationID() string { return "test-invocation-id" } +func (m *mockToolContext) AgentName() string { return "test-agent" } +func (m *mockToolContext) ReadonlyState() session.ReadonlyState { return m.state } +func (m *mockToolContext) UserID() string { return "test-user" } +func (m *mockToolContext) AppName() string { return "test-app" } +func (m *mockToolContext) SessionID() string { return "test-session" } +func (m *mockToolContext) Branch() string { return "" } +func (m *mockToolContext) Artifacts() agent.Artifacts { return &mockArtifacts{} } +func (m *mockToolContext) State() session.State { return m.state } + +type mockArtifacts struct{} + +func (m *mockArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { + return nil, nil +} +func (m *mockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } +func (m *mockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { return nil, nil } +func (m *mockArtifacts) LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error) { + return nil, nil +} + +type mockState struct { + data map[string]any +} + +func (m *mockState) Get(key string) (any, error) { + if v, ok := m.data[key]; ok { + return v, nil + } + return nil, session.ErrStateKeyNotExist +} + +func (m *mockState) Set(key string, val any) error { + if m.data == nil { + m.data = make(map[string]any) + } + m.data[key] = val + return nil +} + +func (m *mockState) All() iter.Seq2[string, any] { + return func(yield func(string, any) bool) { + for k, v := range m.data { + if !yield(k, v) { + return + } + } + } +} + +func TestA2AMetadataProvider(t *testing.T) { + testCases := []struct { + name string + a2aMeta *adka2a.A2AMetadata + forwardKeys []string + want map[string]any + }{ + { + name: "no a2a metadata in context", + a2aMeta: nil, + want: nil, + }, + { + name: "forward all metadata", + a2aMeta: &adka2a.A2AMetadata{ + TaskID: "task-123", + ContextID: "ctx-456", + RequestMetadata: map[string]any{"trace_id": "trace-789"}, + MessageMetadata: map[string]any{"correlation_id": "corr-abc"}, + }, + forwardKeys: nil, + want: map[string]any{ + "a2a:task_id": "task-123", + "a2a:context_id": "ctx-456", + "trace_id": "trace-789", + "correlation_id": "corr-abc", + }, + }, + { + name: "forward only specific keys", + a2aMeta: &adka2a.A2AMetadata{ + TaskID: "task-123", + ContextID: "ctx-456", + RequestMetadata: map[string]any{"trace_id": "trace-789", "ignored": "value"}, + MessageMetadata: map[string]any{"correlation_id": "corr-abc"}, + }, + forwardKeys: []string{"trace_id"}, + want: map[string]any{ + "a2a:task_id": "task-123", + "a2a:context_id": "ctx-456", + "trace_id": "trace-789", + }, + }, + { + name: "task id only", + a2aMeta: &adka2a.A2AMetadata{ + TaskID: "task-123", + }, + forwardKeys: nil, + want: map[string]any{ + "a2a:task_id": "task-123", + }, + }, + { + name: "empty metadata returns nil", + a2aMeta: &adka2a.A2AMetadata{}, + forwardKeys: nil, + want: nil, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + if tc.a2aMeta != nil { + ctx = adka2a.ContextWithA2AMetadata(ctx, tc.a2aMeta) + } + + mockCtx := &mockToolContext{Context: ctx, state: &mockState{}} + provider := adka2a.A2AMetadataProvider(tc.forwardKeys) + got := provider(mockCtx) + + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Errorf("A2AMetadataProvider() mismatch (-want +got):\n%s", diff) + } + }) + } +} From b556c629c47e3c989abb1ae0a87d7b0abe0b2e30 Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 16:16:30 +0000 Subject: [PATCH 06/20] resolve conflicts --- server/adka2a/metadata_context.go | 2 +- server/adka2a/metadata_provider.go | 16 +-- tool/mcptoolset/providers.go | 92 ------------ tool/mcptoolset/providers_test.go | 217 ----------------------------- 4 files changed, 3 insertions(+), 324 deletions(-) delete mode 100644 tool/mcptoolset/providers.go delete mode 100644 tool/mcptoolset/providers_test.go diff --git a/server/adka2a/metadata_context.go b/server/adka2a/metadata_context.go index 4d630364e..70b44bd3b 100644 --- a/server/adka2a/metadata_context.go +++ b/server/adka2a/metadata_context.go @@ -21,7 +21,7 @@ import ( type a2aMetadataCtxKey struct{} // A2AMetadata contains metadata from A2A requests that can be propagated -// through the context to downstream services like MCP tool servers. +// through the context to downstream services. type A2AMetadata struct { // TaskID is the A2A task identifier. TaskID string diff --git a/server/adka2a/metadata_provider.go b/server/adka2a/metadata_provider.go index 73dd255f9..b706f62c2 100644 --- a/server/adka2a/metadata_provider.go +++ b/server/adka2a/metadata_provider.go @@ -29,18 +29,6 @@ import ( // // If forwardKeys is nil or empty, all request and message metadata keys are forwarded. // If forwardKeys is non-empty, only the specified keys are forwarded. -// -// Example usage: -// -// mcptoolset.New(mcptoolset.Config{ -// Transport: transport, -// MetadataProvider: adka2a.A2AMetadataProvider(nil), // forward all -// }) -// -// mcptoolset.New(mcptoolset.Config{ -// Transport: transport, -// MetadataProvider: adka2a.A2AMetadataProvider([]string{"trace_id"}), // forward specific keys -// }) func A2AMetadataProvider(forwardKeys []string) func(tool.Context) map[string]any { keySet := make(map[string]bool) for _, k := range forwardKeys { @@ -57,10 +45,10 @@ func A2AMetadataProvider(forwardKeys []string) func(tool.Context) map[string]any // Always include task and context IDs if present if a2aMeta.TaskID != "" { - result["a2a:task_id"] = a2aMeta.TaskID + result["task_id"] = a2aMeta.TaskID } if a2aMeta.ContextID != "" { - result["a2a:context_id"] = a2aMeta.ContextID + result["context_id"] = a2aMeta.ContextID } // Forward selected or all metadata keys diff --git a/tool/mcptoolset/providers.go b/tool/mcptoolset/providers.go deleted file mode 100644 index ce4c03ebf..000000000 --- a/tool/mcptoolset/providers.go +++ /dev/null @@ -1,92 +0,0 @@ -// Copyright 2025 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package mcptoolset - -import ( - "google.golang.org/adk/tool" -) - -// SessionStateMetadataProvider creates a MetadataProvider that reads metadata -// from session state keys. This is useful for non-A2A scenarios where metadata -// is stored in session state. -// -// The stateKeys map specifies which session state keys to read and how to name -// them in the MCP metadata. For example: -// -// mcptoolset.New(mcptoolset.Config{ -// Transport: transport, -// MetadataProvider: mcptoolset.SessionStateMetadataProvider(map[string]string{ -// "temp:trace_id": "x-trace-id", -// "temp:request_id": "x-request-id", -// }), -// }) -// -// This would read "temp:trace_id" from state and forward it as "x-trace-id" in MCP metadata. -func SessionStateMetadataProvider(stateKeys map[string]string) MetadataProvider { - return func(ctx tool.Context) map[string]any { - if len(stateKeys) == 0 { - return nil - } - - result := make(map[string]any) - state := ctx.ReadonlyState() - - for stateKey, metaKey := range stateKeys { - if val, err := state.Get(stateKey); err == nil { - result[metaKey] = val - } - } - - if len(result) == 0 { - return nil - } - return result - } -} - -// ChainMetadataProviders combines multiple MetadataProviders into one. -// Each provider is called in order, and later providers can override -// keys set by earlier providers. -// -// Example usage: -// -// mcptoolset.New(mcptoolset.Config{ -// Transport: transport, -// MetadataProvider: mcptoolset.ChainMetadataProviders( -// adka2a.A2AMetadataProvider(nil), -// mcptoolset.SessionStateMetadataProvider(map[string]string{ -// "temp:custom_field": "custom-field", -// }), -// ), -// }) -func ChainMetadataProviders(providers ...MetadataProvider) MetadataProvider { - return func(ctx tool.Context) map[string]any { - result := make(map[string]any) - for _, p := range providers { - if p == nil { - continue - } - if meta := p(ctx); meta != nil { - for k, v := range meta { - result[k] = v - } - } - } - if len(result) == 0 { - return nil - } - return result - } -} diff --git a/tool/mcptoolset/providers_test.go b/tool/mcptoolset/providers_test.go deleted file mode 100644 index 671080464..000000000 --- a/tool/mcptoolset/providers_test.go +++ /dev/null @@ -1,217 +0,0 @@ -// Copyright 2025 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package mcptoolset_test - -import ( - "context" - "iter" - "testing" - - "github.com/google/go-cmp/cmp" - "google.golang.org/genai" - - "google.golang.org/adk/agent" - "google.golang.org/adk/artifact" - "google.golang.org/adk/memory" - "google.golang.org/adk/session" - "google.golang.org/adk/tool" - "google.golang.org/adk/tool/mcptoolset" -) - -// mockToolContext implements tool.Context for testing -type mockToolContext struct { - context.Context - state *mockState -} - -var _ tool.Context = (*mockToolContext)(nil) - -func (m *mockToolContext) FunctionCallID() string { return "test-function-call-id" } -func (m *mockToolContext) Actions() *session.EventActions { return nil } -func (m *mockToolContext) SearchMemory(context.Context, string) (*memory.SearchResponse, error) { return nil, nil } -func (m *mockToolContext) UserContent() *genai.Content { return nil } -func (m *mockToolContext) InvocationID() string { return "test-invocation-id" } -func (m *mockToolContext) AgentName() string { return "test-agent" } -func (m *mockToolContext) ReadonlyState() session.ReadonlyState { return m.state } -func (m *mockToolContext) UserID() string { return "test-user" } -func (m *mockToolContext) AppName() string { return "test-app" } -func (m *mockToolContext) SessionID() string { return "test-session" } -func (m *mockToolContext) Branch() string { return "" } -func (m *mockToolContext) Artifacts() agent.Artifacts { return &mockArtifacts{} } -func (m *mockToolContext) State() session.State { return m.state } - -type mockArtifacts struct{} - -func (m *mockArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { - return nil, nil -} -func (m *mockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } -func (m *mockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { return nil, nil } -func (m *mockArtifacts) LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error) { - return nil, nil -} - -type mockState struct { - data map[string]any -} - -func (m *mockState) Get(key string) (any, error) { - if v, ok := m.data[key]; ok { - return v, nil - } - return nil, session.ErrStateKeyNotExist -} - -func (m *mockState) Set(key string, val any) error { - if m.data == nil { - m.data = make(map[string]any) - } - m.data[key] = val - return nil -} - -func (m *mockState) All() iter.Seq2[string, any] { - return func(yield func(string, any) bool) { - for k, v := range m.data { - if !yield(k, v) { - return - } - } - } -} - -func TestSessionStateMetadataProvider(t *testing.T) { - testCases := []struct { - name string - stateData map[string]any - stateKeys map[string]string - want map[string]any - }{ - { - name: "empty state keys", - stateKeys: nil, - want: nil, - }, - { - name: "read from state", - stateData: map[string]any{ - "temp:trace_id": "trace-123", - "temp:request_id": "req-456", - }, - stateKeys: map[string]string{ - "temp:trace_id": "x-trace-id", - "temp:request_id": "x-request-id", - }, - want: map[string]any{ - "x-trace-id": "trace-123", - "x-request-id": "req-456", - }, - }, - { - name: "missing state key ignored", - stateData: map[string]any{ - "temp:trace_id": "trace-123", - }, - stateKeys: map[string]string{ - "temp:trace_id": "x-trace-id", - "temp:missing": "x-missing", - }, - want: map[string]any{ - "x-trace-id": "trace-123", - }, - }, - { - name: "no matching keys returns nil", - stateData: map[string]any{}, - stateKeys: map[string]string{ - "temp:missing": "x-missing", - }, - want: nil, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - mockCtx := &mockToolContext{ - Context: context.Background(), - state: &mockState{data: tc.stateData}, - } - - provider := mcptoolset.SessionStateMetadataProvider(tc.stateKeys) - got := provider(mockCtx) - - if diff := cmp.Diff(tc.want, got); diff != "" { - t.Errorf("SessionStateMetadataProvider() mismatch (-want +got):\n%s", diff) - } - }) - } -} - -func TestChainMetadataProviders(t *testing.T) { - testCases := []struct { - name string - providers []mcptoolset.MetadataProvider - want map[string]any - }{ - { - name: "no providers", - providers: nil, - want: nil, - }, - { - name: "nil provider in chain", - providers: []mcptoolset.MetadataProvider{ - nil, - func(ctx tool.Context) map[string]any { - return map[string]any{"key": "value"} - }, - }, - want: map[string]any{"key": "value"}, - }, - { - name: "later provider overrides earlier", - providers: []mcptoolset.MetadataProvider{ - func(ctx tool.Context) map[string]any { - return map[string]any{"key1": "first", "key2": "first"} - }, - func(ctx tool.Context) map[string]any { - return map[string]any{"key2": "second", "key3": "second"} - }, - }, - want: map[string]any{"key1": "first", "key2": "second", "key3": "second"}, - }, - { - name: "all nil returns nil", - providers: []mcptoolset.MetadataProvider{ - func(ctx tool.Context) map[string]any { return nil }, - func(ctx tool.Context) map[string]any { return nil }, - }, - want: nil, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - mockCtx := &mockToolContext{Context: context.Background(), state: &mockState{}} - - provider := mcptoolset.ChainMetadataProviders(tc.providers...) - got := provider(mockCtx) - - if diff := cmp.Diff(tc.want, got); diff != "" { - t.Errorf("ChainMetadataProviders() mismatch (-want +got):\n%s", diff) - } - }) - } -} From b0fb26953eb04e7a7e6b7a4fa8194557000d4c8b Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 16:16:45 +0000 Subject: [PATCH 07/20] Tidy --- tool/mcptoolset/metadata_provider.go | 92 +++++++++ tool/mcptoolset/metadata_provider_test.go | 217 ++++++++++++++++++++++ 2 files changed, 309 insertions(+) create mode 100644 tool/mcptoolset/metadata_provider.go create mode 100644 tool/mcptoolset/metadata_provider_test.go diff --git a/tool/mcptoolset/metadata_provider.go b/tool/mcptoolset/metadata_provider.go new file mode 100644 index 000000000..ce4c03ebf --- /dev/null +++ b/tool/mcptoolset/metadata_provider.go @@ -0,0 +1,92 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcptoolset + +import ( + "google.golang.org/adk/tool" +) + +// SessionStateMetadataProvider creates a MetadataProvider that reads metadata +// from session state keys. This is useful for non-A2A scenarios where metadata +// is stored in session state. +// +// The stateKeys map specifies which session state keys to read and how to name +// them in the MCP metadata. For example: +// +// mcptoolset.New(mcptoolset.Config{ +// Transport: transport, +// MetadataProvider: mcptoolset.SessionStateMetadataProvider(map[string]string{ +// "temp:trace_id": "x-trace-id", +// "temp:request_id": "x-request-id", +// }), +// }) +// +// This would read "temp:trace_id" from state and forward it as "x-trace-id" in MCP metadata. +func SessionStateMetadataProvider(stateKeys map[string]string) MetadataProvider { + return func(ctx tool.Context) map[string]any { + if len(stateKeys) == 0 { + return nil + } + + result := make(map[string]any) + state := ctx.ReadonlyState() + + for stateKey, metaKey := range stateKeys { + if val, err := state.Get(stateKey); err == nil { + result[metaKey] = val + } + } + + if len(result) == 0 { + return nil + } + return result + } +} + +// ChainMetadataProviders combines multiple MetadataProviders into one. +// Each provider is called in order, and later providers can override +// keys set by earlier providers. +// +// Example usage: +// +// mcptoolset.New(mcptoolset.Config{ +// Transport: transport, +// MetadataProvider: mcptoolset.ChainMetadataProviders( +// adka2a.A2AMetadataProvider(nil), +// mcptoolset.SessionStateMetadataProvider(map[string]string{ +// "temp:custom_field": "custom-field", +// }), +// ), +// }) +func ChainMetadataProviders(providers ...MetadataProvider) MetadataProvider { + return func(ctx tool.Context) map[string]any { + result := make(map[string]any) + for _, p := range providers { + if p == nil { + continue + } + if meta := p(ctx); meta != nil { + for k, v := range meta { + result[k] = v + } + } + } + if len(result) == 0 { + return nil + } + return result + } +} diff --git a/tool/mcptoolset/metadata_provider_test.go b/tool/mcptoolset/metadata_provider_test.go new file mode 100644 index 000000000..671080464 --- /dev/null +++ b/tool/mcptoolset/metadata_provider_test.go @@ -0,0 +1,217 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcptoolset_test + +import ( + "context" + "iter" + "testing" + + "github.com/google/go-cmp/cmp" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/artifact" + "google.golang.org/adk/memory" + "google.golang.org/adk/session" + "google.golang.org/adk/tool" + "google.golang.org/adk/tool/mcptoolset" +) + +// mockToolContext implements tool.Context for testing +type mockToolContext struct { + context.Context + state *mockState +} + +var _ tool.Context = (*mockToolContext)(nil) + +func (m *mockToolContext) FunctionCallID() string { return "test-function-call-id" } +func (m *mockToolContext) Actions() *session.EventActions { return nil } +func (m *mockToolContext) SearchMemory(context.Context, string) (*memory.SearchResponse, error) { return nil, nil } +func (m *mockToolContext) UserContent() *genai.Content { return nil } +func (m *mockToolContext) InvocationID() string { return "test-invocation-id" } +func (m *mockToolContext) AgentName() string { return "test-agent" } +func (m *mockToolContext) ReadonlyState() session.ReadonlyState { return m.state } +func (m *mockToolContext) UserID() string { return "test-user" } +func (m *mockToolContext) AppName() string { return "test-app" } +func (m *mockToolContext) SessionID() string { return "test-session" } +func (m *mockToolContext) Branch() string { return "" } +func (m *mockToolContext) Artifacts() agent.Artifacts { return &mockArtifacts{} } +func (m *mockToolContext) State() session.State { return m.state } + +type mockArtifacts struct{} + +func (m *mockArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { + return nil, nil +} +func (m *mockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } +func (m *mockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { return nil, nil } +func (m *mockArtifacts) LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error) { + return nil, nil +} + +type mockState struct { + data map[string]any +} + +func (m *mockState) Get(key string) (any, error) { + if v, ok := m.data[key]; ok { + return v, nil + } + return nil, session.ErrStateKeyNotExist +} + +func (m *mockState) Set(key string, val any) error { + if m.data == nil { + m.data = make(map[string]any) + } + m.data[key] = val + return nil +} + +func (m *mockState) All() iter.Seq2[string, any] { + return func(yield func(string, any) bool) { + for k, v := range m.data { + if !yield(k, v) { + return + } + } + } +} + +func TestSessionStateMetadataProvider(t *testing.T) { + testCases := []struct { + name string + stateData map[string]any + stateKeys map[string]string + want map[string]any + }{ + { + name: "empty state keys", + stateKeys: nil, + want: nil, + }, + { + name: "read from state", + stateData: map[string]any{ + "temp:trace_id": "trace-123", + "temp:request_id": "req-456", + }, + stateKeys: map[string]string{ + "temp:trace_id": "x-trace-id", + "temp:request_id": "x-request-id", + }, + want: map[string]any{ + "x-trace-id": "trace-123", + "x-request-id": "req-456", + }, + }, + { + name: "missing state key ignored", + stateData: map[string]any{ + "temp:trace_id": "trace-123", + }, + stateKeys: map[string]string{ + "temp:trace_id": "x-trace-id", + "temp:missing": "x-missing", + }, + want: map[string]any{ + "x-trace-id": "trace-123", + }, + }, + { + name: "no matching keys returns nil", + stateData: map[string]any{}, + stateKeys: map[string]string{ + "temp:missing": "x-missing", + }, + want: nil, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + mockCtx := &mockToolContext{ + Context: context.Background(), + state: &mockState{data: tc.stateData}, + } + + provider := mcptoolset.SessionStateMetadataProvider(tc.stateKeys) + got := provider(mockCtx) + + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Errorf("SessionStateMetadataProvider() mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestChainMetadataProviders(t *testing.T) { + testCases := []struct { + name string + providers []mcptoolset.MetadataProvider + want map[string]any + }{ + { + name: "no providers", + providers: nil, + want: nil, + }, + { + name: "nil provider in chain", + providers: []mcptoolset.MetadataProvider{ + nil, + func(ctx tool.Context) map[string]any { + return map[string]any{"key": "value"} + }, + }, + want: map[string]any{"key": "value"}, + }, + { + name: "later provider overrides earlier", + providers: []mcptoolset.MetadataProvider{ + func(ctx tool.Context) map[string]any { + return map[string]any{"key1": "first", "key2": "first"} + }, + func(ctx tool.Context) map[string]any { + return map[string]any{"key2": "second", "key3": "second"} + }, + }, + want: map[string]any{"key1": "first", "key2": "second", "key3": "second"}, + }, + { + name: "all nil returns nil", + providers: []mcptoolset.MetadataProvider{ + func(ctx tool.Context) map[string]any { return nil }, + func(ctx tool.Context) map[string]any { return nil }, + }, + want: nil, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + mockCtx := &mockToolContext{Context: context.Background(), state: &mockState{}} + + provider := mcptoolset.ChainMetadataProviders(tc.providers...) + got := provider(mockCtx) + + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Errorf("ChainMetadataProviders() mismatch (-want +got):\n%s", diff) + } + }) + } +} From 9774bc6ef7fb6c23f25ccbb5a7e41058a4cfb0ef Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 17:40:36 +0000 Subject: [PATCH 08/20] Remove a2a metadata from sdk --- examples/a2a_mcp_metadata/main.go | 70 ++++++++++ server/adka2a/metadata_context.go | 51 ------- server/adka2a/metadata_context_test.go | 77 ----------- server/adka2a/metadata_provider.go | 71 ---------- server/adka2a/metadata_provider_test.go | 172 ------------------------ 5 files changed, 70 insertions(+), 371 deletions(-) delete mode 100644 server/adka2a/metadata_context.go delete mode 100644 server/adka2a/metadata_context_test.go delete mode 100644 server/adka2a/metadata_provider.go delete mode 100644 server/adka2a/metadata_provider_test.go diff --git a/examples/a2a_mcp_metadata/main.go b/examples/a2a_mcp_metadata/main.go index 9e2f53c78..35c79fd16 100644 --- a/examples/a2a_mcp_metadata/main.go +++ b/examples/a2a_mcp_metadata/main.go @@ -74,6 +74,76 @@ import ( "google.golang.org/adk/tool/mcptoolset" ) +type a2aMetadataCtxKey struct{} + +type A2AMetadata struct { + // RequestMetadata contains arbitrary metadata from the A2A request. + RequestMetadata map[string]any + // MessageMetadata contains metadata from the A2A message. + MessageMetadata map[string]any +} + +// ContextWithA2AMetadata returns a new context with A2A metadata attached. +// This can be used in BeforeExecuteCallback to attach A2A request metadata +// to the context for downstream propagation to MCP tool calls. +func ContextWithA2AMetadata(ctx context.Context, meta *A2AMetadata) context.Context { + return context.WithValue(ctx, a2aMetadataCtxKey{}, meta) +} + +// A2AMetadataFromContext retrieves A2A metadata from the context. +// Returns nil if no metadata is present. +func A2AMetadataFromContext(ctx context.Context) *A2AMetadata { + meta, ok := ctx.Value(a2aMetadataCtxKey{}).(*A2AMetadata) + if !ok { + return nil + } + return meta +} + +// A2AMetadataProvider creates a function that extracts A2A request metadata +// from the context. The returned function can be used as mcptoolset.MetadataProvider +// to forward A2A metadata to MCP tool calls. +// +// The forwarded metadata includes: +// - "a2a:task_id": The A2A task ID (if present) +// - "a2a:context_id": The A2A context ID (if present) +// - Any keys specified in forwardKeys from request/message metadata +// +// If forwardKeys is nil or empty, all request and message metadata keys are forwarded. +// If forwardKeys is non-empty, only the specified keys are forwarded. +func A2AMetadataProvider(forwardKeys []string) func(tool.Context) map[string]any { + keySet := make(map[string]bool) + for _, k := range forwardKeys { + keySet[k] = true + } + + return func(ctx tool.Context) map[string]any { + a2aMeta := A2AMetadataFromContext(ctx) + if a2aMeta == nil { + return nil + } + + result := make(map[string]any) + + // Forward selected or all metadata keys + forwardMetadata := func(source map[string]any) { + for k, v := range source { + if len(keySet) == 0 || keySet[k] { + result[k] = v + } + } + } + + forwardMetadata(a2aMeta.RequestMetadata) + forwardMetadata(a2aMeta.MessageMetadata) + + if len(result) == 0 { + return nil + } + return result + } +} + // EchoInput defines the input schema for the echo tool. type EchoInput struct { Message string `json:"message" jsonschema:"The message to echo back"` diff --git a/server/adka2a/metadata_context.go b/server/adka2a/metadata_context.go deleted file mode 100644 index 70b44bd3b..000000000 --- a/server/adka2a/metadata_context.go +++ /dev/null @@ -1,51 +0,0 @@ -// Copyright 2025 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package adka2a - -import ( - "context" -) - -type a2aMetadataCtxKey struct{} - -// A2AMetadata contains metadata from A2A requests that can be propagated -// through the context to downstream services. -type A2AMetadata struct { - // TaskID is the A2A task identifier. - TaskID string - // ContextID is the A2A context identifier. - ContextID string - // RequestMetadata contains arbitrary metadata from the A2A request. - RequestMetadata map[string]any - // MessageMetadata contains metadata from the A2A message. - MessageMetadata map[string]any -} - -// ContextWithA2AMetadata returns a new context with A2A metadata attached. -// This can be used in BeforeExecuteCallback to attach A2A request metadata -// to the context for downstream propagation to MCP tool calls. -func ContextWithA2AMetadata(ctx context.Context, meta *A2AMetadata) context.Context { - return context.WithValue(ctx, a2aMetadataCtxKey{}, meta) -} - -// A2AMetadataFromContext retrieves A2A metadata from the context. -// Returns nil if no metadata is present. -func A2AMetadataFromContext(ctx context.Context) *A2AMetadata { - meta, ok := ctx.Value(a2aMetadataCtxKey{}).(*A2AMetadata) - if !ok { - return nil - } - return meta -} diff --git a/server/adka2a/metadata_context_test.go b/server/adka2a/metadata_context_test.go deleted file mode 100644 index 634479745..000000000 --- a/server/adka2a/metadata_context_test.go +++ /dev/null @@ -1,77 +0,0 @@ -// Copyright 2025 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package adka2a - -import ( - "context" - "testing" - - "github.com/google/go-cmp/cmp" -) - -func TestA2AMetadataContext(t *testing.T) { - testCases := []struct { - name string - meta *A2AMetadata - }{ - { - name: "full metadata", - meta: &A2AMetadata{ - TaskID: "task-123", - ContextID: "ctx-456", - RequestMetadata: map[string]any{"trace_id": "trace-789"}, - MessageMetadata: map[string]any{"correlation_id": "corr-abc"}, - }, - }, - { - name: "task id only", - meta: &A2AMetadata{ - TaskID: "task-123", - }, - }, - { - name: "empty metadata", - meta: &A2AMetadata{}, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - ctx := context.Background() - - // Initially no metadata - if got := A2AMetadataFromContext(ctx); got != nil { - t.Errorf("A2AMetadataFromContext() = %v, want nil", got) - } - - // Add metadata - ctx = ContextWithA2AMetadata(ctx, tc.meta) - - // Should retrieve the same metadata - got := A2AMetadataFromContext(ctx) - if diff := cmp.Diff(tc.meta, got); diff != "" { - t.Errorf("A2AMetadataFromContext() mismatch (-want +got):\n%s", diff) - } - }) - } -} - -func TestA2AMetadataFromContext_NoMetadata(t *testing.T) { - ctx := context.Background() - got := A2AMetadataFromContext(ctx) - if got != nil { - t.Errorf("A2AMetadataFromContext() = %v, want nil", got) - } -} diff --git a/server/adka2a/metadata_provider.go b/server/adka2a/metadata_provider.go deleted file mode 100644 index b706f62c2..000000000 --- a/server/adka2a/metadata_provider.go +++ /dev/null @@ -1,71 +0,0 @@ -// Copyright 2025 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package adka2a - -import ( - "google.golang.org/adk/tool" -) - -// A2AMetadataProvider creates a function that extracts A2A request metadata -// from the context. The returned function can be used as mcptoolset.MetadataProvider -// to forward A2A metadata to MCP tool calls. -// -// The forwarded metadata includes: -// - "a2a:task_id": The A2A task ID (if present) -// - "a2a:context_id": The A2A context ID (if present) -// - Any keys specified in forwardKeys from request/message metadata -// -// If forwardKeys is nil or empty, all request and message metadata keys are forwarded. -// If forwardKeys is non-empty, only the specified keys are forwarded. -func A2AMetadataProvider(forwardKeys []string) func(tool.Context) map[string]any { - keySet := make(map[string]bool) - for _, k := range forwardKeys { - keySet[k] = true - } - - return func(ctx tool.Context) map[string]any { - a2aMeta := A2AMetadataFromContext(ctx) - if a2aMeta == nil { - return nil - } - - result := make(map[string]any) - - // Always include task and context IDs if present - if a2aMeta.TaskID != "" { - result["task_id"] = a2aMeta.TaskID - } - if a2aMeta.ContextID != "" { - result["context_id"] = a2aMeta.ContextID - } - - // Forward selected or all metadata keys - forwardMetadata := func(source map[string]any) { - for k, v := range source { - if len(keySet) == 0 || keySet[k] { - result[k] = v - } - } - } - - forwardMetadata(a2aMeta.RequestMetadata) - forwardMetadata(a2aMeta.MessageMetadata) - - if len(result) == 0 { - return nil - } - return result - } -} diff --git a/server/adka2a/metadata_provider_test.go b/server/adka2a/metadata_provider_test.go deleted file mode 100644 index 480b548e4..000000000 --- a/server/adka2a/metadata_provider_test.go +++ /dev/null @@ -1,172 +0,0 @@ -// Copyright 2025 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package adka2a_test - -import ( - "context" - "iter" - "testing" - - "github.com/google/go-cmp/cmp" - "google.golang.org/genai" - - "google.golang.org/adk/agent" - "google.golang.org/adk/artifact" - "google.golang.org/adk/memory" - "google.golang.org/adk/server/adka2a" - "google.golang.org/adk/session" - "google.golang.org/adk/tool" -) - -// mockToolContext implements tool.Context for testing -type mockToolContext struct { - context.Context - state *mockState -} - -var _ tool.Context = (*mockToolContext)(nil) - -func (m *mockToolContext) FunctionCallID() string { return "test-function-call-id" } -func (m *mockToolContext) Actions() *session.EventActions { return nil } -func (m *mockToolContext) SearchMemory(context.Context, string) (*memory.SearchResponse, error) { return nil, nil } -func (m *mockToolContext) UserContent() *genai.Content { return nil } -func (m *mockToolContext) InvocationID() string { return "test-invocation-id" } -func (m *mockToolContext) AgentName() string { return "test-agent" } -func (m *mockToolContext) ReadonlyState() session.ReadonlyState { return m.state } -func (m *mockToolContext) UserID() string { return "test-user" } -func (m *mockToolContext) AppName() string { return "test-app" } -func (m *mockToolContext) SessionID() string { return "test-session" } -func (m *mockToolContext) Branch() string { return "" } -func (m *mockToolContext) Artifacts() agent.Artifacts { return &mockArtifacts{} } -func (m *mockToolContext) State() session.State { return m.state } - -type mockArtifacts struct{} - -func (m *mockArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { - return nil, nil -} -func (m *mockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } -func (m *mockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { return nil, nil } -func (m *mockArtifacts) LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error) { - return nil, nil -} - -type mockState struct { - data map[string]any -} - -func (m *mockState) Get(key string) (any, error) { - if v, ok := m.data[key]; ok { - return v, nil - } - return nil, session.ErrStateKeyNotExist -} - -func (m *mockState) Set(key string, val any) error { - if m.data == nil { - m.data = make(map[string]any) - } - m.data[key] = val - return nil -} - -func (m *mockState) All() iter.Seq2[string, any] { - return func(yield func(string, any) bool) { - for k, v := range m.data { - if !yield(k, v) { - return - } - } - } -} - -func TestA2AMetadataProvider(t *testing.T) { - testCases := []struct { - name string - a2aMeta *adka2a.A2AMetadata - forwardKeys []string - want map[string]any - }{ - { - name: "no a2a metadata in context", - a2aMeta: nil, - want: nil, - }, - { - name: "forward all metadata", - a2aMeta: &adka2a.A2AMetadata{ - TaskID: "task-123", - ContextID: "ctx-456", - RequestMetadata: map[string]any{"trace_id": "trace-789"}, - MessageMetadata: map[string]any{"correlation_id": "corr-abc"}, - }, - forwardKeys: nil, - want: map[string]any{ - "a2a:task_id": "task-123", - "a2a:context_id": "ctx-456", - "trace_id": "trace-789", - "correlation_id": "corr-abc", - }, - }, - { - name: "forward only specific keys", - a2aMeta: &adka2a.A2AMetadata{ - TaskID: "task-123", - ContextID: "ctx-456", - RequestMetadata: map[string]any{"trace_id": "trace-789", "ignored": "value"}, - MessageMetadata: map[string]any{"correlation_id": "corr-abc"}, - }, - forwardKeys: []string{"trace_id"}, - want: map[string]any{ - "a2a:task_id": "task-123", - "a2a:context_id": "ctx-456", - "trace_id": "trace-789", - }, - }, - { - name: "task id only", - a2aMeta: &adka2a.A2AMetadata{ - TaskID: "task-123", - }, - forwardKeys: nil, - want: map[string]any{ - "a2a:task_id": "task-123", - }, - }, - { - name: "empty metadata returns nil", - a2aMeta: &adka2a.A2AMetadata{}, - forwardKeys: nil, - want: nil, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - ctx := context.Background() - if tc.a2aMeta != nil { - ctx = adka2a.ContextWithA2AMetadata(ctx, tc.a2aMeta) - } - - mockCtx := &mockToolContext{Context: ctx, state: &mockState{}} - provider := adka2a.A2AMetadataProvider(tc.forwardKeys) - got := provider(mockCtx) - - if diff := cmp.Diff(tc.want, got); diff != "" { - t.Errorf("A2AMetadataProvider() mismatch (-want +got):\n%s", diff) - } - }) - } -} From e3b74d36885b9941598c29ff422d83e12e9d83a5 Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 17:42:24 +0000 Subject: [PATCH 09/20] Move skd a2a to app --- examples/a2a_mcp_metadata/main.go | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/examples/a2a_mcp_metadata/main.go b/examples/a2a_mcp_metadata/main.go index 35c79fd16..406562e92 100644 --- a/examples/a2a_mcp_metadata/main.go +++ b/examples/a2a_mcp_metadata/main.go @@ -210,7 +210,7 @@ func createAgent(ctx context.Context) agent.Agent { // MetadataProvider extracts A2A metadata from the context and forwards it to MCP tools. // A2AMetadataProvider(nil) forwards all metadata fields. // You can also specify specific keys to forward: A2AMetadataProvider([]string{"trace_id"}) - MetadataProvider: adka2a.A2AMetadataProvider(nil), + MetadataProvider: A2AMetadataProvider(nil), }) if err != nil { log.Fatalf("Failed to create MCP tool set: %v", err) @@ -268,10 +268,7 @@ func startA2AServer(ctx context.Context) string { log.Printf("[A2A Server] Received request with TaskID: %s, ContextID: %s", reqCtx.TaskID, reqCtx.ContextID) // Extract metadata from the A2A request - meta := &adka2a.A2AMetadata{ - TaskID: string(reqCtx.TaskID), - ContextID: reqCtx.ContextID, - } + meta := &A2AMetadata{} // Include reqest-level metadata if reqCtx.Metadata != nil { @@ -285,7 +282,7 @@ func startA2AServer(ctx context.Context) string { log.Printf("[A2A Server] Message metadata: %v", reqCtx.Message.Metadata) } - return adka2a.ContextWithA2AMetadata(ctx, meta), nil + return ContextWithA2AMetadata(ctx, meta), nil }, }) From c3481d00f3fc280e9505e861c90f231388fa0dff Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 18:13:34 +0000 Subject: [PATCH 10/20] Tidy example --- examples/a2a_mcp_metadata/main.go | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/a2a_mcp_metadata/main.go b/examples/a2a_mcp_metadata/main.go index 406562e92..0cc841d20 100644 --- a/examples/a2a_mcp_metadata/main.go +++ b/examples/a2a_mcp_metadata/main.go @@ -267,7 +267,6 @@ func startA2AServer(ctx context.Context) string { BeforeExecuteCallback: func(ctx context.Context, reqCtx *a2asrv.RequestContext) (context.Context, error) { log.Printf("[A2A Server] Received request with TaskID: %s, ContextID: %s", reqCtx.TaskID, reqCtx.ContextID) - // Extract metadata from the A2A request meta := &A2AMetadata{} // Include reqest-level metadata From 8c9736cb6e5211fa56eaccd286877d029e1a8e18 Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 18:15:20 +0000 Subject: [PATCH 11/20] Remove example --- examples/a2a_mcp_metadata/main.go | 324 ------------------------------ 1 file changed, 324 deletions(-) delete mode 100644 examples/a2a_mcp_metadata/main.go diff --git a/examples/a2a_mcp_metadata/main.go b/examples/a2a_mcp_metadata/main.go deleted file mode 100644 index 0cc841d20..000000000 --- a/examples/a2a_mcp_metadata/main.go +++ /dev/null @@ -1,324 +0,0 @@ -// Copyright 2025 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// This example demonstrates forwarding A2A request metadata to MCP tool calls. -// -// This example shows how to: -// 1. Create an A2A server that exposes an ADK agent with MCP tools -// 2. Use BeforeExecuteCallback to attach A2A metadata to the context -// 3. Configure the MCP toolset to forward metadata to MCP tool calls -// 4. Access the forwarded metadata in MCP tool handlers -// - -// ┌─────────────┐ A2A Request ┌─────────────────┐ MCP CallTool ┌────────────────┐ -// │ A2A Client │ ──────────────────▶ │ ADK Agent │ ─────────────────▶ │ MCP Server │ -// │ │ (with metadata) │ (A2A + MCP) │ (with metadata) │ (echo tool) │ -// └─────────────┘ └─────────────────┘ └────────────────┘ -// │ -// │ BeforeExecuteCallback -// │ attaches A2A metadata -// │ to context -// ▼ -// ┌─────────────────┐ -// │ A2AMetadata in │ -// │ context.Context │ -// └─────────────────┘ -// │ -// │ MetadataProvider -// │ extracts metadata -// │ from context -// ▼ -// ┌─────────────────┐ -// │ mcp.CallTool- │ -// │ Params.Meta │ -// └─────────────────┘ - -package main - -import ( - "context" - "fmt" - "log" - "net" - "net/http" - "net/url" - "os" - "os/signal" - - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2asrv" - "github.com/modelcontextprotocol/go-sdk/mcp" - "google.golang.org/genai" - - "google.golang.org/adk/agent" - "google.golang.org/adk/agent/llmagent" - "google.golang.org/adk/agent/remoteagent" - "google.golang.org/adk/cmd/launcher" - "google.golang.org/adk/cmd/launcher/full" - "google.golang.org/adk/model/gemini" - "google.golang.org/adk/runner" - "google.golang.org/adk/server/adka2a" - "google.golang.org/adk/session" - "google.golang.org/adk/tool" - "google.golang.org/adk/tool/mcptoolset" -) - -type a2aMetadataCtxKey struct{} - -type A2AMetadata struct { - // RequestMetadata contains arbitrary metadata from the A2A request. - RequestMetadata map[string]any - // MessageMetadata contains metadata from the A2A message. - MessageMetadata map[string]any -} - -// ContextWithA2AMetadata returns a new context with A2A metadata attached. -// This can be used in BeforeExecuteCallback to attach A2A request metadata -// to the context for downstream propagation to MCP tool calls. -func ContextWithA2AMetadata(ctx context.Context, meta *A2AMetadata) context.Context { - return context.WithValue(ctx, a2aMetadataCtxKey{}, meta) -} - -// A2AMetadataFromContext retrieves A2A metadata from the context. -// Returns nil if no metadata is present. -func A2AMetadataFromContext(ctx context.Context) *A2AMetadata { - meta, ok := ctx.Value(a2aMetadataCtxKey{}).(*A2AMetadata) - if !ok { - return nil - } - return meta -} - -// A2AMetadataProvider creates a function that extracts A2A request metadata -// from the context. The returned function can be used as mcptoolset.MetadataProvider -// to forward A2A metadata to MCP tool calls. -// -// The forwarded metadata includes: -// - "a2a:task_id": The A2A task ID (if present) -// - "a2a:context_id": The A2A context ID (if present) -// - Any keys specified in forwardKeys from request/message metadata -// -// If forwardKeys is nil or empty, all request and message metadata keys are forwarded. -// If forwardKeys is non-empty, only the specified keys are forwarded. -func A2AMetadataProvider(forwardKeys []string) func(tool.Context) map[string]any { - keySet := make(map[string]bool) - for _, k := range forwardKeys { - keySet[k] = true - } - - return func(ctx tool.Context) map[string]any { - a2aMeta := A2AMetadataFromContext(ctx) - if a2aMeta == nil { - return nil - } - - result := make(map[string]any) - - // Forward selected or all metadata keys - forwardMetadata := func(source map[string]any) { - for k, v := range source { - if len(keySet) == 0 || keySet[k] { - result[k] = v - } - } - } - - forwardMetadata(a2aMeta.RequestMetadata) - forwardMetadata(a2aMeta.MessageMetadata) - - if len(result) == 0 { - return nil - } - return result - } -} - -// EchoInput defines the input schema for the echo tool. -type EchoInput struct { - Message string `json:"message" jsonschema:"The message to echo back"` -} - -// EchoOutput defines the output schema for the echo tool. -type EchoOutput struct { - Echo string `json:"echo" jsonschema:"The echoed message"` - Metadata map[string]any `json:"metadata" jsonschema:"Metadata received from the request"` -} - -// EchoWithMetadata is an MCP tool that echoes back the input message along with -// any metadata that was forwarded from the A2A request. -func EchoWithMetadata(ctx context.Context, req *mcp.CallToolRequest, input EchoInput) (*mcp.CallToolResult, EchoOutput, error) { - // The metadata forwarded from A2A is available in req.Params.Meta - metadata := make(map[string]any) - if req.Params.Meta != nil { - metadata = req.Params.Meta - } - - // Log the received metadata for demonstration - log.Printf("[MCP Tool] Received metadata: %v", metadata) - - return nil, EchoOutput{ - Echo: fmt.Sprintf("You said: %s", input.Message), - Metadata: metadata, - }, nil -} - -// createMCPTransport creates an in-memory MCP server with the echo tool. -func createMCPTransport(ctx context.Context) mcp.Transport { - clientTransport, serverTransport := mcp.NewInMemoryTransports() - - server := mcp.NewServer(&mcp.Implementation{ - Name: "metadata_demo_server", - Version: "v1.0.0", - }, nil) - - mcp.AddTool(server, &mcp.Tool{ - Name: "echo_with_metadata", - Description: "Echoes back the message along with any metadata from the request. Use this to verify metadata forwarding.", - }, EchoWithMetadata) - - _, err := server.Connect(ctx, serverTransport, nil) - if err != nil { - log.Fatalf("Failed to connect MCP server: %v", err) - } - - return clientTransport -} - -// createAgent creates an LLM agent with MCP tools configured to forward A2A metadata. -func createAgent(ctx context.Context) agent.Agent { - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ - APIKey: os.Getenv("GOOGLE_API_KEY"), - }) - if err != nil { - log.Fatalf("Failed to create model: %v", err) - } - - // Create MCP toolset with metadata forwarding enabled - mcpToolSet, err := mcptoolset.New(mcptoolset.Config{ - Transport: createMCPTransport(ctx), - // MetadataProvider extracts A2A metadata from the context and forwards it to MCP tools. - // A2AMetadataProvider(nil) forwards all metadata fields. - // You can also specify specific keys to forward: A2AMetadataProvider([]string{"trace_id"}) - MetadataProvider: A2AMetadataProvider(nil), - }) - if err != nil { - log.Fatalf("Failed to create MCP tool set: %v", err) - } - - a, err := llmagent.New(llmagent.Config{ - Name: "metadata_demo_agent", - Model: model, - Description: "An agent that demonstrates A2A to MCP metadata forwarding.", - Instruction: `You are a helpful assistant that can echo messages back to users. -When the user asks you to echo something or test metadata, use the echo_with_metadata tool. -The tool will show any metadata that was forwarded from the A2A request.`, - Toolsets: []tool.Toolset{mcpToolSet}, - }) - if err != nil { - log.Fatalf("Failed to create agent: %v", err) - } - - return a -} - -// startA2AServer starts an HTTP server exposing the agent via A2A protocol. -// It uses BeforeExecuteCallback to attach A2A metadata to the context. -func startA2AServer(ctx context.Context) string { - listener, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - log.Fatalf("Failed to bind to a port: %v", err) - } - - baseURL := &url.URL{Scheme: "http", Host: listener.Addr().String()} - log.Printf("Starting A2A server on %s", baseURL.String()) - - go func() { - a := createAgent(ctx) - agentPath := "/invoke" - - agentCard := &a2a.AgentCard{ - Name: a.Name(), - Skills: adka2a.BuildAgentSkills(a), - PreferredTransport: a2a.TransportProtocolJSONRPC, - URL: baseURL.JoinPath(agentPath).String(), - Capabilities: a2a.AgentCapabilities{Streaming: true}, - } - - executor := adka2a.NewExecutor(adka2a.ExecutorConfig{ - RunnerConfig: runner.Config{ - AppName: a.Name(), - Agent: a, - SessionService: session.InMemoryService(), - }, - // BeforeExecuteCallback is called before each agent execution. - // Here we extract A2A request metadata and attach it to the context - // so it can be forwarded to MCP tool calls. - BeforeExecuteCallback: func(ctx context.Context, reqCtx *a2asrv.RequestContext) (context.Context, error) { - log.Printf("[A2A Server] Received request with TaskID: %s, ContextID: %s", reqCtx.TaskID, reqCtx.ContextID) - - meta := &A2AMetadata{} - - // Include reqest-level metadata - if reqCtx.Metadata != nil { - meta.RequestMetadata = reqCtx.Metadata - log.Printf("[A2A Server] Request metadata: %v", reqCtx.Metadata) - } - - // Include message-level metadata - if reqCtx.Message != nil && reqCtx.Message.Metadata != nil { - meta.MessageMetadata = reqCtx.Message.Metadata - log.Printf("[A2A Server] Message metadata: %v", reqCtx.Message.Metadata) - } - - return ContextWithA2AMetadata(ctx, meta), nil - }, - }) - - mux := http.NewServeMux() - mux.Handle(a2asrv.WellKnownAgentCardPath, a2asrv.NewStaticAgentCardHandler(agentCard)) - mux.Handle(agentPath, a2asrv.NewJSONRPCHandler(a2asrv.NewHandler(executor))) - - if err := http.Serve(listener, mux); err != nil { - log.Printf("A2A server stopped: %v", err) - } - }() - - return baseURL.String() -} - -func main() { - ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) - defer stop() - - // Start A2A server with the agent - a2aServerAddress := startA2AServer(ctx) - - // Create a remote agent that connects to the A2A server - remoteAgent, err := remoteagent.NewA2A(remoteagent.A2AConfig{ - Name: "Remote Metadata Demo Agent", - AgentCardSource: a2aServerAddress, - }) - if err != nil { - log.Fatalf("Failed to create remote agent: %v", err) - } - - config := &launcher.Config{ - AgentLoader: agent.NewSingleLoader(remoteAgent), - } - - l := full.NewLauncher() - if err = l.Execute(ctx, config, os.Args[1:]); err != nil { - log.Fatalf("Run failed: %v\n\n%s", err, l.CommandLineSyntax()) - } -} From d18f5247a9d9bce9dfaef5c874ed57c29e739313 Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 18:16:53 +0000 Subject: [PATCH 12/20] Headers --- tool/mcptoolset/metadata_provider_test.go | 34 +++++++++++++---------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/tool/mcptoolset/metadata_provider_test.go b/tool/mcptoolset/metadata_provider_test.go index 671080464..bb3f9bc44 100644 --- a/tool/mcptoolset/metadata_provider_test.go +++ b/tool/mcptoolset/metadata_provider_test.go @@ -38,27 +38,31 @@ type mockToolContext struct { var _ tool.Context = (*mockToolContext)(nil) -func (m *mockToolContext) FunctionCallID() string { return "test-function-call-id" } -func (m *mockToolContext) Actions() *session.EventActions { return nil } -func (m *mockToolContext) SearchMemory(context.Context, string) (*memory.SearchResponse, error) { return nil, nil } -func (m *mockToolContext) UserContent() *genai.Content { return nil } -func (m *mockToolContext) InvocationID() string { return "test-invocation-id" } -func (m *mockToolContext) AgentName() string { return "test-agent" } -func (m *mockToolContext) ReadonlyState() session.ReadonlyState { return m.state } -func (m *mockToolContext) UserID() string { return "test-user" } -func (m *mockToolContext) AppName() string { return "test-app" } -func (m *mockToolContext) SessionID() string { return "test-session" } -func (m *mockToolContext) Branch() string { return "" } -func (m *mockToolContext) Artifacts() agent.Artifacts { return &mockArtifacts{} } -func (m *mockToolContext) State() session.State { return m.state } +func (m *mockToolContext) FunctionCallID() string { return "test-function-call-id" } +func (m *mockToolContext) Actions() *session.EventActions { return nil } +func (m *mockToolContext) SearchMemory(context.Context, string) (*memory.SearchResponse, error) { + return nil, nil +} +func (m *mockToolContext) UserContent() *genai.Content { return nil } +func (m *mockToolContext) InvocationID() string { return "test-invocation-id" } +func (m *mockToolContext) AgentName() string { return "test-agent" } +func (m *mockToolContext) ReadonlyState() session.ReadonlyState { return m.state } +func (m *mockToolContext) UserID() string { return "test-user" } +func (m *mockToolContext) AppName() string { return "test-app" } +func (m *mockToolContext) SessionID() string { return "test-session" } +func (m *mockToolContext) Branch() string { return "" } +func (m *mockToolContext) Artifacts() agent.Artifacts { return &mockArtifacts{} } +func (m *mockToolContext) State() session.State { return m.state } type mockArtifacts struct{} func (m *mockArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { return nil, nil } -func (m *mockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } -func (m *mockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { return nil, nil } +func (m *mockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } +func (m *mockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { + return nil, nil +} func (m *mockArtifacts) LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error) { return nil, nil } From d61e5c9b163b51ffae5b6cce94b60150bd111422 Mon Sep 17 00:00:00 2001 From: Philipp Lang Date: Mon, 12 Jan 2026 18:19:50 +0000 Subject: [PATCH 13/20] Dont need providers --- tool/mcptoolset/metadata_provider.go | 92 --------- tool/mcptoolset/metadata_provider_test.go | 221 ---------------------- 2 files changed, 313 deletions(-) delete mode 100644 tool/mcptoolset/metadata_provider.go delete mode 100644 tool/mcptoolset/metadata_provider_test.go diff --git a/tool/mcptoolset/metadata_provider.go b/tool/mcptoolset/metadata_provider.go deleted file mode 100644 index ce4c03ebf..000000000 --- a/tool/mcptoolset/metadata_provider.go +++ /dev/null @@ -1,92 +0,0 @@ -// Copyright 2025 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package mcptoolset - -import ( - "google.golang.org/adk/tool" -) - -// SessionStateMetadataProvider creates a MetadataProvider that reads metadata -// from session state keys. This is useful for non-A2A scenarios where metadata -// is stored in session state. -// -// The stateKeys map specifies which session state keys to read and how to name -// them in the MCP metadata. For example: -// -// mcptoolset.New(mcptoolset.Config{ -// Transport: transport, -// MetadataProvider: mcptoolset.SessionStateMetadataProvider(map[string]string{ -// "temp:trace_id": "x-trace-id", -// "temp:request_id": "x-request-id", -// }), -// }) -// -// This would read "temp:trace_id" from state and forward it as "x-trace-id" in MCP metadata. -func SessionStateMetadataProvider(stateKeys map[string]string) MetadataProvider { - return func(ctx tool.Context) map[string]any { - if len(stateKeys) == 0 { - return nil - } - - result := make(map[string]any) - state := ctx.ReadonlyState() - - for stateKey, metaKey := range stateKeys { - if val, err := state.Get(stateKey); err == nil { - result[metaKey] = val - } - } - - if len(result) == 0 { - return nil - } - return result - } -} - -// ChainMetadataProviders combines multiple MetadataProviders into one. -// Each provider is called in order, and later providers can override -// keys set by earlier providers. -// -// Example usage: -// -// mcptoolset.New(mcptoolset.Config{ -// Transport: transport, -// MetadataProvider: mcptoolset.ChainMetadataProviders( -// adka2a.A2AMetadataProvider(nil), -// mcptoolset.SessionStateMetadataProvider(map[string]string{ -// "temp:custom_field": "custom-field", -// }), -// ), -// }) -func ChainMetadataProviders(providers ...MetadataProvider) MetadataProvider { - return func(ctx tool.Context) map[string]any { - result := make(map[string]any) - for _, p := range providers { - if p == nil { - continue - } - if meta := p(ctx); meta != nil { - for k, v := range meta { - result[k] = v - } - } - } - if len(result) == 0 { - return nil - } - return result - } -} diff --git a/tool/mcptoolset/metadata_provider_test.go b/tool/mcptoolset/metadata_provider_test.go deleted file mode 100644 index bb3f9bc44..000000000 --- a/tool/mcptoolset/metadata_provider_test.go +++ /dev/null @@ -1,221 +0,0 @@ -// Copyright 2025 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package mcptoolset_test - -import ( - "context" - "iter" - "testing" - - "github.com/google/go-cmp/cmp" - "google.golang.org/genai" - - "google.golang.org/adk/agent" - "google.golang.org/adk/artifact" - "google.golang.org/adk/memory" - "google.golang.org/adk/session" - "google.golang.org/adk/tool" - "google.golang.org/adk/tool/mcptoolset" -) - -// mockToolContext implements tool.Context for testing -type mockToolContext struct { - context.Context - state *mockState -} - -var _ tool.Context = (*mockToolContext)(nil) - -func (m *mockToolContext) FunctionCallID() string { return "test-function-call-id" } -func (m *mockToolContext) Actions() *session.EventActions { return nil } -func (m *mockToolContext) SearchMemory(context.Context, string) (*memory.SearchResponse, error) { - return nil, nil -} -func (m *mockToolContext) UserContent() *genai.Content { return nil } -func (m *mockToolContext) InvocationID() string { return "test-invocation-id" } -func (m *mockToolContext) AgentName() string { return "test-agent" } -func (m *mockToolContext) ReadonlyState() session.ReadonlyState { return m.state } -func (m *mockToolContext) UserID() string { return "test-user" } -func (m *mockToolContext) AppName() string { return "test-app" } -func (m *mockToolContext) SessionID() string { return "test-session" } -func (m *mockToolContext) Branch() string { return "" } -func (m *mockToolContext) Artifacts() agent.Artifacts { return &mockArtifacts{} } -func (m *mockToolContext) State() session.State { return m.state } - -type mockArtifacts struct{} - -func (m *mockArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { - return nil, nil -} -func (m *mockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } -func (m *mockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { - return nil, nil -} -func (m *mockArtifacts) LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error) { - return nil, nil -} - -type mockState struct { - data map[string]any -} - -func (m *mockState) Get(key string) (any, error) { - if v, ok := m.data[key]; ok { - return v, nil - } - return nil, session.ErrStateKeyNotExist -} - -func (m *mockState) Set(key string, val any) error { - if m.data == nil { - m.data = make(map[string]any) - } - m.data[key] = val - return nil -} - -func (m *mockState) All() iter.Seq2[string, any] { - return func(yield func(string, any) bool) { - for k, v := range m.data { - if !yield(k, v) { - return - } - } - } -} - -func TestSessionStateMetadataProvider(t *testing.T) { - testCases := []struct { - name string - stateData map[string]any - stateKeys map[string]string - want map[string]any - }{ - { - name: "empty state keys", - stateKeys: nil, - want: nil, - }, - { - name: "read from state", - stateData: map[string]any{ - "temp:trace_id": "trace-123", - "temp:request_id": "req-456", - }, - stateKeys: map[string]string{ - "temp:trace_id": "x-trace-id", - "temp:request_id": "x-request-id", - }, - want: map[string]any{ - "x-trace-id": "trace-123", - "x-request-id": "req-456", - }, - }, - { - name: "missing state key ignored", - stateData: map[string]any{ - "temp:trace_id": "trace-123", - }, - stateKeys: map[string]string{ - "temp:trace_id": "x-trace-id", - "temp:missing": "x-missing", - }, - want: map[string]any{ - "x-trace-id": "trace-123", - }, - }, - { - name: "no matching keys returns nil", - stateData: map[string]any{}, - stateKeys: map[string]string{ - "temp:missing": "x-missing", - }, - want: nil, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - mockCtx := &mockToolContext{ - Context: context.Background(), - state: &mockState{data: tc.stateData}, - } - - provider := mcptoolset.SessionStateMetadataProvider(tc.stateKeys) - got := provider(mockCtx) - - if diff := cmp.Diff(tc.want, got); diff != "" { - t.Errorf("SessionStateMetadataProvider() mismatch (-want +got):\n%s", diff) - } - }) - } -} - -func TestChainMetadataProviders(t *testing.T) { - testCases := []struct { - name string - providers []mcptoolset.MetadataProvider - want map[string]any - }{ - { - name: "no providers", - providers: nil, - want: nil, - }, - { - name: "nil provider in chain", - providers: []mcptoolset.MetadataProvider{ - nil, - func(ctx tool.Context) map[string]any { - return map[string]any{"key": "value"} - }, - }, - want: map[string]any{"key": "value"}, - }, - { - name: "later provider overrides earlier", - providers: []mcptoolset.MetadataProvider{ - func(ctx tool.Context) map[string]any { - return map[string]any{"key1": "first", "key2": "first"} - }, - func(ctx tool.Context) map[string]any { - return map[string]any{"key2": "second", "key3": "second"} - }, - }, - want: map[string]any{"key1": "first", "key2": "second", "key3": "second"}, - }, - { - name: "all nil returns nil", - providers: []mcptoolset.MetadataProvider{ - func(ctx tool.Context) map[string]any { return nil }, - func(ctx tool.Context) map[string]any { return nil }, - }, - want: nil, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - mockCtx := &mockToolContext{Context: context.Background(), state: &mockState{}} - - provider := mcptoolset.ChainMetadataProviders(tc.providers...) - got := provider(mockCtx) - - if diff := cmp.Diff(tc.want, got); diff != "" { - t.Errorf("ChainMetadataProviders() mismatch (-want +got):\n%s", diff) - } - }) - } -} From f5b9a910057e8e0a845158bb2dd1a49a77c1297f Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Sat, 28 Feb 2026 08:58:21 +0800 Subject: [PATCH 14/20] resolve conflicts rebased from main branch --- tool/mcptoolset/set.go | 4 +- tool/mcptoolset/set_test.go | 115 ++++++++++++++++++++++++++++++++++++ tool/mcptoolset/tool.go | 19 ++++-- 3 files changed, 132 insertions(+), 6 deletions(-) diff --git a/tool/mcptoolset/set.go b/tool/mcptoolset/set.go index 8b51be8bf..c4304ced5 100644 --- a/tool/mcptoolset/set.go +++ b/tool/mcptoolset/set.go @@ -62,6 +62,7 @@ func New(cfg Config) (tool.Toolset, error) { toolFilter: cfg.ToolFilter, requireConfirmation: cfg.RequireConfirmation, requireConfirmationProvider: cfg.RequireConfirmationProvider, + metadataProvider: cfg.MetadataProvider, }, nil } @@ -103,6 +104,7 @@ type set struct { toolFilter tool.Predicate requireConfirmation bool requireConfirmationProvider ConfirmationProvider + metadataProvider MetadataProvider } func (*set) Name() string { @@ -126,7 +128,7 @@ func (s *set) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { var adkTools []tool.Tool for _, mcpTool := range mcpTools { - t, err := convertTool(mcpTool, s.mcpClient, s.requireConfirmation, s.requireConfirmationProvider) + t, err := convertTool(mcpTool, s.mcpClient, s.requireConfirmation, s.requireConfirmationProvider, s.metadataProvider) if err != nil { return nil, fmt.Errorf("failed to convert MCP tool %q to adk tool: %w", mcpTool.Name, err) } diff --git a/tool/mcptoolset/set_test.go b/tool/mcptoolset/set_test.go index 0b4c1cffd..c6d40df39 100644 --- a/tool/mcptoolset/set_test.go +++ b/tool/mcptoolset/set_test.go @@ -790,3 +790,118 @@ func TestNewToolSet_RequireConfirmationProvider_Validation(t *testing.T) { }) } } + +func TestMetadataProvider(t *testing.T) { + var receivedMeta map[string]any + var toolCalled bool + echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { + toolCalled = true + receivedMeta = req.Params.Meta + return nil, struct{ Message string }{Message: "ok"}, nil + } + + 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 + } + + result := runMetadataTest(t, metadataProvider, echoToolFunc) + if result == nil { + t.Fatal("Expected non-nil result") + } + + if !toolCalled { + t.Fatal("Tool was not called") + } + + if diff := cmp.Diff(testMetadata, receivedMeta); diff != "" { + t.Errorf("metadata mismatch (-want +got):\n%s", diff) + } +} + +func TestMetadataProviderNil(t *testing.T) { + var toolCalled bool + echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { + toolCalled = true + 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) + if !toolCalled { + t.Fatal("Tool was not called") + } +} + +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 + } + + metadataProvider := func(ctx tool.Context) map[string]any { + return nil + } + + _ = runMetadataTest(t, metadataProvider, echoToolFunc) + + if !metaCalled { + t.Fatal("Tool was not called") + } + + if receivedMeta != nil { + t.Errorf("Expected nil metadata when provider returns nil, got %v", receivedMeta) + } +} + +func runMetadataTest[In, Out any](t *testing.T, provider mcptoolset.MetadataProvider, toolFunc mcp.ToolHandlerFor[In, Out]) map[string]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, nil) + result, err := fnTool.Run(toolCtx, map[string]any{}) + if err != nil { + t.Fatalf("Failed to run tool: %v", err) + } + return result +} diff --git a/tool/mcptoolset/tool.go b/tool/mcptoolset/tool.go index 640c0cfc3..523772d66 100644 --- a/tool/mcptoolset/tool.go +++ b/tool/mcptoolset/tool.go @@ -28,7 +28,7 @@ import ( "google.golang.org/adk/tool" ) -func convertTool(t *mcp.Tool, client MCPClient, requireConfirmation bool, requireConfirmationProvider ConfirmationProvider) (tool.Tool, error) { +func convertTool(t *mcp.Tool, client MCPClient, requireConfirmation bool, requireConfirmationProvider ConfirmationProvider, metadataProvider MetadataProvider) (tool.Tool, error) { mcp := &mcpTool{ name: t.Name, description: t.Description, @@ -39,6 +39,7 @@ func convertTool(t *mcp.Tool, client MCPClient, requireConfirmation bool, requir mcpClient: client, requireConfirmation: requireConfirmation, requireConfirmationProvider: requireConfirmationProvider, + metadataProvider: metadataProvider, } // Since t.InputSchema and t.OutputSchema are pointers (*jsonschema.Schema) and the destination ResponseJsonSchema @@ -65,6 +66,8 @@ type mcpTool struct { requireConfirmation bool requireConfirmationProvider ConfirmationProvider + + metadataProvider MetadataProvider } // Name implements the tool.Tool. @@ -115,12 +118,18 @@ func (t *mcpTool) Run(ctx tool.Context, args any) (map[string]any, error) { return nil, fmt.Errorf("error tool %q requires confirmation, please approve or reject", t.Name()) } } - - // TODO: add auth - res, err := t.mcpClient.CallTool(ctx, &mcp.CallToolParams{ + params := &mcp.CallToolParams{ Name: t.name, Arguments: args, - }) + } + + if t.metadataProvider != nil { + if meta := t.metadataProvider(ctx); meta != nil { + params.Meta = meta + } + } + // TODO: add auth + res, err := t.mcpClient.CallTool(ctx, params) if err != nil { return nil, fmt.Errorf("failed to call MCP tool %q with err: %w", t.name, err) } From de458365233276b2c910a856065397c4f84bc2f4 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Sat, 28 Feb 2026 10:04:23 +0800 Subject: [PATCH 15/20] follow suggestions of gemini to rename the variable: metaCalled to toolCalled --- tool/mcptoolset/set_test.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tool/mcptoolset/set_test.go b/tool/mcptoolset/set_test.go index c6d40df39..76aa4538b 100644 --- a/tool/mcptoolset/set_test.go +++ b/tool/mcptoolset/set_test.go @@ -840,10 +840,10 @@ func TestMetadataProviderNil(t *testing.T) { func TestMetadataProviderReturnsNil(t *testing.T) { var receivedMeta map[string]any - var metaCalled bool + var toolCalled bool echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { - metaCalled = true + toolCalled = true receivedMeta = req.Params.Meta return nil, struct{ Message string }{Message: "ok"}, nil } @@ -854,7 +854,7 @@ func TestMetadataProviderReturnsNil(t *testing.T) { _ = runMetadataTest(t, metadataProvider, echoToolFunc) - if !metaCalled { + if !toolCalled { t.Fatal("Tool was not called") } From d92b864ac8a24ecb53448940f83a8d6b1233cf96 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Sat, 28 Feb 2026 10:44:48 +0800 Subject: [PATCH 16/20] follow suggestions of gemini to reduce nesting and simplify the codes --- tool/mcptoolset/set_test.go | 8 +++++--- tool/mcptoolset/tool.go | 4 +--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tool/mcptoolset/set_test.go b/tool/mcptoolset/set_test.go index 76aa4538b..5f6961430 100644 --- a/tool/mcptoolset/set_test.go +++ b/tool/mcptoolset/set_test.go @@ -824,18 +824,20 @@ func TestMetadataProvider(t *testing.T) { } func TestMetadataProviderNil(t *testing.T) { + var receivedMeta map[string]any var toolCalled bool echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { toolCalled = true - if req.Params.Meta != nil { - t.Errorf("Expected nil metadata, got %v", req.Params.Meta) - } + receivedMeta = req.Params.Meta return nil, struct{ Message string }{Message: "ok"}, nil } _ = runMetadataTest(t, nil, echoToolFunc) if !toolCalled { t.Fatal("Tool was not called") } + if receivedMeta != nil { + t.Errorf("Expected nil metadata, got %v", receivedMeta) + } } func TestMetadataProviderReturnsNil(t *testing.T) { diff --git a/tool/mcptoolset/tool.go b/tool/mcptoolset/tool.go index 523772d66..8c050cb60 100644 --- a/tool/mcptoolset/tool.go +++ b/tool/mcptoolset/tool.go @@ -124,9 +124,7 @@ func (t *mcpTool) Run(ctx tool.Context, args any) (map[string]any, error) { } if t.metadataProvider != nil { - if meta := t.metadataProvider(ctx); meta != nil { - params.Meta = meta - } + params.Meta = t.metadataProvider(ctx) } // TODO: add auth res, err := t.mcpClient.CallTool(ctx, params) From 05d9def1bdc39c0a15dc3c26dcbbad1f8222f568 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Sat, 28 Feb 2026 10:52:33 +0800 Subject: [PATCH 17/20] follow suggestions of gemini to merge three test functions into one function --- tool/mcptoolset/set_test.go | 105 ++++++++++++++++-------------------- 1 file changed, 45 insertions(+), 60 deletions(-) diff --git a/tool/mcptoolset/set_test.go b/tool/mcptoolset/set_test.go index 5f6961430..5eb6f49a2 100644 --- a/tool/mcptoolset/set_test.go +++ b/tool/mcptoolset/set_test.go @@ -791,77 +791,62 @@ func TestNewToolSet_RequireConfirmationProvider_Validation(t *testing.T) { } } -func TestMetadataProvider(t *testing.T) { - var receivedMeta map[string]any - var toolCalled bool - echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { - toolCalled = true - receivedMeta = req.Params.Meta - return nil, struct{ Message string }{Message: "ok"}, nil - } - +func TestMetadata(t *testing.T) { 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 - } - - result := runMetadataTest(t, metadataProvider, echoToolFunc) - if result == nil { - t.Fatal("Expected non-nil result") - } - - if !toolCalled { - t.Fatal("Tool was not called") - } - - if diff := cmp.Diff(testMetadata, receivedMeta); diff != "" { - t.Errorf("metadata mismatch (-want +got):\n%s", diff) - } -} -func TestMetadataProviderNil(t *testing.T) { - var receivedMeta map[string]any - var toolCalled bool - echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { - toolCalled = true - receivedMeta = req.Params.Meta - return nil, struct{ Message string }{Message: "ok"}, nil - } - _ = runMetadataTest(t, nil, echoToolFunc) - if !toolCalled { - t.Fatal("Tool was not called") - } - if receivedMeta != nil { - t.Errorf("Expected nil metadata, got %v", receivedMeta) - } -} - -func TestMetadataProviderReturnsNil(t *testing.T) { - var receivedMeta map[string]any - var toolCalled bool - - echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { - toolCalled = true - receivedMeta = req.Params.Meta - return nil, struct{ Message string }{Message: "ok"}, nil + testCases := []struct { + name string + provider mcptoolset.MetadataProvider + wantMetadata map[string]any + }{ + { + name: "provider returns metadata", + provider: func(ctx tool.Context) map[string]any { + return testMetadata + }, + wantMetadata: testMetadata, + }, + { + name: "provider is nil", + provider: nil, + wantMetadata: nil, + }, + { + name: "provider returns nil", + provider: func(ctx tool.Context) map[string]any { + return nil + }, + wantMetadata: nil, + }, } - metadataProvider := func(ctx tool.Context) map[string]any { - return nil - } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var receivedMeta map[string]any + var toolCalled bool + echoToolFunc := func(ctx context.Context, req *mcp.CallToolRequest, input struct{}) (*mcp.CallToolResult, struct{ Message string }, error) { + toolCalled = true + receivedMeta = req.Params.Meta + return nil, struct{ Message string }{Message: "ok"}, nil + } - _ = runMetadataTest(t, metadataProvider, echoToolFunc) + result := runMetadataTest(t, tc.provider, echoToolFunc) + if result == nil { + t.Fatal("Expected non-nil result") + } - if !toolCalled { - t.Fatal("Tool was not called") - } + if !toolCalled { + t.Fatal("Tool was not called") + } - if receivedMeta != nil { - t.Errorf("Expected nil metadata when provider returns nil, got %v", receivedMeta) + if diff := cmp.Diff(tc.wantMetadata, receivedMeta); diff != "" { + t.Errorf("metadata mismatch (-want +got):\n%s", diff) + } + }) } } From f841fca7d32e89452042f0dee0c314674cb00d78 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Sat, 28 Feb 2026 11:07:51 +0800 Subject: [PATCH 18/20] change the test function name & check the result in runMetadataTest instead of return the result to calling functions --- tool/mcptoolset/set_test.go | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/tool/mcptoolset/set_test.go b/tool/mcptoolset/set_test.go index 5eb6f49a2..64f418384 100644 --- a/tool/mcptoolset/set_test.go +++ b/tool/mcptoolset/set_test.go @@ -791,7 +791,7 @@ func TestNewToolSet_RequireConfirmationProvider_Validation(t *testing.T) { } } -func TestMetadata(t *testing.T) { +func TestMetadataProvider(t *testing.T) { testMetadata := map[string]any{ "request_id": "test-123", "user_id": "user-456", @@ -834,11 +834,7 @@ func TestMetadata(t *testing.T) { return nil, struct{ Message string }{Message: "ok"}, nil } - result := runMetadataTest(t, tc.provider, echoToolFunc) - if result == nil { - t.Fatal("Expected non-nil result") - } - + runMetadataTest(t, tc.provider, echoToolFunc) if !toolCalled { t.Fatal("Tool was not called") } @@ -850,7 +846,7 @@ func TestMetadata(t *testing.T) { } } -func runMetadataTest[In, Out any](t *testing.T, provider mcptoolset.MetadataProvider, toolFunc mcp.ToolHandlerFor[In, Out]) map[string]any { +func runMetadataTest[In, Out any](t *testing.T, provider mcptoolset.MetadataProvider, toolFunc mcp.ToolHandlerFor[In, Out]) { t.Helper() clientTransport, serverTransport := mcp.NewInMemoryTransports() @@ -890,5 +886,7 @@ func runMetadataTest[In, Out any](t *testing.T, provider mcptoolset.MetadataProv if err != nil { t.Fatalf("Failed to run tool: %v", err) } - return result + if result == nil { + t.Fatal("Expected non-nil result from tool run") + } } From da61ec7c60e6dc9439d54a2504bca1ebe0c4b202 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Sat, 28 Feb 2026 11:16:36 +0800 Subject: [PATCH 19/20] remove the result checking and only use the error checking in runMetadataTest --- tool/mcptoolset/set_test.go | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tool/mcptoolset/set_test.go b/tool/mcptoolset/set_test.go index 64f418384..56821b043 100644 --- a/tool/mcptoolset/set_test.go +++ b/tool/mcptoolset/set_test.go @@ -882,11 +882,8 @@ func runMetadataTest[In, Out any](t *testing.T, provider mcptoolset.MetadataProv } toolCtx := toolinternal.NewToolContext(invCtx, "", nil, nil) - result, err := fnTool.Run(toolCtx, map[string]any{}) + _, 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 from tool run") - } } From ce3f6579df32c04c6c1dbe21d4ca2f49ac578bc9 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Sat, 28 Feb 2026 11:25:43 +0800 Subject: [PATCH 20/20] add recover to make the function more robust --- tool/mcptoolset/tool.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tool/mcptoolset/tool.go b/tool/mcptoolset/tool.go index 8c050cb60..d54e73c70 100644 --- a/tool/mcptoolset/tool.go +++ b/tool/mcptoolset/tool.go @@ -17,6 +17,7 @@ package mcptoolset import ( "errors" "fmt" + "runtime/debug" "strings" "github.com/modelcontextprotocol/go-sdk/mcp" @@ -93,7 +94,13 @@ func (t *mcpTool) Declaration() *genai.FunctionDeclaration { return t.funcDeclaration } -func (t *mcpTool) Run(ctx tool.Context, args any) (map[string]any, error) { +func (t *mcpTool) Run(ctx tool.Context, args any) (result map[string]any, err error) { + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("panic in tool %q: %v\nstack: %s", t.Name(), r, debug.Stack()) + } + }() + if confirmation := ctx.ToolConfirmation(); confirmation != nil { if !confirmation.Confirmed { return nil, fmt.Errorf("error tool %q call is rejected", t.Name())