diff --git a/python/packages/core/agent_framework/_harness/_todo.py b/python/packages/core/agent_framework/_harness/_todo.py index 7aaeeecb2a..b083f6ab0f 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 2b76c1b49c..e8a08ee64c 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: