From 71f95d042a46dc517b4c0b6b99fe1a3eea1e6a5e Mon Sep 17 00:00:00 2001
From: Jiayuan Zhang <forrestchang7@gmail.com>
Date: Sun, 15 Feb 2026 21:31:36 +0800
Subject: [PATCH] fix(agent): prevent double-save of assistant message on abort

Track the last assistant message saved by the message_end event handler
and skip saving it again in the abort handler. This prevents the
duplicate assistant entries in session.jsonl that caused the
"tool_call_id is not found" bug.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 packages/core/src/agent/runner.ts | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/packages/core/src/agent/runner.ts b/packages/core/src/agent/runner.ts
index da3112a1..1d5491a7 100644
--- a/packages/core/src/agent/runner.ts
+++ b/packages/core/src/agent/runner.ts
@@ -153,6 +153,8 @@ export class Agent {
   private _internalRun = false;
   private _isRunning = false;
   private _aborted = false;
+  /** Last assistant message saved by the message_end event handler */
+  private _lastEventSavedAssistant: AgentMessage | undefined;
   private _runMutex: Promise<void> = Promise.resolve();
   private _compactionPromise: Promise<void> = Promise.resolve();
   private currentUserDisplayPrompt: string | undefined;
@@ -681,15 +683,17 @@ export class Agent {
     } finally {
       // On abort, persist any partial messages that pi-agent-core appended
       // via appendMessage() (no message_end event fires for those).
+      // Skip if message_end already fired for this message (avoids duplicates).
       if (this._aborted) {
         const messages = this.agent.state.messages;
         const lastMsg = messages[messages.length - 1];
-        if (lastMsg?.role === "assistant") {
+        if (lastMsg?.role === "assistant" && lastMsg !== this._lastEventSavedAssistant) {
           this.session.saveMessage(lastMsg);
         }
       }
       this._isRunning = false;
       this._aborted = false;
+      this._lastEventSavedAssistant = undefined;
       this.currentUserDisplayPrompt = undefined;
       this.currentUserSource = undefined;
       this.runLog.flush().catch(() => {});
@@ -829,6 +833,9 @@ export class Agent {
         saveOptions.source = this.currentUserSource;
       }
       this.session.saveMessage(message, Object.keys(saveOptions).length > 0 ? saveOptions : undefined);
+      if (message.role === "assistant") {
+        this._lastEventSavedAssistant = message;
+      }
       // Skip compaction during internal runs — internal messages will be
       // rolled back from memory afterwards, so compacting now would be incorrect.
       if (message.role === "assistant" && !this._internalRun) {