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

🔴 Required Information

Describe the Bug:
When parallel tool calls write the same dict-valued state key, and one call removes keys from the dict another call wrote, the removed keys come back in the persisted session state.

In the snippet, add_items sets state["cart"] = {"apple": 1, "pear": 2}. remove_pear runs after it in the same parallel batch, reads that cart, removes pear and writes {"apple": 1} back. While the batch runs, the tools see {"apple": 1}. The merged function-response event, and therefore the session, still stores {"apple": 1, "pear": 2}.

Root cause: _apply_latest_state_writes in src/google/adk/flows/llm_flows/tools/_batch_executor.py finds that remove_pear's write is the latest value of cart. It then calls deep_merge_dicts(merged_delta, latest_writes), which merges that value into the union of both writes instead of replacing it, so pear survives. Lists are fine, because deep_merge_dicts overwrites lists. This is the dict side of #5190.

Steps to Reproduce:

1. pip install google-adk==2.11.0
2. Run the snippet below. It uses a fake model, so no API key is needed.

Expected Behavior:

cart seen by the tools after remove_pear: {'apple': 1}
cart persisted in the session: {'apple': 1}

Observed Behavior:

cart seen by the tools after remove_pear: {'apple': 1}
cart persisted in the session: {'apple': 1, 'pear': 2}

Environment Details:

• ADK Library Version: 2.11.0 (also reproduced on main at d324744)
• Desktop OS: macOS
• Python Version: 3.12.13

Model Information:

• Are you using LiteLLM: No
• Which model is being used: N/A (fake BaseLlm in the snippet)

🟡 Optional Information

Minimal Reproduction Code:

import asyncio
from typing import AsyncGenerator

from google.adk.agents.llm_agent import LlmAgent
from google.adk.models.base_llm import BaseLlm
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.runners import InMemoryRunner
from google.adk.tools.tool_context import ToolContext
from google.genai import types


class TwoParallelCalls(BaseLlm):
  """Asks for add_items and remove_pear in one turn, then says done."""

  model: str = "fake"

  async def generate_content_async(
      self, llm_request: LlmRequest, stream: bool = False
  ) -> AsyncGenerator[LlmResponse, None]:
    if llm_request.contents[-1].parts[0].function_response:
      yield LlmResponse(content=types.ModelContent("done"))
      return
    yield LlmResponse(
        content=types.Content(
            role="model",
            parts=[
                types.Part.from_function_call(name="add_items", args={}),
                types.Part.from_function_call(name="remove_pear", args={}),
            ],
        )
    )


async def add_items(tool_context: ToolContext) -> str:
  tool_context.state["cart"] = {"apple": 1, "pear": 2}
  return "ok"


async def remove_pear(tool_context: ToolContext) -> str:
  await asyncio.sleep(0.01)  # runs after add_items has written the cart
  cart = dict(tool_context.state["cart"])
  cart.pop("pear")
  tool_context.state["cart"] = cart
  print("cart seen by the tools after remove_pear:", tool_context.state["cart"])
  return "ok"


async def main():
  agent = LlmAgent(
      name="root", model=TwoParallelCalls(), tools=[add_items, remove_pear]
  )
  runner = InMemoryRunner(agent=agent, app_name="app")
  session = await runner.session_service.create_session(
      app_name="app", user_id="u"
  )
  async for _ in runner.run_async(
      user_id="u", session_id=session.id, new_message=types.UserContent("go")
  ):
    pass
  session = await runner.session_service.get_session(
      app_name="app", user_id="u", session_id=session.id
  )
  print("cart persisted in the session:", session.state["cart"])


asyncio.run(main())

How often has this issue occurred?:

• Always (100%)

Proposed fix (looking for guidance before I open a PR)

The simple fix is to replace the key with the latest value instead of deep-merging it. That breaks an existing, intended behaviour. test_parallel_writes_of_separate_dict_keys_are_all_kept (tests/unittests/flows/llm_flows/tools/test_functions_parallel.py) has two parallel calls each write state["drafts"] = {draft_id: "body"} without reading it, and expects the persisted value to be the union {"a": "body", "b": "body"}. From the merge's point of view the two cases look the same: in both, the latest value is a dict that lacks a key another call wrote. The difference is whether the latest writer saw the other write before writing.

Options:

1. Track whether each call read the key before writing it (my preference). When a call reads a dict or list key through tool_context.state, record the object it got. In _apply_latest_state_writes, replace the key with the latest value when the latest writer had read another call's write of that key (it built on it, so missing keys were removed on purpose). Otherwise keep today's union. This fixes the cart case and keeps the drafts test passing, because those calls never read drafts. The cost is a little read bookkeeping in State or the tool context.
2. Heuristic: replace when the latest dict shares keys with another call's write. This needs no new tracking and covers both tests above. It guesses, though: two calls that independently write {"a": 1, "updated": True} and {"b": 1, "updated": True} share a key, so "a" would be dropped.
3. Always last-write-wins for dicts, the same as lists. This is the simplest and matches what the tools see during the run. It changes the tested union behaviour for separate writes, so callers that rely on it would need to read-modify-write instead.

Questions:

• Is the union for independent dict writes behaviour you want to keep? If so, does option 1 fit, or is there a place you'd prefer to track this?
• If last-write-wins (option 3) is acceptable, I can do that with a small change and update the existing test.

I can open a PR with a regression test as soon as a direction is chosen.