Skip to content

FIX SQLite cancellation cleanup - #2982

Open
Roman Lutz (romanlutz) wants to merge 3 commits into
microsoft:mainfrom
romanlutz:romanlutz-sqlite-cancellation-cleanup
Open

Roman Lutz (romanlutz) wants to merge 3 commits into
microsoft:mainfrom
romanlutz:romanlutz-sqlite-cancellation-cleanup

Conversation

@romanlutz

@romanlutz Roman Lutz (romanlutz) commented Oct 4, 2026 •

Copy link
Copy Markdown
Contributor

Description

Separate, main-based follow-up for #2856. Cancelling a score-validation SELECT after BEGIN IMMEDIATE can leave an active native SQLite cursor holding the writer lock even after its connection is closed. The attack's subsequent error-result INSERT then fails with database table is locked, producing an ExceptionGroup instead of preserving the scorer failure.

This change finalizes live cursors through SQLite's supported native connection factory and drains interrupted connection/session cleanup before returning or releasing in-memory transaction access. Weak cursor tracking also covers SQLAlchemy 2.0.41, whose adapted cursor close does not finalize the interrupted native cursor.

Critical self-review found another failure path: a failed session rollback could leave its connection open and replace cancellation. The fix now discards only the affected connection while the transaction is still owned, preserves the original cancellation object, and attaches cleanup failures as its cause. Ordinary errors retain their SQLAlchemy exception and invalidation information.

The production change is confined to sqlite_memory.py. There are no public memory API, schema, SQL Server, dependency, or scoring-policy changes; no generic lock retries or installed-driver patches. #2856's task-draining design remains unchanged.

The user also requested a test-only fix here for the pre-existing generator deadline failure exposed by Python 3.12 CI. The two deadline-focused tests now mock system-prompt setup and retry-history database operations, so unrelated memory latency cannot consume their outer timeout before the intended send. Their original deadlines, cancellation-object identity, error-cause identity, and cleanup checks remain unchanged. New assertions require setup, send, and reset to be awaited. Real-memory generation and retry tests are unchanged.

Tests and Documentation

Added 31 real-database resource-lifetime cases across in-memory and file SQLite, including executed-but-unfetched SELECTs, queued native worker operations, slow cleanup, repeated cancellation, rollback failures, and original-error preservation. A deterministic attack regression verifies that the auxiliary scorer's original error is preserved and its error result is persisted. SQLite backend/session documentation is updated in the owning module.

Validation of the SQLite implementation was performed on Windows with Python 3.11.15. The current environment uses SQLAlchemy 2.1.1 and aiosqlite 0.22.1.

Command Result
uv run pytest -q tests\unit\memory tests\unit\executor\attack\test_scoring_expectation_transport.py tests\unit\prompt_target\test_batch_helper.py --tb=short 1,140 passed, 1 existing skipped
uv run --with 'sqlalchemy==2.0.41' --with 'aiosqlite==0.21.0' pytest -q tests\unit\memory\test_sqlite_cancellation.py tests\unit\executor\attack\test_scoring_expectation_transport.py --tb=short 72 passed
uv run --with 'sqlalchemy==2.1.1' --with 'aiosqlite==0.21.0' pytest -q tests\unit\memory\test_sqlite_cancellation.py --tb=short 31 passed
uv run --with 'sqlalchemy==2.0.41' --with 'aiosqlite==0.22.1' pytest -q tests\unit\memory\test_sqlite_cancellation.py --tb=short 31 passed
uv run ruff format --check pyrit\memory\sqlite_memory.py tests\unit\memory\test_sqlite_cancellation.py tests\unit\executor\attack\test_scoring_expectation_transport.py Passed
uv run ruff check pyrit\memory\sqlite_memory.py tests\unit\memory\test_sqlite_cancellation.py tests\unit\executor\attack\test_scoring_expectation_transport.py Passed
uv run ty check pyrit\memory\sqlite_memory.py tests\unit\memory\test_sqlite_cancellation.py tests\unit\executor\attack\test_scoring_expectation_transport.py Passed
uv run python -m build_scripts.check_async_suffix Passed
git diff --cached --check Passed
uv run --frozen --extra all --link-mode=copy ty check pyrit Passed as the commit hook

An isolated temporary overlay of #2856's exact runtime modules and scorer/batch regression tests at bb29c8e5281274aea948f0b7965abf2d8bb3afb0 passed 102 cases with this fix, on both current and minimum database dependencies. Imported module paths were verified; neither branch's source files were replaced. The pre-fix control reproduced the scorer-plus-SQLite-lock failure at both dependency pairs. Temporary overlay artifacts were removed after validation.

The generator test fix was additionally validated on Windows with Python 3.11.15 and Python 3.12.13. Python 3.12 used a separate uv-managed environment via UV_PROJECT_ENVIRONMENT.

Command Result
uv run pytest -q tests\unit\executor\promptgen\test_target_objective_generator.py --tb=short 51 passed on Python 3.11
uv run --frozen --python 3.12 pytest -q -n 4 --dist=loadfile tests\unit\memory\test_sqlite_cancellation.py tests\unit\executor\attack\test_scoring_expectation_transport.py tests\unit\executor\promptgen\test_target_objective_generator.py --tb=short 123 passed on Python 3.12
uv run ruff format --check tests\unit\executor\promptgen\test_target_objective_generator.py Passed
uv run ruff check tests\unit\executor\promptgen\test_target_objective_generator.py Passed
uv run ty check tests\unit\executor\promptgen\test_target_objective_generator.py Passed
git diff --check Passed

A temporary event-controlled diagnostic blocked every async database operation. Before the test fix, it reproduced the exact CancelledError identity failure seen in CI. After the fix, all four deadline cases passed on both Python versions without awaiting any database operation, while still exercising the actual generator and cleanup timeouts. The diagnostic and its temporary Python 3.12 environment were removed.

JupyText was not run because no notebooks or examples changed. The memory suite includes Azure SQL unit/semantic coverage, but no live SQL Server or full Linux PyRIT run was performed locally. Linux CI is still needed.

Roman Lutz (romanlutz) and others added 3 commits October 3, 2026 20:28
Finalize native cursors before disconnecting and drain connection and session cleanup under repeated cancellation before releasing transaction access. Preserve original cancellation causes and cover real database rollback, queued worker operations, and attack error persistence.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Discard the interrupted connection while its transaction is still owned, drain cleanup before releasing memory access, and retain rollback failures as cancellation causes. Exercise body cancellation, cancellation during close, repeated cancellation, and ordinary rollback failures with real SQLite transactions.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Mock system-prompt setup and retry-history operations in the two deadline unit tests so unrelated SQLite work cannot replace the intended cancellation. Keep the original deadlines and outcome identity assertions, and assert that setup, send, and cleanup are all awaited.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant