From a7d9c9879bf9e156b895a66f0f9865ab3b830063 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Wed, 15 Jul 2026 09:41:32 +0800 Subject: [PATCH] core/cli+agent: run TTY confirms off the event loop Max-rounds and destructive-tool prompts use asyncio.to_thread with cancel paused on the main thread so multi-tool turns no longer look like user cancel. Hooks may be sync or async. --- CLAUDE.md | 2 +- src/plyngent/agent/chat.py | 4 +- src/plyngent/agent/loop.py | 20 ++++---- src/plyngent/agent/tools.py | 10 +++- src/plyngent/cli/interrupt.py | 27 ++++++++--- src/plyngent/cli/limits.py | 78 +++++++++++++++++++++----------- src/plyngent/cli/retry.py | 5 +- src/plyngent/cli/state.py | 8 ++-- tests/test_agent/test_loop.py | 42 +++++++++++++++++ tests/test_cli/test_interrupt.py | 29 +++++++++++- 10 files changed, 170 insertions(+), 55 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 9320034..696e433 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -78,7 +78,7 @@ Click app + readline REPL. Entry: `plyngent` / `python -m plyngent`. - **`plyngent chat`**: provider/model selection (flags or interactive), SQLite sessions via config `[database]` (file DB under user data if unset/`:memory:`), sessions bound to workspace dir; resumes latest **for cwd/`--workspace`** by default (`--new` / `--session`). - Slash: `/history`, `/sessions`, `/resume`, `/rounds`, `/retry`, … - Explicit `/resume` or `--session` from another workspace prompts: **keep** session path, **update** binding to current, or **abort**. -- Failed/cancelled turns: not written to DB; Ctrl+C cancels the in-flight turn task; **TTY confirms** (max-rounds / destructive tools) pause cancel so prompts work; auto-retry 10s/20s/30s; `/retry` manual. +- Failed/cancelled turns: not written to DB; Ctrl+C cancels the in-flight turn task; **TTY confirms** (max-rounds / destructive tools) run off-loop via `asyncio.to_thread` + pause cancel so prompts do not abort the turn; auto-retry 10s/20s/30s; `/retry` manual. - **`plyngent providers`**: list config providers. - **`plyngent config path|edit`**: print or open config in `$EDITOR` (`shlex`-split, e.g. `codium --wait`). - If no providers and `$EDITOR` is set, chat/providers prompt to edit config then reload. diff --git a/src/plyngent/agent/chat.py b/src/plyngent/agent/chat.py index ac035ef..648186e 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -8,7 +8,7 @@ from .budget import DEFAULT_CONTEXT_MAX_CHARS, DEFAULT_TOOL_RESULT_MAX_CHARS from .loop import DEFAULT_MAX_ROUNDS, run_chat_loop if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable, Sequence + from collections.abc import AsyncIterator, Awaitable, Callable, Sequence from plyngent.lmproto.openai_compatible.model import AnyChatMessage from plyngent.memory import MemoryStore @@ -17,7 +17,7 @@ if TYPE_CHECKING: from .events import AgentEvent from .tools import ToolRegistry - type LimitContinueHook = Callable[[str], bool] + type LimitContinueHook = Callable[[str], bool | Awaitable[bool]] class ChatAgent: diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index 5a913d6..b1fe7d0 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import inspect from typing import TYPE_CHECKING, cast from msgspec import UNSET @@ -33,24 +34,24 @@ from .events import ( ) if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable, Sequence + from collections.abc import AsyncIterator, Awaitable, Callable, Sequence from plyngent.lmproto.openai_compatible.model import AnyChatMessage, AnyToolItem from .client import ChatClient from .tools import ToolRegistry - type LimitContinueHook = Callable[[str], bool] + type LimitContinueHook = Callable[[str], bool | Awaitable[bool]] DEFAULT_MAX_ROUNDS = 32 -def _raise_if_cancelled() -> None: - """Cooperative cancel points between model/tool steps.""" - task = asyncio.current_task() - if task is not None and task.cancelling(): - raise asyncio.CancelledError +async def _call_on_limit(on_limit: LimitContinueHook, reason: str) -> bool: + result = on_limit(reason) + if inspect.isawaitable(result): + return bool(await result) + return bool(result) async def _run_one_tool( @@ -59,7 +60,6 @@ async def _run_one_tool( *, max_result_chars: int, ) -> tuple[ToolChatMessage, ErrorEvent | None]: - _raise_if_cancelled() if isinstance(call, AssistantFunctionToolCall): try: result_text = await tools.execute(call.function.name, call.function.arguments) @@ -94,7 +94,6 @@ async def _execute_tool_calls( for call in tool_calls: yield ToolCallEvent(tool_call=call) - _raise_if_cancelled() if parallel and len(tool_calls) > 1: results = await asyncio.gather( *[_run_one_tool(tools, call, max_result_chars=max_result_chars) for call in tool_calls] @@ -197,7 +196,6 @@ async def run_chat_loop( while True: while rounds_used < allowance: rounds_used += 1 - _raise_if_cancelled() request_messages = compact_messages_for_request( messages, max_chars=max_context_chars, @@ -235,7 +233,7 @@ async def run_chat_loop( yield event reason = f"tool loop reached {allowance} rounds (used {rounds_used})" - if on_limit is not None and on_limit(reason): + if on_limit is not None and await _call_on_limit(on_limit, reason): yield MaxRoundsEvent(rounds=allowance, continued=True) allowance += max_rounds continue diff --git a/src/plyngent/agent/tools.py b/src/plyngent/agent/tools.py index 786c119..428f3ad 100644 --- a/src/plyngent/agent/tools.py +++ b/src/plyngent/agent/tools.py @@ -13,7 +13,10 @@ from plyngent.typedef import JSONSchema # noqa: TC001 type ToolHandler = Callable[..., Any | Awaitable[Any]] type DangerClassifier = Callable[[str, Mapping[str, object]], str | None] -type ToolConfirmHook = Callable[[str, Mapping[str, object], str], bool] +type ToolConfirmHook = Callable[ + [str, Mapping[str, object], str], + bool | Awaitable[bool], +] _PRIMITIVE_SCHEMA: dict[type, JSONSchema] = { str: {"type": "string"}, @@ -215,7 +218,10 @@ class ToolRegistry: reason = self._danger(name, args) if reason is None: return None - if self._on_confirm(name, args, reason): + allowed = self._on_confirm(name, args, reason) + if inspect.isawaitable(allowed): + allowed = await allowed + if allowed: return None return f"error: user denied tool {name!r} ({reason})" diff --git a/src/plyngent/cli/interrupt.py b/src/plyngent/cli/interrupt.py index cfcf289..bdbb77b 100644 --- a/src/plyngent/cli/interrupt.py +++ b/src/plyngent/cli/interrupt.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import contextlib import signal from contextlib import contextmanager @@ -30,7 +31,8 @@ def set_sigint_reinstall(callback: Callable[[], None] | None) -> None: 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 + 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. """ token = _allow_task_cancel.set(False) @@ -38,21 +40,34 @@ def pause_task_cancel_for_prompt() -> Generator[None]: 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) + try: + previous = signal.getsignal(signal.SIGINT) + _ = signal.signal(signal.SIGINT, signal.default_int_handler) + except ValueError: + # Not on the main thread — skip OS signal rebinding. + previous = signal.SIG_DFL yield finally: - _ = signal.signal(signal.SIGINT, previous) # type: ignore[arg-type] + with contextlib.suppress(ValueError): + _ = 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) + + +async def run_in_prompt_thread[**P, R](func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R: + """Run a blocking TTY prompt off the event loop with cancel paused. + + Pause/SIGINT rebinding happens on the main thread; only the prompt body + runs in a worker so the asyncio loop stays free and turn cancel is disabled. + """ + with pause_task_cancel_for_prompt(): + return await asyncio.to_thread(func, *args, **kwargs) diff --git a/src/plyngent/cli/limits.py b/src/plyngent/cli/limits.py index 5850a43..aaf6755 100644 --- a/src/plyngent/cli/limits.py +++ b/src/plyngent/cli/limits.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Literal import click -from plyngent.cli.interrupt import pause_task_cancel_for_prompt +from plyngent.cli.interrupt import pause_task_cancel_for_prompt, run_in_prompt_thread from plyngent.tools.process.pty_session import PtyManager if TYPE_CHECKING: @@ -13,35 +13,52 @@ if TYPE_CHECKING: type WorkspaceMismatchChoice = Literal["keep", "rebind", "abort"] -def prompt_continue_limit(reason: str) -> bool: - """Ask the user whether to raise a limit and continue (TTY).""" +def _prompt_continue_limit_sync(reason: str) -> bool: click.echo() click.secho(f"[limit] {reason}", fg="yellow") + try: + return bool(click.confirm("Raise limit and continue?", default=True)) + except (click.Abort, KeyboardInterrupt): + return False + + +def prompt_continue_limit(reason: str) -> bool: + """Ask the user whether to raise a limit and continue (TTY, sync).""" with pause_task_cancel_for_prompt(): - try: - return bool(click.confirm("Raise limit and continue?", default=True)) - except (click.Abort, KeyboardInterrupt): - return False + return _prompt_continue_limit_sync(reason) + + +async def prompt_continue_limit_async(reason: str) -> bool: + """Async variant: confirm off the event loop so the turn is not cancelled.""" + return await run_in_prompt_thread(_prompt_continue_limit_sync, reason) + + +def _prompt_confirm_tool_sync(name: str, args: Mapping[str, object], reason: str) -> bool: + del args + click.echo() + click.secho(f"[confirm] tool {name!r}: {reason}", fg="yellow") + try: + return bool(click.confirm("Allow this tool call?", default=False)) + except (click.Abort, KeyboardInterrupt): + return False def prompt_confirm_tool(name: str, args: Mapping[str, object], reason: str) -> bool: """Ask whether to allow a destructive tool call (TTY). Default is deny.""" - del args - click.echo() - click.secho(f"[confirm] tool {name!r}: {reason}", fg="yellow") with pause_task_cancel_for_prompt(): - try: - return bool(click.confirm("Allow this tool call?", default=False)) - except (click.Abort, KeyboardInterrupt): - return False + return _prompt_confirm_tool_sync(name, args, reason) -def prompt_workspace_mismatch( +async def prompt_confirm_tool_async(name: str, args: Mapping[str, object], reason: str) -> bool: + """Async variant: confirm off the event loop.""" + return await run_in_prompt_thread(_prompt_confirm_tool_sync, name, args, reason) + + +def _prompt_workspace_mismatch_sync( session_id: int, session_workspace: str, current_workspace: str, ) -> WorkspaceMismatchChoice: - """Ask how to handle resuming a session bound to a different directory.""" click.echo() click.secho(f"[workspace] session {session_id} is bound to a different directory:", fg="yellow") click.echo(f" session: {session_workspace}") @@ -49,16 +66,15 @@ def prompt_workspace_mismatch( click.echo(" k = keep session workspace (switch tools root to session path)") click.echo(" u = update binding to current workspace") click.echo(" a = abort resume") - with pause_task_cancel_for_prompt(): - try: - raw = click.prompt( - "Choice", - type=click.Choice(["k", "u", "a"], case_sensitive=False), - default="k", - show_choices=True, - ) - except (click.Abort, KeyboardInterrupt): - return "abort" + try: + raw = click.prompt( + "Choice", + type=click.Choice(["k", "u", "a"], case_sensitive=False), + default="k", + show_choices=True, + ) + except (click.Abort, KeyboardInterrupt): + return "abort" key = str(raw).strip().lower() if key == "u": return "rebind" @@ -67,6 +83,16 @@ def prompt_workspace_mismatch( return "keep" +def prompt_workspace_mismatch( + session_id: int, + session_workspace: str, + current_workspace: str, +) -> WorkspaceMismatchChoice: + """Ask how to handle resuming a session bound to a different directory.""" + with pause_task_cancel_for_prompt(): + return _prompt_workspace_mismatch_sync(session_id, session_workspace, current_workspace) + + def install_cli_limit_hooks() -> None: """Register interactive continue hooks for process-global tool limits.""" PtyManager.set_limit_continue_hook(prompt_continue_limit) diff --git a/src/plyngent/cli/retry.py b/src/plyngent/cli/retry.py index 1885edc..2efe61f 100644 --- a/src/plyngent/cli/retry.py +++ b/src/plyngent/cli/retry.py @@ -78,7 +78,10 @@ async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T: finally: set_sigint_reinstall(None) if installed: - _ = loop.remove_signal_handler(signal.SIGINT) + with contextlib.suppress(NotImplementedError, RuntimeError, ValueError): + _ = loop.remove_signal_handler(signal.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(): _ = task.cancel() with contextlib.suppress(asyncio.CancelledError): diff --git a/src/plyngent/cli/state.py b/src/plyngent/cli/state.py index d420f38..6024c88 100644 --- a/src/plyngent/cli/state.py +++ b/src/plyngent/cli/state.py @@ -49,7 +49,7 @@ class ReplState: def _tool_registry(self) -> ToolRegistry | None: if not self.tools_enabled: return None - from plyngent.cli.limits import prompt_confirm_tool + from plyngent.cli.limits import prompt_confirm_tool_async from plyngent.tools.danger import classify_danger agent_cfg = self.config.agent_config @@ -57,12 +57,12 @@ class ReplState: return ToolRegistry( list(DEFAULT_TOOLS), danger=classify_danger, - on_confirm=prompt_confirm_tool, + on_confirm=prompt_confirm_tool_async, ) return ToolRegistry(list(DEFAULT_TOOLS)) def _make_agent(self) -> ChatAgent: - from plyngent.cli.limits import prompt_continue_limit + from plyngent.cli.limits import prompt_continue_limit_async agent_cfg = self.config.agent_config system_prompt = agent_cfg.system_prompt or None @@ -73,7 +73,7 @@ class ReplState: memory=self.memory, session_id=self.session_id, max_rounds=self.max_rounds, - on_limit=prompt_continue_limit, + on_limit=prompt_continue_limit_async, stream=True, system_prompt=system_prompt, max_tool_result_chars=agent_cfg.max_tool_result_chars, diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py index 43451b3..42d9fad 100644 --- a/tests/test_agent/test_loop.py +++ b/tests/test_agent/test_loop.py @@ -280,6 +280,48 @@ async def test_max_rounds_continue_hook() -> None: assert len(client.calls) == 3 # noqa: PLR2004 +async def test_max_rounds_async_continue_hook() -> None: + @tool + def ping() -> str: + return "pong" + + registry = ToolRegistry([ping]) + forever = _response( + AssistantChatMessage( + content="", + tool_calls=[ + AssistantFunctionToolCall( + id="c", + function=AssistantFunctionTool(name="ping", arguments="{}"), + ) + ], + ) + ) + final = _response(AssistantChatMessage(content="done")) + client = ScriptedClient([forever, forever, final]) + messages: list[AnyChatMessage] = [UserChatMessage(content="x")] + asks: list[str] = [] + + async def on_limit(reason: str) -> bool: + asks.append(reason) + return True + + events = [ + e + async for e in run_chat_loop( + client, + messages, + model="m", + tools=registry, + max_rounds=2, + on_limit=on_limit, + ) + ] + assert len(asks) == 1 + assert any(isinstance(e, MaxRoundsEvent) and e.continued for e in events) + assert any(isinstance(e, TextDeltaEvent) and e.content == "done" for e in events) + + async def test_default_max_rounds_is_generous() -> None: from plyngent.agent.loop import DEFAULT_MAX_ROUNDS diff --git a/tests/test_cli/test_interrupt.py b/tests/test_cli/test_interrupt.py index f715c71..ac415d3 100644 --- a/tests/test_cli/test_interrupt.py +++ b/tests/test_cli/test_interrupt.py @@ -2,8 +2,12 @@ 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 +from plyngent.cli.interrupt import ( + allow_task_cancel, + pause_task_cancel_for_prompt, + run_in_prompt_thread, +) +from plyngent.cli.limits import prompt_continue_limit, prompt_continue_limit_async if TYPE_CHECKING: import pytest @@ -23,3 +27,24 @@ def test_prompt_continue_limit_under_pause(monkeypatch: pytest.MonkeyPatch) -> N monkeypatch.setattr("click.confirm", _confirm) 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" + + # 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 + + +async def test_prompt_continue_limit_async(monkeypatch: pytest.MonkeyPatch) -> None: + def _confirm(*_a: object, **_k: object) -> bool: + return True + + monkeypatch.setattr("click.confirm", _confirm) + assert await prompt_continue_limit_async("too many rounds") is True