From 1d322911c3e82ac273f1a9ebae81939e76d5ba4f Mon Sep 17 00:00:00 2001 From: ninglinLiu Date: Fri, 30 Oct 2026 12:23:39 +0800 Subject: [PATCH] fix(python): share todo mutation locks across providers on one store TodoProvider kept per-session locks on the provider instance, so two providers over the same TodoStore and session could each load the same next_id, append colliding items, and drop the other save. Move the lock to a store/session/source_id map keyed per event loop, matching the file-access _store_write_lock pattern. Fixes #9868 --- .../core/agent_framework/_harness/_todo.py | 42 +++++++++--- .../core/tests/core/test_harness_todo.py | 65 ++++++++++++++++--- 2 files changed, 89 insertions(+), 18 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_todo.py b/python/packages/core/agent_framework/_harness/_todo.py index 7aaeeecb2a3..6744fdeaa8c 100644 --- a/python/packages/core/agent_framework/_harness/_todo.py +++ b/python/packages/core/agent_framework/_harness/_todo.py @@ -411,6 +411,37 @@ def _save_state_sync(state_path: Path, payload: str) -> None: _atomic_write_text(state_path, payload) + +# Locks handed out by ``_todo_mutation_lock``. Values are held weakly, so an entry lasts +# only while a mutation holds or awaits its lock, and the map never grows with the number +# of sessions ever mutated. +_TODO_MUTATION_LOCKS: weakref.WeakValueDictionary[ + tuple[int, asyncio.AbstractEventLoop, int, str], asyncio.Lock +] = weakref.WeakValueDictionary() + + +def _todo_mutation_lock(store: TodoStore, session: AgentSession, source_id: str) -> asyncio.Lock: + """Return the lock that serializes todo mutations for one store/session/source_id. + + Mutating tools hold it across their read-modify-write. It belongs to the store + instance rather than to one provider, so two providers sharing a store cannot both + load the same state before either saves and silently drop one of the edits. + + Locks are 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 key holds + ``id(store)`` and ``id(session)`` rather than the objects, so neither needs to be + hashable and is never kept alive by this map; those ids cannot be reused while their + lock is in use, because whoever holds or awaits the lock still references the store + and session. + """ + key = (id(store), asyncio.get_running_loop(), id(session), source_id) + lock = _TODO_MUTATION_LOCKS.get(key) + if lock is None: + lock = asyncio.Lock() + _TODO_MUTATION_LOCKS[key] = lock + return lock + + class TodoProvider(ContextProvider): """Provide todo management tools and instructions to an agent. @@ -446,17 +477,10 @@ 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 + """Return the store-level lock for read-modify-write todo operations.""" + return _todo_mutation_lock(self.store, session, self.source_id) async def before_run( self, diff --git a/python/packages/core/tests/core/test_harness_todo.py b/python/packages/core/tests/core/test_harness_todo.py index 2b76c1b49ce..5e14d88190b 100644 --- a/python/packages/core/tests/core/test_harness_todo.py +++ b/python/packages/core/tests/core/test_harness_todo.py @@ -199,18 +199,27 @@ async def test_todo_stores_clamp_next_id_to_avoid_collisions(tmp_path: Path) -> assert items == [TodoItem(id=7, title="Seeded")] -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.""" - import gc +async def test_todo_mutation_locks_are_shared_across_providers() -> None: + """Two providers over one store and session should share the mutation lock.""" + from agent_framework._harness._todo import ( # pyright: ignore[reportPrivateUsage] + _TODO_MUTATION_LOCKS, + _todo_mutation_lock, + ) - provider = TodoProvider() + store = TodoSessionStore() + provider_a = TodoProvider(store=store) + provider_b = TodoProvider(store=store) + provider_other_source = TodoProvider(source_id="other", store=store) session = AgentSession(session_id="session-1") - provider._mutation_lock(session) # pyright: ignore[reportPrivateUsage] - assert len(provider._mutation_locks) == 1 # pyright: ignore[reportPrivateUsage] - del session - gc.collect() - assert len(provider._mutation_locks) == 0 # pyright: ignore[reportPrivateUsage] + lock_a = provider_a._mutation_lock(session) # pyright: ignore[reportPrivateUsage] + lock_b = provider_b._mutation_lock(session) # pyright: ignore[reportPrivateUsage] + other_lock = provider_other_source._mutation_lock(session) # pyright: ignore[reportPrivateUsage] + + assert lock_a is lock_b + assert other_lock is not lock_a + assert _todo_mutation_lock(store, session, "todo") is lock_a + assert (id(store), asyncio.get_running_loop(), id(session), "todo") in _TODO_MUTATION_LOCKS async def test_todo_file_store_rejects_session_path_traversal(tmp_path: Path) -> None: @@ -368,6 +377,44 @@ async def test_todo_provider_serializes_concurrent_mutations( assert {item["id"] for item in payload if item["is_complete"]} == {1, 2, 3, 4, 5} + +async def test_todo_providers_sharing_file_store_serialize_mutations( + tmp_path: Path, chat_client_base: SupportsChatGetResponse +) -> None: + """Two providers over one file store must not drop items under concurrent adds.""" + session = AgentSession(session_id="session-shared") + store = TodoFileStore(tmp_path) + provider_a = TodoProvider(store=store) + provider_b = TodoProvider(store=store) + agent_a = Agent(client=chat_client_base, context_providers=[provider_a]) + agent_b = Agent(client=chat_client_base, context_providers=[provider_b]) + + _, options_a = await agent_a._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Track A"])], + ) + _, options_b = await agent_b._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Track B"])], + ) + add_a = _tool_by_name(options_a["tools"], "todos_add") # type: ignore[index] + add_b = _tool_by_name(options_b["tools"], "todos_add") # type: ignore[index] + get_all = _tool_by_name(options_a["tools"], "todos_get_all") # type: ignore[index] + + await asyncio.gather( + add_a.invoke(arguments={"todos": [{"title": "From A"}]}), # type: ignore[attr-defined] + add_b.invoke(arguments={"todos": [{"title": "From B"}]}), # type: ignore[attr-defined] + ) + + get_all_result = await get_all.invoke() # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + payload = json.loads(get_all_result[0].text) + titles = {item["title"] for item in payload} + ids = [item["id"] for item in payload] + assert titles == {"From A", "From B"} + assert sorted(ids) == [1, 2] + assert len(ids) == len(set(ids)) + + def test_todo_harness_graduated_classes_are_not_experimental() -> None: """Graduated todo harness types should carry no experimental metadata; TodoFileStore stays experimental.""" for graduated in (TodoStore, TodoItem, TodoInput, TodoSessionStore, TodoProvider):