Skip to content

Commit f15b685

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
fix: replay session_events into a scratch session and delete it afterwards
`AdkApp.async_stream_query(session_events=...)` raised `AttributeError` in the `vertexai` template and, in every template, left behind the managed session it created for the replayed events. The session is now read back before the replay and deleted when the query ends. Fixes #7118, fixes #7119 PiperOrigin-RevId: 987536032
1 parent 4661183 commit f15b685

4 files changed

Lines changed: 340 additions & 38 deletions

File tree

‎agentplatform/frameworks/adk.py‎

Lines changed: 35 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1209,6 +1209,8 @@ async def async_stream_query(
12091209
session_events (Optional[List[Dict[str, Any]]]):
12101210
Optional. The session events to use for the query. This will be
12111211
used to initialize the session if `session_id` is not provided.
1212+
That session exists only for this query and is deleted once the
1213+
stream ends.
12121214
run_config (Optional[Dict[str, Any]]):
12131215
Optional. The run config to use for the query. If you want to
12141216
pass in a `run_config` pydantic object, you can pass in a dict
@@ -1244,18 +1246,23 @@ async def async_stream_query(
12441246
raise ValueError(
12451247
"Only one of session_id and session_events should be specified."
12461248
)
1249+
scratch_session_id = None
12471250
if not session_id:
12481251
session = await self.async_create_session(user_id=user_id)
12491252
session_id = session["id"]
12501253
if session_events is not None:
1254+
scratch_session_id = session_id
1255+
1256+
try:
1257+
if scratch_session_id:
12511258
# We allow for session_events to be an empty list.
12521259
from google.adk.events.event import Event
12531260

12541261
session_service = self._tmpl_attrs.get("session_service")
12551262
session_obj = await session_service.get_session(
12561263
app_name=self._app_name(),
12571264
user_id=user_id,
1258-
session_id=session_id,
1265+
session_id=scratch_session_id,
12591266
)
12601267
for event in session_events:
12611268
if not isinstance(event, Event):
@@ -1265,28 +1272,38 @@ async def async_stream_query(
12651272
event=event,
12661273
)
12671274

1268-
run_config = _validate_run_config(run_config)
1269-
if run_config:
1270-
events_async = self._tmpl_attrs.get("runner").run_async(
1271-
user_id=user_id,
1272-
session_id=session_id,
1273-
new_message=content,
1274-
run_config=run_config,
1275-
**kwargs,
1276-
)
1277-
else:
1278-
events_async = self._tmpl_attrs.get("runner").run_async(
1279-
user_id=user_id,
1280-
session_id=session_id,
1281-
new_message=content,
1282-
**kwargs,
1283-
)
1275+
run_config = _validate_run_config(run_config)
1276+
if run_config:
1277+
events_async = self._tmpl_attrs.get("runner").run_async(
1278+
user_id=user_id,
1279+
session_id=session_id,
1280+
new_message=content,
1281+
run_config=run_config,
1282+
**kwargs,
1283+
)
1284+
else:
1285+
events_async = self._tmpl_attrs.get("runner").run_async(
1286+
user_id=user_id,
1287+
session_id=session_id,
1288+
new_message=content,
1289+
**kwargs,
1290+
)
12841291

1285-
try:
12861292
async for event in events_async:
12871293
# Yield the event data as a dictionary
12881294
yield _runtimes_utils.dump_event_for_json(event)
12891295
finally:
1296+
# The caller never sees the id of a session created for
1297+
# session_events, so nothing else can delete it.
1298+
if scratch_session_id:
1299+
try:
1300+
await self._tmpl_attrs.get("session_service").delete_session(
1301+
app_name=self._app_name(),
1302+
user_id=user_id,
1303+
session_id=scratch_session_id,
1304+
)
1305+
except Exception as e: # pylint: disable=broad-exception-caught
1306+
_warn(f"Failed to delete scratch session {scratch_session_id}: {e}")
12901307
# Avoid telemetry data loss having to do with CPU throttling on instance turndown
12911308
_ = await _force_flush_otel(
12921309
tracing_enabled=self._tracing_enabled(),

‎tests/unit/agentplatform/frameworks/test_frameworks_adk.py‎

Lines changed: 61 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -612,7 +612,7 @@ async def test_async_stream_query_with_empty_session_events(
612612
events.append(event)
613613
assert app._tmpl_attrs.get("session_service") is not None
614614
sessions = app.list_sessions(user_id=_TEST_USER_ID)
615-
assert len(sessions.sessions) == 1
615+
assert not sessions.sessions
616616

617617
@pytest.mark.asyncio
618618
async def test_async_stream_query_with_session_events(
@@ -633,7 +633,66 @@ async def test_async_stream_query_with_session_events(
633633
events.append(event)
634634
assert app._tmpl_attrs.get("session_service") is not None
635635
sessions = app.list_sessions(user_id=_TEST_USER_ID)
636-
assert len(sessions.sessions) == 1
636+
assert not sessions.sessions
637+
638+
@pytest.mark.asyncio
639+
async def test_async_stream_query_deletes_session_when_replay_fails(
640+
self,
641+
default_instrumentor_builder_mock: mock.Mock,
642+
get_project_id_mock: mock.Mock,
643+
):
644+
app = adk_template.AdkApp(agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL))
645+
app.set_up()
646+
app._tmpl_attrs["runner"] = _MockRunner()
647+
with pytest.raises(ValueError):
648+
async for _ in app.async_stream_query(
649+
user_id=_TEST_USER_ID,
650+
session_events=[123],
651+
message="test message",
652+
):
653+
pass
654+
assert not app.list_sessions(user_id=_TEST_USER_ID).sessions
655+
656+
@pytest.mark.asyncio
657+
async def test_async_stream_query_keeps_session_without_session_events(
658+
self,
659+
default_instrumentor_builder_mock: mock.Mock,
660+
get_project_id_mock: mock.Mock,
661+
):
662+
app = adk_template.AdkApp(agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL))
663+
app.set_up()
664+
app._tmpl_attrs["runner"] = _MockRunner()
665+
async for _ in app.async_stream_query(
666+
user_id=_TEST_USER_ID,
667+
message="test message",
668+
):
669+
pass
670+
assert len(app.list_sessions(user_id=_TEST_USER_ID).sessions) == 1
671+
672+
@pytest.mark.asyncio
673+
async def test_async_stream_query_warns_when_session_delete_fails(
674+
self,
675+
default_instrumentor_builder_mock: mock.Mock,
676+
get_project_id_mock: mock.Mock,
677+
):
678+
app = adk_template.AdkApp(agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL))
679+
app.set_up()
680+
app._tmpl_attrs["runner"] = _MockRunner()
681+
with mock.patch.object(
682+
app._tmpl_attrs["session_service"],
683+
"delete_session",
684+
side_effect=RuntimeError("delete failed"),
685+
), mock.patch.object(adk_template, "_warn") as warn_mock:
686+
events = [
687+
event
688+
async for event in app.async_stream_query(
689+
user_id=_TEST_USER_ID,
690+
session_events=[],
691+
message="test message",
692+
)
693+
]
694+
assert len(events) == 1
695+
warn_mock.assert_called_once()
637696

638697
@pytest.mark.asyncio
639698
@mock.patch.dict(

‎tests/unit/vertex_adk/test_agent_engine_templates_adk.py‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,83 @@ def __init__(self, name: str, model: str):
7070
_TEST_USER_ID = "test_user_id"
7171
_TEST_AGENT_NAME = "test_agent"
7272
_TEST_AGENT = Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL)
73+
_TEST_SESSION_EVENTS = [
74+
{
75+
"author": "user",
76+
"content": {
77+
"parts": [
78+
{
79+
"text": "What is the exchange rate from US dollars to "
80+
"Swedish krona on 2025-09-25?"
81+
}
82+
],
83+
"role": "user",
84+
},
85+
"id": "8967297909049524224",
86+
"invocationId": "e-308f65d7-a99f-41e3-b80d-40feb5f1b065",
87+
"timestamp": 1765832134.629513,
88+
},
89+
{
90+
"author": "currency_exchange_agent",
91+
"content": {
92+
"parts": [
93+
{
94+
"functionCall": {
95+
"args": {
96+
"currency_date": "2025-09-25",
97+
"currency_from": "USD",
98+
"currency_to": "SEK",
99+
},
100+
"id": "adk-136738ad-9e57-4cfb-8e23-b0f3e50a37d7",
101+
"name": "get_exchange_rate",
102+
}
103+
}
104+
],
105+
"role": "model",
106+
},
107+
"id": "3155402589927899136",
108+
"invocationId": "e-308f65d7-a99f-41e3-b80d-40feb5f1b065",
109+
"timestamp": 1765832134.723713,
110+
},
111+
{
112+
"author": "currency_exchange_agent",
113+
"content": {
114+
"parts": [
115+
{
116+
"functionResponse": {
117+
"id": "adk-136738ad-9e57-4cfb-8e23-b0f3e50a37d7",
118+
"name": "get_exchange_rate",
119+
"response": {
120+
"amount": 1,
121+
"base": "USD",
122+
"date": "2025-09-25",
123+
"rates": {"SEK": 9.4118},
124+
},
125+
}
126+
}
127+
],
128+
"role": "user",
129+
},
130+
"id": "1678221912150376448",
131+
"invocationId": "e-308f65d7-a99f-41e3-b80d-40feb5f1b065",
132+
"timestamp": 1765832135.764961,
133+
},
134+
{
135+
"author": "currency_exchange_agent",
136+
"content": {
137+
"parts": [
138+
{
139+
"text": "The exchange rate from US dollars to Swedish "
140+
"krona on 2025-09-25 is 1 USD to 9.4118 SEK."
141+
}
142+
],
143+
"role": "model",
144+
},
145+
"id": "2470855446567583744",
146+
"invocationId": "e-308f65d7-a99f-41e3-b80d-40feb5f1b065",
147+
"timestamp": 1765832135.853299,
148+
},
149+
]
73150
_TEST_SESSION = {
74151
"id": "ca18c25a-644b-4e13-9b24-78c150ec3eb9",
75152
"app_name": "default_app_name",
@@ -317,6 +394,23 @@ async def run_async(self, *args, **kwargs):
317394
)
318395

319396

397+
class _SessionReadingRunner:
398+
"""Records which events the session holds when the run starts."""
399+
400+
def __init__(self, session_service, app_name):
401+
self._session_service = session_service
402+
self._app_name = app_name
403+
self.seen_event_ids = None
404+
405+
async def run_async(self, *, user_id, session_id, **kwargs):
406+
session = await self._session_service.get_session(
407+
app_name=self._app_name, user_id=user_id, session_id=session_id
408+
)
409+
self.seen_event_ids = [event.id for event in session.events]
410+
async for event in _MockRunner().run_async():
411+
yield event
412+
413+
320414
@pytest.mark.usefixtures("google_auth_mock")
321415
class TestAdkApp:
322416
def test_adk_version(self):
@@ -436,6 +530,116 @@ async def test_async_stream_query(
436530
events.append(event)
437531
assert len(events) == 1
438532

533+
@pytest.mark.asyncio
534+
async def test_async_stream_query_replays_session_events_then_deletes_session(
535+
self,
536+
default_instrumentor_builder_mock: mock.Mock,
537+
get_project_id_mock: mock.Mock,
538+
):
539+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
540+
app.set_up()
541+
runner = _SessionReadingRunner(
542+
app._tmpl_attrs["session_service"], app._app_name()
543+
)
544+
app._tmpl_attrs["runner"] = runner
545+
events = [
546+
event
547+
async for event in app.async_stream_query(
548+
user_id=_TEST_USER_ID,
549+
session_events=_TEST_SESSION_EVENTS,
550+
message="on the day after that?",
551+
)
552+
]
553+
assert len(events) == 1
554+
assert runner.seen_event_ids == [e["id"] for e in _TEST_SESSION_EVENTS]
555+
assert not app.list_sessions(user_id=_TEST_USER_ID).sessions
556+
557+
@pytest.mark.asyncio
558+
async def test_async_stream_query_deletes_session_when_replay_fails(
559+
self,
560+
default_instrumentor_builder_mock: mock.Mock,
561+
get_project_id_mock: mock.Mock,
562+
):
563+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
564+
app.set_up()
565+
app._tmpl_attrs["runner"] = _MockRunner()
566+
with pytest.raises(ValueError):
567+
async for _ in app.async_stream_query(
568+
user_id=_TEST_USER_ID,
569+
session_events=[123],
570+
message="test message",
571+
):
572+
pass
573+
assert not app.list_sessions(user_id=_TEST_USER_ID).sessions
574+
575+
@pytest.mark.asyncio
576+
async def test_async_stream_query_keeps_session_without_session_events(
577+
self,
578+
default_instrumentor_builder_mock: mock.Mock,
579+
get_project_id_mock: mock.Mock,
580+
):
581+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
582+
app.set_up()
583+
app._tmpl_attrs["runner"] = _MockRunner()
584+
async for _ in app.async_stream_query(
585+
user_id=_TEST_USER_ID,
586+
message="test message",
587+
):
588+
pass
589+
assert len(app.list_sessions(user_id=_TEST_USER_ID).sessions) == 1
590+
591+
@pytest.mark.asyncio
592+
async def test_async_stream_query_warns_when_session_delete_fails(
593+
self,
594+
default_instrumentor_builder_mock: mock.Mock,
595+
get_project_id_mock: mock.Mock,
596+
):
597+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
598+
app.set_up()
599+
app._tmpl_attrs["runner"] = _MockRunner()
600+
with mock.patch.object(
601+
app._tmpl_attrs["session_service"],
602+
"delete_session",
603+
side_effect=RuntimeError("delete failed"),
604+
), mock.patch.object(adk_template, "_warn") as warn_mock:
605+
events = [
606+
event
607+
async for event in app.async_stream_query(
608+
user_id=_TEST_USER_ID,
609+
session_events=[],
610+
message="test message",
611+
)
612+
]
613+
assert len(events) == 1
614+
warn_mock.assert_called_once()
615+
616+
@pytest.mark.asyncio
617+
async def test_async_stream_query_passes_run_config(
618+
self,
619+
default_instrumentor_builder_mock: mock.Mock,
620+
get_project_id_mock: mock.Mock,
621+
):
622+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
623+
app.set_up()
624+
runner_kwargs = {}
625+
626+
async def run_async(**kwargs):
627+
runner_kwargs.update(kwargs)
628+
async for event in _MockRunner().run_async():
629+
yield event
630+
631+
app._tmpl_attrs["runner"] = mock.Mock(run_async=run_async)
632+
events = [
633+
event
634+
async for event in app.async_stream_query(
635+
user_id=_TEST_USER_ID,
636+
message="test message",
637+
run_config={"max_llm_calls": 5},
638+
)
639+
]
640+
assert len(events) == 1
641+
assert runner_kwargs["run_config"].max_llm_calls == 5
642+
439643
def test_set_up_runner_auto_create_session_enabled(
440644
self,
441645
default_instrumentor_builder_mock: mock.Mock,

0 commit comments

Comments
 (0)