mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 15:37:55 +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,
|
resolve_config_path,
|
||||||
)
|
)
|
||||||
from plyngent.cli.exit_codes import EXIT_CANCELLED, EXIT_OK, EXIT_TURN_FAILED
|
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.limits import install_cli_limit_hooks
|
||||||
from plyngent.cli.repl import run_repl
|
from plyngent.cli.repl import run_repl
|
||||||
from plyngent.cli.retry import run_user_text_with_retries
|
from plyngent.cli.retry import run_user_text_with_retries
|
||||||
@@ -484,6 +485,13 @@ def chat_cmd(
|
|||||||
raise click.ClickException(msg)
|
raise click.ClickException(msg)
|
||||||
|
|
||||||
root = workspace if workspace is not None else Path.cwd()
|
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(
|
code = asyncio.run(
|
||||||
_run_chat(
|
_run_chat(
|
||||||
config_path=config_path,
|
config_path=config_path,
|
||||||
|
|||||||
@@ -29,6 +29,33 @@ def set_sigint_reinstall(callback: Callable[[], None] | None) -> None:
|
|||||||
_reinstall_holder[0] = callback
|
_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
|
@contextmanager
|
||||||
def pause_task_cancel_for_prompt() -> Generator[None]:
|
def pause_task_cancel_for_prompt() -> Generator[None]:
|
||||||
"""Disable turn-task cancel during blocking TTY prompts (confirm, etc.).
|
"""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)
|
task: asyncio.Task[T] = asyncio.create_task(coro)
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
installed = False
|
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:
|
def _on_sigint() -> None:
|
||||||
# allow_task_cancel() uses a process-level depth counter (not ContextVar)
|
# 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:
|
if installed:
|
||||||
with contextlib.suppress(NotImplementedError, RuntimeError, ValueError):
|
with contextlib.suppress(NotImplementedError, RuntimeError, ValueError):
|
||||||
_ = loop.remove_signal_handler(signal.SIGINT)
|
_ = 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).
|
# Only cancel if still running (e.g. KeyboardInterrupt path above).
|
||||||
# Do not cancel a finished task — that would mask success.
|
# Do not cancel a finished task — that would mask success.
|
||||||
if not task.done():
|
if not task.done():
|
||||||
|
|||||||
@@ -8,7 +8,9 @@ import pytest
|
|||||||
|
|
||||||
from plyngent.cli.interrupt import (
|
from plyngent.cli.interrupt import (
|
||||||
allow_task_cancel,
|
allow_task_cancel,
|
||||||
|
install_keyboard_interrupt_sigint,
|
||||||
pause_task_cancel_for_prompt,
|
pause_task_cancel_for_prompt,
|
||||||
|
raise_keyboard_interrupt_sigint,
|
||||||
run_in_prompt_thread,
|
run_in_prompt_thread,
|
||||||
set_sigint_reinstall,
|
set_sigint_reinstall,
|
||||||
)
|
)
|
||||||
@@ -113,3 +115,71 @@ async def test_sigint_cancels_after_prompt_pause() -> None:
|
|||||||
assert cancelled.is_set()
|
assert cancelled.is_set()
|
||||||
# Cleanup any leftover reinstall hook from run_cancellable
|
# Cleanup any leftover reinstall hook from run_cancellable
|
||||||
set_sigint_reinstall(None)
|
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