Skip to content

FIX: Read manual XPIA input off the event loop - #3775

Open
Lanre Shittu (Shizoqua) wants to merge 3 commits into
microsoft:mainfrom
Shizoqua:fix/xpia-manual-input-off-event-loop
Open

Lanre Shittu (Shizoqua) wants to merge 3 commits into
microsoft:mainfrom
Shizoqua:fix/xpia-manual-input-off-event-loop

Conversation

@Shizoqua

@Shizoqua Lanre Shittu (Shizoqua) commented Nov 1, 2026 •

Copy link
Copy Markdown
Contributor

Description

Fixes #3719.

XPIAManualProcessingWorkflow sets a processing callback that called input() directly inside an async function, so the event loop stopped for as long as the operator took to trigger the processing target and paste its output. Any other async work, such as other attacks, scorers or target calls, waited with it.

The callback now reads the console through a small _ConsoleLineReader:

  • Each read runs input() on a daemon thread, so the event loop keeps running and a pending read never delays interpreter shutdown.
  • Callers take turns, so concurrent workflows never receive the same line.
  • If a caller is cancelled while the operator is typing, the line goes to the next caller instead of being dropped. A caller cancelled while still queued passes its turn on only after the caller ahead of it finishes.
  • A read that fails, for example on EOF, raises to its caller and is not reused.

The prompt text and the returned value are unchanged.

Tests and Documentation

Added TestXPIAManualProcessingWorkflow to tests/unit/executor/workflow/test_xpia.py:

  • The callback returns the operator's input and shows the same prompt as before.
  • The patched input() only returns after a task on the event loop releases it, so the test fails if the read blocks the loop.
  • A read cancelled mid-input hands its line to the next read, and input() is called once.
  • Concurrent reads receive separate lines in order.
  • A queued read that is cancelled does not let a later read share the pending line.
  • A failed read is raised and not reused.

No docs or notebooks changed, so JupyText was not needed.

@hannahwestra25 hannahwestra25 self-assigned this Nov 2, 2026
@hannahwestra25

Copy link
Copy Markdown
Contributor

Cancelling this read (timeout, TaskGroup, Ctrl+C) leaves input() running in the default executor: asyncio.run hangs until Enter and the next paste is silently dropped (reproduced). Suggest a daemon-thread reader that serializes prompts and hands a cancelled read's line to the next caller, plus a regression test:

--- a/pyrit/executor/workflow/xpia.py
+++ b/pyrit/executor/workflow/xpia.py
@@ -3,5 +3,8 @@
 
 import asyncio
+import concurrent.futures
 import logging
+import queue
+import threading
 import uuid
 from dataclasses import dataclass, field
@@ -590,4 +593,46 @@
 
 
+class _ConsoleLineReader:
+    """
+    Read console lines on one daemon thread, one prompt at a time.
+
+    A caller cancelled mid-read hands the operator's line to the next caller instead of dropping it,
+    and the daemon thread never delays event-loop or interpreter shutdown.
+    """
+
+    def __init__(self) -> None:
+        self._requests: queue.SimpleQueue[tuple[str, concurrent.futures.Future[str]]] = queue.SimpleQueue()
+        self._start_lock = threading.Lock()
+        self._thread: threading.Thread | None = None
+
+    async def read_line_async(self, *, prompt: str) -> str:
+        with self._start_lock:
+            if self._thread is None:
+                self._thread = threading.Thread(target=self._serve_requests, name="pyrit-console-input", daemon=True)
+                self._thread.start()
+        line: concurrent.futures.Future[str] = concurrent.futures.Future()
+        self._requests.put((prompt, line))
+        return await asyncio.wrap_future(line)
+
+    def _serve_requests(self) -> None:
+        while True:
+            prompt, line = self._requests.get()
+            if line.cancelled():
+                continue
+            try:
+                outcome: str | Exception = input(prompt)
+            except Exception as exc:
+                outcome = exc
+            while not line.set_running_or_notify_cancel():
+                _, line = self._requests.get()
+            if isinstance(outcome, Exception):
+                line.set_exception(outcome)
+            else:
+                line.set_result(outcome)
+
+
+_console_line_reader = _ConsoleLineReader()
+
+
 class XPIAManualProcessingWorkflow(XPIAWorkflow):
     """
@@ -663,9 +708,9 @@
         """
 
-        # Create the manual input callback. The console read runs in a worker thread so other
+        # 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 await asyncio.to_thread(input, prompt)
+            return await _console_line_reader.read_line_async(prompt=prompt)
 
         # Set the processing callback on the context
--- a/tests/unit/executor/workflow/test_xpia.py
+++ b/tests/unit/executor/workflow/test_xpia.py
@@ -5,7 +5,9 @@
 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, wait_for_completion_async, wait_for_worker_ready_async
 from unit.mocks import get_mock_prompt_normalizer
 
@@ -16,4 +18,5 @@
     XPIAStatus,
     XPIAWorkflow,
+    _ConsoleLineReader,
 )
 from pyrit.memory import MemoryInterface
@@ -714,4 +717,9 @@
     """Tests for the manual processing callback created during setup."""
 
+    @pytest.fixture(autouse=True)
+    def fresh_console_reader(self) -> Iterator[None]:
+        with patch("pyrit.executor.workflow.xpia._console_line_reader", _ConsoleLineReader()):
+            yield
+
     async def test_setup_async_sets_callback_that_returns_operator_input(
         self, manual_workflow: XPIAManualProcessingWorkflow, valid_message: Message
@@ -749,2 +757,30 @@
 
         assert response == "processed output"
+
+    async def test_cancelled_manual_input_hands_the_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
+        loop = asyncio.get_running_loop()
+        prompted = asyncio.Event()
+        release_line = threading.Event()
+
+        def gated_input(prompt: str) -> str:
+            loop.call_soon_threadsafe(prompted.set)
+            assert release_line.wait(timeout=TEST_WATCHDOG_SECONDS), "Operator input was not released"
+            return "processed output"
+
+        with patch("builtins.input", side_effect=gated_input) as mock_input:
+            cancelled_read = asyncio.create_task(context.processing_callback())
+            await wait_for_worker_ready_async(started=prompted, operation=cancelled_read)
+            cancelled_read.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await cancelled_read
+
+            retried_read = asyncio.create_task(context.processing_callback())
+            release_line.set()
+            assert await wait_for_completion_async(future=retried_read) == "processed output"
+
+        mock_input.assert_called_once()

Copilot AI balanced review requested due to automatic review settings November 3, 2026 20:28

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Cancelling a queued waiter prematurely releases its turn, allowing two workflows to consume the same operator response and potentially misassociate input.

1 open issue
What changed in this PR

Addresses #3719 by moving manual XPIA console input off the event loop, allowing other PyRIT async work to continue while awaiting operator input.

Changes:

  • Adds a shared daemon-thread console reader with queued reads and pending-input reuse after cancellation.
  • Preserves the existing prompt and returned input.
  • Adds tests for responsiveness, cancellation, concurrent reads, and input errors.
File Description
tests/​unit/​executor/​workflow/​test_xpia.py Adds manual-processing callback and console-reader lifecycle tests.
pyrit/​executor/​workflow/​xpia.py Routes manual input through an asynchronous console-reader helper.

Review effort: Balanced

Comment thread pyrit/executor/workflow/xpia.py Resolved
@Shizoqua

Copy link
Copy Markdown
Contributor Author

Thanks for catching this, I reproduced it too. 4e1cd2a moves the console read to a daemon thread reader that serializes prompts and hands a cancelled read's line to the next caller instead of dropping it, with regression tests for the cancelled read, concurrent reads and a failed read.

@Shizoqua

Copy link
Copy Markdown
Contributor Author

The failing build-book check isn't from this PR. The same docs warnings show up when building from main, and this PR only touches xpia.py and its tests.

Comment on lines +719 to +722
with self._lock:
if self._unclaimed_line is None:
self._unclaimed_line = self._start_read(prompt=prompt)
line = self._unclaimed_line

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If a read is cancelled and the operator then answers its prompt, that line waits in _unclaimed_line; the next unrelated workflow receives it instantly, without prompting, and saves and scores it as its own result. Hand off only still-pending reads, re-prompting; discard completed ones.

Suggested change
with self._lock:
if self._unclaimed_line is None:
self._unclaimed_line = self._start_read(prompt=prompt)
line = self._unclaimed_line
with self._lock:
if self._unclaimed_line is not None and self._unclaimed_line.done():
# Nobody was waiting when this line arrived, so it answers a cancelled caller's prompt.
logger.warning("Discarding console input entered after its manual XPIA read was cancelled.")
self._unclaimed_line = None
if self._unclaimed_line is None:
self._unclaimed_line = self._start_read(prompt=prompt)
else:
print(prompt, flush=True)
line = self._unclaimed_line
Docstrings and tests
@@ class _ConsoleLineReader
     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.
+    shutdown. A pending read passes to the next caller, which shows its own prompt; a line typed
+    while no caller is waiting is discarded. Callers take turns, so concurrent reads never receive
+    the same line.
@@ _ConsoleLineReader.read_line_async
-            prompt (str): The prompt shown when a new read starts.
+            prompt (str): The prompt shown to the operator.
@@ _ConsoleLineReader.read_line_async
             asyncio.CancelledError: If the caller is cancelled before the line arrives. The read
-                keeps running and its line goes to the next caller.
+                keeps running; its line goes to a later caller only if one is waiting when it arrives.
@@ test_cancelled_read_hands_its_line_to_the_next_read
-        with patch("builtins.input", side_effect=operator_input) as mock_input:
+        prompt_shown_again = asyncio.Event()
+
+        with (
+            patch("builtins.input", side_effect=operator_input) as mock_input,
+            patch("builtins.print", side_effect=lambda *args, **kwargs: prompt_shown_again.set()),
+        ):
             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())
+            await asyncio.wait_for(prompt_shown_again.wait(), timeout=TEST_WATCHDOG_SECONDS)
             release_input.set()
    async def test_line_entered_after_its_read_was_cancelled_is_not_given_to_a_later_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_stale_line = threading.Event()
        lines = iter(["stale output", "fresh output"])

        def operator_input(prompt: str) -> str:
            line = next(lines)
            if line == "stale output":
                input_started.set()
                assert release_stale_line.wait(timeout=TEST_WATCHDOG_SECONDS)
            return line

        with patch("builtins.input", side_effect=operator_input) as mock_input:
            cancelled_read = asyncio.create_task(context.processing_callback())
            assert await asyncio.to_thread(input_started.wait, TEST_WATCHDOG_SECONDS)
            stale_reader = next(t for t in threading.enumerate() if t.name == "pyrit-console-reader")
            cancelled_read.cancel()
            with pytest.raises(asyncio.CancelledError):
                await cancelled_read
            release_stale_line.set()
            await asyncio.to_thread(stale_reader.join, TEST_WATCHDOG_SECONDS)

            response = await asyncio.wait_for(context.processing_callback(), timeout=TEST_WATCHDOG_SECONDS)

        assert stale_reader.daemon, "A pending console read must never delay interpreter shutdown"
        assert response == "fresh output"
        assert mock_input.call_count == 2

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Manual XPIA processing blocks the whole event loop while waiting for input

3 participants