fix(adkrest): render missing toolset tools and sub-agent connections in agent graph - #1730

The root-level case from #1726 works: I removed each half of teh fix in turn and TestAgentGraphWithToolsetAndSubAgent went red both times. Outside that case,the new sub-agent edge breaks the graph for any coordinator with a workflow sub-agent, and resolving toolsets lets a single toolset fail the whole graph.

I ran probe tests against GetAgentGraph and the build_graph_image handler and compared them with main at 49904c0:

• An llmagent root with a sequential, loop or parallel sub-agent gets the DOT body error: pipeline is not a node or a subgraph, served with HTTP 200. Main renders it.
• A toolset whose Tools reads ctx.UserID(), or a tool.FilterToolset predicate that reads ctx.ReadonlyState(), panics with a nil pointer dereference and the client gets EOF. main renders it.
• An mcptoolset with an unreachable server, or with per-user auth/gcp auth, turns the whole graph into a 500. Main renders it.
• An llmagent inside a sequential agent gets no edges to its toolset tools or sub-agents.

Parity with adk-python 2.11.0 (src/google/adk/cli/agent_graph.py, src/google/adk/agents/llm_agent.py): the root-level behavior matches agent_graph.py:289-301. Python skips the edge when the sub-agent is a workflow agent and conects the parent into the cluster instead (agent_graph.py:159-160, 177-178, 196-199, 292-297). It also logs and skips a toolset that fails to load, so the rest of the graph still renders (llm_agent.py:224-245).

Two things need to change before this merges. Toolset tool names go into the DOT source unquoted, so an MCP toolset with a hyphenated tool name now makes the whole graph unrenderable. The new failure paths also have no tests.

The four cases from the earlier review render correctly on this head. I drove each one through BuildGraphImageHandler and GetAgentGraph. An llmagent with a sequential, loop or parallel sub-agent gets an edge into the cluster with lhead. Toolsets that panic on ctx.UserID(), filter on ReadonlyState() or return an error are skipped and the rest of the graph renders. An llmagent inside a sequential agent gets edges to its tools, its toolset tools and its sub-agent. TestAgentGraphWithToolsetAndSubAgent fails on the PR's base (49904c0) and passes on this head (290c5dc).

Must change before merge:

• Toolset tool names reach the DOT output as bare IDs. The new loop at agentgraphgenerator.go#L394-L402 hands each toolset tool to drawNode and drawEdge, which use the name as the node ID without quoting (L321, L343). MCP servers pick these names, and hyphens are common (Context7's resolve-library-id, for example). For an llmagent with one direct tool and a toolset returning resolve-library-id, the generated DOT contains root->resolve-library-id[...], and Graphviz rejects it with syntax error in line 5 near '-'. The web UI then has nothing it can draw. Before this change the same agent rendered, just without the toolset tool. A name containing " fails the same way. A directly configured tool with a hyphen already fails like this on main, so the gap is older than this PR, but the PR now sends every MCP tool name through it. Whatever the fix, every node ID the generator emits has to be valid DOT for any tool name, on edges as well as nodes.

• The new failure paths have no tests. Under the package's tests, the recover block at L381-L383 and the skip branch at L388-L392 never run, so deleting either one leaves the suite green. No test builds an llmagent with a workflow sub-agent either, which is the drawClusterEdge call at L411-L416, and TestAgentGraphWithToolsetAndSubAgent is the only test that uses Toolsets at all, so a toolset on an llmagent nested inside a workflow agent is not covered. The cases worth pinning are an erroring toolset, a panicking toolset, an llmagent with a sequential sub-agent, and a toolset on an llmagent inside a sequential agent. Parsing the output with the existing parseGraph helper instead of strings.Contains would also catch the ID problem above: gographviz.ParseString rejects the hyphenated DOT.

package services

import (
	"context"
	"testing"

	"github.com/awalterschulze/gographviz"
	"google.golang.org/adk/v2/agent"
	"google.golang.org/adk/v2/agent/llmagent"
	"google.golang.org/adk/v2/tool"
)

type namedTool struct{ n string }

func (t namedTool) Name() string        { return t.n }
func (t namedTool) Description() string { return "" }
func (t namedTool) IsLongRunning() bool { return false }

type mcpLikeToolset struct{}

func (mcpLikeToolset) Name() string { return "mcp" }
func (mcpLikeToolset) Tools(agent.ReadonlyContext) ([]tool.Tool, error) {
	return []tool.Tool{namedTool{"resolve-library-id"}}, nil
}

func TestHyphenatedToolsetName(t *testing.T) {
	root, err := llmagent.New(llmagent.Config{Name: "root", Model: &dummyLLM{}, Toolsets: []tool.Toolset{mcpLikeToolset{}}})
	if err != nil {
		t.Fatal(err)
	}
	dot, err := GetAgentGraph(context.Background(), root, nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Log(dot)
	if _, err := gographviz.ParseString(dot); err != nil {
		t.Errorf("generated DOT does not parse: %v", err)
	}
}

The latest commit (c0ab228) adds two problems that need fixing before this PR merges. Building the graph can now leave a stdio MCP toolset broken for every later run of the agent. The comment that justifies the new 5s budget for building the graph also does not hold for stdio servers.

Both points from my earlier review are fixed. I rendered 40 names that are not plain identifiers as toolset tools, direct tools, sub-agents and workflow agents. They cover hyphens, quotes, trailing backslashes, DOT keywords, colons as in ns:tool, newlines, the empty string and invalid UTF-8. An agent cannot have an empty name, so the 40 names in four positions make 158 graphs rather than 160. On this head Graphviz accepts all of them. The 154 that are valid UTF-8 have the expected nodes and edges. On the previous head (290c5dc), 129 of the 158 failed to parse or had the wrong nodes and edges. The failure paths have tests now, and a tree that uses only plain names gives byte-identical DOT to the previous head. The checks I listed last time all pass again.

Must change before merge:

• A graph request that runs out of budget while a stdio MCP server is starting breaks its toolset until the ADK server restarts. The 5s budget from agentgraphgenerator.go#L547-L549 at c0ab228 sets the deadline on the context mcptoolset connects with (client.go#L174-L194). When the MCP server needs longer than 5s to answer initialize, the MCP SDK closes the session and kills the server's process (go-sdk v1.8.0 client.go#L390-L393). mcptoolset keeps the same CommandTransport, whose exec.Cmd cannot be started a second time (go-sdk v1.8.0 cmd.go#L29-L46). Every later run of any agent using that toolset then fails with failed to extract tools from the tool set "mcp_tool_set": ... exec: Stdout already set (tools_processor.go#L40-L44).

  The toolset breaks before anyone runs the agent, because the dev UI requests the graph as soon as an app is selected. CommandTransport is the setup the mcptoolset docs show (set.go#L63-L65). A server launched through npx -y or uvx downloads its package before it starts, so its first initialize waits on the network. The graph build connects each toolset in turn under one 5s budget. A server that starts after another one gets what is left.

  On the previous head the reproduction below passes, because the graph build waits for initialize, and later agent runs reuse the MCP session that the graph build opened. The underlying problem is older than this PR. mcptoolset reuses a CommandTransport whose process was already killed, and on main a first agent run that is cancelled while the server starts breaks the toolset the same way. The new budget breaks the toolset this way whenever a server needs more than 5s to start, without anyone cancelling anything. Any fix has to leave later runs unaffected by building the graph.

• The budget comment's timing is wrong for stdio servers. agentgraphgenerator.go#L518-L522 says an MCP toolset spends about 5s more after its context ends, which keeps the response under the launcher's 15s write timeout. For a stdio server the SDK waits 5s after closing stdin and another 5s after SIGTERM before it kills the process (go-sdk v1.8.0 cmd.go#L69-L107). The 5s budget plus up to 10s of shutdown reaches the 15s write timeout. In my test the server never answers and ignores SIGTERM. A handler calling GetAgentGraphWithTheme ran behind an http.Server with the launcher's 15s WriteTimeout, and after 15s the client got EOF with no body.

package services

import (
	"context"
	"iter"
	"os"
	"os/exec"
	"testing"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	"google.golang.org/genai"

	"google.golang.org/adk/v2/agent"
	"google.golang.org/adk/v2/agent/llmagent"
	"google.golang.org/adk/v2/model"
	"google.golang.org/adk/v2/runner"
	"google.golang.org/adk/v2/session"
	"google.golang.org/adk/v2/tool"
	"google.golang.org/adk/v2/tool/mcptoolset"
)

// Re-executed as a stdio MCP server that needs 7s before it answers.
func TestSlowStdioServerHelper(t *testing.T) {
	if os.Getenv("SLOW_MCP_SERVER") != "1" {
		t.Skip("helper process")
	}
	time.Sleep(7 * time.Second)
	srv := mcp.NewServer(&mcp.Implementation{Name: "slow", Version: "v0"}, nil)
	mcp.AddTool(srv, &mcp.Tool{Name: "lookup"}, func(context.Context, *mcp.CallToolRequest, struct{}) (*mcp.CallToolResult, struct{}, error) {
		return &mcp.CallToolResult{}, struct{}{}, nil
	})
	_ = srv.Run(context.Background(), &mcp.StdioTransport{})
	os.Exit(0)
}

type okModel struct{}

func (okModel) Name() string { return "ok" }
func (okModel) GenerateContent(context.Context, *model.LLMRequest, bool) iter.Seq2[*model.LLMResponse, error] {
	return func(yield func(*model.LLMResponse, error) bool) {
		yield(&model.LLMResponse{Content: genai.NewContentFromText("ok", genai.RoleModel)}, nil)
	}
}

func TestAgentRunAfterGraphWithSlowStdioServer(t *testing.T) {
	cmd := exec.Command(os.Args[0], "-test.run=^TestSlowStdioServerHelper$")
	cmd.Env = append(os.Environ(), "SLOW_MCP_SERVER=1")
	ts, err := mcptoolset.New(mcptoolset.Config{Transport: &mcp.CommandTransport{Command: cmd}})
	if err != nil {
		t.Fatal(err)
	}
	root, err := llmagent.New(llmagent.Config{Name: "root", Model: okModel{}, Toolsets: []tool.Toolset{ts}})
	if err != nil {
		t.Fatal(err)
	}

	// The dev UI requests the graph when the app is opened, before any run.
	if _, err := GetAgentGraph(context.Background(), root, nil); err != nil {
		t.Fatal(err)
	}

	r, err := runner.New(runner.Config{AppName: "app", Agent: root, SessionService: session.InMemoryService(), AutoCreateSession: true})
	if err != nil {
		t.Fatal(err)
	}
	for _, err := range r.Run(context.Background(), "user", "s1", genai.NewContentFromText("hi", genai.RoleUser), agent.RunConfig{}) {
		if err != nil {
			t.Fatalf("agent run after the graph: %v", err)
		}
	}
}

On your logging questions

WARN fits the level table in AGENTS.md (AGENTS.md#L163-L169), since the graph still renders. Logging the error's type rather than its text follows the AGENTS.md rule against URLs in log lines (AGENTS.md#L134-L143), because an MCP error can carry the server's URL. The //nolint:forbidigo follows the same pattern as exp/appinfo, which also skips a toolset that cannot list its tools and logs only the error's type (appinfo.go#L309).

My change request from my earlier review stands. A graph request can still leave an MCP toolset unusable for every later run and block its Close, which is what I asked the fix to rule out.

What the new commit fixes:

• With a stdio server that needs 7s before it answers, the graph now returns in 5s without that server's tools. Two agent runs started after the graph request now succeed, and the next graph request shows the server's tool. On the previous head (c0ab228) both runs failed with exec: Stdout already set.
• The graph's response now stays within its 5s budget. With a stdio server that never answers and ignores SIGTERM, a graph endpoint served under the launcher's 15s write timeout returns 200 in 5s.
• The toolset name for the log is now read under a recover.
• toolsearch now lists its own tools in the graph.

The build, lint and test checks from my earlier reviews pass again on the current head, 07f79b4, and every line link below is pinned to it.

On the older problems you listed, #2111 already covers mcptoolset reusing a CommandTransport whose process was killed, and the exp/appinfo timeout is worth an issue of its own.

Must change before merge

A graph request against an MCP server that never answers leaves that toolset locked for the life of the process. The listing runs under context.WithoutCancel (agentgraphtoolsets.go#L111), so nothing can end it. While mcptoolset connects, getSession holds the connection's lock for the whole connect (client.go#L211-L231, client.go#L264-L272). The MCP SDK puts no timeout of its own on connecting. Against a server that accepts the connection and never answers, holding the lock through a connect that nothing can cancel has two effects:

• Later runs never reconnect. Every later agent run waits for the lock until its own deadline and fails. For an HTTP server, later runs keep failing after the server recovers, because no new connect is ever tried. In the reproduction below the server holds its first request and serves every later one normally. After a graph request, two runs with a 6s deadline both fail, and the server sees one request in total. Without the graph request the second run succeeds. On the previous head the graph's deadline released the lock and both runs succeeded. With Config.Auth set to auth.StaticToken(...), which uses the per-credential session pool, both runs fail the same way on this head.
• Close never returns. It takes the same lock and has no context (client.go#L91-L95), so it blocks for good. A program that defers Close then cannot shut down, and deferring it is what the mcptoolset docs show (set.go#L51-L69). A stdio server that never answers is never stopped either. Close blocks and the server stays alive whether or not the server ignores SIGTERM. With a server that keeps default signal handling, a goroutine dump shows Close waiting while the listing sits in Client.Connect. Without the graph request, Close returns at once. It also returns at once on the merge base with main (be8b100) and on the previous head, with or without a graph request.

The comment at agentgraphtoolsets.go#L40-L47 says that a listing left running finishes in the background and that later agent runs reuse its session. That does not hold when the server never answers. Whatever the fix, building the graph must not leave a toolset in a state that breaks or blocks a later run, a later graph or Close, including when its server never answers.

On the previous head, ending the listing at a deadline killed a stdio server's process. The toolset then stayed broken, because mcptoolset keeps reusing a transport whose process is gone (#2111). So a deadline that frees the lock within seconds brings back the failure from my earlier review for any server slower than that deadline.

package services

import (
	"context"
	"io"
	"iter"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	"google.golang.org/genai"

	"google.golang.org/adk/v2/agent"
	"google.golang.org/adk/v2/agent/llmagent"
	"google.golang.org/adk/v2/model"
	"google.golang.org/adk/v2/runner"
	"google.golang.org/adk/v2/session"
	"google.golang.org/adk/v2/tool"
	"google.golang.org/adk/v2/tool/mcptoolset"
)

type replyModel struct{}

func (replyModel) Name() string { return "reply" }
func (replyModel) GenerateContent(context.Context, *model.LLMRequest, bool) iter.Seq2[*model.LLMResponse, error] {
	return func(yield func(*model.LLMResponse, error) bool) {
		yield(&model.LLMResponse{Content: genai.NewContentFromText("ok", genai.RoleModel)}, nil)
	}
}

// The MCP server holds its first request until the test ends and serves every
// later request normally, like a server that stalls once and recovers.
func TestAgentRunAfterGraphWithStalledServer(t *testing.T) {
	srv := mcp.NewServer(&mcp.Implementation{Name: "stall", Version: "v0"}, nil)
	mcp.AddTool(srv, &mcp.Tool{Name: "lookup"}, func(context.Context, *mcp.CallToolRequest, struct{}) (*mcp.CallToolResult, struct{}, error) {
		return &mcp.CallToolResult{}, struct{}{}, nil
	})
	h := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return srv }, nil)
	stop := make(chan struct{})
	var requests atomic.Int32
	hs := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if requests.Add(1) == 1 {
			select {
			case <-stop:
			case <-r.Context().Done():
			}
			return
		}
		h.ServeHTTP(w, r)
	}))
	t.Cleanup(func() { close(stop); hs.CloseClientConnections(); hs.Close() })

	ts, err := mcptoolset.New(mcptoolset.Config{Endpoint: hs.URL})
	if err != nil {
		t.Fatal(err)
	}
	root, err := llmagent.New(llmagent.Config{Name: "root", Model: replyModel{}, Toolsets: []tool.Toolset{ts}})
	if err != nil {
		t.Fatal(err)
	}

	// Remove this call and the second run succeeds.
	if _, err := GetAgentGraph(context.Background(), root, nil); err != nil {
		t.Fatal(err)
	}

	r, err := runner.New(runner.Config{AppName: "app", Agent: root, SessionService: session.InMemoryService(), AutoCreateSession: true})
	if err != nil {
		t.Fatal(err)
	}
	var runErr error
	for _, sid := range []string{"s1", "s2"} {
		ctx, cancel := context.WithTimeout(context.Background(), 6*time.Second)
		runErr = nil
		for _, err := range r.Run(ctx, "user", sid, genai.NewContentFromText("hi", genai.RoleUser), agent.RunConfig{}) {
			if err != nil {
				runErr = err
				break
			}
		}
		cancel()
	}
	if runErr != nil {
		t.Errorf("second run, after the server recovered: %v (server saw %d requests)", runErr, requests.Load())
	}

	closed := make(chan struct{})
	go func() { _ = ts.(io.Closer).Close(); close(closed) }()
	select {
	case <-closed:
	case <-time.After(10 * time.Second):
		t.Error("Close still blocked after 10s")
	}
}