diff --git a/src/plyngent/cli/interrupt.py b/src/plyngent/cli/interrupt.py index 52eb3cd..36e39a3 100644 --- a/src/plyngent/cli/interrupt.py +++ b/src/plyngent/cli/interrupt.py @@ -4,7 +4,6 @@ import asyncio import contextlib import signal from contextlib import contextmanager -from contextvars import ContextVar from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -13,13 +12,16 @@ if TYPE_CHECKING: type SigHandler = Callable[[int, FrameType | None], None] | int | None -_allow_task_cancel: ContextVar[bool] = ContextVar("allow_task_cancel", default=True) +# Depth of nested pause_task_cancel_for_prompt. Must be a plain int (not ContextVar): +# asyncio.add_signal_handler freezes the ContextVar snapshot at install time, so a +# ContextVar reset after reinstall would leave SIGINT permanently non-cancelling. +_prompt_depth: int = 0 _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() + return _prompt_depth == 0 def set_sigint_reinstall(callback: Callable[[], None] | None) -> None: @@ -34,8 +36,14 @@ def pause_task_cancel_for_prompt() -> Generator[None]: Must run on the main thread (signal handlers are main-thread only). Restores the default SIGINT handler so ``click.confirm`` can receive KeyboardInterrupt / Abort instead of the asyncio turn being cancelled. + + Nested prompts are supported via a depth counter. The asyncio SIGINT + handler is reinstalled only after depth returns to 0, and only after + cancel is re-enabled, so the handler never freezes ``allow=False``. """ - token = _allow_task_cancel.set(False) + + global _prompt_depth # noqa: PLW0603 + _prompt_depth += 1 loop_handler_removed = False previous: SigHandler = signal.SIG_DFL try: @@ -56,11 +64,15 @@ def pause_task_cancel_for_prompt() -> Generator[None]: finally: with contextlib.suppress(ValueError): _ = signal.signal(signal.SIGINT, previous) - 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) + _prompt_depth = max(0, _prompt_depth - 1) + # Reinstall only when fully out of prompts and cancel is allowed again. + # Order matters: decrement depth first so allow_task_cancel() is True + # before reinstall (and handlers never freeze allow=False). + if _prompt_depth == 0 and loop_handler_removed: + reinstall = _reinstall_holder[0] + if reinstall is not None: + with contextlib.suppress(RuntimeError, NotImplementedError, ValueError): + reinstall() async def run_in_prompt_thread[**P, R](func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R: diff --git a/src/plyngent/cli/retry.py b/src/plyngent/cli/retry.py index 1e1f180..0e4bafb 100644 --- a/src/plyngent/cli/retry.py +++ b/src/plyngent/cli/retry.py @@ -66,6 +66,9 @@ async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T: installed = False def _on_sigint() -> None: + # allow_task_cancel() uses a process-level depth counter (not ContextVar) + # so this remains correct even if the handler was installed under a + # frozen context (asyncio signal handles capture contextvars). if allow_task_cancel() and not task.done(): _ = task.cancel() diff --git a/tests/test_cli/test_interrupt.py b/tests/test_cli/test_interrupt.py index ac415d3..c8be061 100644 --- a/tests/test_cli/test_interrupt.py +++ b/tests/test_cli/test_interrupt.py @@ -1,16 +1,22 @@ from __future__ import annotations +import asyncio +import signal from typing import TYPE_CHECKING +import pytest + from plyngent.cli.interrupt import ( allow_task_cancel, pause_task_cancel_for_prompt, run_in_prompt_thread, + set_sigint_reinstall, ) from plyngent.cli.limits import prompt_continue_limit, prompt_continue_limit_async +from plyngent.cli.retry import run_cancellable if TYPE_CHECKING: - import pytest + pass def test_pause_task_cancel_for_prompt() -> None: @@ -20,6 +26,16 @@ def test_pause_task_cancel_for_prompt() -> None: assert allow_task_cancel() is True +def test_nested_pause_depth() -> None: + assert allow_task_cancel() is True + with pause_task_cancel_for_prompt(): + assert allow_task_cancel() is False + with pause_task_cancel_for_prompt(): + assert allow_task_cancel() is False + 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 @@ -36,7 +52,6 @@ async def test_run_in_prompt_thread_pauses_cancel() -> None: def work() -> str: return "ok" - # ContextVar may not propagate to worker threads; assert pause around the call. result = await run_in_prompt_thread(work) assert result == "ok" assert allow_task_cancel() is True @@ -48,3 +63,53 @@ async def test_prompt_continue_limit_async(monkeypatch: pytest.MonkeyPatch) -> N monkeypatch.setattr("click.confirm", _confirm) assert await prompt_continue_limit_async("too many rounds") is True + + +async def test_sigint_cancels_after_prompt_pause() -> None: + """Regression: after a mid-turn prompt, SIGINT must still cancel the turn. + + Previously, reinstalling the asyncio SIGINT handler while + allow_task_cancel was still False froze that value into the callback + forever (ContextVar snapshot at add_signal_handler). + """ + loop = asyncio.get_running_loop() + try: + loop.add_signal_handler(signal.SIGINT, lambda: None) + loop.remove_signal_handler(signal.SIGINT) + except NotImplementedError, RuntimeError, ValueError: + pytest.skip("asyncio signal handlers not available on this platform") + + started = asyncio.Event() + cancelled = asyncio.Event() + + async def hang() -> None: + started.set() + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + cancelled.set() + raise + + async def turn() -> None: + # Simulate soft-confirm pause mid-turn, then keep streaming. + with pause_task_cancel_for_prompt(): + assert allow_task_cancel() is False + assert allow_task_cancel() is True + await hang() + + task = asyncio.create_task(run_cancellable(turn())) + await started.wait() + # Give run_cancellable a tick to install the handler after the pause reinstall. + await asyncio.sleep(0) + assert allow_task_cancel() is True + # Deliver SIGINT the same way Ctrl+C does under asyncio. + loop.call_soon(lambda: None) # ensure loop is processing + # Invoke the process signal: raise SIGINT to this process. + import os + + os.kill(os.getpid(), signal.SIGINT) + with pytest.raises(asyncio.CancelledError): + await task + assert cancelled.is_set() + # Cleanup any leftover reinstall hook from run_cancellable + set_sigint_reinstall(None)