diff --git a/src/plyngent/cli/interrupt.py b/src/plyngent/cli/interrupt.py new file mode 100644 index 0000000..cfcf289 --- /dev/null +++ b/src/plyngent/cli/interrupt.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import contextlib +import signal +from contextlib import contextmanager +from contextvars import ContextVar +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + from types import FrameType + + type SigHandler = Callable[[int, FrameType | None], None] | int | None + +_allow_task_cancel: ContextVar[bool] = ContextVar("allow_task_cancel", default=True) +_reinstall_holder: list[Callable[[], None] | None] = [None] + + +def allow_task_cancel() -> bool: + """Whether the CLI SIGINT handler should cancel the in-flight turn task.""" + return _allow_task_cancel.get() + + +def set_sigint_reinstall(callback: Callable[[], None] | None) -> None: + """Register how to re-bind SIGINT after a TTY prompt (set by run_cancellable).""" + _reinstall_holder[0] = callback + + +@contextmanager +def pause_task_cancel_for_prompt() -> Generator[None]: + """Disable turn-task cancel during blocking TTY prompts (confirm, etc.). + + Also restores the default SIGINT handler so ``click.confirm`` can receive + KeyboardInterrupt / Abort instead of the asyncio turn being cancelled. + """ + token = _allow_task_cancel.set(False) + loop_handler_removed = False + previous: SigHandler = signal.SIG_DFL + try: + try: + import asyncio + + loop = asyncio.get_running_loop() + _ = loop.remove_signal_handler(signal.SIGINT) + loop_handler_removed = True + except (RuntimeError, NotImplementedError, ValueError): + loop_handler_removed = False + + previous = signal.getsignal(signal.SIGINT) + _ = signal.signal(signal.SIGINT, signal.default_int_handler) + yield + finally: + _ = signal.signal(signal.SIGINT, previous) # type: ignore[arg-type] + reinstall = _reinstall_holder[0] + if loop_handler_removed and reinstall is not None: + with contextlib.suppress(RuntimeError, NotImplementedError, ValueError): + reinstall() + _allow_task_cancel.reset(token) diff --git a/src/plyngent/cli/limits.py b/src/plyngent/cli/limits.py index 64585e8..d171ed2 100644 --- a/src/plyngent/cli/limits.py +++ b/src/plyngent/cli/limits.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING import click +from plyngent.cli.interrupt import pause_task_cancel_for_prompt from plyngent.tools.process.pty_session import PtyManager if TYPE_CHECKING: @@ -14,10 +15,11 @@ def prompt_continue_limit(reason: str) -> bool: """Ask the user whether to raise a limit and continue (TTY).""" click.echo() click.secho(f"[limit] {reason}", fg="yellow") - try: - return bool(click.confirm("Raise limit and continue?", default=True)) - except click.Abort: - return False + with pause_task_cancel_for_prompt(): + try: + return bool(click.confirm("Raise limit and continue?", default=True)) + except (click.Abort, KeyboardInterrupt): + return False def prompt_confirm_tool(name: str, args: Mapping[str, object], reason: str) -> bool: @@ -25,10 +27,11 @@ def prompt_confirm_tool(name: str, args: Mapping[str, object], reason: str) -> b del args click.echo() click.secho(f"[confirm] tool {name!r}: {reason}", fg="yellow") - try: - return bool(click.confirm("Allow this tool call?", default=False)) - except click.Abort: - return False + with pause_task_cancel_for_prompt(): + try: + return bool(click.confirm("Allow this tool call?", default=False)) + except (click.Abort, KeyboardInterrupt): + return False def install_cli_limit_hooks() -> None: diff --git a/src/plyngent/cli/retry.py b/src/plyngent/cli/retry.py index b19d910..1885edc 100644 --- a/src/plyngent/cli/retry.py +++ b/src/plyngent/cli/retry.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING import click from plyngent.cli.display import render_events +from plyngent.cli.interrupt import allow_task_cancel, set_sigint_reinstall if TYPE_CHECKING: from collections.abc import AsyncIterator, Callable, Coroutine @@ -39,20 +40,30 @@ async def sleep_cancellable(seconds: float) -> bool: async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T: """Await ``coro`` as a task; Ctrl+C / SIGINT cancels the task. + Interactive prompts (max-rounds / destructive confirm) temporarily disable + task cancel so the user can answer instead of aborting the whole turn. + Raises: asyncio.CancelledError: If the task was cancelled (including via SIGINT). """ - task = asyncio.create_task(coro) + task: asyncio.Task[T] = asyncio.create_task(coro) loop = asyncio.get_running_loop() installed = False def _on_sigint() -> None: - if not task.done(): + if allow_task_cancel() and not task.done(): _ = task.cancel() + def _reinstall() -> None: + try: + loop.add_signal_handler(signal.SIGINT, _on_sigint) + except (NotImplementedError, RuntimeError, ValueError): + return + try: loop.add_signal_handler(signal.SIGINT, _on_sigint) installed = True + set_sigint_reinstall(_reinstall) except (NotImplementedError, RuntimeError, ValueError): installed = False @@ -65,6 +76,7 @@ async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T: await task raise asyncio.CancelledError from None finally: + set_sigint_reinstall(None) if installed: _ = loop.remove_signal_handler(signal.SIGINT) if not task.done(): diff --git a/tests/test_cli/test_interrupt.py b/tests/test_cli/test_interrupt.py new file mode 100644 index 0000000..f715c71 --- /dev/null +++ b/tests/test_cli/test_interrupt.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from plyngent.cli.interrupt import allow_task_cancel, pause_task_cancel_for_prompt +from plyngent.cli.limits import prompt_continue_limit + +if TYPE_CHECKING: + import pytest + + +def test_pause_task_cancel_for_prompt() -> None: + assert allow_task_cancel() is True + with pause_task_cancel_for_prompt(): + assert allow_task_cancel() is False + assert allow_task_cancel() is True + + +def test_prompt_continue_limit_under_pause(monkeypatch: pytest.MonkeyPatch) -> None: + def _confirm(*_a: object, **_k: object) -> bool: + assert allow_task_cancel() is False + return True + + monkeypatch.setattr("click.confirm", _confirm) + assert prompt_continue_limit("too many rounds") is True