fix(memory): enforce closed state in MongoDBSession

docs/sessions/index.md

@@ -448,7 +448,7 @@ await session.close()
 
 Notes:
 
--   `from_uri(...)` creates and owns the `AsyncMongoClient` and closes it on `session.close()`. If your application already manages a client, construct `MongoDBSession(...)` directly with `client=...`; in that case `session.close()` is a no-op and lifecycle stays with the caller.
+-   `from_uri(...)` creates and owns the `AsyncMongoClient` and closes it on `session.close()`. An owned-client session is terminal after `close()`, and subsequent session operations raise `RuntimeError`. If your application already manages a client, construct `MongoDBSession(...)` directly with `client=...`; in that case `session.close()` is a no-op, and lifecycle plus session usability stay with the caller.
 -   Connect to [MongoDB Atlas](https://www.mongodb.com/products/platform) by passing an `mongodb+srv://user:password@cluster.example.mongodb.net` URI to `from_uri(...)` with no other changes.
 -   Two collections are used and both names are configurable via `sessions_collection=` (default `agent_sessions`) and `messages_collection=` (default `agent_messages`). Indexes are created automatically on first use. Each message document carries a monotonically increasing `seq` counter that preserves ordering across concurrent writers and processes.
 -   Use `await session.ping()` to verify connectivity before your first run.

src/agents/extensions/memory/mongodb_session.py

@@ -139,6 +139,7 @@ def __init__(
         )
         self._client = client
         self._owns_client = False
+        self._closed = False
 
         client.append_metadata(_DRIVER_INFO)
 
@@ -219,14 +220,25 @@ def _mark_init_done(self) -> None:
                 weakref.finalize(self._client, self._init_state.pop, self._client_id, None)
             per_client[self._init_sub_key] = True
 
+    def _check_not_closed(self) -> None:
+        """Raise if the session has already been closed."""
+        if self._closed:
+            raise RuntimeError("MongoDBSession is closed")
+
     async def _ensure_indexes(self) -> None:
         """Create required indexes the first time this (client, sub_key) is accessed.
 
         ``create_index`` is idempotent on the server side, so concurrent calls
         from different coroutines or event loops are safe — at most a redundant
         round-trip is issued.  The threading-lock-guarded boolean prevents that
         extra round-trip after the first call completes.
+
+        Session operations that require index initialization go through here, so
+        this is also where they reject a closed session. The empty ``add_items``
+        fast path and ``ping`` check the closed state directly.
         """
+        self._check_not_closed()
+
         if self._is_init_done():
             return
 
@@ -312,6 +324,10 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
         Args:
             items: List of input items to append to the session.
         """
+        # Checked before the empty-list fast path, which would otherwise return
+        # successfully on a closed session.
+        self._check_not_closed()
+
         if not items:
             return
 
@@ -385,18 +401,32 @@ async def close(self) -> None:
         """Close the underlying MongoDB connection.
 
         Only closes the client if this session owns it (i.e. it was created
-        via :meth:`from_uri`).  If the client was injected externally the
-        caller is responsible for managing its lifecycle.
+        via :meth:`from_uri`).  In that case the session becomes terminal and
+        subsequent operations raise ``RuntimeError``.  If the client was injected
+        externally the caller is responsible for managing its lifecycle and this
+        is a no-op.
+
+        The session is terminal from the first close attempt.  If releasing the
+        client fails or is cancelled, operations still raise and a later close()
+        retries the release, which ``AsyncMongoClient.close`` allows.
         """
-        if self._owns_client:
-            await self._client.close()
+        if not self._owns_client:
+            return
+
+        self._closed = True
+        await self._client.close()
 
     async def ping(self) -> bool:
         """Test MongoDB connectivity.
 
         Returns:
             ``True`` if the server is reachable, ``False`` otherwise.
+
+        Raises:
+            RuntimeError: If the session owns its client and has been closed.
         """
+        # Checked outside the try block; the except clause below would swallow it.
+        self._check_not_closed()
         try:
             await self._client.admin.command("ping")
             return True

tests/extensions/memory/test_mongodb_session.py

@@ -8,12 +8,13 @@
 
 from __future__ import annotations
 
+import asyncio
 import sys
 import types
 from collections import defaultdict
 from datetime import datetime, timezone
 from typing import Any
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 
@@ -729,6 +730,115 @@ async def test_close_owned_client_is_closed() -> None:
         assert fake_client._closed
 
 
+def _make_owned_session(session_id: str = "owned") -> MongoDBSession:
+    """Create a from_uri session, which is the case where close() owns the client."""
+    MongoDBSession._init_state.clear()
+    with patch(
+        "agents.extensions.memory.mongodb_session.AsyncMongoClient",
+        return_value=FakeAsyncMongoClient(),
+    ):
+        return MongoDBSession.from_uri(session_id, uri="mongodb://localhost:27017", database="t")
+
+
+async def test_closed_operations_raise_runtime_error() -> None:
+    """Operations on a closed session must fail instead of running against a released client."""
+    session = _make_owned_session()
+    await session.add_items([{"role": "user", "content": "hi"}])
+    await session.close()
+
+    with pytest.raises(RuntimeError, match="^MongoDBSession is closed$"):
+        await session.get_items()
+    with pytest.raises(RuntimeError, match="^MongoDBSession is closed$"):
+        await session.add_items([{"role": "user", "content": "after close"}])
+    with pytest.raises(RuntimeError, match="^MongoDBSession is closed$"):
+        await session.pop_item()
+    with pytest.raises(RuntimeError, match="^MongoDBSession is closed$"):
+        await session.clear_session()
+
+
+async def test_closed_rejects_empty_add_items() -> None:
+    """add_items([]) must not bypass the closed check through the empty-list fast path."""
+    session = _make_owned_session()
+    await session.close()
+
+    with pytest.raises(RuntimeError, match="^MongoDBSession is closed$"):
+        await session.add_items([])
+
+
+async def test_close_before_use_is_terminal() -> None:
+    """close() before the first operation must still be terminal."""
+    session = _make_owned_session()
+    await session.close()
+
+    with pytest.raises(RuntimeError, match="^MongoDBSession is closed$"):
+        await session.get_items()
+
+
+async def test_repeated_close_remains_safe() -> None:
+    """Repeated close() calls must remain safe for callers."""
+    session = _make_owned_session()
+
+    await session.close()
+    await session.close()
+
+    with pytest.raises(RuntimeError, match="^MongoDBSession is closed$"):
+        await session.get_items()
+
+
+async def test_failed_close_is_terminal_and_can_be_retried() -> None:
+    """A failed client release must leave the session terminal and cleanup retryable."""
+    session = _make_owned_session()
+    close_mock = AsyncMock(side_effect=[ConnectionError("close failed"), None])
+
+    with patch.object(session._client, "close", close_mock):
+        with pytest.raises(ConnectionError, match="close failed"):
+            await session.close()
+
+        with pytest.raises(RuntimeError, match="^MongoDBSession is closed$"):
+            await session.get_items()
+
+        await session.close()
+
+    assert close_mock.await_count == 2
+
+
+async def test_cancelled_close_is_terminal_and_can_be_retried() -> None:
+    """A cancelled client release must leave the session terminal and cleanup retryable."""
+    session = _make_owned_session()
+    close_mock = AsyncMock(side_effect=[asyncio.CancelledError(), None])
+
+    with patch.object(session._client, "close", close_mock):
+        with pytest.raises(asyncio.CancelledError):
+            await session.close()
+
+        with pytest.raises(RuntimeError, match="^MongoDBSession is closed$"):
+            await session.get_items()
+
+        await session.close()
+
+    assert close_mock.await_count == 2
+
+
+async def test_ping_on_closed_session_raises() -> None:
+    """ping() swallows connectivity errors, so the closed check runs outside its try."""
+    session = _make_owned_session()
+    await session.close()
+
+    with pytest.raises(RuntimeError, match="^MongoDBSession is closed$"):
+        await session.ping()
+
+
+async def test_external_client_session_stays_usable_after_close() -> None:
+    """An injected client is the caller's to manage, so close() must not be terminal."""
+    session = _make_session()
+    assert session._owns_client is False
+
+    await session.close()
+
+    await session.add_items([{"role": "user", "content": "still works"}])
+    assert len(await session.get_items()) == 1
+
+
 # ---------------------------------------------------------------------------
 # Runner integration
 # ---------------------------------------------------------------------------