diff --git a/CLAUDE.md b/CLAUDE.md index 9bf9326..9fe3eb0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -55,7 +55,7 @@ Async SQLAlchemy + aiosqlite. `MemoryStore`: schema init, default local user, se - **`ChatClient`** Protocol for `chat_completions`. - **`@tool` / `ToolRegistry`**: decorator infers JSON Schema from type hints; execute tools by name. - **`run_chat_loop`**: multi-round tool loop, yields `AgentEvent` stream; optional `on_limit` to continue past max rounds. -- **`ChatAgent`**: wrapper with optional `MemoryStore` bind (load/append messages). +- **`ChatAgent`**: wrapper with optional `MemoryStore` bind (load/append messages on success only); `pending_retry_text` + `retry()` after failed turns. ### Tools (`tools/`) @@ -73,7 +73,8 @@ Module-level `@tool` handlers. Call `set_workspace_root()` before use. 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`, … +- Slash: `/history`, `/sessions`, `/resume`, `/rounds`, `/retry`, … +- Failed turns: not written to DB; 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 ce6a09e..39c35f0 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -31,6 +31,7 @@ class ChatAgent: temperature: float | None on_limit: LimitContinueHook | None messages: list[AnyChatMessage] + pending_retry_text: str | None def __init__( # noqa: PLR0913 self, @@ -54,6 +55,7 @@ class ChatAgent: self.temperature = temperature self.on_limit = on_limit self.messages = list(messages) if messages is not None else [] + self.pending_retry_text = None async def load_history(self) -> None: """Replace in-memory messages from the bound memory session.""" @@ -61,6 +63,7 @@ class ChatAgent: msg = "load_history requires memory and session_id" raise RuntimeError(msg) self.messages = await self.memory.list_messages(self.session_id) + self.pending_retry_text = None async def bind_session(self, session_id: int, *, load: bool = True) -> None: """Attach a memory session id; optionally load existing messages.""" @@ -75,24 +78,48 @@ class ChatAgent: if self.memory is not None and self.session_id is not None: _ = await self.memory.append_message(self.session_id, message) + 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 + if pre_len < 0 or self.messages[pre_len] is not user_msg: + pre_len = len(self.messages) + self.messages.append(user_msg) + + try: + async for event in run_chat_loop( + self.client, + self.messages, + model=self.model, + tools=self.tools, + max_rounds=self.max_rounds, + temperature=self.temperature, + 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 + raise + + for message in self.messages[pre_len:]: + await self._persist(message) + self.pending_retry_text = None + async def run(self, user_text: str) -> AsyncIterator[AgentEvent]: - """Append a user message, run the tool loop, yield events, persist new messages.""" + """Append a user message, run the tool loop, yield events, persist only on success.""" user_msg = UserChatMessage(content=user_text) self.messages.append(user_msg) - await self._persist(user_msg) - - start_len = len(self.messages) - async for event in run_chat_loop( - self.client, - self.messages, - model=self.model, - tools=self.tools, - max_rounds=self.max_rounds, - temperature=self.temperature, - on_limit=self.on_limit, - ): + async for event in self._run_from_user_message(user_msg): yield event - # Persist messages appended by the loop after the user message. - for message in self.messages[start_len:]: - await self._persist(message) + async def retry(self) -> AsyncIterator[AgentEvent]: + """Re-run the last failed user turn (no duplicate user message in history/DB).""" + text = self.pending_retry_text + if text is None: + msg = "nothing to retry" + raise RuntimeError(msg) + user_msg = UserChatMessage(content=text) + self.messages.append(user_msg) + async for event in self._run_from_user_message(user_msg): + yield event diff --git a/src/plyngent/cli/readline_setup.py b/src/plyngent/cli/readline_setup.py index 5932221..919217d 100644 --- a/src/plyngent/cli/readline_setup.py +++ b/src/plyngent/cli/readline_setup.py @@ -28,6 +28,7 @@ SLASH_COMMANDS: tuple[str, ...] = ( "/model", "/tools", "/rounds", + "/retry", ) _TOOLS_ARGS: tuple[str, ...] = ("on", "off") diff --git a/src/plyngent/cli/repl.py b/src/plyngent/cli/repl.py index 7a57074..e505552 100644 --- a/src/plyngent/cli/repl.py +++ b/src/plyngent/cli/repl.py @@ -7,8 +7,8 @@ from typing import TYPE_CHECKING import click from msgspec import UNSET -from plyngent.cli.display import render_events from plyngent.cli.readline_setup import setup_readline +from plyngent.cli.retry import retry_pending_with_retries, run_user_text_with_retries from plyngent.cli.selection import select_model, select_provider from plyngent.lmproto.openai_compatible.model import ( AssistantChatMessage, @@ -35,6 +35,10 @@ Commands: /model [id] Show or switch model /tools [on|off] Show or toggle tools /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). Tab completes slash commands and some arguments (provider, model, tools). Use --session ID or /resume to continue a prior chat after restart. @@ -141,6 +145,7 @@ def _cmd_rounds(state: ReplState, arg: str) -> None: def _cmd_clear(state: ReplState) -> None: state.agent.messages.clear() + state.agent.pending_retry_text = None click.echo("conversation cleared (in-memory only; DB history kept)") @@ -197,6 +202,15 @@ def _cmd_history(state: ReplState, arg: str) -> None: click.echo(f"session={state.session_id} messages={len(messages)} showing={len(messages) - start}") for offset, message in enumerate(messages[start:]): click.echo(_format_history_message(start + offset, message)) + if state.agent.pending_retry_text is not None: + click.secho( + f"(pending retry) user: {_preview_content(state.agent.pending_retry_text)}", + fg="yellow", + ) + + +async def _cmd_retry(state: ReplState) -> None: + _ = await retry_pending_with_retries(state.agent) async def _dispatch_slash(state: ReplState, command: str, arg: str) -> bool: @@ -214,6 +228,7 @@ async def _dispatch_slash(state: ReplState, command: str, arg: str) -> bool: "model": lambda: _cmd_model(state, arg), "tools": lambda: _cmd_tools(state, arg), "rounds": lambda: _cmd_rounds(state, arg), + "retry": lambda: _cmd_retry(state), } handler = handlers.get(command) if handler is None: @@ -269,8 +284,4 @@ async def run_repl(state: ReplState) -> None: click.secho("user: ", fg="green", nl=False) click.echo(line) - try: - await render_events(state.agent.run(line)) - except Exception as exc: # noqa: BLE001 — show API/runtime errors in REPL - click.secho(f"error: {exc}", fg="red") - click.echo() # match render_events spacing before next prompt + _ = await run_user_text_with_retries(state.agent, line) diff --git a/src/plyngent/cli/retry.py b/src/plyngent/cli/retry.py new file mode 100644 index 0000000..2a0f9d4 --- /dev/null +++ b/src/plyngent/cli/retry.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import asyncio +import time +from typing import TYPE_CHECKING + +import click + +from plyngent.cli.display import render_events + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Callable + + from plyngent.agent import AgentEvent + from plyngent.agent.chat import ChatAgent + +# Wait before retry attempt 1, 2, and 3 (after the first failure). +DEFAULT_RETRY_DELAYS_SECONDS: tuple[float, ...] = (10.0, 20.0, 30.0) +_PREVIEW_LEN = 80 + + +async def sleep_cancellable(seconds: float) -> bool: + """Sleep in short steps so Ctrl+C can cancel. Returns False if interrupted.""" + deadline = time.monotonic() + seconds + try: + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + return True + await asyncio.sleep(min(0.5, remaining)) + except asyncio.CancelledError: + return False + except KeyboardInterrupt: + return False + + +async def run_turn_with_retries( + agent: ChatAgent, + *, + starter: Callable[[], AsyncIterator[AgentEvent]], + delays: tuple[float, ...] = DEFAULT_RETRY_DELAYS_SECONDS, +) -> bool: + """Run a chat turn with automatic retries on failure. + + ``starter`` produces the event stream for the current attempt (``agent.run`` + or ``agent.retry``). Returns True if the turn completed successfully. + """ + max_retries = len(delays) + attempt = 0 + while True: + try: + await render_events(starter()) + except KeyboardInterrupt: + click.echo() + click.secho("interrupted", fg="yellow") + return False + except Exception as exc: # noqa: BLE001 — surface and optionally retry + click.secho(f"error: {exc}", fg="red") + if attempt >= max_retries: + if agent.pending_retry_text is not None: + click.secho( + "auto-retry exhausted; use /retry to try again, or send a new message", + fg="yellow", + ) + click.echo() + return False + + wait = delays[attempt] + attempt += 1 + click.secho( + f"auto-retry {attempt}/{max_retries} in {wait:g}s " + f"(Ctrl+C to cancel; then /retry later)", + fg="yellow", + ) + try: + ok = await sleep_cancellable(wait) + except KeyboardInterrupt: + ok = False + if not ok: + click.secho("auto-retry cancelled; use /retry to try again", fg="yellow") + click.echo() + return False + click.secho(f"retrying ({attempt}/{max_retries})…", fg="yellow") + else: + return True + + +async def run_user_text_with_retries(agent: ChatAgent, text: str) -> bool: + """Send a new user message with auto-retry.""" + return await run_turn_with_retries(agent, starter=lambda: agent.run(text)) + + +async def retry_pending_with_retries(agent: ChatAgent) -> bool: + """Retry the last failed user turn with auto-retry.""" + if agent.pending_retry_text is None: + click.echo("nothing to retry") + return False + preview = agent.pending_retry_text + if len(preview) > _PREVIEW_LEN: + preview = preview[:_PREVIEW_LEN] + "…" + click.echo(f"retrying: {preview}") + return await run_turn_with_retries(agent, starter=agent.retry) diff --git a/src/plyngent/cli/state.py b/src/plyngent/cli/state.py index 76cc249..a47e752 100644 --- a/src/plyngent/cli/state.py +++ b/src/plyngent/cli/state.py @@ -66,6 +66,7 @@ class ReplState: session = await self.memory.create_session(name=name) self.session_id = session.sid self.agent = self._make_agent() + self.agent.pending_retry_text = None async def resume_session(self, session_id: int) -> None: row = await self.memory.get_session(session_id) diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py index 46dbb36..991cd7f 100644 --- a/tests/test_agent/test_loop.py +++ b/tests/test_agent/test_loop.py @@ -2,6 +2,8 @@ from __future__ import annotations from typing import TYPE_CHECKING, Literal, overload +import pytest + from plyngent.agent import ( AssistantMessageEvent, ChatAgent, @@ -225,3 +227,91 @@ async def test_chat_agent_memory_roundtrip() -> None: await agent2.load_history() assert len(agent2.messages) == 2 # noqa: PLR2004 await store.close() + + +async def test_chat_agent_failed_turn_not_persisted() -> None: + store = await MemoryStore.open(DatabaseConfig()) + session = await store.create_session(name="t") + + class BoomClient: + @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 + msg = "network down" + raise RuntimeError(msg) + + agent = ChatAgent(BoomClient(), model="m", memory=store, session_id=session.sid) + with pytest.raises(RuntimeError, match="network down"): + _ = [e async for e in agent.run("hello")] + + assert agent.messages == [] + assert agent.pending_retry_text == "hello" + loaded = await store.list_messages(session.sid) + assert loaded == [] + await store.close() + + +async def test_chat_agent_retry_after_failure() -> None: + store = await MemoryStore.open(DatabaseConfig()) + session = await store.create_session(name="t") + + class FlakyClient: + calls: int + + def __init__(self) -> None: + self.calls = 0 + + @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 + self.calls += 1 + if self.calls == 1: + msg = "temporary" + raise RuntimeError(msg) + if stream: + + async def empty() -> AsyncIterator[ChatCompletionChunk]: + empty_chunks: list[ChatCompletionChunk] = [] + for chunk in empty_chunks: + yield chunk + + return empty() + return _response(AssistantChatMessage(content="recovered")) + + client = FlakyClient() + agent = ChatAgent(client, model="m", memory=store, session_id=session.sid) + with pytest.raises(RuntimeError, match="temporary"): + _ = [e async for e in agent.run("ping")] + assert agent.pending_retry_text == "ping" + assert await store.list_messages(session.sid) == [] + + events = [e async for e in agent.retry()] + assert any(isinstance(e, TextDeltaEvent) and e.content == "recovered" for e in events) + assert agent.pending_retry_text is None + loaded = await store.list_messages(session.sid) + assert len(loaded) == 2 # noqa: PLR2004 + assert isinstance(loaded[0], UserChatMessage) + assert loaded[0].content == "ping" + await store.close() diff --git a/tests/test_cli/test_retry.py b/tests/test_cli/test_retry.py new file mode 100644 index 0000000..d9459b8 --- /dev/null +++ b/tests/test_cli/test_retry.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal, overload + +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 +from plyngent.lmproto.openai_compatible.model import ( + AssistantChatMessage, + ChatCompletionChoice, + ChatCompletionChunk, + ChatCompletionResponse, + ChatCompletionsParam, + UserChatMessage, +) +from plyngent.memory import MemoryStore + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + import pytest + + +def _response(message: AssistantChatMessage) -> ChatCompletionResponse: + return ChatCompletionResponse( + id="1", + object="chat.completion", + created=0, + model="test", + choices=[ + ChatCompletionChoice( + index=0, + message=message, + logprobs={}, + finish_reason="stop", + ) + ], + system_fingerprint="", + usage={}, + ) + + +class FlakyClient: + calls: int + fail_times: int + + def __init__(self, fail_times: int) -> None: + self.calls = 0 + self.fail_times = fail_times + + @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 + self.calls += 1 + if self.calls <= self.fail_times: + msg = f"fail-{self.calls}" + raise RuntimeError(msg) + if stream: + + async def empty() -> AsyncIterator[ChatCompletionChunk]: + empty_chunks: list[ChatCompletionChunk] = [] + for chunk in empty_chunks: + yield chunk + + return empty() + return _response(AssistantChatMessage(content="ok")) + + +async def test_sleep_cancellable_completes() -> None: + assert await sleep_cancellable(0.01) is True + + +async def test_auto_retry_eventually_succeeds(monkeypatch: pytest.MonkeyPatch) -> None: + async def no_wait(_seconds: float) -> bool: + return True + + monkeypatch.setattr("plyngent.cli.retry.sleep_cancellable", no_wait) + store = await MemoryStore.open(DatabaseConfig()) + session = await store.create_session(name="t") + client = FlakyClient(fail_times=2) + agent = ChatAgent(client, model="m", memory=store, session_id=session.sid) + + ok = await run_turn_with_retries( + agent, + starter=lambda: agent.run("hello"), + delays=(0.01, 0.01, 0.01), + ) + assert ok is True + assert client.calls == 3 # noqa: PLR2004 + loaded = await store.list_messages(session.sid) + assert len(loaded) == 2 # noqa: PLR2004 + assert isinstance(loaded[0], UserChatMessage) + assert loaded[0].content == "hello" + await store.close() + + +async def test_manual_retry_after_exhausted(monkeypatch: pytest.MonkeyPatch) -> None: + async def no_wait(_seconds: float) -> bool: + return True + + monkeypatch.setattr("plyngent.cli.retry.sleep_cancellable", no_wait) + store = await MemoryStore.open(DatabaseConfig()) + session = await store.create_session(name="t") + client = FlakyClient(fail_times=5) + agent = ChatAgent(client, model="m", memory=store, session_id=session.sid) + + ok = await run_turn_with_retries( + agent, + starter=lambda: agent.run("hold-me"), + delays=(0.01,), # one auto-retry only → 2 attempts total, both fail + ) + assert ok is False + assert agent.pending_retry_text == "hold-me" + assert await store.list_messages(session.sid) == [] + + # Now succeed on manual /retry path + client.fail_times = client.calls # next call succeeds + ok2 = await retry_pending_with_retries(agent) + # retry_pending uses default delays; force zero-wait already patched + # But fail_times set so first attempt of manual retry succeeds immediately + assert ok2 is True + assert agent.pending_retry_text is None + loaded = await store.list_messages(session.sid) + assert len(loaded) == 2 # noqa: PLR2004 + assert isinstance(loaded[0], UserChatMessage) + assert loaded[0].content == "hold-me" + await store.close()