From 4aad60760265953cf0621af1ed2ae446b3a296ee Mon Sep 17 00:00:00 2001 From: Henry Su Date: Mon, 27 Jul 2026 08:36:43 -0500 Subject: [PATCH 1/3] fix(memory): prevent close() race and enforce closed state in AsyncSQLiteSession --- .../extensions/memory/async_sqlite_session.py | 16 ++++++--- .../memory/test_async_sqlite_session.py | 36 +++++++++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) 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..b5fa2b1e39 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 @@ -409,3 +410,38 @@ 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() From a19a13aa88d29e0b8ee31793dbb008517b5469d9 Mon Sep 17 00:00:00 2001 From: Henry Su Date: Mon, 27 Jul 2026 08:56:48 -0500 Subject: [PATCH 2/3] test(memory): lock down AsyncSQLiteSession close interleaving Add a controlled regression for an operation waiting behind close() so it raises RuntimeError and does not recreate the connection. --- .../memory/test_async_sqlite_session.py | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index b5fa2b1e39..040f60a930 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -13,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 @@ -445,3 +446,66 @@ async def test_async_sqlite_session_closed_operations_raise_runtime_error(): 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 + close_started = asyncio.Event() + release_close = 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 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] + monkeypatch.setattr(aiosqlite, "connect", tracking_connect) + + close_task = asyncio.create_task(session.close()) + await close_started.wait() + assert session._lock.locked() + assert session._closed is True + + get_task = asyncio.create_task(session.get_items()) + # Wait until get_items() is parked on the session lock held by close(). + for _ in range(100): + waiters = session._lock._waiters + if waiters: + break + assert not get_task.done() + await asyncio.sleep(0) + else: + pytest.fail("get_items() did not wait on the session lock held by close()") + assert not get_task.done() + + release_close.set() + await close_task + + with pytest.raises(RuntimeError, match="AsyncSQLiteSession is closed"): + await get_task + + assert session._closed is True + assert session._connection is None + assert connect_calls == 0 From 77aa05344ed2f0aa875f312cd8c855b75bfbd8e9 Mon Sep 17 00:00:00 2001 From: Henry Su Date: Mon, 27 Jul 2026 10:38:55 -0500 Subject: [PATCH 3/3] test(memory): harden AsyncSQLiteSession close interleaving cleanup Guarantee release/task cleanup with try/finally, signal waiters via a test-owned lock event, and bound synchronization waits so assertion failures cannot hang the test. --- .../memory/test_async_sqlite_session.py | 58 +++++++++++-------- 1 file changed, 35 insertions(+), 23 deletions(-) diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index 040f60a930..be04b15092 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -465,8 +465,10 @@ async def test_async_sqlite_session_close_rejects_operation_waiting_on_lock( 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 @@ -475,37 +477,47 @@ async def paused_close(*args: Any, **kwargs: Any) -> None: 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()) - await close_started.wait() - assert session._lock.locked() - assert session._closed is True - - get_task = asyncio.create_task(session.get_items()) - # Wait until get_items() is parked on the session lock held by close(). - for _ in range(100): - waiters = session._lock._waiters - if waiters: - break + 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() - await asyncio.sleep(0) - else: - pytest.fail("get_items() did not wait on the session lock held by close()") - assert not get_task.done() - - release_close.set() - await close_task - - with pytest.raises(RuntimeError, match="AsyncSQLiteSession is closed"): - await get_task - assert session._closed is True - assert session._connection is None - assert connect_calls == 0 + 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)