FIX: Read manual XPIA input off the event loop

pyrit/executor/workflow/xpia.py

@@ -1,7 +1,10 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+import asyncio
+import concurrent.futures
 import logging
+import threading
 import uuid
 from dataclasses import dataclass, field
 from enum import Enum
@@ -661,11 +664,92 @@ async def _setup_async(self, *, context: XPIAContext) -> None:
                 modified to include the manual processing callback.
         """
 
-        # Create the manual input callback
+        # Create the manual input callback. The console read runs off the event loop so other
+        # async work keeps running while the operator triggers the processing target.
         async def manual_input_async() -> str:
             prompt = "Please trigger the processing target's execution and paste the output here:"
-            return input(prompt)
+            return await _CONSOLE_READER.read_line_async(prompt=prompt)
 
         # Set the processing callback on the context
         context.processing_callback = manual_input_async
         return await super()._setup_async(context=context)
+
+
+class _ConsoleLineReader:
+    """
+    Read operator input from the console off the event loop, one caller at a time.
+
+    ``input()`` cannot be interrupted, so a read keeps waiting for the operator even after its
+    caller is cancelled. Each read runs on a daemon thread so it never delays interpreter
+    shutdown, and a line typed for a cancelled caller is handed to the next caller instead of
+    being dropped. Callers take turns, so concurrent reads never receive the same line.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._turn: concurrent.futures.Future[None] = concurrent.futures.Future()
+        self._turn.set_result(None)
+        self._unclaimed_line: concurrent.futures.Future[str] | None = None
+
+    async def read_line_async(self, *, prompt: str) -> str:
+        """
+        Return the next console line, reusing a read left behind by a cancelled caller.
+
+        Args:
+            prompt (str): The prompt shown when a new read starts.
+
+        Returns:
+            str: The line the operator entered.
+
+        Raises:
+            asyncio.CancelledError: If the caller is cancelled before the line arrives. The read
+                keeps running and its line goes to the next caller.
+        """
+        with self._lock:
+            previous_turn = self._turn
+            self._turn = turn = concurrent.futures.Future()
+        try:
+            await asyncio.shield(asyncio.wrap_future(previous_turn))
+        except BaseException:
+            # A caller cancelled while queued never held the turn, so it passes the turn on only once
+            # the caller ahead of it is done. Releasing it now would let a later caller share the pending line.
+            previous_turn.add_done_callback(lambda _: turn.set_result(None))
+            raise
+        try:
+            with self._lock:
+                if self._unclaimed_line is None:
+                    self._unclaimed_line = self._start_read(prompt=prompt)
+                line = self._unclaimed_line
+            try:
+                result = await asyncio.shield(asyncio.wrap_future(line))
+            except asyncio.CancelledError:
+                # Leave the read in place so the line the operator types reaches the next caller.
+                raise
+            except BaseException:
+                self._release_line(line=line)
+                raise
+            self._release_line(line=line)
+            return result
+        finally:
+            turn.set_result(None)
+
+    def _release_line(self, *, line: concurrent.futures.Future[str]) -> None:
+        with self._lock:
+            if self._unclaimed_line is line:
+                self._unclaimed_line = None
+
+    @staticmethod
+    def _start_read(*, prompt: str) -> concurrent.futures.Future[str]:
+        line: concurrent.futures.Future[str] = concurrent.futures.Future()
+
+        def read() -> None:
+            try:
+                line.set_result(input(prompt))
+            except BaseException as exc:
+                line.set_exception(exc)
+
+        threading.Thread(target=read, name="pyrit-console-reader", daemon=True).start()
+        return line
+
+
+_CONSOLE_READER = _ConsoleLineReader()

tests/unit/executor/workflow/test_xpia.py

@@ -1,14 +1,20 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+import asyncio
+import threading
 import uuid
+from collections.abc import Iterator
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from unit.async_utils import TEST_WATCHDOG_SECONDS
 from unit.mocks import get_mock_prompt_normalizer
 
+from pyrit.executor.workflow import xpia as xpia_module
 from pyrit.executor.workflow.xpia import (
     XPIAContext,
+    XPIAManualProcessingWorkflow,
     XPIAResult,
     XPIAStatus,
     XPIAWorkflow,
@@ -698,3 +704,149 @@ async def test_xpia_test_setup_raises_when_processing_prompt_is_none(
         assert context.processing_callback is not None
         with pytest.raises(RuntimeError, match="context.processing_prompt is not initialized"):
             await context.processing_callback()
+
+
+@pytest.fixture
+def manual_workflow(mock_attack_setup_target: MagicMock, mock_scorer: MagicMock) -> XPIAManualProcessingWorkflow:
+    """Create a manual processing XPIA workflow for testing."""
+    return XPIAManualProcessingWorkflow(attack_setup_target=mock_attack_setup_target, scorer=mock_scorer)
+
+
+@pytest.fixture
+def fresh_console_reader() -> Iterator[None]:
+    """Give each test its own console reader so a read left pending cannot leak between tests."""
+    with patch.object(xpia_module, "_CONSOLE_READER", xpia_module._ConsoleLineReader()):
+        yield
+
+
+@pytest.mark.usefixtures("patch_central_database", "fresh_console_reader")
+class TestXPIAManualProcessingWorkflow:
+    """Tests for the manual processing callback created during setup."""
+
+    async def test_setup_async_sets_callback_that_returns_operator_input(
+        self, manual_workflow: XPIAManualProcessingWorkflow, valid_message: Message
+    ) -> None:
+        context = XPIAContext(attack_content=valid_message)
+        await manual_workflow._setup_async(context=context)
+        assert context.processing_callback is not None
+
+        with patch("builtins.input", return_value="processed output") as mock_input:
+            response = await context.processing_callback()
+
+        assert response == "processed output"
+        mock_input.assert_called_once_with(
+            "Please trigger the processing target's execution and paste the output here:"
+        )
+
+    async def test_manual_input_does_not_block_the_event_loop(
+        self, manual_workflow: XPIAManualProcessingWorkflow, valid_message: Message
+    ) -> None:
+        context = XPIAContext(attack_content=valid_message)
+        await manual_workflow._setup_async(context=context)
+        assert context.processing_callback is not None
+        event_loop_released_input = threading.Event()
+
+        def wait_for_event_loop(prompt: str) -> str:
+            # Only a task on the event loop sets this, so the read can finish only if the loop keeps running.
+            assert event_loop_released_input.wait(timeout=TEST_WATCHDOG_SECONDS)
+            return "processed output"
+
+        async def release_input_async() -> None:
+            event_loop_released_input.set()
+
+        with patch("builtins.input", side_effect=wait_for_event_loop):
+            response, _ = await asyncio.gather(context.processing_callback(), release_input_async())
+
+        assert response == "processed output"
+
+    async def test_cancelled_read_hands_its_line_to_the_next_read(
+        self, manual_workflow: XPIAManualProcessingWorkflow, valid_message: Message
+    ) -> None:
+        context = XPIAContext(attack_content=valid_message)
+        await manual_workflow._setup_async(context=context)
+        assert context.processing_callback is not None
+        input_started = threading.Event()
+        release_input = threading.Event()
+
+        def operator_input(prompt: str) -> str:
+            input_started.set()
+            assert release_input.wait(timeout=TEST_WATCHDOG_SECONDS)
+            return "processed output"
+
+        with patch("builtins.input", side_effect=operator_input) as mock_input:
+            first_read = asyncio.create_task(context.processing_callback())
+            assert await asyncio.to_thread(input_started.wait, TEST_WATCHDOG_SECONDS)
+            first_read.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await first_read
+
+            second_read = asyncio.create_task(context.processing_callback())
+            release_input.set()
+            response = await asyncio.wait_for(second_read, timeout=TEST_WATCHDOG_SECONDS)
+
+        assert response == "processed output"
+        mock_input.assert_called_once()
+
+    async def test_cancelled_queued_read_keeps_later_reads_waiting_their_turn(
+        self, manual_workflow: XPIAManualProcessingWorkflow, valid_message: Message
+    ) -> None:
+        context = XPIAContext(attack_content=valid_message)
+        await manual_workflow._setup_async(context=context)
+        assert context.processing_callback is not None
+        first_input_started = threading.Event()
+        release_first_input = threading.Event()
+        lines = iter(["first output", "second output"])
+
+        def operator_input(prompt: str) -> str:
+            line = next(lines)
+            if line == "first output":
+                first_input_started.set()
+                assert release_first_input.wait(timeout=TEST_WATCHDOG_SECONDS)
+            return line
+
+        with patch("builtins.input", side_effect=operator_input) as mock_input:
+            reading = asyncio.create_task(context.processing_callback())
+            assert await asyncio.to_thread(first_input_started.wait, TEST_WATCHDOG_SECONDS)
+            cancelled_while_queued = asyncio.create_task(context.processing_callback())
+            queued = asyncio.create_task(context.processing_callback())
+            await asyncio.sleep(0)
+            cancelled_while_queued.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await cancelled_while_queued
+            # Give a read that wrongly skipped the queue the chance to claim the pending line.
+            for _ in range(10):
+                await asyncio.sleep(0)
+
+            release_first_input.set()
+            responses = await asyncio.wait_for(asyncio.gather(reading, queued), timeout=TEST_WATCHDOG_SECONDS)
+
+        assert responses == ["first output", "second output"]
+        assert mock_input.call_count == 2
+
+    async def test_concurrent_reads_receive_separate_lines_in_order(
+        self, manual_workflow: XPIAManualProcessingWorkflow, valid_message: Message
+    ) -> None:
+        context = XPIAContext(attack_content=valid_message)
+        await manual_workflow._setup_async(context=context)
+        assert context.processing_callback is not None
+        lines = iter(["first output", "second output"])
+
+        with patch("builtins.input", side_effect=lambda prompt: next(lines)):
+            responses = await asyncio.gather(context.processing_callback(), context.processing_callback())
+
+        assert responses == ["first output", "second output"]
+
+    async def test_failed_read_is_raised_and_not_reused(
+        self, manual_workflow: XPIAManualProcessingWorkflow, valid_message: Message
+    ) -> None:
+        context = XPIAContext(attack_content=valid_message)
+        await manual_workflow._setup_async(context=context)
+        assert context.processing_callback is not None
+
+        with patch("builtins.input", side_effect=[EOFError(), "processed output"]) as mock_input:
+            with pytest.raises(EOFError):
+                await context.processing_callback()
+            response = await context.processing_callback()
+
+        assert response == "processed output"
+        assert mock_input.call_count == 2