fix(memory): roll back every failed SQLite session write

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,),

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,),

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()

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()