diff --git a/src/agents/extensions/memory/async_sqlite_session.py b/src/agents/extensions/memory/async_sqlite_session.py index 215a668902..9c088efaf7 100644 --- a/src/agents/extensions/memory/async_sqlite_session.py +++ b/src/agents/extensions/memory/async_sqlite_session.py @@ -2,6 +2,7 @@ import asyncio import json +import sqlite3 from collections.abc import AsyncIterator from contextlib import asynccontextmanager from pathlib import Path @@ -134,7 +135,7 @@ async def _get_connection(self) -> aiosqlite.Connection: raise assert connection is not None try: - await connection.execute("PRAGMA journal_mode=WAL") + await self._configure_connection(connection) await self._init_db_for_connection(connection) except BaseException as initialization_error: close_error = await self._close_owned_connection(connection) @@ -147,6 +148,24 @@ async def _get_connection(self) -> aiosqlite.Connection: return self._connection + @staticmethod + async def _configure_connection(conn: aiosqlite.Connection) -> None: + """Enable WAL, retrying its transient initialization lock.""" + async with conn.execute("PRAGMA busy_timeout") as cursor: + timeout_row = await cursor.fetchone() + timeout_seconds = (timeout_row[0] if timeout_row is not None else 0) / 1000 + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_seconds + while True: + try: + async with conn.execute("PRAGMA journal_mode=WAL") as cursor: + await cursor.fetchone() + return + except sqlite3.OperationalError as exc: + if "locked" not in str(exc).lower() or loop.time() >= deadline: + raise + await asyncio.sleep(min(0.01, max(0, deadline - loop.time()))) + def _check_not_closed(self) -> None: """Raise if the session has already been closed.""" if self._closed: diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index 6b4d616c0f..2213f88858 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -730,6 +730,131 @@ async def _init_db_for_connection(self, conn: Any) -> None: await session.captured_connection.close() +async def test_initialization_retries_transient_wal_lock( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Concurrent first connections can fail WAL setup before busy_timeout elapses.""" + import aiosqlite + + real_execute = aiosqlite.Connection.execute + wal_attempts = 0 + + def execute(self: Any, sql: str, *args: Any, **kwargs: Any) -> Any: + nonlocal wal_attempts + if sql == "PRAGMA journal_mode=WAL": + wal_attempts += 1 + if wal_attempts == 1: + raise sqlite3.OperationalError("database is locked") + return real_execute(self, sql, *args, **kwargs) + + monkeypatch.setattr(aiosqlite.Connection, "execute", execute) + session = AsyncSQLiteSession("startup", tmp_path / "startup.db") + items: list[TResponseInputItem] = [{"role": "user", "content": "hello"}] + try: + await session.add_items(items) + assert await session.get_items() == items + assert wal_attempts == 2 + finally: + await session.close() + + +async def test_initialization_releases_wal_cursor_before_schema_write(tmp_path: Path) -> None: + """A competing schema write must not leave startup using an old WAL snapshot.""" + import aiosqlite + + class ConcurrentSchemaSession(AsyncSQLiteSession): + async def _init_db_for_connection(self, conn: aiosqlite.Connection) -> None: + # Coordinate a real writer between WAL setup and the session schema write. + async with aiosqlite.connect(self.db_path) as peer: + await peer.execute("CREATE TABLE competing_writer (id INTEGER)") + await peer.commit() + await super()._init_db_for_connection(conn) + + session = ConcurrentSchemaSession("startup", tmp_path / "startup.db") + items: list[TResponseInputItem] = [{"role": "user", "content": "hello"}] + try: + await session.add_items(items) + assert await session.get_items() == items + finally: + await session.close() + + +@pytest.mark.parametrize( + ("error", "timeout"), + [("database is locked", 0), ("database is locked", 0.01), ("disk I/O error", 5)], + ids=["busy-timeout-disabled", "busy-timeout-elapsed", "unrelated-error"], +) +async def test_initialization_does_not_retry_unrecoverable_wal_errors( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, error: str, timeout: float +) -> None: + """Disabled retries and unrelated failures still close the startup connection.""" + import aiosqlite + + real_connect = aiosqlite.connect + real_execute = aiosqlite.Connection.execute + connections: list[Any] = [] + wal_attempts = 0 + + def connect(database: str) -> Any: + conn = real_connect(database, timeout=timeout) + connections.append(conn) + return conn + + def execute(self: Any, sql: str, *args: Any, **kwargs: Any) -> Any: + nonlocal wal_attempts + if sql == "PRAGMA journal_mode=WAL": + wal_attempts += 1 + raise sqlite3.OperationalError(error) + return real_execute(self, sql, *args, **kwargs) + + monkeypatch.setattr(aiosqlite, "connect", connect) + monkeypatch.setattr(aiosqlite.Connection, "execute", execute) + session = AsyncSQLiteSession("startup", tmp_path / "startup.db") + try: + with pytest.raises(sqlite3.OperationalError, match=error): + await asyncio.wait_for(session.get_items(), timeout=1) + if timeout == 0 or error != "database is locked": + assert wal_attempts == 1 + assert len(connections) == 1 + assert connections[0]._running is False + finally: + await session.close() + + +async def test_cancelled_wal_retry_closes_startup_connection( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Cancelling while WAL remains busy must not retain an initialization worker.""" + import aiosqlite + + real_execute = aiosqlite.Connection.execute + retry_started = asyncio.Event() + connections: list[Any] = [] + + def execute(self: Any, sql: str, *args: Any, **kwargs: Any) -> Any: + if sql == "PRAGMA journal_mode=WAL": + connections.append(self) + if len(connections) == 2: + retry_started.set() + raise sqlite3.OperationalError("database is locked") + return real_execute(self, sql, *args, **kwargs) + + monkeypatch.setattr(aiosqlite.Connection, "execute", execute) + session = AsyncSQLiteSession("startup", tmp_path / "startup.db") + task = asyncio.create_task(session.get_items()) + try: + await asyncio.wait_for(retry_started.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert connections[0]._running is False + assert session._connection is None + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await session.close() + + async def test_cancelled_add_items_rolls_back_write_transaction( tmp_path: Path, monkeypatch: pytest.MonkeyPatch,