From 365d6b741985e3006b70cab1046ce7680b570ef3 Mon Sep 17 00:00:00 2001 From: victor-omarchy Date: Thu, 5 Nov 2026 09:45:45 +0100 Subject: [PATCH] MAINT: Clean up sibling token-region conversions on failure (P311-11) Selected token regions ran under a bare asyncio.gather, so when one region failed the error reached the caller while the other regions kept running unowned. Those regions can still call a target or write a file after the caller has already handled the failure or retried. Await them with gather_with_cleanup_async, which cancels and drains the siblings before re-raising the first error. Region order, error identity, caller cancellation and the caller's cancellation count are unchanged. --- pyrit/converter/converter.py | 6 +- tests/unit/converter/test_converter.py | 98 +++++++++++++++++++ .../test_selective_text_converter.py | 29 ++++++ 3 files changed, 131 insertions(+), 2 deletions(-) diff --git a/pyrit/converter/converter.py b/pyrit/converter/converter.py index 1ac51a3dc8..532b21d96d 100644 --- a/pyrit/converter/converter.py +++ b/pyrit/converter/converter.py @@ -4,7 +4,6 @@ from __future__ import annotations import abc -import asyncio import functools import inspect import re @@ -13,6 +12,7 @@ from pyrit import converter from pyrit.common.random_context import get_random_generator, random_execution +from pyrit.common.task_utils import gather_with_cleanup_async from pyrit.models import ComponentIdentifier, ConverterIdentifier, Identifiable, PromptDataType from pyrit.prompt_target.common.target_requirements import TargetRequirements @@ -254,6 +254,8 @@ async def convert_tokens_async( start and end delimiters form flat pairs. Without delimiters, the entire prompt is converted, including non-text inputs. Selected regions require text input and text output. All delimiters are validated before any conversion is invoked. + If a region fails, the regions still running are cancelled and awaited before its + error is raised. Args: prompt (str): The input prompt containing text to be converted. @@ -311,7 +313,7 @@ async def _convert_token_regions_async( convert_region_async = convert_text_async or self._replace_text_match_async tasks = [convert_region_async(prompt[start + len(start_token) : end - len(end_token)]) for start, end in spans] - converted_parts = await asyncio.gather(*tasks) + converted_parts = await gather_with_cleanup_async(tasks) parts: list[str] = [] previous_end = 0 diff --git a/tests/unit/converter/test_converter.py b/tests/unit/converter/test_converter.py index a108635119..303ee1ffdc 100644 --- a/tests/unit/converter/test_converter.py +++ b/tests/unit/converter/test_converter.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, call, patch import pytest +from unit.async_utils import settle_tasks_async, wait_for_completion_async, wait_for_event_async from unit.mocks import MockPromptTarget from pyrit.converter import ( @@ -252,6 +253,103 @@ async def convert_region_async(*, prompt: str, input_type: PromptDataType = "tex assert result.output_text == "⟪ONE⟫ and ⟪TWO⟫" +async def test_convert_tokens_out_of_order_completion_keeps_region_order_async() -> None: + converter = Base64Converter() + second_finished = asyncio.Event() + + async def convert_region_async(*, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + if prompt == "one": + await second_finished.wait() + else: + second_finished.set() + return ConverterResult(output_text=prompt.upper(), output_type=input_type) + + with patch.object(converter, "convert_async", new_callable=AsyncMock, side_effect=convert_region_async): + result = await converter.convert_tokens_async(prompt="⟪one⟫ and ⟪two⟫") + + assert result.output_text == "ONE and TWO" + + +async def test_convert_tokens_failed_region_cancels_in_flight_regions_async() -> None: + converter = Base64Converter() + error = RuntimeError("rate limited") + error.__cause__ = OSError("connection reset") + sibling_started = asyncio.Event() + cleaned_up: list[str] = [] + + async def convert_region_async(*, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + if prompt == "bad": + await sibling_started.wait() + raise error + sibling_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + # Cleanup that yields must still finish before the caller sees the failure. + await asyncio.sleep(0) + cleaned_up.append(prompt) + raise + raise AssertionError("sibling region was not cancelled") + + with ( + patch.object(converter, "convert_async", new_callable=AsyncMock, side_effect=convert_region_async), + pytest.raises(RuntimeError) as exc_info, + ): + await converter.convert_tokens_async(prompt="⟪slow⟫ and ⟪bad⟫") + + assert exc_info.value is error + assert isinstance(exc_info.value.__cause__, OSError) + assert cleaned_up == ["slow"] + + +async def test_convert_tokens_failure_preserves_caller_cancellation_count_async() -> None: + converter = Base64Converter() + caller = asyncio.current_task() + assert caller is not None + initial_cancellations = caller.cancelling() + error = RuntimeError("rate limited") + + with patch.object(converter, "convert_async", new_callable=AsyncMock, side_effect=error): + for _ in range(2): + with pytest.raises(RuntimeError) as exc_info: + await converter.convert_tokens_async(prompt="⟪one⟫ and ⟪two⟫") + assert exc_info.value is error + assert caller.cancelling() == initial_cancellations + + +async def test_convert_tokens_caller_cancellation_drains_regions_async() -> None: + converter = Base64Converter() + started: set[str] = set() + all_started = asyncio.Event() + cancelled: set[str] = set() + + async def convert_region_async(*, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + started.add(prompt) + if len(started) == 2: + all_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.add(prompt) + raise + raise AssertionError("region was not cancelled") + + with patch.object(converter, "convert_async", new_callable=AsyncMock, side_effect=convert_region_async): + task = asyncio.create_task(converter.convert_tokens_async(prompt="⟪one⟫ and ⟪two⟫")) + try: + await wait_for_event_async(event=all_started, tasks=[task]) + task.cancel("stop converting") + with pytest.raises(asyncio.CancelledError, match="stop converting"): + await wait_for_completion_async(future=task) + assert cancelled == {"one", "two"} + assert task.cancelled() + assert task.cancelling() == 1 + finally: + if not task.done(): + task.cancel() + await settle_tasks_async(task) + + async def test_convert_tokens_nested_parents_wait_for_next_call_async() -> None: converter = SuffixAppendConverter(suffix="tail") first = await converter.convert_tokens_async(prompt="keep ⟪outer ⟪one⟫ and ⟪two⟫ end⟫") diff --git a/tests/unit/converter/test_selective_text_converter.py b/tests/unit/converter/test_selective_text_converter.py index 046ca0eb8d..05ed1ef20e 100644 --- a/tests/unit/converter/test_selective_text_converter.py +++ b/tests/unit/converter/test_selective_text_converter.py @@ -857,3 +857,32 @@ def test_identifier_ignores_word_separator_for_char_level_strategies() -> None: ) assert "word_separator" not in space.get_identifier().params assert space.get_identifier().hash == comma.get_identifier().hash + + +async def test_token_selection_failed_region_cancels_in_flight_regions_async() -> None: + sub_converter = Base64Converter() + converter = SelectiveTextConverter(sub_converter=sub_converter, selection_strategy=TokenSelectionStrategy()) + error = RuntimeError("rate limited") + sibling_started = asyncio.Event() + cancelled: list[str] = [] + + async def convert_region_async(*, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + if prompt == "bad": + await sibling_started.wait() + raise error + sibling_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.append(prompt) + raise + raise AssertionError("sibling region was not cancelled") + + with ( + patch.object(sub_converter, "convert_async", new_callable=AsyncMock, side_effect=convert_region_async), + pytest.raises(RuntimeError) as exc_info, + ): + await converter.convert_tokens_async(prompt="⟪slow⟫ and ⟪bad⟫") + + assert exc_info.value is error + assert cancelled == ["slow"]