FIX SQLite cancellation cleanup

pyrit/memory/sqlite_memory.py

@@ -6,20 +6,26 @@
 import threading
 import uuid
 import weakref
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from contextlib import closing
+from contextvars import ContextVar
 from datetime import datetime
 from pathlib import Path
+from sqlite3 import Connection as SQLiteConnection
+from sqlite3 import Cursor as SQLiteCursor
+from types import TracebackType
 from typing import TYPE_CHECKING, Any, Literal
 
-from sqlalchemy import and_, case, create_engine, exists, func, or_, select, text
+from sqlalchemy import and_, case, create_engine, event, exists, func, or_, select, text
+from sqlalchemy.engine import AdaptedConnection, ExceptionContext
 from sqlalchemy.engine.base import Engine
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
 from sqlalchemy.orm import InstrumentedAttribute, sessionmaker
 from sqlalchemy.orm.session import Session
 from sqlalchemy.pool import StaticPool
 from sqlalchemy.sql.expression import TextClause
+from sqlalchemy.util.concurrency import greenlet_spawn
 
 from pyrit.common.path import DB_DATA_PATH
 from pyrit.common.singleton import Singleton
@@ -38,20 +44,137 @@
     from sqlalchemy.engine import Connection
 
 logger = logging.getLogger(__name__)
+_sqlite_session_cleanup: ContextVar[bool] = ContextVar("sqlite_session_cleanup", default=False)
 
 
-class _SerializedAsyncSession(AsyncSession):
-    def __init__(self, *, engine: AsyncEngine, release: Callable[[], None]) -> None:
+class _CursorClosingSQLiteConnection(SQLiteConnection):
+    """A native SQLite connection that finalizes live cursors before disconnecting."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._cursors: weakref.WeakSet[SQLiteCursor] = weakref.WeakSet()
+
+    def cursor(self, *args: Any, **kwargs: Any) -> SQLiteCursor:
+        return self._track_cursor(super().cursor(*args, **kwargs))
+
+    def execute(self, *args: Any, **kwargs: Any) -> SQLiteCursor:
+        return self._track_cursor(super().execute(*args, **kwargs))
+
+    def executemany(self, *args: Any, **kwargs: Any) -> SQLiteCursor:
+        return self._track_cursor(super().executemany(*args, **kwargs))
+
+    def executescript(self, *args: Any, **kwargs: Any) -> SQLiteCursor:
+        return self._track_cursor(super().executescript(*args, **kwargs))
+
+    def close(self) -> None:
+        for cursor in tuple(self._cursors):
+            cursor.close()
+        self._cursors.clear()
+        super().close()
+
+    def _track_cursor(self, cursor: SQLiteCursor) -> SQLiteCursor:
+        self._cursors.add(cursor)
+        return cursor
+
+
+async def _finish_sqlite_cleanup_async(cleanup: Awaitable[None]) -> asyncio.CancelledError | None:
+    """
+    Drain SQLite cleanup and retain cancellation received while waiting.
+
+    Returns:
+        asyncio.CancelledError | None: The first cancellation received during cleanup.
+    """
+    task = asyncio.ensure_future(cleanup)
+    cancellation: asyncio.CancelledError | None = None
+    while not task.done():
+        try:
+            await asyncio.shield(task)
+        except asyncio.CancelledError as error:
+            if cancellation is None:
+                cancellation = error
+        except Exception:
+            break
+    try:
+        task.result()
+    except Exception as error:
+        if cancellation is not None:
+            raise cancellation from error
+        raise
+    return cancellation
+
+
+def _cleanup_interrupted_sqlite_connection(context: ExceptionContext) -> None:
+    execution_context, connection = context.execution_context, context.connection
+    cancelled = isinstance(context.original_exception, asyncio.CancelledError)
+    if connection is None or not (cancelled or _sqlite_session_cleanup.get()):
+        return
+    # The public interface has empty slots, but implementations expose these writable flags.
+    context.is_disconnect = True  # type: ignore[ty:missing-slot]
+    context.invalidate_pool_on_disconnect = False  # type: ignore[ty:missing-slot]
+    dbapi_connection = connection.connection.dbapi_connection
+    if not isinstance(dbapi_connection, AdaptedConnection):
+        raise TypeError("Async SQLite memory requires an adapted driver connection.")
+    cursor = execution_context.cursor if execution_context is not None else None
+
+    def close_and_invalidate() -> None:
+        # SQLAlchemy skips cursor cleanup on cancellation. SQLite keeps an active
+        # statement's transaction lock even after its connection is closed.
+        if cursor is not None:
+            cursor.close()
+        connection.invalidate(context.original_exception)
+
+    try:
+        dbapi_connection.run_async(lambda _: _finish_sqlite_cleanup_async(greenlet_spawn(close_and_invalidate)))
+    except (asyncio.CancelledError, Exception) as error:
+        if cancelled:
+            cause = (
+                error.__cause__ if isinstance(error, asyncio.CancelledError) and error.__cause__ is not None else error
+            )
+            raise context.original_exception from cause
+        raise
+
+
+class _SQLiteAsyncSession(AsyncSession):
+    def __init__(self, *, engine: AsyncEngine, release: Callable[[], None] | None = None) -> None:
         super().__init__(bind=engine, sync_session_class=MemorySession)
         self._release: Callable[[], None] | None = release
 
     async def close(self) -> None:  # pyrit-async-suffix-exempt
         try:
-            await super().close()
+            cancellation = await _finish_sqlite_cleanup_async(self._close_session_async())
         finally:
             if self._release is not None:
                 release, self._release = self._release, None
                 release()
+        if cancellation is not None:
+            raise cancellation
+
+    async def __aexit__(
+        self,
+        type_: type[BaseException] | None,
+        value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        try:
+            await self.close()
+        except asyncio.CancelledError as error:
+            if isinstance(value, asyncio.CancelledError):
+                cause = error.__cause__ if error.__cause__ is not None else value.__cause__
+                raise value from cause
+            raise
+        except Exception as error:
+            if isinstance(value, asyncio.CancelledError):
+                raise value from error
+            raise
+
+    async def _close_session_async(self) -> None:
+        # A failed rollback must discard its connection before the ORM drops
+        # the transaction reference and before exclusive access is released.
+        token = _sqlite_session_cleanup.set(True)
+        try:
+            await super().close()
+        finally:
+            _sqlite_session_cleanup.reset(token)
 
 
 class SQLiteMemory(MemoryInterface, metaclass=Singleton):
@@ -61,6 +184,11 @@ class SQLiteMemory(MemoryInterface, metaclass=Singleton):
     This class provides functionality to insert, query, and manage conversation data
     using SQLite. It supports both file-based and in-memory databases.
 
+    Cancellation finalizes active cursors and closes interrupted connections before
+    returning to the caller. Session cleanup also finishes under repeated cancellation.
+    Failed disconnects and session rollbacks preserve the original cancellation
+    and expose cleanup failures as its cause.
+
     Note: this is replacing the old DuckDB implementation.
     """
 
@@ -124,17 +252,21 @@ def _initialize_schema(self) -> None:
 
     def _create_async_engine(self) -> AsyncEngine:
         database = self._memory_uri if self.db_path == ":memory:" else str(self.db_path)
-        kwargs: dict[str, Any] = {}
+        kwargs: dict[str, Any] = {"connect_args": {"factory": _CursorClosingSQLiteConnection}}
         if self.db_path == ":memory:":
             kwargs["poolclass"] = StaticPool
-        return create_async_engine(f"sqlite+aiosqlite:///{database}", echo=self._verbose, **kwargs)
+        engine = create_async_engine(f"sqlite+aiosqlite:///{database}", echo=self._verbose, **kwargs)
+        event.listen(engine.sync_engine, "handle_error", _cleanup_interrupted_sqlite_connection)
+        return engine
 
     async def get_session_async(self) -> AsyncSession:
         """
-        Create a session with exclusive access to the shared in-memory database.
+        Create a session with cancellation-safe SQLite cleanup.
+
+        In-memory sessions also have exclusive access to the shared database.
 
         Returns:
-            AsyncSession: A session that releases exclusive access when closed.
+            AsyncSession: A session that finishes cleanup before releasing exclusive access.
 
         Raises:
             NotImplementedError: If a custom sync session hook has not been migrated.
@@ -146,11 +278,11 @@ async def get_session_async(self) -> AsyncSession:
             raise RuntimeError("Close the synchronous memory session before opening an async session on this thread.")
         connection_lock = self._transaction_lock
         if connection_lock is None:
-            return await super().get_session_async()
+            return _SQLiteAsyncSession(engine=self._get_async_engine())
         while not connection_lock.acquire(blocking=False):
             await asyncio.sleep(0.01)
         try:
-            return _SerializedAsyncSession(engine=self._get_async_engine(), release=connection_lock.release)
+            return _SQLiteAsyncSession(engine=self._get_async_engine(), release=connection_lock.release)
         except BaseException:
             connection_lock.release()
             raise

tests/unit/executor/attack/test_scoring_expectation_transport.py

@@ -4,11 +4,14 @@
 """Per-execution expectations reach real outcome scorers without becoming attack prompts."""
 
 import asyncio
+from collections.abc import Iterable
 from contextlib import nullcontext
 from typing import Any, Literal, cast
 from unittest.mock import AsyncMock, patch
 
 import pytest
+from aiosqlite import Cursor
+from sqlalchemy import event
 from unit.mocks import MockPromptTarget, store_message_async
 
 from pyrit.exceptions import (
@@ -588,6 +591,62 @@ async def test_attack_error_names_failing_scorer_async(
         assert raised.value.__cause__.__cause__ is original
         assert get_execution_context() is None
 
+    async def test_attack_error_persists_after_cancelled_score_validation_async(
+        self, sqlite_instance: SQLiteMemory
+    ) -> None:
+        objective, auxiliary = _RecordingScorer(), _RecordingScorer(value=False)
+        attack = PromptSendingAttack(
+            objective_target=MockPromptTarget(),
+            attack_scoring_config=AttackScoringConfig(objective_scorer=objective, auxiliary_scorers=[auxiliary]),
+        )
+        write_started, validation_started, release_read = asyncio.Event(), asyncio.Event(), asyncio.Event()
+        original = ValueError("scoring failed")
+
+        def record_write(*args: Any) -> None:
+            if args[2] == "BEGIN IMMEDIATE":
+                write_started.set()
+
+        async def fail_auxiliary_async(**_kwargs: Any) -> list[Score]:
+            await validation_started.wait()
+            raise original
+
+        fetchall = Cursor.fetchall
+
+        async def delayed_fetchall_async(cursor: Cursor) -> Iterable[Any]:
+            if write_started.is_set() and not validation_started.is_set():
+                validation_started.set()
+                await release_read.wait()
+            return await fetchall(cursor)
+
+        engine = sqlite_instance._get_async_engine()
+        event.listen(engine.sync_engine, "after_cursor_execute", record_write)
+        try:
+            with (
+                patch.object(auxiliary, "_score_scorable_async", side_effect=fail_auxiliary_async),
+                patch.object(Cursor, "fetchall", new=delayed_fetchall_async),
+            ):
+                task = asyncio.create_task(
+                    attack.execute_async(objective="attack objective", expectation=_expectation())
+                )
+                try:
+                    with pytest.raises(RuntimeError, match="Strategy execution failed for auxiliary_scorer") as raised:
+                        await asyncio.wait_for(task, timeout=5)
+                finally:
+                    release_read.set()
+                    if not task.done():
+                        task.cancel()
+                    await asyncio.gather(task, return_exceptions=True)
+        finally:
+            event.remove(engine.sync_engine, "after_cursor_execute", record_write)
+
+        assert validation_started.is_set()
+        assert raised.value.__cause__.__cause__ is original
+        [stored] = await sqlite_instance.get_attack_results_async(objective="attack objective")
+        assert stored.outcome == AttackOutcome.ERROR
+        assert stored.error_message is not None
+        assert "scoring failed" in stored.error_message
+        assert await sqlite_instance.get_scores_async() == []
+
     @pytest.mark.parametrize("duplicate", [False, True], ids=["initial_node", "duplicated_node"])
     @pytest.mark.parametrize("explicit_expectation", [False, True], ids=["default_objective", "scoring_objective"])
     async def test_tap_node_preserves_execution_expectation_async(

tests/unit/executor/promptgen/test_target_objective_generator.py

@@ -21,10 +21,10 @@
 
 
 @contextmanager
-def _isolate_pre_send_io(generator: TargetObjectiveGenerator) -> Iterator[None]:
+def _isolate_pre_send_io(generator: TargetObjectiveGenerator) -> Iterator[AsyncMock]:
     """Mock setup and retry-history I/O so timeout tests reach their intended await."""
     with (
-        patch.object(generator._target, "set_system_prompt_async", new_callable=AsyncMock),
+        patch.object(generator._target, "set_system_prompt_async", new_callable=AsyncMock) as setup,
         patch.object(generator._normalizer.memory, "get_message_pieces_async", new_callable=AsyncMock, return_value=[]),
         patch.object(
             generator._normalizer.memory,
@@ -33,7 +33,7 @@ def _isolate_pre_send_io(generator: TargetObjectiveGenerator) -> Iterator[None]:
             return_value=0,
         ),
     ):
-        yield
+        yield setup
 
 
 @pytest.mark.usefixtures("patch_central_database")
@@ -249,7 +249,7 @@ async def reset_async(*, conversation_id: str) -> None:
                 cleanup_cancelled.set()
 
         with (
-            _isolate_pre_send_io(generator),
+            _isolate_pre_send_io(generator) as setup,
             patch.object(generator, "_CLEANUP_TIMEOUT_SECONDS", 0.01),
             patch.object(
                 generator._normalizer, "send_prompt_async", new_callable=AsyncMock, side_effect=wait_forever_async
@@ -262,6 +262,7 @@ async def reset_async(*, conversation_id: str) -> None:
         send.assert_awaited_once()
         assert cancelled.is_set()
         assert cleanup_cancelled.is_set()
+        setup.assert_awaited_once()
         reset.assert_awaited_once()
         assert "Timed out resetting generation conversation" in caplog.text
 
@@ -323,15 +324,15 @@ async def reset_async(*, conversation_id: str) -> None:
                 cleanup_cancelled.set()
 
         with (
-            _isolate_pre_send_io(generator),
+            _isolate_pre_send_io(generator) as setup,
             patch.object(generator, "_CLEANUP_TIMEOUT_SECONDS", 0.01),
             patch.object(
                 generator._normalizer,
                 "send_prompt_async",
                 new_callable=AsyncMock,
                 return_value=response,
                 side_effect=failure,
-            ),
+            ) as send,
             patch.object(generator._target, "reset_conversation_async", side_effect=reset_async) as reset,
         ):
             async with asyncio.timeout(1):
@@ -348,6 +349,8 @@ async def reset_async(*, conversation_id: str) -> None:
                     assert generation_error.value.__cause__ is failure
         reset.assert_awaited_once()
         assert cleanup_cancelled.is_set()
+        setup.assert_awaited_once()
+        send.assert_awaited_once()
         assert "Timed out resetting generation conversation" in caplog.text
 
     async def test_cancellation_propagates_async(self) -> None: