Repository navigation
fix(flows): persist the dict key removal of a later parallel tool call - #7703
danieldutu wants to merge 1 commit into
Conversation
Somay-kousis
left a comment
There was a problem hiding this comment.
What I ran (head 62190b9, merge base 48d18a6):
pytest tests/unittests/flows/llm_flows/tools: 310 passed on head.tests/unittests/flows/llm_flowsplustests/unittests/sessions/test_state.py: 1182 passed.- The new test
test_parallel_dict_write_that_removes_a_key_replaces_the_union: passes on head, fails withsrc/reverted to the merge base. So it does pin the cart fix. - A probe with three parallel calls of one read-modify-write tool, on head and on the base.
Finding (blocking): an independent write is dropped when the latest writer read a different call's write.
_apply_latest_state_writes (_batch_executor.py:153-156) deletes the merged value if the latest writer read any other call's write of the key, then keeps only the latest value. That is right when the latest writer has seen every other write. It is wrong when a third call wrote the key without the latest writer ever seeing it.
Probe tool, save_draft(draft_id), does drafts = dict(state.get('drafts', {})); drafts[draft_id] = 'body'; state['drafts'] = drafts, with awaits between the read and the write. Interleaving:
- a and b each read before the other wrote. a writes
{'a'}, b writes{'b'}. - c reads after b's write and writes last:
{'b', 'c'}.
Nobody removed a key. Stored drafts:
- base:
{'a': 'body', 'b': 'body', 'c': 'body'} - head:
{'b': 'body', 'c': 'body'}. a's draft is lost.
On head, c's _state_reads holds b's write, so line 153 matches and the union, including a's key, is thrown away. This differs from the case the description lists as not covered. There, the latest writer writes an unrelated value. Here, c's value builds on b's exactly as intended, and the loss comes from a write c never saw. Read-modify-write is the usual way to write such a tool, so this is likely to happen whenever three or more parallel calls touch one dict key.
Suggested direction: build the result from the latest value plus the writes the latest writer did not read, instead of replacing the whole key. For example, deep_merge_dicts the unread writes in call order, then apply latest over them. The cart test still passes under that rule (the add write was read, so it is excluded) and this probe keeps a.
Minor (nit): State._read_values (state.py:119) holds a reference to every dict or list read through that State for its lifetime, including on callback and invocation contexts that never take part in a parallel batch. A comment saying why the values, not just ids, are kept (to stop id reuse) would help, since it looks like a leak at first read.
What I did not check: live mode, reads through to_dict() (the description already notes these are untracked), and lists.
Probe and raw output: probes/probe_7703.py, review_7703_evidence.txt.
Two parallel tool calls wrote one dict-valued state key. The later call read the earlier call's value and removed a key. The tools saw {'apple': 1} and the session stored {'apple': 1, 'pear': 2} (google#7661).
_apply_latest_state_writes (_batch_executor.py:144) deep-merged the session's latest value into the union of all writes, so a key present in any write survived. Blind writes depend on that union: test_parallel_writes_of_separate_dict_keys_are_all_kept (d42a3be) expects {'a': 'body', 'b': 'body'} from two calls that never read the key. The two cases differ in one fact: whether the latest writer read another call's write before writing.
State now records the dict and list objects each read returns. _build_response_event passes them to the response event. When the latest writer read another call's write of a key, _apply_latest_state_writes (_batch_executor.py:154-159) merges only the writes it did not read, in call order, and applies the latest write over them. Otherwise the union stays.
A first version replaced the whole key in that case. Three calls of a read-modify-write tool lost a write: a and b read before either wrote, c read after b wrote and wrote last. The stored drafts were {b, c} instead of {a, b, c}. test_parallel_dict_write_keeps_independent_write_the_latest_did_not_read pins this: it fails with that version (1 failed, 8 deselected) and passes here.
Not covered: reads outside State.__getitem__ (to_dict, session.state) are not tracked and keep the union. A call that reads another call's write and then writes an unrelated value replaces the writes it read.
Evidence: tests/unittests/flows and tests/unittests/sessions/test_state.py: 1183 passed. test_functions_parallel.py: 9 passed.
Fixes google#7661
62190b9 to
43cda38
Compare
|
Thanks for the probe.
Not covered: reads through |
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
2. Or, if no issue exists, describe the change:
Problem:
Two parallel tool calls wrote the dict-valued state key
cart.add_itemswrote{'apple': 1, 'pear': 2}.remove_pearread that value and wrote{'apple': 1}. The tools saw{'apple': 1}. The session stored{'apple': 1, 'pear': 2}. The issue repro prints both values._apply_latest_state_writes(src/google/adk/flows/llm_flows/tools/_batch_executor.py:144onmain) callsdeep_merge_dicts(merged_delta, latest_writes). That merges the session's latest value into the union of all writes, so a key present in any write survives.state['drafts'] = {draft_id: 'body'}without reading the key. The union{'a': 'body', 'b': 'body'}is intended.test_parallel_writes_of_separate_dict_keys_are_all_keptasserts it. It comes from d42a3be ("Part of Parallel tool calls: state_delta list values silently overwritten during merge #5190"), whose message says the merge re-applies the session's latest value for list and dict keys that several calls wrote.Both cases show the same shape to the merge: the latest value is a dict that lacks a key another call wrote. They differ in whether the latest writer read the other call's write first. I kept the drafts test unchanged and used that difference.
Solution:
State.__getitem__records each dict or list a read returns, by key andid, in a privateState._read_values(sessions/state.py)._build_response_eventsetsevent._state_readsto that record (tools/_caller.py)._apply_latest_state_writesfinds the event whose write is the session's latest value. If that event read another call's write of the key, the latest write replaces the merged value. Otherwisedeep_merge_dictskeeps the union (tools/_batch_executor.py).Diff: 4 files, +45/-3, 24 source lines added outside tests. No public method, signature or export changed.
Not covered:
State.__getitem__(state.to_dict(),session.state) are not tracked. A call that reads this way keeps the union, as before.deep_merge_dictsalready overwrote them.Alternatives considered:
{'a': 1, 'updated': True}with{'b': 1, 'updated': True}.Open question: The seam is a private attribute on
Event, set in_caller.py. Is there a place you prefer for tracking which call read which write?Testing Plan
Unit Tests:
Added
test_parallel_dict_write_that_removes_a_key_replaces_the_uniontotests/unittests/flows/llm_flows/tools/test_functions_parallel.py. It runsadd_items/remove_pear-shaped calls, ordered withasyncio.Event, and asserts the storedcartis{'apple': 1}.Before (
main's_batch_executor.py):After:
Manual End-to-End (E2E) Tests:
The issue's repro script (
repro_issue.py, verbatim from #7661, fakeBaseLlm+InMemoryRunner, no API key) run withPYTHONPATH=src uv run python repro_issue.py. Before:cart persisted in the session: {'apple': 1, 'pear': 2}. After:cart persisted in the session: {'apple': 1}. Logs are in the Unit Tests block above. Noadk webrun: the change has no UI surface.Checklist
Additional context
Tested on macOS, Python 3.12,
uv. The drafts test from d42a3be stays unchanged and passes. CLA signed.Review follow-up (2026-10-24)
After the review on this PR, the merge rule changed.
_apply_latest_state_writesnow builds the stored value from the writes the latest writer did not read, merged in call order, with the latest write applied on top. A write that the latest writer read is replaced by the latest value; a write it never saw is kept. The cart test still passes (the add write was read), and the new testtest_parallel_dict_write_keeps_independent_write_the_latest_did_not_readintests/unittests/flows/llm_flows/tools/test_functions_parallel.pypins the three-call case where an unread write used to be dropped.State._read_valueshas a comment that explains why the values are kept.Checks after the change:
uv run pytest tests/unittests/flows tests/unittests/sessions/test_state.pypasses (1183 tests), anduv run pre-commit runon the touched files passes.