From ee73a5c76ef65ef2ba42a598251dc58b1992b1c8 Mon Sep 17 00:00:00 2001 From: chelsealong Date: Wed, 12 Aug 2026 16:21:12 +0000 Subject: [PATCH] fix(workflow): keep session identity across single-turn node contexts Fixes #6691 prepare_llm_agent_context() replaced a single-turn LlmAgent node's InvocationContext.session with a shallow copy. Session.events and Session.state are mutable containers so the copy shared them by reference, but Session.last_update_time and the private storage revision marker were duplicated by value. When such a node triggers mid-invocation token-threshold compaction, the compaction event is appended through that per-node session copy, advancing the marker there but not on the parent/session object other nodes derive their own copies from. DatabaseSessionService then rejects the next append made through the stale parent copy with "The session has been modified in storage since it was loaded." Dropping the copy keeps every node's InvocationContext pointing at the same Session object, so the revision marker stays consistent for the rest of the Workflow. --- src/google/adk/workflow/_llm_agent_wrapper.py | 7 +- .../apps/test_compaction_runner_e2e.py | 80 +++++++++++++++++++ 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/src/google/adk/workflow/_llm_agent_wrapper.py b/src/google/adk/workflow/_llm_agent_wrapper.py index f1209781ef3..3bc550ee130 100644 --- a/src/google/adk/workflow/_llm_agent_wrapper.py +++ b/src/google/adk/workflow/_llm_agent_wrapper.py @@ -294,7 +294,12 @@ def prepare_llm_agent_context(agent: LlmAgent, ctx: Context) -> Context: ) agent_ctx.isolation_scope = ctx.isolation_scope - ic.session = ic.session.model_copy(deep=False) + # Do not copy `ic.session`: it stays the same object shared with the + # parent context. A shallow copy here duplicates `last_update_time` and + # the storage revision marker, so a DatabaseSessionService write made + # through this node's session (e.g. mid-invocation compaction) would + # never be reflected on the parent's session object, causing the next + # node's write through the parent copy to be rejected as stale. return agent_ctx diff --git a/tests/unittests/apps/test_compaction_runner_e2e.py b/tests/unittests/apps/test_compaction_runner_e2e.py index a1f6a515d70..c30450d1f0c 100644 --- a/tests/unittests/apps/test_compaction_runner_e2e.py +++ b/tests/unittests/apps/test_compaction_runner_e2e.py @@ -24,7 +24,10 @@ from google.adk.apps.llm_event_summarizer import LlmEventSummarizer from google.adk.events.event import Event from google.adk.runners import Runner +from google.adk.sessions.database_session_service import DatabaseSessionService from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.workflow import START +from google.adk.workflow._workflow import Workflow from google.genai import types from google.genai.types import Content from google.genai.types import Part @@ -210,3 +213,80 @@ async def test_runner_appends_sliding_window_compaction_event(): assert ( compaction_events ), "runner did not append the sliding-window compaction event" + + +@pytest.mark.asyncio +async def test_mid_workflow_compaction_does_not_stale_later_node_append(): + """Compacting a non-last Workflow node must not stale-fail later nodes. + + Each single-turn LlmAgent node in a Workflow runs against its own copy of + the InvocationContext (see ``prepare_llm_agent_context``). Token-threshold + compaction writes through that per-node context's session, so a + ``DatabaseSessionService`` (which rejects an ``append_event`` whose + in-memory revision marker is behind storage) must still accept later + writes made through the shared session object: they should see the marker + the compaction write left behind, not a stale copy of it. + """ + agent1_model = testing_utils.MockModel.create( + responses=["agent1 turn 1", "agent1 turn 2"] + ) + agent2_model = testing_utils.MockModel.create( + responses=["agent2 turn 1", "agent2 turn 2"] + ) + agent1 = Agent(name="agent1", model=agent1_model, mode="single_turn") + agent2 = Agent(name="agent2", model=agent2_model, mode="single_turn") + workflow = Workflow( + name="wf", + edges=[(START, agent1), (agent1, agent2)], + ) + app = App( + name="test_app", + root_agent=workflow, + events_compaction_config=EventsCompactionConfig( + token_threshold=100, + event_retention_size=0, + summarizer=LlmEventSummarizer( + llm=testing_utils.MockModel.create(responses=["summary"]) + ), + ), + ) + session_service = DatabaseSessionService("sqlite+aiosqlite:///:memory:") + await session_service.create_session( + app_name="test_app", user_id="u1", session_id="s1" + ) + runner = Runner(app=app, session_service=session_service) + + # Turn 1: short message, well below the token threshold estimate. No + # compaction triggered. + async for _ in runner.run_async( + user_id="u1", + session_id="s1", + new_message=Content(role="user", parts=[Part(text="hi")]), + ): + pass + + # Turn 2: a long message pushes agent1's estimated prompt token count + # above the threshold, so its compaction request-processor compacts + # mid-invocation, before agent2 runs. + long_message = "lorem ipsum dolor sit amet " * 40 + events = [ + event + async for event in runner.run_async( + user_id="u1", + session_id="s1", + new_message=Content(role="user", parts=[Part(text=long_message)]), + ) + ] + + agent2_events = [event for event in events if event.author == "agent2"] + assert agent2_events, "agent2's response was not produced/persisted" + + refreshed = await session_service.get_session( + app_name="test_app", user_id="u1", session_id="s1" + ) + persisted_agent2_events = [ + event for event in refreshed.events if event.author == "agent2" + ] + assert ( + len(persisted_agent2_events) == 2 + ), "agent2's response was not persisted to storage"