Skip to content

Commit 847df16

Browse files
GWealecopybara-github
authored andcommitted
fix: handle App instances returned by agent_loader.load_agent
The `agent_loader.load_agent` method can now return an `App` object. This change unwraps the `App` to get its `root_agent` before passing it to the graph builder, makes sure a `BaseAgent` instance is always used PiperOrigin-RevId: 817209601
1 parent 55aa6f6 commit 847df16

File tree

2 files changed

+60
-1
lines changed

2 files changed

+60
-1
lines changed

src/google/adk/cli/adk_web_server.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1429,7 +1429,14 @@ async def get_event_graph(
14291429

14301430
function_calls = event.get_function_calls()
14311431
function_responses = event.get_function_responses()
1432-
root_agent = self.agent_loader.load_agent(app_name)
1432+
agent_or_app = self.agent_loader.load_agent(app_name)
1433+
# The loader may return an App; unwrap to its root agent so the graph builder
1434+
# receives a BaseAgent instance.
1435+
root_agent = (
1436+
agent_or_app.root_agent
1437+
if isinstance(agent_or_app, App)
1438+
else agent_or_app
1439+
)
14331440
dot_graph = None
14341441
if function_calls:
14351442
function_call_highlights = []

tests/unittests/cli/test_fast_api.py

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from fastapi.testclient import TestClient
3030
from google.adk.agents.base_agent import BaseAgent
3131
from google.adk.agents.run_config import RunConfig
32+
from google.adk.apps.app import App
3233
from google.adk.cli.fast_api import get_fast_api_app
3334
from google.adk.evaluation.eval_case import EvalCase
3435
from google.adk.evaluation.eval_case import Invocation
@@ -39,6 +40,7 @@
3940
from google.adk.events.event_actions import EventActions
4041
from google.adk.runners import Runner
4142
from google.adk.sessions.base_session_service import ListSessionsResponse
43+
from google.adk.sessions.session import Session
4244
from google.genai import types
4345
from pydantic import BaseModel
4446
import pytest
@@ -1007,6 +1009,56 @@ def test_debug_trace(test_app):
10071009
logger.info("Debug trace test completed successfully")
10081010

10091011

1012+
def test_get_event_graph_returns_dot_src_for_app_agent():
1013+
"""Ensure graph endpoint unwraps App instances before building the graph."""
1014+
from google.adk.cli.adk_web_server import AdkWebServer
1015+
1016+
root_agent = DummyAgent(name="dummy_agent")
1017+
app_agent = App(name="test_app", root_agent=root_agent)
1018+
1019+
class Loader:
1020+
1021+
def load_agent(self, app_name):
1022+
return app_agent
1023+
1024+
def list_agents(self):
1025+
return [app_agent.name]
1026+
1027+
session_service = AsyncMock()
1028+
session = Session(
1029+
id="session_id",
1030+
app_name="test_app",
1031+
user_id="user",
1032+
state={},
1033+
events=[Event(author="dummy_agent")],
1034+
)
1035+
event_id = session.events[0].id
1036+
session_service.get_session.return_value = session
1037+
1038+
adk_web_server = AdkWebServer(
1039+
agent_loader=Loader(),
1040+
session_service=session_service,
1041+
memory_service=MagicMock(),
1042+
artifact_service=MagicMock(),
1043+
credential_service=MagicMock(),
1044+
eval_sets_manager=MagicMock(),
1045+
eval_set_results_manager=MagicMock(),
1046+
agents_dir=".",
1047+
)
1048+
1049+
fast_api_app = adk_web_server.get_fast_api_app(
1050+
setup_observer=lambda _observer, _server: None,
1051+
tear_down_observer=lambda _observer, _server: None,
1052+
)
1053+
1054+
client = TestClient(fast_api_app)
1055+
response = client.get(
1056+
f"/apps/test_app/users/user/sessions/session_id/events/{event_id}/graph"
1057+
)
1058+
assert response.status_code == 200
1059+
assert "dotSrc" in response.json()
1060+
1061+
10101062
@pytest.mark.skipif(
10111063
sys.version_info < (3, 10), reason="A2A requires Python 3.10+"
10121064
)

0 commit comments

Comments
 (0)