fix: address Codex review feedback on MongoDBSession

- P1: Use client.close() (sync) instead of aclose() when closing an
  owned AsyncMongoClient; PyMongo does not expose aclose().
- P1: Include id(client) in _init_key so that different AsyncMongoClient
  instances pointing at different clusters each run their own index-
  creation pass rather than sharing a single guard.
- P2: Add TypeError to the except tuple in get_items() and pop_item() so
  non-string BSON message_data values (e.g. int/object) are silently
  skipped rather than aborting history retrieval.

Adds two new tests: test_non_string_message_data_is_skipped (P2) and
test_different_clients_each_run_index_init (P1).

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

Author: alexbevi

src/agents/extensions/memory/mongodb_session.py

@@ -67,14 +67,16 @@ class MongoDBSession(SessionABC):
     collection.  A lightweight ``sessions`` collection tracks metadata
     (creation time, last-updated time) for each session.
 
-    Indexes are created once per ``(database, sessions_collection,
+    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.
     """
 
     # Class-level registry so index creation only runs once per unique key.
-    _initialized_keys: set[tuple[str, str, str]] = set()
-    _init_locks: dict[tuple[str, str, str], asyncio.Lock] = {}
+    # 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()
 
     session_settings: SessionSettings | None = None
@@ -118,7 +120,7 @@ def __init__(
         self._sessions: AsyncCollection = db[sessions_collection]
         self._messages: AsyncCollection = db[messages_collection]
 
-        self._init_key = (database, sessions_collection, messages_collection)
+        self._init_key = (id(client), database, sessions_collection, messages_collection)
 
     # ------------------------------------------------------------------
     # Convenience constructors
@@ -250,8 +252,8 @@ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
         for doc in docs:
             try:
                 items.append(await self._deserialize_item(doc["message_data"]))
-            except (json.JSONDecodeError, KeyError):
-                # Skip corrupted or malformed documents.
+            except (json.JSONDecodeError, KeyError, TypeError):
+                # Skip corrupted or malformed documents (including non-string BSON values).
                 continue
 
         return items
@@ -302,7 +304,7 @@ async def pop_item(self) -> TResponseInputItem | None:
 
         try:
             return await self._deserialize_item(doc["message_data"])
-        except (json.JSONDecodeError, KeyError):
+        except (json.JSONDecodeError, KeyError, TypeError):
             return None
 
     async def clear_session(self) -> None:
@@ -323,7 +325,7 @@ async def close(self) -> None:
         caller is responsible for managing its lifecycle.
         """
         if self._owns_client:
-            await self._client.aclose()
+            self._client.close()
 
     async def ping(self) -> bool:
         """Test MongoDB connectivity.

tests/extensions/memory/test_mongodb_session.py

@@ -191,10 +191,14 @@ def append_metadata(self, driver_info: FakeDriverInfo) -> None:
         """Record append_metadata calls for test assertions."""
         self._metadata_calls.append(driver_info)
 
-    async def aclose(self) -> None:
+    def close(self) -> None:
+        """Synchronous close — matches PyMongo's AsyncMongoClient.close() signature."""
         self._closed = True
         self.admin._closed = True
 
+    async def aclose(self) -> None:
+        self.close()
+
 
 # ---------------------------------------------------------------------------
 # Inject fake pymongo into sys.modules before importing the module under test
@@ -484,6 +488,19 @@ async def test_missing_message_data_field_is_skipped(session: MongoDBSession) ->
     assert len(items) == 1
 
 
+async def test_non_string_message_data_is_skipped(session: MongoDBSession) -> None:
+    """Documents whose message_data is a non-string BSON type are silently skipped."""
+    await session.add_items([{"role": "user", "content": "valid"}])
+
+    # Inject a document where message_data is an integer — json.loads raises TypeError.
+    bad_doc = {"_id": FakeObjectId(), "session_id": session.session_id, "message_data": 42}
+    session._messages._docs[id(bad_doc["_id"])] = bad_doc  # type: ignore[attr-defined]
+
+    items = await session.get_items()
+    assert len(items) == 1
+    assert items[0].get("content") == "valid"
+
+
 # ---------------------------------------------------------------------------
 # Index initialisation (idempotency)
 # ---------------------------------------------------------------------------
@@ -513,6 +530,40 @@ async def counting(*args: Any, **kwargs: Any) -> str:
     session._sessions.create_index = original_sessions  # type: ignore[attr-defined]
 
 
+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()
+
+    client_a = FakeAsyncMongoClient()
+    client_b = FakeAsyncMongoClient()
+
+    call_counts: dict[str, int] = {"a": 0, "b": 0}
+
+    async def counting_a(*args: Any, **kwargs: Any) -> str:
+        call_counts["a"] += 1
+        return "fake_index"
+
+    async def counting_b(*args: Any, **kwargs: Any) -> str:
+        call_counts["b"] += 1
+        return "fake_index"
+
+    s_a = MongoDBSession("x", client=client_a, database="agents_test")  # type: ignore[arg-type]
+    s_b = MongoDBSession("x", client=client_b, database="agents_test")  # type: ignore[arg-type]
+
+    s_a._messages.create_index = counting_a  # type: ignore[attr-defined]
+    s_a._sessions.create_index = counting_a  # type: ignore[attr-defined]
+    s_b._messages.create_index = counting_b  # type: ignore[attr-defined]
+    s_b._sessions.create_index = counting_b  # type: ignore[attr-defined]
+
+    await s_a._ensure_indexes()
+    await s_b._ensure_indexes()
+
+    # Each client must trigger its own index creation (2 calls = sessions + messages).
+    assert call_counts["a"] == 2
+    assert call_counts["b"] == 2
+
+
 # ---------------------------------------------------------------------------
 # Connectivity and lifecycle
 # ---------------------------------------------------------------------------