diff --git a/docs/sessions/index.md b/docs/sessions/index.md index dd66c087a4..95f66172d0 100644 --- a/docs/sessions/index.md +++ b/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. diff --git a/src/agents/extensions/memory/mongodb_session.py b/src/agents/extensions/memory/mongodb_session.py index 3886c7b853..b2ba601ab0 100644 --- a/src/agents/extensions/memory/mongodb_session.py +++ b/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,6 +220,11 @@ 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. @@ -226,7 +232,13 @@ async def _ensure_indexes(self) -> None: 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 diff --git a/tests/extensions/memory/test_mongodb_session.py b/tests/extensions/memory/test_mongodb_session.py index d463f458ba..3bd8f7c034 100644 --- a/tests/extensions/memory/test_mongodb_session.py +++ b/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 # ---------------------------------------------------------------------------