Merge pull request #205 from multica-ai/forrestchang/fix-orphan-tool-call

fix: prevent orphaned tool calls from session abort

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

packages/core/src/agent/session/session-transcript-repair.test.ts

@@ -113,6 +113,78 @@ describe("sanitizeToolUseResultPairing", () => {
     expect(out.map((m) => m.role)).toEqual(["user", "assistant"]);
   });
 
+  it("drops duplicate assistant messages from abort double-save", () => {
+    // Reproduces the bug: session abort saves the same assistant message twice,
+    // leaving the second copy with a tool call that has no matching tool result.
+    const input = [
+      {
+        role: "assistant",
+        content: [
+          { type: "text", text: "Let me write a file" },
+          { type: "toolCall", id: "tool_ABC", name: "write", arguments: { path: "/tmp/test.txt", content: "hello" } },
+        ],
+      },
+      // Duplicate from abort handler saving the same message again
+      {
+        role: "assistant",
+        content: [
+          { type: "text", text: "Let me write a file" },
+          { type: "toolCall", id: "tool_ABC", name: "write", arguments: { path: "/tmp/test.txt", content: "hello" } },
+        ],
+      },
+      // User sends a new message after the abort
+      { role: "user", content: "hello" },
+    ] as AgentMessage[];
+
+    const out = sanitizeToolUseResultPairing(input);
+
+    // Should have: assistant, synthetic toolResult, user
+    // The duplicate assistant should be removed
+    const assistants = out.filter((m) => m.role === "assistant");
+    expect(assistants).toHaveLength(1);
+
+    const toolResults = out.filter((m) => m.role === "toolResult") as Array<{ toolCallId?: string }>;
+    expect(toolResults).toHaveLength(1);
+    expect(toolResults[0]?.toolCallId).toBe("tool_ABC");
+
+    const users = out.filter((m) => m.role === "user");
+    expect(users).toHaveLength(1);
+
+    // Verify ordering: assistant, toolResult, user
+    expect(out.map((m) => m.role)).toEqual(["assistant", "toolResult", "user"]);
+  });
+
+  it("drops duplicate assistant followed by error assistant", () => {
+    // Full reproduction: duplicate assistant + user + error assistant
+    const input = [
+      {
+        role: "assistant",
+        content: [
+          { type: "toolCall", id: "tool_ABC", name: "write", arguments: { path: "/tmp/test.txt", content: "hello" } },
+        ],
+      },
+      {
+        role: "assistant",
+        content: [
+          { type: "toolCall", id: "tool_ABC", name: "write", arguments: { path: "/tmp/test.txt", content: "hello" } },
+        ],
+      },
+      { role: "user", content: "continue" },
+      { role: "assistant", content: [] },
+      { role: "user", content: "how are you" },
+      { role: "assistant", content: [] },
+    ] as AgentMessage[];
+
+    const out = sanitizeToolUseResultPairing(input);
+
+    // The duplicate assistant should be removed; error assistants are kept (no tool calls)
+    const assistants = out.filter((m) => m.role === "assistant");
+    expect(assistants).toHaveLength(3); // original + 2 error assistants
+
+    const toolResults = out.filter((m) => m.role === "toolResult");
+    expect(toolResults).toHaveLength(1);
+  });
+
   it("drops tool results with empty tool call id", () => {
     const input = [
       {

packages/core/src/agent/session/session-transcript-repair.ts

@@ -254,6 +254,17 @@ export function repairToolUseResultPairing(messages: AgentMessage[]): ToolUseRep
       }
     }
 
+    // Drop duplicate assistant messages whose tool calls all already have results.
+    // This happens when a session abort persists the same assistant message twice.
+    if (toolCalls.every((call) => seenToolResultIds.has(call.id))) {
+      for (const rem of remainder) {
+        out.push(rem);
+      }
+      changed = true;
+      i = j - 1;
+      continue;
+    }
+
     out.push(msg);
 
     if (spanResultsById.size > 0 && remainder.length > 0) {