Python: share todo mutation locks across providers on one store

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,

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):