diff --git a/plugin/retryandreflect/plugin.go b/plugin/retryandreflect/plugin.go index 31fe520c7..393dda4fa 100644 --- a/plugin/retryandreflect/plugin.go +++ b/plugin/retryandreflect/plugin.go @@ -145,8 +145,8 @@ func (r *retryAndReflect) onToolError(ctx agent.Context, tool tool.Tool, args ma } func (r *retryAndReflect) handleToolError(ctx agent.Context, failedTool tool.Tool, args map[string]any, err error) (map[string]any, error) { - // skip if the error is tool.ErrConfirmationRequired. - if errors.Is(err, tool.ErrConfirmationRequired) || errors.Is(err, tool.ErrConfirmationRejected) { + // Skip HITL sentinels: these pause the run for user input, not for retry. + if errors.Is(err, tool.ErrConfirmationRequired) || errors.Is(err, tool.ErrConfirmationRejected) || errors.Is(err, tool.ErrCredentialRequired) { return nil, nil } diff --git a/tool/mcptoolset/credential_test.go b/tool/mcptoolset/credential_test.go new file mode 100644 index 000000000..413e4abaf --- /dev/null +++ b/tool/mcptoolset/credential_test.go @@ -0,0 +1,174 @@ +// Copyright 2026 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 ( + "context" + "errors" + "testing" + + "github.com/modelcontextprotocol/go-sdk/mcp" + + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/auth" + icontext "google.golang.org/adk/v2/internal/context" + "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/tool" + "google.golang.org/adk/v2/tool/authconsent" +) + +// First invocation with a provider that needs consent: the tool requests +// credential consent, marks the pending sentinel, and does not call the server. +func TestMcpToolRun_RequestsConsent(t *testing.T) { + client := &fakeCredClient{} + tl := &mcpTool{name: "weather", mcpClient: client, auth: consentProvider()} + actions := &session.EventActions{} + + _, err := tl.Run(newToolCtx(t, actions), map[string]any{}) + if !errors.Is(err, tool.ErrCredentialRequired) { + t.Fatalf("Run() error = %v, want tool.ErrCredentialRequired", err) + } + got, ok := actions.RequestedCredentials["fn1"] + if !ok { + t.Fatalf("RequestedCredentials not set; got %+v", actions.RequestedCredentials) + } + if got.AuthURI != "https://consent.example" || got.Key != "k1" { + t.Errorf("RequestedCredentials[fn1] = %+v, want AuthURI/Key from ConsentRequiredError", got) + } + if client.called != 0 { + t.Errorf("CallTool called %d times, want 0 (paused for consent)", client.called) + } +} + +// After consent (AuthResponse present) and the provider now resolving, the tool +// proceeds to call the server. +func TestMcpToolRun_ProceedsAfterConsent(t *testing.T) { + client := &fakeCredClient{} + tl := &mcpTool{name: "weather", mcpClient: client, auth: okProvider()} + + ctx := newToolCtx(t, &session.EventActions{}). + WithDelta(&agent.CommonContextDelta{CredentialResponse: &authconsent.Response{Token: "tok"}}) + + res, err := tl.Run(ctx, map[string]any{}) + if err != nil { + t.Fatalf("Run() error = %v, want nil", err) + } + if res["output"] != "ok" { + t.Errorf("Run() output = %v, want %q", res["output"], "ok") + } + if client.called != 1 { + t.Errorf("CallTool called %d times, want 1", client.called) + } +} + +// Resumed after consent but the provider still needs consent: fail instead of +// re-requesting (which would loop). +func TestMcpToolRun_ConsentGuardNoLoop(t *testing.T) { + client := &fakeCredClient{} + tl := &mcpTool{name: "weather", mcpClient: client, auth: consentProvider()} + actions := &session.EventActions{} + + ctx := agent.NewToolContext( + icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}), + "fn1", actions, nil, + ).WithDelta(&agent.CommonContextDelta{CredentialResponse: &authconsent.Response{}}) + + _, err := tl.Run(ctx, map[string]any{}) + if err == nil || errors.Is(err, tool.ErrCredentialRequired) { + t.Fatalf("Run() error = %v, want a non-pending failure", err) + } + if len(actions.RequestedCredentials) != 0 { + t.Errorf("RequestedCredentials = %+v, want empty (no re-request on resume)", actions.RequestedCredentials) + } + if client.called != 0 { + t.Errorf("CallTool called %d times, want 0", client.called) + } +} + +// With no auth provider the pre-flight probe is skipped entirely. +func TestMcpToolRun_NoAuthProvider(t *testing.T) { + client := &fakeCredClient{} + tl := &mcpTool{name: "weather", mcpClient: client} + + res, err := tl.Run(newToolCtx(t, &session.EventActions{}), map[string]any{}) + if err != nil { + t.Fatalf("Run() error = %v, want nil", err) + } + if res["output"] != "ok" || client.called != 1 { + t.Errorf("Run() output=%v called=%d, want output=ok called=1", res["output"], client.called) + } +} + +// A non-consent probe error fails fast: the tool surfaces the real resolution +// error and does not call the server (the RoundTripper would fail the same way, +// but the MCP SDK would mangle the error chain). +func TestMcpToolRun_NonConsentErrorFailsFast(t *testing.T) { + client := &fakeCredClient{} + tl := &mcpTool{name: "weather", mcpClient: client, auth: errProvider()} + actions := &session.EventActions{} + + _, err := tl.Run(newToolCtx(t, actions), map[string]any{}) + if err == nil { + t.Fatal("Run() error = nil, want the resolution error surfaced") + } + if errors.Is(err, tool.ErrCredentialRequired) { + t.Errorf("Run() error = %v, want a non-consent failure (not ErrCredentialRequired)", err) + } + if client.called != 0 { + t.Errorf("CallTool called %d times, want 0 (fail fast before the server call)", client.called) + } + if len(actions.RequestedCredentials) != 0 { + t.Errorf("RequestedCredentials = %+v, want empty (no consent for a non-consent error)", actions.RequestedCredentials) + } +} + +type fakeCredClient struct { + called int +} + +func (f *fakeCredClient) CallTool(context.Context, *mcp.CallToolParams) (*mcp.CallToolResult, error) { + f.called++ + return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "ok"}}}, nil +} + +func (f *fakeCredClient) ListTools(context.Context) ([]*mcp.Tool, error) { return nil, nil } + +// consentProvider always reports interactive consent is required. +func consentProvider() auth.CredentialProvider { + return auth.ProviderFunc(func(context.Context) (*auth.Credential, error) { + return nil, &auth.ConsentRequiredError{AuthURI: "https://consent.example", Nonce: "n1", Key: "k1"} + }) +} + +// okProvider always resolves a credential. +func okProvider() auth.CredentialProvider { + return auth.ProviderFunc(func(context.Context) (*auth.Credential, error) { + return &auth.Credential{}, nil + }) +} + +// errProvider always fails with a plain (non-consent) error. +func errProvider() auth.CredentialProvider { + return auth.ProviderFunc(func(context.Context) (*auth.Credential, error) { + return nil, errors.New("boom") + }) +} + +// newToolCtx builds a tool context for function call id "fn1". +func newToolCtx(t *testing.T, actions *session.EventActions) agent.Context { + t.Helper() + inv := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) + return agent.NewToolContext(inv, "fn1", actions, nil) +} diff --git a/tool/mcptoolset/set.go b/tool/mcptoolset/set.go index 1f264486c..a97759a12 100644 --- a/tool/mcptoolset/set.go +++ b/tool/mcptoolset/set.go @@ -58,6 +58,7 @@ func New(cfg Config) (tool.Toolset, error) { toolFilter: cfg.ToolFilter, requireConfirmation: cfg.RequireConfirmation, requireConfirmationProvider: cfg.RequireConfirmationProvider, + auth: cfg.Auth, }, nil } @@ -111,6 +112,12 @@ type Config struct { // HTTP transport: set Endpoint, or pass a *mcp.StreamableClientTransport. // Combining Auth with a non-HTTP transport (e.g. a stdio command) is a // configuration error. See package google.golang.org/adk/v2/auth. + // + // Interactive 3-legged consent (a provider returning auth.ConsentRequiredError) + // is driven only during tool execution, which can pause the run for consent. + // Tool listing (Tools) runs before any tool call and cannot pause, so a server + // that authenticates listing must use a non-interactive or already-consented + // (cached) credential; otherwise listing fails. Auth auth.CredentialProvider // Deprecated: use tool.FilterToolset instead. @@ -142,6 +149,7 @@ type set struct { toolFilter tool.Predicate requireConfirmation bool requireConfirmationProvider tool.ConfirmationProvider + auth auth.CredentialProvider } func (*set) Name() string { @@ -165,7 +173,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.auth) 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/tool.go b/tool/mcptoolset/tool.go index ec73d8f67..a5482b930 100644 --- a/tool/mcptoolset/tool.go +++ b/tool/mcptoolset/tool.go @@ -23,13 +23,15 @@ import ( "google.golang.org/genai" "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/auth" "google.golang.org/adk/v2/internal/toolinternal" "google.golang.org/adk/v2/model" "google.golang.org/adk/v2/tool" + "google.golang.org/adk/v2/tool/authconsent" "google.golang.org/adk/v2/tool/toolutils" ) -func convertTool(t *mcp.Tool, client MCPClient, requireConfirmation bool, requireConfirmationProvider tool.ConfirmationProvider) (tool.Tool, error) { +func convertTool(t *mcp.Tool, client MCPClient, requireConfirmation bool, requireConfirmationProvider tool.ConfirmationProvider, authProvider auth.CredentialProvider) (tool.Tool, error) { mcp := &mcpTool{ name: t.Name, description: t.Description, @@ -40,6 +42,7 @@ func convertTool(t *mcp.Tool, client MCPClient, requireConfirmation bool, requir mcpClient: client, requireConfirmation: requireConfirmation, requireConfirmationProvider: requireConfirmationProvider, + auth: authProvider, } // Since t.InputSchema and t.OutputSchema are pointers (*jsonschema.Schema) and the destination ResponseJsonSchema @@ -66,6 +69,12 @@ type mcpTool struct { requireConfirmation bool requireConfirmationProvider tool.ConfirmationProvider + + // auth, when set, resolves the per-request credential. It mirrors the + // provider wired into the transport's RoundTripper and is used here for a + // pre-flight probe that can initiate interactive consent (which the + // RoundTripper cannot: it has no function call id). + auth auth.CredentialProvider } // Name implements the tool.Tool. @@ -117,6 +126,37 @@ func (t *mcpTool) Run(ctx agent.Context, args any) (map[string]any, error) { } } + // Interactive auth pre-flight: the transport's RoundTripper applies the + // credential per request but cannot start a consent flow (it has no function + // call id, and the MCP SDK drops the underlying error chain). So probe the + // provider here; if it needs interactive consent, request it and pause. On + // success the provider caches the credential, so the RoundTripper reuses it + // on the actual call below. + if t.auth != nil { + if _, err := t.auth.Credential(ctx); err != nil { + var consent *auth.ConsentRequiredError + if !errors.As(err, &consent) { + // A non-consent resolution failure means the RoundTripper would fail + // the same way on the call below, but the MCP SDK mangles that error + // chain — so surface the real cause now instead of a wasted call. + return nil, fmt.Errorf("mcp tool %q: resolve credential: %w", t.Name(), err) + } + if ctx.AuthResponse() != nil { + // Resumed after consent but the credential is still unavailable: + // fail rather than request consent again (which would loop). + return nil, fmt.Errorf("mcp tool %q: consent completed but credential unavailable: %w", t.Name(), err) + } + if rerr := ctx.RequestCredential(authconsent.Request{ + AuthURI: consent.AuthURI, + Nonce: consent.Nonce, + Key: consent.Key, + }); rerr != nil { + return nil, fmt.Errorf("mcp tool %q: request credential: %w", t.Name(), rerr) + } + return nil, fmt.Errorf("mcp tool %q: %w", t.Name(), tool.ErrCredentialRequired) + } + } + res, err := t.mcpClient.CallTool(ctx, &mcp.CallToolParams{ Name: t.name, Arguments: args, diff --git a/tool/tool.go b/tool/tool.go index 6ed1ad9a6..0f2f1e4ad 100644 --- a/tool/tool.go +++ b/tool/tool.go @@ -34,6 +34,12 @@ var ErrConfirmationRequired = errors.New("requires confirmation, please approve // ErrConfirmationRejected indicated that the tool call confirmation rejected. var ErrConfirmationRejected = errors.New("call is rejected") +// ErrCredentialRequired indicates that the tool needs interactive (3-legged) +// OAuth consent before it can proceed. It is the credential analog of +// ErrConfirmationRequired: the tool has called agent.Context.RequestCredential, +// and the run pauses until the user completes consent. +var ErrCredentialRequired = errors.New("requires interactive credential consent") + // Tool defines the interface for a callable tool. type Tool interface { // Name returns the name of the tool.