ToolContext race

[jf-li00]

Environment

• SDK version: 0.1.5
• Python version: (e.g., 3.13.1)
• OS / platform: (e.g., macOS ARM64, Linux x86_64)
• Install method: pip install google-antigravity

Description

ToolContext provides a per-conversation key/value store (get_state / set_state) backed by a plain dict (tool_context.py:63). The docstring states one ToolContext is "created per session and shared across all tools" (tool_context.py:47).

ToolRunner.process_tool_calls executes a batch of tool calls concurrently via asyncio.gather (tool_runner.py:323), and synchronous tools are dispatched to real OS threads through asyncio.to_thread (tool_runner.py:232).

Because a single ToolContext is injected into every tool that requests one (tool_runner.py:256-258), concurrent tool calls — including multiple synchronous tools running simultaneously in worker threads — read and write the same _state dict with no lock. A tool that does ctx.set_state("k", ctx.get_state("k", 0) + 1) loses updates when two such calls run at once. The GIL does not make the compound read-modify-write atomic.

process_tool_calls warns about ordering of side effects (tool_runner.py:294-295) but is silent on atomicity of the shared state store, so callers have no signal that they must take their own lock.

Steps to Reproduce

1. Save the following as poc_toolcontext_race.py:

"""PoC: ToolContext shared-state race in process_tool_calls."""

import asyncio
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from unittest import mock

from absl.testing import absltest

from google.antigravity import types as sdk_types
from google.antigravity.conversation import conversation as conversation_module
from google.antigravity.tools import tool_context
from google.antigravity.tools import tool_runner


def _make_ctx() -> tool_context.ToolContext:
  conv = mock.MagicMock(spec=conversation_module.Conversation)
  conv.conversation_id = "poc-conv"
  return tool_context.ToolContext(conv)


class ToolContextRacePoC(absltest.TestCase):

  def test_barrier_all_threads_read_same_value(self):
    """N sync tools all read the same stale value via a barrier, then write."""
    NUM = 20
    ctx = _make_ctx()
    barrier = threading.Barrier(NUM)

    def increment(ctx: tool_context.ToolContext) -> int:
      current = ctx.get_state("counter", 0)
      barrier.wait(timeout=10)
      ctx.set_state("counter", current + 1)
      return current + 1

    runner = tool_runner.ToolRunner([increment])
    runner.set_context(ctx)

    loop = asyncio.new_event_loop()
    executor = ThreadPoolExecutor(max_workers=NUM)
    loop.set_default_executor(executor)
    try:
      calls = [sdk_types.ToolCall(name="increment", args={})] * NUM
      loop.run_until_complete(runner.process_tool_calls(calls))
    finally:
      loop.close()
      executor.shutdown(wait=True)

    final = ctx.get_state("counter", 0)
    self.assertEqual(
        final, 1,
        f"Expected final=1 (all {NUM} read 0 via barrier), got {final}",
    )
    print(
        f"\n[BARRIER] {NUM} concurrent sync tools. "
        f"Expected counter={NUM}, got counter={final}. "
        f"Lost {NUM - final} updates."
    )

  def test_sleep_gap_loses_updates(self):
    """100 sync tools with a 5 ms gap between get and set lose updates."""
    NUM = 100
    ctx = _make_ctx()

    def increment(ctx: tool_context.ToolContext) -> int:
      current = ctx.get_state("counter", 0)
      time.sleep(0.005)
      ctx.set_state("counter", current + 1)
      return current + 1

    runner = tool_runner.ToolRunner([increment])
    runner.set_context(ctx)

    calls = [sdk_types.ToolCall(name="increment", args={})] * NUM
    asyncio.run(runner.process_tool_calls(calls))

    final = ctx.get_state("counter", 0)
    self.assertLess(
        final, NUM,
        f"Expected lost updates (final < {NUM}), got {final}.",
    )
    print(
        f"\n[SLEEP] {NUM} concurrent sync tools with 5 ms gap. "
        f"Expected counter={NUM}, got counter={final}. "
        f"Lost {NUM - final} updates."
    )

  def test_sync_tools_use_real_threads(self):
    """Proves sync tools run on genuine OS threads, not the event loop."""
    main_thread = threading.get_ident()
    seen: set[int] = set()
    lock = threading.Lock()

    def thread_probe(ctx: tool_context.ToolContext) -> int:
      tid = threading.get_ident()
      with lock:
        seen.add(tid)
      return tid

    ctx = _make_ctx()
    runner = tool_runner.ToolRunner([thread_probe])
    runner.set_context(ctx)

    calls = [sdk_types.ToolCall(name="thread_probe", args={})] * 20
    asyncio.run(runner.process_tool_calls(calls))

    self.assertNotIn(main_thread, seen)
    self.assertGreaterEqual(len(seen), 2)
    print(
        f"\n[THREADS] Sync tools ran on {len(seen)} distinct "
        f"worker threads (event-loop thread excluded)."
    )


if __name__ == "__main__":
  absltest.main()

2. Run:

pip install google-antigravity
python poc_toolcontext_race.py

Expected Behavior

When 20 sync tools each call ctx.set_state("counter", ctx.get_state("counter", 0) + 1) concurrently, the final value of counter should be 20. State access on the shared ToolContext should be atomic or the documentation should warn that tools must take their own lock.

Actual Behavior

[BARRIER] 20 concurrent sync tools. Expected counter=20, got counter=1. Lost 19 updates.
[SLEEP] 100 concurrent sync tools with 5 ms gap. Expected counter=100, got counter=4. Lost 96 updates.
[THREADS] Sync tools ran on 5 distinct worker threads (event-loop thread excluded).

The barrier test deterministically shows all 20 threads read counter=0 before any write, so the final value is 1 instead of 20 — 19 updates lost. The sleep test (100 calls with a 5 ms gap) loses 96 updates. The thread test confirms sync tools run on real OS worker threads (5 distinct, none the event-loop thread), which is what makes the race real rather than theoretical.

Additional Context

Affected files:

• google/antigravity/tools/tool_context.py — ToolContext._state is a bare dict; get_state is dict.get, set_state is __setitem__. No lock.
• google/antigravity/tools/tool_runner.py:229-238 — _execute_fn runs sync tools via asyncio.to_thread.
• google/antigravity/tools/tool_runner.py:284-323 — process_tool_calls runs all tool calls concurrently via asyncio.gather.

Suggested fix: Add a threading.Lock to ToolContext and guard get_state / set_state. Alternatively, document that ToolContext state access is not atomic and tools must take their own lock if they perform compound read-modify-write operations across a concurrency boundary.