From 7d711a88c7f3639e57d305c508e383de2c4ea8eb Mon Sep 17 00:00:00 2001 From: abhay-codes07 Date: Wed, 5 Aug 2026 15:20:24 +0530 Subject: [PATCH 1/2] fix(memory): roll back every failed SQLite session write #4163 established that a write failing partway through leaves an open transaction on the cached connection, and that an open write transaction holds the SQLite write lock for the lifetime of that connection and blocks every later writer. That fix reached only SQLiteSession.add_items. The same defect remained in SQLiteSession.pop_item and clear_session, and in all three AsyncSQLiteSession write paths. clear_session is the clearest case: it issues two DELETEs under one commit, so a failure on the second leaves the first applied inside an open transaction. AsyncSQLiteSession is the most damaging, because it holds one connection for the whole session, so the lock stays held until the session is closed. The rollback obligation belongs to the connection rather than to any single method, so add a _rollback_on_failure(conn) guard per module and apply it at every _locked_connection() write site, including the add_items path that previously inlined it. Commit points are unchanged; only the failure path differs. The guard catches BaseException so an interrupted write cannot strand the lock either. --- .../extensions/memory/async_sqlite_session.py | 22 ++++- src/agents/memory/sqlite_session.py | 35 ++++--- .../memory/test_async_sqlite_session.py | 91 +++++++++++++++++++ tests/memory/test_session.py | 68 ++++++++++++++ 4 files changed, 199 insertions(+), 17 deletions(-) diff --git a/src/agents/extensions/memory/async_sqlite_session.py b/src/agents/extensions/memory/async_sqlite_session.py index 06d0cc1755..18f709b303 100644 --- a/src/agents/extensions/memory/async_sqlite_session.py +++ b/src/agents/extensions/memory/async_sqlite_session.py @@ -18,6 +18,22 @@ ) +@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: + await conn.rollback() + raise + + class AsyncSQLiteSession(SessionABC): """Async SQLite-based implementation of session storage. @@ -206,7 +222,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 +255,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 +300,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..1a668528ee 100644 --- a/src/agents/memory/sqlite_session.py +++ b/src/agents/memory/sqlite_session.py @@ -14,6 +14,22 @@ 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: + conn.rollback() + raise + + class SQLiteSession(SessionABC): """SQLite-based implementation of session storage. @@ -287,18 +303,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 +317,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 +367,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..8c45ac1d0a 100644 --- a/tests/extensions/memory/test_async_sqlite_session.py +++ b/tests/extensions/memory/test_async_sqlite_session.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import sqlite3 import tempfile from collections.abc import Sequence from datetime import datetime @@ -495,3 +496,93 @@ 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() 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() From 1cf970724bf82dba48148f30f641a0f867bfea2e Mon Sep 17 00:00:00 2001 From: abhay-codes07 Date: Wed, 5 Aug 2026 15:54:57 +0530 Subject: [PATCH 2/2] fix(memory): keep a failing rollback from masking the original error Review follow-up. Rollback is cleanup, so a connection that is already closed or otherwise unusable would previously replace the failure the caller needs to see. Attempt it best-effort and always re-raise the original. Also add a regression test for cancellation mid-write, which is the case the guard most needs to cover on the async backend: the session holds one connection, so a transaction left open by a cancelled write holds the SQLite write lock until the session is closed. --- .../extensions/memory/async_sqlite_session.py | 8 +++- src/agents/memory/sqlite_session.py | 8 +++- .../memory/test_async_sqlite_session.py | 41 +++++++++++++++++++ 3 files changed, 55 insertions(+), 2 deletions(-) diff --git a/src/agents/extensions/memory/async_sqlite_session.py b/src/agents/extensions/memory/async_sqlite_session.py index 18f709b303..cf07f55b6b 100644 --- a/src/agents/extensions/memory/async_sqlite_session.py +++ b/src/agents/extensions/memory/async_sqlite_session.py @@ -30,7 +30,13 @@ async def _rollback_on_failure(conn: aiosqlite.Connection) -> AsyncIterator[None try: yield except BaseException: - await conn.rollback() + # 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 diff --git a/src/agents/memory/sqlite_session.py b/src/agents/memory/sqlite_session.py index 1a668528ee..3df2f6e855 100644 --- a/src/agents/memory/sqlite_session.py +++ b/src/agents/memory/sqlite_session.py @@ -26,7 +26,13 @@ def _rollback_on_failure(conn: sqlite3.Connection) -> Iterator[None]: try: yield except BaseException: - conn.rollback() + # 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 diff --git a/tests/extensions/memory/test_async_sqlite_session.py b/tests/extensions/memory/test_async_sqlite_session.py index 8c45ac1d0a..8ec48208c5 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 sqlite3 import tempfile @@ -586,3 +587,43 @@ async def test_failed_pop_item_releases_write_lock(): 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()