core/agent: keep executed tool results when a failed turn is rolled back

A tool batch counted as committed only once every call had its result, so a
batch cut short was dropped whole and the retry ran tools that had already
had effects. Keep the calls that did report a result: their assistant row
stays as the pairing carrier (a request must answer every tool_calls id),
without the text of the unfinished turn, and the calls without a result are
absent so the retry re-issues them instead of assuming they ran. The kept
rows go to the store right away, and a failed rollback write is not allowed
to mask the failure being rolled back.

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:
+    async def _rollback_uncommitted(self, user_index: int) -> None:
-        """Drop incomplete suffix; keep committed tool rounds after the user."""
+        """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.
+        # Already-committed rows stay; write any the failed rollback could not.
-        self._persist_from = len(self.messages)
+        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

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")

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)) == []