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.
This commit is contained in:
2026-07-15 09:41:32 +08:00
parent b80ebf783a
commit a7d9c9879b
10 changed files with 170 additions and 55 deletions
+1 -1
View File
@@ -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`). - **`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`, … - Slash: `/history`, `/sessions`, `/resume`, `/rounds`, `/retry`, …
- Explicit `/resume` or `--session` from another workspace prompts: **keep** session path, **update** binding to current, or **abort**. - 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 providers`**: list config providers.
- **`plyngent config path|edit`**: print or open config in `$EDITOR` (`shlex`-split, e.g. `codium --wait`). - **`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. - If no providers and `$EDITOR` is set, chat/providers prompt to edit config then reload.
+2 -2
View File
@@ -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 from .loop import DEFAULT_MAX_ROUNDS, run_chat_loop
if TYPE_CHECKING: 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.lmproto.openai_compatible.model import AnyChatMessage
from plyngent.memory import MemoryStore from plyngent.memory import MemoryStore
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
from .events import AgentEvent from .events import AgentEvent
from .tools import ToolRegistry from .tools import ToolRegistry
type LimitContinueHook = Callable[[str], bool] type LimitContinueHook = Callable[[str], bool | Awaitable[bool]]
class ChatAgent: class ChatAgent:
+9 -11
View File
@@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import inspect
from typing import TYPE_CHECKING, cast from typing import TYPE_CHECKING, cast
from msgspec import UNSET from msgspec import UNSET
@@ -33,24 +34,24 @@ from .events import (
) )
if TYPE_CHECKING: 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 plyngent.lmproto.openai_compatible.model import AnyChatMessage, AnyToolItem
from .client import ChatClient from .client import ChatClient
from .tools import ToolRegistry from .tools import ToolRegistry
type LimitContinueHook = Callable[[str], bool] type LimitContinueHook = Callable[[str], bool | Awaitable[bool]]
DEFAULT_MAX_ROUNDS = 32 DEFAULT_MAX_ROUNDS = 32
def _raise_if_cancelled() -> None: async def _call_on_limit(on_limit: LimitContinueHook, reason: str) -> bool:
"""Cooperative cancel points between model/tool steps.""" result = on_limit(reason)
task = asyncio.current_task() if inspect.isawaitable(result):
if task is not None and task.cancelling(): return bool(await result)
raise asyncio.CancelledError return bool(result)
async def _run_one_tool( async def _run_one_tool(
@@ -59,7 +60,6 @@ async def _run_one_tool(
*, *,
max_result_chars: int, max_result_chars: int,
) -> tuple[ToolChatMessage, ErrorEvent | None]: ) -> tuple[ToolChatMessage, ErrorEvent | None]:
_raise_if_cancelled()
if isinstance(call, AssistantFunctionToolCall): if isinstance(call, AssistantFunctionToolCall):
try: try:
result_text = await tools.execute(call.function.name, call.function.arguments) 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: for call in tool_calls:
yield ToolCallEvent(tool_call=call) yield ToolCallEvent(tool_call=call)
_raise_if_cancelled()
if parallel and len(tool_calls) > 1: if parallel and len(tool_calls) > 1:
results = await asyncio.gather( results = await asyncio.gather(
*[_run_one_tool(tools, call, max_result_chars=max_result_chars) for call in tool_calls] *[_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 True:
while rounds_used < allowance: while rounds_used < allowance:
rounds_used += 1 rounds_used += 1
_raise_if_cancelled()
request_messages = compact_messages_for_request( request_messages = compact_messages_for_request(
messages, messages,
max_chars=max_context_chars, max_chars=max_context_chars,
@@ -235,7 +233,7 @@ async def run_chat_loop(
yield event yield event
reason = f"tool loop reached {allowance} rounds (used {rounds_used})" 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) yield MaxRoundsEvent(rounds=allowance, continued=True)
allowance += max_rounds allowance += max_rounds
continue continue
+8 -2
View File
@@ -13,7 +13,10 @@ from plyngent.typedef import JSONSchema # noqa: TC001
type ToolHandler = Callable[..., Any | Awaitable[Any]] type ToolHandler = Callable[..., Any | Awaitable[Any]]
type DangerClassifier = Callable[[str, Mapping[str, object]], str | None] 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] = { _PRIMITIVE_SCHEMA: dict[type, JSONSchema] = {
str: {"type": "string"}, str: {"type": "string"},
@@ -215,7 +218,10 @@ class ToolRegistry:
reason = self._danger(name, args) reason = self._danger(name, args)
if reason is None: if reason is None:
return 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 None
return f"error: user denied tool {name!r} ({reason})" return f"error: user denied tool {name!r} ({reason})"
+18 -3
View File
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import contextlib import contextlib
import signal import signal
from contextlib import contextmanager 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]: 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.).
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. KeyboardInterrupt / Abort instead of the asyncio turn being cancelled.
""" """
token = _allow_task_cancel.set(False) token = _allow_task_cancel.set(False)
@@ -38,21 +40,34 @@ def pause_task_cancel_for_prompt() -> Generator[None]:
previous: SigHandler = signal.SIG_DFL previous: SigHandler = signal.SIG_DFL
try: try:
try: try:
import asyncio
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
_ = loop.remove_signal_handler(signal.SIGINT) _ = loop.remove_signal_handler(signal.SIGINT)
loop_handler_removed = True loop_handler_removed = True
except (RuntimeError, NotImplementedError, ValueError): except (RuntimeError, NotImplementedError, ValueError):
loop_handler_removed = False loop_handler_removed = False
try:
previous = signal.getsignal(signal.SIGINT) previous = signal.getsignal(signal.SIGINT)
_ = signal.signal(signal.SIGINT, signal.default_int_handler) _ = signal.signal(signal.SIGINT, signal.default_int_handler)
except ValueError:
# Not on the main thread — skip OS signal rebinding.
previous = signal.SIG_DFL
yield yield
finally: finally:
with contextlib.suppress(ValueError):
_ = signal.signal(signal.SIGINT, previous) # type: ignore[arg-type] _ = signal.signal(signal.SIGINT, previous) # type: ignore[arg-type]
reinstall = _reinstall_holder[0] reinstall = _reinstall_holder[0]
if loop_handler_removed and reinstall is not None: if loop_handler_removed and reinstall is not None:
with contextlib.suppress(RuntimeError, NotImplementedError, ValueError): with contextlib.suppress(RuntimeError, NotImplementedError, ValueError):
reinstall() reinstall()
_allow_task_cancel.reset(token) _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)
+36 -10
View File
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Literal
import click 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 from plyngent.tools.process.pty_session import PtyManager
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -13,35 +13,52 @@ if TYPE_CHECKING:
type WorkspaceMismatchChoice = Literal["keep", "rebind", "abort"] type WorkspaceMismatchChoice = Literal["keep", "rebind", "abort"]
def prompt_continue_limit(reason: str) -> bool: def _prompt_continue_limit_sync(reason: str) -> bool:
"""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")
with pause_task_cancel_for_prompt():
try: try:
return bool(click.confirm("Raise limit and continue?", default=True)) return bool(click.confirm("Raise limit and continue?", default=True))
except (click.Abort, KeyboardInterrupt): except (click.Abort, KeyboardInterrupt):
return False return False
def prompt_confirm_tool(name: str, args: Mapping[str, object], reason: str) -> bool: def prompt_continue_limit(reason: str) -> bool:
"""Ask whether to allow a destructive tool call (TTY). Default is deny.""" """Ask the user whether to raise a limit and continue (TTY, sync)."""
with pause_task_cancel_for_prompt():
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 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")
with pause_task_cancel_for_prompt():
try: try:
return bool(click.confirm("Allow this tool call?", default=False)) return bool(click.confirm("Allow this tool call?", default=False))
except (click.Abort, KeyboardInterrupt): except (click.Abort, KeyboardInterrupt):
return False return False
def prompt_workspace_mismatch( 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."""
with pause_task_cancel_for_prompt():
return _prompt_confirm_tool_sync(name, args, reason)
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_id: int,
session_workspace: str, session_workspace: str,
current_workspace: str, current_workspace: str,
) -> WorkspaceMismatchChoice: ) -> WorkspaceMismatchChoice:
"""Ask how to handle resuming a session bound to a different directory."""
click.echo() click.echo()
click.secho(f"[workspace] session {session_id} is bound to a different directory:", fg="yellow") click.secho(f"[workspace] session {session_id} is bound to a different directory:", fg="yellow")
click.echo(f" session: {session_workspace}") click.echo(f" session: {session_workspace}")
@@ -49,7 +66,6 @@ def prompt_workspace_mismatch(
click.echo(" k = keep session workspace (switch tools root to session path)") click.echo(" k = keep session workspace (switch tools root to session path)")
click.echo(" u = update binding to current workspace") click.echo(" u = update binding to current workspace")
click.echo(" a = abort resume") click.echo(" a = abort resume")
with pause_task_cancel_for_prompt():
try: try:
raw = click.prompt( raw = click.prompt(
"Choice", "Choice",
@@ -67,6 +83,16 @@ def prompt_workspace_mismatch(
return "keep" 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: def install_cli_limit_hooks() -> None:
"""Register interactive continue hooks for process-global tool limits.""" """Register interactive continue hooks for process-global tool limits."""
PtyManager.set_limit_continue_hook(prompt_continue_limit) PtyManager.set_limit_continue_hook(prompt_continue_limit)
+3
View File
@@ -78,7 +78,10 @@ async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T:
finally: finally:
set_sigint_reinstall(None) set_sigint_reinstall(None)
if installed: if installed:
with contextlib.suppress(NotImplementedError, RuntimeError, ValueError):
_ = loop.remove_signal_handler(signal.SIGINT) _ = 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(): if not task.done():
_ = task.cancel() _ = task.cancel()
with contextlib.suppress(asyncio.CancelledError): with contextlib.suppress(asyncio.CancelledError):
+4 -4
View File
@@ -49,7 +49,7 @@ class ReplState:
def _tool_registry(self) -> ToolRegistry | None: def _tool_registry(self) -> ToolRegistry | None:
if not self.tools_enabled: if not self.tools_enabled:
return None 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 from plyngent.tools.danger import classify_danger
agent_cfg = self.config.agent_config agent_cfg = self.config.agent_config
@@ -57,12 +57,12 @@ class ReplState:
return ToolRegistry( return ToolRegistry(
list(DEFAULT_TOOLS), list(DEFAULT_TOOLS),
danger=classify_danger, danger=classify_danger,
on_confirm=prompt_confirm_tool, on_confirm=prompt_confirm_tool_async,
) )
return ToolRegistry(list(DEFAULT_TOOLS)) return ToolRegistry(list(DEFAULT_TOOLS))
def _make_agent(self) -> ChatAgent: 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 agent_cfg = self.config.agent_config
system_prompt = agent_cfg.system_prompt or None system_prompt = agent_cfg.system_prompt or None
@@ -73,7 +73,7 @@ class ReplState:
memory=self.memory, memory=self.memory,
session_id=self.session_id, session_id=self.session_id,
max_rounds=self.max_rounds, max_rounds=self.max_rounds,
on_limit=prompt_continue_limit, on_limit=prompt_continue_limit_async,
stream=True, stream=True,
system_prompt=system_prompt, system_prompt=system_prompt,
max_tool_result_chars=agent_cfg.max_tool_result_chars, max_tool_result_chars=agent_cfg.max_tool_result_chars,
+42
View File
@@ -280,6 +280,48 @@ async def test_max_rounds_continue_hook() -> None:
assert len(client.calls) == 3 # noqa: PLR2004 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: async def test_default_max_rounds_is_generous() -> None:
from plyngent.agent.loop import DEFAULT_MAX_ROUNDS from plyngent.agent.loop import DEFAULT_MAX_ROUNDS
+27 -2
View File
@@ -2,8 +2,12 @@ from __future__ import annotations
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from plyngent.cli.interrupt import allow_task_cancel, pause_task_cancel_for_prompt from plyngent.cli.interrupt import (
from plyngent.cli.limits import prompt_continue_limit 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: if TYPE_CHECKING:
import pytest import pytest
@@ -23,3 +27,24 @@ def test_prompt_continue_limit_under_pause(monkeypatch: pytest.MonkeyPatch) -> N
monkeypatch.setattr("click.confirm", _confirm) monkeypatch.setattr("click.confirm", _confirm)
assert prompt_continue_limit("too many rounds") is True 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