mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 22:44:59 +08:00
core/cli: level-1 cancel of in-flight turns via Ctrl+C
Run turns as asyncio tasks; SIGINT cancels the task; agent rolls back and keeps pending_retry_text for /retry.
This commit is contained in:
@@ -74,7 +74,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:`), resumes latest session 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:`), resumes latest session by default (`--new` / `--session`).
|
||||||
- Slash: `/history`, `/sessions`, `/resume`, `/rounds`, `/retry`, …
|
- Slash: `/history`, `/sessions`, `/resume`, `/rounds`, `/retry`, …
|
||||||
- Failed turns: not written to DB; auto-retry 10s/20s/30s (Ctrl+C cancels wait); `/retry` manual.
|
- Failed/cancelled turns: not written to DB; Ctrl+C cancels the in-flight turn task; auto-retry 10s/20s/30s (Ctrl+C cancels wait); `/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.
|
||||||
|
|||||||
@@ -78,6 +78,10 @@ class ChatAgent:
|
|||||||
if self.memory is not None and self.session_id is not None:
|
if self.memory is not None and self.session_id is not None:
|
||||||
_ = await self.memory.append_message(self.session_id, message)
|
_ = await self.memory.append_message(self.session_id, message)
|
||||||
|
|
||||||
|
def _rollback_turn(self, pre_len: int, user_text: str) -> None:
|
||||||
|
del self.messages[pre_len:]
|
||||||
|
self.pending_retry_text = user_text
|
||||||
|
|
||||||
async def _run_from_user_message(self, user_msg: UserChatMessage) -> AsyncIterator[AgentEvent]:
|
async def _run_from_user_message(self, user_msg: UserChatMessage) -> AsyncIterator[AgentEvent]:
|
||||||
"""Run the tool loop for an already-appended user message; persist only on success."""
|
"""Run the tool loop for an already-appended user message; persist only on success."""
|
||||||
pre_len = len(self.messages) - 1
|
pre_len = len(self.messages) - 1
|
||||||
@@ -85,6 +89,7 @@ class ChatAgent:
|
|||||||
pre_len = len(self.messages)
|
pre_len = len(self.messages)
|
||||||
self.messages.append(user_msg)
|
self.messages.append(user_msg)
|
||||||
|
|
||||||
|
completed = False
|
||||||
try:
|
try:
|
||||||
async for event in run_chat_loop(
|
async for event in run_chat_loop(
|
||||||
self.client,
|
self.client,
|
||||||
@@ -96,10 +101,11 @@ class ChatAgent:
|
|||||||
on_limit=self.on_limit,
|
on_limit=self.on_limit,
|
||||||
):
|
):
|
||||||
yield event
|
yield event
|
||||||
except Exception:
|
completed = True
|
||||||
# Roll back the whole turn so a failed user message is not left half-applied.
|
except BaseException:
|
||||||
del self.messages[pre_len:]
|
# Includes CancelledError / KeyboardInterrupt paths via task cancel.
|
||||||
self.pending_retry_text = user_msg.content
|
if not completed:
|
||||||
|
self._rollback_turn(pre_len, user_msg.content)
|
||||||
raise
|
raise
|
||||||
|
|
||||||
for message in self.messages[pre_len:]:
|
for message in self.messages[pre_len:]:
|
||||||
|
|||||||
@@ -37,8 +37,9 @@ Commands:
|
|||||||
/rounds [n] Show or set max tool-loop rounds
|
/rounds [n] Show or set max tool-loop rounds
|
||||||
/retry Retry the last failed user turn (after errors)
|
/retry Retry the last failed user turn (after errors)
|
||||||
|
|
||||||
On network/API errors, the user turn is not saved to the DB. Auto-retry
|
On network/API errors or Ctrl+C during a turn, the turn is not saved to
|
||||||
waits 10s, 20s, then 30s (Ctrl+C cancels waits; use /retry later).
|
the DB. Auto-retry waits 10s, 20s, then 30s (Ctrl+C cancels waits or the
|
||||||
|
in-flight turn; use /retry later).
|
||||||
|
|
||||||
Tab completes slash commands and some arguments (provider, model, tools).
|
Tab completes slash commands and some arguments (provider, model, tools).
|
||||||
Use --session ID or /resume to continue a prior chat after restart.
|
Use --session ID or /resume to continue a prior chat after restart.
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import contextlib
|
||||||
|
import signal
|
||||||
import time
|
import time
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
@@ -9,7 +11,7 @@ import click
|
|||||||
from plyngent.cli.display import render_events
|
from plyngent.cli.display import render_events
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator, Callable
|
from collections.abc import AsyncIterator, Callable, Coroutine
|
||||||
|
|
||||||
from plyngent.agent import AgentEvent
|
from plyngent.agent import AgentEvent
|
||||||
from plyngent.agent.chat import ChatAgent
|
from plyngent.agent.chat import ChatAgent
|
||||||
@@ -34,6 +36,43 @@ async def sleep_cancellable(seconds: float) -> bool:
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T:
|
||||||
|
"""Await ``coro`` as a task; Ctrl+C / SIGINT cancels the task.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
asyncio.CancelledError: If the task was cancelled (including via SIGINT).
|
||||||
|
"""
|
||||||
|
task = asyncio.create_task(coro)
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
installed = False
|
||||||
|
|
||||||
|
def _on_sigint() -> None:
|
||||||
|
if not task.done():
|
||||||
|
_ = task.cancel()
|
||||||
|
|
||||||
|
try:
|
||||||
|
loop.add_signal_handler(signal.SIGINT, _on_sigint)
|
||||||
|
installed = True
|
||||||
|
except (NotImplementedError, RuntimeError, ValueError):
|
||||||
|
installed = False
|
||||||
|
|
||||||
|
try:
|
||||||
|
return await task
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
if not task.done():
|
||||||
|
_ = task.cancel()
|
||||||
|
with contextlib.suppress(asyncio.CancelledError):
|
||||||
|
await task
|
||||||
|
raise asyncio.CancelledError from None
|
||||||
|
finally:
|
||||||
|
if installed:
|
||||||
|
_ = loop.remove_signal_handler(signal.SIGINT)
|
||||||
|
if not task.done():
|
||||||
|
_ = task.cancel()
|
||||||
|
with contextlib.suppress(asyncio.CancelledError):
|
||||||
|
await task
|
||||||
|
|
||||||
|
|
||||||
async def run_turn_with_retries(
|
async def run_turn_with_retries(
|
||||||
agent: ChatAgent,
|
agent: ChatAgent,
|
||||||
*,
|
*,
|
||||||
@@ -44,15 +83,27 @@ async def run_turn_with_retries(
|
|||||||
|
|
||||||
``starter`` produces the event stream for the current attempt (``agent.run``
|
``starter`` produces the event stream for the current attempt (``agent.run``
|
||||||
or ``agent.retry``). Returns True if the turn completed successfully.
|
or ``agent.retry``). Returns True if the turn completed successfully.
|
||||||
|
|
||||||
|
Ctrl+C during a model/tool turn cancels the in-flight task (level-1 cancel).
|
||||||
|
The agent rolls back the turn and keeps ``pending_retry_text`` for ``/retry``.
|
||||||
"""
|
"""
|
||||||
max_retries = len(delays)
|
max_retries = len(delays)
|
||||||
attempt = 0
|
attempt = 0
|
||||||
while True:
|
while True:
|
||||||
try:
|
try:
|
||||||
await render_events(starter())
|
await run_cancellable(render_events(starter()))
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
click.echo()
|
||||||
|
click.secho(
|
||||||
|
"cancelled (turn not saved); use /retry to try again",
|
||||||
|
fg="yellow",
|
||||||
|
)
|
||||||
|
click.echo()
|
||||||
|
return False
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
click.echo()
|
click.echo()
|
||||||
click.secho("interrupted", fg="yellow")
|
click.secho("interrupted", fg="yellow")
|
||||||
|
click.echo()
|
||||||
return False
|
return False
|
||||||
except Exception as exc: # noqa: BLE001 — surface and optionally retry
|
except Exception as exc: # noqa: BLE001 — surface and optionally retry
|
||||||
click.secho(f"error: {exc}", fg="red")
|
click.secho(f"error: {exc}", fg="red")
|
||||||
|
|||||||
@@ -2,6 +2,8 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from typing import TYPE_CHECKING, Literal, overload
|
from typing import TYPE_CHECKING, Literal, overload
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
from plyngent.agent import ChatAgent
|
from plyngent.agent import ChatAgent
|
||||||
from plyngent.cli.retry import retry_pending_with_retries, run_turn_with_retries, sleep_cancellable
|
from plyngent.cli.retry import retry_pending_with_retries, run_turn_with_retries, sleep_cancellable
|
||||||
from plyngent.config.models import DatabaseConfig
|
from plyngent.config.models import DatabaseConfig
|
||||||
@@ -18,8 +20,6 @@ from plyngent.memory import MemoryStore
|
|||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator
|
from collections.abc import AsyncIterator
|
||||||
|
|
||||||
import pytest
|
|
||||||
|
|
||||||
|
|
||||||
def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
|
def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
|
||||||
return ChatCompletionResponse(
|
return ChatCompletionResponse(
|
||||||
@@ -105,6 +105,64 @@ async def test_auto_retry_eventually_succeeds(monkeypatch: pytest.MonkeyPatch) -
|
|||||||
await store.close()
|
await store.close()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_run_cancellable_cancels_task() -> None:
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from plyngent.cli.retry import run_cancellable
|
||||||
|
|
||||||
|
started = asyncio.Event()
|
||||||
|
|
||||||
|
async def hang() -> None:
|
||||||
|
started.set()
|
||||||
|
await asyncio.sleep(60)
|
||||||
|
|
||||||
|
task = asyncio.create_task(run_cancellable(hang()))
|
||||||
|
_ = await started.wait()
|
||||||
|
_ = task.cancel()
|
||||||
|
with pytest.raises(asyncio.CancelledError):
|
||||||
|
await task
|
||||||
|
|
||||||
|
|
||||||
|
async def test_cancel_turn_rolls_back_and_sets_pending() -> None:
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from plyngent.cli.display import render_events
|
||||||
|
from plyngent.cli.retry import run_cancellable
|
||||||
|
|
||||||
|
store = await MemoryStore.open(DatabaseConfig())
|
||||||
|
session = await store.create_session(name="t")
|
||||||
|
|
||||||
|
class HangClient:
|
||||||
|
@overload
|
||||||
|
async def chat_completions(
|
||||||
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
||||||
|
) -> ChatCompletionResponse: ...
|
||||||
|
|
||||||
|
@overload
|
||||||
|
async def chat_completions(
|
||||||
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
||||||
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
||||||
|
|
||||||
|
async def chat_completions(
|
||||||
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
||||||
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
||||||
|
del param, stream
|
||||||
|
await asyncio.sleep(60)
|
||||||
|
return _response(AssistantChatMessage(content="never"))
|
||||||
|
|
||||||
|
agent = ChatAgent(HangClient(), model="m", memory=store, session_id=session.sid)
|
||||||
|
turn = asyncio.create_task(run_cancellable(render_events(agent.run("cancel-me"))))
|
||||||
|
await asyncio.sleep(0.05)
|
||||||
|
_ = turn.cancel()
|
||||||
|
with pytest.raises(asyncio.CancelledError):
|
||||||
|
await turn
|
||||||
|
|
||||||
|
assert agent.pending_retry_text == "cancel-me"
|
||||||
|
assert agent.messages == []
|
||||||
|
assert await store.list_messages(session.sid) == []
|
||||||
|
await store.close()
|
||||||
|
|
||||||
|
|
||||||
async def test_manual_retry_after_exhausted(monkeypatch: pytest.MonkeyPatch) -> None:
|
async def test_manual_retry_after_exhausted(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
async def no_wait(_seconds: float) -> bool:
|
async def no_wait(_seconds: float) -> bool:
|
||||||
return True
|
return True
|
||||||
|
|||||||
Reference in New Issue
Block a user