🤖 AI-generated review. I looked through them to make sure they weren't crazy and they all seem worth looking at.

Comment left by Claude on behalf of @EItanya.

Reviewed at 14fe59397bbddb01103294af0b0e5f105219d4a2.

Making startup and resume independent of MCP discovery is the right goal. I think most of the new machinery can go away, though. mcptoolset already connects lazily ("MCP session is created lazily on the first request to LLM"), so the startup half is mostly a deletion. The added complexity is all in getting MCP Apps _meta classification, and the ADK and go-sdk can give us that without a second connection or a cache.

1. Token passthrough: every caller is limited to the first caller's tool list

mcpAppToolset.classify caches the classification (a StringPredicate of tool names) on the toolset for the life of the process. That classification is made with the first caller's credentials. Every later caller's tool list is then filtered by those names. If a passthrough server returns different tools per token, later users only get the overlap. On main, passthrough servers had no name filter, so this is a regression for them. The cache also never refreshes, so tools a server adds later stay hidden until the runtime restarts.

Repro: one shared MCP session, with a server that filters tools/list by the request's Authorization header. Alice gets [common, board]. Bob should also get admin_only, but gets [common, board].

package mcp

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"

	mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
	"github.com/stretchr/testify/require"

	"github.com/kagent-dev/kagent/go/adk/pkg/models"
)

// perRequestAuthServer filters tools/list per request Authorization header on one shared MCP session.
func perRequestAuthServer(t *testing.T) *httptest.Server {
	s := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "p", Version: "1"}, nil)
	h := func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) {
		return nil, nil, nil
	}
	mcpsdk.AddTool(s, &mcpsdk.Tool{Name: "common"}, h)
	mcpsdk.AddTool(s, &mcpsdk.Tool{Name: "board", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://b"}}}, h)
	mcpsdk.AddTool(s, &mcpsdk.Tool{Name: "refresh", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://b", "visibility": []any{"app"}}}}, h)
	mcpsdk.AddTool(s, &mcpsdk.Tool{Name: "admin_only"}, h)
	s.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler {
		return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) {
			res, err := next(ctx, method, req)
			if lr, ok := res.(*mcpsdk.ListToolsResult); ok && req.GetExtra().Header.Get("Authorization") != "Bearer bob" {
				out := lr.Tools[:0:0]
				for _, tl := range lr.Tools {
					if tl.Name != "admin_only" {
						out = append(out, tl)
					}
				}
				lr.Tools = out
			}
			return res, err
		}
	})
	srv := httptest.NewServer(mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return s }, nil))
	t.Cleanup(srv.Close)
	t.Cleanup(srv.CloseClientConnections)
	return srv
}

func TestProbePRCachesFirstCallerToolList(t *testing.T) {
	srv := perRequestAuthServer(t)
	ts, err := initializeToolSet(t.Context(), mcpServerParams{URL: srv.URL, ServerType: "http", TokenPassthrough: true}, nil)
	require.NoError(t, err)
	names := func(tok string) []string {
		ctx := WithDiscoveryScope(context.WithValue(t.Context(), models.BearerTokenKey, tok))
		tools, err := ts.Tools(testReadonlyContext{Context: ctx})
		require.NoError(t, err)
		var out []string
		for _, tl := range tools {
			out = append(out, tl.Name())
		}
		return out
	}
	require.ElementsMatch(t, []string{"common", "board"}, names("alice"))
	require.ElementsMatch(t, []string{"common", "board", "admin_only"}, names("bob")) // fails: admin_only missing
}

2. Classify from ADK's own tools/list instead of a separate discovery path

Right now classification opens its own connection and repeats tools/list with its own pagination and cursor-loop check. For stdio servers that starts a second process. The classify loop (mutex, wait channel, retry loop, omitted-check before and after locking) and WithDiscoveryScope exist only to share and cache that extra call. WithDiscoveryScope is also a new public API passed implicitly through ctx, and every runner.Run caller has to remember to add it.

Everything needed for a stateless version already exists:

• mcptoolset.Config.Client lets us pass our own *mcpsdk.Client. This PR already does.
• (*mcpsdk.Client).AddSendingMiddleware sees the raw *ListToolsResult, _meta included, for ADK's own tools/list calls, including every page. It runs with the caller's ctx, so results can be collected per call even though the MCP session is shared.
• ADK rebuilds the tool list before every model step (toolProcessor), and ProcessRequest puts each tool object into req.Tools[name]. scopedTool in baggage.go already uses this to swap in a wrapper.

Sketch:

func newMCPAppClient() *mcpsdk.Client {
	c := mcpsdk.NewClient(impl, &mcpsdk.ClientOptions{Capabilities: mcpUIClientCapabilities()})
	c.AddSendingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler {
		return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) {
			res, err := next(ctx, method, req)
			if kinds, ok := ctx.Value(toolKindsKey{}).(*toolKinds); ok {
				if list, ok := res.(*mcpsdk.ListToolsResult); ok {
					kinds.record(list.Tools) // name -> mcpToolKindOf(t.Meta)
				}
			}
			return res, err
		}
	})
	return c
}

func (m *mcpAppToolset) Tools(ctx adkagent.ReadonlyContext) ([]tool.Tool, error) {
	kinds := newToolKinds()
	tools, err := m.inner.Tools(readonlyValuesContext{ctx, context.WithValue(ctx, toolKindsKey{}, kinds)})
	if err != nil {
		return nil, err
	}
	visible := tools[:0:0]
	for _, t := range tools {
		switch kind, ok := kinds.get(t.Name()); {
		case !ok || kind == mcpToolKindAppInternal:
			continue // never classified, so hidden
		case kind == mcpToolKindApp:
			visible = append(visible, appTool{t.(functionTool)}) // re-registers itself in req.Tools, like scopedTool
		default:
			visible = append(visible, t)
		}
	}
	return visible, nil
}

// Before-model callback: no toolset closure, no shared lookups.
if _, ok := req.Tools[part.FunctionResponse.Name].(appTool); ok { compact(...) }

The configured tool allowlist can stay as tool.FilterToolset on the inner toolset.

This removes:

• the second connection (and second stdio process), plus the hand-written pagination and cursor guard
• the cache, mutex, wait channel and retry loop
• WithDiscoveryScope, including its wiring in executor.go and the oneshot example
• MCPAppToolNamesFromToolsets

It also fixes #1 by construction, because classification comes from each caller's own response. It stays fail-closed, since a tool can't be exposed without its _meta having been read. I prototyped the middleware against the real mcptoolset with token passthrough: Alice gets [common, board] with board detected as an app tool and the app-only refresh hidden, and Bob also gets admin_only. I did not build the appTool wrapper or the compaction callback end to end.

One separate decision: today the PR leaves a server out for the rest of an execution when classification fails, but an ADK discovery error after classification still fails the run. With one merged call, this becomes a single policy question: should an unreachable MCP server fail the run or be skipped? If skipping is wanted, it's a few lines in mcpAppToolset.Tools. I'd suggest making that its own change rather than bundling it here.

3. Small cleanup

MakeMCPAppModelResultCallback is now only used by tests. Production uses the new private version. The redesign above removes the need for either.

Checks run

• The PR's tests for adk/pkg/mcp, adk/pkg/agent and adk/pkg/a2a/... pass with -race -count=2.
• The repro above fails at 14fe593.
• The middleware prototype passes.

🤖 AI-generated review. I looked through them to make sure they weren't crazy and they all seem worth looking at.

Comment left by Claude on behalf of @EItanya.

Follow-up on 4a3ec2907e0b8023291d663b5ffad8f1d0bb2210.

Thanks for switching to middleware. Filtering each raw tools/list page there, so the SDK caches already-filtered results, is a better approach than what I sketched: my per-call version would have hidden every tool on SDK cache hits, which skip middleware. Agreed on passthrough too. With one execution per Actor, per-caller visibility is a policy question, not an SDK one.

A few things still look like more machinery than this change needs:

1. Omitting failed toolsets is still a policy change

Thanks for removing the sticky classified flag. A failed listing is now omitted for the execution whether or not that toolset worked before, and I confirmed that on 4a3ec29 with the same repro. That fixes the inconsistency.

It is still a behavior change from main, though, where a discovery error fails the step. Whether an unreachable server should be skipped or fail the run seems worth deciding in its own PR. I'd drop WithDiscoveryScope and the omission logic here, so this PR only defers discovery.

2. The lock and extra state guard a case that doesn't happen

listingMu serializes every Tools() call on a toolset, held across the network call, to stop overlapping listings from mixing app classifications. The Go runtime admits one execution per Actor, and each agent and subagent builds its own toolsets, so overlapping listings on one toolset effectively don't occur. The only thing that needs to survive SDK cache hits is which names are app tools, for compaction.

Suggestion: have the middleware write name → isApp into one map per toolset on each fresh listing, and have the compaction callback read it. Remove listingMu and the second app-name map.

3. Cursor rewriting is out of scope

ADK's connectionRefresher.ListTools has no pagination-cycle check, and main has the same gap. Encoding every visited cursor into the opaque cursor (which grows with each page) and rebuilding the outgoing ClientRequest is a lot to carry in a PR about deferring discovery. I'd report it to ADK upstream or fix it separately.

4. Docs

The new mcp.md section mostly describes internals: locking, cursor encoding, and when app names are published. I'd trim it to what users and operators see. Discovery is lazy, app-only tools are hidden, and app results are compacted for the model.

With 1–3, discovery.go should shrink to roughly the middleware plus one map, and most of discovery_test.go with it.

Checks: the PR's tests for adk/pkg/mcp, adk/pkg/agent and adk/pkg/a2a/... pass with -race -count=2. Checked on 4a3ec29.