mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 15:37:55 +08:00
cli/interrupt: route SIGINT through a session handler
A SIGINT delivered while the event loop is parked in its selector is raised in the loop's own frame: it escapes run_until_complete, the Runner cancels the main task, and the REPL is left half-shut-down (no prompt, and the next line the human types is eaten by an abandoned prompt reader) until the following Ctrl+C kills the process. Off-loop prompts (ask_user_* / wait / tool confirms run in a worker thread via asyncio.to_thread) and the auto-retry countdown are the usual windows. Keep one asyncio-level SIGINT handler installed for the whole session instead: it cancels the innermost registered target (run_cancellable registers the turn task, sleep_cancellable the retry countdown, off_loop_prompt the waiting prompt) and never raises. Reads that block on the main thread inside a coroutine (the REPL entry prompt, synchronous click prompts) suspend the router and take SIGINT as a KeyboardInterrupt. Off-loop readline prompts cannot be interrupted without abandoning their reader thread on stdin, so they ignore SIGINT with a one-line hint.
This commit is contained in:
+208
-115
@@ -1,124 +1,67 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import signal
|
||||
import threading
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
|
||||
from plyngent.cli.interrupt import (
|
||||
allow_task_cancel,
|
||||
install_keyboard_interrupt_sigint,
|
||||
install_sigint_router,
|
||||
off_loop_prompt,
|
||||
pause_task_cancel_for_prompt,
|
||||
raise_keyboard_interrupt_sigint,
|
||||
run_in_prompt_thread,
|
||||
set_sigint_reinstall,
|
||||
sigint_cancels,
|
||||
sigint_routed,
|
||||
uninstall_sigint_router,
|
||||
)
|
||||
from plyngent.cli.limits import prompt_continue_limit, prompt_continue_limit_async
|
||||
from plyngent.cli.retry import run_cancellable
|
||||
from plyngent.cli.limits import prompt_continue_limit
|
||||
from plyngent.cli.retry import run_cancellable, sleep_cancellable
|
||||
from plyngent.prompting import temporary_backend
|
||||
from tests.test_prompting import ScriptedBackend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
import pytest as pytest_module
|
||||
|
||||
|
||||
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
|
||||
class FakeTarget:
|
||||
"""Minimal ``InterruptTarget`` recording its cancels."""
|
||||
|
||||
def __init__(self, log: list[str], name: str = "cancel") -> None:
|
||||
self.log = log
|
||||
self.name = name
|
||||
|
||||
def cancel(self) -> None:
|
||||
self.log.append(self.name)
|
||||
|
||||
|
||||
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:
|
||||
backend = ScriptedBackend(["y"])
|
||||
|
||||
def _read_line(prompt: str, *, default: str | None = None, completions: object = None) -> str:
|
||||
del prompt, default, completions
|
||||
assert allow_task_cancel() is False
|
||||
return "y"
|
||||
|
||||
monkeypatch.setattr(backend, "read_line", _read_line)
|
||||
with temporary_backend(backend):
|
||||
assert prompt_continue_limit("too many rounds") is True
|
||||
|
||||
|
||||
async def test_run_in_prompt_thread_pauses_cancel() -> None:
|
||||
"""Cancel is paused on the main thread for the whole to_thread call."""
|
||||
assert allow_task_cancel() is True
|
||||
|
||||
def work() -> str:
|
||||
return "ok"
|
||||
|
||||
result = await run_in_prompt_thread(work)
|
||||
assert result == "ok"
|
||||
assert allow_task_cancel() is True
|
||||
|
||||
|
||||
async def test_prompt_continue_limit_async(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
del monkeypatch
|
||||
with temporary_backend(ScriptedBackend(["y"])):
|
||||
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()
|
||||
@pytest.fixture
|
||||
async def router() -> AsyncIterator[None]:
|
||||
"""Install the session SIGINT router; restore the CLI handler afterwards."""
|
||||
install_keyboard_interrupt_sigint()
|
||||
_ = install_sigint_router()
|
||||
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")
|
||||
yield
|
||||
finally:
|
||||
uninstall_sigint_router()
|
||||
signal.signal(signal.SIGINT, signal.default_int_handler)
|
||||
|
||||
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
|
||||
async def _deliver_sigint() -> None:
|
||||
"""Send SIGINT to this process and let the loop run its signal callback.
|
||||
|
||||
asyncio's signal wakeup is only noticed while the loop parks in its
|
||||
selector, so a real (short) wait is needed — ``sleep(0)`` never parks and
|
||||
the callback would not run at all.
|
||||
"""
|
||||
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)
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
|
||||
def test_install_keyboard_interrupt_sigint() -> None:
|
||||
@@ -148,8 +91,6 @@ def test_asyncio_run_keeps_keyboard_interrupt_sigint() -> None:
|
||||
KeyboardInterrupt, which also cancelled ``memory.close()`` on a Ctrl+D exit.
|
||||
With a non-default handler installed first, Runner leaves SIGINT alone.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
install_keyboard_interrupt_sigint()
|
||||
try:
|
||||
|
||||
@@ -162,28 +103,180 @@ def test_asyncio_run_keeps_keyboard_interrupt_sigint() -> None:
|
||||
signal.signal(signal.SIGINT, signal.default_int_handler)
|
||||
|
||||
|
||||
async def test_run_cancellable_restores_previous_sigint() -> None:
|
||||
"""Regression: after a turn SIGINT is restored, not left as SIG_DFL.
|
||||
async def test_sigint_cancels_registered_target(router: None) -> None:
|
||||
"""SIGINT cancels the target registered around the awaited turn task."""
|
||||
del router
|
||||
cancelled = asyncio.Event()
|
||||
|
||||
async def hang() -> None:
|
||||
try:
|
||||
await asyncio.sleep(3600)
|
||||
except asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
raise
|
||||
|
||||
task = asyncio.create_task(hang())
|
||||
await asyncio.sleep(0)
|
||||
with sigint_cancels(task):
|
||||
os.kill(os.getpid(), signal.SIGINT)
|
||||
# Waiting on the effect (not on a fixed sleep) keeps this deterministic.
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=5)
|
||||
assert cancelled.is_set()
|
||||
|
||||
|
||||
async def test_sigint_while_loop_is_idle_does_not_escape(router: None) -> None:
|
||||
"""Regression: a stray Ctrl+C while the loop is parked in its selector.
|
||||
|
||||
The interrupt used to be raised in the loop's own frame (not inside any
|
||||
coroutine), escaping ``run_until_complete``: the Runner cancelled the main
|
||||
task and the REPL was left half-shut-down, so the next Ctrl+C killed the
|
||||
process instead of returning to the prompt.
|
||||
"""
|
||||
del router
|
||||
|
||||
def _send() -> None:
|
||||
time.sleep(0.05)
|
||||
os.kill(os.getpid(), signal.SIGINT)
|
||||
|
||||
sender = threading.Thread(target=_send)
|
||||
sender.start()
|
||||
try:
|
||||
await asyncio.sleep(0.2) # parked in the selector when SIGINT lands
|
||||
finally:
|
||||
sender.join()
|
||||
task = asyncio.current_task()
|
||||
assert task is not None
|
||||
assert task.cancelling() == 0
|
||||
assert not task.cancelled()
|
||||
|
||||
|
||||
async def test_sigint_cancels_innermost_target_only(router: None) -> None:
|
||||
"""Nested targets: the innermost (the active turn or prompt) wins."""
|
||||
del router
|
||||
outer: list[str] = []
|
||||
inner: list[str] = []
|
||||
with sigint_cancels(FakeTarget(outer, "outer")):
|
||||
await _deliver_sigint()
|
||||
assert outer == ["outer"]
|
||||
with sigint_cancels(FakeTarget(inner, "inner")):
|
||||
await _deliver_sigint()
|
||||
assert inner == ["inner"]
|
||||
assert outer == ["outer"] # the nested signal never reached it
|
||||
|
||||
|
||||
async def test_sigint_target_stops_after_context_exit(router: None) -> None:
|
||||
"""A popped target no longer sees SIGINT (and an empty stack ignores it)."""
|
||||
del router
|
||||
calls: list[str] = []
|
||||
with sigint_cancels(FakeTarget(calls)):
|
||||
await _deliver_sigint()
|
||||
assert calls == ["cancel"]
|
||||
await _deliver_sigint()
|
||||
assert calls == ["cancel"]
|
||||
|
||||
|
||||
async def test_off_loop_prompt_keeps_turn_alive(
|
||||
router: None,
|
||||
capsys: pytest_module.CaptureFixture[str],
|
||||
) -> None:
|
||||
"""A prompt's SIGINT target shields the turn; afterwards the turn cancels.
|
||||
|
||||
Regression (badcb3c, revisited): the router resolves the innermost target
|
||||
per signal instead of freezing cancel state into the handler, so a mid-turn
|
||||
prompt neither aborts the turn nor freezes turn-cancel.
|
||||
"""
|
||||
del router
|
||||
started = asyncio.Event()
|
||||
cancelled = asyncio.Event()
|
||||
|
||||
async def hang() -> None:
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.sleep(3600)
|
||||
except asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
raise
|
||||
|
||||
task = asyncio.create_task(hang())
|
||||
await started.wait()
|
||||
with sigint_cancels(task):
|
||||
with off_loop_prompt():
|
||||
await _deliver_sigint()
|
||||
assert not task.cancelling()
|
||||
err = capsys.readouterr().err
|
||||
assert "Ctrl+C ignored while a prompt is waiting" in err
|
||||
# The prompt returned (outer context still active): SIGINT cancels now.
|
||||
os.kill(os.getpid(), signal.SIGINT)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=5)
|
||||
assert cancelled.is_set()
|
||||
|
||||
|
||||
async def test_off_loop_prompt_cancel_hook_is_preferred(router: None) -> None:
|
||||
"""A poll-based prompt read gets cancelled instead of the enclosing turn."""
|
||||
del router
|
||||
hooks: list[str] = []
|
||||
turn_calls: list[str] = []
|
||||
with sigint_cancels(FakeTarget(turn_calls, "turn")), off_loop_prompt(lambda: hooks.append("read")):
|
||||
await _deliver_sigint()
|
||||
assert hooks == ["read"]
|
||||
assert turn_calls == []
|
||||
|
||||
|
||||
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
|
||||
routed = signal.getsignal(signal.SIGINT)
|
||||
assert routed is not raise_keyboard_interrupt_sigint
|
||||
with pause_task_cancel_for_prompt():
|
||||
assert signal.getsignal(signal.SIGINT) is raise_keyboard_interrupt_sigint
|
||||
with pause_task_cancel_for_prompt():
|
||||
assert signal.getsignal(signal.SIGINT) is raise_keyboard_interrupt_sigint
|
||||
# Inner exit must not resume the router while the outer read blocks.
|
||||
assert signal.getsignal(signal.SIGINT) is raise_keyboard_interrupt_sigint
|
||||
assert signal.getsignal(signal.SIGINT) is routed
|
||||
|
||||
|
||||
async def test_prompt_continue_limit_under_pause(
|
||||
router: None,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Synchronous prompts read on the main thread, where SIGINT is a KI."""
|
||||
del router
|
||||
backend = ScriptedBackend(["y"])
|
||||
|
||||
def _read_line(prompt: str, *, default: str | None = None, completions: object = None) -> str:
|
||||
del prompt, default, completions
|
||||
assert signal.getsignal(signal.SIGINT) is raise_keyboard_interrupt_sigint
|
||||
return "y"
|
||||
|
||||
monkeypatch.setattr(backend, "read_line", _read_line)
|
||||
with temporary_backend(backend):
|
||||
assert prompt_continue_limit("too many rounds") is True
|
||||
|
||||
|
||||
async def test_run_cancellable_keeps_router_installed(router: None) -> None:
|
||||
"""Regression: after a turn SIGINT must not be left as SIG_DFL.
|
||||
|
||||
``loop.remove_signal_handler`` leaves SIG_DFL, which would terminate the
|
||||
process on a stray Ctrl+C between turns instead of letting the REPL catch
|
||||
a KeyboardInterrupt.
|
||||
process on a stray Ctrl+C between turns; the session router stays installed.
|
||||
"""
|
||||
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")
|
||||
|
||||
previous = signal.getsignal(signal.SIGINT)
|
||||
del router
|
||||
|
||||
async def quick() -> None:
|
||||
return None
|
||||
|
||||
try:
|
||||
await run_cancellable(quick())
|
||||
assert signal.getsignal(signal.SIGINT) is previous
|
||||
assert signal.getsignal(signal.SIGINT) is not signal.SIG_DFL
|
||||
finally:
|
||||
signal.signal(signal.SIGINT, signal.default_int_handler)
|
||||
await run_cancellable(quick())
|
||||
assert sigint_routed()
|
||||
assert signal.getsignal(signal.SIGINT) is not signal.SIG_DFL
|
||||
assert signal.getsignal(signal.SIGINT) is not signal.default_int_handler
|
||||
|
||||
|
||||
async def test_sleep_cancellable_cancelled_by_sigint(router: None) -> None:
|
||||
"""The auto-retry countdown is cancelled by Ctrl+C (was a teardown window)."""
|
||||
del router
|
||||
waiter = asyncio.create_task(sleep_cancellable(30))
|
||||
await asyncio.sleep(0)
|
||||
os.kill(os.getpid(), signal.SIGINT)
|
||||
assert await asyncio.wait_for(waiter, timeout=5) is False
|
||||
|
||||
Reference in New Issue
Block a user