[Víctor Fernández Ayllón (vayllon301)]

Description

Part of #3166 (P311-11): the selected-region case in Converter._convert_token_regions_async.

Selected regions ran under a bare asyncio.gather. When one region failed, the error reached the caller while the other regions kept running unowned. A region can call a target or write a file, so a request could still be sent, or a file written, after the caller had handled the failure or started a retry. This applies to convert_tokens_async and to SelectiveTextConverter, which share this code path.

This PR awaits the regions with the existing gather_with_cleanup_async helper, so the first failure cancels the regions still running and waits for their cleanup before the original error is raised.

I used the helper rather than a new asyncio.TaskGroup because it already gives the behaviour this item asks for:

• region order: results stay in input order
• error identity: the original exception is re-raised, with no ExceptionGroup to unwrap
• external cancellation: regions are cancelled and drained, and the cancellation message is kept
• Python 3.11/3.12 cancellation counts: no group cancels the caller's task, so no wrapper task is needed to isolate the count

A plain TaskGroup in this method fails the new cancellation-count test on Python 3.11. Scope is this one site. pyrit/common/task_utils.py and other gather call sites are unchanged. The result is the same as before when no region fails.

Conversions already running in a worker thread cannot be interrupted, so cancellation only stops the awaiting task in that case.

Reproduction

import asyncio

from pyrit.converter import Converter, ConverterResult


class SlowTargetConverter(Converter):
    SUPPORTED_INPUT_TYPES = ("text",)
    SUPPORTED_OUTPUT_TYPES = ("text",)

    async def convert_async(self, *, prompt, input_type="text"):
        if prompt == "bad":
            await asyncio.sleep(0)
            raise RuntimeError("rate limited")
        await asyncio.sleep(0.05)
        print(f"region {prompt!r} finished: target called, file written")
        return ConverterResult(output_text=prompt, output_type="text")


async def main():
    try:
        await SlowTargetConverter().convert_tokens_async(prompt="⟪slow⟫ and ⟪bad⟫")
    except RuntimeError as error:
        print(f"caller got {error!r}, regions still running: {len(asyncio.all_tasks()) - 1}")
    await asyncio.sleep(0.1)


asyncio.run(main())

Before:

caller got RuntimeError('rate limited'), regions still running: 1
region 'slow' finished: target called, file written

After:

caller got RuntimeError('rate limited'), regions still running: 0

Tests and Documentation

New tests:

• test_convert_tokens_failed_region_cancels_in_flight_regions_async: a failing region cancels its sibling, the sibling's cleanup finishes before the caller sees the error, and the exact original exception (with its __cause__) is raised. Fails on main.
• test_token_selection_failed_region_cancels_in_flight_regions_async: a failing region cancels its sibling when conversion runs through SelectiveTextConverter, and the original exception is raised. Fails on main.
• test_convert_tokens_out_of_order_completion_keeps_region_order_async: regions that finish in reverse order are still assembled in source order.
• test_convert_tokens_failure_preserves_caller_cancellation_count_async: repeated failures leave Task.cancelling() unchanged.
• test_convert_tokens_caller_cancellation_drains_regions_async: external cancellation cancels every region, keeps the cancellation message, and leaves a count of 1.

The last three pass on main too and pin behaviour that must not change. To check that the tests catch real regressions, I temporarily swapped in four wrong implementations on Python 3.11, and each failed at least two of the new tests:

• plain TaskGroup, and a TaskGroup that unwraps the group: failed the cancellation-count test and the failure test
• cancel without drain: failed both failure tests
• results in completion order: failed the order test and the caller-cancellation test, among others

Results on the commit before rebasing onto current main, which only added an unrelated TemplateSegmentConverter change. After the rebase I reran test_converter.py, test_selective_text_converter.py and test_async_utils.py on Python 3.11: 344 passed, 2 skipped.

• Python 3.11.17: pytest tests/unit/converter tests/unit/prompt_normalizer tests/unit/test_async_utils.py tests/unit/backend/test_converter_service.py tests/unit/backend/test_message_send_service.py -n 2 --dist=loadfile: 3587 passed, 48 skipped.
• Python 3.11.17: pytest tests/unit/common tests/unit/registry tests/unit/executor/promptgen -n 2 --dist=loadfile: 2214 passed, 17 skipped, 4 xfailed.
• Python 3.12.14 and 3.14.8: pytest tests/unit/converter/test_converter.py tests/unit/converter/test_selective_text_converter.py tests/unit/test_async_utils.py: 344 passed, 2 skipped on each.
• ruff check and ruff format --check on the changed files: passed.
• ty check on the changed files: passed, without optional extras.
• pre-commit run --files on the changed files: passed, with the ty-check hook skipped because it needs --extra all.

Not run: the full make unit-test suite, make ty with --extra all, and tests that need optional dependencies, which account for the skips. No live target was used.

Documentation: the convert_tokens_async docstring now states the failure behaviour. The converter notebooks do not describe failures and are unchanged.

JupyText: not run. No notebook or doc page changed, and conversion output is unchanged when every region succeeds.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Víctor Fernández Ayllón (vayllon301)]

@microsoft-github-policy-service agree