From 27bd7b46f7dde23eb1a2c8a442bcfb4beec969de Mon Sep 17 00:00:00 2001 From: worldmozara Date: Tue, 22 Sep 2026 19:28:03 +0800 Subject: [PATCH] prompting: let Ctrl+C cancel poll-based prompt reads The wait tool's prompt runs off-loop on a worker thread, but its read polls select/msvcrt, so it can act on a cancel: the SIGINT target now sets a flag the reader checks between polls and raises PromptCancelledError with, which the tool reports as "cancelled by user" while the turn continues. run_prompt_async / run_cancellable_prompt_async hand the matching hook to the CLI guard; a stale cancel never leaks into the next prompt read. --- src/plyngent/prompting.py | 90 +++++++++++++++++++++++++---- src/plyngent/tools/chat/wait.py | 17 ++++-- tests/test_cli/test_interrupt.py | 39 ++++++++++++- tests/test_prompting.py | 57 ++++++++++++++++++ tests/test_tools/test_chat_tools.py | 18 +++++- 5 files changed, 204 insertions(+), 17 deletions(-) diff --git a/src/plyngent/prompting.py b/src/plyngent/prompting.py index 3eeed51..561a2aa 100644 --- a/src/plyngent/prompting.py +++ b/src/plyngent/prompting.py @@ -4,6 +4,7 @@ import asyncio import contextlib import select import sys +import threading import time from contextlib import AbstractContextManager from dataclasses import dataclass @@ -19,6 +20,30 @@ class NonInteractiveError(RuntimeError): """Raised when a prompt is required but no interactive backend is available.""" +class PromptCancelledError(RuntimeError): + """Raised by a poll-based prompt read when the human cancels it (Ctrl+C). + + Only poll-based reads (``read_line_with_timeout``) can observe the cancel: + they check :func:`prompt_cancelled` between polls, so the reader thread + returns and nothing is left blocked on stdin. + """ + + +# Set by the SIGINT target registered for a poll-based prompt read (see +# cli.interrupt.off_loop_prompt); cleared around each off-loop prompt. +_prompt_cancel = threading.Event() + + +def prompt_cancelled() -> bool: + """Whether the in-flight poll-based prompt read was cancelled by SIGINT.""" + return _prompt_cancel.is_set() + + +def cancel_prompt_read() -> None: + """Ask the in-flight poll-based prompt read to return (thread-safe).""" + _prompt_cancel.set() + + @dataclass(frozen=True, slots=True) class ChoiceOption: """One selectable option for :func:`choose`.""" @@ -117,6 +142,8 @@ def _read_line_with_timeout_posix( Returns the line (without trailing newline), ``""`` for an empty Enter, or ``None`` when the timeout elapses before any line is complete (EOF too). + Raises :class:`PromptCancelledError` when :func:`prompt_cancelled` turns true + while waiting (the poll loop notices within one interval). """ _ = sys.stdout.write(prompt) _ = sys.stdout.flush() @@ -125,6 +152,8 @@ def _read_line_with_timeout_posix( remaining = deadline - time.monotonic() if remaining <= 0: return None + if prompt_cancelled(): + raise PromptCancelledError ready, _, _ = poll([stream], [], [], min(remaining, 0.1)) if not ready: continue @@ -143,6 +172,8 @@ def _read_line_with_timeout_windows(prompt: str, timeout: float) -> str | None: deadline = time.monotonic() + timeout buf: list[str] = [] while time.monotonic() < deadline: + if prompt_cancelled(): + raise PromptCancelledError if msvcrt.kbhit(): ch = msvcrt.getwch() if ch in ("\r", "\n"): @@ -171,7 +202,9 @@ def read_line_with_timeout(prompt: str, timeout: float) -> str | None: Returns the typed line (without trailing newline), ``""`` for an empty Enter, or ``None`` when the timeout elapses first (EOF too). Ctrl+C - propagates as :class:`KeyboardInterrupt`. POSIX polls stdin via ``select``; + propagates as :class:`KeyboardInterrupt`; :func:`cancel_prompt_read` (the + SIGINT target of an off-loop prompt read) raises + :class:`PromptCancelledError` instead. POSIX polls stdin via ``select``; Windows uses ``msvcrt`` char polling (best-effort). """ if timeout < 0: @@ -263,16 +296,18 @@ class NonInteractiveBackend: _backend: PromptBackend = ClickPromptBackend() -_pause_factory: Callable[[], AbstractContextManager[None]] | None = None +# CLI hook: context manager around an off-loop prompt. It receives the read's +# cancel hook (poll-based reads) or None (reads that cannot be interrupted). +_pause_factory: Callable[[Callable[[], None] | None], AbstractContextManager[None]] | None = None _prompt_lock = asyncio.Lock() def configure_prompting( *, backend: PromptBackend | None = None, - pause_factory: Callable[[], AbstractContextManager[None]] | None = None, + pause_factory: Callable[[Callable[[], None] | None], AbstractContextManager[None]] | None = None, ) -> None: - """Install process-wide prompt backend and optional cancel-pause context.""" + """Install process-wide prompt backend and optional SIGINT guard for prompts.""" global _backend, _pause_factory # noqa: PLW0603 if backend is not None: _backend = backend @@ -284,10 +319,11 @@ def get_prompt_backend() -> PromptBackend: def reset_prompting() -> None: - """Restore default Click backend and clear pause hook (tests).""" + """Restore default Click backend, clear the pause hook and cancel flag (tests).""" global _backend, _pause_factory # noqa: PLW0603 _backend = ClickPromptBackend() _pause_factory = None + _prompt_cancel.clear() def _normalize_options( @@ -464,13 +500,45 @@ def form( backend.echo("Starting over…") -async def run_prompt_async[**P, R](func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R: - """Run a blocking prompt off the event loop, serialized, with optional SIGINT pause.""" +async def _prompt_off_loop[R]( + func: Callable[..., R], + *args: object, + cancel_hook: Callable[[], None] | None, + **kwargs: object, +) -> R: + """Shared body: run *func* off the event loop under the CLI's SIGINT guard.""" async with _prompt_lock: - if _pause_factory is not None: - with _pause_factory(): - return await asyncio.to_thread(func, *args, **kwargs) - return await asyncio.to_thread(func, *args, **kwargs) + _prompt_cancel.clear() + try: + if _pause_factory is not None: + with _pause_factory(cancel_hook): + return await asyncio.to_thread(func, *args, **kwargs) + return await asyncio.to_thread(func, *args, **kwargs) + finally: + _prompt_cancel.clear() + + +async def run_prompt_async[**P, R](func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R: + """Run a blocking prompt off the event loop, serialized, SIGINT-guarded. + + While the prompt waits, a Ctrl+C must not cancel the enclosing turn: the CLI + guard registers the prompt as the SIGINT target. readline-based reads cannot + be interrupted without abandoning the reader thread on stdin, so they ignore + the interrupt with a one-line hint (see :func:`run_cancellable_prompt_async` + for poll-based reads). + """ + return await _prompt_off_loop(func, *args, cancel_hook=None, **kwargs) + + +async def run_cancellable_prompt_async[**P, R](func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R: + """Run a *poll-based* prompt read off the event loop; Ctrl+C cancels it. + + Only for reads that observe :func:`prompt_cancelled` + (``read_line_with_timeout``): the SIGINT raises :class:`PromptCancelledError` + in the reader thread, so it returns to the caller instead of leaving a thread + blocked on stdin (and without cancelling the enclosing turn). + """ + return await _prompt_off_loop(func, *args, cancel_hook=cancel_prompt_read, **kwargs) async def ask_async(prompt: str, *, default: str | None = None) -> str: diff --git a/src/plyngent/tools/chat/wait.py b/src/plyngent/tools/chat/wait.py index bf29b35..dfc9a75 100644 --- a/src/plyngent/tools/chat/wait.py +++ b/src/plyngent/tools/chat/wait.py @@ -7,7 +7,12 @@ import sys from click import style from plyngent.agent import ToolTag, tool -from plyngent.prompting import get_prompt_backend, read_line_with_timeout, run_prompt_async +from plyngent.prompting import ( + PromptCancelledError, + get_prompt_backend, + read_line_with_timeout, + run_cancellable_prompt_async, +) from plyngent.tools.chat.shape import first_error, number_arg, string_error _ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;]*m") @@ -36,7 +41,10 @@ async def _interactive_wait(seconds: float, reason: str | None) -> str: """Show the two-line wait prompt and report how the wait ended.""" label = f"{seconds:g}" prompt = _wait_prompt(seconds, reason=reason) - line = await run_prompt_async(read_line_with_timeout, prompt, seconds) + try: + line = await run_cancellable_prompt_async(read_line_with_timeout, prompt, seconds) + except PromptCancelledError: + return "cancelled by user" if line is None: return f"waited {label}s" text = line.strip() @@ -51,8 +59,9 @@ async def wait(duration: int, *, reason: str | None = None) -> str: Interactive sessions show a two-line prompt: a status line, then an optional-reason input — pressing Enter (optionally after typing a reason) - "disturbs" the wait so the turn continues immediately. Non-interactive runs - simply sleep the full duration. + "disturbs" the wait so the turn continues immediately, and Ctrl+C cancels + the wait (the turn continues). Non-interactive runs simply sleep the full + duration. """ seconds, duration_error = number_arg("wait", "duration", duration, example="5") reason_error = None if reason is None else string_error("wait", "reason", reason) diff --git a/tests/test_cli/test_interrupt.py b/tests/test_cli/test_interrupt.py index 1e0db9a..d09e4e6 100644 --- a/tests/test_cli/test_interrupt.py +++ b/tests/test_cli/test_interrupt.py @@ -21,7 +21,13 @@ from plyngent.cli.interrupt import ( ) from plyngent.cli.limits import prompt_continue_limit from plyngent.cli.retry import run_cancellable, sleep_cancellable -from plyngent.prompting import temporary_backend +from plyngent.prompting import ( + PromptCancelledError, + configure_prompting, + prompt_cancelled, + run_cancellable_prompt_async, + temporary_backend, +) from tests.test_prompting import ScriptedBackend if TYPE_CHECKING: @@ -224,6 +230,37 @@ async def test_off_loop_prompt_cancel_hook_is_preferred(router: None) -> None: assert turn_calls == [] +async def test_sigint_cancels_off_loop_prompt_read(router: None) -> None: + """Wiring: SIGINT aborts a poll-based read (not the turn) and clears the flag.""" + del router + turn_calls: list[str] = [] + seen: list[str] = [] + + def _poll_read(prompt: str, timeout: float) -> str: + """Stand-in for ``read_line_with_timeout``: polls the cancel flag.""" + del prompt + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if prompt_cancelled(): + seen.append("cancelled") + raise PromptCancelledError + time.sleep(0.01) + seen.append("timed out") + return "" + + with temporary_backend(ScriptedBackend([])): + configure_prompting(pause_factory=off_loop_prompt) + with sigint_cancels(FakeTarget(turn_calls, "turn")): + reader = asyncio.create_task(run_cancellable_prompt_async(_poll_read, "p> ", 30.0)) + await asyncio.sleep(0.05) + os.kill(os.getpid(), signal.SIGINT) + with pytest.raises(PromptCancelledError): + await asyncio.wait_for(reader, timeout=5) + assert seen == ["cancelled"] + assert turn_calls == [] + assert not prompt_cancelled() # cleared for the next prompt + + async def test_pause_task_cancel_for_prompt_swaps_sigint(router: None) -> None: """A main-thread blocking read takes SIGINT as a KeyboardInterrupt.""" del router diff --git a/tests/test_prompting.py b/tests/test_prompting.py index 0280c43..3708173 100644 --- a/tests/test_prompting.py +++ b/tests/test_prompting.py @@ -1,6 +1,8 @@ from __future__ import annotations +import contextlib import io +from typing import TYPE_CHECKING import pytest @@ -9,15 +11,33 @@ from plyngent.prompting import ( FormField, NonInteractiveBackend, NonInteractiveError, + PromptCancelledError, _read_line_with_timeout_posix, ask, + cancel_prompt_read, choose, + configure_prompting, confirm, form, + prompt_cancelled, read_line_with_timeout, + reset_prompting, + run_cancellable_prompt_async, + run_prompt_async, temporary_backend, ) +if TYPE_CHECKING: + from collections.abc import Generator, Iterator + + +@pytest.fixture(autouse=True) +def reset_prompting_state() -> Iterator[None]: + """Prompting state (backend, SIGINT guard, cancel flag) is process-global.""" + reset_prompting() + yield + reset_prompting() + def _ready_poll(r, w, x, t): return (r, w, x) @@ -51,6 +71,43 @@ def test_read_line_with_timeout_negative_timeout() -> None: assert read_line_with_timeout("p> ", -1) is None +def test_read_line_with_timeout_cancelled() -> None: + """A cancel request (SIGINT during an off-loop prompt) ends the read.""" + stream = io.StringIO() + cancel_prompt_read() + try: + assert prompt_cancelled() + with pytest.raises(PromptCancelledError): + _read_line_with_timeout_posix("p> ", 5.0, stream=stream, poll=_never_poll) + finally: + stream.close() + + +async def test_run_prompt_async_clears_cancel_flag() -> None: + """A stale cancel request must not leak into the next prompt read.""" + cancel_prompt_read() + with temporary_backend(NonInteractiveBackend()): + assert await run_prompt_async(lambda: "ok") == "ok" + assert not prompt_cancelled() + + +async def test_run_prompt_async_passes_cancel_hook_only_when_cancellable() -> None: + """The CLI guard gets the read's cancel hook for poll-based reads only.""" + seen: list[object] = [] + + @contextlib.contextmanager + def _guard(cancel: object) -> Generator[None]: + seen.append(cancel) + yield + + with temporary_backend(NonInteractiveBackend()): + # temporary_backend resets the guard, so install it inside. + configure_prompting(pause_factory=_guard) + assert await run_cancellable_prompt_async(lambda: "a") == "a" + assert await run_prompt_async(lambda: "b") == "b" + assert seen == [cancel_prompt_read, None] + + class ScriptedBackend: """Deterministic backend for unit tests.""" diff --git a/tests/test_tools/test_chat_tools.py b/tests/test_tools/test_chat_tools.py index c5ee950..c21fad8 100644 --- a/tests/test_tools/test_chat_tools.py +++ b/tests/test_tools/test_chat_tools.py @@ -3,8 +3,10 @@ from __future__ import annotations import importlib import json +import pytest + from plyngent.agent import ToolRegistry -from plyngent.prompting import NonInteractiveBackend, temporary_backend +from plyngent.prompting import NonInteractiveBackend, PromptCancelledError, temporary_backend from plyngent.tools.chat import CHAT_TOOLS, ask_user, choose_user, form_user, wait from tests.test_prompting import ScriptedBackend @@ -146,6 +148,20 @@ async def test_wait_tool_times_out(monkeypatch) -> None: assert out == "waited 3s" +async def test_wait_tool_cancelled_by_user(monkeypatch: pytest.MonkeyPatch) -> None: + """Ctrl+C during the wait prompt cancels the wait; the turn continues.""" + + def _cancelled(prompt: str, timeout: float) -> str: + del prompt, timeout + raise PromptCancelledError + + monkeypatch.setattr(wait_module, "read_line_with_timeout", _cancelled) + with temporary_backend(ScriptedBackend([])): + registry = ToolRegistry([wait]) + out = await registry.execute("wait", '{"duration": 30}') + assert out == "cancelled by user" + + async def test_wait_tool_prompt_two_lines_with_reason(monkeypatch) -> None: prompts: list[str] = [] monkeypatch.setattr(