From 13a757d2678e9086756206fc99b3f770fe680ae9 Mon Sep 17 00:00:00 2001 From: Shashank <21697719+shashankram@users.noreply.github.com> Date: Fri, 6 Nov 2026 23:17:53 +0000 Subject: [PATCH] adk/mcp: defer tool discovery until invocation (#3560) # Description Slow MCP discovery can exceed router timeouts and fail agent resumes, holding worker reservations and capacity while startup waits. Separate classification also duplicates MCP connections and paginated listings. This fixes it by: - Deferring HTTP, SSE, and stdio MCP connections until invocation preparation. - Classifying and filtering ADK's own discovery responses before the SDK caches them, preserving server TTLs and catalog notification handling. - Omitting failed toolsets for the execution regardless of prior discovery success, keeping healthy tools available and retrying on later executions without exposing unclassified tools. - Serializing discovery and classification publication per toolset, with cancellation support for queued invocations. - Preserving MCP App result compaction through a single callback, alongside rich stored results, credentials, approvals, and tracing. Tools reclassified as ordinary keep their real output; removed app names remain available for historical compaction. - Documenting discovery, SDK caching, pagination guards, and failure handling in `docs/architecture/mcp.md`. # Change Type ``` /kind fix ``` # Testing - Added unit coverage for zero startup requests, one paginated discovery, filtered SDK cache hits, TTL expiry, catalog notifications, partial failures, execution-scoped retries after initial and later failures, cancellation, concurrent discovery, reclassification, and stdio startup. - Added app-result coverage for cached tools, approvals, and historical results after catalog removal, preserving the original rich payloads. - Added E2E coverage for slow MCP discovery across repeated resumes and initial discovery failure and recovery; extended app-history E2E coverage to verify SDK caching across model steps and rediscovery after restore. # Changelog ```release-note MCP discovery no longer delays startup or resume and caches filtered tools safely. ``` --------- Signed-off-by: Shashank Ram (cherry picked from commit ccb2dd350e06be990618f3e464fbb07d0d67ca3d) Signed-off-by: Eitan Yarmush --- docs/architecture/mcp.md | 27 + go/adk/examples/oneshot/main.go | 3 +- go/adk/pkg/a2a/executor.go | 2 + go/adk/pkg/a2a/executor_mcp_test.go | 118 +++ go/adk/pkg/a2a/executor_test.go | 4 + .../pkg/a2a/server/lazy_mcp_startup_test.go | 76 ++ go/adk/pkg/agent/agent.go | 12 +- go/adk/pkg/agent/mcp_apps.go | 10 +- go/adk/pkg/agent/mcp_apps_test.go | 156 +++- go/adk/pkg/mcp/discovery.go | 160 +++++ go/adk/pkg/mcp/discovery_test.go | 675 ++++++++++++++++++ go/adk/pkg/mcp/mcp_ui.go | 92 --- go/adk/pkg/mcp/passthrough_test.go | 8 +- go/adk/pkg/mcp/registry.go | 42 +- go/adk/pkg/mcp/registry_test.go | 6 +- go/core/test/e2e/mcp_discovery_test.go | 440 ++++++++++++ 16 files changed, 1678 insertions(+), 153 deletions(-) create mode 100644 go/adk/pkg/a2a/executor_mcp_test.go create mode 100644 go/adk/pkg/a2a/server/lazy_mcp_startup_test.go create mode 100644 go/adk/pkg/mcp/discovery.go create mode 100644 go/adk/pkg/mcp/discovery_test.go create mode 100644 go/core/test/e2e/mcp_discovery_test.go diff --git a/docs/architecture/mcp.md b/docs/architecture/mcp.md index 661026b4c3..0ec7af648b 100644 --- a/docs/architecture/mcp.md +++ b/docs/architecture/mcp.md @@ -75,3 +75,30 @@ not protected: registering a name or URI core already uses replaces it. Authorization inside an added tool stays with the program that adds it. The implementation is in [`go/core/internal/mcp`](../../go/core/internal/mcp). + +## Go ADK runtime discovery + +The Go ADK connects to MCP servers and discovers tools during invocation +preparation. Startup and Actor resume do not wait for HTTP or SSE connections +or start stdio MCP processes. This also applies to subagents. + +Discovery filters app-only tools and applies configured allowlists before tools +reach the model. The MCP SDK caches filtered responses and owns TTL expiry and +catalog-notification invalidation. Cached tools remain filtered. Credentials, +approvals, and tool-call authorization continue to apply. + +If discovery fails, that server's toolset stays hidden for the current A2A +execution, including subsequent model steps and subagent calls. Healthy toolsets +remain available. The next execution retries, including an approval continuation. +There is no fallback to unclassified tools. Discovery failures are logged as +warnings. Cancellation and tool-call errors retain their existing handling. + +MCP App results are compacted for the model while stored results remain intact +for UI rendering. Removed app names remain recognized for historical compaction +during the process lifetime. Tools reclassified as ordinary keep their full +output, including historical results. A new process rediscovers app metadata +during invocation preparation. + +Standalone runners can use `mcp.WithDiscoveryScope(ctx)` for each `runner.Run` +to get the same omission behavior. Without it, discovery errors propagate. +Controller-side catalog discovery used by the UI is independent. diff --git a/go/adk/examples/oneshot/main.go b/go/adk/examples/oneshot/main.go index 4e7709ef1d..ff4fb37f6f 100644 --- a/go/adk/examples/oneshot/main.go +++ b/go/adk/examples/oneshot/main.go @@ -27,6 +27,7 @@ import ( "github.com/kagent-dev/kagent/go/adk/pkg/agent" "github.com/kagent-dev/kagent/go/adk/pkg/config" + "github.com/kagent-dev/kagent/go/adk/pkg/mcp" "github.com/kagent-dev/kagent/go/pkg/logging" ) @@ -106,7 +107,7 @@ func main() { streamMode = adkagent.StreamingModeSSE } - for ev, err := range r.Run(ctx, "user", sess.Session.ID(), content, adkagent.RunConfig{ + for ev, err := range r.Run(mcp.WithDiscoveryScope(ctx), "user", sess.Session.ID(), content, adkagent.RunConfig{ StreamingMode: streamMode, }) { if err != nil { diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index dd6469b7cf..03b4c845b3 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -22,6 +22,7 @@ import ( "google.golang.org/genai" "github.com/kagent-dev/kagent/go/adk/pkg/auth" + "github.com/kagent-dev/kagent/go/adk/pkg/mcp" "github.com/kagent-dev/kagent/go/adk/pkg/models" "github.com/kagent-dev/kagent/go/adk/pkg/telemetry" apia2a "github.com/kagent-dev/kagent/go/api/a2a" @@ -264,6 +265,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon ctx = context.WithValue(ctx, publicContextIDKey{}, reqCtx.ContextID) ctx = withBearerToken(ctx) + ctx = mcp.WithDiscoveryScope(ctx) // The synthetic ADK user ID is only a native session lookup key. Memory // and outgoing credentials must receive only the passed-through caller. ctx = auth.WithUserID(ctx, trustedUserID) diff --git a/go/adk/pkg/a2a/executor_mcp_test.go b/go/adk/pkg/a2a/executor_mcp_test.go new file mode 100644 index 0000000000..2e23f70fb7 --- /dev/null +++ b/go/adk/pkg/a2a/executor_mcp_test.go @@ -0,0 +1,118 @@ +package a2a + +import ( + "context" + "errors" + "log/slog" + "net/http" + "net/http/httptest" + "slices" + "sync/atomic" + "testing" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/require" + adkagent "google.golang.org/adk/v2/agent" + + "github.com/kagent-dev/kagent/go/adk/pkg/mcp" + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" +) + +func TestKAgentExecutorScopesMCPDiscoveryFailures(t *testing.T) { + // Given a temporarily unavailable MCP catalog, when execution continues and + // later resumes with reused IDs, assert that omission lasts only one execution. + var unavailable atomic.Bool + unavailable.Store(true) + var listings atomic.Int32 + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "execution-discovery", Version: "1"}, &mcpsdk.ServerOptions{ + SupportedProtocolVersions: []string{"2025-11-25"}, + }) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "available_tool"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + return nil, map[string]any{"result": "ok"}, nil + }) + server.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + if method == "tools/list" { + listings.Add(1) + if unavailable.Load() { + return nil, errors.New("catalog temporarily unavailable") + } + } + return next(ctx, method, req) + } + }) + httpServer := httptest.NewServer(mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, &mcpsdk.StreamableHTTPOptions{Stateless: true})) + t.Cleanup(httpServer.Close) + healthy := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "healthy-discovery"}, &mcpsdk.ServerOptions{SupportedProtocolVersions: []string{"2025-11-25"}}) + mcpsdk.AddTool(healthy, &mcpsdk.Tool{Name: "healthy_tool"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + return nil, map[string]any{"result": "healthy"}, nil + }) + healthyHTTP := httptest.NewServer(mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return healthy }, &mcpsdk.StreamableHTTPOptions{Stateless: true})) + t.Cleanup(healthyHTTP.Close) + toolsets := mcp.CreateToolsets(t.Context(), []*configv1.ADKHTTPMCPServerConfig{{ + Params: &configv1.ADKHTTPConnectionParams{Url: httpServer.URL}, + }, { + Params: &configv1.ADKHTTPConnectionParams{Url: healthyHTTP.URL}, + }}, nil, nil, false, nil) + require.Len(t, toolsets, 2) + + var visible [][]string + builtin := &recordingExecutor{onExecute: func(ctx context.Context) { + readonly := &executionDiscoveryContext{StrictContextMock: adkagent.NewStrictContextMock(ctx)} + for range 2 { + var names []string + for _, toolset := range toolsets { + tools, err := toolset.Tools(readonly) + require.NoError(t, err, "discovery failures must leave the execution usable") + for _, tool := range tools { + names = append(names, tool.Name()) + } + } + slices.Sort(names) + visible = append(visible, names) + unavailable.Store(false) + } + }} + executor := &KAgentExecutor{builtin: builtin, logger: slog.New(slog.DiscardHandler)} + request := &a2asrv.ExecutorContext{ + TaskID: "reused-task", ContextID: "reused-session", + Message: a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("continue")), + } + for _, err := range executor.Execute(t.Context(), request) { + require.NoError(t, err) + } + require.Equal(t, [][]string{{"healthy_tool"}, {"healthy_tool"}}, visible, "recovery during an execution must not retry its omitted server") + require.EqualValues(t, 1, listings.Load()) + + visible = nil + for _, err := range executor.Execute(t.Context(), request) { + require.NoError(t, err) + } + require.Equal(t, [][]string{{"available_tool", "healthy_tool"}, {"available_tool", "healthy_tool"}}, visible, "a later execution with the same IDs must retry discovery") + require.Greater(t, listings.Load(), int32(1)) + + // Given prior successful discovery, when the server goes down again, assert + // that healthy tools remain available and a fresh execution retries recovery. + unavailable.Store(true) + before := listings.Load() + visible = nil + for _, err := range executor.Execute(t.Context(), request) { + require.NoError(t, err) + } + require.Equal(t, [][]string{{"healthy_tool"}, {"healthy_tool"}}, visible) + require.Equal(t, before+1, listings.Load(), "the failed catalog must stay omitted across model steps") + visible = nil + for _, err := range executor.Execute(t.Context(), request) { + require.NoError(t, err) + } + require.Equal(t, [][]string{{"available_tool", "healthy_tool"}, {"available_tool", "healthy_tool"}}, visible) +} + +type executionDiscoveryContext struct { + adkagent.StrictContextMock +} + +func (*executionDiscoveryContext) InvocationID() string { return "reused-invocation" } +func (*executionDiscoveryContext) SessionID() string { return "reused-session" } diff --git a/go/adk/pkg/a2a/executor_test.go b/go/adk/pkg/a2a/executor_test.go index 4c557f9b4c..3d8e8f8251 100644 --- a/go/adk/pkg/a2a/executor_test.go +++ b/go/adk/pkg/a2a/executor_test.go @@ -40,6 +40,7 @@ type recordingExecutor struct { userID string cleanupCalled bool events []a2atype.Event + onExecute func(context.Context) // span, when set, is a span the turn emits, the way ADK does. span string } @@ -48,6 +49,9 @@ func (e *recordingExecutor) Execute(ctx context.Context, reqCtx *a2asrv.Executor return func(yield func(a2atype.Event, error) bool) { e.message = reqCtx.Message e.userID = auth.UserIDFromContext(ctx) + if e.onExecute != nil { + e.onExecute(ctx) + } if e.span != "" { _, span := otel.Tracer("adk-test").Start(ctx, e.span) span.End() diff --git a/go/adk/pkg/a2a/server/lazy_mcp_startup_test.go b/go/adk/pkg/a2a/server/lazy_mcp_startup_test.go new file mode 100644 index 0000000000..e0a9dd8160 --- /dev/null +++ b/go/adk/pkg/a2a/server/lazy_mcp_startup_test.go @@ -0,0 +1,76 @@ +package server + +import ( + "context" + "io" + "log/slog" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync/atomic" + "testing" + "time" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/types/known/durationpb" + + kagentagent "github.com/kagent-dev/kagent/go/adk/pkg/agent" + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" +) + +func TestReadinessDoesNotWaitForMCPDiscovery(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "test-key") + t.Run("Given slow MCP servers and MCP-bound subagents, when the runtime starts, assert that readiness precedes discovery", func(t *testing.T) { + var requests atomic.Int32 + mcpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + select { + case <-time.After(500 * time.Millisecond): + case <-r.Context().Done(): + } + http.Error(w, "slow server", http.StatusServiceUnavailable) + })) + t.Cleanup(mcpServer.Close) + marker := filepath.Join(t.TempDir(), "stdio-started") + config := func(name string) *configv1.ADKAgentConfig { + return &configv1.ADKAgentConfig{ + Name: name, Instruction: "Use the configured tools.", SkillsDirectory: t.TempDir(), + Model: &configv1.ADKModelConfig{Model: "test-model", Provider: &configv1.ADKModelConfig_Openai{Openai: &configv1.ADKOpenAI{}}}, + HttpTools: []*configv1.ADKHTTPMCPServerConfig{{Params: &configv1.ADKHTTPConnectionParams{Url: mcpServer.URL, Timeout: durationpb.New(time.Second)}}}, + SseTools: []*configv1.ADKSSEMCPServerConfig{{Params: &configv1.ADKSSEConnectionParams{Url: mcpServer.URL, Timeout: durationpb.New(time.Second)}}}, + } + } + rootConfig := config("root") + rootConfig.StdioTools = []*configv1.ADKStdioMCPServerConfig{{Command: "touch", Args: []string{marker}}} + rootConfig.SubAgents = []*configv1.ADKAgentConfig{config("child")} + _, err := kagentagent.CreateGoogleADKAgent(t.Context(), &configv1.ADKConfig{Agent: rootConfig}, "root", nil) + require.NoError(t, err) + srv, err := NewA2AServer(a2atype.AgentCard{}, substrateExecutor{}, slog.New(slog.DiscardHandler), ServerConfig{Host: "127.0.0.1", Port: "0"}) + require.NoError(t, err) + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + srv.readyServer.Addr = listener.Addr().String() + require.NoError(t, listener.Close()) + require.NoError(t, srv.Start()) + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + require.NoError(t, srv.httpServer.Shutdown(ctx)) + require.NoError(t, srv.readyServer.Shutdown(ctx)) + srv.grpcServer.Stop() + }) + client := &http.Client{Timeout: time.Second} + resp, err := client.Get("http://" + srv.readyServer.Addr + "/readyz") + require.NoError(t, err) + _, err = io.Copy(io.Discard, resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Zero(t, requests.Load(), "root and subagent MCP connections must be deferred") + _, err = os.Stat(marker) + require.ErrorIs(t, err, os.ErrNotExist, "stdio process must not start before readiness") + }) +} diff --git a/go/adk/pkg/agent/agent.go b/go/adk/pkg/agent/agent.go index 64413eacba..b5363b6288 100644 --- a/go/adk/pkg/agent/agent.go +++ b/go/adk/pkg/agent/agent.go @@ -98,7 +98,6 @@ func createGoogleADKAgent(ctx context.Context, agentConfig *configv1.ADKAgentCon log.InfoContext(ctx, "wired local skills", "skills_directory", skillsDirectory, "skill_count", len(skills)) } } - mcpAppToolNames := mcp.MCPAppToolNamesFromToolsets(toolsets) var remoteAgentTools []tool.Tool for _, remoteAgent := range agentConfig.RemoteAgents { @@ -150,12 +149,11 @@ func createGoogleADKAgent(ctx context.Context, agentConfig *configv1.ADKAgentCon beforeToolCallbacks := []llmagent.BeforeToolCallback{} beforeModelCallbacks := []llmagent.BeforeModelCallback{} - if len(mcpAppToolNames) > 0 { - // For MCP App-capable tools, keep rich tool payloads in chat history for UI rendering, - // but compact what is sent back to the model to avoid redundant polling/tool churn. - log.InfoContext(ctx, "wiring MCP App model result callback", "tool_count", len(mcpAppToolNames)) - beforeModelCallbacks = append(beforeModelCallbacks, MakeMCPAppModelResultCallback(mcpAppToolNames)) - } + // Tools are classified before each model request is prepared. Resolve app + // names here rather than capturing an empty snapshot during construction. + beforeModelCallbacks = append(beforeModelCallbacks, makeMCPAppModelResultCallback(func(ctx agent.Context) map[string]bool { + return mcp.MCPAppToolNamesFromToolsets(ctx, toolsets) + })) beforeToolCallbacks = append(beforeToolCallbacks, makeBeforeToolCallback(log)) llmAgentConfig := llmagent.Config{ diff --git a/go/adk/pkg/agent/mcp_apps.go b/go/adk/pkg/agent/mcp_apps.go index 38d10c4ff5..50617b6b8d 100644 --- a/go/adk/pkg/agent/mcp_apps.go +++ b/go/adk/pkg/agent/mcp_apps.go @@ -19,20 +19,22 @@ import ( // tool carrying a UI resourceUri, independent of the tool's name or payload keys. const mcpAppRenderedNotice = "The interactive UI for this tool has been rendered to the user and now updates live inside the app. Treat this as complete and do not call this tool again unless the user explicitly asks for it." -// MakeMCPAppModelResultCallback replaces what the model sees for MCP App +// makeMCPAppModelResultCallback replaces what the model sees for MCP App // (UI-rendering) tool results: instead of the heavy render payload it receives a // terminal directive (see mcpAppRenderedNotice). The full result is still // streamed to the UI separately, so this only changes the model's view and // prevents the model from looping on the rendering tool. Errors are passed // through so the model can still react to and recover from failures. -func MakeMCPAppModelResultCallback(appToolNames map[string]bool) llmagent.BeforeModelCallback { - return func(_ agent.Context, req *adkmodel.LLMRequest) (*adkmodel.LLMResponse, error) { +// Tool names are resolved during model preparation so discovery stays lazy. +func makeMCPAppModelResultCallback(appToolNames func(agent.Context) map[string]bool) llmagent.BeforeModelCallback { + return func(ctx agent.Context, req *adkmodel.LLMRequest) (*adkmodel.LLMResponse, error) { + names := appToolNames(ctx) for _, content := range req.Contents { if content == nil { continue } for _, part := range content.Parts { - if part == nil || part.FunctionResponse == nil || !appToolNames[part.FunctionResponse.Name] { + if part == nil || part.FunctionResponse == nil || !names[part.FunctionResponse.Name] { continue } part.FunctionResponse.Response = compactMCPAppModelResponse(part.FunctionResponse.Response) diff --git a/go/adk/pkg/agent/mcp_apps_test.go b/go/adk/pkg/agent/mcp_apps_test.go index 65b8e2ffd0..23d1e0a242 100644 --- a/go/adk/pkg/agent/mcp_apps_test.go +++ b/go/adk/pkg/agent/mcp_apps_test.go @@ -1,13 +1,26 @@ package agent import ( + "context" + "net/http" + "net/http/httptest" + "sync/atomic" "testing" + "time" + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/require" + adkagent "google.golang.org/adk/v2/agent" adkmodel "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/tool" + "google.golang.org/adk/v2/tool/toolconfirmation" "google.golang.org/genai" + + "github.com/kagent-dev/kagent/go/adk/pkg/mcp" + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" ) -func TestMakeMCPAppModelResultCallbackReplacesRenderPayloadWithNotice(t *testing.T) { +func TestMCPAppModelResultReplacesRenderPayloadWithNotice(t *testing.T) { t.Parallel() req := &adkmodel.LLMRequest{ @@ -37,7 +50,9 @@ func TestMakeMCPAppModelResultCallbackReplacesRenderPayloadWithNotice(t *testing }}, } - callback := MakeMCPAppModelResultCallback(map[string]bool{"jenkins_monitor_build": true}) + callback := makeMCPAppModelResultCallback(func(adkagent.Context) map[string]bool { + return map[string]bool{"jenkins_monitor_build": true} + }) if _, err := callback(nil, req); err != nil { t.Fatalf("callback returned error: %v", err) } @@ -70,7 +85,7 @@ func TestMakeMCPAppModelResultCallbackReplacesRenderPayloadWithNotice(t *testing } } -func TestMakeMCPAppModelResultCallbackPreservesIsError(t *testing.T) { +func TestMCPAppModelResultPreservesIsError(t *testing.T) { t.Parallel() req := &adkmodel.LLMRequest{ @@ -91,7 +106,9 @@ func TestMakeMCPAppModelResultCallbackPreservesIsError(t *testing.T) { }}, } - callback := MakeMCPAppModelResultCallback(map[string]bool{"jenkins_monitor_build": true}) + callback := makeMCPAppModelResultCallback(func(adkagent.Context) map[string]bool { + return map[string]bool{"jenkins_monitor_build": true} + }) if _, err := callback(nil, req); err != nil { t.Fatalf("callback returned error: %v", err) } @@ -110,7 +127,7 @@ func TestMakeMCPAppModelResultCallbackPreservesIsError(t *testing.T) { } } -func TestMakeMCPAppModelResultCallbackLeavesNonAppToolsAlone(t *testing.T) { +func TestMCPAppModelResultLeavesNonAppToolsAlone(t *testing.T) { t.Parallel() original := map[string]any{ @@ -131,7 +148,9 @@ func TestMakeMCPAppModelResultCallbackLeavesNonAppToolsAlone(t *testing.T) { }}, } - callback := MakeMCPAppModelResultCallback(map[string]bool{"some_app_tool": true}) + callback := makeMCPAppModelResultCallback(func(adkagent.Context) map[string]bool { + return map[string]bool{"some_app_tool": true} + }) if _, err := callback(nil, req); err != nil { t.Fatalf("callback returned error: %v", err) } @@ -143,3 +162,128 @@ func TestMakeMCPAppModelResultCallbackLeavesNonAppToolsAlone(t *testing.T) { t.Fatalf("non-app tool response modified: %#v", got) } } + +func TestMCPAppCompactionPreservesCacheApprovalAndHistoricalResults(t *testing.T) { + // Given real MCP tools with approvals, when fresh and cached catalogs + // prepare model results, assert that app compaction preserves rich results + // and approval checks, including history after the app tool is removed. + var lists, calls atomic.Int32 + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "app-compaction", Version: "1"}, &mcpsdk.ServerOptions{ + SupportedProtocolVersions: []string{"2026-07-28"}, + SetCacheable: func(_ context.Context, _ mcpsdk.Request, c *mcpsdk.Cacheable) { c.TTLMs = 60000 }, + }) + for _, descriptor := range []*mcpsdk.Tool{ + {Name: "ordinary"}, + {Name: "panel", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://panel"}}}, + {Name: "app_only", Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}, + } { + mcpsdk.AddTool(server, descriptor, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + calls.Add(1) + return nil, map[string]any{"payload": "rich result"}, nil + }) + } + server.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + } + return next(ctx, method, req) + } + }) + srv := httptest.NewServer(mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, &mcpsdk.StreamableHTTPOptions{Stateless: true})) + t.Cleanup(func() { srv.CloseClientConnections(); srv.Close() }) + sets := mcp.CreateToolsets(t.Context(), []*configv1.ADKHTTPMCPServerConfig{{Params: &configv1.ADKHTTPConnectionParams{Url: srv.URL}, RequireApproval: true}}, nil, nil, false, nil) + require.Len(t, sets, 1) + ctx := &appCompactionContext{StrictContextMock: adkagent.NewStrictContextMock(t.Context()), confirmed: true} + callback := makeMCPAppModelResultCallback(func(ctx adkagent.Context) map[string]bool { + return mcp.MCPAppToolNamesFromToolsets(ctx, sets) + }) + var historical map[string]any + for range 2 { + listed, err := sets[0].Tools(ctx) + require.NoError(t, err) + require.Len(t, listed, 2) + req := &adkmodel.LLMRequest{} + original := make(map[string]map[string]any) + for _, candidate := range listed { + fn := candidate.(interface { + ProcessRequest(adkagent.Context, *adkmodel.LLMRequest) error + Run(adkagent.Context, any) (map[string]any, error) + }) + require.NoError(t, fn.ProcessRequest(ctx, req)) + response, err := fn.Run(ctx, map[string]any{}) + require.NoError(t, err) + original[candidate.Name()] = response + req.Contents = append(req.Contents, &genai.Content{Parts: []*genai.Part{{FunctionResponse: &genai.FunctionResponse{Name: candidate.Name(), Response: response}}}}) + } + require.NotContains(t, req.Tools, "app_only") + _, err = callback(ctx, req) + require.NoError(t, err) + for _, content := range req.Contents { + response := content.Parts[0].FunctionResponse + if response.Name == "panel" { + require.Equal(t, mcpAppRenderedNotice, response.Response["content"].([]any)[0].(map[string]any)["text"]) + } else { + require.Equal(t, original["ordinary"], response.Response) + } + } + historical = original["panel"] + require.Equal(t, map[string]any{"output": map[string]any{"payload": "rich result"}}, historical, "history retains the original rich result") + ctx.confirmed = false + _, err = req.Tools["panel"].(interface { + Run(adkagent.Context, any) (map[string]any, error) + }).Run(ctx, map[string]any{}) + require.ErrorIs(t, err, tool.ErrConfirmationRejected) + ctx.confirmed = true + } + require.EqualValues(t, 1, lists.Load(), "cached model preparation makes no extra tools/list call") + require.EqualValues(t, 4, calls.Load(), "rejected approvals must not execute MCP tools") + server.RemoveTools("panel") + require.Eventually(t, func() bool { + listed, err := sets[0].Tools(ctx) + require.NoError(t, err) + return len(listed) == 1 && listed[0].Name() == "ordinary" + }, 2*time.Second, 10*time.Millisecond) + req := &adkmodel.LLMRequest{Contents: []*genai.Content{{Parts: []*genai.Part{{FunctionResponse: &genai.FunctionResponse{Name: "panel", Response: historical}}}}}} + _, err := callback(ctx, req) + require.NoError(t, err) + require.Equal(t, mcpAppRenderedNotice, req.Contents[0].Parts[0].FunctionResponse.Response["content"].([]any)[0].(map[string]any)["text"]) + require.Equal(t, map[string]any{"output": map[string]any{"payload": "rich result"}}, historical) + + // Given the retained app name, when the tool returns as ordinary, assert + // that its new output reaches the model without the rendered notice. + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "panel"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + return nil, map[string]any{"payload": "ordinary result"}, nil + }) + var ordinary tool.Tool + require.Eventually(t, func() bool { + listed, err := sets[0].Tools(ctx) + require.NoError(t, err) + for _, candidate := range listed { + if candidate.Name() == "panel" { + ordinary = candidate + return !mcp.MCPAppToolNamesFromToolsets(ctx, sets)["panel"] + } + } + return false + }, time.Second, 10*time.Millisecond) + response, err := ordinary.(interface { + Run(adkagent.Context, any) (map[string]any, error) + }).Run(ctx, map[string]any{}) + require.NoError(t, err) + req = &adkmodel.LLMRequest{Contents: []*genai.Content{{Parts: []*genai.Part{{FunctionResponse: &genai.FunctionResponse{Name: "panel", Response: response}}}}}} + _, err = callback(ctx, req) + require.NoError(t, err) + require.Equal(t, map[string]any{"output": map[string]any{"payload": "ordinary result"}}, req.Contents[0].Parts[0].FunctionResponse.Response) +} + +type appCompactionContext struct { + adkagent.StrictContextMock + confirmed bool +} + +func (c *appCompactionContext) ToolConfirmation() *toolconfirmation.ToolConfirmation { + return &toolconfirmation.ToolConfirmation{Confirmed: c.confirmed} +} + +func (*appCompactionContext) FunctionCallID() string { return "app-compaction-call" } diff --git a/go/adk/pkg/mcp/discovery.go b/go/adk/pkg/mcp/discovery.go new file mode 100644 index 0000000000..d79692f483 --- /dev/null +++ b/go/adk/pkg/mcp/discovery.go @@ -0,0 +1,160 @@ +package mcp + +import ( + "context" + "fmt" + "sync" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + adkagent "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/tool" + + "github.com/kagent-dev/kagent/go/pkg/logging" +) + +type discoveryScopeKey struct{} + +// WithDiscoveryScope remembers unavailable MCP toolsets for one runner.Run or +// A2A execution, including all its model steps and sub-agent calls. Start a new +// scope for each execution, even when an approval resume reuses InvocationID. +// The scope contains no connections and needs no explicit cleanup. +func WithDiscoveryScope(ctx context.Context) context.Context { + return context.WithValue(ctx, discoveryScopeKey{}, &discoveryScope{omitted: make(map[*mcpAppToolset]bool)}) +} + +type discoveryScope struct { + mu sync.Mutex + omitted map[*mcpAppToolset]bool +} + +func (d *discoveryScope) isOmitted(ts *mcpAppToolset) bool { + d.mu.Lock() + defer d.mu.Unlock() + return d.omitted[ts] +} + +func (d *discoveryScope) omit(ts *mcpAppToolset) { + d.mu.Lock() + defer d.mu.Unlock() + d.omitted[ts] = true +} + +// Raw metadata survives SDK cache hits, including pages from incomplete lists. +// Published app names change only after successful discovery and are retained +// after removal for historical result compaction. Visibility is enforced on +// each raw tools/list response, never by these maps. +type mcpAppDiscovery struct { + mu sync.Mutex + listedAppNames map[string]bool + appToolNames map[string]bool +} + +type mcpAppToolset struct { + inner tool.Toolset + discovery *mcpAppDiscovery +} + +var _ tool.Toolset = (*mcpAppToolset)(nil) + +func (m *mcpAppToolset) Name() string { return m.inner.Name() } + +func (m *mcpAppToolset) Tools(ctx adkagent.ReadonlyContext) ([]tool.Tool, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + scope, _ := ctx.Value(discoveryScopeKey{}).(*discoveryScope) + if scope != nil && scope.isOmitted(m) { + return nil, nil + } + tools, err := m.inner.Tools(ctx) + if ctx.Err() != nil { + return nil, ctx.Err() + } + if err != nil { + if scope == nil { + return nil, err + } + scope.omit(m) + logging.FromContext(ctx).WarnContext(ctx, "failed to discover MCP tools; omitting toolset for this invocation", "error", err) + return nil, nil + } + m.discovery.mu.Lock() + for _, t := range tools { + if m.discovery.listedAppNames[t.Name()] { + m.discovery.appToolNames[t.Name()] = true + } else { + delete(m.discovery.appToolNames, t.Name()) + } + } + m.discovery.mu.Unlock() + if scope != nil && scope.isOmitted(m) { + return nil, nil + } + return tools, nil +} + +// MCPAppToolNamesFromToolsets returns successfully discovered MCP App-capable +// tool names, including names retained for historical model-result compaction. +// Call it during model preparation, after Tools has run. +func MCPAppToolNamesFromToolsets(ctx context.Context, toolsets []tool.Toolset) map[string]bool { + out := make(map[string]bool) + scope, _ := ctx.Value(discoveryScopeKey{}).(*discoveryScope) + for _, ts := range toolsets { + aware, ok := ts.(*mcpAppToolset) + if !ok || scope != nil && scope.isOmitted(aware) { + continue + } + aware.discovery.mu.Lock() + for name := range aware.discovery.appToolNames { + out[name] = true + } + aware.discovery.mu.Unlock() + } + return out +} + +func newMCPAppClient(configuredFilter map[string]bool, discovery *mcpAppDiscovery) *mcpsdk.Client { + client := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "kagent-adk"}, &mcpsdk.ClientOptions{ + Capabilities: mcpUIClientCapabilities(), + // New-protocol servers subscribe only when a handler is configured. + // The SDK itself invalidates its list cache on these notifications. + ToolListChangedHandler: func(context.Context, *mcpsdk.ToolListChangedRequest) {}, + }) + client.AddSendingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + if method != "tools/list" { + return next(ctx, method, req) + } + result, err := next(ctx, method, req) + if err != nil { + return nil, err + } + page, ok := result.(*mcpsdk.ListToolsResult) + if !ok || page == nil { + return nil, fmt.Errorf("invalid MCP tools/list result") + } + visible := make([]*mcpsdk.Tool, 0, len(page.Tools)) + discovery.mu.Lock() + for _, t := range page.Tools { + if t == nil || t.Name == "" || len(configuredFilter) > 0 && !configuredFilter[t.Name] { + continue + } + switch mcpToolKindOf(t.Meta) { + case mcpToolKindAppInternal: + continue + case mcpToolKindApp: + discovery.listedAppNames[t.Name] = true + default: + delete(discovery.listedAppNames, t.Name) + } + visible = append(visible, t) + } + discovery.mu.Unlock() + page.Tools = visible + // ListTools caches this already-filtered response after middleware + // returns. Preserve its TTL and all other response metadata. + return page, nil + } + }) + return client +} diff --git a/go/adk/pkg/mcp/discovery_test.go b/go/adk/pkg/mcp/discovery_test.go new file mode 100644 index 0000000000..35fbc5019a --- /dev/null +++ b/go/adk/pkg/mcp/discovery_test.go @@ -0,0 +1,675 @@ +package mcp + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/require" + "google.golang.org/adk/v2/tool" + + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" +) + +func TestCreateToolsetsDoesNotContactSlowServers(t *testing.T) { + t.Run("Given a slow MCP server, when toolsets are constructed, assert that no request delays startup", func(t *testing.T) { + var requests atomic.Int32 + server := newDiscoveryServer(t, 0, nil, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + select { + case <-time.After(500 * time.Millisecond): + case <-r.Context().Done(): + } + })) + sets := CreateToolsets(t.Context(), []*configv1.ADKHTTPMCPServerConfig{{Params: &configv1.ADKHTTPConnectionParams{Url: server.URL}}}, nil, nil, false, nil) + require.Len(t, sets, 1) + require.Zero(t, requests.Load()) + }) +} + +func TestFailedClassificationIsOmittedUntilNextRun(t *testing.T) { + // Given a transient discovery failure, when multiple model steps run, assert + // that the server stays omitted until a new execution scope retries it. + var unavailable atomic.Bool + unavailable.Store(true) + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + if unavailable.Load() { + return nil, errors.New("catalog unavailable") + } + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + ctx := testReadonlyContext{Context: WithDiscoveryScope(t.Context())} + listed, err := ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.EqualValues(t, 1, lists.Load(), "ADK must not discover after classification fails") + unavailable.Store(false) + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.EqualValues(t, 1, lists.Load()) + + // Same InvocationID, new run (including approval resumes). + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + require.Contains(t, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts}), "c_app") + // Another run's successful classification must not change the omitted run. + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.Empty(t, MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts}), "another run's classification must not tag omitted tools in this run") +} + +func TestClassificationIncludesEveryPageAndPreservesAllowlist(t *testing.T) { + // Given one tool per page and a configured allowlist, when discovery runs, + // assert that later app tools survive and app-only tools remain excluded. + var lists atomic.Int32 + server := newDiscoveryServer(t, 1, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + require.Contains(t, req.GetSession().(*mcpsdk.ServerSession).InitializeParams().Capabilities.Extensions, mcpUIExtensionName) + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, map[string]bool{"b_internal": true, "c_app": true}) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"c_app"}, discoveryNames(listed)) + require.EqualValues(t, 3, lists.Load(), "classification must use ADK's own paginated discovery") +} + +func TestPartialClassificationDoesNotPublishTools(t *testing.T) { + // Given a later page fails, when discovery runs, assert that no partial + // classification or unclassified ADK fallback becomes model-visible. + var pages atomic.Int32 + server := newDiscoveryServer(t, 1, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && pages.Add(1) == 2 { + return nil, errors.New("second page unavailable") + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + require.Empty(t, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + require.EqualValues(t, 2, pages.Load()) +} + +func TestClassificationWithoutScopeReturnsFailure(t *testing.T) { + // Given a standalone caller, when classification fails, assert that an error + // is returned without exposing ADK's unclassified catalog. + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + return nil, errors.New("catalog unavailable") + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + listed, err := ts.Tools(testReadonlyContext{Context: t.Context()}) + require.ErrorContains(t, err, "catalog unavailable") + require.Empty(t, listed) + require.EqualValues(t, 1, lists.Load()) +} + +func TestDiscoveryFailureAfterSuccessfulDiscoveryIsOmittedUntilNextRun(t *testing.T) { + // Given one successful discovery, when the next listing fails, assert that + // omission lasts for that execution and a new execution can recover. + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && lists.Add(1) == 2 { + return nil, errors.New("ADK catalog unavailable") + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + ctx := testReadonlyContext{Context: WithDiscoveryScope(t.Context())} + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.Empty(t, MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts})) + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.EqualValues(t, 2, lists.Load(), "later model steps must not retry the omitted toolset") + + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + require.EqualValues(t, 3, lists.Load()) + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed, "another execution's recovery must not undo omission") + require.EqualValues(t, 3, lists.Load()) + canceled, cancel := context.WithCancel(ctx) + cancel() + _, err = ts.Tools(testReadonlyContext{Context: canceled}) + require.ErrorIs(t, err, context.Canceled, "omission must not swallow cancellation") +} + +func TestDiscoveryFailureWithoutScopePropagatesAfterSuccess(t *testing.T) { + // Given a previously discovered standalone toolset, when its catalog fails, + // assert that discovery returns the error and the next call can retry. + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && lists.Add(1) == 2 { + return nil, errors.New("standalone catalog unavailable") + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + ctx := testReadonlyContext{Context: t.Context()} + listed, err := ts.Tools(ctx) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + listed, err = ts.Tools(ctx) + require.ErrorContains(t, err, "standalone catalog unavailable") + require.Empty(t, listed) + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + require.EqualValues(t, 3, lists.Load()) +} + +func TestEmptyLegacyCatalogRemainsClassified(t *testing.T) { + // Given an empty legacy catalog, when successive runs discover, assert that + // each model step lists once and a successful empty catalog remains valid. + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + result, err := next(ctx, method, req) + if method == "tools/list" && err == nil { + lists.Add(1) + result.(*mcpsdk.ListToolsResult).Tools = nil + } + return result, err + }, nil) + ts := discoveryToolset(t, server.URL, nil) + for range 2 { + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + } + require.EqualValues(t, 2, lists.Load(), "one discovery per legacy-protocol model step") +} + +func TestClassificationWaiterCanBeCanceled(t *testing.T) { + // Given discovery is already in flight, when another invocation is canceled, + // assert that it exits promptly and leaves the original discovery usable. + entered, release := make(chan struct{}), make(chan struct{}) + var once sync.Once + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + once.Do(func() { close(entered); <-release }) + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + done := make(chan error, 1) + go func() { + _, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + done <- err + }() + <-entered + ctx, cancel := context.WithCancel(WithDiscoveryScope(t.Context())) + waiter := &observedReadonlyContext{testReadonlyContext: testReadonlyContext{Context: ctx}, waiting: make(chan struct{})} + waiterDone := make(chan error, 1) + go func() { _, err := ts.Tools(waiter); waiterDone <- err }() + <-waiter.waiting + cancel() + require.ErrorIs(t, <-waiterDone, context.Canceled) + close(release) + require.NoError(t, <-done) +} + +func TestCanceledClassificationCanRecover(t *testing.T) { + // Given listing waits for cancellation, when the caller cancels discovery, + // assert that cancellation propagates and a later run can recover. + var canceled atomic.Bool + entered := make(chan struct{}) + release := make(chan struct{}) + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && !canceled.Load() { + close(entered) + <-release + } + return next(ctx, method, req) + }, nil) + ts := discoveryToolset(t, server.URL, nil) + ctx, cancel := context.WithCancel(WithDiscoveryScope(t.Context())) + done := make(chan error, 1) + go func() { _, err := ts.Tools(testReadonlyContext{Context: ctx}); done <- err }() + <-entered + canceled.Store(true) + cancel() + close(release) + require.ErrorIs(t, <-done, context.Canceled) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) +} + +func TestSDKCachePreservesFilteredToolsAndAppMetadata(t *testing.T) { + // Given a cacheable paginated catalog, when multiple executions prepare + // tools, assert that cached pages remain filtered and keep app metadata. + var lists, discoveries atomic.Int32 + server := newCacheableDiscoveryServer(1, 60000) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + } + if method == "server/discover" { + discoveries.Add(1) + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, map[string]bool{"b_internal": true, "c_app": true}) + for range 3 { + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"c_app"}, discoveryNames(listed)) + require.Equal(t, map[string]bool{"c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + } + require.EqualValues(t, 3, lists.Load(), "cached pages must not repeat discovery or bypass filtering") + require.EqualValues(t, 1, discoveries.Load(), "classification must share ADK's connection") +} + +func TestSDKCachePreservesEmptyFilteredCatalog(t *testing.T) { + // Given only an app-only tool is allowed, when cached tools are requested, + // assert that an empty filtered response remains empty without rediscovery. + var lists atomic.Int32 + server := newCacheableDiscoveryServer(0, 60000) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, map[string]bool{"b_internal": true}) + for range 2 { + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + } + require.EqualValues(t, 1, lists.Load()) +} + +func TestSDKCacheExpiryDiscoversNewTools(t *testing.T) { + // Given a cacheable catalog without notifications, when its TTL expires, + // assert that a newly advertised tool becomes visible through ADK discovery. + var updated atomic.Bool + var lists atomic.Int32 + server := newCacheableDiscoveryServer(0, 40) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + result, err := next(ctx, method, req) + if method == "tools/list" && err == nil { + lists.Add(1) + if updated.Load() { + page := result.(*mcpsdk.ListToolsResult) + page.Tools = append(page.Tools, &mcpsdk.Tool{Name: "d_new", InputSchema: map[string]any{"type": "object"}}) + } + } + return result, err + }, nil, true).URL + ts := discoveryToolset(t, url, nil) + listed, err := ts.Tools(testReadonlyContext{Context: t.Context()}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + updated.Store(true) + require.Eventually(t, func() bool { + listed, err := ts.Tools(testReadonlyContext{Context: t.Context()}) + require.NoError(t, err) + return slices.Contains(discoveryNames(listed), "d_new") + }, time.Second, 10*time.Millisecond) + require.GreaterOrEqual(t, lists.Load(), int32(2)) +} + +func TestSDKNotificationRefreshesVisibilityAndRetainsHistoricalAppNames(t *testing.T) { + // Given a long-lived cache, when the catalog changes, assert that the SDK + // refreshes visibility while retaining app names used to compact old results. + server := newCacheableDiscoveryServer(0, 3600000) + url := serveDiscoveryServer(t, server, nil, nil, true).URL + ts := discoveryToolset(t, url, nil) + listed, err := ts.Tools(testReadonlyContext{Context: t.Context()}) + require.NoError(t, err) + require.Equal(t, []string{"a_plain", "c_app"}, discoveryNames(listed)) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_plain", Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}, discoveryToolHandler) + server.RemoveTools("c_app") + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "d_new"}, discoveryToolHandler) + require.Eventually(t, func() bool { + listed, err := ts.Tools(testReadonlyContext{Context: t.Context()}) + require.NoError(t, err) + return slices.Equal(discoveryNames(listed), []string{"d_new"}) + }, 2*time.Second, 10*time.Millisecond) + require.Equal(t, map[string]bool{"c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) +} + +func TestPartialSDKCacheRetainsMetadataForLaterRecovery(t *testing.T) { + // Given an app page is cached before a later page fails, when a new + // execution recovers, assert that cached app metadata is still available. + var unavailable atomic.Bool + unavailable.Store(true) + var firstPages atomic.Int32 + server := newCacheableDiscoveryServer(1, 60000) + server.RemoveTools("a_plain") + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_app", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://first"}}}, discoveryToolHandler) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + if req.GetParams().(*mcpsdk.ListToolsParams).Cursor == "" { + firstPages.Add(1) + } else if unavailable.Load() { + return nil, errors.New("later page unavailable") + } + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, nil) + ctx := testReadonlyContext{Context: WithDiscoveryScope(t.Context())} + listed, err := ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed) + require.Empty(t, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + unavailable.Store(false) + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_app", "c_app"}, discoveryNames(listed)) + require.Equal(t, map[string]bool{"a_app": true, "c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + require.EqualValues(t, 1, firstPages.Load(), "the SDK must reuse its cached first page") + listed, err = ts.Tools(ctx) + require.NoError(t, err) + require.Empty(t, listed, "recovery cannot undo another execution's omission") +} + +func TestFailedDiscoveryDoesNotPublishStaleAppNamesAfterRecovery(t *testing.T) { + // Given an app page followed by a failure, when the server changes that + // tool to ordinary before recovery, assert that partial metadata is not published. + var unavailable atomic.Bool + unavailable.Store(true) + server := newCacheableDiscoveryServer(1, 0) + server.RemoveTools("a_plain") + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_app", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://first"}}}, discoveryToolHandler) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && req.GetParams().(*mcpsdk.ListToolsParams).Cursor != "" && unavailable.Load() { + return nil, errors.New("later page unavailable") + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, nil) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + require.Empty(t, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_app"}, discoveryToolHandler) + unavailable.Store(false) + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_app", "c_app"}, discoveryNames(listed)) + require.Equal(t, map[string]bool{"c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) +} + +func TestFailedDiscoveryDoesNotRetainRemovedUnpublishedAppNames(t *testing.T) { + // Given a published ordinary tool, when an app reclassification listing + // fails and recovery removes it, assert that its app name is never published. + var unavailable atomic.Bool + server := newCacheableDiscoveryServer(1, 0) + server.RemoveTools("a_plain") + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_changing"}, discoveryToolHandler) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" && req.GetParams().(*mcpsdk.ListToolsParams).Cursor != "" && unavailable.Load() { + return nil, errors.New("later page unavailable") + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, nil) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"a_changing", "c_app"}, discoveryNames(listed)) + require.Equal(t, map[string]bool{"c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts})) + + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "a_changing", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://changing"}}}, discoveryToolHandler) + unavailable.Store(true) + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + server.RemoveTools("a_changing") + unavailable.Store(false) + listed, err = ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"c_app"}, discoveryNames(listed)) + require.Equal(t, map[string]bool{"c_app": true}, MCPAppToolNamesFromToolsets(t.Context(), []tool.Toolset{ts}), "failed partial metadata must not compact old ordinary results") +} + +func TestSuccessfulDiscoveryClearsReclassifiedAppNames(t *testing.T) { + // Given a published app tool, when a complete listing changes it to an + // ordinary tool, assert that only removed app names remain for compaction. + for _, protocol := range []string{"2025-11-25", "2026-07-28"} { + t.Run(protocol, func(t *testing.T) { + var lists atomic.Int32 + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "reclassification"}, &mcpsdk.ServerOptions{ + SupportedProtocolVersions: []string{protocol}, + SetCacheable: func(_ context.Context, _ mcpsdk.Request, c *mcpsdk.Cacheable) { c.TTLMs = 60000 }, + }) + addDiscoveryTools(server) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "removed_app", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://removed"}}}, discoveryToolHandler) + url := serveDiscoveryServer(t, server, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + if method == "tools/list" { + lists.Add(1) + } + return next(ctx, method, req) + }, nil, true).URL + ts := discoveryToolset(t, url, nil) + ctx := testReadonlyContext{Context: WithDiscoveryScope(t.Context())} + _, err := ts.Tools(ctx) + require.NoError(t, err) + require.Equal(t, map[string]bool{"c_app": true, "removed_app": true}, MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts})) + + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "c_app"}, discoveryToolHandler) + server.RemoveTools("removed_app") + require.Eventually(t, func() bool { + listed, err := ts.Tools(ctx) + require.NoError(t, err) + return slices.Equal(discoveryNames(listed), []string{"a_plain", "c_app"}) && + !MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts})["c_app"] + }, time.Second, 10*time.Millisecond) + require.Equal(t, map[string]bool{"removed_app": true}, MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts})) + before := lists.Load() + _, err = ts.Tools(ctx) + require.NoError(t, err) + require.NotContains(t, MCPAppToolNamesFromToolsets(ctx, []tool.Toolset{ts}), "c_app") + if protocol == "2026-07-28" { + require.Equal(t, before, lists.Load(), "cached ordinary tools must not regain app classification") + } + }) + } +} + +func TestConcurrentDiscoveryKeepsEachResponseFiltered(t *testing.T) { + // Given concurrent uncached listings with different visibility, when both + // finish, assert that one response's classification never filters the other. + entered, release := make(chan struct{}), make(chan struct{}) + var lists atomic.Int32 + server := newDiscoveryServer(t, 0, func(ctx context.Context, method string, req mcpsdk.Request, next mcpsdk.MethodHandler) (mcpsdk.Result, error) { + result, err := next(ctx, method, req) + if method == "tools/list" && err == nil { + page := result.(*mcpsdk.ListToolsResult) + if lists.Add(1) == 1 { + close(entered) + <-release + page.Tools = []*mcpsdk.Tool{{Name: "same_name", InputSchema: map[string]any{"type": "object"}}} + } else { + page.Tools = []*mcpsdk.Tool{{Name: "same_name", InputSchema: map[string]any{"type": "object"}, Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}} + } + } + return result, err + }, nil) + ts := discoveryToolset(t, server.URL, nil) + var releaseOnce sync.Once + releaseFirst := func() { releaseOnce.Do(func() { close(release) }) } + t.Cleanup(releaseFirst) + done := make(chan []tool.Tool, 1) + go func() { + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + if err != nil { + t.Error(err) + } + done <- listed + }() + <-entered + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, listed) + releaseFirst() + require.Equal(t, []string{"same_name"}, discoveryNames(<-done)) +} + +func TestStdioDiscoveryStartsOnlyOneProcess(t *testing.T) { + // Given a stdio MCP server, when construction and discovery run, assert + // that startup starts no process and classification uses ADK's one process. + marker := filepath.Join(t.TempDir(), "process-starts") + ts, err := initializeToolSet(t.Context(), mcpServerParams{URL: "helper", ServerType: "stdio", Command: os.Args[0], Args: []string{"-test.run=^TestStdioDiscoveryHelper$"}, Env: map[string]string{"KAGENT_MCP_DISCOVERY_CHILD": marker}}, nil) + require.NoError(t, err) + t.Cleanup(func() { + if err := ts.inner.(io.Closer).Close(); err != nil && !strings.Contains(err.Error(), "signal: killed") { + t.Errorf("close stdio: %v", err) + } + }) + _, err = os.Stat(marker) + require.ErrorIs(t, err, os.ErrNotExist) + listed, err := ts.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Equal(t, []string{"stdio_tool"}, discoveryNames(listed)) + content, err := os.ReadFile(marker) + require.NoError(t, err) + require.Equal(t, "started\n", string(content)) +} + +func TestStdioDiscoveryHelper(t *testing.T) { + marker := os.Getenv("KAGENT_MCP_DISCOVERY_CHILD") + if marker == "" { + t.Skip("subprocess helper") + } + f, err := os.OpenFile(marker, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600) + if err != nil { + os.Exit(2) + } + if _, err := f.WriteString("started\n"); err != nil { + os.Exit(3) + } + if err := f.Close(); err != nil { + os.Exit(4) + } + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "stdio-discovery"}, &mcpsdk.ServerOptions{SupportedProtocolVersions: []string{"2025-11-25"}}) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "stdio_tool"}, discoveryToolHandler) + if err := server.Run(context.Background(), &mcpsdk.StdioTransport{}); err != nil { + os.Exit(5) + } + os.Exit(0) +} + +type discoveryMiddleware func(context.Context, string, mcpsdk.Request, mcpsdk.MethodHandler) (mcpsdk.Result, error) + +type observedReadonlyContext struct { + testReadonlyContext + waiting chan struct{} + once sync.Once +} + +func (o *observedReadonlyContext) Done() <-chan struct{} { + o.once.Do(func() { close(o.waiting) }) + return o.Context.Done() +} + +func newDiscoveryServer(t *testing.T, pageSize int, middleware discoveryMiddleware, intercept http.Handler) *httptest.Server { + t.Helper() + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "discovery-test", Version: "1"}, &mcpsdk.ServerOptions{PageSize: pageSize, SupportedProtocolVersions: []string{"2025-11-25"}}) + addDiscoveryTools(server) + return serveDiscoveryServer(t, server, middleware, intercept, false) +} + +func newCacheableDiscoveryServer(pageSize, ttl int) *mcpsdk.Server { + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "cached-discovery-test", Version: "1"}, &mcpsdk.ServerOptions{ + PageSize: pageSize, SupportedProtocolVersions: []string{"2026-07-28"}, + SetCacheable: func(_ context.Context, _ mcpsdk.Request, c *mcpsdk.Cacheable) { c.TTLMs = ttl }, + }) + addDiscoveryTools(server) + return server +} + +func addDiscoveryTools(server *mcpsdk.Server) { + for _, descriptor := range []*mcpsdk.Tool{ + {Name: "a_plain"}, + {Name: "b_internal", Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}, + {Name: "c_app", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://discovery"}}}, + } { + mcpsdk.AddTool(server, descriptor, discoveryToolHandler) + } +} + +func discoveryToolHandler(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + return nil, map[string]any{"result": "ok"}, nil +} + +func serveDiscoveryServer(t *testing.T, server *mcpsdk.Server, middleware discoveryMiddleware, intercept http.Handler, stateless bool) *httptest.Server { + t.Helper() + if middleware != nil { + server.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + return middleware(ctx, method, req, next) + } + }) + } + handler := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, &mcpsdk.StreamableHTTPOptions{Stateless: stateless}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if intercept != nil { + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + r.Body = io.NopCloser(bytes.NewReader(body)) + intercept.ServeHTTP(w, r) + } + handler.ServeHTTP(w, r) + })) + t.Cleanup(func() { srv.CloseClientConnections(); srv.Close() }) + return srv +} + +func discoveryToolset(t *testing.T, url string, filter map[string]bool) *mcpAppToolset { + t.Helper() + ts, err := initializeToolSet(t.Context(), mcpServerParams{URL: url, ServerType: "http"}, filter) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, ts.inner.(io.Closer).Close()) }) + return ts +} + +func discoveryNames(tools []tool.Tool) []string { + var names []string + for _, t := range tools { + names = append(names, t.Name()) + } + slices.Sort(names) + return names +} diff --git a/go/adk/pkg/mcp/mcp_ui.go b/go/adk/pkg/mcp/mcp_ui.go index 404a9831e0..515ef3bd2f 100644 --- a/go/adk/pkg/mcp/mcp_ui.go +++ b/go/adk/pkg/mcp/mcp_ui.go @@ -17,12 +17,7 @@ package mcp // in one place and is easy to track as the extension evolves. import ( - "context" - "fmt" - mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" - adkagent "google.golang.org/adk/v2/agent" - "google.golang.org/adk/v2/tool" ) const ( @@ -43,44 +38,6 @@ func mcpUIClientCapabilities() *mcpsdk.ClientCapabilities { return caps } -// mcpAppToolset wraps an MCP toolset and records which model-visible tools -// render as MCP App widgets. Classification happens during ListTools inside -// agentVisibleToolFilter because mcpsdk.Tool.Meta is not preserved on the ADK -// tool.Tool values the toolset exposes later. -type mcpAppToolset struct { - inner tool.Toolset - // appToolNames is this server's set of model-visible tool names whose - // results render as interactive MCP App (UI) widgets (the tool declares a - // `_meta.ui.resourceUri`). Used as a set, so only key presence is meaningful. - appToolNames map[string]bool -} - -func (m *mcpAppToolset) Name() string { - return m.inner.Name() -} - -func (m *mcpAppToolset) Tools(ctx adkagent.ReadonlyContext) ([]tool.Tool, error) { - return m.inner.Tools(ctx) -} - -// MCPAppToolNamesFromToolsets returns the union of MCP App-capable tool names -// recorded on toolsets built by CreateToolsets. The agent attaches the -// model-result compaction callback (see agent.MakeMCPAppModelResultCallback) -// only to these tools. -func MCPAppToolNamesFromToolsets(toolsets []tool.Toolset) map[string]bool { - out := make(map[string]bool) - for _, ts := range toolsets { - aware, ok := ts.(*mcpAppToolset) - if !ok { - continue - } - for name := range aware.appToolNames { - out[name] = true - } - } - return out -} - // mcpToolKind classifies how an MCP tool is surfaced, following the MCP Apps // extension. It is derived from the tool's `_meta.ui` block: presence of a // `resourceUri` means the tool renders an app, and the `visibility` field @@ -191,52 +148,3 @@ func normalizeVisibility(value any) []string { } return nil } - -// agentVisibleToolFilter lists tools from the MCP server, filters out app-only -// tools and any not in the configured allow-list, and returns a predicate the -// toolset can apply plus the MCP App-capable tool names discovered on this -// server. -// -// Classification must happen here because MCP Apps metadata lives on -// mcpsdk.Tool.Meta, which ADK mcptoolset drops when converting to tool.Tool. -func agentVisibleToolFilter(ctx context.Context, params mcpServerParams, configuredFilter map[string]bool) (tool.Predicate, map[string]bool, error) { - mcpTransport, err := createTransport(ctx, params) - if err != nil { - return nil, nil, fmt.Errorf("failed to create transport for %s: %w", params.URL, err) - } - - client := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "kagent-adk"}, &mcpsdk.ClientOptions{Capabilities: mcpUIClientCapabilities()}) - session, err := client.Connect(ctx, mcpTransport, nil) - if err != nil { - return nil, nil, fmt.Errorf("failed to connect MCP client for %s: %w", params.URL, err) - } - defer session.Close() - - result, err := session.ListTools(ctx, &mcpsdk.ListToolsParams{}) - if err != nil { - return nil, nil, fmt.Errorf("failed to list MCP tools for %s: %w", params.URL, err) - } - - allowedTools := make([]string, 0, len(result.Tools)) - appToolNames := make(map[string]bool) - for _, t := range result.Tools { - if t == nil || t.Name == "" { - continue - } - if len(configuredFilter) > 0 && !configuredFilter[t.Name] { - continue - } - switch mcpToolKindOf(t.Meta) { - case mcpToolKindAppInternal: - // Hidden from the model; only the rendered MCP App calls it. - continue - case mcpToolKindApp: - allowedTools = append(allowedTools, t.Name) - appToolNames[t.Name] = true - default: // mcpToolKindModel - allowedTools = append(allowedTools, t.Name) - } - } - - return tool.StringPredicate(allowedTools), appToolNames, nil -} diff --git a/go/adk/pkg/mcp/passthrough_test.go b/go/adk/pkg/mcp/passthrough_test.go index 6cb1199ba6..b945a01c32 100644 --- a/go/adk/pkg/mcp/passthrough_test.go +++ b/go/adk/pkg/mcp/passthrough_test.go @@ -87,11 +87,13 @@ func TestTokenPassthroughToolsetDiscoversWithCallerToken(t *testing.T) { require.Empty(t, seen, "startup must not dial without a caller token") mu.Unlock() - _, err = toolset.Tools(testReadonlyContext{Context: t.Context()}) - require.Error(t, err) + // Given no caller token, when discovery runs, assert that no tools are exposed. + tools, err := toolset.Tools(testReadonlyContext{Context: WithDiscoveryScope(t.Context())}) + require.NoError(t, err) + require.Empty(t, tools) ctx := context.WithValue(t.Context(), models.BearerTokenKey, "alice") - tools, err := toolset.Tools(testReadonlyContext{Context: ctx}) + tools, err = toolset.Tools(testReadonlyContext{Context: WithDiscoveryScope(ctx)}) require.NoError(t, err) require.Len(t, tools, 1) require.Equal(t, "whoami", tools[0].Name()) diff --git a/go/adk/pkg/mcp/registry.go b/go/adk/pkg/mcp/registry.go index 600fcc401b..e7fd478941 100644 --- a/go/adk/pkg/mcp/registry.go +++ b/go/adk/pkg/mcp/registry.go @@ -371,31 +371,18 @@ func (rt *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, erro return rt.base.RoundTrip(req) } -// initializeToolSet fetches tools from an MCP server using Google ADK's -// mcptoolset and wraps the result with any MCP App-capable tool names found -// during classification. +// initializeToolSet prepares a lazy MCP toolset. Neither classification nor +// ADK discovery may connect before the runtime reports ready. func initializeToolSet(ctx context.Context, params mcpServerParams, toolFilter map[string]bool) (*mcpAppToolset, error) { mcpTransport, err := createTransport(ctx, params) if err != nil { return nil, fmt.Errorf("failed to create transport for %s: %w", params.URL, err) } - var toolPredicate tool.Predicate - var appToolNames map[string]bool - if params.TokenPassthrough { - // Startup has no caller token, so discover lazily on the first invocation. - toolPredicate = configuredToolPredicate(toolFilter) - } else { - toolPredicate, appToolNames, err = agentVisibleToolFilter(ctx, params, toolFilter) - if err != nil { - logging.FromContext(ctx).ErrorContext(ctx, "failed to classify MCP tools; falling back to lazy tool discovery", "error", err, "url", params.URL) - appToolNames = nil - toolPredicate = configuredToolPredicate(toolFilter) - } - } - + discovery := &mcpAppDiscovery{listedAppNames: make(map[string]bool), appToolNames: make(map[string]bool)} cfg := mcptoolset.Config{ Transport: mcpTransport, + Client: newMCPAppClient(toolFilter, discovery), } toolset, err := mcptoolset.New(cfg) @@ -403,24 +390,5 @@ func initializeToolSet(ctx context.Context, params mcpServerParams, toolFilter m return nil, fmt.Errorf("failed to create MCP toolset for %s: %w", params.URL, err) } - visibleTools := tool.Toolset(toolset) - if toolPredicate != nil { - visibleTools = tool.FilterToolset(toolset, toolPredicate) - } - return &mcpAppToolset{inner: visibleTools, appToolNames: appToolNames}, nil -} - -// configuredToolPredicate selects the configured tool names, or every tool -// when none are configured. -func configuredToolPredicate(toolFilter map[string]bool) tool.Predicate { - if len(toolFilter) == 0 { - return nil - } - allowedTools := make([]string, 0, len(toolFilter)) - for name, enabled := range toolFilter { - if enabled { - allowedTools = append(allowedTools, name) - } - } - return tool.StringPredicate(allowedTools) + return &mcpAppToolset{inner: toolset, discovery: discovery}, nil } diff --git a/go/adk/pkg/mcp/registry_test.go b/go/adk/pkg/mcp/registry_test.go index 2620816588..fc04f3aa7c 100644 --- a/go/adk/pkg/mcp/registry_test.go +++ b/go/adk/pkg/mcp/registry_test.go @@ -218,12 +218,12 @@ func TestMCPAppToolNamesFromToolsets(t *testing.T) { inner := &stubToolset{name: "mcp-server"} toolsets := []tool.Toolset{ - &mcpAppToolset{inner: inner, appToolNames: map[string]bool{"show_board": true}}, - &mcpAppToolset{inner: inner, appToolNames: map[string]bool{"refresh": true}}, + &mcpAppToolset{inner: inner, discovery: &mcpAppDiscovery{appToolNames: map[string]bool{"show_board": true}}}, + &mcpAppToolset{inner: inner, discovery: &mcpAppDiscovery{appToolNames: map[string]bool{"refresh": true}}}, inner, } - got := MCPAppToolNamesFromToolsets(toolsets) + got := MCPAppToolNamesFromToolsets(t.Context(), toolsets) if len(got) != 2 || !got["show_board"] || !got["refresh"] { t.Fatalf("MCPAppToolNamesFromToolsets() = %#v, want show_board and refresh", got) } diff --git a/go/core/test/e2e/mcp_discovery_test.go b/go/core/test/e2e/mcp_discovery_test.go new file mode 100644 index 0000000000..bd4369a2a3 --- /dev/null +++ b/go/core/test/e2e/mcp_discovery_test.go @@ -0,0 +1,440 @@ +package e2e_test + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "net/http/httptest" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv" + "github.com/google/uuid" + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + kagenta2a "github.com/kagent-dev/kagent/go/api/a2a" + "github.com/kagent-dev/kagent/go/api/v1alpha3" +) + +func TestMCPClassificationFailurePreservesHealthyToolsAndRecovers(t *testing.T) { + // Given healthy and unavailable catalogs, when an execution continues and + // another execution resumes, assert that failed tools stay hidden until recovery. + target := interactionTarget(t) + var unavailable atomic.Bool + unavailable.Store(true) + var runtimeLists, healthyCalls, recoveredCalls, hiddenCalls atomic.Int32 + healthy := newMCPDiscoveryServer() + mcpsdk.AddTool(healthy, &mcpsdk.Tool{Name: "healthy_tool"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + healthyCalls.Add(1) + // Recovery between model steps must not undo this execution's omission. + unavailable.Store(false) + return nil, map[string]any{"result": "healthy"}, nil + }) + recovering := newMCPDiscoveryServer() + mcpsdk.AddTool(recovering, &mcpsdk.Tool{Name: "recovered_tool"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + recoveredCalls.Add(1) + return nil, map[string]any{"result": "recovered"}, nil + }) + mcpsdk.AddTool(recovering, &mcpsdk.Tool{Name: "app_only_tool", Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + hiddenCalls.Add(1) + return nil, map[string]any{"result": "hidden"}, nil + }) + recovering.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + if method == "tools/list" && req.GetSession().(*mcpsdk.ServerSession).InitializeParams().ClientInfo.Name == "kagent-adk" { + runtimeLists.Add(1) + if unavailable.Load() { + return nil, errors.New("runtime catalog temporarily unavailable") + } + } + return next(ctx, method, req) + } + }) + var modelRequests atomic.Int32 + modelURL := startMCPDiscoveryModel(t, func(request discoveryModelRequest) (string, string, error) { + modelRequests.Add(1) + recovered := strings.Contains(request.lastUser(), "recovered") + names := request.toolNames() + if !slices.Contains(names, "healthy_tool") || slices.Contains(names, "app_only_tool") || slices.Contains(names, "recovered_tool") != recovered { + return "", "", fmt.Errorf("unexpected tool visibility for recovered=%t: %v", recovered, names) + } + toolName, result := "healthy_tool", "HEALTHY_DISCOVERY_OK" + if recovered { + toolName, result = "recovered_tool", "RECOVERED_DISCOVERY_OK" + } + if _, ok := request.toolResults(true)[toolName]; !ok { + return toolName, "", nil + } + return "", result, nil + }) + template := createMCPDiscoveryTemplate(t, modelURL, startMCPDiscoveryServer(t, healthy), startMCPDiscoveryServer(t, recovering)) + require.Zero(t, runtimeLists.Load(), "runtime startup must not classify remote tools") + fixture := newInteractionFixtureForTemplate(t, testHarness{name: "kagent"}, target, template) + _, _, first := fixture.send(t, "Use the healthy tool.") + require.Equal(t, a2atype.TaskStateCompleted, first.Status.State, taskText(first)) + require.Contains(t, taskText(first), "HEALTHY_DISCOVERY_OK") + require.EqualValues(t, 1, healthyCalls.Load()) + require.Zero(t, recoveredCalls.Load()) + require.EqualValues(t, 1, runtimeLists.Load(), "an omitted catalog must not be rediscovered in later model steps") + + assertActorSuspended(t, fixture) + _, _, second := fixture.send(t, "Use the recovered tool.") + require.Equal(t, a2atype.TaskStateCompleted, second.Status.State, taskText(second)) + require.Contains(t, taskText(second), "RECOVERED_DISCOVERY_OK") + require.EqualValues(t, 1, recoveredCalls.Load()) + require.EqualValues(t, 1, healthyCalls.Load()) + require.Zero(t, hiddenCalls.Load()) + require.Greater(t, runtimeLists.Load(), int32(1)) + require.EqualValues(t, 4, modelRequests.Load(), "visibility must be checked before and after both tool calls") +} + +func TestMCPAppCompactionPreservesHistoryAcrossResume(t *testing.T) { + // Given an app tool and an ordinary tool, when a restored execution prepares + // model history, assert that only app results are compacted and public history stays rich. + target := interactionTarget(t) + renderPayload := map[string]any{"panel": map[string]any{"title": "Discovery test panel", "rows": []any{"RICH_MCP_RENDER_PAYLOAD", strings.Repeat("render data ", 128)}}} + ordinaryPayload := map[string]any{"note": "ORDINARY_MCP_RESULT"} + var appCalls, ordinaryCalls, hiddenCalls, runtimeLists atomic.Int32 + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "cached-app-discovery", Version: "1"}, &mcpsdk.ServerOptions{ + SupportedProtocolVersions: []string{"2026-07-28"}, + SetCacheable: func(_ context.Context, _ mcpsdk.Request, c *mcpsdk.Cacheable) { c.TTLMs = 60000 }, + }) + server.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + if method == "tools/list" && req.(*mcpsdk.ListToolsRequest).ClientInfo().Name == "kagent-adk" { + runtimeLists.Add(1) + } + return next(ctx, method, req) + } + }) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "show_panel", Meta: mcpsdk.Meta{"ui": map[string]any{"resourceUri": "ui://discovery/panel"}}}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + appCalls.Add(1) + return nil, renderPayload, nil + }) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "ordinary_tool"}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + ordinaryCalls.Add(1) + return nil, ordinaryPayload, nil + }) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "app_only_tool", Meta: mcpsdk.Meta{"ui": map[string]any{"visibility": []string{"app"}}}}, func(context.Context, *mcpsdk.CallToolRequest, map[string]any) (*mcpsdk.CallToolResult, map[string]any, error) { + hiddenCalls.Add(1) + return nil, map[string]any{"result": "hidden"}, nil + }) + var restoredRequests atomic.Int32 + modelURL := startMCPDiscoveryModel(t, func(request discoveryModelRequest) (string, string, error) { + names := request.toolNames() + if !slices.Contains(names, "show_panel") || !slices.Contains(names, "ordinary_tool") || slices.Contains(names, "app_only_tool") { + return "", "", fmt.Errorf("unexpected model tools: %v", names) + } + results := request.toolResults(false) + appResult, rendered := results["show_panel"] + if !rendered { + return "show_panel", "", nil + } + if appResult != "The interactive UI for this tool has been rendered to the user and now updates live inside the app. Treat this as complete and do not call this tool again unless the user explicitly asks for it." { + return "", "", fmt.Errorf("app result was not compacted: %s", appResult) + } + ordinaryResult, called := results["ordinary_tool"] + if !called { + return "ordinary_tool", "", nil + } + var output struct { + Output map[string]any `json:"output"` + } + if err := json.Unmarshal([]byte(ordinaryResult), &output); err != nil || len(output.Output) != len(ordinaryPayload) || output.Output["note"] != ordinaryPayload["note"] { + return "", "", fmt.Errorf("ordinary result was modified: %s", ordinaryResult) + } + if strings.Contains(request.lastUser(), "without calling tools") { + restoredRequests.Add(1) + return "", "MCP_APP_HISTORY_RESTORED", nil + } + return "", "MCP_APP_RENDERED", nil + }) + handler := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, &mcpsdk.StreamableHTTPOptions{Stateless: true}) + serverURL := reachableServerURL(t, startMCPDiscoveryHTTP(t, handler), "/mcp") + template := createMCPDiscoveryTemplate(t, modelURL, serverURL) + fixture := newInteractionFixtureForTemplate(t, testHarness{name: "kagent"}, target, template) + _, _, first := fixture.send(t, "Render the panel and call the ordinary tool.") + require.Equal(t, a2atype.TaskStateCompleted, first.Status.State, taskText(first)) + require.Contains(t, taskText(first), "MCP_APP_RENDERED") + require.EqualValues(t, 1, runtimeLists.Load(), "SDK caching must preserve filtering and app metadata across model steps") + assertPersistedMCPResult(t, fixture, first.ID, "show_panel", renderPayload) + assertPersistedMCPResult(t, fixture, first.ID, "ordinary_tool", ordinaryPayload) + + assertActorSuspended(t, fixture) + _, _, second := fixture.send(t, "Summarize the previous results without calling tools.") + require.Equal(t, a2atype.TaskStateCompleted, second.Status.State, taskText(second)) + require.Contains(t, taskText(second), "MCP_APP_HISTORY_RESTORED") + require.EqualValues(t, 1, restoredRequests.Load(), "the first restored model request must already contain compacted history") + require.EqualValues(t, 1, appCalls.Load(), "restoring history must not invoke the app again") + require.EqualValues(t, 1, ordinaryCalls.Load()) + require.Zero(t, hiddenCalls.Load()) + require.EqualValues(t, 2, runtimeLists.Load(), "a restored process must discover once after readiness") + assertPersistedMCPResult(t, fixture, first.ID, "show_panel", renderPayload) + assertPersistedMCPResult(t, fixture, first.ID, "ordinary_tool", ordinaryPayload) +} + +func TestSlowMCPDiscoveryDoesNotBlockResume(t *testing.T) { + for _, slow := range []string{"tools/list", "GET"} { + t.Run("Given slow "+strings.ReplaceAll(slow, "/", " ")+", when an agent becomes ready and resumes repeatedly, assert that discovery waits until invocation", func(t *testing.T) { + target := interactionTarget(t) + harness := testHarness{name: "kagent", runtimeLabel: "kagent"} + url, runtimeDiscovery := startSlowMCPServer(t, slow) + template, _ := createMCPInteractionTemplate(t, harness, url, false, nil) + require.Zero(t, runtimeDiscovery.Load(), "golden runtime readiness must not depend on MCP discovery") + fixture := newInteractionFixtureForTemplate(t, harness, target, template) + for turn := range 2 { + if turn > 0 { + assertActorSuspended(t, fixture) + } + // Use one SendMessage call: retries would hide a parking timeout. + _, request := newMessageRequest(t, "Add 3 and 5 using the configured MCP server.") + request.Message.ContextId = fixture.sessionID + request.Tenant = fixture.tenant + response, err := fixture.client.SendMessage(fixture.ctx, request) + require.NoError(t, err, "slow MCP discovery must not make the router reject a restore") + result, err := pbconv.FromProtoSendMessageResponse(response) + require.NoError(t, err) + task, ok := result.(*a2atype.Task) + require.True(t, ok, "A2A response = %T, want Task", result) + require.Equal(t, a2atype.TaskStateCompleted, task.Status.State, taskText(task)) + require.Contains(t, taskText(task), "result is 8") + } + require.Positive(t, runtimeDiscovery.Load(), "successful invocations must actually discover the slow server") + }) + } +} + +func startSlowMCPServer(t *testing.T, slow string) (string, *atomic.Int32) { + t.Helper() + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "slow-startup", Version: "1"}, &mcpsdk.ServerOptions{SupportedProtocolVersions: []string{"2025-11-25"}}) + mcpsdk.AddTool(server, &mcpsdk.Tool{Name: "add_numbers", Description: "Add two numbers"}, func(_ context.Context, _ *mcpsdk.CallToolRequest, input struct { + A int `json:"a"` + B int `json:"b"` + }) (*mcpsdk.CallToolResult, map[string]any, error) { + return &mcpsdk.CallToolResult{Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: "The result is 8."}}}, map[string]any{"result": input.A + input.B}, nil + }) + var clients sync.Map + runtimeDiscovery := new(atomic.Int32) + server.AddReceivingMiddleware(func(next mcpsdk.MethodHandler) mcpsdk.MethodHandler { + return func(ctx context.Context, method string, req mcpsdk.Request) (mcpsdk.Result, error) { + session := req.GetSession().(*mcpsdk.ServerSession) + if method == "initialize" { + clients.Store(session.ID(), req.GetParams().(*mcpsdk.InitializeParams).ClientInfo.Name) + } + if method == "tools/list" && session.InitializeParams().ClientInfo.Name == "kagent-adk" { + runtimeDiscovery.Add(1) + if slow == method { + if err := waitForSlowDiscovery(ctx); err != nil { + return nil, err + } + } + } + return next(ctx, method, req) + } + }) + handler := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, nil) + baseURL := startMCPDiscoveryHTTP(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + if client, ok := clients.Load(r.Header.Get("Mcp-Session-Id")); ok && client == "kagent-adk" { + runtimeDiscovery.Add(1) + if slow == r.Method { + if err := waitForSlowDiscovery(r.Context()); err != nil { + return + } + } + } + } + handler.ServeHTTP(w, r) + })) + return reachableServerURL(t, baseURL, "/mcp"), runtimeDiscovery +} + +func waitForSlowDiscovery(ctx context.Context) error { + // Longer than the E2E router's five-second parking budget, while fitting + // within the configured MCP request timeout. Workload requests may take longer. + timer := time.NewTimer(8 * time.Second) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func newMCPDiscoveryServer() *mcpsdk.Server { + return mcpsdk.NewServer(&mcpsdk.Implementation{Name: "discovery-e2e", Version: "1"}, &mcpsdk.ServerOptions{SupportedProtocolVersions: []string{"2025-11-25"}}) +} + +func startMCPDiscoveryServer(t *testing.T, server *mcpsdk.Server) string { + t.Helper() + handler := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return server }, nil) + return reachableServerURL(t, startMCPDiscoveryHTTP(t, handler), "/mcp") +} + +func startMCPDiscoveryHTTP(t *testing.T, handler http.Handler) string { + t.Helper() + server := httptest.NewUnstartedServer(handler) + require.NoError(t, server.Listener.Close()) + var err error + server.Listener, err = net.Listen("tcp", "0.0.0.0:0") + require.NoError(t, err) + server.Start() + t.Cleanup(func() { server.CloseClientConnections(); server.Close() }) + return server.URL +} + +func createMCPDiscoveryTemplate(t *testing.T, modelURL string, serverURLs ...string) string { + t.Helper() + kube := interactionKubeClient(t) + harness := testHarness{name: "kagent", runtimeLabel: "kagent"} + model := &v1alpha3.ModelConfig{ + ObjectMeta: metav1.ObjectMeta{GenerateName: "mcp-discovery-", Namespace: "kagent"}, + Spec: v1alpha3.ModelConfigSpec{Provider: v1alpha3.ModelProviderOpenAI, Model: "gpt-4.1-mini", Stream: new(false), OpenAI: &v1alpha3.OpenAIConfig{BaseURL: modelURL}}, + } + require.NoError(t, kube.Create(t.Context(), model)) + t.Cleanup(func() { + if err := kube.Delete(context.Background(), model); err != nil && !apierrors.IsNotFound(err) { + t.Errorf("delete discovery model: %v", err) + } + }) + template := &v1alpha3.AgentTemplate{ + ObjectMeta: metav1.ObjectMeta{GenerateName: "mcp-discovery-", Namespace: "kagent", Labels: harness.labels()}, + Spec: v1alpha3.AgentTemplateSpec{Description: "MCP discovery E2E fixture", SystemPrompt: "Use the configured MCP tools.", ModelConfig: &v1alpha3.ObjectReference{Name: model.Name}}, + } + for _, serverURL := range serverURLs { + server := &v1alpha3.RemoteMCPServer{ + ObjectMeta: metav1.ObjectMeta{GenerateName: "mcp-discovery-", Namespace: "kagent"}, + Spec: v1alpha3.RemoteMCPServerSpec{Description: "MCP discovery E2E fixture", Protocol: v1alpha3.RemoteMCPServerProtocolStreamableHttp, URL: serverURL}, + } + require.NoError(t, kube.Create(t.Context(), server)) + t.Cleanup(func() { + if err := kube.Delete(context.Background(), server); err != nil && !apierrors.IsNotFound(err) { + t.Errorf("delete MCP server: %v", err) + } + }) + template.Spec.Tools = append(template.Spec.Tools, v1alpha3.ToolBinding{MCP: &v1alpha3.MCPToolBinding{Server: corev1.TypedObjectReference{Kind: "RemoteMCPServer", Name: server.Name}}}) + } + createAndWaitInteractionTemplateForHarness(t, kube, template, harness.name) + return template.Name +} + +func assertPersistedMCPResult(t *testing.T, fixture *interactionFixture, taskID a2atype.TaskID, toolName string, payload map[string]any) { + t.Helper() + request, err := pbconv.ToProtoGetTaskRequest(&a2atype.GetTaskRequest{Tenant: fixture.tenant, ID: taskID}) + require.NoError(t, err) + response, err := fixture.client.GetTask(fixture.ctx, request) + require.NoError(t, err) + task, err := pbconv.FromProtoTask(response) + require.NoError(t, err) + var responses []any + // The runtime publishes tool results as A2A artifacts. These are the + // persisted chat timeline payloads; Task.History contains input messages. + for _, artifact := range task.Artifacts { + for _, part := range artifact.Parts { + activity, ok := kagenta2a.ParseToolActivity(part) + if ok && activity.Kind == kagenta2a.ToolResultKind && activity.Name == toolName { + responses = append(responses, activity.Response) + } + } + } + require.Len(t, responses, 1, "persisted artifacts must contain the original %s result", toolName) + require.Equal(t, map[string]any{"output": payload}, responses[0]) +} + +type discoveryModelRequest struct { + Messages []struct { + Role string `json:"role"` + Content string `json:"content"` + ToolCallID string `json:"tool_call_id"` + ToolCalls []struct { + ID string `json:"id"` + Function struct { + Name string `json:"name"` + } `json:"function"` + } `json:"tool_calls"` + } `json:"messages"` + Tools []struct { + Function struct { + Name string `json:"name"` + } `json:"function"` + } `json:"tools"` +} + +func (r discoveryModelRequest) lastUser() string { + for _, message := range slices.Backward(r.Messages) { + if message.Role == "user" { + return message.Content + } + } + return "" +} + +func (r discoveryModelRequest) toolNames() []string { + var names []string + for _, tool := range r.Tools { + names = append(names, tool.Function.Name) + } + return names +} + +func (r discoveryModelRequest) toolResults(currentTurnOnly bool) map[string]string { + names := make(map[string]string) + start := 0 + for i, message := range r.Messages { + for _, call := range message.ToolCalls { + names[call.ID] = call.Function.Name + } + if currentTurnOnly && message.Role == "user" { + start = i + } + } + results := make(map[string]string) + for _, message := range r.Messages[start:] { + if message.Role == "tool" { + results[names[message.ToolCallID]] = message.Content + } + } + return results +} + +func startMCPDiscoveryModel(t *testing.T, respond func(discoveryModelRequest) (toolName, text string, err error)) string { + t.Helper() + baseURL := startMCPDiscoveryHTTP(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request discoveryModelRequest + if err := json.NewDecoder(r.Body).Decode(&request); err != nil { + t.Logf("model mock could not decode request: %v", err) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + toolName, text, err := respond(request) + if err != nil { + t.Logf("model mock rejected request: %v", err) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + message := map[string]any{"role": "assistant", "content": text} + finishReason := "stop" + if toolName != "" { + message["tool_calls"] = []any{map[string]any{"id": uuid.NewString(), "type": "function", "function": map[string]any{"name": toolName, "arguments": "{}"}}} + finishReason = "tool_calls" + } + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(map[string]any{"id": uuid.NewString(), "object": "chat.completion", "model": "test-model", "choices": []any{map[string]any{"index": 0, "message": message, "finish_reason": finishReason}}}); err != nil { + t.Errorf("write model response: %v", err) + } + })) + return reachableModelURL(t, baseURL) +}