fix(agent-core): handle incomplete tool_calls on session resume (#660)

.changeset/fix-session-resume-incomplete-tool-calls.md

@@ -0,0 +1,7 @@
+---
+"@moonshot-ai/agent-core": patch
+"@moonshot-ai/kimi-code": patch
+
+---
+
+Fix session resume failing with 400 error when previous turn was interrupted mid-tool-call.

packages/agent-core/src/agent/context/index.ts

@@ -245,6 +245,12 @@ export class ContextMemory {
         return;
       }
       case 'tool.call': {
+        // Skip stale tool_call_ids from previous incomplete turns.
+        // These are identified during replay pre-scan and would otherwise
+        // pollute pendingToolResultIds, causing deferred user messages.
+        if (this.agent.staleToolCallIds.has(event.toolCallId)) {
+          return;
+        }
         const openStep = this.openSteps.get(event.stepUuid);
         if (openStep === undefined) {
           throw new Error(
@@ -298,6 +304,79 @@ export class ContextMemory {
     return this.pendingToolResultIds.size > 0;
   }
 
+  /**
+   * Remove stale entries from `pendingToolResultIds` and trim orphaned
+   * assistant messages from `_history`.  This happens when a session is
+   * killed mid-tool-call and later resumed — the tool.call events are
+   * replayed but the tool.result events never arrived.  Without this
+   * cleanup, `hasOpenToolExchange()` would remain true (deferring new
+   * messages) and the orphaned assistant would still be sent to the
+   * provider on the next turn, causing a 400 error.
+   */
+  cleanupOrphanedToolCalls(): void {
+    // Clear stale pendingToolResultIds.
+    this.pendingToolResultIds.clear();
+
+    // Find assistant messages that have unanswered tool_calls.
+    // Check positionally: a tool_call_id is "answered" only if there
+    // is a tool result AFTER the assistant in history, before the next
+    // assistant. This prevents false matches when toolCallIds are reused
+    // across turns.
+    const assistantsToRemove = new Set<number>();
+    for (let i = 0; i < this._history.length; i++) {
+      const message = this._history[i];
+      if (message === undefined || message.role !== 'assistant' || message.toolCalls.length === 0) continue;
+
+      const allAnswered = message.toolCalls.every((tc) => {
+        for (let j = i + 1; j < this._history.length; j++) {
+          const later = this._history[j];
+          if (later !== undefined && later.role === 'tool' && later.toolCallId === tc.id) return true;
+          if (later !== undefined && later.role === 'assistant') break;
+        }
+        return false;
+      });
+
+      if (!allAnswered) {
+        assistantsToRemove.add(i);
+      }
+    }
+
+    if (assistantsToRemove.size > 0) {
+      // Build a set of indices to remove: each removed assistant AND all
+      // tool messages that follow it (up to the next assistant or end of
+      // history). This avoids globally removing tool messages by ID, which
+      // could incorrectly remove valid tool results from earlier turns
+      // when toolCallIds are reused.
+      const indicesToRemove = new Set<number>();
+      for (const idx of assistantsToRemove) {
+        indicesToRemove.add(idx);
+        // Remove tool messages after this assistant up to the next assistant.
+        for (let j = idx + 1; j < this._history.length; j++) {
+          const later = this._history[j];
+          if (later !== undefined && later.role === 'assistant') break;
+          if (later !== undefined && later.role === 'tool') {
+            indicesToRemove.add(j);
+          }
+        }
+      }
+
+      const removedMessages = new Set<ContextMessage>();
+      this._history = this._history.filter((message, index) => {
+        if (indicesToRemove.has(index)) {
+          removedMessages.add(message);
+          return false;
+        }
+        return true;
+      });
+      // Also remove from replay builder so ResumeSessionResult doesn't
+      // include stale orphaned messages.
+      this.agent.replayBuilder.removeLastMessages(removedMessages);
+    }
+
+    // Flush any deferred messages that were blocked by the stale pending set.
+    this.flushDeferredMessagesIfToolExchangeClosed();
+  }
+
   private pushHistory(...messages: ContextMessage[]): void {
     this._history.push(...messages);
     for (const message of messages) {

packages/agent-core/src/agent/context/projector.ts

@@ -12,7 +12,11 @@ export function project(history: readonly ContextMessage[]): Message[] {
       !(message.role === 'assistant' && message.content.length === 0 && message.toolCalls.length === 0)
     );
   });
-  return mergeAdjacentUserMessages(usable);
+  // Trim any trailing assistant message whose tool_calls were never answered
+  // (e.g. the session was killed mid-tool-call).  Sending an assistant
+  // message with open tool_calls violates the API contract and causes a
+  // 400 error on resume.
+  return trimTrailingOpenToolExchange(mergeAdjacentUserMessages(usable));
 }
 
 function mergeAdjacentUserMessages(history: readonly ContextMessage[]): Message[] {
@@ -77,8 +81,11 @@ export function trimTrailingOpenToolExchange(history: readonly Message[]): Messa
     lastNonToolIndex -= 1;
   }
 
+  // No assistant message found — nothing to trim.
+  if (lastNonToolIndex < 0) return [...history];
+
   const assistant = history[lastNonToolIndex];
-  if (assistant === undefined) return [];
+  if (assistant === undefined) return [...history];
   if (assistant.role !== 'assistant' || assistant.toolCalls.length === 0) return [...history];
 
   const trailingToolCallIds = new Set(

packages/agent-core/src/agent/index.ts

@@ -118,6 +118,8 @@ export class Agent {
   readonly fullCompaction: FullCompaction;
   readonly microCompaction: MicroCompaction;
   readonly context: ContextMemory;
+  /** Tool_call_ids from stale incomplete turns, identified during replay. */
+  staleToolCallIds: Set<string> = new Set();
   readonly config: ConfigState;
   readonly turn: TurnFlow;
   readonly injection: InjectionManager;
@@ -296,10 +298,15 @@ export class Agent {
 
   async resume(): Promise<{ warning?: string }> {
     const result = await this.records.replay();
+    this.staleToolCallIds = this.records.staleToolCallIds;
     this.goal.normalizeAfterReplay();
     await this.background.loadFromDisk();
     await this.background.reconcile();
     await this.cron?.loadFromDisk();
+    // Clean up any tool_call IDs that were never answered (session killed
+    // mid-tool-call).  Without this, new user messages would be silently
+    // deferred because `hasOpenToolExchange()` would remain true.
+    this.context.cleanupOrphanedToolCalls();
     this.turn.finishResume();
     return result;
   }

packages/agent-core/src/agent/records/index.ts

@@ -130,6 +130,8 @@ export interface RestoringContext {
 export class AgentRecords {
   private _restoring: RestoringContext | null = null;
   private metadataInitialized = false;
+  /** Tool_call_ids identified as stale during pre-scan of replay records. */
+  staleToolCallIds: Set<string> = new Set();
 
   constructor(
     private readonly agent: Agent,
@@ -174,12 +176,21 @@ export class AgentRecords {
 
   async replay(): Promise<{ warning?: string }> {
     if (!this.persistence) throw new Error('No persistence provided for AgentRecords');
+
+    // First pass: collect all records and pre-scan for stale tool_call_ids.
+    const allRecords: AgentRecord[] = [];
+    for await (const record of this.persistence.read()) {
+      allRecords.push(record as AgentRecord);
+    }
+    this.staleToolCallIds = this.findStaleToolCallIds(allRecords);
+
+    // Second pass: process records.
     let migrations: readonly WireMigration[] = [];
     let hasMetadata = false;
     let shouldRewrite = false;
     let warning: string | undefined;
     const replayedRecords: AgentRecord[] = [];
-    for await (const record of this.persistence.read()) {
+    for (const record of allRecords) {
       if (!hasMetadata) {
         if (record.type !== 'metadata') {
           throw new Error('AgentRecords replay expected metadata as the first record');
@@ -234,6 +245,62 @@ export class AgentRecords {
     return { warning };
   }
 
+  /**
+   * Pre-scan replay records to identify stale tool_call_ids — tool_calls
+   * that have no matching tool.result before the next context.append_message.
+   * These are from sessions killed mid-tool-call.  The stale set is stored
+   * on the agent so appendLoopEvent can skip them during replay, preventing
+   * them from polluting pendingToolResultIds.
+   */
+  findStaleToolCallIds(records: readonly AgentRecord[]): Set<string> {
+    const stale = new Set<string>();
+    const pendingByAssistant = new Map<string, Set<string>>();
+    let hasOpenToolExchange = false;
+    for (const record of records) {
+      if (record.type === 'context.append_loop_event') {
+        const event = record.event;
+        if (event.type === 'tool.call') {
+          if (!pendingByAssistant.has(event.stepUuid)) {
+            pendingByAssistant.set(event.stepUuid, new Set());
+          }
+          pendingByAssistant.get(event.stepUuid)!.add(event.toolCallId);
+          hasOpenToolExchange = true;
+        } else if (event.type === 'tool.result') {
+          for (const pending of pendingByAssistant.values()) {
+            pending.delete(event.toolCallId);
+          }
+          // Check if the exchange is now closed.
+          hasOpenToolExchange = false;
+          for (const pending of pendingByAssistant.values()) {
+            if (pending.size > 0) {
+              hasOpenToolExchange = true;
+              break;
+            }
+          }
+        }
+      } else if (record.type === 'context.append_message' && record.message.role === 'user') {
+        // A user message that arrives while a tool exchange is still open
+        // is a deferred same-turn message, not a turn boundary.  Only
+        // treat it as a turn boundary if the exchange is closed.
+        if (!hasOpenToolExchange) {
+          for (const pending of pendingByAssistant.values()) {
+            for (const id of pending) {
+              stale.add(id);
+            }
+          }
+          pendingByAssistant.clear();
+        }
+      }
+    }
+    // Any remaining pending IDs at the end of records are also stale.
+    for (const pending of pendingByAssistant.values()) {
+      for (const id of pending) {
+        stale.add(id);
+      }
+    }
+    return stale;
+  }
+
   async flush(): Promise<void> {
     await this.persistence?.flush();
   }

packages/agent-core/test/agent/context.test.ts

@@ -357,7 +357,7 @@
       variant: 'host',
     });
 
     expect(ctx.agent.context.messages.map((message) => message.role)).toEqual([
       'assistant',
       'user',
       'assistant',
@@ -776,6 +776,105 @@
     expect(textOf(messages[1]!)).toBe('No origin prompt');
     expect(textOf(messages[2]!)).toBe('Third real prompt');
   });
+
+  it('project() trims trailing assistant message with unanswered tool_calls', () => {
+    const history: ContextMessage[] = [
+      userMessage('hello'),
+      {
+        role: 'assistant',
+        content: [{ type: 'text', text: 'I will run a tool' }],
+        toolCalls: [{ type: 'function', id: 'call_1', name: 'Bash', arguments: '{}' }],
+      },
+      // No tool result for call_1 — session was killed.
+    ];
+    const messages = project(history);
+    // The assistant message with open tool_calls should be trimmed.
+    expect(messages).toHaveLength(1);
+    expect(messages[0]!.role).toBe('user');
+  });
+
+  it('project() keeps assistant message when all tool_calls are answered', () => {
+    const history: ContextMessage[] = [
+      userMessage('hello'),
+      {
+        role: 'assistant',
+        content: [{ type: 'text', text: 'I will run a tool' }],
+        toolCalls: [{ type: 'function', id: 'call_1', name: 'Bash', arguments: '{}' }],
+      },
+      {
+        role: 'tool',
+        content: [{ type: 'text', text: 'tool output' }],
+        toolCalls: [],
+        toolCallId: 'call_1',
+      },
+    ];
+    const messages = project(history);
+    // All three messages should be present.
+    expect(messages).toHaveLength(3);
+    expect(messages[1]!.role).toBe('assistant');
+    expect(messages[2]!.role).toBe('tool');
+  });
+
+  it('cleanupOrphanedToolCalls removes entire assistant and sibling tool messages', () => {
+    const ctx = testAgent();
+    ctx.configure();
+
+    // Simulate an assistant with two tool calls: call_A (orphaned) and call_B (answered).
+    ctx.dispatch({
+      type: 'context.append_loop_event',
+      event: { type: 'step.begin', uuid: 'step-multi', turnId: '', step: 1 },
+    });
+    ctx.dispatch({
+      type: 'context.append_loop_event',
+      event: {
+        type: 'tool.call',
+        uuid: 'tool-a',
+        stepUuid: 'step-multi',
+        turnId: '',
+        step: 1,
+        toolCallId: 'call_A',
+        name: 'Bash',
+        args: {},
+      },
+    });
+    ctx.dispatch({
+      type: 'context.append_loop_event',
+      event: {
+        type: 'tool.call',
+        uuid: 'tool-b',
+        stepUuid: 'step-multi',
+        turnId: '',
+        step: 1,
+        toolCallId: 'call_B',
+        name: 'Read',
+        args: {},
+      },
+    });
+    // Only call_B got a result before the crash.
+    ctx.dispatch({
+      type: 'context.append_loop_event',
+      event: {
+        type: 'tool.result',
+        parentUuid: 'tool-b',
+        toolCallId: 'call_B',
+        result: { output: 'file content' },
+      },
+    });
+
+    // The assistant and both tool messages should be in history.
+    const assistantBefore = ctx.agent.context.history.filter((m) => m.role === 'assistant');
+    const toolsBefore = ctx.agent.context.history.filter((m) => m.role === 'tool');
+    expect(assistantBefore).toHaveLength(1);
+    expect(toolsBefore).toHaveLength(1);
+
+    // Cleanup should remove the assistant AND the answered sibling tool message.
+    ctx.agent.context.cleanupOrphanedToolCalls();
+
+    const assistantAfter = ctx.agent.context.history.filter((m) => m.role === 'assistant');
+    const toolsAfter = ctx.agent.context.history.filter((m) => m.role === 'tool');
+    expect(assistantAfter).toHaveLength(0);
+    expect(toolsAfter).toHaveLength(0);
+  });
 });
 
 function userMessage(text: string, origin?: ContextMessage['origin']): ContextMessage {