From 74b29b14d184914d292bd2a019505fda91e81c92 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Sat, 17 Oct 2026 18:16:55 -0700 Subject: [PATCH 1/2] TEST: Cover atomic branch contention and cancellation Exercise controlled SQLite writer contention without changing production lock acquisition, retries, or correctness-test timeouts. Verify SQLITE_BUSY propagation, rollback/commit membership ordering, and cancellation cleanup for branch registration and promotion. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 53404e18-54fd-4013-aae7-97be8c012f4d --- .../memory/test_conversation_branching.py | 150 +++++++++++++++++- 1 file changed, 148 insertions(+), 2 deletions(-) diff --git a/tests/unit/memory/test_conversation_branching.py b/tests/unit/memory/test_conversation_branching.py index 03e6b3bff7..c2ca3cadf6 100644 --- a/tests/unit/memory/test_conversation_branching.py +++ b/tests/unit/memory/test_conversation_branching.py @@ -4,15 +4,17 @@ """Atomic branch insertion and reference updates on real SQLite sessions.""" import asyncio +import sqlite3 import uuid from collections.abc import AsyncGenerator from pathlib import Path +from typing import Any from unittest.mock import MagicMock import pytest -from sqlalchemy import create_engine, select +from sqlalchemy import create_engine, event, select from sqlalchemy.dialects import mssql, sqlite -from sqlalchemy.exc import IntegrityError +from sqlalchemy.exc import IntegrityError, OperationalError from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.orm import Session, sessionmaker @@ -36,6 +38,7 @@ MessagePiece, TargetIdentifier, ) +from unit.async_utils import wait_for_completion_async from unit.mocks import run_memory_session_async @@ -75,6 +78,21 @@ async def _copy_async( return source.model_copy(update={"conversation_id": conversation_id}, deep=True), list(pieces) +async def _change_attack_async( + *, memory: MemoryInterface, attack: AttackResult, source: Conversation, branch: Conversation, promote: bool +) -> bool: + if promote: + return await memory.promote_attack_conversation_async( + attack_result_id=attack.attack_result_id, conversation_id=branch.conversation_id + ) + return await memory.add_conversation_branches_to_attack_async( + attack_result_id=attack.attack_result_id, + source_conversation=source, + conversations=[branch], + message_pieces=[], + ) + + @pytest.mark.usefixtures("patch_central_database") class TestAtomicConversationBranching: def test_history_snapshot_locks_sql_server_rows_and_ranges(self) -> None: @@ -405,6 +423,134 @@ async def promote_async(*, conversation_id: str) -> bool: assert current.conversation_id in {branch.conversation_id for branch in initial} assert len(current.get_pruned_conversation_ids()) == 10 + @pytest.mark.parametrize("promote", [False, True]) + @pytest.mark.parametrize("commit_holder", [False, True]) + async def test_contended_write_preserves_error_and_references( + self, *, file_memory: SQLiteMemory, promote: bool, commit_holder: bool + ) -> None: + attack, source = await _store_attack_async(file_memory) + initial = [Conversation(conversation_id=str(uuid.uuid4())) for _ in range(2)] + assert await file_memory.add_conversation_branches_to_attack_async( + attack_result_id=attack.attack_result_id, conversations=initial, message_pieces=[] + ) + branch = initial[1] if promote else Conversation(conversation_id=str(uuid.uuid4())) + + def load_state(session: Session) -> tuple[str, list[str]]: + entry = session.get(AttackResultEntry, uuid.UUID(attack.attack_result_id)) + assert entry is not None + assert isinstance(entry.conversation_id, str) + return entry.conversation_id, list(entry.pruned_conversation_ids or []) + + def disable_busy_wait(*, dbapi_connection: Any, **kwargs: Any) -> None: + cursor = dbapi_connection.cursor() + try: + cursor.execute("PRAGMA busy_timeout=0") + finally: + cursor.close() + + engine = file_memory._get_async_engine().sync_engine + async with await file_memory.get_session_async() as holder: + entry = await holder.run_sync( + lambda session: file_memory._get_locked_attack_result( + session=session, attack_result_id=attack.attack_result_id + ) + ) + assert entry is not None + entry.conversation_id = initial[0].conversation_id + entry.pruned_conversation_ids = [initial[1].conversation_id, source.conversation_id] + await holder.flush() + event.listen(engine, "checkout", disable_busy_wait, named=True) + try: + # Force native contention without racing the holder against a wall-clock timeout. + with pytest.raises(OperationalError, match="database is locked") as locked: + await _change_attack_async( + memory=file_memory, attack=attack, source=source, branch=branch, promote=promote + ) + assert isinstance(locked.value.orig, sqlite3.OperationalError) + assert locked.value.orig.sqlite_errorcode == sqlite3.SQLITE_BUSY + assert locked.value.statement is not None + assert locked.value.statement.startswith('UPDATE "AttackResultEntries"') + assert await run_memory_session_async(memory=file_memory, operation=load_state) == ( + source.conversation_id, + [item.conversation_id for item in initial], + ) + if commit_holder: + await holder.commit() + else: + await holder.rollback() + finally: + event.remove(engine, "checkout", disable_busy_wait) + await holder.rollback() + + primary, pruned = await run_memory_session_async(memory=file_memory, operation=load_state) + assert (primary, pruned) == ( + (initial[0].conversation_id, [initial[1].conversation_id, source.conversation_id]) + if commit_holder + else (source.conversation_id, [item.conversation_id for item in initial]) + ) + if not promote: + assert await file_memory.get_conversation_metadata_async(conversation_id=branch.conversation_id) is None + assert await file_memory.get_message_pieces_async(conversation_id=branch.conversation_id) == [] + assert await _change_attack_async( + memory=file_memory, attack=attack, source=source, branch=branch, promote=promote + ) + assert await run_memory_session_async(memory=file_memory, operation=load_state) == ( + (branch.conversation_id, [value for value in pruned if value != branch.conversation_id] + [primary]) + if promote + else (primary, [*pruned, branch.conversation_id]) + ) + + @pytest.mark.parametrize("promote", [False, True]) + async def test_cancelled_contended_write_does_not_change_membership( + self, *, file_memory: SQLiteMemory, promote: bool + ) -> None: + attack, source = await _store_attack_async(file_memory) + branch = Conversation(conversation_id=str(uuid.uuid4())) + if promote: + assert await file_memory.add_conversation_branches_to_attack_async( + attack_result_id=attack.attack_result_id, conversations=[branch], message_pieces=[] + ) + [before] = await file_memory.get_attack_results_async(attack_result_ids=[attack.attack_result_id]) + started = asyncio.Event() + + def mark_guard_started(*, statement: str, **kwargs: Any) -> None: + if statement.startswith('UPDATE "AttackResultEntries"'): + started.set() + + engine = file_memory._get_async_engine().sync_engine + async with await file_memory.get_session_async() as holder: + await holder.run_sync( + lambda session: file_memory._get_locked_attack_result( + session=session, attack_result_id=attack.attack_result_id + ) + ) + event.listen(engine, "before_cursor_execute", mark_guard_started, named=True) + task = asyncio.create_task( + _change_attack_async(memory=file_memory, attack=attack, source=source, branch=branch, promote=promote) + ) + ready = asyncio.create_task(started.wait()) + try: + await wait_for_completion_async(future=ready) + assert not task.done() + task.cancel("cancel branch mutation") + await holder.rollback() + with pytest.raises(asyncio.CancelledError, match="cancel branch mutation"): + await wait_for_completion_async(future=task) + finally: + event.remove(engine, "before_cursor_execute", mark_guard_started) + await holder.rollback() + for pending in (task, ready): + if not pending.done(): + pending.cancel() + await asyncio.gather(task, ready, return_exceptions=True) + [after] = await file_memory.get_attack_results_async(attack_result_ids=[attack.attack_result_id]) + assert after == before + if not promote: + assert await file_memory.get_conversation_metadata_async(conversation_id=branch.conversation_id) is None + assert await _change_attack_async( + memory=file_memory, attack=attack, source=source, branch=branch, promote=promote + ) + async def test_promotion_preserves_diagnostic_and_prior_pruned_references( self, sqlite_instance: SQLiteMemory ) -> None: From 466b7e38e7d47d1bcd859331fa9d8ccecae033bd Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Sat, 17 Oct 2026 22:11:12 -0700 Subject: [PATCH 2/2] TEST: Exercise production SQLite branch waits Keep the production async connection factory and cleanup hooks in the branching fixture while preserving its existing busy timeout. Observe SQLite's native UPDATE execution before releasing a competing writer or cancelling, and verify that completed and cancelled calls preserve the holder's committed branch membership without retrying. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 53404e18-54fd-4013-aae7-97be8c012f4d --- .../memory/test_conversation_branching.py | 70 ++++++++++++++----- 1 file changed, 53 insertions(+), 17 deletions(-) diff --git a/tests/unit/memory/test_conversation_branching.py b/tests/unit/memory/test_conversation_branching.py index c2ca3cadf6..d0cfe7360c 100644 --- a/tests/unit/memory/test_conversation_branching.py +++ b/tests/unit/memory/test_conversation_branching.py @@ -15,7 +15,6 @@ from sqlalchemy import create_engine, event, select from sqlalchemy.dialects import mssql, sqlite from sqlalchemy.exc import IntegrityError, OperationalError -from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.orm import Session, sessionmaker from pyrit.memory import MemoryInterface, SQLiteMemory @@ -49,9 +48,15 @@ async def file_memory(*, tmp_path: Path) -> AsyncGenerator[SQLiteMemory, None]: memory.engine.dispose() memory.engine = create_engine(memory.engine.url, connect_args={"timeout": 30}) memory.SessionFactory = sessionmaker(bind=memory.engine, class_=MemorySession) - memory._async_engines[asyncio.get_running_loop()] = create_async_engine( - memory.engine.url.set(drivername="sqlite+aiosqlite"), connect_args={"timeout": 30} - ) + + def configure_busy_timeout(*, dbapi_connection: Any, **kwargs: Any) -> None: + cursor = dbapi_connection.cursor() + try: + cursor.execute("PRAGMA busy_timeout=30000") + finally: + cursor.close() + + event.listen(memory._get_async_engine().sync_engine, "connect", configure_busy_timeout, named=True) try: Base.metadata.create_all(memory.engine) yield memory @@ -500,9 +505,10 @@ def disable_busy_wait(*, dbapi_connection: Any, **kwargs: Any) -> None: else (primary, [*pruned, branch.conversation_id]) ) + @pytest.mark.parametrize("cancel", [False, True]) @pytest.mark.parametrize("promote", [False, True]) - async def test_cancelled_contended_write_does_not_change_membership( - self, *, file_memory: SQLiteMemory, promote: bool + async def test_contended_write_waits_or_cancels_without_losing_committed_membership( + self, *, file_memory: SQLiteMemory, promote: bool, cancel: bool ) -> None: attack, source = await _store_attack_async(file_memory) branch = Conversation(conversation_id=str(uuid.uuid4())) @@ -511,20 +517,36 @@ async def test_cancelled_contended_write_does_not_change_membership( attack_result_id=attack.attack_result_id, conversations=[branch], message_pieces=[] ) [before] = await file_memory.get_attack_results_async(attack_result_ids=[attack.attack_result_id]) + holder_branch = Conversation(conversation_id=str(uuid.uuid4()), attack_result_id=attack.attack_result_id) + before.related_conversations.add( + ConversationReference( + conversation_id=holder_branch.conversation_id, conversation_type=ConversationType.PRUNED + ) + ) + loop = asyncio.get_running_loop() started = asyncio.Event() + statements: list[str] = [] - def mark_guard_started(*, statement: str, **kwargs: Any) -> None: + def trace_statement(statement: str) -> None: + statements.append(statement) if statement.startswith('UPDATE "AttackResultEntries"'): - started.set() + loop.call_soon_threadsafe(started.set) + + def install_trace(*, dbapi_connection: Any, **kwargs: Any) -> None: + dbapi_connection.run_async(lambda driver: driver.set_trace_callback(trace_statement)) engine = file_memory._get_async_engine().sync_engine async with await file_memory.get_session_async() as holder: - await holder.run_sync( + entry = await holder.run_sync( lambda session: file_memory._get_locked_attack_result( session=session, attack_result_id=attack.attack_result_id ) ) - event.listen(engine, "before_cursor_execute", mark_guard_started, named=True) + assert entry is not None + entry.pruned_conversation_ids = [*(entry.pruned_conversation_ids or []), holder_branch.conversation_id] + holder.add(ConversationEntry(conversation=holder_branch)) + await holder.flush() + event.listen(engine, "checkout", install_trace, named=True) task = asyncio.create_task( _change_attack_async(memory=file_memory, attack=attack, source=source, branch=branch, promote=promote) ) @@ -532,20 +554,34 @@ def mark_guard_started(*, statement: str, **kwargs: Any) -> None: try: await wait_for_completion_async(future=ready) assert not task.done() - task.cancel("cancel branch mutation") - await holder.rollback() - with pytest.raises(asyncio.CancelledError, match="cancel branch mutation"): - await wait_for_completion_async(future=task) + assert not any(statement.startswith("SELECT") for statement in statements) + if cancel: + task.cancel("cancel branch mutation") + await holder.commit() + if cancel: + with pytest.raises(asyncio.CancelledError, match="cancel branch mutation"): + await wait_for_completion_async(future=task) + else: + assert await wait_for_completion_async(future=task) + assert ( + sum('SET timestamp="AttackResultEntries".timestamp' in statement for statement in statements) == 1 + ) finally: - event.remove(engine, "before_cursor_execute", mark_guard_started) + event.remove(engine, "checkout", install_trace) await holder.rollback() for pending in (task, ready): if not pending.done(): pending.cancel() await asyncio.gather(task, ready, return_exceptions=True) [after] = await file_memory.get_attack_results_async(attack_result_ids=[attack.attack_result_id]) - assert after == before - if not promote: + if cancel: + assert after == before + else: + assert after.get_active_conversation_ids() == before.get_active_conversation_ids() | { + branch.conversation_id + } + assert after.conversation_id == (branch.conversation_id if promote else source.conversation_id) + if cancel and not promote: assert await file_memory.get_conversation_metadata_async(conversation_id=branch.conversation_id) is None assert await _change_attack_async( memory=file_memory, attack=attack, source=source, branch=branch, promote=promote