mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 07:17:54 +08:00
cli: keep Ctrl+C a plain KeyboardInterrupt under asyncio.run
asyncio.run (Runner) replaces SIGINT with its own handler when the current one is signal.default_int_handler: the first Ctrl+C silently cancels the main task (invisible at the REPL prompt, leaving the task marked cancelled) and the second raises KeyboardInterrupt, so a Ctrl+D exit after a cancel was cancelled mid-memory.close() (SQLAlchemy CancelledError traceback) and the CLI died with "Aborted!". Install a distinct KeyboardInterrupt-raising handler before asyncio.run so Runner leaves SIGINT alone, and have run_cancellable restore the previous handler after each turn (remove_signal_handler leaves SIG_DFL, which would kill the process on a stray Ctrl+C between turns).
This commit is contained in:
@@ -17,6 +17,7 @@ from plyngent.cli.editor import (
|
||||
resolve_config_path,
|
||||
)
|
||||
from plyngent.cli.exit_codes import EXIT_CANCELLED, EXIT_OK, EXIT_TURN_FAILED
|
||||
from plyngent.cli.interrupt import install_keyboard_interrupt_sigint
|
||||
from plyngent.cli.limits import install_cli_limit_hooks
|
||||
from plyngent.cli.repl import run_repl
|
||||
from plyngent.cli.retry import run_user_text_with_retries
|
||||
@@ -484,6 +485,13 @@ def chat_cmd(
|
||||
raise click.ClickException(msg)
|
||||
|
||||
root = workspace if workspace is not None else Path.cwd()
|
||||
# asyncio.run installs its own SIGINT handler when the current one is the
|
||||
# default: the first Ctrl+C then silently cancels the main task (invisible
|
||||
# at the REPL prompt) and the second raises KeyboardInterrupt, which also
|
||||
# cancels memory.close() on a Ctrl+D exit ("Aborted!"). Pre-install a plain
|
||||
# KeyboardInterrupt-raising handler so Ctrl+C stays a normal KeyboardInterrupt
|
||||
# that the REPL / turn-retry / oneshot code handles deliberately.
|
||||
install_keyboard_interrupt_sigint()
|
||||
code = asyncio.run(
|
||||
_run_chat(
|
||||
config_path=config_path,
|
||||
|
||||
@@ -29,6 +29,33 @@ def set_sigint_reinstall(callback: Callable[[], None] | None) -> None:
|
||||
_reinstall_holder[0] = callback
|
||||
|
||||
|
||||
def raise_keyboard_interrupt_sigint(signum: int, frame: FrameType | None) -> None:
|
||||
"""SIGINT handler: raise KeyboardInterrupt (same effect as the default).
|
||||
|
||||
A distinct function object — deliberately *not*
|
||||
``signal.default_int_handler`` — so that :func:`asyncio.run`'s Runner does
|
||||
not recognise it and install its own silent-cancel handler on top.
|
||||
"""
|
||||
del signum, frame
|
||||
raise KeyboardInterrupt
|
||||
|
||||
|
||||
def install_keyboard_interrupt_sigint() -> None:
|
||||
"""Install :func:`raise_keyboard_interrupt_sigint` as the SIGINT handler.
|
||||
|
||||
``asyncio.run`` (Runner) replaces the SIGINT handler with its own when the
|
||||
current one *is* ``signal.default_int_handler``: the first Ctrl+C then
|
||||
silently cancels the main task (invisible at the REPL prompt, leaving the
|
||||
task marked cancelled) and only the second raises KeyboardInterrupt. A
|
||||
later Ctrl+D (EOF) is then cancelled mid-``memory.close()`` and the CLI
|
||||
exits with ``Aborted!``. Installing a plain KeyboardInterrupt-raising
|
||||
handler first keeps every Ctrl+C a normal KeyboardInterrupt that the CLI
|
||||
(prompt, turn-retry, oneshot) already handles deliberately.
|
||||
"""
|
||||
with contextlib.suppress(ValueError):
|
||||
_ = signal.signal(signal.SIGINT, raise_keyboard_interrupt_sigint)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def pause_task_cancel_for_prompt() -> Generator[None]:
|
||||
"""Disable turn-task cancel during blocking TTY prompts (confirm, etc.).
|
||||
|
||||
@@ -64,6 +64,11 @@ async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T:
|
||||
task: asyncio.Task[T] = asyncio.create_task(coro)
|
||||
loop = asyncio.get_running_loop()
|
||||
installed = False
|
||||
# What to restore once the turn-cancel handler is removed. remove_signal_handler
|
||||
# leaves SIG_DFL, which would terminate the process on a stray Ctrl+C between
|
||||
# turns; restore the handler from before (the CLI installs a
|
||||
# KeyboardInterrupt-raising one so the REPL can catch it).
|
||||
previous_sigint = signal.getsignal(signal.SIGINT)
|
||||
|
||||
def _on_sigint() -> None:
|
||||
# allow_task_cancel() uses a process-level depth counter (not ContextVar)
|
||||
@@ -98,6 +103,8 @@ async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T:
|
||||
if installed:
|
||||
with contextlib.suppress(NotImplementedError, RuntimeError, ValueError):
|
||||
_ = loop.remove_signal_handler(signal.SIGINT)
|
||||
with contextlib.suppress(ValueError):
|
||||
_ = signal.signal(signal.SIGINT, previous_sigint)
|
||||
# Only cancel if still running (e.g. KeyboardInterrupt path above).
|
||||
# Do not cancel a finished task — that would mask success.
|
||||
if not task.done():
|
||||
|
||||
@@ -8,7 +8,9 @@ import pytest
|
||||
|
||||
from plyngent.cli.interrupt import (
|
||||
allow_task_cancel,
|
||||
install_keyboard_interrupt_sigint,
|
||||
pause_task_cancel_for_prompt,
|
||||
raise_keyboard_interrupt_sigint,
|
||||
run_in_prompt_thread,
|
||||
set_sigint_reinstall,
|
||||
)
|
||||
@@ -113,3 +115,71 @@ async def test_sigint_cancels_after_prompt_pause() -> None:
|
||||
assert cancelled.is_set()
|
||||
# Cleanup any leftover reinstall hook from run_cancellable
|
||||
set_sigint_reinstall(None)
|
||||
|
||||
|
||||
def test_install_keyboard_interrupt_sigint() -> None:
|
||||
"""Installs a non-default SIGINT handler that raises KeyboardInterrupt.
|
||||
|
||||
The handler must be a distinct function object: ``asyncio.run`` only skips
|
||||
installing its own silent-cancel handler when the current handler is not
|
||||
``signal.default_int_handler``.
|
||||
"""
|
||||
install_keyboard_interrupt_sigint()
|
||||
try:
|
||||
handler = signal.getsignal(signal.SIGINT)
|
||||
assert handler is not signal.default_int_handler
|
||||
assert handler is not signal.SIG_DFL
|
||||
assert callable(handler)
|
||||
with pytest.raises(KeyboardInterrupt):
|
||||
handler(signal.SIGINT, None) # type: ignore[call-arg]
|
||||
finally:
|
||||
signal.signal(signal.SIGINT, signal.default_int_handler)
|
||||
|
||||
|
||||
def test_asyncio_run_keeps_keyboard_interrupt_sigint() -> None:
|
||||
"""Regression: Runner must not replace our handler with its own.
|
||||
|
||||
Previously the first Ctrl+C under ``asyncio.run`` silently cancelled the
|
||||
main task (invisible at the REPL prompt) and the second raised
|
||||
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:
|
||||
|
||||
async def quick() -> int:
|
||||
return 7
|
||||
|
||||
assert asyncio.run(quick()) == 7
|
||||
assert signal.getsignal(signal.SIGINT) is raise_keyboard_interrupt_sigint
|
||||
finally:
|
||||
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.
|
||||
|
||||
``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.
|
||||
"""
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user