diff --git a/CLAUDE.md b/CLAUDE.md index 9fe3eb0..ba96ec1 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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`). - 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 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 39c35f0..39536b4 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -78,6 +78,10 @@ class ChatAgent: if self.memory is not None and self.session_id is not None: _ = 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]: """Run the tool loop for an already-appended user message; persist only on success.""" pre_len = len(self.messages) - 1 @@ -85,6 +89,7 @@ class ChatAgent: pre_len = len(self.messages) self.messages.append(user_msg) + completed = False try: async for event in run_chat_loop( self.client, @@ -96,10 +101,11 @@ class ChatAgent: on_limit=self.on_limit, ): yield event - except Exception: - # Roll back the whole turn so a failed user message is not left half-applied. - del self.messages[pre_len:] - self.pending_retry_text = user_msg.content + completed = True + except BaseException: + # Includes CancelledError / KeyboardInterrupt paths via task cancel. + if not completed: + self._rollback_turn(pre_len, user_msg.content) raise for message in self.messages[pre_len:]: diff --git a/src/plyngent/cli/repl.py b/src/plyngent/cli/repl.py index e505552..03849fe 100644 --- a/src/plyngent/cli/repl.py +++ b/src/plyngent/cli/repl.py @@ -37,8 +37,9 @@ Commands: /rounds [n] Show or set max tool-loop rounds /retry Retry the last failed user turn (after errors) -On network/API errors, the user turn is not saved to the DB. Auto-retry -waits 10s, 20s, then 30s (Ctrl+C cancels waits; use /retry later). +On network/API errors or Ctrl+C during a turn, the turn is not saved to +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). Use --session ID or /resume to continue a prior chat after restart. diff --git a/src/plyngent/cli/retry.py b/src/plyngent/cli/retry.py index 2a0f9d4..8c34c43 100644 --- a/src/plyngent/cli/retry.py +++ b/src/plyngent/cli/retry.py @@ -1,6 +1,8 @@ from __future__ import annotations import asyncio +import contextlib +import signal import time from typing import TYPE_CHECKING @@ -9,7 +11,7 @@ import click from plyngent.cli.display import render_events if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable + from collections.abc import AsyncIterator, Callable, Coroutine from plyngent.agent import AgentEvent from plyngent.agent.chat import ChatAgent @@ -34,6 +36,43 @@ async def sleep_cancellable(seconds: float) -> bool: 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( agent: ChatAgent, *, @@ -44,15 +83,27 @@ async def run_turn_with_retries( ``starter`` produces the event stream for the current attempt (``agent.run`` 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) attempt = 0 while True: 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: click.echo() click.secho("interrupted", fg="yellow") + click.echo() return False except Exception as exc: # noqa: BLE001 — surface and optionally retry click.secho(f"error: {exc}", fg="red") diff --git a/tests/test_cli/test_retry.py b/tests/test_cli/test_retry.py index d9459b8..0e1c61b 100644 --- a/tests/test_cli/test_retry.py +++ b/tests/test_cli/test_retry.py @@ -2,6 +2,8 @@ from __future__ import annotations from typing import TYPE_CHECKING, Literal, overload +import pytest + from plyngent.agent import ChatAgent from plyngent.cli.retry import retry_pending_with_retries, run_turn_with_retries, sleep_cancellable from plyngent.config.models import DatabaseConfig @@ -18,8 +20,6 @@ from plyngent.memory import MemoryStore if TYPE_CHECKING: from collections.abc import AsyncIterator - import pytest - def _response(message: AssistantChatMessage) -> ChatCompletionResponse: return ChatCompletionResponse( @@ -105,6 +105,64 @@ async def test_auto_retry_eventually_succeeds(monkeypatch: pytest.MonkeyPatch) - 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 no_wait(_seconds: float) -> bool: return True