fix(mongodb): address Codex P1/P2 review feedback (round 3)

- Replace asyncio.Lock process-wide guard with threading.Lock so
  _get_or_create_init_lock() is safe to call from any event loop
- Replace id(client)-keyed dicts with WeakKeyDictionary so entries are
  pruned when clients are GC'd, preventing stale id reuse from skipping
  index creation on a new client
- Add explicit seq field (monotonic per-session counter via $inc) to
  every message document and sort by seq instead of _id; ObjectId is
  only second-level accurate across processes and not reliably monotonic
- Update FakeAsyncCollection in tests with find_one_and_update + $inc
  support; replace _initialized_keys/_init_locks clears with _init_state

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: alexbevi

src/agents/extensions/memory/mongodb_session.py

@@ -33,6 +33,8 @@
 
 import asyncio
 import json
+import threading
+import weakref
 from typing import Any
 
 try:
@@ -70,14 +72,29 @@ class MongoDBSession(SessionABC):
     Indexes are created once per ``(client, database, sessions_collection,
     messages_collection)`` combination on the first call to any of the
     session protocol methods.  Subsequent calls skip the setup entirely.
+
+    Each message document carries a ``seq`` field — an integer assigned by
+    atomically incrementing a counter on the session metadata document.  This
+    guarantees a strictly monotonic insertion order that is safe across
+    multiple writers and processes, unlike sorting by ``_id`` / ObjectId which
+    is only second-level accurate and non-monotonic across machines.
     """
 
-    # Class-level registry so index creation only runs once per unique key.
-    # The key includes id(client) so that different AsyncMongoClient instances
-    # (e.g. pointing at different clusters) each get their own init pass.
-    _initialized_keys: set[tuple[int, str, str, str]] = set()
-    _init_locks: dict[tuple[int, str, str, str], asyncio.Lock] = {}
-    _init_locks_guard: asyncio.Lock = asyncio.Lock()
+    # Class-level registry so index creation runs only once per unique
+    # (client, database, sessions_collection, messages_collection) combination.
+    #
+    # WeakKeyDictionary lets the client object itself be the outer key so
+    # that the entry is automatically pruned when the client is garbage-
+    # collected.  This prevents id(client) reuse from causing a new client to
+    # skip index creation because it shares an address with a dead one.
+    #
+    # _init_guard is a plain threading.Lock so it is safe to acquire from any
+    # event loop — an asyncio.Lock binds to the loop that first contends on
+    # it and raises RuntimeError when a different loop tries to acquire it.
+    _init_state: weakref.WeakKeyDictionary[
+        Any, dict[tuple[str, str, str], asyncio.Lock | bool]
+    ] = weakref.WeakKeyDictionary()
+    _init_guard: threading.Lock = threading.Lock()
 
     session_settings: SessionSettings | None = None
 
@@ -120,7 +137,9 @@ def __init__(
         self._sessions: AsyncCollection[Any] = db[sessions_collection]
         self._messages: AsyncCollection[Any] = db[messages_collection]
 
-        self._init_key = (id(client), database, sessions_collection, messages_collection)
+        # Key within the per-client mapping — no id() needed because the client
+        # object itself is the outer WeakKeyDictionary key.
+        self._init_sub_key = (database, sessions_collection, messages_collection)
 
     # ------------------------------------------------------------------
     # Convenience constructors
@@ -172,33 +191,59 @@ def from_uri(
     # Index initialisation
     # ------------------------------------------------------------------
 
-    async def _get_init_lock(self) -> asyncio.Lock:
-        """Return (creating if necessary) the per-init-key asyncio Lock."""
-        async with self._init_locks_guard:
-            lock = self._init_locks.get(self._init_key)
-            if lock is None:
+    def _get_or_create_init_lock(self) -> tuple[asyncio.Lock, bool]:
+        """Return (lock, already_done) for this session's (client, sub_key) pair.
+
+        Uses a threading.Lock for the registry mutation so it is safe to call
+        from any event loop without risking a cross-loop RuntimeError.
+        """
+        with self._init_guard:
+            per_client = self._init_state.get(self._client)
+            if per_client is None:
+                per_client = {}
+                self._init_state[self._client] = per_client
+
+            entry = per_client.get(self._init_sub_key)
+            if entry is True:
+                # Already initialised.
+                return asyncio.Lock(), True
+            if entry is None:
                 lock = asyncio.Lock()
-                self._init_locks[self._init_key] = lock
-            return lock
+                per_client[self._init_sub_key] = lock
+                return lock, False
+            # entry is an asyncio.Lock — initialisation is in progress.
+            assert isinstance(entry, asyncio.Lock)
+            return entry, False
+
+    def _mark_init_done(self) -> None:
+        """Record that index creation is complete for this (client, sub_key)."""
+        with self._init_guard:
+            per_client = self._init_state.get(self._client)
+            if per_client is not None:
+                per_client[self._init_sub_key] = True
 
     async def _ensure_indexes(self) -> None:
-        """Create required indexes the first time this key is accessed."""
-        if self._init_key in self._initialized_keys:
+        """Create required indexes the first time this (client, sub_key) is accessed."""
+        lock, done = self._get_or_create_init_lock()
+        if done:
             return
 
-        lock = await self._get_init_lock()
         async with lock:
             # Double-checked locking: another coroutine may have finished first.
-            if self._init_key in self._initialized_keys:
+            _, done = self._get_or_create_init_lock()
+            if done:
                 return
 
             # sessions: unique index on session_id.
             await self._sessions.create_index("session_id", unique=True)
 
             # messages: compound index for efficient per-session retrieval and sorting.
-            await self._messages.create_index([("session_id", 1), ("_id", 1)])
+            # seq provides a strictly monotonic insertion-order tie-breaker that is
+            # reliable across multiple writers (unlike ObjectId which is only
+            # second-level accurate).
+            await self._messages.create_index([("session_id", 1), ("seq", 1)])
 
-            self._initialized_keys.add(self._init_key)
+            self._mark_init_done()
 
     # ------------------------------------------------------------------
     # Serialization helpers
@@ -239,12 +284,12 @@ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
         query = {"session_id": self.session_id}
 
         if session_limit is None:
-            cursor = self._messages.find(query).sort("_id", 1)
+            cursor = self._messages.find(query).sort("seq", 1)
             docs = await cursor.to_list()
         else:
             # Fetch the latest N documents in reverse order, then reverse the
             # list to restore chronological order.
-            cursor = self._messages.find(query).sort("_id", -1).limit(session_limit)
+            cursor = self._messages.find(query).sort("seq", -1).limit(session_limit)
             docs = await cursor.to_list()
             docs.reverse()
 
@@ -269,19 +314,27 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
 
         await self._ensure_indexes()
 
-        # Upsert the session metadata document.
-        await self._sessions.update_one(
+        # Atomically reserve a block of sequence numbers for this batch.
+        # $inc returns the new value, so subtract len(items) to get the first
+        # number in the block.
+        result = await self._sessions.find_one_and_update(
             {"session_id": self.session_id},
-            {"$setOnInsert": {"session_id": self.session_id}},
+            {
+                "$setOnInsert": {"session_id": self.session_id},
+                "$inc": {"_seq": len(items)},
+            },
             upsert=True,
+            return_document=True,
         )
+        next_seq: int = (result["_seq"] if result else len(items)) - len(items)
 
         payload = [
             {
                 "session_id": self.session_id,
+                "seq": next_seq + i,
                 "message_data": await self._serialize_item(item),
             }
-            for item in items
+            for i, item in enumerate(items)
         ]
 
         await self._messages.insert_many(payload, ordered=True)
@@ -296,7 +349,7 @@ async def pop_item(self) -> TResponseInputItem | None:
 
         doc = await self._messages.find_one_and_delete(
             {"session_id": self.session_id},
-            sort=[("_id", -1)],
+            sort=[("seq", -1)],
         )
 
         if doc is None:

tests/extensions/memory/test_mongodb_session.py

@@ -115,19 +115,41 @@ async def insert_many(
                 doc["_id"] = FakeObjectId()
             self._docs[id(doc["_id"])] = dict(doc)
 
-    async def update_one(
+    async def find_one_and_update(
         self,
         query: dict[str, Any],
         update: dict[str, Any],
         upsert: bool = False,
-    ) -> None:
+        return_document: bool = False,
+    ) -> dict[str, Any] | None:
         for doc in self._docs.values():
             if self._matches(doc, query):
-                return  # Exists — $setOnInsert is a no-op on existing docs.
+                # Apply $inc fields.
+                for field, delta in update.get("$inc", {}).items():
+                    doc[field] = doc.get(field, 0) + delta
+                return dict(doc) if return_document else None
         if upsert:
             new_doc: dict[str, Any] = {"_id": FakeObjectId()}
             new_doc.update(update.get("$setOnInsert", {}))
+            for field, delta in update.get("$inc", {}).items():
+                new_doc[field] = new_doc.get(field, 0) + delta
             self._docs[id(new_doc["_id"])] = new_doc
+            return dict(new_doc) if return_document else None
+        return None
+
+    async def update_one(
+        self,
+        query: dict[str, Any],
+        update: dict[str, Any],
+        upsert: bool = False,
+    ) -> None:
+        for doc in self._docs.values():
+            if self._matches(doc, query):
+                return  # Exists — $setOnInsert is a no-op on existing docs.
+        if upsert:
+            new_doc2: dict[str, Any] = {"_id": FakeObjectId()}
+            new_doc2.update(update.get("$setOnInsert", {}))
+            self._docs[id(new_doc2["_id"])] = new_doc2
 
     async def delete_many(self, query: dict[str, Any]) -> None:
         to_remove = [k for k, d in self._docs.items() if self._matches(d, query)]
@@ -235,8 +257,7 @@ def _make_fake_pymongo_modules() -> None:
 def _make_session(session_id: str = "test-session", **kwargs: Any) -> MongoDBSession:
     """Create a MongoDBSession backed by a FakeAsyncMongoClient."""
     client = FakeAsyncMongoClient()
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
     return MongoDBSession(
         session_id,
         client=client,  # type: ignore[arg-type]
@@ -353,8 +374,7 @@ async def test_get_items_limit_exceeds_count(session: MongoDBSession) -> None:
 
 async def test_session_settings_limit_used_as_default() -> None:
     """session_settings.limit is applied when no explicit limit is given."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
     s = MongoDBSession(
         "ls-test",
         client=FakeAsyncMongoClient(),  # type: ignore[arg-type]
@@ -371,8 +391,7 @@ async def test_session_settings_limit_used_as_default() -> None:
 
 async def test_explicit_limit_overrides_session_settings() -> None:
     """An explicit limit passed to get_items must override session_settings.limit."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
     s = MongoDBSession(
         "override-test",
         client=FakeAsyncMongoClient(),  # type: ignore[arg-type]
@@ -394,8 +413,7 @@ async def test_explicit_limit_overrides_session_settings() -> None:
 
 async def test_sessions_are_isolated() -> None:
     """Two sessions with different IDs must not share data."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
     client = FakeAsyncMongoClient()
     s1 = MongoDBSession("alice", client=client, database="agents_test")  # type: ignore[arg-type]
     s2 = MongoDBSession("bob", client=client, database="agents_test")  # type: ignore[arg-type]
@@ -409,8 +427,7 @@ async def test_sessions_are_isolated() -> None:
 
 async def test_clear_does_not_affect_other_sessions() -> None:
     """Clearing one session must leave sibling sessions untouched."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
     client = FakeAsyncMongoClient()
     s1 = MongoDBSession("s1", client=client, database="agents_test")  # type: ignore[arg-type]
     s2 = MongoDBSession("s2", client=client, database="agents_test")  # type: ignore[arg-type]
@@ -529,8 +546,7 @@ async def counting(*args: Any, **kwargs: Any) -> str:
 
 async def test_different_clients_each_run_index_init() -> None:
     """Each distinct AsyncMongoClient gets its own index-creation pass."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
 
     client_a = FakeAsyncMongoClient()
     client_b = FakeAsyncMongoClient()
@@ -585,8 +601,7 @@ async def _fail(*args: Any, **kwargs: Any) -> dict[str, Any]:
 
 async def test_close_external_client_not_closed() -> None:
     """close() must NOT close a client that was injected externally."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
     client = FakeAsyncMongoClient()
     s = MongoDBSession("x", client=client, database="agents_test")  # type: ignore[arg-type]
     assert s._owns_client is False
@@ -597,8 +612,7 @@ async def test_close_external_client_not_closed() -> None:
 
 async def test_close_owned_client_is_closed() -> None:
     """close() must close a client created by from_uri."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
     fake_client = FakeAsyncMongoClient()
     with patch(
         "agents.extensions.memory.mongodb_session.AsyncMongoClient",
@@ -636,8 +650,7 @@ async def test_runner_integration(agent: Agent) -> None:
 
 async def test_runner_session_isolation(agent: Agent) -> None:
     """Two independent sessions must not bleed history into each other."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
     client = FakeAsyncMongoClient()
     s1 = MongoDBSession("user-a", client=client, database="agents_test")  # type: ignore[arg-type]
     s2 = MongoDBSession("user-b", client=client, database="agents_test")  # type: ignore[arg-type]
@@ -659,8 +672,7 @@ async def test_runner_with_session_settings_limit(agent: Agent) -> None:
     """RunConfig.session_settings.limit must cap the history sent to the model."""
     from agents import RunConfig
 
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
     session = MongoDBSession(
         "limit-test",
         client=FakeAsyncMongoClient(),  # type: ignore[arg-type]
@@ -694,8 +706,7 @@ async def test_runner_with_session_settings_limit(agent: Agent) -> None:
 
 async def test_injected_client_receives_append_metadata() -> None:
     """Append_metadata is called on a caller-supplied client."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
     client = FakeAsyncMongoClient()
 
     MongoDBSession("meta-test", client=client, database="agents_test")  # type: ignore[arg-type]
@@ -707,8 +718,7 @@ async def test_injected_client_receives_append_metadata() -> None:
 
 async def test_from_uri_passes_driver_info_to_constructor() -> None:
     """driver=_DRIVER_INFO is forwarded to AsyncMongoClient via from_uri."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
 
     captured_kwargs: dict[str, Any] = {}
 
@@ -728,8 +738,7 @@ def _fake_client(uri: str, **kwargs: Any) -> FakeAsyncMongoClient:
 
 async def test_caller_supplied_driver_info_is_not_overwritten() -> None:
     """A caller-supplied driver kwarg must not be silently replaced."""
-    MongoDBSession._initialized_keys.clear()
-    MongoDBSession._init_locks.clear()
+    MongoDBSession._init_state.clear()
 
     captured_kwargs: dict[str, Any] = {}
     custom_info = FakeDriverInfo(name="MyApp")