FIX SQLite cancellation cleanup - #2982
Open
Roman Lutz (romanlutz) wants to merge 3 commits into
Open
Roman Lutz (romanlutz) wants to merge 3 commits into
Roman Lutz (romanlutz) wants to merge 3 commits into
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Separate, main-based follow-up for #2856. Cancelling a score-validation SELECT after
BEGIN IMMEDIATEcan 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 withdatabase table is locked, producing anExceptionGroupinstead 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.
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=shortuv 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=shortuv run --with 'sqlalchemy==2.1.1' --with 'aiosqlite==0.21.0' pytest -q tests\unit\memory\test_sqlite_cancellation.py --tb=shortuv run --with 'sqlalchemy==2.0.41' --with 'aiosqlite==0.22.1' pytest -q tests\unit\memory\test_sqlite_cancellation.py --tb=shortuv run ruff format --check pyrit\memory\sqlite_memory.py tests\unit\memory\test_sqlite_cancellation.py tests\unit\executor\attack\test_scoring_expectation_transport.pyuv run ruff check pyrit\memory\sqlite_memory.py tests\unit\memory\test_sqlite_cancellation.py tests\unit\executor\attack\test_scoring_expectation_transport.pyuv run ty check pyrit\memory\sqlite_memory.py tests\unit\memory\test_sqlite_cancellation.py tests\unit\executor\attack\test_scoring_expectation_transport.pyuv run python -m build_scripts.check_async_suffixgit diff --cached --checkuv run --frozen --extra all --link-mode=copy ty check pyritAn isolated temporary overlay of #2856's exact runtime modules and scorer/batch regression tests at
bb29c8e5281274aea948f0b7965abf2d8bb3afb0passed 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.uv run pytest -q tests\unit\executor\promptgen\test_target_objective_generator.py --tb=shortuv 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=shortuv run ruff format --check tests\unit\executor\promptgen\test_target_objective_generator.pyuv run ruff check tests\unit\executor\promptgen\test_target_objective_generator.pyuv run ty check tests\unit\executor\promptgen\test_target_objective_generator.pygit diff --checkA temporary event-controlled diagnostic blocked every async database operation. Before the test fix, it reproduced the exact
CancelledErroridentity 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.