diff --git a/src/agents/extensions/memory/async_sqlite_session.py b/src/agents/extensions/memory/async_sqlite_session.py index 63ae77081b..0b2c9b0505 100644 --- a/src/agents/extensions/memory/async_sqlite_session.py +++ b/src/agents/extensions/memory/async_sqlite_session.py @@ -59,6 +59,7 @@ def __init__( self._connection: aiosqlite.Connection | None = None self._lock = asyncio.Lock() self._init_lock = asyncio.Lock() + self._closed = False async def _init_db_for_connection(self, conn: aiosqlite.Connection) -> None: """Initialize the database schema for a specific connection.""" @@ -96,10 +97,15 @@ async def _init_db_for_connection(self, conn: aiosqlite.Connection) -> None: async def _get_connection(self) -> aiosqlite.Connection: """Get or create a database connection.""" + if self._closed: + raise RuntimeError("AsyncSQLiteSession is closed") + if self._connection is not None: return self._connection async with self._init_lock: + if self._closed: + raise RuntimeError("AsyncSQLiteSession is closed") if self._connection is None: self._connection = await aiosqlite.connect(str(self.db_path)) await self._connection.execute("PRAGMA journal_mode=WAL") @@ -264,8 +270,10 @@ async def clear_session(self) -> None: async def close(self) -> None: """Close the database connection.""" - if self._connection is None: - return async with self._lock: - await self._connection.close() - self._connection = None + if self._closed: + return + self._closed = True + if self._connection is not None: + await self._connection.close() + self._connection = None diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index 6ab3d9feb4..be04b15092 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import json import tempfile from collections.abc import Sequence @@ -12,6 +13,7 @@ import pytest pytest.importorskip("aiosqlite") # Skip tests if aiosqlite is not installed +import aiosqlite from agents import Agent, Runner, TResponseInputItem from agents.extensions.memory import AsyncSQLiteSession @@ -409,3 +411,113 @@ async def test_async_sqlite_session_pop_item_same_timestamp_returns_latest(): assert _item_ids(remaining) == ["rs_pop_same_ts"] await session.close() + + +async def test_async_sqlite_session_concurrent_close(): + """Test that concurrent and repeated calls to close() are safe and idempotent.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "async_close_race.db" + session = AsyncSQLiteSession("close_race_test", db_path) + await session.add_items([{"role": "user", "content": "hello"}]) + + # Multiple concurrent close calls should succeed without raising AttributeError + await asyncio.gather(session.close(), session.close(), session.close()) + + # Additional sequential close call should also be a safe no-op + await session.close() + + +async def test_async_sqlite_session_closed_operations_raise_runtime_error(): + """Test that operations on a closed AsyncSQLiteSession raise RuntimeError.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "async_closed_ops.db" + session = AsyncSQLiteSession("closed_ops_test", db_path) + await session.add_items([{"role": "user", "content": "hello"}]) + await session.close() + + with pytest.raises(RuntimeError, match="AsyncSQLiteSession is closed"): + await session.get_items() + + with pytest.raises(RuntimeError, match="AsyncSQLiteSession is closed"): + await session.add_items([{"role": "user", "content": "more"}]) + + with pytest.raises(RuntimeError, match="AsyncSQLiteSession is closed"): + await session.pop_item() + + with pytest.raises(RuntimeError, match="AsyncSQLiteSession is closed"): + await session.clear_session() + + +async def test_async_sqlite_session_close_rejects_operation_waiting_on_lock( + monkeypatch: pytest.MonkeyPatch, +): + """An operation queued behind close() must fail closed without reconnecting. + + Controlled interleaving: close() acquires the session lock first, an + operation starts while close is paused inside the lock, then close + completes. The waiter must raise RuntimeError and must not recreate the + connection. + """ + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "async_close_interleave.db" + session = AsyncSQLiteSession("close_interleave_test", db_path) + await session.add_items([{"role": "user", "content": "hello"}]) + + assert session._connection is not None + original_conn_close = session._connection.close + original_acquire = session._lock.acquire + close_started = asyncio.Event() + release_close = asyncio.Event() + operation_waiting = asyncio.Event() + connect_calls = 0 + real_connect = aiosqlite.connect + + async def paused_close(*args: Any, **kwargs: Any) -> None: + close_started.set() + await release_close.wait() + await original_conn_close(*args, **kwargs) + + async def acquire_and_signal(*args: Any, **kwargs: Any) -> bool: + # Signal before waiting whenever the lock is already held by close(). + if session._lock.locked(): + operation_waiting.set() + return await original_acquire(*args, **kwargs) + + async def tracking_connect(*args: Any, **kwargs: Any) -> Any: + nonlocal connect_calls + connect_calls += 1 + return await real_connect(*args, **kwargs) + + session._connection.close = paused_close # type: ignore[method-assign] + session._lock.acquire = acquire_and_signal # type: ignore[method-assign] + monkeypatch.setattr(aiosqlite, "connect", tracking_connect) + + close_task = asyncio.create_task(session.close()) + get_task: asyncio.Task[Any] | None = None + try: + await asyncio.wait_for(close_started.wait(), timeout=1) + assert session._lock.locked() + assert session._closed is True + + get_task = asyncio.create_task(session.get_items()) + await asyncio.wait_for(operation_waiting.wait(), timeout=1) + assert not get_task.done() + + release_close.set() + await asyncio.wait_for(close_task, timeout=1) + + with pytest.raises(RuntimeError, match="AsyncSQLiteSession is closed"): + await asyncio.wait_for(get_task, timeout=1) + + assert session._closed is True + assert session._connection is None + assert connect_calls == 0 + finally: + release_close.set() + tasks = [task for task in (close_task, get_task) if task is not None] + if tasks: + _done, pending = await asyncio.wait(tasks, timeout=1) + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True)