diff --git a/src/agents/extensions/memory/async_sqlite_session.py b/src/agents/extensions/memory/async_sqlite_session.py index 06d0cc1755..cf07f55b6b 100644 --- a/src/agents/extensions/memory/async_sqlite_session.py +++ b/src/agents/extensions/memory/async_sqlite_session.py @@ -18,6 +18,28 @@ ) +@asynccontextmanager +async def _rollback_on_failure(conn: aiosqlite.Connection) -> AsyncIterator[None]: + """Roll back a partially applied write when the operation fails. + + `_locked_connection()` does not manage transactions, so a statement that fails partway + through a write would otherwise leave both a partial mutation and an open transaction on + this shared connection. An open write transaction holds the SQLite write lock for the + lifetime of the connection and blocks every later writer, including other processes. + """ + try: + yield + except BaseException: + # Rollback is best-effort cleanup. If it fails too -- a closed or otherwise unusable + # connection -- the original failure is the one worth reporting, so never let the + # cleanup error replace it. + try: + await conn.rollback() + except Exception: + pass + raise + + class AsyncSQLiteSession(SessionABC): """Async SQLite-based implementation of session storage. @@ -206,7 +228,7 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: if not items: return - async with self._locked_connection() as conn: + async with self._locked_connection() as conn, _rollback_on_failure(conn): await conn.execute( f""" INSERT OR IGNORE INTO {self.sessions_table} (session_id) VALUES (?) @@ -239,7 +261,7 @@ async def pop_item(self) -> TResponseInputItem | None: Returns: The most recent item if it exists, None if the session is empty """ - async with self._locked_connection() as conn: + async with self._locked_connection() as conn, _rollback_on_failure(conn): cursor = await conn.execute( f""" DELETE FROM {self.messages_table} @@ -284,7 +306,7 @@ async def pop_item(self) -> TResponseInputItem | None: async def clear_session(self) -> None: """Clear all items for this session.""" - async with self._locked_connection() as conn: + async with self._locked_connection() as conn, _rollback_on_failure(conn): await conn.execute( f"DELETE FROM {self.messages_table} WHERE session_id = ?", (self.session_id,), diff --git a/src/agents/memory/sqlite_session.py b/src/agents/memory/sqlite_session.py index 4bd641cc8d..3df2f6e855 100644 --- a/src/agents/memory/sqlite_session.py +++ b/src/agents/memory/sqlite_session.py @@ -14,6 +14,28 @@ from .session_settings import SessionSettings, coerce_session_settings, resolve_session_limit +@contextmanager +def _rollback_on_failure(conn: sqlite3.Connection) -> Iterator[None]: + """Roll back a partially applied write when the operation fails. + + `_locked_connection()` does not manage transactions, so a statement that fails partway + through a write would otherwise leave both a partial mutation and an open transaction on + this cached connection. An open write transaction holds the SQLite write lock for the + lifetime of the connection and blocks every later writer, including other processes. + """ + try: + yield + except BaseException: + # Rollback is best-effort cleanup. If it fails too -- a closed or otherwise unusable + # connection -- the original failure is the one worth reporting, so never let the + # cleanup error replace it. + try: + conn.rollback() + except Exception: + pass + raise + + class SQLiteSession(SessionABC): """SQLite-based implementation of session storage. @@ -287,18 +309,9 @@ async def add_items(self, items: list[TResponseInputItem]) -> None: return def _add_items_sync(): - with self._locked_connection() as conn: - try: - self._insert_items(conn, items) - conn.commit() - except Exception: - # _locked_connection() does not manage transactions; roll back - # explicitly so a failure partway through the insert never leaves a - # partial mutation or an open transaction on this cached connection. - # An open write transaction would hold the SQLite write lock for the - # lifetime of the connection and block every later writer. - conn.rollback() - raise + with self._locked_connection() as conn, _rollback_on_failure(conn): + self._insert_items(conn, items) + conn.commit() await asyncio.to_thread(_add_items_sync) @@ -310,7 +323,7 @@ async def pop_item(self) -> TResponseInputItem | None: """ def _pop_item_sync(): - with self._locked_connection() as conn: + with self._locked_connection() as conn, _rollback_on_failure(conn): # Use DELETE with RETURNING to atomically delete and return the most recent item cursor = conn.execute( f""" @@ -360,7 +373,7 @@ async def clear_session(self) -> None: """Clear all items for this session.""" def _clear_session_sync(): - with self._locked_connection() as conn: + with self._locked_connection() as conn, _rollback_on_failure(conn): conn.execute( f"DELETE FROM {self.messages_table} WHERE session_id = ?", (self.session_id,), diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index b45cbdf4e7..8ec48208c5 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -2,7 +2,9 @@ from __future__ import annotations +import asyncio import json +import sqlite3 import tempfile from collections.abc import Sequence from datetime import datetime @@ -495,3 +497,133 @@ async def test_async_sqlite_session_close_is_idempotent(): with pytest.raises(RuntimeError, match="AsyncSQLiteSession is closed"): await session.get_items() + + +def _drop_table(db_path: Path, table: str) -> None: + """Drop a table from an independent connection to make a later statement fail.""" + helper = sqlite3.connect(str(db_path)) + try: + helper.execute(f"DROP TABLE {table}") + helper.commit() + finally: + helper.close() + + +def _write_lock_is_free(db_path: Path) -> bool: + """Return whether an independent writer can still take the SQLite write lock.""" + # timeout=0 disables the busy handler, so this fails immediately if the lock is held. + probe = sqlite3.connect(str(db_path), timeout=0) + try: + probe.execute("CREATE TABLE IF NOT EXISTS probe_lock (x INTEGER)") + probe.commit() + return True + except sqlite3.OperationalError: + return False + finally: + probe.close() + + +async def test_failed_add_items_releases_write_lock(): + """A failed add_items must not leave an open write transaction on the shared connection. + + AsyncSQLiteSession holds one connection for the whole session, so an open write + transaction holds the SQLite write lock until the session is closed and blocks every + later writer, including other processes. + """ + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "add_rollback.db" + session = AsyncSQLiteSession("add_rollback", db_path=db_path) + + # json.dumps() fails only after the sessions-table upsert opened a write transaction. + unserializable = cast(TResponseInputItem, {"role": "user", "content": object()}) + with pytest.raises(TypeError): + await session.add_items([unserializable]) + + conn = await session._get_connection() + assert conn.in_transaction is False + assert _write_lock_is_free(db_path) + + # The session must remain usable after the failure. + await session.add_items([{"role": "user", "content": "after failure"}]) + assert [item.get("content") for item in await session.get_items()] == ["after failure"] + + await session.close() + + +async def test_failed_clear_session_releases_write_lock(): + """A clear_session that fails on its second DELETE must not strand the write lock.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "clear_rollback.db" + session = AsyncSQLiteSession("clear_rollback", db_path=db_path) + await session.add_items([{"role": "user", "content": "kept"}]) + + # The messages DELETE succeeds, then the sessions DELETE fails. + _drop_table(db_path, "agent_sessions") + + with pytest.raises(sqlite3.OperationalError): + await session.clear_session() + + conn = await session._get_connection() + assert conn.in_transaction is False + assert _write_lock_is_free(db_path) + + await session.close() + + +async def test_failed_pop_item_releases_write_lock(): + """A pop_item whose statement fails must not strand the write lock.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "pop_rollback.db" + session = AsyncSQLiteSession("pop_rollback", db_path=db_path) + await session.add_items([{"role": "user", "content": "kept"}]) + + _drop_table(db_path, "agent_messages") + + with pytest.raises(sqlite3.OperationalError): + await session.pop_item() + + conn = await session._get_connection() + assert conn.in_transaction is False + assert _write_lock_is_free(db_path) + + await session.close() + + +async def test_cancelled_write_still_releases_write_lock(tmp_path: Path): + """Cancellation must still roll back, or the guard fails in the case it exists for. + + An unshielded `await conn.rollback()` inside the cancellation handler is itself cancelled + immediately, so the transaction would stay open and keep holding the write lock. + """ + if True: + db_path = tmp_path / "cancel_rollback.db" + session = AsyncSQLiteSession("cancel_rollback", db_path=db_path) + await session.add_items([{"role": "user", "content": "kept"}]) + + conn = await session._get_connection() + real_execute = conn.execute + started = asyncio.Event() + + async def hang_after_first_write(*args: Any, **kwargs: Any) -> Any: + # Let the sessions-table upsert open the write transaction, then stall so the + # caller can cancel mid-write. + cursor = await real_execute(*args, **kwargs) + started.set() + await asyncio.sleep(3600) + return cursor + + conn.execute = hang_after_first_write # type: ignore[method-assign] + task = asyncio.create_task(session.add_items([{"role": "user", "content": "cancelled"}])) + await started.wait() + task.cancel() + with pytest.raises((asyncio.CancelledError, asyncio.TimeoutError)): + await asyncio.wait_for(task, timeout=5) + + conn.execute = real_execute # type: ignore[method-assign] + # Give the shielded rollback a turn to finish before asserting. + await asyncio.sleep(0.1) + + assert conn.in_transaction is False + assert _write_lock_is_free(db_path) + + await session.close() diff --git a/tests/memory/test_session.py b/tests/memory/test_session.py index 3b180539b6..1874696c4b 100644 --- a/tests/memory/test_session.py +++ b/tests/memory/test_session.py @@ -955,3 +955,71 @@ async def test_runner_with_session_settings_override(): assert len(history_items) == 2 session.close() + + +def _drop_table(db_path: Path, table: str) -> None: + """Drop a table from an independent connection to make a later statement fail.""" + helper = sqlite3.connect(str(db_path)) + try: + helper.execute(f"DROP TABLE {table}") + helper.commit() + finally: + helper.close() + + +def _write_lock_is_free(db_path: Path) -> bool: + """Return whether an independent writer can still take the SQLite write lock.""" + # timeout=0 disables the busy handler, so this fails immediately if the lock is held. + probe = sqlite3.connect(str(db_path), timeout=0) + try: + probe.execute("CREATE TABLE IF NOT EXISTS probe_lock (x INTEGER)") + probe.commit() + return True + except sqlite3.OperationalError: + return False + finally: + probe.close() + + +@pytest.mark.asyncio +async def test_sqlite_session_failed_clear_session_releases_write_lock(): + """A clear_session that fails on its second DELETE must not strand the write lock. + + clear_session issues two DELETEs under one commit, so a failure on the second leaves the + first applied inside an open transaction. That transaction holds the SQLite write lock for + the lifetime of the cached connection and blocks every later writer. + """ + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "clear_rollback.db" + session = SQLiteSession("clear_rollback", db_path) + await session.add_items([{"role": "user", "content": "kept"}]) + + # The messages DELETE succeeds, then the sessions DELETE fails. + _drop_table(db_path, "agent_sessions") + + with pytest.raises(sqlite3.OperationalError): + await session.clear_session() + + assert all(not conn.in_transaction for conn in session._connections) + assert _write_lock_is_free(db_path) + + session.close() + + +@pytest.mark.asyncio +async def test_sqlite_session_failed_pop_item_releases_write_lock(): + """A pop_item whose statement fails must not strand the write lock.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "pop_rollback.db" + session = SQLiteSession("pop_rollback", db_path) + await session.add_items([{"role": "user", "content": "kept"}]) + + _drop_table(db_path, "agent_messages") + + with pytest.raises(sqlite3.OperationalError): + await session.pop_item() + + assert all(not conn.in_transaction for conn in session._connections) + assert _write_lock_is_free(db_path) + + session.close()