core/agent+cli: persist user immediately; /retry reuses DB turn

Keep incomplete user messages in memory/DB; roll back only
assistant/tool on failure. /retry and resume of orphan-user sessions
re-run without duplicating the user message.
This commit is contained in:
2026-07-15 11:21:01 +08:00
parent b1223c8189
commit 00a8ed3a9c
6 changed files with 111 additions and 36 deletions
+2 -2
View File
@@ -55,7 +55,7 @@ Async SQLAlchemy + aiosqlite. `MemoryStore`: schema init (+ lightweight SQLite `
- **`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; default **streaming** text deltas + stream tool-call merge; parallel tools; tool-result char budget; soft context compact on request (**API-calibrated** after first usage when available); cooperative cancel points; optional `on_limit`. - **`run_chat_loop`**: multi-round tool loop; default **streaming** text deltas + stream tool-call merge; parallel tools; tool-result char budget; soft context compact on request (**API-calibrated** after first usage when available); cooperative cancel points; optional `on_limit`.
- **`ChatAgent`**: optional `MemoryStore` (persist on success only); `stream`; system prompt; `pending_retry_text` + `retry()`. - **`ChatAgent`**: optional `MemoryStore` (user message persisted immediately; assistant/tools on success); `stream`; system prompt; `pending_retry_text` + `retry()` for incomplete turns (including orphan user after resume).
- **`/compact`**: soft-compact tool dumps → model summary (no tools) → **new** session seeded with summary message. - **`/compact`**: soft-compact tool dumps → model summary (no tools) → **new** session seeded with summary message.
- Events: text_delta, assistant_message, tool_call/result, max_rounds, **error** (`retryable`/`source`), **cancelled** (`reason`), **usage** (`TokenUsage`). - Events: text_delta, assistant_message, tool_call/result, max_rounds, **error** (`retryable`/`source`), **cancelled** (`reason`), **usage** (`TokenUsage`).
- Usage: API `usage` from completions (stream with `include_usage`); **char≈token fallback** (~4 chars/token) when omitted; **context size** = last request ``prompt_tokens`` (API preferred); `last_turn_usage` / `session_usage` are **billed sums** (tool rounds re-send history); CLI end-of-turn + `/status`. - Usage: API `usage` from completions (stream with `include_usage`); **char≈token fallback** (~4 chars/token) when omitted; **context size** = last request ``prompt_tokens`` (API preferred); `last_turn_usage` / `session_usage` are **billed sums** (tool rounds re-send history); CLI end-of-turn + `/status`.
@@ -81,7 +81,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 **most recently updated** session 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 **most recently updated** session for cwd/`--workspace` by default (`--new` / `--session`).
- Slash: `/history`, `/sessions` (newest first), `/resume [id]`, `/compact`, `/status` (incl. context char estimate), `/rounds`, `/retry`, … - Slash: `/history`, `/sessions` (newest first), `/resume [id]`, `/compact`, `/status` (incl. context char estimate), `/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) run off-loop via `asyncio.to_thread` + pause cancel so prompts do not abort the turn; auto-retry 10s/20s/30s; `/retry` manual. - Failed/cancelled turns: user message kept in DB; partial assistant/tool rolled back; Ctrl+C cancels in-flight turn; **TTY confirms** off-loop; auto-retry 10s/20s/30s then `/retry` (no duplicate user message).
- **`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.
+56 -12
View File
@@ -87,6 +87,7 @@ class ChatAgent:
self.last_request_usage = TokenUsage() self.last_request_usage = TokenUsage()
self.last_turn_rounds = 0 self.last_turn_rounds = 0
self._ensure_system_prompt() self._ensure_system_prompt()
self._sync_pending_from_orphan_user()
@property @property
def context_tokens(self) -> int: def context_tokens(self) -> int:
@@ -115,14 +116,21 @@ class ChatAgent:
return return
self.messages.insert(0, SystemChatMessage(content=self.system_prompt)) self.messages.insert(0, SystemChatMessage(content=self.system_prompt))
def _sync_pending_from_orphan_user(self) -> None:
"""If history ends with a user message, that turn is incomplete → retryable."""
if self.messages and isinstance(self.messages[-1], UserChatMessage):
self.pending_retry_text = self.messages[-1].content
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."""
if self.memory is None or self.session_id is None: if self.memory is None or self.session_id is None:
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
self._ensure_system_prompt() self._ensure_system_prompt()
self._sync_pending_from_orphan_user()
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."""
@@ -137,16 +145,31 @@ 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: def _user_index(self, user_msg: UserChatMessage) -> int:
del self.messages[pre_len:] for i in range(len(self.messages) - 1, -1, -1):
if self.messages[i] is user_msg:
return i
# Fallback: last matching content user message
for i in range(len(self.messages) - 1, -1, -1):
msg = self.messages[i]
if isinstance(msg, UserChatMessage) and msg.content == user_msg.content:
return i
msg = "user message not found in history"
raise RuntimeError(msg)
def _rollback_partial(self, user_index: int, user_text: str) -> None:
"""Drop assistant/tool messages after the user; keep user for retry/DB."""
del self.messages[user_index + 1 :]
self.pending_retry_text = user_text 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.
pre_len = len(self.messages) - 1
if pre_len < 0 or self.messages[pre_len] is not user_msg: On success, persists assistant/tool messages produced after the user.
pre_len = len(self.messages) On failure, keeps the user message (already in memory/DB) and sets
self.messages.append(user_msg) ``pending_retry_text`` so ``retry()`` can re-run without duplicating it.
"""
user_index = self._user_index(user_msg)
completed = False completed = False
turn_usage = TokenUsage() turn_usage = TokenUsage()
@@ -176,32 +199,53 @@ class ChatAgent:
completed = True completed = True
except BaseException: except BaseException:
if not completed: if not completed:
self._rollback_turn(pre_len, user_msg.content) self._rollback_partial(user_index, user_msg.content)
raise raise
self.last_turn_usage = turn_usage self.last_turn_usage = turn_usage
self.last_request_usage = last_request self.last_request_usage = last_request
self.last_turn_rounds = turn_rounds self.last_turn_rounds = turn_rounds
for message in self.messages[pre_len:]: for message in self.messages[user_index + 1 :]:
await self._persist(message) await self._persist(message)
self.pending_retry_text = None self.pending_retry_text = None
async def run(self, user_text: str) -> AsyncIterator[AgentEvent]: async def run(self, user_text: str) -> AsyncIterator[AgentEvent]:
"""Append a user message, run the tool loop, yield events, persist only on success.""" """Append a user message (persist immediately), run the tool loop, yield events."""
self._ensure_system_prompt() self._ensure_system_prompt()
user_msg = UserChatMessage(content=user_text) user_msg = UserChatMessage(content=user_text)
self.messages.append(user_msg) self.messages.append(user_msg)
await self._persist(user_msg)
async for event in self._run_from_user_message(user_msg): async for event in self._run_from_user_message(user_msg):
yield event yield event
async def retry(self) -> AsyncIterator[AgentEvent]: async def retry(self) -> AsyncIterator[AgentEvent]:
"""Re-run the last failed user turn (no duplicate user message in history/DB).""" """Re-run the incomplete last user turn without appending a new user message.
Useful after a failed/cancelled turn (user already in memory/DB) or after
``load_history`` when the session ends with an orphan user message.
"""
text = self.pending_retry_text
if text is None:
self._sync_pending_from_orphan_user()
text = self.pending_retry_text text = self.pending_retry_text
if text is None: if text is None:
msg = "nothing to retry" msg = "nothing to retry"
raise RuntimeError(msg) raise RuntimeError(msg)
self._ensure_system_prompt() self._ensure_system_prompt()
if self.messages and isinstance(self.messages[-1], UserChatMessage):
user_msg = self.messages[-1]
if user_msg.content != text:
# Pending text out of sync: replace trailing user
user_msg = UserChatMessage(content=text)
self.messages[-1] = user_msg
else:
user_msg = UserChatMessage(content=text) user_msg = UserChatMessage(content=text)
self.messages.append(user_msg) self.messages.append(user_msg)
# Already persisted on the original run when memory is bound; only
# persist if this is a reconstructed orphan without DB (no memory).
if self.memory is None or self.session_id is None:
await self._persist(user_msg)
async for event in self._run_from_user_message(user_msg): async for event in self._run_from_user_message(user_msg):
yield event yield event
+4 -4
View File
@@ -36,12 +36,12 @@ 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) /retry Re-run incomplete last user turn (DB/orphan user; no retype)
/status Show session/provider/tools/rounds status /status Show session/provider/tools/rounds status
On network/API errors or Ctrl+C during a turn, the turn is not saved to User messages are saved immediately. On API errors or Ctrl+C, partial
the DB. Auto-retry waits 10s, 20s, then 30s (Ctrl+C cancels waits or the assistant/tool output is discarded but the user message stays (so /retry
in-flight turn; use /retry later). Streaming is on by default. works after resume, not only via readline history). Auto-retry: 10s/20s/30s.
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.
+10 -6
View File
@@ -139,22 +139,24 @@ async def run_turn_with_retries(
) -> bool: ) -> bool:
"""Run a chat turn with automatic retries on failure. """Run a chat turn with automatic retries on failure.
``starter`` produces the event stream for the current attempt (``agent.run`` ``starter`` produces the event stream for the first attempt (usually
or ``agent.retry``). Returns True if the turn completed successfully. ``agent.run``). After a failure that sets ``pending_retry_text``, further
attempts use ``agent.retry`` so the user message is not duplicated in
history/DB (it is already persisted).
Ctrl+C during a model/tool turn cancels the in-flight task (level-1 cancel). Ctrl+C cancels the in-flight task; user message stays in DB for ``/retry``.
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
current: Callable[[], AsyncIterator[AgentEvent]] = starter
while True: while True:
try: try:
await run_cancellable(render_events(starter())) await run_cancellable(render_events(current()))
except asyncio.CancelledError: except asyncio.CancelledError:
# Do not auto-retry cancelled turns — user intent was stop, not retry. # Do not auto-retry cancelled turns — user intent was stop, not retry.
click.echo() click.echo()
click.secho( click.secho(
"cancelled (turn not saved); use /retry to try again", "cancelled; user message kept — use /retry to try again",
fg="yellow", fg="yellow",
) )
click.echo() click.echo()
@@ -166,6 +168,8 @@ async def run_turn_with_retries(
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")
if agent.pending_retry_text is not None:
current = agent.retry
if attempt >= max_retries: if attempt >= max_retries:
if agent.pending_retry_text is not None: if agent.pending_retry_text is not None:
click.secho( click.secho(
+29 -4
View File
@@ -505,7 +505,7 @@ async def test_chat_agent_memory_roundtrip() -> None:
await store.close() await store.close()
async def test_chat_agent_failed_turn_not_persisted() -> None: async def test_chat_agent_failed_turn_keeps_user_in_db() -> None:
store = await MemoryStore.open(DatabaseConfig()) store = await MemoryStore.open(DatabaseConfig())
session = await store.create_session(name="t") session = await store.create_session(name="t")
@@ -531,10 +531,13 @@ async def test_chat_agent_failed_turn_not_persisted() -> None:
with pytest.raises(RuntimeError, match="network down"): with pytest.raises(RuntimeError, match="network down"):
_ = [e async for e in agent.run("hello")] _ = [e async for e in agent.run("hello")]
assert agent.messages == [] assert len(agent.messages) == 1
assert isinstance(agent.messages[0], UserChatMessage)
assert agent.pending_retry_text == "hello" assert agent.pending_retry_text == "hello"
loaded = await store.list_messages(session.sid) loaded = await store.list_messages(session.sid)
assert loaded == [] assert len(loaded) == 1
assert isinstance(loaded[0], UserChatMessage)
assert loaded[0].content == "hello"
await store.close() await store.close()
@@ -581,7 +584,8 @@ async def test_chat_agent_retry_after_failure() -> None:
with pytest.raises(RuntimeError, match="temporary"): with pytest.raises(RuntimeError, match="temporary"):
_ = [e async for e in agent.run("ping")] _ = [e async for e in agent.run("ping")]
assert agent.pending_retry_text == "ping" assert agent.pending_retry_text == "ping"
assert await store.list_messages(session.sid) == [] # User already in DB after first attempt
assert len(await store.list_messages(session.sid)) == 1
events = [e async for e in agent.retry()] events = [e async for e in agent.retry()]
assert any(isinstance(e, TextDeltaEvent) and e.content == "recovered" for e in events) assert any(isinstance(e, TextDeltaEvent) and e.content == "recovered" for e in events)
@@ -590,9 +594,30 @@ async def test_chat_agent_retry_after_failure() -> None:
assert len(loaded) == 2 assert len(loaded) == 2
assert isinstance(loaded[0], UserChatMessage) assert isinstance(loaded[0], UserChatMessage)
assert loaded[0].content == "ping" assert loaded[0].content == "ping"
# Single user message (no duplicate on retry)
assert sum(1 for m in loaded if isinstance(m, UserChatMessage)) == 1
await store.close() await store.close()
async def test_retry_after_resume_orphan_user() -> None:
store = await MemoryStore.open(DatabaseConfig())
session = await store.create_session(name="t")
_ = await store.append_message(session.sid, UserChatMessage(content="left hanging"))
client = ScriptedClient([_response(AssistantChatMessage(content="ok now"))])
agent = ChatAgent(client, model="m", memory=store, session_id=session.sid)
await agent.load_history()
assert agent.pending_retry_text == "left hanging"
events = [e async for e in agent.retry()]
assert any(isinstance(e, TextDeltaEvent) and e.content == "ok now" for e in events)
loaded = await store.list_messages(session.sid)
assert len(loaded) == 2
assert sum(1 for m in loaded if isinstance(m, UserChatMessage)) == 1
await store.close()
async def test_chat_agent_system_prompt_prepended() -> None: async def test_chat_agent_system_prompt_prepended() -> None:
from plyngent.lmproto.openai_compatible.model import SystemChatMessage from plyngent.lmproto.openai_compatible.model import SystemChatMessage
+8 -6
View File
@@ -173,8 +173,11 @@ async def test_cancel_turn_rolls_back_and_sets_pending() -> None:
await turn await turn
assert agent.pending_retry_text == "cancel-me" assert agent.pending_retry_text == "cancel-me"
assert agent.messages == [] assert len(agent.messages) == 1
assert await store.list_messages(session.sid) == [] assert isinstance(agent.messages[0], UserChatMessage)
loaded = await store.list_messages(session.sid)
assert len(loaded) == 1
assert isinstance(loaded[0], UserChatMessage)
await store.close() await store.close()
@@ -195,17 +198,16 @@ async def test_manual_retry_after_exhausted(monkeypatch: pytest.MonkeyPatch) ->
) )
assert ok is False assert ok is False
assert agent.pending_retry_text == "hold-me" assert agent.pending_retry_text == "hold-me"
assert await store.list_messages(session.sid) == [] assert len(await store.list_messages(session.sid)) == 1
# Now succeed on manual /retry path # Now succeed on manual /retry path (no second user message)
client.fail_times = client.calls # next call succeeds client.fail_times = client.calls # next call succeeds
ok2 = await retry_pending_with_retries(agent) 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 ok2 is True
assert agent.pending_retry_text is None assert agent.pending_retry_text is None
loaded = await store.list_messages(session.sid) loaded = await store.list_messages(session.sid)
assert len(loaded) == 2 assert len(loaded) == 2
assert isinstance(loaded[0], UserChatMessage) assert isinstance(loaded[0], UserChatMessage)
assert loaded[0].content == "hold-me" assert loaded[0].content == "hold-me"
assert sum(1 for m in loaded if isinstance(m, UserChatMessage)) == 1
await store.close() await store.close()