From 3ca215119bcbe3ac2cf8cd5df4f212df5dedd215 Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:55:14 +0530 Subject: [PATCH] Python: scope the todo mutation lock to the store and session The lock lived on the provider instance, so two providers sharing a todo store could both load the same state before either saved, silently dropping one write, and a session reused across event loops could raise from the loop bound lock. The lock registry is now keyed by store, session, and event loop, mirroring the file access write lock. Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../core/agent_framework/_harness/_todo.py | 44 +++++++---- .../core/tests/core/test_harness_todo.py | 76 +++++++++++++++++-- 2 files changed, 101 insertions(+), 19 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_todo.py b/python/packages/core/agent_framework/_harness/_todo.py index 7aaeeecb2a3..b083f6ab0ff 100644 --- a/python/packages/core/agent_framework/_harness/_todo.py +++ b/python/packages/core/agent_framework/_harness/_todo.py @@ -287,6 +287,33 @@ async def save_state(self, session: AgentSession, items: list[TodoItem], *, next provider_state["next_id"] = _safe_next_id(items, next_id) +_TODO_SESSION_LOCKS: weakref.WeakKeyDictionary[ + AgentSession, dict[tuple[int, asyncio.AbstractEventLoop], asyncio.Lock] +] = weakref.WeakKeyDictionary() + + +def _todo_mutation_lock(store: TodoStore, session: AgentSession) -> asyncio.Lock: + """Return the lock that serializes read-modify-write todo operations for one store and session. + + The lock belongs to the store and session rather than to one provider, so + two providers sharing a todo store cannot both load the same state before + either saves it back and silently drop one of the edits. Locks are also per + event loop, because an :class:`asyncio.Lock` that has had a waiter on one + loop raises when a caller has to wait for it on another. The outer map is + weak so locks for garbage collected sessions are evicted rather than + accumulating in long-running services; the key holds ``id(store)`` so a + store need not be hashable, mirroring ``_store_write_lock`` in the file + access layer. + """ + per_session = _TODO_SESSION_LOCKS.setdefault(session, {}) + key = (id(store), asyncio.get_running_loop()) + lock = per_session.get(key) + if lock is None: + lock = asyncio.Lock() + per_session[key] = lock + return lock + + @experimental(feature_id=ExperimentalFeature.HARNESS) class TodoFileStore(TodoStore): """Store todo state in one JSON file per session and source ID. @@ -446,17 +473,6 @@ def __init__( super().__init__(source_id) self.instructions = instructions or DEFAULT_TODO_INSTRUCTIONS self.store = store or TodoSessionStore() - # WeakKeyDictionary so per-session locks are evicted automatically when the session is GC'd - # rather than accumulating in long-running services that create many sessions. - self._mutation_locks: weakref.WeakKeyDictionary[AgentSession, asyncio.Lock] = weakref.WeakKeyDictionary() - - def _mutation_lock(self, session: AgentSession) -> asyncio.Lock: - """Return the per-session lock for read-modify-write todo operations.""" - lock = self._mutation_locks.get(session) - if lock is None: - lock = asyncio.Lock() - self._mutation_locks[session] = lock - return lock async def before_run( self, @@ -476,7 +492,7 @@ async def todos_add(todos: list[_TodoAddItemSchema]) -> str: if not todos: raise ValueError("todos must contain at least one item.") - async with self._mutation_lock(session): + async with _todo_mutation_lock(self.store, session): existing_items, next_id = await self.store.load_state(session, source_id=self.source_id) created_items: list[TodoItem] = [] for raw_todo in todos: @@ -505,7 +521,7 @@ async def todos_complete(items: list[_TodoCompleteItemSchema]) -> str: parsed = [_coerce_todo_complete_input(entry) for entry in items] ids = [entry.id for entry in parsed] - async with self._mutation_lock(session): + async with _todo_mutation_lock(self.store, session): existing_items, next_id = await self.store.load_state(session, source_id=self.source_id) id_set = set(ids) completed_count = 0 @@ -534,7 +550,7 @@ async def todos_remove(ids: list[int]) -> str: if not ids: raise ValueError("ids must contain at least one todo ID.") - async with self._mutation_lock(session): + async with _todo_mutation_lock(self.store, session): items, next_id = await self.store.load_state(session, source_id=self.source_id) remaining_items = [item for item in items if item.id not in set(ids)] removed_count = len(items) - len(remaining_items) diff --git a/python/packages/core/tests/core/test_harness_todo.py b/python/packages/core/tests/core/test_harness_todo.py index 2b76c1b49ce..e8a08ee64ce 100644 --- a/python/packages/core/tests/core/test_harness_todo.py +++ b/python/packages/core/tests/core/test_harness_todo.py @@ -6,6 +6,7 @@ import json import os from pathlib import Path +from typing import Any, cast import pytest @@ -200,17 +201,20 @@ async def test_todo_stores_clamp_next_id_to_avoid_collisions(tmp_path: Path) -> async def test_todo_provider_evicts_locks_when_session_is_garbage_collected() -> None: - """The provider should not retain mutation locks for sessions that have been GC'd.""" + """Mutation locks should not be retained for sessions that have been GC'd.""" import gc - provider = TodoProvider() + from agent_framework._harness._todo import _TODO_SESSION_LOCKS, _todo_mutation_lock # pyright: ignore[reportPrivateUsage] + + store = TodoSessionStore() session = AgentSession(session_id="session-1") - provider._mutation_lock(session) # pyright: ignore[reportPrivateUsage] - assert len(provider._mutation_locks) == 1 # pyright: ignore[reportPrivateUsage] + _todo_mutation_lock(store, session) # pyright: ignore[reportPrivateUsage] + _todo_mutation_lock(store, session) # pyright: ignore[reportPrivateUsage] + assert len(_TODO_SESSION_LOCKS) == 1 # pyright: ignore[reportPrivateUsage] del session gc.collect() - assert len(provider._mutation_locks) == 0 # pyright: ignore[reportPrivateUsage] + assert len(_TODO_SESSION_LOCKS) == 0 # pyright: ignore[reportPrivateUsage] async def test_todo_file_store_rejects_session_path_traversal(tmp_path: Path) -> None: @@ -321,6 +325,68 @@ async def test_todo_provider_tools_manage_session_state( ] +class _ToolCapturingContext: + """Minimal context that records what the todo provider injects.""" + + def __init__(self) -> None: + self.tools: list[object] = [] + + def extend_instructions(self, source_id: str, instructions: list[str]) -> None: + del source_id, instructions + + def extend_tools(self, source_id: str, tools: list[object]) -> None: + del source_id + self.tools.extend(tools) + + def extend_messages(self, source_id: str, messages: list[object]) -> None: + del source_id, messages + + +async def test_todo_provider_serializes_mutations_across_providers_sharing_a_store(tmp_path: Path) -> None: + """Two providers over one store must serialize like one provider does. + + The mutation lock belongs to the store and session rather than to one + provider, so concurrent tool calls from agents that share a todo store + cannot both load the same state before either saves and silently drop one + of the writes. + """ + store = TodoFileStore(tmp_path) + session = AgentSession(session_id="session-1") + + tools_per_provider: list[dict[str, object]] = [] + for _ in range(2): + context = _ToolCapturingContext() + await TodoProvider(store=store).before_run(agent=None, session=session, context=context, state={}) # type: ignore[arg-type] + tools_per_provider.append({tool.name: tool for tool in cast(list[Any], context.tools) if hasattr(tool, "name")}) # type: ignore[attr-defined] + + # Gate loads only after tool capture: before_run itself loads the todo + # state to inject it into the run. + gate = asyncio.Event() + original_load = store.load_state + + async def gated_load(session: AgentSession, *, source_id: str) -> tuple[list[TodoItem], int]: + await gate.wait() + return await original_load(session, source_id=source_id) + + store.load_state = gated_load + + async def add(provider_tools: dict[str, object], title: str) -> None: + add_tool = cast(Any, provider_tools["todos_add"]) + await add_tool.invoke(arguments={"todos": [{"title": title}]}) + + first = asyncio.create_task(add(tools_per_provider[0], "From provider A")) + second = asyncio.create_task(add(tools_per_provider[1], "From provider B")) + await asyncio.sleep(0) + gate.set() + await asyncio.wait_for(asyncio.gather(first, second), timeout=10) + + get_all_tool = cast(Any, tools_per_provider[0]["todos_get_all"]) + get_all_result = await get_all_tool.invoke() + payload = json.loads(get_all_result[0].text) + assert sorted(item["title"] for item in payload) == ["From provider A", "From provider B"] + assert sorted(item["id"] for item in payload) == [1, 2] + + async def test_todo_provider_serializes_concurrent_mutations( chat_client_base: SupportsChatGetResponse, ) -> None: