mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/cli: pause turn cancel during TTY confirm prompts
Max-rounds and destructive-tool confirms no longer cancel the in-flight turn; SIGINT during the prompt aborts only the confirm.
This commit is contained in:
@@ -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)
|
||||||
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING
|
|||||||
|
|
||||||
import click
|
import click
|
||||||
|
|
||||||
|
from plyngent.cli.interrupt import pause_task_cancel_for_prompt
|
||||||
from plyngent.tools.process.pty_session import PtyManager
|
from plyngent.tools.process.pty_session import PtyManager
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
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)."""
|
"""Ask the user whether to raise a limit and continue (TTY)."""
|
||||||
click.echo()
|
click.echo()
|
||||||
click.secho(f"[limit] {reason}", fg="yellow")
|
click.secho(f"[limit] {reason}", fg="yellow")
|
||||||
try:
|
with pause_task_cancel_for_prompt():
|
||||||
return bool(click.confirm("Raise limit and continue?", default=True))
|
try:
|
||||||
except click.Abort:
|
return bool(click.confirm("Raise limit and continue?", default=True))
|
||||||
return False
|
except (click.Abort, KeyboardInterrupt):
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
def prompt_confirm_tool(name: str, args: Mapping[str, object], reason: str) -> bool:
|
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
|
del args
|
||||||
click.echo()
|
click.echo()
|
||||||
click.secho(f"[confirm] tool {name!r}: {reason}", fg="yellow")
|
click.secho(f"[confirm] tool {name!r}: {reason}", fg="yellow")
|
||||||
try:
|
with pause_task_cancel_for_prompt():
|
||||||
return bool(click.confirm("Allow this tool call?", default=False))
|
try:
|
||||||
except click.Abort:
|
return bool(click.confirm("Allow this tool call?", default=False))
|
||||||
return False
|
except (click.Abort, KeyboardInterrupt):
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
def install_cli_limit_hooks() -> None:
|
def install_cli_limit_hooks() -> None:
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING
|
|||||||
import click
|
import click
|
||||||
|
|
||||||
from plyngent.cli.display import render_events
|
from plyngent.cli.display import render_events
|
||||||
|
from plyngent.cli.interrupt import allow_task_cancel, set_sigint_reinstall
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator, Callable, Coroutine
|
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:
|
async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T:
|
||||||
"""Await ``coro`` as a task; Ctrl+C / SIGINT cancels the task.
|
"""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:
|
Raises:
|
||||||
asyncio.CancelledError: If the task was cancelled (including via SIGINT).
|
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()
|
loop = asyncio.get_running_loop()
|
||||||
installed = False
|
installed = False
|
||||||
|
|
||||||
def _on_sigint() -> None:
|
def _on_sigint() -> None:
|
||||||
if not task.done():
|
if allow_task_cancel() and not task.done():
|
||||||
_ = task.cancel()
|
_ = task.cancel()
|
||||||
|
|
||||||
|
def _reinstall() -> None:
|
||||||
|
try:
|
||||||
|
loop.add_signal_handler(signal.SIGINT, _on_sigint)
|
||||||
|
except (NotImplementedError, RuntimeError, ValueError):
|
||||||
|
return
|
||||||
|
|
||||||
try:
|
try:
|
||||||
loop.add_signal_handler(signal.SIGINT, _on_sigint)
|
loop.add_signal_handler(signal.SIGINT, _on_sigint)
|
||||||
installed = True
|
installed = True
|
||||||
|
set_sigint_reinstall(_reinstall)
|
||||||
except (NotImplementedError, RuntimeError, ValueError):
|
except (NotImplementedError, RuntimeError, ValueError):
|
||||||
installed = False
|
installed = False
|
||||||
|
|
||||||
@@ -65,6 +76,7 @@ async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T:
|
|||||||
await task
|
await task
|
||||||
raise asyncio.CancelledError from None
|
raise asyncio.CancelledError from None
|
||||||
finally:
|
finally:
|
||||||
|
set_sigint_reinstall(None)
|
||||||
if installed:
|
if installed:
|
||||||
_ = loop.remove_signal_handler(signal.SIGINT)
|
_ = loop.remove_signal_handler(signal.SIGINT)
|
||||||
if not task.done():
|
if not task.done():
|
||||||
|
|||||||
@@ -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
|
||||||
Reference in New Issue
Block a user