Skip to content

Commit 066fbce

Browse files
GWealecopybara-github
authored andcommitted
fix: apply run_config custom_metadata to user event in node runtime
The node-based runtime path (LlmAgent chat mode and BaseNode) appended the user event without the run-level custom_metadata that the legacy agent path and v1 both apply, so it was dropped from session history. Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 934433213
1 parent a40f199 commit 066fbce

3 files changed

Lines changed: 61 additions & 0 deletions

File tree

src/google/adk/runners.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -736,6 +736,7 @@ async def _append_user_event(
736736
iso = _find_active_task_isolation_scope(ic.session)
737737
if iso is not None:
738738
event.isolation_scope = iso
739+
_apply_run_config_custom_metadata(event, ic.run_config)
739740
return await self.session_service.append_event(
740741
session=ic.session, event=event
741742
)

tests/unittests/runners/test_runner_node.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from google.adk.agents.callback_context import CallbackContext
2828
from google.adk.agents.context import Context
2929
from google.adk.agents.llm_agent import LlmAgent
30+
from google.adk.agents.run_config import RunConfig
3031
from google.adk.events.event import Event
3132
from google.adk.runners import Runner
3233
from google.adk.sessions.in_memory_session_service import InMemorySessionService
@@ -256,6 +257,30 @@ async def test_multiple_invocations_accumulate_events():
256257
assert outputs == ['Echo: first', 'Echo: second', 'Echo: third']
257258

258259

260+
@pytest.mark.asyncio
261+
async def test_run_config_custom_metadata_stamps_user_event():
262+
"""The node path stamps the user event with run-level custom_metadata."""
263+
ss = InMemorySessionService()
264+
runner = Runner(
265+
app_name='test', node=_EchoNode(name='echo'), session_service=ss
266+
)
267+
session = await ss.create_session(app_name='test', user_id='u')
268+
269+
async for _ in runner.run_async(
270+
user_id='u',
271+
session_id=session.id,
272+
new_message=_user_message('hi'),
273+
run_config=RunConfig(custom_metadata={'turn_id': 't-1'}),
274+
):
275+
pass
276+
277+
updated = await ss.get_session(
278+
app_name='test', user_id='u', session_id=session.id
279+
)
280+
user_event = next(e for e in updated.events if e.author == 'user')
281+
assert user_event.custom_metadata == {'turn_id': 't-1'}
282+
283+
259284
# ---------------------------------------------------------------------------
260285
# yield_user_message
261286
# ---------------------------------------------------------------------------

tests/unittests/test_runners.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -972,6 +972,41 @@ async def test_run_config_custom_metadata_propagates_to_events():
972972
assert user_event.custom_metadata == {"request_id": "req-1"}
973973

974974

975+
@pytest.mark.asyncio
976+
async def test_run_config_custom_metadata_stamps_user_event_in_chat_mode():
977+
"""LlmAgent chat path stamps the user event with run-level custom_metadata."""
978+
session_service = InMemorySessionService()
979+
980+
def _before_agent_callback(callback_context) -> types.Content:
981+
del callback_context # Unused; short-circuits the model call.
982+
return types.Content(role="model", parts=[types.Part(text="hi back")])
983+
984+
agent = LlmAgent(
985+
name="chat_agent", before_agent_callback=_before_agent_callback
986+
)
987+
runner = Runner(
988+
app_name=TEST_APP_ID, agent=agent, session_service=session_service
989+
)
990+
await session_service.create_session(
991+
app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
992+
)
993+
994+
run_config = RunConfig(custom_metadata={"turn_id": "t-1"})
995+
async for _ in runner.run_async(
996+
user_id=TEST_USER_ID,
997+
session_id=TEST_SESSION_ID,
998+
new_message=types.Content(role="user", parts=[types.Part(text="hi")]),
999+
run_config=run_config,
1000+
):
1001+
pass
1002+
1003+
session = await session_service.get_session(
1004+
app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
1005+
)
1006+
user_event = next(event for event in session.events if event.author == "user")
1007+
assert user_event.custom_metadata == {"turn_id": "t-1"}
1008+
1009+
9751010
class TestRunnerWithPlugins:
9761011
"""Tests for Runner with plugins."""
9771012

0 commit comments

Comments
 (0)