diff --git a/tool/mcptoolset/client.go b/tool/mcptoolset/client.go index a154ef7d2..fbe3498ca 100644 --- a/tool/mcptoolset/client.go +++ b/tool/mcptoolset/client.go @@ -53,10 +53,10 @@ var refreshableErrors = []error{ } // newConnectionRefresher creates a new connectionRefresher with the given client and transport. -// If client is nil, a default MCP client will be created. -func newConnectionRefresher(client *mcp.Client, transport mcp.Transport) *connectionRefresher { +// If client is nil, a default MCP client will be created with the given options. +func newConnectionRefresher(client *mcp.Client, transport mcp.Transport, options *mcp.ClientOptions) *connectionRefresher { if client == nil { - client = mcp.NewClient(&mcp.Implementation{Name: "adk-mcp-client", Version: version.Version}, nil) + client = mcp.NewClient(&mcp.Implementation{Name: "adk-mcp-client", Version: version.Version}, options) } return &connectionRefresher{ client: client, diff --git a/tool/mcptoolset/set.go b/tool/mcptoolset/set.go index 228d777d5..e01b41ae7 100644 --- a/tool/mcptoolset/set.go +++ b/tool/mcptoolset/set.go @@ -16,6 +16,7 @@ package mcptoolset import ( + "context" "fmt" "github.com/modelcontextprotocol/go-sdk/mcp" @@ -47,8 +48,32 @@ import ( // }, // }) func New(cfg Config) (tool.Toolset, error) { + var clientOptions *mcp.ClientOptions + if cfg.ElicitationHandler != nil || cfg.ElicitationCompleteHandler != nil { + if cfg.Client != nil { + return nil, fmt.Errorf("mcptoolset: ElicitationHandler and ElicitationCompleteHandler cannot be combined with a custom Client; set them in the client's mcp.ClientOptions instead") + } + if cfg.ElicitationHandler == nil { + return nil, fmt.Errorf("mcptoolset: ElicitationCompleteHandler requires ElicitationHandler to be set; the client cannot service an elicitation without it") + } + clientOptions = &mcp.ClientOptions{ + ElicitationHandler: cfg.ElicitationHandler, + ElicitationCompleteHandler: cfg.ElicitationCompleteHandler, + // The capability inferred from ElicitationHandler alone covers only + // form mode; URL mode must be declared explicitly. RootsV2 preserves + // the default roots capability, which setting Capabilities would + // otherwise disable. + Capabilities: &mcp.ClientCapabilities{ + Elicitation: &mcp.ElicitationCapabilities{ + Form: &mcp.FormElicitationCapabilities{}, + URL: &mcp.URLElicitationCapabilities{}, + }, + RootsV2: &mcp.RootCapabilities{ListChanged: true}, + }, + } + } return &set{ - mcpClient: newConnectionRefresher(cfg.Client, cfg.Transport), + mcpClient: newConnectionRefresher(cfg.Client, cfg.Transport, clientOptions), toolFilter: cfg.ToolFilter, requireConfirmation: cfg.RequireConfirmation, requireConfirmationProvider: cfg.RequireConfirmationProvider, @@ -72,6 +97,23 @@ type Config struct { // a Human-in-the-Loop (HITL) confirmation request when a tool is invoked. RequireConfirmation bool + // ElicitationHandler handles elicitation/create requests from the MCP + // server, including URL-mode elicitations that servers use for + // out-of-band interactions such as auth challenges. Setting it makes the + // client advertise the elicitation capability. + // It can only be set when Client is nil; for a custom Client, set the + // handler in the client's mcp.ClientOptions instead. + ElicitationHandler func(context.Context, *mcp.ElicitRequest) (*mcp.ElicitResult, error) + + // ElicitationCompleteHandler handles notifications/elicitation/complete + // notifications, which servers send when an out-of-band (URL-mode) + // elicitation has been completed. It requires ElicitationHandler to also + // be set, since a completion notification cannot arrive unless an + // elicitation was created first. + // It can only be set when Client is nil; for a custom Client, set the + // handler in the client's mcp.ClientOptions instead. + ElicitationCompleteHandler func(context.Context, *mcp.ElicitationCompleteNotificationRequest) + // RequireConfirmationProvider allows for dynamic determination of whether // user confirmation is needed. This field is a function called at runtime to decide if // a confirmation request should be sent. The function takes the toolName and tool's input parameters as arguments. diff --git a/tool/mcptoolset/set_test.go b/tool/mcptoolset/set_test.go index 5f7ca9aa0..54be96863 100644 --- a/tool/mcptoolset/set_test.go +++ b/tool/mcptoolset/set_test.go @@ -16,6 +16,7 @@ package mcptoolset_test import ( "context" + "encoding/json" "errors" "fmt" "io" @@ -796,3 +797,182 @@ func TestNewToolSet_RequireConfirmationProvider_Validation(t *testing.T) { }) } } + +func TestCallToolMetaPassthrough(t *testing.T) { + challengeMeta := map[string]any{ + "example.com/auth": map[string]any{"url": "https://idp.example.com/login"}, + } + + tests := []struct { + name string + handler mcp.ToolHandler + want map[string]any + }{ + { + name: "text result with meta", + handler: func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Meta: mcp.Meta(challengeMeta), + Content: []mcp.Content{&mcp.TextContent{Text: "login required"}}, + }, nil + }, + want: map[string]any{ + "output": "login required", + "_meta": challengeMeta, + }, + }, + { + name: "structured result with meta", + handler: func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Meta: mcp.Meta(challengeMeta), + Content: []mcp.Content{&mcp.TextContent{Text: "login required"}}, + StructuredContent: map[string]any{"status": "unauthenticated"}, + }, nil + }, + want: map[string]any{ + "output": map[string]any{"status": "unauthenticated"}, + "_meta": challengeMeta, + }, + }, + { + name: "text result without meta", + handler: func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "all good"}}, + }, nil + }, + want: map[string]any{ + "output": "all good", + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + clientTransport, serverTransport := mcp.NewInMemoryTransports() + + server := mcp.NewServer(&mcp.Implementation{Name: "test_server", Version: "v1.0.0"}, nil) + server.AddTool(&mcp.Tool{ + Name: "gateway_tool", + Description: "returns a result with metadata", + InputSchema: json.RawMessage(`{"type":"object"}`), + }, tc.handler) + if _, err := server.Connect(t.Context(), serverTransport, nil); err != nil { + t.Fatal(err) + } + + ts, err := mcptoolset.New(mcptoolset.Config{Transport: clientTransport}) + if err != nil { + t.Fatalf("Failed to create MCP tool set: %v", err) + } + + invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) + tools, err := ts.Tools(icontext.NewReadonlyContext(invCtx)) + if err != nil { + t.Fatalf("Tools call failed: %v", err) + } + toolCtx := agent.NewToolContext(invCtx, "", nil, nil) + + fnTool := tools[0].(toolinternal.FunctionTool) + result, err := fnTool.Run(toolCtx, map[string]any{}) + if err != nil { + t.Fatalf("Tool call failed: %v", err) + } + + if diff := cmp.Diff(tc.want, result); diff != "" { + t.Errorf("Tool result mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestElicitationHandler(t *testing.T) { + const loginURL = "https://idp.example.com/login" + + server := mcp.NewServer(&mcp.Implementation{Name: "test_server", Version: "v1.0.0"}, nil) + server.AddTool(&mcp.Tool{ + Name: "gateway_tool", + Description: "elicits a login before responding", + InputSchema: json.RawMessage(`{"type":"object"}`), + }, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + res, err := req.Session.Elicit(ctx, &mcp.ElicitParams{ + Mode: "url", + Message: "please log in", + URL: loginURL, + ElicitationID: "elicitation-1", + }) + if err != nil { + return nil, err + } + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "elicitation " + res.Action + "ed"}}, + }, nil + }) + + clientTransport, serverTransport := mcp.NewInMemoryTransports() + if _, err := server.Connect(t.Context(), serverTransport, nil); err != nil { + t.Fatal(err) + } + + var gotURL string + ts, err := mcptoolset.New(mcptoolset.Config{ + Transport: clientTransport, + ElicitationHandler: func(ctx context.Context, req *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + gotURL = req.Params.URL + return &mcp.ElicitResult{Action: "accept"}, nil + }, + }) + if err != nil { + t.Fatalf("Failed to create MCP tool set: %v", err) + } + + invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) + tools, err := ts.Tools(icontext.NewReadonlyContext(invCtx)) + if err != nil { + t.Fatalf("Tools call failed: %v", err) + } + toolCtx := agent.NewToolContext(invCtx, "", nil, nil) + + fnTool := tools[0].(toolinternal.FunctionTool) + result, err := fnTool.Run(toolCtx, map[string]any{}) + if err != nil { + t.Fatalf("Tool call failed: %v", err) + } + + if gotURL != loginURL { + t.Errorf("Elicitation handler got URL %q, want %q", gotURL, loginURL) + } + want := map[string]any{"output": "elicitation accepted"} + if diff := cmp.Diff(want, result); diff != "" { + t.Errorf("Tool result mismatch (-want +got):\n%s", diff) + } +} + +func TestNewRejectsElicitationHandlerWithCustomClient(t *testing.T) { + clientTransport, _ := mcp.NewInMemoryTransports() + + _, err := mcptoolset.New(mcptoolset.Config{ + Client: mcp.NewClient(&mcp.Implementation{Name: "custom", Version: "v1.0.0"}, nil), + Transport: clientTransport, + ElicitationHandler: func(ctx context.Context, req *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + return &mcp.ElicitResult{Action: "decline"}, nil + }, + }) + if err == nil { + t.Fatal("expected error when combining ElicitationHandler with a custom Client, got nil") + } +} + +func TestNewRejectsElicitationCompleteHandlerWithoutElicitationHandler(t *testing.T) { + clientTransport, _ := mcp.NewInMemoryTransports() + + _, err := mcptoolset.New(mcptoolset.Config{ + Transport: clientTransport, + ElicitationCompleteHandler: func(context.Context, *mcp.ElicitationCompleteNotificationRequest) { + }, + }) + if err == nil { + t.Fatal("expected error when setting ElicitationCompleteHandler without ElicitationHandler, got nil") + } +} diff --git a/tool/mcptoolset/tool.go b/tool/mcptoolset/tool.go index 10633a28b..622bcf4a0 100644 --- a/tool/mcptoolset/tool.go +++ b/tool/mcptoolset/tool.go @@ -147,9 +147,7 @@ func (t *mcpTool) Run(ctx agent.Context, args any) (map[string]any, error) { } if res.StructuredContent != nil { - return map[string]any{ - "output": res.StructuredContent, - }, nil + return functionResponse(res, res.StructuredContent), nil } textResponse := strings.Builder{} @@ -169,9 +167,24 @@ func (t *mcpTool) Run(ctx agent.Context, args any) (map[string]any, error) { return nil, errors.New("no text content in tool response") } - return map[string]any{ - "output": textResponse.String(), - }, nil + return functionResponse(res, textResponse.String()), nil +} + +// functionResponse builds the function response map for a tool result. +// The result's _meta field is preserved under the "_meta" key, mirroring the +// raw MCP serialization, so that metadata attached by the server (e.g. auth +// challenges from MCP gateways) is not silently dropped. This map is the +// function response returned to the model, so _meta reaches the LLM and is +// persisted to session and traces, in addition to being available to +// callbacks and the embedding application. +func functionResponse(res *mcp.CallToolResult, output any) map[string]any { + response := map[string]any{ + "output": output, + } + if len(res.Meta) > 0 { + response["_meta"] = map[string]any(res.Meta) + } + return response } var (