Skip to content

fix(flows): persist the dict key removal of a later parallel tool call - #7703

Open
danieldutu wants to merge 1 commit into
google:mainfrom
danieldutu:fix/flows-parallel-dict-key-removal
Open

danieldutu wants to merge 1 commit into
google:mainfrom
danieldutu:fix/flows-parallel-dict-key-removal

Conversation

@danieldutu

@danieldutu danieldutu commented Oct 23, 2026 •

Copy link
Copy Markdown

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_items wrote {'apple': 1, 'pear': 2}. remove_pear read 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:144 on main) calls deep_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.

  • Cart case: the latest writer read another call's write, then removed a key. The removal is intended, so the latest write must replace the union.
  • Drafts case: two calls each write 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_kept asserts 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 and id, in a private State._read_values (sessions/state.py).
  • _build_response_event sets event._state_reads to that record (tools/_caller.py).
  • _apply_latest_state_writes finds 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. Otherwise deep_merge_dicts keeps 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:

  • Reads outside State.__getitem__ (state.to_dict(), session.state) are not tracked. A call that reads this way keeps the union, as before.
  • A call that reads another call's write and then writes an unrelated value replaces the union and drops the other call's keys.
  • Lists are unchanged: deep_merge_dicts already overwrote them.
  • Not tested: platforms other than macOS (Python 3.12), live mode.

Alternatives considered:

  • Replace the key with the latest write always: fails the drafts test and reverses the intent of d42a3be.
  • Replace when the latest dict shares a key with another write: mishandles {'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:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Added test_parallel_dict_write_that_removes_a_key_replaces_the_union to tests/unittests/flows/llm_flows/tools/test_functions_parallel.py. It runs add_items/remove_pear-shaped calls, ordered with asyncio.Event, and asserts the stored cart is {'apple': 1}.

Before (main's _batch_executor.py):

$ uv run pytest tests/unittests/flows/llm_flows/tools/test_functions_parallel.py -q
FAILED ...::test_parallel_dict_write_that_removes_a_key_replaces_the_union
1 failed, 7 passed in 2.03s
$ PYTHONPATH=src uv run python repro_issue.py
cart seen by the tools after remove_pear: {'apple': 1}
cart persisted in the session: {'apple': 1, 'pear': 2}

After:

$ uv run pytest tests/unittests/flows/llm_flows/tools/test_functions_parallel.py -q
8 passed in 1.46s
$ PYTHONPATH=src uv run python repro_issue.py
cart seen by the tools after remove_pear: {'apple': 1}
cart persisted in the session: {'apple': 1}
$ uv run pytest tests/unittests/flows -q -n auto
1152 passed, 145 warnings in 21.04s
$ uv run pytest tests/unittests/flows/llm_flows/tools/test_batch_executor.py tests/unittests/flows/llm_flows/tools/test_functions_parallel.py tests/unittests/sessions/test_state.py -q
63 passed in 3.22s
$ uv run pre-commit run --files <the 4 touched files>
ruff, isort, pyink, addlicense, compliance checks, codespell: Passed

Manual End-to-End (E2E) Tests:

The issue's repro script (repro_issue.py, verbatim from #7661, fake BaseLlm + InMemoryRunner, no API key) run with PYTHONPATH=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. No adk web run: the change has no UI surface.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

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_writes now 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 test test_parallel_dict_write_keeps_independent_write_the_latest_did_not_read in tests/unittests/flows/llm_flows/tools/test_functions_parallel.py pins the three-call case where an unread write used to be dropped. State._read_values has a comment that explains why the values are kept.

Checks after the change: uv run pytest tests/unittests/flows tests/unittests/sessions/test_state.py passes (1183 tests), and uv run pre-commit run on the touched files passes.

@Somay-kousis Somay-kousis left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What I ran (head 62190b9, merge base 48d18a6):

  • pytest tests/unittests/flows/llm_flows/tools: 310 passed on head. tests/unittests/flows/llm_flows plus tests/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 with src/ 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
@danieldutu
danieldutu force-pushed the fix/flows-parallel-dict-key-removal branch from 62190b9 to 43cda38 Compare October 24, 2026 15:59
@danieldutu

Copy link
Copy Markdown
Author

Thanks for the probe.

  1. Independent write dropped when the latest writer read a different call's write: reproduced as test_parallel_dict_write_keeps_independent_write_the_latest_did_not_read (test_functions_parallel.py:230). Before: stored drafts {'b','c'}. After: {'a','b','c'}. Change at _batch_executor.py:150-159: the result is the writes the latest writer did not read, merged in call order, with the latest write applied over them, as suggested. The cart test and the maintainer drafts test pass unchanged.
  2. State._read_values comment: added at state.py:119-120. It states the values are kept so an id cannot be reused by a different object.

Not covered: reads through to_dict() and session.state stay untracked. Lists are unchanged. A call that reads one write and then writes an unrelated value replaces the writes it read.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Parallel tool calls: a dict state key removed by a later call comes back in the persisted state

3 participants