core/agent+cli: defer DB write until turn success; auto/manual retry

Failed turns roll back memory and keep pending_retry_text; persist only
after success. CLI auto-retries 10s/20s/30s (cancellable) and /retry.
This commit is contained in:
2026-07-14 20:56:08 +08:00
parent f079ab73df
commit e0916a6352
8 changed files with 395 additions and 24 deletions
+3 -2
View File
@@ -55,7 +55,7 @@ Async SQLAlchemy + aiosqlite. `MemoryStore`: schema init, default local user, se
- **`ChatClient`** Protocol for `chat_completions`. - **`ChatClient`** Protocol for `chat_completions`.
- **`@tool` / `ToolRegistry`**: decorator infers JSON Schema from type hints; execute tools by name. - **`@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. - **`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/`) ### 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`. 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`, … - 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 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.
+34 -7
View File
@@ -31,6 +31,7 @@ class ChatAgent:
temperature: float | None temperature: float | None
on_limit: LimitContinueHook | None on_limit: LimitContinueHook | None
messages: list[AnyChatMessage] messages: list[AnyChatMessage]
pending_retry_text: str | None
def __init__( # noqa: PLR0913 def __init__( # noqa: PLR0913
self, self,
@@ -54,6 +55,7 @@ class ChatAgent:
self.temperature = temperature self.temperature = temperature
self.on_limit = on_limit self.on_limit = on_limit
self.messages = list(messages) if messages is not None else [] self.messages = list(messages) if messages is not None else []
self.pending_retry_text = None
async def load_history(self) -> None: async def load_history(self) -> None:
"""Replace in-memory messages from the bound memory session.""" """Replace in-memory messages from the bound memory session."""
@@ -61,6 +63,7 @@ class ChatAgent:
msg = "load_history requires memory and session_id" msg = "load_history requires memory and session_id"
raise RuntimeError(msg) raise RuntimeError(msg)
self.messages = await self.memory.list_messages(self.session_id) 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: async def bind_session(self, session_id: int, *, load: bool = True) -> None:
"""Attach a memory session id; optionally load existing messages.""" """Attach a memory session id; optionally load existing messages."""
@@ -75,13 +78,14 @@ 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)
async def run(self, user_text: str) -> AsyncIterator[AgentEvent]: async def _run_from_user_message(self, user_msg: UserChatMessage) -> AsyncIterator[AgentEvent]:
"""Append a user message, run the tool loop, yield events, persist new messages.""" """Run the tool loop for an already-appended user message; persist only on success."""
user_msg = UserChatMessage(content=user_text) 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) self.messages.append(user_msg)
await self._persist(user_msg)
start_len = len(self.messages) try:
async for event in run_chat_loop( async for event in run_chat_loop(
self.client, self.client,
self.messages, self.messages,
@@ -92,7 +96,30 @@ class ChatAgent:
on_limit=self.on_limit, on_limit=self.on_limit,
): ):
yield event 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
# Persist messages appended by the loop after the user message. for message in self.messages[pre_len:]:
for message in self.messages[start_len:]:
await self._persist(message) 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 only on success."""
user_msg = UserChatMessage(content=user_text)
self.messages.append(user_msg)
async for event in self._run_from_user_message(user_msg):
yield event
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
+1
View File
@@ -28,6 +28,7 @@ SLASH_COMMANDS: tuple[str, ...] = (
"/model", "/model",
"/tools", "/tools",
"/rounds", "/rounds",
"/retry",
) )
_TOOLS_ARGS: tuple[str, ...] = ("on", "off") _TOOLS_ARGS: tuple[str, ...] = ("on", "off")
+17 -6
View File
@@ -7,8 +7,8 @@ from typing import TYPE_CHECKING
import click import click
from msgspec import UNSET from msgspec import UNSET
from plyngent.cli.display import render_events
from plyngent.cli.readline_setup import setup_readline 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.cli.selection import select_model, select_provider
from plyngent.lmproto.openai_compatible.model import ( from plyngent.lmproto.openai_compatible.model import (
AssistantChatMessage, AssistantChatMessage,
@@ -35,6 +35,10 @@ Commands:
/model [id] Show or switch model /model [id] Show or switch model
/tools [on|off] Show or toggle tools /tools [on|off] Show or toggle tools
/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)
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). 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.
@@ -141,6 +145,7 @@ def _cmd_rounds(state: ReplState, arg: str) -> None:
def _cmd_clear(state: ReplState) -> None: def _cmd_clear(state: ReplState) -> None:
state.agent.messages.clear() state.agent.messages.clear()
state.agent.pending_retry_text = None
click.echo("conversation cleared (in-memory only; DB history kept)") 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}") click.echo(f"session={state.session_id} messages={len(messages)} showing={len(messages) - start}")
for offset, message in enumerate(messages[start:]): for offset, message in enumerate(messages[start:]):
click.echo(_format_history_message(start + offset, message)) 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: 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), "model": lambda: _cmd_model(state, arg),
"tools": lambda: _cmd_tools(state, arg), "tools": lambda: _cmd_tools(state, arg),
"rounds": lambda: _cmd_rounds(state, arg), "rounds": lambda: _cmd_rounds(state, arg),
"retry": lambda: _cmd_retry(state),
} }
handler = handlers.get(command) handler = handlers.get(command)
if handler is None: if handler is None:
@@ -269,8 +284,4 @@ async def run_repl(state: ReplState) -> None:
click.secho("user: ", fg="green", nl=False) click.secho("user: ", fg="green", nl=False)
click.echo(line) click.echo(line)
try: _ = await run_user_text_with_retries(state.agent, line)
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
+102
View File
@@ -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)
+1
View File
@@ -66,6 +66,7 @@ class ReplState:
session = await self.memory.create_session(name=name) session = await self.memory.create_session(name=name)
self.session_id = session.sid self.session_id = session.sid
self.agent = self._make_agent() self.agent = self._make_agent()
self.agent.pending_retry_text = None
async def resume_session(self, session_id: int) -> None: async def resume_session(self, session_id: int) -> None:
row = await self.memory.get_session(session_id) row = await self.memory.get_session(session_id)
+90
View File
@@ -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 ( from plyngent.agent import (
AssistantMessageEvent, AssistantMessageEvent,
ChatAgent, ChatAgent,
@@ -225,3 +227,91 @@ async def test_chat_agent_memory_roundtrip() -> None:
await agent2.load_history() await agent2.load_history()
assert len(agent2.messages) == 2 # noqa: PLR2004 assert len(agent2.messages) == 2 # noqa: PLR2004
await store.close() 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()
+138
View File
@@ -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()