MAINT: Clean up sibling token-region conversions on failure (P311-11)

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

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⟫")

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"]