diff --git a/src/plyngent/agent/chat.py b/src/plyngent/agent/chat.py index c5a72a2..10c67cc 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -1,9 +1,11 @@ from __future__ import annotations +import contextlib from contextvars import ContextVar, Token from pathlib import Path from typing import TYPE_CHECKING +import msgspec from msgspec import UNSET from plyngent.lmproto.openai_compatible.model import ( @@ -37,7 +39,7 @@ from .usage import TokenUsage if TYPE_CHECKING: from collections.abc import AsyncIterator, Awaitable, Callable, Sequence - from plyngent.lmproto.openai_compatible.model import AnyChatMessage + from plyngent.lmproto.openai_compatible.model import AnyAssistantToolCall, AnyChatMessage from plyngent.memory import MemoryStore from .events import AgentEvent @@ -206,6 +208,75 @@ def committed_prefix_end(messages: Sequence[AnyChatMessage], user_index: int) -> return end +def rollback_tail(messages: Sequence[AnyChatMessage], start: int) -> list[AnyChatMessage]: + """Rows to keep from *start* when a failed turn is rolled back. + + The trailing tool batch is committed only once every call has its result + (:func:`committed_prefix_end`), so a batch cut short — cancelled before the + last result landed, or ids that do not pair up exactly — would be dropped + whole and the retry would run calls that already had external effects. Keep + the executed ones: their assistant message stays as the carrier for the + pairing (API messages must answer every ``tool_calls`` id), without the text + of the unfinished turn, and the calls that never reported a result are + absent so the retry re-issues them instead of assuming they ran. + + Developer rows (host notices, directive checkpoints) are already durable, so + they stay in place; unfinished assistant text and stray rows go. + """ + body = list(messages[start:]) + kept: list[AnyChatMessage] = [] + index = 0 + while index < len(body): + message = body[index] + if isinstance(message, DeveloperChatMessage): + kept.append(message) + index += 1 + continue + if isinstance(message, AssistantChatMessage) and message.tool_calls is not UNSET and message.tool_calls: + results: list[ToolChatMessage] = [] + follow = index + 1 + while follow < len(body): + following = body[follow] + if not isinstance(following, ToolChatMessage): + break + results.append(following) + follow += 1 + kept.extend(_executed_tool_pair(message, message.tool_calls, results)) + index = follow + continue + index += 1 + return kept + + +def _executed_tool_pair( + assistant: AssistantChatMessage, + calls: Sequence[AnyAssistantToolCall], + results: Sequence[ToolChatMessage], +) -> list[AnyChatMessage]: + """Assistant call record plus the results of the calls that produced one. + + Pairs by call id (first result per id), so the history keeps the API's + one-result-per-call shape even when a model repeats an id. Text and + reasoning go with the rest of the unfinished turn. + """ + by_id: dict[str, ToolChatMessage] = {} + for result in results: + _ = by_id.setdefault(result.tool_call_id, result) + executed: list[AnyAssistantToolCall] = [] + paired: list[AnyChatMessage] = [] + seen: set[str] = set() + for call in calls: + if call.id in seen or call.id not in by_id: + continue + seen.add(call.id) + executed.append(call) + paired.append(by_id[call.id]) + if not executed: + return [] + record = msgspec.structs.replace(assistant, content=UNSET, reasoning_content=UNSET, tool_calls=executed) + return [record, *paired] + + class ChatAgent: """Thin wrapper: chat client + optional tools + optional memory bind.""" @@ -475,10 +546,20 @@ class ChatAgent: msg = "nothing to retry" raise RuntimeError(msg) - def _rollback_uncommitted(self, user_index: int) -> None: - """Drop incomplete suffix; keep committed tool rounds after the user.""" + async def _rollback_uncommitted(self, user_index: int) -> None: + """Drop the unfinished suffix; keep committed rounds and executed tools. + + Rows before the checkpoint cursor are already stored, so only the kept + suffix is written: a restart has to continue from the same tool results + the model is about to see, or the retry re-runs their effects. + """ end = committed_prefix_end(self.messages, user_index) - del self.messages[end:] + self.messages[end:] = rollback_tail(self.messages, end) + if len(self.messages) > self._persist_from: + # A store error must not mask the failure being rolled back; the + # rows stay in memory, and retry() writes what is still unstored. + with contextlib.suppress(Exception): + await self._persist_range(self._persist_from, len(self.messages)) def _developer_tail_end(self, start: int) -> int: """Extend *start* through trailing developer checkpoint messages.""" @@ -492,7 +573,8 @@ class ChatAgent: After each completed tool batch, commits assistant+tool messages to the DB and keeps them on failure so :meth:`retry` continues without redoing - side-effecting tools. Unfinished assistant/stream suffix is rolled back. + side-effecting tools. Without such a batch, only the unfinished text is + dropped: calls that did report a result survive (:func:`rollback_tail`). """ user_index = self._user_index(user_msg) if self.todo_stack is not None: @@ -550,7 +632,7 @@ class ChatAgent: completed = True except BaseException: if not completed: - self._rollback_uncommitted(user_index) + await self._rollback_uncommitted(user_index) raise self.last_turn_usage = turn_usage @@ -687,7 +769,8 @@ class ChatAgent: user_msg = self.messages[user_index] assert isinstance(user_msg, UserChatMessage) self._ensure_system_prompt() - # Already-committed messages stay; only new rounds are persisted. - self._persist_from = len(self.messages) + # Already-committed rows stay; write any the failed rollback could not. + if self._persist_from < len(self.messages): + await self._persist_range(self._persist_from, len(self.messages)) async for event in self._run_from_user_message(user_msg): yield event diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py index 189e5f8..d59e86d 100644 --- a/tests/test_agent/test_loop.py +++ b/tests/test_agent/test_loop.py @@ -639,6 +639,77 @@ async def test_chat_agent_failed_turn_keeps_user_in_db() -> None: await store.close() +async def test_rollback_keeps_executed_tool_result_in_store() -> None: + """Rolling back a cut-short batch keeps its executed result, in memory and on disk.""" + store = await MemoryStore.open(DatabaseConfig()) + session = await store.create_session(name="t") + ran: list[str] = [] + + @tool(register=False) + def first() -> str: + ran.append("first") + return "first ran" + + client = ScriptedClient([_response(AssistantChatMessage(content="finished"))]) + agent = ChatAgent( + client, + model="m", + tools=ToolRegistry([first]), + memory=store, + session_id=session.sid, + stream=False, + ) + user = UserChatMessage(content="do both") + await store.append_message(session.sid, user) + agent.replace_messages( + [ + user, + AssistantChatMessage( + content="running both", + tool_calls=[ + AssistantFunctionToolCall( + id="call_1", function=AssistantFunctionTool(name="first", arguments="{}") + ), + AssistantFunctionToolCall( + id="call_2", function=AssistantFunctionTool(name="other", arguments="{}") + ), + ], + ), + ToolChatMessage(tool_call_id="call_1", content="first ran"), + ], + persist_from=1, + ) + + await agent._rollback_uncommitted(0) # the failure path itself + + kept = [message for message in agent.messages if isinstance(message, ToolChatMessage)] + assert [message.content for message in kept] == ["first ran"] + assert agent.pending_retry_text == "do both" + # What the model continues from is on disk too, or a restart re-runs it. + assert agent.persist_from == len(agent.messages) + + stored = await store.list_messages(session.sid) + assert [type(message).__name__ for message in stored] == [ + "UserChatMessage", + "AssistantChatMessage", + "ToolChatMessage", + ] + record = stored[1] + assert isinstance(record, AssistantChatMessage) + assert record.content is UNSET + calls = record.tool_calls + assert calls is not UNSET + assert [call.id for call in calls] == ["call_1"] + + # The retry answers from the stored result instead of running the tool again. + events = [event async for event in agent.retry()] + assert ran == [] + assert any(isinstance(event, TextDeltaEvent) and event.content == "finished" for event in events) + sent = client.calls[-1].messages + assert any(isinstance(message, ToolChatMessage) and message.content == "first ran" for message in sent) + await store.close() + + async def test_chat_agent_retry_after_failure() -> None: store = await MemoryStore.open(DatabaseConfig()) session = await store.create_session(name="t") diff --git a/tests/test_agent/test_rollback.py b/tests/test_agent/test_rollback.py new file mode 100644 index 0000000..c4f68b8 --- /dev/null +++ b/tests/test_agent/test_rollback.py @@ -0,0 +1,99 @@ +"""Rollback of a failed turn keeps the tool work that already ran. + +``committed_prefix_end`` commits a batch only once every call has its result, so +``rollback_tail`` keeps the executed part of a batch that was cut short: the +retry then continues from real results instead of running those tools again. +""" + +from __future__ import annotations + +from msgspec import UNSET + +from plyngent.agent.chat import rollback_tail +from plyngent.lmproto.openai_compatible.model import ( + AnyChatMessage, + AssistantChatMessage, + AssistantFunctionTool, + AssistantFunctionToolCall, + DeveloperChatMessage, + ToolChatMessage, + UserChatMessage, +) + + +def _call(call_id: str, name: str = "some_tool") -> AssistantFunctionToolCall: + return AssistantFunctionToolCall(id=call_id, function=AssistantFunctionTool(name=name, arguments="{}")) + + +def test_rollback_tail_keeps_the_executed_call_of_a_cut_short_batch() -> None: + messages: list[AnyChatMessage] = [ + UserChatMessage(content="do both"), + AssistantChatMessage(content="first, then the other", tool_calls=[_call("call_a"), _call("call_b")]), + ToolChatMessage(tool_call_id="call_a", content="a ran"), + ] + + tail = rollback_tail(messages, 1) + + assert len(tail) == 2 + record, result = tail + assert isinstance(record, AssistantChatMessage) + # The assistant row survives only as the carrier of the pairing. + assert record.content is UNSET + assert record.reasoning_content is UNSET + calls = record.tool_calls + assert calls is not UNSET + assert [call.id for call in calls] == ["call_a"] + assert isinstance(result, ToolChatMessage) + assert result.content == "a ran" + + +def test_rollback_tail_drops_a_batch_that_reported_nothing() -> None: + messages: list[AnyChatMessage] = [ + UserChatMessage(content="go"), + AssistantChatMessage(content="about to call", tool_calls=[_call("call_a")]), + ] + + assert rollback_tail(messages, 1) == [] + + +def test_rollback_tail_drops_unfinished_text_and_keeps_developer_rows() -> None: + messages: list[AnyChatMessage] = [ + UserChatMessage(content="go"), + AssistantChatMessage(content=UNSET, tool_calls=[_call("call_a")]), + ToolChatMessage(tool_call_id="call_a", content="a ran"), + DeveloperChatMessage(content="checkpoint"), + AssistantChatMessage(content="half-written answer"), + ] + + tail = rollback_tail(messages, 1) + + assert [type(message) for message in tail] == [AssistantChatMessage, ToolChatMessage, DeveloperChatMessage] + + +def test_rollback_tail_pairs_a_repeated_call_id_once() -> None: + messages: list[AnyChatMessage] = [ + UserChatMessage(content="go"), + AssistantChatMessage(content=UNSET, tool_calls=[_call("call_a"), _call("call_a")]), + ToolChatMessage(tool_call_id="call_a", content="only once"), + ] + + tail = rollback_tail(messages, 1) + + assert len(tail) == 2 + record, result = tail + assert isinstance(record, AssistantChatMessage) + calls = record.tool_calls + assert calls is not UNSET + assert [call.id for call in calls] == ["call_a"] + assert isinstance(result, ToolChatMessage) + assert result.content == "only once" + + +def test_rollback_tail_ignores_rows_before_start() -> None: + messages: list[AnyChatMessage] = [ + UserChatMessage(content="go"), + AssistantChatMessage(content=UNSET, tool_calls=[_call("call_a")]), + ToolChatMessage(tool_call_id="call_a", content="a ran"), + ] + + assert rollback_tail(messages, len(messages)) == []