fix(memory): prevent close() race and enforce closed state in AsyncSQLiteSession

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

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)