Repository navigation
FIX: Read manual XPIA input off the event loop - #3775
Lanre Shittu (Shizoqua) wants to merge 3 commits into
Conversation
|
Cancelling this read (timeout, TaskGroup, Ctrl+C) leaves --- 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() |
There was a problem hiding this comment.
🟡 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
|
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. |
|
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. |
| with self._lock: | ||
| if self._unclaimed_line is None: | ||
| self._unclaimed_line = self._start_read(prompt=prompt) | ||
| line = self._unclaimed_line |
There was a problem hiding this comment.
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.
| 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
Description
Fixes #3719.
XPIAManualProcessingWorkflowsets a processing callback that calledinput()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:input()on a daemon thread, so the event loop keeps running and a pending read never delays interpreter shutdown.The prompt text and the returned value are unchanged.
Tests and Documentation
Added
TestXPIAManualProcessingWorkflowtotests/unit/executor/workflow/test_xpia.py:input()only returns after a task on the event loop releases it, so the test fails if the read blocks the loop.input()is called once.No docs or notebooks changed, so JupyText was not needed.