fix: keep the latest write when parallel tool calls update a state list

Parallel tool calls share the session state, so a call that reads a list after another call wrote it builds on that write, yet the merged event took each key from the last call in the model's response, letting a call that finished first drop the other calls' entries. The merge now re-applies the session's latest value for list and dict keys that several calls wrote, when that value is one of their writes.

Part of #5190

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 987669949

Author: GWeale

src/google/adk/flows/llm_flows/tools/_batch_executor.py

@@ -116,6 +116,33 @@ def merge_parallel_function_response_events(
   return merged_event
 
 
+def _apply_latest_state_writes(
+    merged_event: Event,
+    function_response_events: list[Event],
+    session_state: dict[str, Any],
+) -> None:
+  """Re-applies the latest value of list and dict keys several calls wrote.
+
+  Calls share the session state, so a call that reads a key after another call
+  wrote it builds on that write, which merging in call order alone can drop.
+  """
+  merged_delta = merged_event.actions.state_delta
+  latest_writes: dict[str, Any] = {}
+  for key in merged_delta:
+    latest = session_state.get(key)
+    if not isinstance(latest, (dict, list)):
+      continue
+    writes = [
+        event.actions.state_delta[key]
+        for event in function_response_events
+        if key in event.actions.state_delta
+    ]
+    # State stores the same object in the session and in the call's delta.
+    if len(writes) > 1 and any(write is latest for write in writes):
+      latest_writes[key] = latest
+  deep_merge_dicts(merged_delta, latest_writes)
+
+
 def _merge_and_trace_function_response_events(
     invocation_context: InvocationContext,
     function_response_events: list[Event],
@@ -124,6 +151,9 @@ def _merge_and_trace_function_response_events(
   merged_event = merge_parallel_function_response_events(
       function_response_events
   )
+  _apply_latest_state_writes(
+      merged_event, function_response_events, invocation_context.session.state
+  )
 
   # this is needed for debug traces of parallel calls
   # individual response with tool.name is traced in __build_response_event

tests/unittests/flows/llm_flows/tools/test_functions_parallel.py

@@ -12,6 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable
+from collections.abc import Callable
+import copy
+from typing import Any
+
 from google.adk.agents.llm_agent import Agent
 from google.adk.events.event_actions import EventActions
 from google.adk.tools.tool_context import ToolContext
@@ -21,6 +29,38 @@
 from .... import testing_utils
 
 
+async def _run_parallel_calls(
+    tool: Callable[..., Awaitable[None]],
+    args_list: list[dict[str, str]],
+    initial_state: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+  """Runs one model turn of parallel calls to `tool`, returns the stored state."""
+  function_calls = [
+      types.Part.from_function_call(name=tool.__name__, args=args)
+      for args in args_list
+  ]
+  agent = Agent(
+      name='root_agent',
+      model=testing_utils.MockModel.create(responses=[function_calls, 'done']),
+      tools=[tool],
+  )
+  runner = testing_utils.TestInMemoryRunner(agent)
+  session = await runner.session_service.create_session(
+      app_name=runner.app_name, user_id='test_user', state=initial_state
+  )
+  async for _ in runner.run_async(
+      user_id='test_user',
+      session_id=session.id,
+      new_message=testing_utils.get_user_content('test'),
+  ):
+    pass
+  stored = await runner.session_service.get_session(
+      app_name=runner.app_name, user_id='test_user', session_id=session.id
+  )
+  assert stored is not None
+  return stored.state
+
+
 @pytest.mark.asyncio
 async def test_parallel_function_calls_with_state_change():
   function_calls = [
@@ -105,3 +145,89 @@ async def transfer_to_agent(
       },
       transfer_to_agent='test_sub_agent',
   )
+
+
+@pytest.mark.asyncio
+async def test_parallel_appends_finishing_out_of_order_keep_every_item():
+  b_written = asyncio.Event()
+
+  async def append_item(item: str, tool_context: ToolContext) -> None:
+    if item == 'a':
+      await asyncio.wait_for(b_written.wait(), timeout=5)
+    tool_context.state['items'] = tool_context.state.get('items', []) + [item]
+    b_written.set()
+
+  state = await _run_parallel_calls(append_item, [{'item': 'a'}, {'item': 'b'}])
+
+  assert state['items'] == ['b', 'a']
+
+
+@pytest.mark.asyncio
+async def test_parallel_appends_finishing_in_order_store_no_duplicates():
+  async def append_item(item: str, tool_context: ToolContext) -> None:
+    tool_context.state['items'] = tool_context.state.get('items', []) + [item]
+
+  state = await _run_parallel_calls(append_item, [{'item': 'a'}, {'item': 'b'}])
+
+  assert state['items'] == ['a', 'b']
+
+
+@pytest.mark.asyncio
+async def test_parallel_snapshot_writes_keep_the_latest_nested_list():
+  b_written = asyncio.Event()
+
+  async def add_label(label: str, tool_context: ToolContext) -> None:
+    if label == 'a':
+      await asyncio.wait_for(b_written.wait(), timeout=5)
+    doc = copy.deepcopy(tool_context.state['doc'])
+    doc['entities'][0]['labels'].append(label)
+    tool_context.state['doc'] = doc
+    b_written.set()
+
+  state = await _run_parallel_calls(
+      add_label,
+      [{'label': 'a'}, {'label': 'b'}],
+      initial_state={'doc': {'entities': [{'name': 'e1', 'labels': []}]}},
+  )
+
+  assert state['doc'] == {'entities': [{'name': 'e1', 'labels': ['b', 'a']}]}
+
+
+@pytest.mark.asyncio
+async def test_parallel_writes_of_separate_dict_keys_are_all_kept():
+  async def save_draft(draft_id: str, tool_context: ToolContext) -> None:
+    tool_context.state['drafts'] = {draft_id: 'body'}
+
+  state = await _run_parallel_calls(
+      save_draft, [{'draft_id': 'a'}, {'draft_id': 'b'}]
+  )
+
+  assert state['drafts'] == {'a': 'body', 'b': 'body'}
+
+
+@pytest.mark.asyncio
+async def test_parallel_list_writes_bypassing_state_keep_call_order():
+  async def set_items(item: str, tool_context: ToolContext) -> None:
+    tool_context.actions.state_delta['items'] = [item]
+
+  state = await _run_parallel_calls(
+      set_items,
+      [{'item': 'a'}, {'item': 'b'}],
+      initial_state={'items': ['old']},
+  )
+
+  assert state['items'] == ['b']
+
+
+@pytest.mark.asyncio
+async def test_parallel_scalar_writes_bypassing_state_keep_call_order():
+  async def set_done(done: str, tool_context: ToolContext) -> None:
+    tool_context.actions.state_delta['done'] = done == 'yes'
+
+  state = await _run_parallel_calls(
+      set_done,
+      [{'done': 'no'}, {'done': 'yes'}],
+      initial_state={'done': False},
+  )
+
+  assert state['done'] is True