From 7d9b979604eb82e89de1c6b53ab19a4e0d14be0d Mon Sep 17 00:00:00 2001
From: worldmozara <worldmozara@163.com>
Date: Tue, 6 Oct 2026 20:33:30 +0800
Subject: [PATCH] core/agent: record every tool result before yielding them

A caller that goes away mid-batch (cancel, display error) left the results of
tools that had already run out of history, so the retry ran them again. Append
the collected results first, then emit the events.
---
 src/plyngent/agent/loop.py    |  6 ++++-
 tests/test_agent/test_loop.py | 44 +++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py
index 4a1c71e..c7f1d09 100644
--- a/src/plyngent/agent/loop.py
+++ b/src/plyngent/agent/loop.py
@@ -124,10 +124,14 @@ async def _execute_tool_calls(
     else:
         results = [await _run_one_tool(tools, call, max_result_chars=max_result_chars) for call in tool_calls]
 
+    # Record every result before yielding any: a caller that goes away mid-batch
+    # (cancel, display error) must not lose work that already ran, or the retry
+    # runs that tool a second time.
+    for tool_msg, _err in results:
+        messages.append(tool_msg)
     for tool_msg, err in results:
         if err is not None:
             yield err
-        messages.append(tool_msg)
         yield ToolResultEvent(message=tool_msg)
 
 
diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py
index d59e86d..5e19447 100644
--- a/tests/test_agent/test_loop.py
+++ b/tests/test_agent/test_loop.py
@@ -710,6 +710,50 @@ async def test_rollback_keeps_executed_tool_result_in_store() -> None:
     await store.close()
 
 
+async def test_cancelled_batch_records_every_computed_result() -> None:
+    """Leaving mid-batch keeps the results that were already computed."""
+    ran: list[str] = []
+
+    @tool(register=False)
+    def first() -> str:
+        ran.append("first")
+        return "first ran"
+
+    @tool(register=False)
+    def second() -> str:
+        ran.append("second")
+        return "second ran"
+
+    client = ScriptedClient(
+        [
+            _response(
+                AssistantChatMessage(
+                    content="running both",
+                    tool_calls=[
+                        AssistantFunctionToolCall(
+                            id="call_1", function=AssistantFunctionTool(name="first", arguments="{}")
+                        ),
+                        AssistantFunctionToolCall(
+                            id="call_2", function=AssistantFunctionTool(name="second", arguments="{}")
+                        ),
+                    ],
+                )
+            )
+        ]
+    )
+    agent = ChatAgent(client, model="m", tools=ToolRegistry([first, second]), stream=False)
+
+    stream = agent.run("do both")
+    async for event in stream:
+        if isinstance(event, ToolResultEvent):
+            break
+    await stream.aclose()
+
+    results = [message for message in agent.messages if isinstance(message, ToolChatMessage)]
+    assert [message.content for message in results] == ["first ran", "second ran"]
+    assert sorted(ran) == ["first", "second"]
+
+
 async def test_chat_agent_retry_after_failure() -> None:
     store = await MemoryStore.open(DatabaseConfig())
     session = await store.create_session(name="t")