fix(workflow): keep session identity across single-turn node contexts - #6697
Open
chelsealong wants to merge 1 commit into
Open
fix(workflow): keep session identity across single-turn node contexts#6697chelsealong wants to merge 1 commit into
chelsealong wants to merge 1 commit into
Conversation
Fixes google#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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #6691
Summary
When a
Workflow's root agent runs multiple single-turnLlmAgentnodesand
EventsCompactionConfig.token_thresholdcompaction fires mid-invocation(not on the last node), a
DatabaseSessionService-backed session breaks with:InMemorySessionServicedoes not reproduce this, because it never checks astorage revision marker.
Root cause
prepare_llm_agent_context()insrc/google/adk/workflow/_llm_agent_wrapper.pygave every single-turn
LlmAgentnode its ownInvocationContext, andadditionally replaced that context's
sessionwithic.session.model_copy(deep=False).Session.eventsandSession.stateare mutable containers, so the shallowcopy shared them by reference with the parent session — but
Session.last_update_timeand the private storage-revision marker used byDatabaseSessionServicefor stale-write detection are scalar/privateattributes, so the copy duplicated them by value.
When such a node's request-time token-threshold compaction runs
(
CompactionRequestProcessor), it appends the compaction event through thatnode's own session copy, advancing the revision marker on the copy only. The
parent context's session — the one later nodes copy from, and the one the
Runner's own event-consumption loop appends to — never learns about that
advance. The next append made through that stale copy is rejected by
DatabaseSessionService.append_eventas a stale write.Fix
Drop the
ic.session = ic.session.model_copy(deep=False)line. Sinceevents/statewere already shared by reference, this copy provided noreal isolation — it only caused the marker/timestamp to diverge. Every node
in the same Workflow invocation now shares the exact same
Sessionobject,so a revision-marker advance from one node's compaction write is visible to
every other node.
Testing plan
Added
test_mid_workflow_compaction_does_not_stale_later_node_appendtotests/unittests/apps/test_compaction_runner_e2e.py. It runs a two-nodeWorkflow(agent1 -> agent2, both single-turn) through a realRunner.run_asyncagainst aDatabaseSessionService(sqlite), across twoturns: a short turn (no compaction), then a long turn whose estimated prompt
token count exceeds
token_threshold, triggering mid-invocation compactionon
agent1beforeagent2runs. It assertsagent2's response is producedand persisted to storage.
error (
ValueError: The session has been modified in storage since it was loaded...), by temporarily reverting only_llm_agent_wrapper.pyandre-running.
pre-commit run(isort, pyink, ruff, addlicense, ADK compliance checks) passeson both changed files.
AI-assistance disclosure
This PR was prepared with the help of an AI coding agent (Claude Code), which
investigated the root cause, implemented the fix, and wrote/verified the
regression test. I reviewed and take responsibility for the change.