From 94eda1e9db4689316d3089e94f534c91270c88d2 Mon Sep 17 00:00:00 2001 From: Francesc Campoy Date: Fri, 30 Oct 2026 08:48:20 -0700 Subject: [PATCH 01/10] Breaking: list each Google ADK Go MCP toolset once per model step ADK v2.5.0 and earlier ask a toolset for its tools once per agent run, and the MCP proxy's ProcessRequest lists again before each model call. ADK v2.6.0 also asks before each model call. The proxy runs a ListMcpTools Activity for each request, so a one-call turn scheduled 3 listings with ADK v2.5.0 and 4 with v2.6.0, and MCP Histories failed replay after an ADK upgrade or rollback. ProcessRequest now reuses the list that Tools took earlier in the same step and lists only when there is none. Each model step lists once on every ADK release. Breaking for replay: open runs that used NewMCPToolset fail replay on the new version, because they recorded more ListMcpTools Activities. Let those runs finish or keep them on the old version with Worker Versioning before upgrading. --- go/googleadk/context.go | 2 + go/googleadk/mcp.go | 83 +++++++++++++++++++++++++++++++++++++--- go/googleadk/mcp_test.go | 8 +++- 3 files changed, 86 insertions(+), 7 deletions(-) diff --git a/go/googleadk/context.go b/go/googleadk/context.go index a1b12120..014af9fa 100644 --- a/go/googleadk/context.go +++ b/go/googleadk/context.go @@ -69,6 +69,8 @@ func NewContext(ctx workflow.Context, opts ...ContextOption) context.Context { base = platform.WithTimeProvider(base, func() time.Time { return workflow.Now(ctx) }) base = platform.WithUUIDProvider(base, newDeterministicUUIDProvider(ctx)) base = platform.WithTaskRunner(base, newWorkflowTaskRunner(ctx, cfg.sequentialToolFanout)) + // Lets each MCP toolset list its tools once per model step. + base = context.WithValue(base, mcpListingsKey{}, &mcpListings{}) return base } diff --git a/go/googleadk/mcp.go b/go/googleadk/mcp.go index c6ce2a2d..fd132546 100644 --- a/go/googleadk/mcp.go +++ b/go/googleadk/mcp.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "sync" "go.temporal.io/sdk/activity" "go.temporal.io/sdk/workflow" @@ -59,8 +60,20 @@ type mcpToolset struct { func (m *mcpToolset) Name() string { return m.name } // Tools lists the remote MCP tools via the ListMcpTools Activity and wraps each -// as a declaration-only proxy tool. +// as a declaration-only proxy tool. It keeps the list for the ProcessRequest +// call that follows in the same model step. func (m *mcpToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { + tools, err := m.list(ctx) + if err != nil { + return nil, err + } + listingsFrom(ctx).put(m, ctx, tools) + return tools, nil +} + +// list runs the ListMcpTools Activity and wraps each declaration as a proxy +// tool. +func (m *mcpToolset) list(ctx agent.ReadonlyContext) ([]tool.Tool, error) { wfCtx, ok := workflowContext(ctx) if !ok { return nil, errMissingContext @@ -81,11 +94,18 @@ func (m *mcpToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { } // ProcessRequest packs every advertised proxy tool's declaration into the model -// request. ADK calls this during toolset preprocessing. +// request. ADK calls it before each model call. ADK v2.5.0 and earlier call +// Tools once per agent run; v2.6.0 and later call it before each model call +// too. So ProcessRequest reuses the list Tools took in the same step and lists +// only when there is none. Each model step then lists the toolset once on every +// ADK release, and a History recorded with one release replays with another. func (m *mcpToolset) ProcessRequest(ctx agent.Context, req *model.LLMRequest) error { - tools, err := m.Tools(ctx) - if err != nil { - return err + tools, ok := listingsFrom(ctx).take(m, ctx) + if !ok { + var err error + if tools, err = m.list(ctx); err != nil { + return err + } } for _, t := range tools { if err := packTool(req, t.(*mcpProxyTool)); err != nil { @@ -95,6 +115,59 @@ func (m *mcpToolset) ProcessRequest(ctx agent.Context, req *model.LLMRequest) er return nil } +// mcpListingsKey is the private context key for the mcpListings NewContext +// installs. +type mcpListingsKey struct{} + +// mcpListings hands the tools an MCP toolset's Tools call listed to the +// ProcessRequest call that follows in the same model step. +type mcpListings struct { + mu sync.Mutex + pending map[mcpListingKey]mcpListing +} + +type mcpListingKey struct { + toolset *mcpToolset + agent string +} + +type mcpListing struct { + invocation string + tools []tool.Tool +} + +// listingsFrom returns the mcpListings NewContext installed, or nil. +func listingsFrom(ctx context.Context) *mcpListings { + l, _ := ctx.Value(mcpListingsKey{}).(*mcpListings) + return l +} + +func (l *mcpListings) put(m *mcpToolset, ctx agent.ReadonlyContext, tools []tool.Tool) { + if l == nil { + return + } + l.mu.Lock() + defer l.mu.Unlock() + if l.pending == nil { + l.pending = map[mcpListingKey]mcpListing{} + } + l.pending[mcpListingKey{m, ctx.AgentName()}] = mcpListing{ctx.InvocationID(), tools} +} + +// take returns and forgets the list Tools kept for this toolset, agent and +// invocation. +func (l *mcpListings) take(m *mcpToolset, ctx agent.ReadonlyContext) ([]tool.Tool, bool) { + if l == nil { + return nil, false + } + l.mu.Lock() + defer l.mu.Unlock() + k := mcpListingKey{m, ctx.AgentName()} + p, ok := l.pending[k] + delete(l.pending, k) + return p.tools, ok && p.invocation == ctx.InvocationID() +} + // mcpProxyTool is a declaration-only stand-in for a remote MCP tool. The model // sees its declaration; when the model calls it, Run (invoked in-workflow by // ADK) dispatches the CallMcpTool Activity so the live, stateful MCP session diff --git a/go/googleadk/mcp_test.go b/go/googleadk/mcp_test.go index e5224c2d..9ed11d1c 100644 --- a/go/googleadk/mcp_test.go +++ b/go/googleadk/mcp_test.go @@ -111,8 +111,12 @@ func TestMcpListToolsCarriesParameters(t *testing.T) { require.NotNil(t, decl.Parameters, "MCP tool declaration carried no parameters") assert.Contains(t, decl.Parameters.Properties, "query") - // The tools were listed (>=1, packed each turn) and called once. - assert.GreaterOrEqual(t, counter.get(googleadk.ListMcpToolsActivityName), 1) + // The tools were listed once per model step, on every ADK release, and + // called once. ADK v2.6.0 asks the toolset for its tools twice per step; + // a second listing would add an Activity, and Histories recorded with + // another ADK release would fail replay. + assert.Equal(t, 2, counter.get(googleadk.InvokeModelActivityName)) + assert.Equal(t, 2, counter.get(googleadk.ListMcpToolsActivityName), "one listing per model step") assert.Equal(t, 1, counter.get(googleadk.CallMcpToolActivityName)) } From 38ffaac1f19799a71cf40045e4957bdca50c3a25 Mon Sep 17 00:00:00 2001 From: Francesc Campoy Date: Fri, 30 Oct 2026 10:17:45 -0700 Subject: [PATCH 02/10] Address review findings --- go/googleadk/mcp.go | 134 +++++++++++---- go/googleadk/mcp_test.go | 359 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 464 insertions(+), 29 deletions(-) diff --git a/go/googleadk/mcp.go b/go/googleadk/mcp.go index fd132546..f6d1d4be 100644 --- a/go/googleadk/mcp.go +++ b/go/googleadk/mcp.go @@ -59,15 +59,21 @@ type mcpToolset struct { func (m *mcpToolset) Name() string { return m.name } -// Tools lists the remote MCP tools via the ListMcpTools Activity and wraps each -// as a declaration-only proxy tool. It keeps the list for the ProcessRequest -// call that follows in the same model step. +// Tools returns the toolset's declaration-only proxy tools. It lists them via +// the ListMcpTools Activity until the toolset's ProcessRequest has run for the +// agent, and keeps that list for the ProcessRequest call that follows in the +// same model step. After that, ProcessRequest lists, and Tools returns the +// latest list without an Activity. See ProcessRequest. func (m *mcpToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { + l := listingsFrom(ctx) + if tools, ok := l.cached(m, ctx); ok { + return tools, nil + } tools, err := m.list(ctx) if err != nil { return nil, err } - listingsFrom(ctx).put(m, ctx, tools) + l.listed(m, ctx, tools, true) return tools, nil } @@ -85,7 +91,7 @@ func (m *mcpToolset) list(ctx agent.ReadonlyContext) ([]tool.Tool, error) { } tools := make([]tool.Tool, 0, len(decls)) for _, d := range decls { - pt := &mcpProxyTool{toolset: m.name, decl: d, activityOptions: m.activityOptions} + pt := &mcpProxyTool{owner: m, toolset: m.name, decl: d, activityOptions: m.activityOptions} if m.filter == nil || m.filter(ctx, pt) { tools = append(tools, pt) } @@ -93,19 +99,28 @@ func (m *mcpToolset) list(ctx agent.ReadonlyContext) ([]tool.Tool, error) { return tools, nil } -// ProcessRequest packs every advertised proxy tool's declaration into the model -// request. ADK calls it before each model call. ADK v2.5.0 and earlier call -// Tools once per agent run; v2.6.0 and later call it before each model call -// too. So ProcessRequest reuses the list Tools took in the same step and lists -// only when there is none. Each model step then lists the toolset once on every -// ADK release, and a History recorded with one release replays with another. +// ProcessRequest packs the toolset's current proxy tools into the model +// request. ADK calls it before each model call, after the request processors +// (instructions and the others) have run. +// +// ADK v2.5.0 and earlier call Tools only in the first model step of an agent +// run, before the request processors. v2.6.0 and later call Tools before the +// request processors in every step. To schedule the same commands in the same +// order on every ADK release, the toolset lists in Tools only until its first +// ProcessRequest for the agent, and in ProcessRequest from then on. Each model +// step lists once: the first through Tools, with ProcessRequest reusing that +// list, and every later step here. The request carries this list alone, so a +// tool the MCP server no longer advertises is neither declared nor dispatched, +// whatever list ADK kept from an earlier Tools call. func (m *mcpToolset) ProcessRequest(ctx agent.Context, req *model.LLMRequest) error { - tools, ok := listingsFrom(ctx).take(m, ctx) + l := listingsFrom(ctx) + tools, ok := l.take(m, ctx) if !ok { var err error if tools, err = m.list(ctx); err != nil { return err } + l.listed(m, ctx, tools, false) } for _, t := range tools { if err := packTool(req, t.(*mcpProxyTool)); err != nil { @@ -119,11 +134,12 @@ func (m *mcpToolset) ProcessRequest(ctx agent.Context, req *model.LLMRequest) er // installs. type mcpListingsKey struct{} -// mcpListings hands the tools an MCP toolset's Tools call listed to the -// ProcessRequest call that follows in the same model step. +// mcpListings records, for each MCP toolset and agent, the lists the toolset +// took, so that Tools and ProcessRequest list once per model step at the same +// point on every ADK release. type mcpListings struct { - mu sync.Mutex - pending map[mcpListingKey]mcpListing + mu sync.Mutex + state map[mcpListingKey]*mcpListingState } type mcpListingKey struct { @@ -131,9 +147,16 @@ type mcpListingKey struct { agent string } -type mcpListing struct { - invocation string - tools []tool.Tool +type mcpListingState struct { + // processed reports that the toolset's ProcessRequest has run for the + // agent. From then on ProcessRequest lists, and Tools does not. + processed bool + // latest is the most recent list, which Tools returns once processed. + latest []tool.Tool + // pending holds, per invocation, the list Tools took for the + // ProcessRequest call that follows in the same model step. Concurrent + // invocations of one agent each keep their own. + pending map[string][]tool.Tool } // listingsFrom returns the mcpListings NewContext installed, or nil. @@ -142,30 +165,75 @@ func listingsFrom(ctx context.Context) *mcpListings { return l } -func (l *mcpListings) put(m *mcpToolset, ctx agent.ReadonlyContext, tools []tool.Tool) { +func (l *mcpListings) entry(m *mcpToolset, ctx agent.ReadonlyContext) *mcpListingState { + if l.state == nil { + l.state = map[mcpListingKey]*mcpListingState{} + } + k := mcpListingKey{m, ctx.AgentName()} + s, ok := l.state[k] + if !ok { + s = &mcpListingState{pending: map[string][]tool.Tool{}} + l.state[k] = s + } + return s +} + +// cached returns the latest list when the toolset's ProcessRequest has run for +// the agent, and so lists in its place. It also drops a list Tools kept for +// this invocation that no ProcessRequest took. +func (l *mcpListings) cached(m *mcpToolset, ctx agent.ReadonlyContext) ([]tool.Tool, bool) { + if l == nil { + return nil, false + } + l.mu.Lock() + defer l.mu.Unlock() + s := l.entry(m, ctx) + delete(s.pending, ctx.InvocationID()) + if !s.processed || s.latest == nil { + return nil, false + } + return s.latest, true +} + +// listed records a list the toolset took. With pending, it keeps the list for +// the ProcessRequest call of the same invocation. +func (l *mcpListings) listed(m *mcpToolset, ctx agent.ReadonlyContext, tools []tool.Tool, pending bool) { if l == nil { return } l.mu.Lock() defer l.mu.Unlock() - if l.pending == nil { - l.pending = map[mcpListingKey]mcpListing{} + s := l.entry(m, ctx) + s.latest = tools + if pending { + s.pending[ctx.InvocationID()] = tools } - l.pending[mcpListingKey{m, ctx.AgentName()}] = mcpListing{ctx.InvocationID(), tools} } -// take returns and forgets the list Tools kept for this toolset, agent and -// invocation. +// take marks the toolset processed for the agent and returns the list Tools +// kept for this invocation, if any. func (l *mcpListings) take(m *mcpToolset, ctx agent.ReadonlyContext) ([]tool.Tool, bool) { if l == nil { return nil, false } l.mu.Lock() defer l.mu.Unlock() - k := mcpListingKey{m, ctx.AgentName()} - p, ok := l.pending[k] - delete(l.pending, k) - return p.tools, ok && p.invocation == ctx.InvocationID() + s := l.entry(m, ctx) + s.processed = true + tools, ok := s.pending[ctx.InvocationID()] + delete(s.pending, ctx.InvocationID()) + return tools, ok +} + +// packedByToolset reports whether the toolset's ProcessRequest has run for the +// agent. Its proxy tools then leave packing to it. +func (l *mcpListings) packedByToolset(m *mcpToolset, ctx agent.ReadonlyContext) bool { + if l == nil { + return false + } + l.mu.Lock() + defer l.mu.Unlock() + return l.entry(m, ctx).processed } // mcpProxyTool is a declaration-only stand-in for a remote MCP tool. The model @@ -173,6 +241,7 @@ func (l *mcpListings) take(m *mcpToolset, ctx agent.ReadonlyContext) ([]tool.Too // ADK) dispatches the CallMcpTool Activity so the live, stateful MCP session // runs worker-side, never in the workflow. type mcpProxyTool struct { + owner *mcpToolset toolset string decl *genai.FunctionDeclaration activityOptions workflow.ActivityOptions @@ -183,7 +252,14 @@ func (t *mcpProxyTool) Description() string { return t.decl. func (t *mcpProxyTool) IsLongRunning() bool { return false } func (t *mcpProxyTool) Declaration() *genai.FunctionDeclaration { return t.decl } +// ProcessRequest packs the proxy tool. ADK calls it for each tool it kept from +// a Tools call. Once the owning toolset's ProcessRequest packs the current list +// for the agent, it does nothing: the tool may come from an older list, which +// ADK v2.5.0 and earlier keep for the whole agent run. func (t *mcpProxyTool) ProcessRequest(ctx agent.Context, req *model.LLMRequest) error { + if listingsFrom(ctx).packedByToolset(t.owner, ctx) { + return nil + } return packTool(req, t) } diff --git a/go/googleadk/mcp_test.go b/go/googleadk/mcp_test.go index 9ed11d1c..4bba900e 100644 --- a/go/googleadk/mcp_test.go +++ b/go/googleadk/mcp_test.go @@ -2,18 +2,31 @@ package googleadk_test import ( "context" + "errors" + "fmt" "iter" "sync" + "sync/atomic" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.temporal.io/sdk/activity" + "go.temporal.io/sdk/converter" "go.temporal.io/sdk/testsuite" + "go.temporal.io/sdk/workflow" "google.golang.org/genai" + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/agent/llmagent" "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/runner" + "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/tool" + "google.golang.org/adk/v2/tool/agenttool" "go.temporal.io/googleadk" ) @@ -168,3 +181,349 @@ func TestMcpToolCallRoundTrip(t *testing.T) { assert.Equal(t, "meaning of life", gotArgs["query"]) mu.Unlock() } + +// ruleModel answers each request with the response its rule picks from the +// request, so concurrent runs that share it each get their own script. +type ruleModel struct { + name string + rule func(req *model.LLMRequest) *model.LLMResponse + + mu sync.Mutex + decls [][]string +} + +func (m *ruleModel) Name() string { return m.name } + +func (m *ruleModel) GenerateContent(_ context.Context, req *model.LLMRequest, _ bool) iter.Seq2[*model.LLMResponse, error] { + var names []string + if req.Config != nil { + for _, gt := range req.Config.Tools { + if gt != nil { + for _, d := range gt.FunctionDeclarations { + names = append(names, d.Name) + } + } + } + } + m.mu.Lock() + m.decls = append(m.decls, names) + m.mu.Unlock() + resp := m.rule(req) + return func(yield func(*model.LLMResponse, error) bool) { yield(resp, nil) } +} + +func (m *ruleModel) factory() googleadk.ModelFactory { + return func(context.Context, string) (model.LLM, error) { return m, nil } +} + +// declared returns the tool names each model request declared, in order. +func (m *ruleModel) declared() [][]string { + m.mu.Lock() + defer m.mu.Unlock() + return m.decls +} + +// functionResponses counts the function responses in a request's contents. +func functionResponses(req *model.LLMRequest) int { + n := 0 + for _, c := range req.Contents { + if c == nil { + continue + } + for _, p := range c.Parts { + if p.FunctionResponse != nil { + n++ + } + } + } + return n +} + +// lastIsFunctionResponse reports whether the request's last content carries a +// function response. +func lastIsFunctionResponse(req *model.LLMRequest) bool { + if len(req.Contents) == 0 { + return false + } + for _, p := range req.Contents[len(req.Contents)-1].Parts { + if p.FunctionResponse != nil { + return true + } + } + return false +} + +// lastUserText returns the text of the request's last user content. +func lastUserText(req *model.LLMRequest) string { + for i := len(req.Contents) - 1; i >= 0; i-- { + c := req.Contents[i] + if c == nil || c.Role != genai.RoleUser { + continue + } + for _, p := range c.Parts { + if p.Text != "" { + return p.Text + } + } + } + return "" +} + +// commandLog records the Activities and timers a test workflow schedules, in +// order. +type commandLog struct { + mu sync.Mutex + log []string +} + +func (l *commandLog) add(s string) { + l.mu.Lock() + defer l.mu.Unlock() + l.log = append(l.log, s) +} + +func (l *commandLog) get() []string { + l.mu.Lock() + defer l.mu.Unlock() + return append([]string(nil), l.log...) +} + +func (l *commandLog) count(s string) int { + n := 0 + for _, e := range l.get() { + if e == s { + n++ + } + } + return n +} + +// watchCommands records env's Activity starts and timers in order. It replaces +// the activityCounter listener newEnv installs. +func watchCommands(env *testsuite.TestWorkflowEnvironment) *commandLog { + l := &commandLog{} + env.SetOnActivityStartedListener(func(info *activity.Info, _ context.Context, _ converter.EncodedValues) { + l.add(info.ActivityType.Name) + }) + env.SetOnTimerScheduledListener(func(string, time.Duration) { l.add("Timer") }) + return l +} + +// sleepingInstructionWorkflow runs two turns of an MCP agent whose +// InstructionProvider sleeps on a workflow timer before each model call, the +// way an instruction that reads workflow state through a command would. +func sleepingInstructionWorkflow(ctx workflow.Context) error { + root, err := llmagent.New(llmagent.Config{ + Name: "assistant", + Model: googleadk.NewModel("fake-model"), + InstructionProvider: func(rctx agent.ReadonlyContext) (string, error) { + wfCtx, ok := googleadk.WorkflowContext(rctx) + if !ok { + return "", errors.New("no workflow context") + } + if err := workflow.Sleep(wfCtx, time.Second); err != nil { + return "", err + } + return "be helpful", nil + }, + Toolsets: []tool.Toolset{googleadk.NewMCPToolset(googleadk.MCPToolsetOptions{Name: "search-server"})}, + }) + if err != nil { + return err + } + r, err := runner.New(runner.Config{AppName: "test-app", Agent: root, SessionService: session.InMemoryService(), AutoCreateSession: true}) + if err != nil { + return err + } + adkCtx := googleadk.NewContext(ctx) + for _, text := range []string{"search", "thanks"} { + for _, rerr := range r.Run(adkCtx, "user-1", "session-1", genai.NewContentFromText(text, genai.RoleUser), agent.RunConfig{}) { + if rerr != nil { + return rerr + } + } + } + return nil +} + +// TestMcpListsAtTheSamePointInEveryStep proves the MCP toolset lists at a point +// that does not depend on the ADK release: in the first model step of the +// first turn before the request processors (ADK calls Tools there on every +// release), and in every later step and turn after them, in ProcessRequest. +// ADK v2.6.0 also calls Tools before the request processors of every step; a +// listing there would put ListMcpTools before the instruction's timer, where +// ADK v2.5.0 and earlier put it after, and Histories would not replay across +// releases. +func TestMcpListsAtTheSamePointInEveryStep(t *testing.T) { + m := &ruleModel{name: "fake-model", rule: func(req *model.LLMRequest) *model.LLMResponse { + if lastUserText(req) == "search" && !lastIsFunctionResponse(req) { + return googleadk.FunctionCallResponse("c1", "search", map[string]any{"query": "temporal"}) + } + return googleadk.TextResponse("done") + }} + server := googleadk.NewFakeMCPServer("search-server").AddTool( + "search", "search the web", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{"hits": 1}, nil }, + ) + + var s testsuite.WorkflowTestSuite + env, _ := newEnv(t, &s, googleadk.Config{ + Models: map[string]googleadk.ModelFactory{"fake-model": m.factory()}, + MCPToolsets: map[string]googleadk.MCPFactory{"search-server": server.Factory()}, + }) + commands := watchCommands(env) + env.RegisterWorkflow(sleepingInstructionWorkflow) + env.ExecuteWorkflow(sleepingInstructionWorkflow) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + list, invoke, call := googleadk.ListMcpToolsActivityName, googleadk.InvokeModelActivityName, googleadk.CallMcpToolActivityName + assert.Equal(t, []string{ + // Turn 1, step 1: listed through Tools, before the instruction. + list, "Timer", invoke, call, + // Turn 1, step 2: listed in ProcessRequest, after the instruction. + "Timer", list, invoke, + // Turn 2: listed in ProcessRequest. + "Timer", list, invoke, + }, commands.get()) + for i, names := range m.declared() { + assert.Equal(t, []string{"search"}, names, "model request %d", i) + } +} + +// hidingToolset hides the "search" tool of the wrapped MCP server once hide is +// set, the way a server drops a tool from its list. +type hidingToolset struct { + tool.Toolset + hide *atomic.Bool +} + +func (h hidingToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { + tools, err := h.Toolset.Tools(ctx) + if err != nil || !h.hide.Load() { + return tools, err + } + var kept []tool.Tool + for _, t := range tools { + if t.Name() != "search" { + kept = append(kept, t) + } + } + return kept, nil +} + +// TestMcpDroppedToolIsNeitherDeclaredNorCalled proves each model request +// declares and dispatches the toolset's current list. ADK v2.5.0 and earlier +// keep the tools from the first Tools call of an agent run; if their proxies +// still packed themselves, a tool the server dropped would stay declared and +// a call to it would run CallMcpTool there, but not with ADK v2.6.0. +func TestMcpDroppedToolIsNeitherDeclaredNorCalled(t *testing.T) { + var hide atomic.Bool + server := googleadk.NewFakeMCPServer("search-server"). + AddTool("search", "search the web", searchSchema(), + func(map[string]any) (map[string]any, error) { + hide.Store(true) + return map[string]any{"hits": 1}, nil + }). + AddTool("fetch", "fetch a page", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{}, nil }) + m := &ruleModel{name: "fake-model", rule: func(req *model.LLMRequest) *model.LLMResponse { + if n := functionResponses(req); n < 2 { + return googleadk.FunctionCallResponse(fmt.Sprintf("c%d", n+1), "search", map[string]any{"query": "temporal"}) + } + return googleadk.TextResponse("done") + }} + + var s testsuite.WorkflowTestSuite + env, counter := newEnv(t, &s, googleadk.Config{ + Models: map[string]googleadk.ModelFactory{"fake-model": m.factory()}, + MCPToolsets: map[string]googleadk.MCPFactory{"search-server": func(context.Context) (tool.Toolset, error) { + return hidingToolset{Toolset: server, hide: &hide}, nil + }}, + }) + env.ExecuteWorkflow(agentRunWorkflow, runInput{ModelName: "fake-model", UserMessage: "search twice", MCPToolset: "search-server"}) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + assert.Equal(t, [][]string{{"search", "fetch"}, {"fetch"}, {"fetch"}}, m.declared()) + assert.Equal(t, 3, counter.get(googleadk.ListMcpToolsActivityName)) + assert.Equal(t, 1, counter.get(googleadk.CallMcpToolActivityName), "the dropped tool ran again") +} + +// parallelHelperWorkflow has the root agent call one helper agent twice in +// parallel through an agent tool. The helper has two MCP toolsets. +func parallelHelperWorkflow(ctx workflow.Context, sequential bool) (runResult, error) { + helper, err := llmagent.New(llmagent.Config{ + Name: "helper", + Description: "searches", + Model: googleadk.NewModel("helper-model"), + Instruction: "search", + Toolsets: []tool.Toolset{ + googleadk.NewMCPToolset(googleadk.MCPToolsetOptions{Name: "search-server"}), + googleadk.NewMCPToolset(googleadk.MCPToolsetOptions{Name: "fetch-server"}), + }, + }) + if err != nil { + return runResult{}, err + } + var opts []googleadk.ContextOption + if sequential { + opts = append(opts, googleadk.WithSequentialToolFanout()) + } + return runAgent(ctx, agentBuild{ + ctxOpts: opts, + modelName: "root-model", + userMessage: "ask the helper twice", + tools: []tool.Tool{agenttool.New(helper, nil)}, + }) +} + +// TestMcpParallelInvocationsKeepTheirOwnListings proves concurrent invocations +// of one agent each reuse the list their own Tools call took. Two helper runs +// of two model steps each, with two MCP toolsets, list 8 times, in parallel as +// in sequence. +func TestMcpParallelInvocationsKeepTheirOwnListings(t *testing.T) { + for _, sequential := range []bool{false, true} { + t.Run(fmt.Sprintf("sequential=%v", sequential), func(t *testing.T) { + root := &ruleModel{name: "root-model", rule: func(req *model.LLMRequest) *model.LLMResponse { + if functionResponses(req) == 0 { + return reproFnCalls("helper", "helper") + } + return googleadk.TextResponse("root done") + }} + helper := &ruleModel{name: "helper-model", rule: func(req *model.LLMRequest) *model.LLMResponse { + if lastIsFunctionResponse(req) { + return googleadk.TextResponse("helper done") + } + return googleadk.FunctionCallResponse("h1", "search", map[string]any{"query": "temporal"}) + }} + search := googleadk.NewFakeMCPServer("search-server").AddTool("search", "search the web", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{"hits": 1}, nil }) + fetch := googleadk.NewFakeMCPServer("fetch-server").AddTool("fetch", "fetch a page", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{}, nil }) + + var s testsuite.WorkflowTestSuite + env, counter := newEnv(t, &s, googleadk.Config{ + Models: map[string]googleadk.ModelFactory{"root-model": root.factory(), "helper-model": helper.factory()}, + MCPToolsets: map[string]googleadk.MCPFactory{ + "search-server": search.Factory(), + "fetch-server": fetch.Factory(), + }, + }) + env.RegisterWorkflow(parallelHelperWorkflow) + env.ExecuteWorkflow(parallelHelperWorkflow, sequential) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + assert.Equal(t, 6, counter.get(googleadk.InvokeModelActivityName)) + assert.Equal(t, 8, counter.get(googleadk.ListMcpToolsActivityName), "one listing per toolset per helper model step") + assert.Equal(t, 2, counter.get(googleadk.CallMcpToolActivityName)) + for i, names := range helper.declared() { + assert.ElementsMatch(t, []string{"search", "fetch"}, names, "helper request %d", i) + } + }) + } +} From 5b5117605bad4a3035d83c636a0cdb21679028fc Mon Sep 17 00:00:00 2001 From: Francesc Campoy Date: Fri, 30 Oct 2026 10:33:18 -0700 Subject: [PATCH 03/10] Address review findings --- go/googleadk/mcp.go | 35 ++++++++++++++++++++-------- go/googleadk/mcp_hitl_test.go | 43 +++++++++++++++++++++++++++++++---- 2 files changed, 65 insertions(+), 13 deletions(-) diff --git a/go/googleadk/mcp.go b/go/googleadk/mcp.go index f6d1d4be..4cf2e1cd 100644 --- a/go/googleadk/mcp.go +++ b/go/googleadk/mcp.go @@ -64,21 +64,40 @@ func (m *mcpToolset) Name() string { return m.name } // agent, and keeps that list for the ProcessRequest call that follows in the // same model step. After that, ProcessRequest lists, and Tools returns the // latest list without an Activity. See ProcessRequest. +// +// ToolFilter runs on each call, not when the list is taken, so a list kept +// from an earlier step or turn returns only the tools the filter allows now. +// ADK resumes a confirmed tool call from this list before any ProcessRequest +// of the resuming agent run. func (m *mcpToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { l := listingsFrom(ctx) if tools, ok := l.cached(m, ctx); ok { - return tools, nil + return m.allowed(ctx, tools), nil } tools, err := m.list(ctx) if err != nil { return nil, err } l.listed(m, ctx, tools, true) - return tools, nil + return m.allowed(ctx, tools), nil +} + +// allowed returns the tools the toolset's ToolFilter allows for ctx. +func (m *mcpToolset) allowed(ctx agent.ReadonlyContext, tools []tool.Tool) []tool.Tool { + if m.filter == nil { + return tools + } + out := make([]tool.Tool, 0, len(tools)) + for _, t := range tools { + if m.filter(ctx, t) { + out = append(out, t) + } + } + return out } // list runs the ListMcpTools Activity and wraps each declaration as a proxy -// tool. +// tool. It does not apply ToolFilter; see allowed. func (m *mcpToolset) list(ctx agent.ReadonlyContext) ([]tool.Tool, error) { wfCtx, ok := workflowContext(ctx) if !ok { @@ -91,10 +110,7 @@ func (m *mcpToolset) list(ctx agent.ReadonlyContext) ([]tool.Tool, error) { } tools := make([]tool.Tool, 0, len(decls)) for _, d := range decls { - pt := &mcpProxyTool{owner: m, toolset: m.name, decl: d, activityOptions: m.activityOptions} - if m.filter == nil || m.filter(ctx, pt) { - tools = append(tools, pt) - } + tools = append(tools, &mcpProxyTool{owner: m, toolset: m.name, decl: d, activityOptions: m.activityOptions}) } return tools, nil } @@ -122,7 +138,7 @@ func (m *mcpToolset) ProcessRequest(ctx agent.Context, req *model.LLMRequest) er } l.listed(m, ctx, tools, false) } - for _, t := range tools { + for _, t := range m.allowed(ctx, tools) { if err := packTool(req, t.(*mcpProxyTool)); err != nil { return err } @@ -151,7 +167,8 @@ type mcpListingState struct { // processed reports that the toolset's ProcessRequest has run for the // agent. From then on ProcessRequest lists, and Tools does not. processed bool - // latest is the most recent list, which Tools returns once processed. + // latest is the most recent list, before ToolFilter, which Tools returns + // once processed. latest []tool.Tool // pending holds, per invocation, the list Tools took for the // ProcessRequest call that follows in the same model step. Concurrent diff --git a/go/googleadk/mcp_hitl_test.go b/go/googleadk/mcp_hitl_test.go index 995662a9..a99ec5d5 100644 --- a/go/googleadk/mcp_hitl_test.go +++ b/go/googleadk/mcp_hitl_test.go @@ -82,13 +82,22 @@ type mcpHitlResult struct { // ADK re-dispatches the original call with the decision attached; the // snapshot carries it to the worker-side tool, which executes on // approval or rejects on denial. -func mcpHitlWorkflow(ctx workflow.Context, confirm bool) (mcpHitlResult, error) { +// +// With excludeBeforeResume, the toolset's ToolFilter stops allowing the guarded +// tool after pass 1, so the approved call must not run on resume. +func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume bool) (mcpHitlResult, error) { + excluded := false root, err := llmagent.New(llmagent.Config{ Name: "assistant", Description: "root assistant", Model: googleadk.NewModel("fake-model"), Instruction: "be helpful", - Toolsets: []tool.Toolset{googleadk.NewMCPToolset(googleadk.MCPToolsetOptions{Name: "guarded-server"})}, + Toolsets: []tool.Toolset{googleadk.NewMCPToolset(googleadk.MCPToolsetOptions{ + Name: "guarded-server", + ToolFilter: func(_ agent.ReadonlyContext, t tool.Tool) bool { + return !excluded || t.Name() != "guarded" + }, + })}, }) if err != nil { return mcpHitlResult{}, err @@ -129,6 +138,7 @@ func mcpHitlWorkflow(ctx workflow.Context, confirm bool) (mcpHitlResult, error) } // --- Pass 2: simulate the human's decision, then resume. --- + excluded = excludeBeforeResume resume := googleadk.ConfirmationResponse(googleadk.ConfirmationDecision{ FunctionCallID: pending[0].FunctionCallID, Confirmed: confirm, @@ -189,7 +199,7 @@ func TestMCPToolConfirmationPausesAndResumes(t *testing.T) { var s testsuite.WorkflowTestSuite env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, true) + env.ExecuteWorkflow(mcpHitlWorkflow, true, false) require.True(t, env.IsWorkflowCompleted()) require.NoError(t, env.GetWorkflowError()) @@ -215,7 +225,7 @@ func TestMCPToolConfirmationDenialBlocksTool(t *testing.T) { var s testsuite.WorkflowTestSuite env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, false) + env.ExecuteWorkflow(mcpHitlWorkflow, false, false) require.True(t, env.IsWorkflowCompleted()) require.NoError(t, env.GetWorkflowError()) @@ -230,3 +240,28 @@ func TestMCPToolConfirmationDenialBlocksTool(t *testing.T) { assert.Equal(t, 2, counter.get(googleadk.CallMcpToolActivityName), "the denial is delivered to the worker-side tool, which rejects without running the handler") } + +// TestMCPToolConfirmationResumeHonorsCurrentToolFilter proves that a tool the +// ToolFilter excludes while its confirmation is pending does not run when the +// human approves: ADK resumes the call from the toolset's current Tools list, +// which the filter governs at that point, not when the tools were listed. +func TestMCPToolConfirmationResumeHonorsCurrentToolFilter(t *testing.T) { + mcpGuardedRuns.Store(0) + + var s testsuite.WorkflowTestSuite + env, counter := newMcpHitlEnv(t, &s) + env.ExecuteWorkflow(mcpHitlWorkflow, true, true) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + var res mcpHitlResult + require.NoError(t, env.GetWorkflowResult(&res)) + + require.True(t, res.Paused, "agent must pause on a pending confirmation") + assert.EqualValues(t, 0, mcpGuardedRuns.Load(), "an excluded tool must not run after approval") + assert.Contains(t, res.ResumedGuardedError, "not found", + "ADK must not find the excluded tool when it resumes the call") + assert.Equal(t, 1, counter.get(googleadk.CallMcpToolActivityName), + "only the pending pass reaches the worker") +} From 10601aac40eaef9a9bc86a1cbec991a61ec94b58 Mon Sep 17 00:00:00 2001 From: Francesc Campoy Date: Fri, 30 Oct 2026 10:59:19 -0700 Subject: [PATCH 04/10] Address review findings --- go/googleadk/mcp.go | 79 +++++++++++++++++++++----- go/googleadk/mcp_test.go | 120 +++++++++++++++++++++++++-------------- 2 files changed, 142 insertions(+), 57 deletions(-) diff --git a/go/googleadk/mcp.go b/go/googleadk/mcp.go index 4cf2e1cd..3df91c4b 100644 --- a/go/googleadk/mcp.go +++ b/go/googleadk/mcp.go @@ -65,21 +65,29 @@ func (m *mcpToolset) Name() string { return m.name } // same model step. After that, ProcessRequest lists, and Tools returns the // latest list without an Activity. See ProcessRequest. // -// ToolFilter runs on each call, not when the list is taken, so a list kept -// from an earlier step or turn returns only the tools the filter allows now. -// ADK resumes a confirmed tool call from this list before any ProcessRequest -// of the resuming agent run. +// ToolFilter runs on the first call of each agent run, not when the list is +// taken, so a list kept from an earlier turn returns only the tools the filter +// allows now. ADK resumes a confirmed tool call from this list before any +// ProcessRequest of the resuming run. Later calls in the run, which ADK v2.6.0 +// and later make before every model step, return the list the filter last +// allowed without running it again: a filter that issues Workflow commands +// then issues them at the same points on every ADK release. func (m *mcpToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { l := listingsFrom(ctx) - if tools, ok := l.cached(m, ctx); ok { - return m.allowed(ctx, tools), nil + if tools, ok := l.filteredBefore(m, ctx); ok { + return tools, nil } - tools, err := m.list(ctx) - if err != nil { - return nil, err + tools, ok := l.cached(m, ctx) + if !ok { + var err error + if tools, err = m.list(ctx); err != nil { + return nil, err + } + l.listed(m, ctx, tools, true) } - l.listed(m, ctx, tools, true) - return m.allowed(ctx, tools), nil + tools = m.allowed(ctx, tools) + l.filtered(m, ctx, tools, true) + return tools, nil } // allowed returns the tools the toolset's ToolFilter allows for ctx. @@ -138,7 +146,9 @@ func (m *mcpToolset) ProcessRequest(ctx agent.Context, req *model.LLMRequest) er } l.listed(m, ctx, tools, false) } - for _, t := range m.allowed(ctx, tools) { + tools = m.allowed(ctx, tools) + l.filtered(m, ctx, tools, false) + for _, t := range tools { if err := packTool(req, t.(*mcpProxyTool)); err != nil { return err } @@ -174,6 +184,22 @@ type mcpListingState struct { // ProcessRequest call that follows in the same model step. Concurrent // invocations of one agent each keep their own. pending map[string][]tool.Tool + // allowed holds, per agent run whose first Tools call ran ToolFilter, the + // list the filter last allowed there. Tools returns it on the later calls + // of the run. + allowed map[mcpRunKey][]tool.Tool +} + +// mcpRunKey identifies an agent run. A run that resumes a paused one keeps +// its invocation ID, so the key adds the user content that started the run: +// ADK passes the same content to every step of a run. +type mcpRunKey struct { + invocation string + content *genai.Content +} + +func runKey(ctx agent.ReadonlyContext) mcpRunKey { + return mcpRunKey{ctx.InvocationID(), ctx.UserContent()} } // listingsFrom returns the mcpListings NewContext installed, or nil. @@ -189,7 +215,7 @@ func (l *mcpListings) entry(m *mcpToolset, ctx agent.ReadonlyContext) *mcpListin k := mcpListingKey{m, ctx.AgentName()} s, ok := l.state[k] if !ok { - s = &mcpListingState{pending: map[string][]tool.Tool{}} + s = &mcpListingState{pending: map[string][]tool.Tool{}, allowed: map[mcpRunKey][]tool.Tool{}} l.state[k] = s } return s @@ -227,6 +253,33 @@ func (l *mcpListings) listed(m *mcpToolset, ctx agent.ReadonlyContext, tools []t } } +// filteredBefore returns the list ToolFilter last allowed when Tools already +// ran it in this agent run. +func (l *mcpListings) filteredBefore(m *mcpToolset, ctx agent.ReadonlyContext) ([]tool.Tool, bool) { + if l == nil { + return nil, false + } + l.mu.Lock() + defer l.mu.Unlock() + tools, ok := l.entry(m, ctx).allowed[runKey(ctx)] + return tools, ok +} + +// filtered records a list ToolFilter allowed in the agent run. Tools records +// the first; ProcessRequest replaces it once Tools has. +func (l *mcpListings) filtered(m *mcpToolset, ctx agent.ReadonlyContext, tools []tool.Tool, inTools bool) { + if l == nil { + return + } + l.mu.Lock() + defer l.mu.Unlock() + s := l.entry(m, ctx) + k := runKey(ctx) + if _, ok := s.allowed[k]; ok || inTools { + s.allowed[k] = tools + } +} + // take marks the toolset processed for the agent and returns the list Tools // kept for this invocation, if any. func (l *mcpListings) take(m *mcpToolset, ctx agent.ReadonlyContext) ([]tool.Tool, bool) { diff --git a/go/googleadk/mcp_test.go b/go/googleadk/mcp_test.go index 4bba900e..ad106225 100644 --- a/go/googleadk/mcp_test.go +++ b/go/googleadk/mcp_test.go @@ -298,35 +298,46 @@ func (l *commandLog) count(s string) int { return n } -// watchCommands records env's Activity starts and timers in order. It replaces -// the activityCounter listener newEnv installs. +// watchCommands records env's Activity starts and timers in order, each timer +// as "Timer ". It replaces the activityCounter listener newEnv +// installs. func watchCommands(env *testsuite.TestWorkflowEnvironment) *commandLog { l := &commandLog{} env.SetOnActivityStartedListener(func(info *activity.Info, _ context.Context, _ converter.EncodedValues) { l.add(info.ActivityType.Name) }) - env.SetOnTimerScheduledListener(func(string, time.Duration) { l.add("Timer") }) + env.SetOnTimerScheduledListener(func(_ string, d time.Duration) { l.add("Timer " + d.String()) }) return l } // sleepingInstructionWorkflow runs two turns of an MCP agent whose -// InstructionProvider sleeps on a workflow timer before each model call, the -// way an instruction that reads workflow state through a command would. -func sleepingInstructionWorkflow(ctx workflow.Context) error { +// InstructionProvider sleeps on a 1s workflow timer before each model call, the +// way an instruction that reads workflow state through a command would. With +// filterSleeps, the toolset's ToolFilter also sleeps, on a 2s timer. +func sleepingInstructionWorkflow(ctx workflow.Context, filterSleeps bool) error { + sleep := func(rctx agent.ReadonlyContext, d time.Duration) error { + wfCtx, ok := googleadk.WorkflowContext(rctx) + if !ok { + return errors.New("no workflow context") + } + return workflow.Sleep(wfCtx, d) + } + opts := googleadk.MCPToolsetOptions{Name: "search-server"} + if filterSleeps { + opts.ToolFilter = func(rctx agent.ReadonlyContext, _ tool.Tool) bool { + return sleep(rctx, 2*time.Second) == nil + } + } root, err := llmagent.New(llmagent.Config{ Name: "assistant", Model: googleadk.NewModel("fake-model"), InstructionProvider: func(rctx agent.ReadonlyContext) (string, error) { - wfCtx, ok := googleadk.WorkflowContext(rctx) - if !ok { - return "", errors.New("no workflow context") - } - if err := workflow.Sleep(wfCtx, time.Second); err != nil { + if err := sleep(rctx, time.Second); err != nil { return "", err } return "be helpful", nil }, - Toolsets: []tool.Toolset{googleadk.NewMCPToolset(googleadk.MCPToolsetOptions{Name: "search-server"})}, + Toolsets: []tool.Toolset{googleadk.NewMCPToolset(opts)}, }) if err != nil { return err @@ -353,42 +364,63 @@ func sleepingInstructionWorkflow(ctx workflow.Context) error { // ADK v2.6.0 also calls Tools before the request processors of every step; a // listing there would put ListMcpTools before the instruction's timer, where // ADK v2.5.0 and earlier put it after, and Histories would not replay across -// releases. +// releases. ToolFilter runs at the same points too: in the first Tools call of +// each turn and in each ProcessRequest, never in the later Tools calls only +// ADK v2.6.0 makes. func TestMcpListsAtTheSamePointInEveryStep(t *testing.T) { - m := &ruleModel{name: "fake-model", rule: func(req *model.LLMRequest) *model.LLMResponse { - if lastUserText(req) == "search" && !lastIsFunctionResponse(req) { - return googleadk.FunctionCallResponse("c1", "search", map[string]any{"query": "temporal"}) - } - return googleadk.TextResponse("done") - }} - server := googleadk.NewFakeMCPServer("search-server").AddTool( - "search", "search the web", searchSchema(), - func(map[string]any) (map[string]any, error) { return map[string]any{"hits": 1}, nil }, - ) + list, invoke, call := googleadk.ListMcpToolsActivityName, googleadk.InvokeModelActivityName, googleadk.CallMcpToolActivityName + instr, filter := "Timer 1s", "Timer 2s" + for _, tc := range []struct { + filterSleeps bool + want []string + }{ + {false, []string{ + // Turn 1, step 1: listed through Tools, before the instruction. + list, instr, invoke, call, + // Turn 1, step 2: listed in ProcessRequest, after the instruction. + instr, list, invoke, + // Turn 2: listed in ProcessRequest. + instr, list, invoke, + }}, + {true, []string{ + // Turn 1, step 1: Tools lists and filters, ProcessRequest filters. + list, filter, instr, filter, invoke, call, + // Turn 1, step 2: ProcessRequest lists and filters. + instr, list, filter, invoke, + // Turn 2: Tools filters the kept list, ProcessRequest lists and + // filters. + filter, instr, list, filter, invoke, + }}, + } { + t.Run(fmt.Sprintf("filterSleeps=%v", tc.filterSleeps), func(t *testing.T) { + m := &ruleModel{name: "fake-model", rule: func(req *model.LLMRequest) *model.LLMResponse { + if lastUserText(req) == "search" && !lastIsFunctionResponse(req) { + return googleadk.FunctionCallResponse("c1", "search", map[string]any{"query": "temporal"}) + } + return googleadk.TextResponse("done") + }} + server := googleadk.NewFakeMCPServer("search-server").AddTool( + "search", "search the web", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{"hits": 1}, nil }, + ) - var s testsuite.WorkflowTestSuite - env, _ := newEnv(t, &s, googleadk.Config{ - Models: map[string]googleadk.ModelFactory{"fake-model": m.factory()}, - MCPToolsets: map[string]googleadk.MCPFactory{"search-server": server.Factory()}, - }) - commands := watchCommands(env) - env.RegisterWorkflow(sleepingInstructionWorkflow) - env.ExecuteWorkflow(sleepingInstructionWorkflow) + var s testsuite.WorkflowTestSuite + env, _ := newEnv(t, &s, googleadk.Config{ + Models: map[string]googleadk.ModelFactory{"fake-model": m.factory()}, + MCPToolsets: map[string]googleadk.MCPFactory{"search-server": server.Factory()}, + }) + commands := watchCommands(env) + env.RegisterWorkflow(sleepingInstructionWorkflow) + env.ExecuteWorkflow(sleepingInstructionWorkflow, tc.filterSleeps) - require.True(t, env.IsWorkflowCompleted()) - require.NoError(t, env.GetWorkflowError()) + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) - list, invoke, call := googleadk.ListMcpToolsActivityName, googleadk.InvokeModelActivityName, googleadk.CallMcpToolActivityName - assert.Equal(t, []string{ - // Turn 1, step 1: listed through Tools, before the instruction. - list, "Timer", invoke, call, - // Turn 1, step 2: listed in ProcessRequest, after the instruction. - "Timer", list, invoke, - // Turn 2: listed in ProcessRequest. - "Timer", list, invoke, - }, commands.get()) - for i, names := range m.declared() { - assert.Equal(t, []string{"search"}, names, "model request %d", i) + assert.Equal(t, tc.want, commands.get()) + for i, names := range m.declared() { + assert.Equal(t, []string{"search"}, names, "model request %d", i) + } + }) } } From a47ce349dc512e6d0635665d772382b13409a5e0 Mon Sep 17 00:00:00 2001 From: Francesc Campoy Date: Fri, 30 Oct 2026 11:27:44 -0700 Subject: [PATCH 05/10] Address review findings --- go/googleadk/mcp.go | 41 +++++++---- go/googleadk/mcp_hitl_test.go | 84 ++++++++++++--------- go/googleadk/mcp_test.go | 135 +++++++++++++++++++++++++++++++--- 3 files changed, 205 insertions(+), 55 deletions(-) diff --git a/go/googleadk/mcp.go b/go/googleadk/mcp.go index 3df91c4b..584ea91a 100644 --- a/go/googleadk/mcp.go +++ b/go/googleadk/mcp.go @@ -2,6 +2,8 @@ package googleadk import ( "context" + "crypto/sha256" + "encoding/json" "errors" "fmt" "sync" @@ -72,6 +74,10 @@ func (m *mcpToolset) Name() string { return m.name } // and later make before every model step, return the list the filter last // allowed without running it again: a filter that issues Workflow commands // then issues them at the same points on every ADK release. +// +// A toolset wrapped in a toolset without ProcessRequest, such as +// tool.FilterToolset, lists and filters on every call, because nothing else +// refreshes its list. func (m *mcpToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { l := listingsFrom(ctx) if tools, ok := l.filteredBefore(m, ctx); ok { @@ -184,22 +190,26 @@ type mcpListingState struct { // ProcessRequest call that follows in the same model step. Concurrent // invocations of one agent each keep their own. pending map[string][]tool.Tool - // allowed holds, per agent run whose first Tools call ran ToolFilter, the - // list the filter last allowed there. Tools returns it on the later calls - // of the run. + // allowed holds, per agent run since processed, the list ToolFilter last + // allowed there. Tools returns it on the later calls of the run. allowed map[mcpRunKey][]tool.Tool } // mcpRunKey identifies an agent run. A run that resumes a paused one keeps // its invocation ID, so the key adds the user content that started the run: -// ADK passes the same content to every step of a run. +// ADK passes the same content to every step of a run. The content is compared +// by value too, because a caller may resume with the same *genai.Content, +// overwritten with the approval. type mcpRunKey struct { invocation string content *genai.Content + digest [sha256.Size]byte } func runKey(ctx agent.ReadonlyContext) mcpRunKey { - return mcpRunKey{ctx.InvocationID(), ctx.UserContent()} + c := ctx.UserContent() + b, _ := json.Marshal(c) + return mcpRunKey{ctx.InvocationID(), c, sha256.Sum256(b)} } // listingsFrom returns the mcpListings NewContext installed, or nil. @@ -253,20 +263,25 @@ func (l *mcpListings) listed(m *mcpToolset, ctx agent.ReadonlyContext, tools []t } } -// filteredBefore returns the list ToolFilter last allowed when Tools already -// ran it in this agent run. +// filteredBefore returns the list ToolFilter last allowed in this agent run, +// once the toolset's ProcessRequest has run for the agent. Until then, nothing +// but Tools refreshes the list. func (l *mcpListings) filteredBefore(m *mcpToolset, ctx agent.ReadonlyContext) ([]tool.Tool, bool) { if l == nil { return nil, false } l.mu.Lock() defer l.mu.Unlock() - tools, ok := l.entry(m, ctx).allowed[runKey(ctx)] + s := l.entry(m, ctx) + if !s.processed { + return nil, false + } + tools, ok := s.allowed[runKey(ctx)] return tools, ok } -// filtered records a list ToolFilter allowed in the agent run. Tools records -// the first; ProcessRequest replaces it once Tools has. +// filtered records a list ToolFilter allowed in the agent run. ProcessRequest +// always records it; Tools only once ProcessRequest has run for the agent. func (l *mcpListings) filtered(m *mcpToolset, ctx agent.ReadonlyContext, tools []tool.Tool, inTools bool) { if l == nil { return @@ -274,10 +289,10 @@ func (l *mcpListings) filtered(m *mcpToolset, ctx agent.ReadonlyContext, tools [ l.mu.Lock() defer l.mu.Unlock() s := l.entry(m, ctx) - k := runKey(ctx) - if _, ok := s.allowed[k]; ok || inTools { - s.allowed[k] = tools + if inTools && !s.processed { + return } + s.allowed[runKey(ctx)] = tools } // take marks the toolset processed for the agent and returns the list Tools diff --git a/go/googleadk/mcp_hitl_test.go b/go/googleadk/mcp_hitl_test.go index a99ec5d5..f54809cd 100644 --- a/go/googleadk/mcp_hitl_test.go +++ b/go/googleadk/mcp_hitl_test.go @@ -2,6 +2,7 @@ package googleadk_test import ( "context" + "fmt" "sync/atomic" "testing" @@ -84,8 +85,11 @@ type mcpHitlResult struct { // approval or rejects on denial. // // With excludeBeforeResume, the toolset's ToolFilter stops allowing the guarded -// tool after pass 1, so the approved call must not run on resume. -func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume bool) (mcpHitlResult, error) { +// tool after pass 1, so the approved call must not run on resume. With +// reuseMessage, pass 2 overwrites pass 1's message with the decision and +// resumes with the same *genai.Content, the way a caller that receives each +// message into one variable does. +func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMessage bool) (mcpHitlResult, error) { excluded := false root, err := llmagent.New(llmagent.Config{ Name: "assistant", @@ -143,6 +147,10 @@ func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume bool) (m FunctionCallID: pending[0].FunctionCallID, Confirmed: confirm, }) + if reuseMessage { + *msg = *resume + resume = msg + } for ev, rerr := range r.Run(adkCtx, "user-1", "session-1", resume, agent.RunConfig{}) { if rerr != nil { return res, rerr @@ -195,25 +203,29 @@ func newMcpHitlEnv(t *testing.T, s *testsuite.WorkflowTestSuite) (*testsuite.Tes // surfaces it, with the original call attached for the human), and the approval // flows back through the snapshot so the tool executes exactly once. func TestMCPToolConfirmationPausesAndResumes(t *testing.T) { - mcpGuardedRuns.Store(0) + for _, reuseMessage := range []bool{false, true} { + t.Run(fmt.Sprintf("reuseMessage=%v", reuseMessage), func(t *testing.T) { + mcpGuardedRuns.Store(0) - var s testsuite.WorkflowTestSuite - env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, true, false) + var s testsuite.WorkflowTestSuite + env, counter := newMcpHitlEnv(t, &s) + env.ExecuteWorkflow(mcpHitlWorkflow, true, false, reuseMessage) - require.True(t, env.IsWorkflowCompleted()) - require.NoError(t, env.GetWorkflowError()) + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) - var res mcpHitlResult - require.NoError(t, env.GetWorkflowResult(&res)) + var res mcpHitlResult + require.NoError(t, env.GetWorkflowResult(&res)) - require.True(t, res.Paused, "agent must pause on a pending confirmation") - assert.Equal(t, "guarded", res.OriginalCallName, "the human must see the original tool call") - assert.NotEmpty(t, res.PendingHint) - assert.EqualValues(t, 1, mcpGuardedRuns.Load(), "guarded MCP tool must run exactly once, only after approval") - assert.Contains(t, res.ResumedToolResponses, "guarded") - assert.Equal(t, 2, counter.get(googleadk.CallMcpToolActivityName), - "one pending pass, one approved execution") + require.True(t, res.Paused, "agent must pause on a pending confirmation") + assert.Equal(t, "guarded", res.OriginalCallName, "the human must see the original tool call") + assert.NotEmpty(t, res.PendingHint) + assert.EqualValues(t, 1, mcpGuardedRuns.Load(), "guarded MCP tool must run exactly once, only after approval") + assert.Contains(t, res.ResumedToolResponses, "guarded") + assert.Equal(t, 2, counter.get(googleadk.CallMcpToolActivityName), + "one pending pass, one approved execution") + }) + } } // TestMCPToolConfirmationDenialBlocksTool proves the denial path: the human's @@ -225,7 +237,7 @@ func TestMCPToolConfirmationDenialBlocksTool(t *testing.T) { var s testsuite.WorkflowTestSuite env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, false, false) + env.ExecuteWorkflow(mcpHitlWorkflow, false, false, false) require.True(t, env.IsWorkflowCompleted()) require.NoError(t, env.GetWorkflowError()) @@ -244,24 +256,30 @@ func TestMCPToolConfirmationDenialBlocksTool(t *testing.T) { // TestMCPToolConfirmationResumeHonorsCurrentToolFilter proves that a tool the // ToolFilter excludes while its confirmation is pending does not run when the // human approves: ADK resumes the call from the toolset's current Tools list, -// which the filter governs at that point, not when the tools were listed. +// which the filter governs at that point, not when the tools were listed. The +// resume keeps the paused run's invocation ID, and with reuseMessage also its +// *genai.Content, so only the content's value tells the runs apart. func TestMCPToolConfirmationResumeHonorsCurrentToolFilter(t *testing.T) { - mcpGuardedRuns.Store(0) + for _, reuseMessage := range []bool{false, true} { + t.Run(fmt.Sprintf("reuseMessage=%v", reuseMessage), func(t *testing.T) { + mcpGuardedRuns.Store(0) - var s testsuite.WorkflowTestSuite - env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, true, true) + var s testsuite.WorkflowTestSuite + env, counter := newMcpHitlEnv(t, &s) + env.ExecuteWorkflow(mcpHitlWorkflow, true, true, reuseMessage) - require.True(t, env.IsWorkflowCompleted()) - require.NoError(t, env.GetWorkflowError()) + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) - var res mcpHitlResult - require.NoError(t, env.GetWorkflowResult(&res)) + var res mcpHitlResult + require.NoError(t, env.GetWorkflowResult(&res)) - require.True(t, res.Paused, "agent must pause on a pending confirmation") - assert.EqualValues(t, 0, mcpGuardedRuns.Load(), "an excluded tool must not run after approval") - assert.Contains(t, res.ResumedGuardedError, "not found", - "ADK must not find the excluded tool when it resumes the call") - assert.Equal(t, 1, counter.get(googleadk.CallMcpToolActivityName), - "only the pending pass reaches the worker") + require.True(t, res.Paused, "agent must pause on a pending confirmation") + assert.EqualValues(t, 0, mcpGuardedRuns.Load(), "an excluded tool must not run after approval") + assert.Contains(t, res.ResumedGuardedError, "not found", + "ADK must not find the excluded tool when it resumes the call") + assert.Equal(t, 1, counter.get(googleadk.CallMcpToolActivityName), + "only the pending pass reaches the worker") + }) + } } diff --git a/go/googleadk/mcp_test.go b/go/googleadk/mcp_test.go index ad106225..60c5d67d 100644 --- a/go/googleadk/mcp_test.go +++ b/go/googleadk/mcp_test.go @@ -5,6 +5,8 @@ import ( "errors" "fmt" "iter" + "slices" + "strings" "sync" "sync/atomic" "testing" @@ -195,6 +197,20 @@ type ruleModel struct { func (m *ruleModel) Name() string { return m.name } func (m *ruleModel) GenerateContent(_ context.Context, req *model.LLMRequest, _ bool) iter.Seq2[*model.LLMResponse, error] { + names := declaredNames(req) + m.mu.Lock() + m.decls = append(m.decls, names) + m.mu.Unlock() + resp := m.rule(req) + return func(yield func(*model.LLMResponse, error) bool) { yield(resp, nil) } +} + +func (m *ruleModel) factory() googleadk.ModelFactory { + return func(context.Context, string) (model.LLM, error) { return m, nil } +} + +// declaredNames returns the names of the tools a model request declares. +func declaredNames(req *model.LLMRequest) []string { var names []string if req.Config != nil { for _, gt := range req.Config.Tools { @@ -205,15 +221,7 @@ func (m *ruleModel) GenerateContent(_ context.Context, req *model.LLMRequest, _ } } } - m.mu.Lock() - m.decls = append(m.decls, names) - m.mu.Unlock() - resp := m.rule(req) - return func(yield func(*model.LLMResponse, error) bool) { yield(resp, nil) } -} - -func (m *ruleModel) factory() googleadk.ModelFactory { - return func(context.Context, string) (model.LLM, error) { return m, nil } + return names } // declared returns the tool names each model request declared, in order. @@ -484,6 +492,115 @@ func TestMcpDroppedToolIsNeitherDeclaredNorCalled(t *testing.T) { assert.Equal(t, 1, counter.get(googleadk.CallMcpToolActivityName), "the dropped tool ran again") } +// wrappedToolsLog records, for TestMcpWrappedToolsetListsOnEveryToolsCall, the +// tools each Tools call of the wrapped MCP toolset returns and the tools each +// model request declares, in order. +var wrappedToolsLog *commandLog + +// loggingToolset records in wrappedToolsLog the tools each Tools call returns. +// Like the ADK wrappers, it has no ProcessRequest. +type loggingToolset struct { + tool.Toolset +} + +func (l loggingToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { + tools, err := l.Toolset.Tools(ctx) + names := make([]string, 0, len(tools)) + for _, t := range tools { + names = append(names, t.Name()) + } + wrappedToolsLog.add(fmt.Sprintf("tools %v", names)) + return tools, err +} + +// wrappedMCPWorkflow runs an agent whose MCP toolset is wrapped in the named +// ADK wrapper, which does not forward the toolset's ProcessRequest. +func wrappedMCPWorkflow(ctx workflow.Context, wrapper string) (runResult, error) { + ts := googleadk.NewMCPToolset(googleadk.MCPToolsetOptions{Name: "search-server"}) + switch wrapper { + case "FilterToolset": + ts = tool.FilterToolset(ts, func(agent.ReadonlyContext, tool.Tool) bool { return true }) + case "WithConfirmation": + ts = tool.WithConfirmation(ts, false, nil) + default: + return runResult{}, fmt.Errorf("unknown wrapper %q", wrapper) + } + return runAgent(ctx, agentBuild{ + modelName: "fake-model", + userMessage: "search twice", + toolsets: []tool.Toolset{loggingToolset{ts}}, + }) +} + +// TestMcpWrappedToolsetListsOnEveryToolsCall proves that an MCP toolset whose +// ProcessRequest ADK never calls, because a wrapper hides it, lists on every +// Tools call, as before the toolset listed once per model step: once per agent +// run with ADK v2.5.0 and earlier, and once per model step with v2.6.0. Each +// model request declares what the latest Tools call returned, so with v2.6.0 a +// tool the server dropped is neither declared nor called again. +func TestMcpWrappedToolsetListsOnEveryToolsCall(t *testing.T) { + for _, wrapper := range []string{"FilterToolset", "WithConfirmation"} { + t.Run(wrapper, func(t *testing.T) { + wrappedToolsLog = &commandLog{} + var hide atomic.Bool + server := googleadk.NewFakeMCPServer("search-server"). + AddTool("search", "search the web", searchSchema(), + func(map[string]any) (map[string]any, error) { + hide.Store(true) + return map[string]any{"hits": 1}, nil + }). + AddTool("fetch", "fetch a page", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{}, nil }) + m := &ruleModel{name: "fake-model", rule: func(req *model.LLMRequest) *model.LLMResponse { + wrappedToolsLog.add(fmt.Sprintf("model %v", declaredNames(req))) + if n := functionResponses(req); n < 2 { + return googleadk.FunctionCallResponse(fmt.Sprintf("c%d", n+1), "search", map[string]any{"query": "temporal"}) + } + return googleadk.TextResponse("done") + }} + + var s testsuite.WorkflowTestSuite + env, counter := newEnv(t, &s, googleadk.Config{ + Models: map[string]googleadk.ModelFactory{"fake-model": m.factory()}, + MCPToolsets: map[string]googleadk.MCPFactory{"search-server": func(context.Context) (tool.Toolset, error) { + return hidingToolset{Toolset: server, hide: &hide}, nil + }}, + }) + env.RegisterWorkflow(wrappedMCPWorkflow) + env.ExecuteWorkflow(wrappedMCPWorkflow, wrapper) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + toolsCalls, latest := 0, "" + for _, e := range wrappedToolsLog.get() { + if after, ok := strings.CutPrefix(e, "tools "); ok { + toolsCalls++ + latest = after + continue + } + assert.Equal(t, "model "+latest, e, "a model request declares the latest Tools list") + } + require.Positive(t, toolsCalls) + assert.Equal(t, toolsCalls, counter.get(googleadk.ListMcpToolsActivityName), "one listing per Tools call") + + declared := m.declared() + require.Len(t, declared, 3) + wantCalls := 0 + for _, names := range declared[:2] { + if slices.Contains(names, "search") { + wantCalls++ + } + } + assert.Equal(t, wantCalls, counter.get(googleadk.CallMcpToolActivityName), + "a search call runs CallMcpTool only when the request declared search") + if toolsCalls == len(declared) { + assert.Equal(t, [][]string{{"search", "fetch"}, {"fetch"}, {"fetch"}}, declared) + } + }) + } +} + // parallelHelperWorkflow has the root agent call one helper agent twice in // parallel through an agent tool. The helper has two MCP toolsets. func parallelHelperWorkflow(ctx workflow.Context, sequential bool) (runResult, error) { From 49c0efa91af5cac072025653d51227ad8b045642 Mon Sep 17 00:00:00 2001 From: Francesc Campoy Date: Fri, 30 Oct 2026 12:06:05 -0700 Subject: [PATCH 06/10] Address review findings --- go/googleadk/mcp.go | 262 +++++++++++++++++----------------- go/googleadk/mcp_hitl_test.go | 60 ++++++-- go/googleadk/mcp_test.go | 91 ++++++++++-- 3 files changed, 253 insertions(+), 160 deletions(-) diff --git a/go/googleadk/mcp.go b/go/googleadk/mcp.go index 584ea91a..9a925790 100644 --- a/go/googleadk/mcp.go +++ b/go/googleadk/mcp.go @@ -2,10 +2,9 @@ package googleadk import ( "context" - "crypto/sha256" - "encoding/json" "errors" "fmt" + "reflect" "sync" "go.temporal.io/sdk/activity" @@ -61,39 +60,36 @@ type mcpToolset struct { func (m *mcpToolset) Name() string { return m.name } -// Tools returns the toolset's declaration-only proxy tools. It lists them via -// the ListMcpTools Activity until the toolset's ProcessRequest has run for the -// agent, and keeps that list for the ProcessRequest call that follows in the -// same model step. After that, ProcessRequest lists, and Tools returns the -// latest list without an Activity. See ProcessRequest. +// Tools returns the toolset's declaration-only proxy tools. Until the +// toolset's ProcessRequest has run in the agent run, it lists them via the +// ListMcpTools Activity, applies ToolFilter, and keeps the list for the +// ProcessRequest call that follows in the same model step. After that, +// ProcessRequest lists, and Tools returns the list ToolFilter last allowed +// there, without an Activity and without running the filter. See +// ProcessRequest. // -// ToolFilter runs on the first call of each agent run, not when the list is -// taken, so a list kept from an earlier turn returns only the tools the filter -// allows now. ADK resumes a confirmed tool call from this list before any -// ProcessRequest of the resuming run. Later calls in the run, which ADK v2.6.0 -// and later make before every model step, return the list the filter last -// allowed without running it again: a filter that issues Workflow commands -// then issues them at the same points on every ADK release. +// So the first Tools call of each agent run lists and filters. ADK resumes a +// confirmed tool call from this list before any ProcessRequest of the run, so +// a tool the filter no longer allows does not run. The later calls, which ADK +// v2.6.0 and later make before every model step, issue no Workflow commands, +// so the toolset and its filter issue them at the same points on every ADK +// release. // // A toolset wrapped in a toolset without ProcessRequest, such as // tool.FilterToolset, lists and filters on every call, because nothing else // refreshes its list. func (m *mcpToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { l := listingsFrom(ctx) - if tools, ok := l.filteredBefore(m, ctx); ok { + run := l.startRun(m, ctx) + if tools, ok := l.processedList(run); ok { return tools, nil } - tools, ok := l.cached(m, ctx) - if !ok { - var err error - if tools, err = m.list(ctx); err != nil { - return nil, err - } - l.listed(m, ctx, tools, true) + tools, err := m.list(ctx) + if err != nil { + return nil, err } - tools = m.allowed(ctx, tools) - l.filtered(m, ctx, tools, true) - return tools, nil + l.keepPending(run, tools) + return m.allowed(ctx, tools), nil } // allowed returns the tools the toolset's ToolFilter allows for ctx. @@ -137,23 +133,23 @@ func (m *mcpToolset) list(ctx agent.ReadonlyContext) ([]tool.Tool, error) { // run, before the request processors. v2.6.0 and later call Tools before the // request processors in every step. To schedule the same commands in the same // order on every ADK release, the toolset lists in Tools only until its first -// ProcessRequest for the agent, and in ProcessRequest from then on. Each model -// step lists once: the first through Tools, with ProcessRequest reusing that -// list, and every later step here. The request carries this list alone, so a -// tool the MCP server no longer advertises is neither declared nor dispatched, -// whatever list ADK kept from an earlier Tools call. +// ProcessRequest of the agent run, and in ProcessRequest from then on. Each +// model step lists once: the first through Tools, with ProcessRequest reusing +// that list, and every later step here. The request carries this list alone, +// so a tool the MCP server no longer advertises is neither declared nor +// dispatched, whatever list ADK kept from an earlier Tools call. func (m *mcpToolset) ProcessRequest(ctx agent.Context, req *model.LLMRequest) error { l := listingsFrom(ctx) - tools, ok := l.take(m, ctx) + run := l.current(m, ctx) + tools, ok := l.takePending(run) if !ok { var err error if tools, err = m.list(ctx); err != nil { return err } - l.listed(m, ctx, tools, false) } tools = m.allowed(ctx, tools) - l.filtered(m, ctx, tools, false) + l.processed(run, tools) for _, t := range tools { if err := packTool(req, t.(*mcpProxyTool)); err != nil { return err @@ -166,50 +162,41 @@ func (m *mcpToolset) ProcessRequest(ctx agent.Context, req *model.LLMRequest) er // installs. type mcpListingsKey struct{} -// mcpListings records, for each MCP toolset and agent, the lists the toolset -// took, so that Tools and ProcessRequest list once per model step at the same -// point on every ADK release. +// mcpListings records, for each MCP toolset and agent run, the lists the +// toolset took, so that Tools and ProcessRequest list once per model step at +// the same point on every ADK release. type mcpListings struct { - mu sync.Mutex - state map[mcpListingKey]*mcpListingState -} - -type mcpListingKey struct { - toolset *mcpToolset - agent string -} - -type mcpListingState struct { - // processed reports that the toolset's ProcessRequest has run for the - // agent. From then on ProcessRequest lists, and Tools does not. - processed bool - // latest is the most recent list, before ToolFilter, which Tools returns - // once processed. - latest []tool.Tool - // pending holds, per invocation, the list Tools took for the - // ProcessRequest call that follows in the same model step. Concurrent - // invocations of one agent each keep their own. - pending map[string][]tool.Tool - // allowed holds, per agent run since processed, the list ToolFilter last - // allowed there. Tools returns it on the later calls of the run. - allowed map[mcpRunKey][]tool.Tool + mu sync.Mutex + runs map[mcpAgentKey]*mcpRunState } -// mcpRunKey identifies an agent run. A run that resumes a paused one keeps -// its invocation ID, so the key adds the user content that started the run: -// ADK passes the same content to every step of a run. The content is compared -// by value too, because a caller may resume with the same *genai.Content, -// overwritten with the approval. -type mcpRunKey struct { +// mcpAgentKey identifies a toolset of one agent in one invocation. Agent names +// are unique within an agent tree, and each runner.Run call has its own +// invocation ID, except one that resumes or retries a paused run. Concurrent +// agent-tool calls each make their own runner.Run call. +type mcpAgentKey struct { + toolset *mcpToolset invocation string - content *genai.Content - digest [sha256.Size]byte + agent string } -func runKey(ctx agent.ReadonlyContext) mcpRunKey { - c := ctx.UserContent() - b, _ := json.Marshal(c) - return mcpRunKey{ctx.InvocationID(), c, sha256.Sum256(b)} +// mcpRunState is the state of a toolset in one agent run: one agent in one +// runner.Run call. +type mcpRunState struct { + // config is the RunConfig ADK allocates for each runner.Run call, a resume + // or retry included, and passes to every agent the call runs. It tells + // apart the runs of one invocation. + config *agent.RunConfig + // pending is the list Tools took for the ProcessRequest call that follows + // in the same model step, when hasPending is set. + pending []tool.Tool + hasPending bool + // processed reports that the toolset's ProcessRequest has run in the agent + // run. From then on ProcessRequest lists, and Tools does not. + processed bool + // allowed is the list ToolFilter allowed in the last ProcessRequest, which + // Tools returns once processed. + allowed []tool.Tool } // listingsFrom returns the mcpListings NewContext installed, or nil. @@ -218,107 +205,116 @@ func listingsFrom(ctx context.Context) *mcpListings { return l } -func (l *mcpListings) entry(m *mcpToolset, ctx agent.ReadonlyContext) *mcpListingState { - if l.state == nil { - l.state = map[mcpListingKey]*mcpListingState{} +// runConfigOf returns the RunConfig of the runner.Run call ctx belongs to. ADK +// passes Tools a ReadonlyContext that holds its invocation in the exported +// InvocationContext field. +func runConfigOf(ctx agent.ReadonlyContext) *agent.RunConfig { + v := reflect.ValueOf(ctx) + if v.Kind() != reflect.Pointer || v.IsNil() || v.Elem().Kind() != reflect.Struct { + return nil + } + f := v.Elem().FieldByName("InvocationContext") + if !f.IsValid() || !f.CanInterface() { + return nil } - k := mcpListingKey{m, ctx.AgentName()} - s, ok := l.state[k] + ic, ok := f.Interface().(agent.InvocationContext) if !ok { - s = &mcpListingState{pending: map[string][]tool.Tool{}, allowed: map[mcpRunKey][]tool.Tool{}} - l.state[k] = s + return nil } - return s + return ic.RunConfig() } -// cached returns the latest list when the toolset's ProcessRequest has run for -// the agent, and so lists in its place. It also drops a list Tools kept for -// this invocation that no ProcessRequest took. -func (l *mcpListings) cached(m *mcpToolset, ctx agent.ReadonlyContext) ([]tool.Tool, bool) { +// startRun returns the state of the toolset in ctx's agent run, for Tools. It +// starts a new state when ctx belongs to another runner.Run call than the +// state kept for the agent and invocation. When ctx does not tell its +// runner.Run call, it drops the kept state and returns nil. +// +// Without a state, the toolset lists and filters on every Tools and +// ProcessRequest call, and its proxy tools pack themselves. +func (l *mcpListings) startRun(m *mcpToolset, ctx agent.ReadonlyContext) *mcpRunState { if l == nil { - return nil, false + return nil } + cfg := runConfigOf(ctx) + k := mcpAgentKey{m, ctx.InvocationID(), ctx.AgentName()} l.mu.Lock() defer l.mu.Unlock() - s := l.entry(m, ctx) - delete(s.pending, ctx.InvocationID()) - if !s.processed || s.latest == nil { - return nil, false + if cfg == nil { + delete(l.runs, k) + return nil + } + if l.runs == nil { + l.runs = map[mcpAgentKey]*mcpRunState{} } - return s.latest, true + s, ok := l.runs[k] + if !ok || s.config != cfg { + s = &mcpRunState{config: cfg} + l.runs[k] = s + } + return s } -// listed records a list the toolset took. With pending, it keeps the list for -// the ProcessRequest call of the same invocation. -func (l *mcpListings) listed(m *mcpToolset, ctx agent.ReadonlyContext, tools []tool.Tool, pending bool) { +// current returns the state the toolset's last Tools call for ctx's agent and +// invocation kept, or nil. ADK calls Tools before ProcessRequest in the first +// model step of every agent run. +func (l *mcpListings) current(m *mcpToolset, ctx agent.ReadonlyContext) *mcpRunState { if l == nil { - return + return nil } l.mu.Lock() defer l.mu.Unlock() - s := l.entry(m, ctx) - s.latest = tools - if pending { - s.pending[ctx.InvocationID()] = tools - } + return l.runs[mcpAgentKey{m, ctx.InvocationID(), ctx.AgentName()}] } -// filteredBefore returns the list ToolFilter last allowed in this agent run, -// once the toolset's ProcessRequest has run for the agent. Until then, nothing -// but Tools refreshes the list. -func (l *mcpListings) filteredBefore(m *mcpToolset, ctx agent.ReadonlyContext) ([]tool.Tool, bool) { - if l == nil { +// processedList returns the list ToolFilter last allowed in the run, once the +// toolset's ProcessRequest has run there. +func (l *mcpListings) processedList(s *mcpRunState) ([]tool.Tool, bool) { + if s == nil { return nil, false } l.mu.Lock() defer l.mu.Unlock() - s := l.entry(m, ctx) - if !s.processed { - return nil, false - } - tools, ok := s.allowed[runKey(ctx)] - return tools, ok + return s.allowed, s.processed } -// filtered records a list ToolFilter allowed in the agent run. ProcessRequest -// always records it; Tools only once ProcessRequest has run for the agent. -func (l *mcpListings) filtered(m *mcpToolset, ctx agent.ReadonlyContext, tools []tool.Tool, inTools bool) { - if l == nil { +// keepPending keeps a list Tools took for the ProcessRequest call that follows. +func (l *mcpListings) keepPending(s *mcpRunState, tools []tool.Tool) { + if s == nil { return } l.mu.Lock() defer l.mu.Unlock() - s := l.entry(m, ctx) - if inTools && !s.processed { - return - } - s.allowed[runKey(ctx)] = tools + s.pending, s.hasPending = tools, true } -// take marks the toolset processed for the agent and returns the list Tools -// kept for this invocation, if any. -func (l *mcpListings) take(m *mcpToolset, ctx agent.ReadonlyContext) ([]tool.Tool, bool) { - if l == nil { +// takePending returns and drops the list Tools kept, if any. +func (l *mcpListings) takePending(s *mcpRunState) ([]tool.Tool, bool) { + if s == nil { return nil, false } l.mu.Lock() defer l.mu.Unlock() - s := l.entry(m, ctx) - s.processed = true - tools, ok := s.pending[ctx.InvocationID()] - delete(s.pending, ctx.InvocationID()) + tools, ok := s.pending, s.hasPending + s.pending, s.hasPending = nil, false return tools, ok } -// packedByToolset reports whether the toolset's ProcessRequest has run for the -// agent. Its proxy tools then leave packing to it. -func (l *mcpListings) packedByToolset(m *mcpToolset, ctx agent.ReadonlyContext) bool { - if l == nil { - return false +// processed marks the toolset's ProcessRequest as run in the agent run, and +// records the list ToolFilter allowed there. +func (l *mcpListings) processed(s *mcpRunState, allowed []tool.Tool) { + if s == nil { + return } l.mu.Lock() defer l.mu.Unlock() - return l.entry(m, ctx).processed + s.processed, s.allowed = true, allowed +} + +// packedByToolset reports whether the toolset's ProcessRequest has run in the +// agent run. Its proxy tools then leave packing to it. +func (l *mcpListings) packedByToolset(m *mcpToolset, ctx agent.ReadonlyContext) bool { + _, ok := l.processedList(l.current(m, ctx)) + return ok } // mcpProxyTool is a declaration-only stand-in for a remote MCP tool. The model @@ -339,7 +335,7 @@ func (t *mcpProxyTool) Declaration() *genai.FunctionDeclaration { return t.decl // ProcessRequest packs the proxy tool. ADK calls it for each tool it kept from // a Tools call. Once the owning toolset's ProcessRequest packs the current list -// for the agent, it does nothing: the tool may come from an older list, which +// in the agent run, it does nothing: the tool may come from an older list, which // ADK v2.5.0 and earlier keep for the whole agent run. func (t *mcpProxyTool) ProcessRequest(ctx agent.Context, req *model.LLMRequest) error { if listingsFrom(ctx).packedByToolset(t.owner, ctx) { diff --git a/go/googleadk/mcp_hitl_test.go b/go/googleadk/mcp_hitl_test.go index f54809cd..2e04b132 100644 --- a/go/googleadk/mcp_hitl_test.go +++ b/go/googleadk/mcp_hitl_test.go @@ -2,6 +2,7 @@ package googleadk_test import ( "context" + "errors" "fmt" "sync/atomic" "testing" @@ -53,6 +54,21 @@ func guardedMCPFactory() googleadk.MCPFactory { } } +// flakyToolset is a toolset without tools whose Tools call fails while fail is +// set, the way a toolset fails to reach its server. +type flakyToolset struct { + fail *bool +} + +func (flakyToolset) Name() string { return "flaky" } + +func (f flakyToolset) Tools(agent.ReadonlyContext) ([]tool.Tool, error) { + if *f.fail { + return nil, errors.New("flaky toolset unavailable") + } + return nil, nil +} + // mcpHitlResult is the serializable output of mcpHitlWorkflow. type mcpHitlResult struct { // Paused reports that the first Run pass paused on a pending confirmation. @@ -88,9 +104,12 @@ type mcpHitlResult struct { // tool after pass 1, so the approved call must not run on resume. With // reuseMessage, pass 2 overwrites pass 1's message with the decision and // resumes with the same *genai.Content, the way a caller that receives each -// message into one variable does. -func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMessage bool) (mcpHitlResult, error) { - excluded := false +// message into one variable does. With retryResume, the first resume fails +// because another toolset cannot list its tools, after the MCP toolset listed +// and filtered; the filter changes only then, and the workflow retries the +// resume with the same message. +func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMessage, retryResume bool) (mcpHitlResult, error) { + excluded, flaky := false, false root, err := llmagent.New(llmagent.Config{ Name: "assistant", Description: "root assistant", @@ -101,7 +120,7 @@ func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMe ToolFilter: func(_ agent.ReadonlyContext, t tool.Tool) bool { return !excluded || t.Name() != "guarded" }, - })}, + }), flakyToolset{&flaky}}, }) if err != nil { return mcpHitlResult{}, err @@ -142,7 +161,6 @@ func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMe } // --- Pass 2: simulate the human's decision, then resume. --- - excluded = excludeBeforeResume resume := googleadk.ConfirmationResponse(googleadk.ConfirmationDecision{ FunctionCallID: pending[0].FunctionCallID, Confirmed: confirm, @@ -151,6 +169,21 @@ func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMe *msg = *resume resume = msg } + if retryResume { + flaky = true + var failed error + for _, rerr := range r.Run(adkCtx, "user-1", "session-1", resume, agent.RunConfig{}) { + if rerr != nil { + failed = rerr + break + } + } + if failed == nil { + return res, errors.New("the first resume did not fail") + } + flaky = false + } + excluded = excludeBeforeResume for ev, rerr := range r.Run(adkCtx, "user-1", "session-1", resume, agent.RunConfig{}) { if rerr != nil { return res, rerr @@ -203,13 +236,13 @@ func newMcpHitlEnv(t *testing.T, s *testsuite.WorkflowTestSuite) (*testsuite.Tes // surfaces it, with the original call attached for the human), and the approval // flows back through the snapshot so the tool executes exactly once. func TestMCPToolConfirmationPausesAndResumes(t *testing.T) { - for _, reuseMessage := range []bool{false, true} { - t.Run(fmt.Sprintf("reuseMessage=%v", reuseMessage), func(t *testing.T) { + for _, tc := range []struct{ reuseMessage, retryResume bool }{{false, false}, {true, false}, {true, true}} { + t.Run(fmt.Sprintf("reuseMessage=%v/retryResume=%v", tc.reuseMessage, tc.retryResume), func(t *testing.T) { mcpGuardedRuns.Store(0) var s testsuite.WorkflowTestSuite env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, true, false, reuseMessage) + env.ExecuteWorkflow(mcpHitlWorkflow, true, false, tc.reuseMessage, tc.retryResume) require.True(t, env.IsWorkflowCompleted()) require.NoError(t, env.GetWorkflowError()) @@ -237,7 +270,7 @@ func TestMCPToolConfirmationDenialBlocksTool(t *testing.T) { var s testsuite.WorkflowTestSuite env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, false, false, false) + env.ExecuteWorkflow(mcpHitlWorkflow, false, false, false, false) require.True(t, env.IsWorkflowCompleted()) require.NoError(t, env.GetWorkflowError()) @@ -258,15 +291,16 @@ func TestMCPToolConfirmationDenialBlocksTool(t *testing.T) { // human approves: ADK resumes the call from the toolset's current Tools list, // which the filter governs at that point, not when the tools were listed. The // resume keeps the paused run's invocation ID, and with reuseMessage also its -// *genai.Content, so only the content's value tells the runs apart. +// *genai.Content. With retryResume, a failed resume already listed and +// filtered before the filter changed, and the retry passes the same message. func TestMCPToolConfirmationResumeHonorsCurrentToolFilter(t *testing.T) { - for _, reuseMessage := range []bool{false, true} { - t.Run(fmt.Sprintf("reuseMessage=%v", reuseMessage), func(t *testing.T) { + for _, tc := range []struct{ reuseMessage, retryResume bool }{{false, false}, {true, false}, {false, true}, {true, true}} { + t.Run(fmt.Sprintf("reuseMessage=%v/retryResume=%v", tc.reuseMessage, tc.retryResume), func(t *testing.T) { mcpGuardedRuns.Store(0) var s testsuite.WorkflowTestSuite env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, true, true, reuseMessage) + env.ExecuteWorkflow(mcpHitlWorkflow, true, true, tc.reuseMessage, tc.retryResume) require.True(t, env.IsWorkflowCompleted()) require.NoError(t, env.GetWorkflowError()) diff --git a/go/googleadk/mcp_test.go b/go/googleadk/mcp_test.go index 60c5d67d..52e42570 100644 --- a/go/googleadk/mcp_test.go +++ b/go/googleadk/mcp_test.go @@ -366,15 +366,14 @@ func sleepingInstructionWorkflow(ctx workflow.Context, filterSleeps bool) error } // TestMcpListsAtTheSamePointInEveryStep proves the MCP toolset lists at a point -// that does not depend on the ADK release: in the first model step of the -// first turn before the request processors (ADK calls Tools there on every -// release), and in every later step and turn after them, in ProcessRequest. -// ADK v2.6.0 also calls Tools before the request processors of every step; a -// listing there would put ListMcpTools before the instruction's timer, where -// ADK v2.5.0 and earlier put it after, and Histories would not replay across -// releases. ToolFilter runs at the same points too: in the first Tools call of -// each turn and in each ProcessRequest, never in the later Tools calls only -// ADK v2.6.0 makes. +// that does not depend on the ADK release: in the first model step of each +// turn before the request processors (ADK calls Tools there on every release), +// and in every later step after them, in ProcessRequest. ADK v2.6.0 also calls +// Tools before the request processors of every step; a listing there would put +// ListMcpTools before the instruction's timer, where ADK v2.5.0 and earlier put +// it after, and Histories would not replay across releases. ToolFilter runs at +// the same points too: in the first Tools call of each turn and in each +// ProcessRequest, never in the later Tools calls only ADK v2.6.0 makes. func TestMcpListsAtTheSamePointInEveryStep(t *testing.T) { list, invoke, call := googleadk.ListMcpToolsActivityName, googleadk.InvokeModelActivityName, googleadk.CallMcpToolActivityName instr, filter := "Timer 1s", "Timer 2s" @@ -387,17 +386,16 @@ func TestMcpListsAtTheSamePointInEveryStep(t *testing.T) { list, instr, invoke, call, // Turn 1, step 2: listed in ProcessRequest, after the instruction. instr, list, invoke, - // Turn 2: listed in ProcessRequest. - instr, list, invoke, + // Turn 2: listed through Tools, before the instruction. + list, instr, invoke, }}, {true, []string{ // Turn 1, step 1: Tools lists and filters, ProcessRequest filters. list, filter, instr, filter, invoke, call, // Turn 1, step 2: ProcessRequest lists and filters. instr, list, filter, invoke, - // Turn 2: Tools filters the kept list, ProcessRequest lists and - // filters. - filter, instr, list, filter, invoke, + // Turn 2: Tools lists and filters, ProcessRequest filters. + list, filter, instr, filter, invoke, }}, } { t.Run(fmt.Sprintf("filterSleeps=%v", tc.filterSleeps), func(t *testing.T) { @@ -601,6 +599,71 @@ func TestMcpWrappedToolsetListsOnEveryToolsCall(t *testing.T) { } } +// sharedProxyWorkflow runs one turn of an agent that uses an MCP toolset +// directly, then one turn of another agent, named secondName, that uses the +// same toolset through tool.FilterToolset. The runners share the bridged +// context. +func sharedProxyWorkflow(ctx workflow.Context, secondName string) error { + ts := googleadk.NewMCPToolset(googleadk.MCPToolsetOptions{Name: "search-server"}) + adkCtx := googleadk.NewContext(ctx) + for _, a := range []struct { + name string + toolset tool.Toolset + }{ + {"assistant", ts}, + {secondName, tool.FilterToolset(ts, func(agent.ReadonlyContext, tool.Tool) bool { return true })}, + } { + ag, err := llmagent.New(llmagent.Config{ + Name: a.name, + Model: googleadk.NewModel("fake-model"), + Toolsets: []tool.Toolset{a.toolset}, + }) + if err != nil { + return err + } + r, err := runner.New(runner.Config{AppName: "test-app", Agent: ag, SessionService: session.InMemoryService(), AutoCreateSession: true}) + if err != nil { + return err + } + for _, rerr := range r.Run(adkCtx, "user-1", "session-1", genai.NewContentFromText("hello", genai.RoleUser), agent.RunConfig{}) { + if rerr != nil { + return rerr + } + } + } + return nil +} + +// TestMcpSameNamedAgentsKeepTheirOwnListings proves that the toolset's +// ProcessRequest running for one agent does not stand in for another agent +// with the same name, whose wrapper hides that ProcessRequest: the second +// agent's request still declares the toolset's tools. +func TestMcpSameNamedAgentsKeepTheirOwnListings(t *testing.T) { + for _, secondName := range []string{"assistant", "other"} { + t.Run(secondName, func(t *testing.T) { + m := &ruleModel{name: "fake-model", rule: func(*model.LLMRequest) *model.LLMResponse { + return googleadk.TextResponse("done") + }} + server := googleadk.NewFakeMCPServer("search-server").AddTool("search", "search the web", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{"hits": 1}, nil }) + + var s testsuite.WorkflowTestSuite + env, counter := newEnv(t, &s, googleadk.Config{ + Models: map[string]googleadk.ModelFactory{"fake-model": m.factory()}, + MCPToolsets: map[string]googleadk.MCPFactory{"search-server": server.Factory()}, + }) + env.RegisterWorkflow(sharedProxyWorkflow) + env.ExecuteWorkflow(sharedProxyWorkflow, secondName) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + assert.Equal(t, [][]string{{"search"}, {"search"}}, m.declared()) + assert.Equal(t, 2, counter.get(googleadk.ListMcpToolsActivityName), "one listing per agent run") + }) + } +} + // parallelHelperWorkflow has the root agent call one helper agent twice in // parallel through an agent tool. The helper has two MCP toolsets. func parallelHelperWorkflow(ctx workflow.Context, sequential bool) (runResult, error) { From 98d4e59c3bf721d73f317442c137cc97291fe924 Mon Sep 17 00:00:00 2001 From: Francesc Campoy Date: Fri, 30 Oct 2026 13:30:08 -0700 Subject: [PATCH 07/10] Address review findings --- go/googleadk/mcp.go | 38 +++++++------- go/googleadk/mcp_hitl_test.go | 98 ++++++++++++++++++++++++++--------- 2 files changed, 93 insertions(+), 43 deletions(-) diff --git a/go/googleadk/mcp.go b/go/googleadk/mcp.go index 9a925790..baaff601 100644 --- a/go/googleadk/mcp.go +++ b/go/googleadk/mcp.go @@ -181,12 +181,13 @@ type mcpAgentKey struct { } // mcpRunState is the state of a toolset in one agent run: one agent in one -// runner.Run call. +// execution of a runner.Run iterator. type mcpRunState struct { - // config is the RunConfig ADK allocates for each runner.Run call, a resume - // or retry included, and passes to every agent the call runs. It tells - // apart the runs of one invocation. - config *agent.RunConfig + // run is the InvocationContext ADK creates for the agent in each execution + // of a runner.Run iterator, a resume, a retry or another pass over a kept + // iterator included, and passes to every Tools call of that agent run. It + // tells apart the runs of one agent in one invocation. + run agent.InvocationContext // pending is the list Tools took for the ProcessRequest call that follows // in the same model step, when hasPending is set. pending []tool.Tool @@ -205,10 +206,11 @@ func listingsFrom(ctx context.Context) *mcpListings { return l } -// runConfigOf returns the RunConfig of the runner.Run call ctx belongs to. ADK -// passes Tools a ReadonlyContext that holds its invocation in the exported -// InvocationContext field. -func runConfigOf(ctx agent.ReadonlyContext) *agent.RunConfig { +// agentRunOf returns the InvocationContext of the agent run ctx belongs to, or +// nil. ADK passes Tools a ReadonlyContext that holds it in the exported +// InvocationContext field. It is not the RunConfig: ADK reuses one RunConfig +// when a runner.Run iterator without an LLM agent at its root runs again. +func agentRunOf(ctx agent.ReadonlyContext) agent.InvocationContext { v := reflect.ValueOf(ctx) if v.Kind() != reflect.Pointer || v.IsNil() || v.Elem().Kind() != reflect.Struct { return nil @@ -218,16 +220,16 @@ func runConfigOf(ctx agent.ReadonlyContext) *agent.RunConfig { return nil } ic, ok := f.Interface().(agent.InvocationContext) - if !ok { + if !ok || reflect.ValueOf(ic).Kind() != reflect.Pointer { return nil } - return ic.RunConfig() + return ic } // startRun returns the state of the toolset in ctx's agent run, for Tools. It -// starts a new state when ctx belongs to another runner.Run call than the -// state kept for the agent and invocation. When ctx does not tell its -// runner.Run call, it drops the kept state and returns nil. +// starts a new state when ctx belongs to another agent run than the state +// kept for the agent and invocation. When ctx does not tell its agent run, it +// drops the kept state and returns nil. // // Without a state, the toolset lists and filters on every Tools and // ProcessRequest call, and its proxy tools pack themselves. @@ -235,11 +237,11 @@ func (l *mcpListings) startRun(m *mcpToolset, ctx agent.ReadonlyContext) *mcpRun if l == nil { return nil } - cfg := runConfigOf(ctx) + run := agentRunOf(ctx) k := mcpAgentKey{m, ctx.InvocationID(), ctx.AgentName()} l.mu.Lock() defer l.mu.Unlock() - if cfg == nil { + if run == nil { delete(l.runs, k) return nil } @@ -247,8 +249,8 @@ func (l *mcpListings) startRun(m *mcpToolset, ctx agent.ReadonlyContext) *mcpRun l.runs = map[mcpAgentKey]*mcpRunState{} } s, ok := l.runs[k] - if !ok || s.config != cfg { - s = &mcpRunState{config: cfg} + if !ok || s.run != run { + s = &mcpRunState{run: run} l.runs[k] = s } return s diff --git a/go/googleadk/mcp_hitl_test.go b/go/googleadk/mcp_hitl_test.go index 2e04b132..25d57a0a 100644 --- a/go/googleadk/mcp_hitl_test.go +++ b/go/googleadk/mcp_hitl_test.go @@ -17,6 +17,7 @@ import ( "google.golang.org/adk/v2/agent" "google.golang.org/adk/v2/agent/llmagent" + "google.golang.org/adk/v2/agent/workflowagents/sequentialagent" "google.golang.org/adk/v2/runner" "google.golang.org/adk/v2/session" "google.golang.org/adk/v2/tool" @@ -100,17 +101,10 @@ type mcpHitlResult struct { // snapshot carries it to the worker-side tool, which executes on // approval or rejects on denial. // -// With excludeBeforeResume, the toolset's ToolFilter stops allowing the guarded -// tool after pass 1, so the approved call must not run on resume. With -// reuseMessage, pass 2 overwrites pass 1's message with the decision and -// resumes with the same *genai.Content, the way a caller that receives each -// message into one variable does. With retryResume, the first resume fails -// because another toolset cannot list its tools, after the MCP toolset listed -// and filtered; the filter changes only then, and the workflow retries the -// resume with the same message. -func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMessage, retryResume bool) (mcpHitlResult, error) { +// See mcpHitlCase for the variations. +func mcpHitlWorkflow(ctx workflow.Context, c mcpHitlCase) (mcpHitlResult, error) { excluded, flaky := false, false - root, err := llmagent.New(llmagent.Config{ + assistant, err := llmagent.New(llmagent.Config{ Name: "assistant", Description: "root assistant", Model: googleadk.NewModel("fake-model"), @@ -125,6 +119,15 @@ func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMe if err != nil { return mcpHitlResult{}, err } + root := assistant + if c.SequentialRoot { + if root, err = sequentialagent.New(sequentialagent.Config{AgentConfig: agent.Config{ + Name: "pipeline", + SubAgents: []agent.Agent{assistant}, + }}); err != nil { + return mcpHitlResult{}, err + } + } r, err := runner.New(runner.Config{ AppName: "test-app", Agent: root, @@ -140,7 +143,8 @@ func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMe // --- Pass 1: run until the agent pauses on a confirmation. --- var pass1 []*session.Event msg := genai.NewContentFromText("run the guarded tool", genai.RoleUser) - for ev, rerr := range r.Run(adkCtx, "user-1", "session-1", msg, agent.RunConfig{}) { + run := r.Run(adkCtx, "user-1", "session-1", msg, agent.RunConfig{}) + for ev, rerr := range run { if rerr != nil { return mcpHitlResult{}, rerr } @@ -163,13 +167,13 @@ func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMe // --- Pass 2: simulate the human's decision, then resume. --- resume := googleadk.ConfirmationResponse(googleadk.ConfirmationDecision{ FunctionCallID: pending[0].FunctionCallID, - Confirmed: confirm, + Confirmed: c.Confirm, }) - if reuseMessage { + if c.ReuseMessage || c.ReuseIterator { *msg = *resume resume = msg } - if retryResume { + if c.RetryResume { flaky = true var failed error for _, rerr := range r.Run(adkCtx, "user-1", "session-1", resume, agent.RunConfig{}) { @@ -183,8 +187,11 @@ func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMe } flaky = false } - excluded = excludeBeforeResume - for ev, rerr := range r.Run(adkCtx, "user-1", "session-1", resume, agent.RunConfig{}) { + excluded = c.ExcludeBeforeResume + if !c.ReuseIterator { + run = r.Run(adkCtx, "user-1", "session-1", resume, agent.RunConfig{}) + } + for ev, rerr := range run { if rerr != nil { return res, rerr } @@ -212,6 +219,31 @@ func mcpHitlWorkflow(ctx workflow.Context, confirm, excludeBeforeResume, reuseMe return res, nil } +// mcpHitlCase selects the variations of mcpHitlWorkflow. +type mcpHitlCase struct { + // Confirm is the human's decision. + Confirm bool + // ExcludeBeforeResume makes the toolset's ToolFilter stop allowing the + // guarded tool after pass 1, so the approved call must not run on resume. + ExcludeBeforeResume bool + // ReuseMessage makes pass 2 overwrite pass 1's message with the decision + // and resume with the same *genai.Content, the way a caller that receives + // each message into one variable does. + ReuseMessage bool + // RetryResume makes the first resume fail because another toolset cannot + // list its tools, after the MCP toolset listed and filtered. The filter + // changes only then, and the workflow retries the resume with the same + // message. + RetryResume bool + // SequentialRoot runs the assistant under a sequential agent, so the + // runner's root is not an LLM agent. + SequentialRoot bool + // ReuseIterator makes pass 2 overwrite pass 1's message with the decision + // and iterate pass 1's runner.Run iterator again instead of calling + // runner.Run. + ReuseIterator bool +} + // newMcpHitlEnv builds a test environment with mcpHitlWorkflow registered and // the guarded MCP toolset wired. func newMcpHitlEnv(t *testing.T, s *testsuite.WorkflowTestSuite) (*testsuite.TestWorkflowEnvironment, *activityCounter) { @@ -236,13 +268,19 @@ func newMcpHitlEnv(t *testing.T, s *testsuite.WorkflowTestSuite) (*testsuite.Tes // surfaces it, with the original call attached for the human), and the approval // flows back through the snapshot so the tool executes exactly once. func TestMCPToolConfirmationPausesAndResumes(t *testing.T) { - for _, tc := range []struct{ reuseMessage, retryResume bool }{{false, false}, {true, false}, {true, true}} { - t.Run(fmt.Sprintf("reuseMessage=%v/retryResume=%v", tc.reuseMessage, tc.retryResume), func(t *testing.T) { + for _, tc := range []mcpHitlCase{ + {}, + {ReuseMessage: true}, + {ReuseMessage: true, RetryResume: true}, + {SequentialRoot: true, ReuseIterator: true}, + } { + tc.Confirm = true + t.Run(fmt.Sprintf("%+v", tc), func(t *testing.T) { mcpGuardedRuns.Store(0) var s testsuite.WorkflowTestSuite env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, true, false, tc.reuseMessage, tc.retryResume) + env.ExecuteWorkflow(mcpHitlWorkflow, tc) require.True(t, env.IsWorkflowCompleted()) require.NoError(t, env.GetWorkflowError()) @@ -270,7 +308,7 @@ func TestMCPToolConfirmationDenialBlocksTool(t *testing.T) { var s testsuite.WorkflowTestSuite env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, false, false, false, false) + env.ExecuteWorkflow(mcpHitlWorkflow, mcpHitlCase{}) require.True(t, env.IsWorkflowCompleted()) require.NoError(t, env.GetWorkflowError()) @@ -290,17 +328,27 @@ func TestMCPToolConfirmationDenialBlocksTool(t *testing.T) { // ToolFilter excludes while its confirmation is pending does not run when the // human approves: ADK resumes the call from the toolset's current Tools list, // which the filter governs at that point, not when the tools were listed. The -// resume keeps the paused run's invocation ID, and with reuseMessage also its -// *genai.Content. With retryResume, a failed resume already listed and +// resume keeps the paused run's invocation ID, and with ReuseMessage also its +// *genai.Content. With RetryResume, a failed resume already listed and // filtered before the filter changed, and the retry passes the same message. +// With ReuseIterator under a sequential root, ADK runs the kept iterator again +// with the same RunConfig. func TestMCPToolConfirmationResumeHonorsCurrentToolFilter(t *testing.T) { - for _, tc := range []struct{ reuseMessage, retryResume bool }{{false, false}, {true, false}, {false, true}, {true, true}} { - t.Run(fmt.Sprintf("reuseMessage=%v/retryResume=%v", tc.reuseMessage, tc.retryResume), func(t *testing.T) { + for _, tc := range []mcpHitlCase{ + {}, + {ReuseMessage: true}, + {RetryResume: true}, + {ReuseMessage: true, RetryResume: true}, + {SequentialRoot: true, ReuseMessage: true}, + {SequentialRoot: true, ReuseIterator: true}, + } { + tc.Confirm, tc.ExcludeBeforeResume = true, true + t.Run(fmt.Sprintf("%+v", tc), func(t *testing.T) { mcpGuardedRuns.Store(0) var s testsuite.WorkflowTestSuite env, counter := newMcpHitlEnv(t, &s) - env.ExecuteWorkflow(mcpHitlWorkflow, true, true, tc.reuseMessage, tc.retryResume) + env.ExecuteWorkflow(mcpHitlWorkflow, tc) require.True(t, env.IsWorkflowCompleted()) require.NoError(t, env.GetWorkflowError()) From a55a102984cd9ce86a0f20930915e8c88340973b Mon Sep 17 00:00:00 2001 From: Francesc Campoy Date: Fri, 30 Oct 2026 14:29:43 -0700 Subject: [PATCH 08/10] Address review findings --- go/googleadk/mcp.go | 66 ++++++++---- go/googleadk/mcp_listings_internal_test.go | 115 +++++++++++++++++++++ 2 files changed, 159 insertions(+), 22 deletions(-) create mode 100644 go/googleadk/mcp_listings_internal_test.go diff --git a/go/googleadk/mcp.go b/go/googleadk/mcp.go index baaff601..828b6387 100644 --- a/go/googleadk/mcp.go +++ b/go/googleadk/mcp.go @@ -6,6 +6,7 @@ import ( "fmt" "reflect" "sync" + "weak" "go.temporal.io/sdk/activity" "go.temporal.io/sdk/workflow" @@ -183,11 +184,14 @@ type mcpAgentKey struct { // mcpRunState is the state of a toolset in one agent run: one agent in one // execution of a runner.Run iterator. type mcpRunState struct { - // run is the InvocationContext ADK creates for the agent in each execution - // of a runner.Run iterator, a resume, a retry or another pass over a kept - // iterator included, and passes to every Tools call of that agent run. It - // tells apart the runs of one agent in one invocation. - run agent.InvocationContext + // run points weakly to the InvocationContext ADK creates for the agent in + // each execution of a runner.Run iterator, a resume, a retry or another + // pass over a kept iterator included, and passes to every Tools call of + // that agent run. It tells apart the runs of one agent in one invocation. + // The pointer is weak so that the store does not keep a finished run's + // InvocationContext, and the session events it holds, alive; once the + // garbage collector reclaims it, startRun drops the state. + run weak.Pointer[byte] // pending is the list Tools took for the ProcessRequest call that follows // in the same model step, when hasPending is set. pending []tool.Tool @@ -206,24 +210,32 @@ func listingsFrom(ctx context.Context) *mcpListings { return l } -// agentRunOf returns the InvocationContext of the agent run ctx belongs to, or -// nil. ADK passes Tools a ReadonlyContext that holds it in the exported -// InvocationContext field. It is not the RunConfig: ADK reuses one RunConfig -// when a runner.Run iterator without an LLM agent at its root runs again. -func agentRunOf(ctx agent.ReadonlyContext) agent.InvocationContext { +// agentRunOf returns a weak pointer to the InvocationContext of the agent run +// ctx belongs to, and whether ctx holds one. ADK passes Tools a ReadonlyContext +// that holds it in the exported InvocationContext field. It is not the +// RunConfig: ADK reuses one RunConfig when a runner.Run iterator without an LLM +// agent at its root runs again. +// +// Weak pointers made from the same InvocationContext are equal, and stay +// distinct from those made from any other, even after it is reclaimed. +func agentRunOf(ctx agent.ReadonlyContext) (weak.Pointer[byte], bool) { v := reflect.ValueOf(ctx) if v.Kind() != reflect.Pointer || v.IsNil() || v.Elem().Kind() != reflect.Struct { - return nil + return weak.Pointer[byte]{}, false } f := v.Elem().FieldByName("InvocationContext") if !f.IsValid() || !f.CanInterface() { - return nil + return weak.Pointer[byte]{}, false } ic, ok := f.Interface().(agent.InvocationContext) - if !ok || reflect.ValueOf(ic).Kind() != reflect.Pointer { - return nil + if !ok { + return weak.Pointer[byte]{}, false + } + p := reflect.ValueOf(ic) + if p.Kind() != reflect.Pointer || p.IsNil() { + return weak.Pointer[byte]{}, false } - return ic + return weak.Make((*byte)(p.UnsafePointer())), true } // startRun returns the state of the toolset in ctx's agent run, for Tools. It @@ -233,26 +245,36 @@ func agentRunOf(ctx agent.ReadonlyContext) agent.InvocationContext { // // Without a state, the toolset lists and filters on every Tools and // ProcessRequest call, and its proxy tools pack themselves. +// +// When it starts a state, it also drops the states of finished agent runs, +// whose InvocationContext the garbage collector has reclaimed. No later call +// can belong to such a run, so a long conversation on one NewContext keeps +// only the states of runs that may still call the toolset. func (l *mcpListings) startRun(m *mcpToolset, ctx agent.ReadonlyContext) *mcpRunState { if l == nil { return nil } - run := agentRunOf(ctx) + run, ok := agentRunOf(ctx) k := mcpAgentKey{m, ctx.InvocationID(), ctx.AgentName()} l.mu.Lock() defer l.mu.Unlock() - if run == nil { + if !ok { delete(l.runs, k) return nil } + if s, ok := l.runs[k]; ok && s.run == run { + return s + } + for key, s := range l.runs { + if s.run.Value() == nil { + delete(l.runs, key) + } + } if l.runs == nil { l.runs = map[mcpAgentKey]*mcpRunState{} } - s, ok := l.runs[k] - if !ok || s.run != run { - s = &mcpRunState{run: run} - l.runs[k] = s - } + s := &mcpRunState{run: run} + l.runs[k] = s return s } diff --git a/go/googleadk/mcp_listings_internal_test.go b/go/googleadk/mcp_listings_internal_test.go new file mode 100644 index 00000000..35f41b48 --- /dev/null +++ b/go/googleadk/mcp_listings_internal_test.go @@ -0,0 +1,115 @@ +package googleadk + +// White-box test for the lifetime of the MCP listing store: mcpListings is +// unexported. + +import ( + "context" + "fmt" + "runtime" + "testing" + + "github.com/nexus-rpc/sdk-go/nexus" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.temporal.io/sdk/testsuite" + "go.temporal.io/sdk/workflow" + "google.golang.org/genai" + + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/agent/llmagent" + "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/runner" + "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/tool" +) + +type listingsTestRegistry struct { + *testsuite.TestWorkflowEnvironment +} + +func (listingsTestRegistry) RegisterNexusService(*nexus.Service) {} + +// listingsAfterTurns is what listingsLifetimeWorkflow saw in the store after +// each turn and a garbage collection. +type listingsAfterTurns struct { + Entries []int + Live []int +} + +// listingsLifetimeWorkflow runs turns turns of an MCP agent on one NewContext, +// and records the store's entries, and those whose agent run is still +// reachable, after each turn and a garbage collection. +func listingsLifetimeWorkflow(ctx workflow.Context, turns int) (listingsAfterTurns, error) { + var out listingsAfterTurns + root, err := llmagent.New(llmagent.Config{ + Name: "assistant", + Model: NewModel("fake-model"), + Toolsets: []tool.Toolset{NewMCPToolset(MCPToolsetOptions{Name: "search-server"})}, + }) + if err != nil { + return out, err + } + r, err := runner.New(runner.Config{AppName: "test-app", Agent: root, SessionService: session.InMemoryService(), AutoCreateSession: true}) + if err != nil { + return out, err + } + adkCtx := NewContext(ctx) + l := listingsFrom(adkCtx) + for i := range turns { + msg := genai.NewContentFromText(fmt.Sprintf("turn %d", i), genai.RoleUser) + for _, rerr := range r.Run(adkCtx, "user-1", "session-1", msg, agent.RunConfig{}) { + if rerr != nil { + return out, rerr + } + } + runtime.GC() + l.mu.Lock() + live := 0 + for _, s := range l.runs { + if s.run.Value() != nil { + live++ + } + } + out.Entries = append(out.Entries, len(l.runs)) + out.Live = append(out.Live, live) + l.mu.Unlock() + } + return out, nil +} + +// TestMcpListingsReleaseFinishedRuns proves the listing store does not keep +// finished turns: each turn on one NewContext is a new invocation, and its +// state, with the InvocationContext and the session events it holds, must go +// once the turn ends. Only the last turn's entry may remain, unreachable, until +// the next agent run starts and drops it. +func TestMcpListingsReleaseFinishedRuns(t *testing.T) { + const turns = 20 + server := NewFakeMCPServer("search-server").AddTool("search", "search the web", nil, + func(map[string]any) (map[string]any, error) { return map[string]any{"hits": 1}, nil }) + fm := NewFakeModel( + FunctionCallResponse("c1", "search", map[string]any{"query": "temporal"}), + TextResponse("done"), + ) + acts, err := NewActivities(Config{ + Models: map[string]ModelFactory{"fake-model": func(context.Context, string) (model.LLM, error) { return fm, nil }}, + MCPToolsets: map[string]MCPFactory{"search-server": server.Factory()}, + }) + require.NoError(t, err) + + var s testsuite.WorkflowTestSuite + env := s.NewTestWorkflowEnvironment() + acts.Register(listingsTestRegistry{env}) + env.RegisterWorkflow(listingsLifetimeWorkflow) + env.ExecuteWorkflow(listingsLifetimeWorkflow, turns) + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + var got listingsAfterTurns + require.NoError(t, env.GetWorkflowResult(&got)) + require.Len(t, got.Entries, turns) + for i := range turns { + assert.LessOrEqual(t, got.Entries[i], 1, "entries after turn %d", i) + assert.Zero(t, got.Live[i], "reachable agent runs after turn %d", i) + } +} From f4a3fe7f5bfba6db72ed2fbb42d3b9a2fddf1371 Mon Sep 17 00:00:00 2001 From: Francesc Campoy Date: Sat, 31 Oct 2026 02:12:23 -0700 Subject: [PATCH 09/10] Address review findings --- go/googleadk/mcp.go | 152 +++++++++++++-------- go/googleadk/mcp_listings_internal_test.go | 4 +- go/googleadk/mcp_test.go | 152 ++++++++++++++------- 3 files changed, 198 insertions(+), 110 deletions(-) diff --git a/go/googleadk/mcp.go b/go/googleadk/mcp.go index 828b6387..e0fe91fa 100644 --- a/go/googleadk/mcp.go +++ b/go/googleadk/mcp.go @@ -168,30 +168,26 @@ type mcpListingsKey struct{} // the same point on every ADK release. type mcpListings struct { mu sync.Mutex - runs map[mcpAgentKey]*mcpRunState + runs map[mcpRunKey]*mcpRunState } -// mcpAgentKey identifies a toolset of one agent in one invocation. Agent names -// are unique within an agent tree, and each runner.Run call has its own -// invocation ID, except one that resumes or retries a paused run. Concurrent -// agent-tool calls each make their own runner.Run call. -type mcpAgentKey struct { - toolset *mcpToolset - invocation string - agent string +// mcpRunKey identifies a toolset in one agent run: one agent in one execution +// of a runner.Run iterator. +type mcpRunKey struct { + toolset *mcpToolset + // run points weakly to the invocation ADK creates each time an agent runs: + // in each execution of a runner.Run iterator, a resume, a retry or another + // pass over a kept iterator included, and in each of the concurrent runs + // of one agent that share an invocation ID, such as a custom agent running + // it on several branches. See agentRunOf. The pointer is weak so that the + // store does not keep a finished run's invocation, and the session events + // it holds, alive; once the garbage collector reclaims it, startRun drops + // the state. + run weak.Pointer[byte] } -// mcpRunState is the state of a toolset in one agent run: one agent in one -// execution of a runner.Run iterator. +// mcpRunState is the state of a toolset in one agent run. type mcpRunState struct { - // run points weakly to the InvocationContext ADK creates for the agent in - // each execution of a runner.Run iterator, a resume, a retry or another - // pass over a kept iterator included, and passes to every Tools call of - // that agent run. It tells apart the runs of one agent in one invocation. - // The pointer is weak so that the store does not keep a finished run's - // InvocationContext, and the session events it holds, alive; once the - // garbage collector reclaims it, startRun drops the state. - run weak.Pointer[byte] // pending is the list Tools took for the ProcessRequest call that follows // in the same model step, when hasPending is set. pending []tool.Tool @@ -210,84 +206,120 @@ func listingsFrom(ctx context.Context) *mcpListings { return l } -// agentRunOf returns a weak pointer to the InvocationContext of the agent run -// ctx belongs to, and whether ctx holds one. ADK passes Tools a ReadonlyContext -// that holds it in the exported InvocationContext field. It is not the -// RunConfig: ADK reuses one RunConfig when a runner.Run iterator without an LLM -// agent at its root runs again. +// agentRunOf returns a weak pointer to the invocation of the agent run ctx +// belongs to, and whether ctx tells it. ADK runs each agent on a new +// invocation, which it passes, wrapped, both to the toolset's Tools and to the +// ProcessRequest calls of the toolset and its tools: +// +// - Tools gets a ReadonlyContext whose exported InvocationContext field holds +// the agent run's context. +// - ProcessRequest gets a tool context, whose unexported context field holds +// a copy of that context. +// +// The agent run's context, and its copy, hold the invocation in an unexported +// invocationContext field, or are the invocation themselves. It is not the +// invocation ID, agent name or RunConfig: concurrent runs of one agent can +// share all three. // -// Weak pointers made from the same InvocationContext are equal, and stay -// distinct from those made from any other, even after it is reclaimed. -func agentRunOf(ctx agent.ReadonlyContext) (weak.Pointer[byte], bool) { +// Weak pointers made from the same invocation are equal, and stay distinct +// from those made from any other, even after it is reclaimed. +func agentRunOf(ctx context.Context) (weak.Pointer[byte], bool) { v := reflect.ValueOf(ctx) - if v.Kind() != reflect.Pointer || v.IsNil() || v.Elem().Kind() != reflect.Struct { - return weak.Pointer[byte]{}, false - } - f := v.Elem().FieldByName("InvocationContext") - if !f.IsValid() || !f.CanInterface() { - return weak.Pointer[byte]{}, false + for depth := range 4 { + if v.Kind() == reflect.Interface { + v = v.Elem() + } + if v.Kind() != reflect.Pointer || v.IsNil() || v.Elem().Kind() != reflect.Struct { + return weak.Pointer[byte]{}, false + } + if ic, ok := interfaceField(v, "invocationContext"); ok { + return weakTo(ic.Elem()) + } + next, ok := interfaceField(v, "InvocationContext") + if !ok { + next, ok = interfaceField(v, "context") + } + if !ok { + if depth == 0 { + return weak.Pointer[byte]{}, false + } + return weakTo(v) + } + v = next } - ic, ok := f.Interface().(agent.InvocationContext) - if !ok { + return weak.Pointer[byte]{}, false +} + +// weakTo returns a weak pointer to what the pointer v points to. +func weakTo(v reflect.Value) (weak.Pointer[byte], bool) { + if v.Kind() != reflect.Pointer || v.IsNil() { return weak.Pointer[byte]{}, false } - p := reflect.ValueOf(ic) - if p.Kind() != reflect.Pointer || p.IsNil() { - return weak.Pointer[byte]{}, false + return weak.Make((*byte)(v.UnsafePointer())), true +} + +// interfaceField returns the non-nil interface field name of the struct v +// points to. +func interfaceField(v reflect.Value, name string) (reflect.Value, bool) { + f := v.Elem().FieldByName(name) + if !f.IsValid() || f.Kind() != reflect.Interface || f.IsNil() { + return reflect.Value{}, false } - return weak.Make((*byte)(p.UnsafePointer())), true + return f, true } -// startRun returns the state of the toolset in ctx's agent run, for Tools. It -// starts a new state when ctx belongs to another agent run than the state -// kept for the agent and invocation. When ctx does not tell its agent run, it -// drops the kept state and returns nil. +// startRun returns the state of the toolset in ctx's agent run, for Tools, +// and starts it if needed. When ctx does not tell its agent run, it returns +// nil. // // Without a state, the toolset lists and filters on every Tools and // ProcessRequest call, and its proxy tools pack themselves. // // When it starts a state, it also drops the states of finished agent runs, -// whose InvocationContext the garbage collector has reclaimed. No later call -// can belong to such a run, so a long conversation on one NewContext keeps -// only the states of runs that may still call the toolset. +// whose invocation the garbage collector has reclaimed. No later call can +// belong to such a run, so a long conversation on one NewContext keeps only +// the states of runs that may still call the toolset. func (l *mcpListings) startRun(m *mcpToolset, ctx agent.ReadonlyContext) *mcpRunState { if l == nil { return nil } run, ok := agentRunOf(ctx) - k := mcpAgentKey{m, ctx.InvocationID(), ctx.AgentName()} - l.mu.Lock() - defer l.mu.Unlock() if !ok { - delete(l.runs, k) return nil } - if s, ok := l.runs[k]; ok && s.run == run { + k := mcpRunKey{m, run} + l.mu.Lock() + defer l.mu.Unlock() + if s, ok := l.runs[k]; ok { return s } - for key, s := range l.runs { - if s.run.Value() == nil { + for key := range l.runs { + if key.run.Value() == nil { delete(l.runs, key) } } if l.runs == nil { - l.runs = map[mcpAgentKey]*mcpRunState{} + l.runs = map[mcpRunKey]*mcpRunState{} } - s := &mcpRunState{run: run} + s := &mcpRunState{} l.runs[k] = s return s } -// current returns the state the toolset's last Tools call for ctx's agent and -// invocation kept, or nil. ADK calls Tools before ProcessRequest in the first -// model step of every agent run. -func (l *mcpListings) current(m *mcpToolset, ctx agent.ReadonlyContext) *mcpRunState { +// current returns the state the toolset's Tools call kept for ctx's agent run, +// or nil. ADK calls Tools before ProcessRequest in the first model step of +// every agent run. +func (l *mcpListings) current(m *mcpToolset, ctx context.Context) *mcpRunState { if l == nil { return nil } + run, ok := agentRunOf(ctx) + if !ok { + return nil + } l.mu.Lock() defer l.mu.Unlock() - return l.runs[mcpAgentKey{m, ctx.InvocationID(), ctx.AgentName()}] + return l.runs[mcpRunKey{m, run}] } // processedList returns the list ToolFilter last allowed in the run, once the diff --git a/go/googleadk/mcp_listings_internal_test.go b/go/googleadk/mcp_listings_internal_test.go index 35f41b48..758f3eb2 100644 --- a/go/googleadk/mcp_listings_internal_test.go +++ b/go/googleadk/mcp_listings_internal_test.go @@ -66,8 +66,8 @@ func listingsLifetimeWorkflow(ctx workflow.Context, turns int) (listingsAfterTur runtime.GC() l.mu.Lock() live := 0 - for _, s := range l.runs { - if s.run.Value() != nil { + for k := range l.runs { + if k.run.Value() != nil { live++ } } diff --git a/go/googleadk/mcp_test.go b/go/googleadk/mcp_test.go index 52e42570..0c6e6280 100644 --- a/go/googleadk/mcp_test.go +++ b/go/googleadk/mcp_test.go @@ -25,6 +25,7 @@ import ( "google.golang.org/adk/v2/agent" "google.golang.org/adk/v2/agent/llmagent" "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/platform" "google.golang.org/adk/v2/runner" "google.golang.org/adk/v2/session" "google.golang.org/adk/v2/tool" @@ -664,10 +665,9 @@ func TestMcpSameNamedAgentsKeepTheirOwnListings(t *testing.T) { } } -// parallelHelperWorkflow has the root agent call one helper agent twice in -// parallel through an agent tool. The helper has two MCP toolsets. -func parallelHelperWorkflow(ctx workflow.Context, sequential bool) (runResult, error) { - helper, err := llmagent.New(llmagent.Config{ +// searchFetchHelper returns a helper agent with two MCP toolsets. +func searchFetchHelper() (agent.Agent, error) { + return llmagent.New(llmagent.Config{ Name: "helper", Description: "searches", Model: googleadk.NewModel("helper-model"), @@ -677,6 +677,15 @@ func parallelHelperWorkflow(ctx workflow.Context, sequential bool) (runResult, e googleadk.NewMCPToolset(googleadk.MCPToolsetOptions{Name: "fetch-server"}), }, }) +} + +// parallelHelperWorkflow runs one helper agent twice, in parallel unless +// sequential. Through an agent tool, the root agent calls it twice, and each +// call is its own invocation. With sharedInvocation, a custom root agent runs +// it on branches root.a and root.b through platform.RunTasks instead, so both +// runs share the root's invocation ID and the helper's name. +func parallelHelperWorkflow(ctx workflow.Context, sequential, sharedInvocation bool) (runResult, error) { + helper, err := searchFetchHelper() if err != nil { return runResult{}, err } @@ -684,58 +693,105 @@ func parallelHelperWorkflow(ctx workflow.Context, sequential bool) (runResult, e if sequential { opts = append(opts, googleadk.WithSequentialToolFanout()) } - return runAgent(ctx, agentBuild{ - ctxOpts: opts, - modelName: "root-model", - userMessage: "ask the helper twice", - tools: []tool.Tool{agenttool.New(helper, nil)}, + if !sharedInvocation { + return runAgent(ctx, agentBuild{ + ctxOpts: opts, + modelName: "root-model", + userMessage: "ask the helper twice", + tools: []tool.Tool{agenttool.New(helper, nil)}, + }) + } + root, err := agent.New(agent.Config{ + Name: "root", + SubAgents: []agent.Agent{helper}, + Run: func(ic agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + stopped := false + var tasks []func(context.Context) + for _, branch := range []string{"root.a", "root.b"} { + tasks = append(tasks, func(c context.Context) { + bctx := agent.PromoteWithDelta(ic, &agent.CommonContextDelta{ + InvocationContextDelta: &agent.InvocationContextDelta{Context: &c, Branch: &branch}, + }) + for ev, err := range helper.Run(bctx) { + if stopped || !yield(ev, err) { + stopped = true + return + } + } + }) + } + platform.RunTasks(ic, tasks) + } + }, }) + if err != nil { + return runResult{}, err + } + r, err := runner.New(runner.Config{AppName: "test-app", Agent: root, SessionService: session.InMemoryService(), AutoCreateSession: true}) + if err != nil { + return runResult{}, err + } + for _, rerr := range r.Run(googleadk.NewContext(ctx, opts...), "user-1", "session-1", genai.NewContentFromText("ask the helper twice", genai.RoleUser), agent.RunConfig{}) { + if rerr != nil { + return runResult{}, rerr + } + } + return runResult{}, nil } -// TestMcpParallelInvocationsKeepTheirOwnListings proves concurrent invocations -// of one agent each reuse the list their own Tools call took. Two helper runs +// TestMcpParallelInvocationsKeepTheirOwnListings proves concurrent runs of one +// agent each reuse the list their own Tools call took, whether they are +// separate invocations or share one invocation and agent name. Two helper runs // of two model steps each, with two MCP toolsets, list 8 times, in parallel as // in sequence. func TestMcpParallelInvocationsKeepTheirOwnListings(t *testing.T) { - for _, sequential := range []bool{false, true} { - t.Run(fmt.Sprintf("sequential=%v", sequential), func(t *testing.T) { - root := &ruleModel{name: "root-model", rule: func(req *model.LLMRequest) *model.LLMResponse { - if functionResponses(req) == 0 { - return reproFnCalls("helper", "helper") + for _, sharedInvocation := range []bool{false, true} { + for _, sequential := range []bool{false, true} { + t.Run(fmt.Sprintf("sharedInvocation=%v/sequential=%v", sharedInvocation, sequential), func(t *testing.T) { + root := &ruleModel{name: "root-model", rule: func(req *model.LLMRequest) *model.LLMResponse { + if functionResponses(req) == 0 { + return reproFnCalls("helper", "helper") + } + return googleadk.TextResponse("root done") + }} + helper := &ruleModel{name: "helper-model", rule: func(req *model.LLMRequest) *model.LLMResponse { + if lastIsFunctionResponse(req) { + return googleadk.TextResponse("helper done") + } + return googleadk.FunctionCallResponse("h1", "search", map[string]any{"query": "temporal"}) + }} + search := googleadk.NewFakeMCPServer("search-server").AddTool("search", "search the web", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{"hits": 1}, nil }) + fetch := googleadk.NewFakeMCPServer("fetch-server").AddTool("fetch", "fetch a page", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{}, nil }) + + var s testsuite.WorkflowTestSuite + env, counter := newEnv(t, &s, googleadk.Config{ + Models: map[string]googleadk.ModelFactory{"root-model": root.factory(), "helper-model": helper.factory()}, + MCPToolsets: map[string]googleadk.MCPFactory{ + "search-server": search.Factory(), + "fetch-server": fetch.Factory(), + }, + }) + env.RegisterWorkflow(parallelHelperWorkflow) + env.ExecuteWorkflow(parallelHelperWorkflow, sequential, sharedInvocation) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + rootCalls := 2 + if sharedInvocation { + rootCalls = 0 } - return googleadk.TextResponse("root done") - }} - helper := &ruleModel{name: "helper-model", rule: func(req *model.LLMRequest) *model.LLMResponse { - if lastIsFunctionResponse(req) { - return googleadk.TextResponse("helper done") + assert.Equal(t, 4+rootCalls, counter.get(googleadk.InvokeModelActivityName)) + assert.Equal(t, 8, counter.get(googleadk.ListMcpToolsActivityName), "one listing per toolset per helper model step") + assert.Equal(t, 2, counter.get(googleadk.CallMcpToolActivityName)) + assert.Len(t, helper.declared(), 4) + for i, names := range helper.declared() { + assert.ElementsMatch(t, []string{"search", "fetch"}, names, "helper request %d", i) } - return googleadk.FunctionCallResponse("h1", "search", map[string]any{"query": "temporal"}) - }} - search := googleadk.NewFakeMCPServer("search-server").AddTool("search", "search the web", searchSchema(), - func(map[string]any) (map[string]any, error) { return map[string]any{"hits": 1}, nil }) - fetch := googleadk.NewFakeMCPServer("fetch-server").AddTool("fetch", "fetch a page", searchSchema(), - func(map[string]any) (map[string]any, error) { return map[string]any{}, nil }) - - var s testsuite.WorkflowTestSuite - env, counter := newEnv(t, &s, googleadk.Config{ - Models: map[string]googleadk.ModelFactory{"root-model": root.factory(), "helper-model": helper.factory()}, - MCPToolsets: map[string]googleadk.MCPFactory{ - "search-server": search.Factory(), - "fetch-server": fetch.Factory(), - }, }) - env.RegisterWorkflow(parallelHelperWorkflow) - env.ExecuteWorkflow(parallelHelperWorkflow, sequential) - - require.True(t, env.IsWorkflowCompleted()) - require.NoError(t, env.GetWorkflowError()) - - assert.Equal(t, 6, counter.get(googleadk.InvokeModelActivityName)) - assert.Equal(t, 8, counter.get(googleadk.ListMcpToolsActivityName), "one listing per toolset per helper model step") - assert.Equal(t, 2, counter.get(googleadk.CallMcpToolActivityName)) - for i, names := range helper.declared() { - assert.ElementsMatch(t, []string{"search", "fetch"}, names, "helper request %d", i) - } - }) + } } } From eb361211907c2a605b2f3df75430c39c58f280b4 Mon Sep 17 00:00:00 2001 From: Francesc Campoy Date: Sat, 31 Oct 2026 02:38:23 -0700 Subject: [PATCH 10/10] Address review findings --- go/googleadk/mcp.go | 49 ++++++------ go/googleadk/mcp_test.go | 160 ++++++++++++++++++++++++++++++++------- 2 files changed, 160 insertions(+), 49 deletions(-) diff --git a/go/googleadk/mcp.go b/go/googleadk/mcp.go index e0fe91fa..b8717a56 100644 --- a/go/googleadk/mcp.go +++ b/go/googleadk/mcp.go @@ -65,16 +65,21 @@ func (m *mcpToolset) Name() string { return m.name } // toolset's ProcessRequest has run in the agent run, it lists them via the // ListMcpTools Activity, applies ToolFilter, and keeps the list for the // ProcessRequest call that follows in the same model step. After that, -// ProcessRequest lists, and Tools returns the list ToolFilter last allowed -// there, without an Activity and without running the filter. See -// ProcessRequest. +// ProcessRequest lists, and Tools returns the list it returned before, without +// an Activity and without running the filter. See ProcessRequest. // // So the first Tools call of each agent run lists and filters. ADK resumes a // confirmed tool call from this list before any ProcessRequest of the run, so // a tool the filter no longer allows does not run. The later calls, which ADK // v2.6.0 and later make before every model step, issue no Workflow commands, // so the toolset and its filter issue them at the same points on every ADK -// release. +// release. They return the list ADK v2.5.0 and earlier keep for the whole +// agent run, even when the MCP server or the filter has changed it since: +// ADK draws a UUID from the workflow's random stream for each tool it keeps +// before every model step, so the IDs it generates, such as tool confirmation +// IDs, stay the same on every release. The proxy tools from that list no +// longer pack themselves, so ProcessRequest alone decides what the model +// request declares and dispatches. // // A toolset wrapped in a toolset without ProcessRequest, such as // tool.FilterToolset, lists and filters on every call, because nothing else @@ -89,8 +94,9 @@ func (m *mcpToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { if err != nil { return nil, err } - l.keepPending(run, tools) - return m.allowed(ctx, tools), nil + allowed := m.allowed(ctx, tools) + l.keepPending(run, tools, allowed) + return allowed, nil } // allowed returns the tools the toolset's ToolFilter allows for ctx. @@ -149,9 +155,8 @@ func (m *mcpToolset) ProcessRequest(ctx agent.Context, req *model.LLMRequest) er return err } } - tools = m.allowed(ctx, tools) - l.processed(run, tools) - for _, t := range tools { + l.processed(run) + for _, t := range m.allowed(ctx, tools) { if err := packTool(req, t.(*mcpProxyTool)); err != nil { return err } @@ -195,9 +200,9 @@ type mcpRunState struct { // processed reports that the toolset's ProcessRequest has run in the agent // run. From then on ProcessRequest lists, and Tools does not. processed bool - // allowed is the list ToolFilter allowed in the last ProcessRequest, which - // Tools returns once processed. - allowed []tool.Tool + // returned is the list Tools last returned after listing, which it returns + // again once processed. + returned []tool.Tool } // listingsFrom returns the mcpListings NewContext installed, or nil. @@ -322,25 +327,26 @@ func (l *mcpListings) current(m *mcpToolset, ctx context.Context) *mcpRunState { return l.runs[mcpRunKey{m, run}] } -// processedList returns the list ToolFilter last allowed in the run, once the -// toolset's ProcessRequest has run there. +// processedList returns the list Tools returned after listing in the run, +// once the toolset's ProcessRequest has run there. func (l *mcpListings) processedList(s *mcpRunState) ([]tool.Tool, bool) { if s == nil { return nil, false } l.mu.Lock() defer l.mu.Unlock() - return s.allowed, s.processed + return s.returned, s.processed } -// keepPending keeps a list Tools took for the ProcessRequest call that follows. -func (l *mcpListings) keepPending(s *mcpRunState, tools []tool.Tool) { +// keepPending keeps a list Tools took for the ProcessRequest call that +// follows, and the list Tools returns. +func (l *mcpListings) keepPending(s *mcpRunState, tools, returned []tool.Tool) { if s == nil { return } l.mu.Lock() defer l.mu.Unlock() - s.pending, s.hasPending = tools, true + s.pending, s.hasPending, s.returned = tools, true, returned } // takePending returns and drops the list Tools kept, if any. @@ -355,15 +361,14 @@ func (l *mcpListings) takePending(s *mcpRunState) ([]tool.Tool, bool) { return tools, ok } -// processed marks the toolset's ProcessRequest as run in the agent run, and -// records the list ToolFilter allowed there. -func (l *mcpListings) processed(s *mcpRunState, allowed []tool.Tool) { +// processed marks the toolset's ProcessRequest as run in the agent run. +func (l *mcpListings) processed(s *mcpRunState) { if s == nil { return } l.mu.Lock() defer l.mu.Unlock() - s.processed, s.allowed = true, allowed + s.processed = true } // packedByToolset reports whether the toolset's ProcessRequest has run in the diff --git a/go/googleadk/mcp_test.go b/go/googleadk/mcp_test.go index 0c6e6280..82489580 100644 --- a/go/googleadk/mcp_test.go +++ b/go/googleadk/mcp_test.go @@ -452,43 +452,149 @@ func (h hidingToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { return kept, nil } +// searchDropped tells catalogChangeWorkflow's ToolFilter to drop "search". +var searchDropped atomic.Bool + +// catalogChangeWorkflow runs one turn of an MCP agent and returns how many +// UUIDs ADK drew from the context. With filterDrops, the toolset's ToolFilter +// drops "search" once searchDropped is set. +func catalogChangeWorkflow(ctx workflow.Context, filterDrops bool) (int, error) { + opts := googleadk.MCPToolsetOptions{Name: "search-server"} + if filterDrops { + opts.ToolFilter = func(_ agent.ReadonlyContext, t tool.Tool) bool { + return t.Name() != "search" || !searchDropped.Load() + } + } + root, err := llmagent.New(llmagent.Config{ + Name: "assistant", + Model: googleadk.NewModel("fake-model"), + Instruction: "be helpful", + Toolsets: []tool.Toolset{googleadk.NewMCPToolset(opts)}, + }) + if err != nil { + return 0, err + } + r, err := runner.New(runner.Config{AppName: "test-app", Agent: root, SessionService: session.InMemoryService(), AutoCreateSession: true}) + if err != nil { + return 0, err + } + draws := 0 + adkCtx := platform.WithUUIDProvider(googleadk.NewContext(ctx), func() string { + draws++ + return fmt.Sprintf("00000000-0000-0000-0000-%012d", draws) + }) + for _, rerr := range r.Run(adkCtx, "user-1", "session-1", genai.NewContentFromText("search, then fetch", genai.RoleUser), agent.RunConfig{}) { + if rerr != nil { + return 0, rerr + } + } + return draws, nil +} + // TestMcpDroppedToolIsNeitherDeclaredNorCalled proves each model request // declares and dispatches the toolset's current list. ADK v2.5.0 and earlier // keep the tools from the first Tools call of an agent run; if their proxies // still packed themselves, a tool the server dropped would stay declared and // a call to it would run CallMcpTool there, but not with ADK v2.6.0. +// +// It also proves a change in the list does not change the IDs ADK generates. +// ADK draws a UUID for each tool it kept from Tools before every model step: +// v2.5.0 and earlier keep the first list of the agent run, and v2.6.0 keeps +// what each step's Tools call returns. If those calls returned the current +// list, a dropped tool would shift the IDs drawn after it, tool confirmation +// IDs included, on v2.6.0 only, and Histories would not replay across releases. func TestMcpDroppedToolIsNeitherDeclaredNorCalled(t *testing.T) { - var hide atomic.Bool - server := googleadk.NewFakeMCPServer("search-server"). - AddTool("search", "search the web", searchSchema(), - func(map[string]any) (map[string]any, error) { - hide.Store(true) - return map[string]any{"hits": 1}, nil - }). - AddTool("fetch", "fetch a page", searchSchema(), - func(map[string]any) (map[string]any, error) { return map[string]any{}, nil }) - m := &ruleModel{name: "fake-model", rule: func(req *model.LLMRequest) *model.LLMResponse { - if n := functionResponses(req); n < 2 { - return googleadk.FunctionCallResponse(fmt.Sprintf("c%d", n+1), "search", map[string]any{"query": "temporal"}) - } - return googleadk.TextResponse("done") - }} + t.Run("not called", func(t *testing.T) { + var hide atomic.Bool + server := googleadk.NewFakeMCPServer("search-server"). + AddTool("search", "search the web", searchSchema(), + func(map[string]any) (map[string]any, error) { + hide.Store(true) + return map[string]any{"hits": 1}, nil + }). + AddTool("fetch", "fetch a page", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{}, nil }) + m := &ruleModel{name: "fake-model", rule: func(req *model.LLMRequest) *model.LLMResponse { + if n := functionResponses(req); n < 2 { + return googleadk.FunctionCallResponse(fmt.Sprintf("c%d", n+1), "search", map[string]any{"query": "temporal"}) + } + return googleadk.TextResponse("done") + }} + + var s testsuite.WorkflowTestSuite + env, counter := newEnv(t, &s, googleadk.Config{ + Models: map[string]googleadk.ModelFactory{"fake-model": m.factory()}, + MCPToolsets: map[string]googleadk.MCPFactory{"search-server": func(context.Context) (tool.Toolset, error) { + return hidingToolset{Toolset: server, hide: &hide}, nil + }}, + }) + env.ExecuteWorkflow(agentRunWorkflow, runInput{ModelName: "fake-model", UserMessage: "search twice", MCPToolset: "search-server"}) - var s testsuite.WorkflowTestSuite - env, counter := newEnv(t, &s, googleadk.Config{ - Models: map[string]googleadk.ModelFactory{"fake-model": m.factory()}, - MCPToolsets: map[string]googleadk.MCPFactory{"search-server": func(context.Context) (tool.Toolset, error) { - return hidingToolset{Toolset: server, hide: &hide}, nil - }}, + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + assert.Equal(t, [][]string{{"search", "fetch"}, {"fetch"}, {"fetch"}}, m.declared()) + assert.Equal(t, 3, counter.get(googleadk.ListMcpToolsActivityName)) + assert.Equal(t, 1, counter.get(googleadk.CallMcpToolActivityName), "the dropped tool ran again") }) - env.ExecuteWorkflow(agentRunWorkflow, runInput{ModelName: "fake-model", UserMessage: "search twice", MCPToolset: "search-server"}) - require.True(t, env.IsWorkflowCompleted()) - require.NoError(t, env.GetWorkflowError()) + t.Run("same generated IDs", func(t *testing.T) { + draws := map[string]int{} + for _, tc := range []struct { + name string + drop, filterDrop bool + want [][]string + }{ + {"unchanged", false, false, [][]string{{"search", "fetch"}, {"search", "fetch"}, {"search", "fetch"}}}, + {"server drops search", true, false, [][]string{{"search", "fetch"}, {"fetch"}, {"fetch"}}}, + {"filter drops search", true, true, [][]string{{"search", "fetch"}, {"fetch"}, {"fetch"}}}, + } { + searchDropped.Store(false) + server := googleadk.NewFakeMCPServer("search-server"). + AddTool("search", "search the web", searchSchema(), + func(map[string]any) (map[string]any, error) { + searchDropped.Store(tc.drop) + return map[string]any{"hits": 1}, nil + }). + AddTool("fetch", "fetch a page", searchSchema(), + func(map[string]any) (map[string]any, error) { return map[string]any{}, nil }) + var ts tool.Toolset = server + if !tc.filterDrop { + ts = hidingToolset{Toolset: server, hide: &searchDropped} + } + m := &ruleModel{name: "fake-model", rule: func(req *model.LLMRequest) *model.LLMResponse { + switch functionResponses(req) { + case 0: + return googleadk.FunctionCallResponse("c1", "search", map[string]any{"query": "temporal"}) + case 1: + return googleadk.FunctionCallResponse("c2", "fetch", map[string]any{"query": "temporal"}) + } + return googleadk.TextResponse("done") + }} + + var s testsuite.WorkflowTestSuite + env, counter := newEnv(t, &s, googleadk.Config{ + Models: map[string]googleadk.ModelFactory{"fake-model": m.factory()}, + MCPToolsets: map[string]googleadk.MCPFactory{"search-server": func(context.Context) (tool.Toolset, error) { + return ts, nil + }}, + }) + env.RegisterWorkflow(catalogChangeWorkflow) + env.ExecuteWorkflow(catalogChangeWorkflow, tc.filterDrop) - assert.Equal(t, [][]string{{"search", "fetch"}, {"fetch"}, {"fetch"}}, m.declared()) - assert.Equal(t, 3, counter.get(googleadk.ListMcpToolsActivityName)) - assert.Equal(t, 1, counter.get(googleadk.CallMcpToolActivityName), "the dropped tool ran again") + require.True(t, env.IsWorkflowCompleted(), tc.name) + require.NoError(t, env.GetWorkflowError(), tc.name) + var n int + require.NoError(t, env.GetWorkflowResult(&n), tc.name) + draws[tc.name] = n + + assert.Equal(t, tc.want, m.declared(), tc.name) + assert.Equal(t, 2, counter.get(googleadk.CallMcpToolActivityName), tc.name) + } + searchDropped.Store(false) + assert.Equal(t, draws["unchanged"], draws["server drops search"], "UUIDs drawn: %v", draws) + assert.Equal(t, draws["unchanged"], draws["filter drops search"], "UUIDs drawn: %v", draws) + }) } // wrappedToolsLog records, for TestMcpWrappedToolsetListsOnEveryToolsCall, the