From b55447cb25ca47e6c16ada202eb80560430b21f0 Mon Sep 17 00:00:00 2001
From: Kastrah <abdulkhaliqakinwunmi@gmail.com>
Date: Wed, 28 Jan 2026 03:24:01 +0100
Subject: [PATCH] fix: preserve tool_use + tool_result pairs during history
 truncation

- Fix limitHistoryTurns to include matching tool_use when truncating tool_result
- Move sanitizeToolUseResultPairing to attempt.ts for all providers
- Prevents 'tool id not found' error when dmHistoryLimit truncates sessions
---
 src/agents/pi-embedded-runner/history.ts     | 77 +++++++++++++++++++-
 src/agents/pi-embedded-runner/run/attempt.ts |  6 +-
 2 files changed, 79 insertions(+), 4 deletions(-)

diff --git a/src/agents/pi-embedded-runner/history.ts b/src/agents/pi-embedded-runner/history.ts
index b43f628e7..eadd38321 100644
--- a/src/agents/pi-embedded-runner/history.ts
+++ b/src/agents/pi-embedded-runner/history.ts
@@ -27,10 +27,43 @@ function isToolResultMessage(msg: AgentMessage): boolean {
   );
 }
 
+function extractToolUseIdsFromAssistant(msg: AgentMessage): string[] {
+  if (msg.role !== "assistant") return [];
+  const content = msg.content;
+  if (!Array.isArray(content)) return [];
+
+  const ids: string[] = [];
+  for (const block of content) {
+    if (!block || typeof block !== "object") continue;
+    const rec = block as { type?: unknown; id?: unknown };
+    if (
+      (rec.type === "toolCall" || rec.type === "toolUse" || rec.type === "functionCall") &&
+      typeof rec.id === "string"
+    ) {
+      ids.push(rec.id);
+    }
+  }
+  return ids;
+}
+
+function extractToolUseIdFromResult(msg: AgentMessage): string | null {
+  if (msg.role !== "user") return null;
+  const content = msg.content;
+  if (!Array.isArray(content) || content.length === 0) return null;
+
+  const block = content[0] as { tool_use_id?: unknown; toolCallId?: unknown };
+  const id = block.tool_use_id ?? block.toolCallId;
+  return typeof id === "string" ? id : null;
+}
+
 /**
  * Limits conversation history to the last N user turns (and their associated
  * assistant responses). This reduces token usage for long-running DM sessions.
  * Tool result messages are not counted as new user turns.
+ *
+ * CRITICAL: When truncating, we must preserve tool_use + tool_result pairs.
+ * A tool_result that follows its tool_use belongs to the same logical turn,
+ * even if they're separated by assistant responses.
  */
 export function limitHistoryTurns(
   messages: AgentMessage[],
@@ -43,16 +76,54 @@ export function limitHistoryTurns(
 
   for (let i = messages.length - 1; i >= 0; i--) {
     const msg = messages[i];
-    // Only count genuine user messages, not tool results
     if (msg.role === "user" && !isToolResultMessage(msg)) {
       userCount++;
       if (userCount > limit) {
-        return messages.slice(lastUserIndex);
+        break;
       }
       lastUserIndex = i;
     }
   }
-  return messages;
+
+  if (lastUserIndex === 0 || lastUserIndex === messages.length) {
+    return messages;
+  }
+
+  const slice = messages.slice(lastUserIndex);
+
+  const positionsToAdd = new Set<number>();
+  for (let i = 0; i < slice.length; i++) {
+    const msg = slice[i];
+    if (isToolResultMessage(msg)) {
+      const toolId = extractToolUseIdFromResult(msg);
+      if (toolId) {
+        let j = lastUserIndex + i - 1;
+        for (; j >= 0; j--) {
+          const prev = messages[j];
+          const toolIds = extractToolUseIdsFromAssistant(prev);
+          if (toolIds.includes(toolId)) {
+            positionsToAdd.add(j);
+            break;
+          }
+        }
+      }
+    }
+  }
+
+  if (positionsToAdd.size === 0) {
+    return slice;
+  }
+
+  const minPositionToAdd = Math.min(...positionsToAdd);
+  const result: AgentMessage[] = [];
+  for (let i = minPositionToAdd; i < messages.length; i++) {
+    const inSlice = i >= lastUserIndex;
+    const inPositionsToAdd = positionsToAdd.has(i);
+    if (inSlice || inPositionsToAdd) {
+      result.push(messages[i]);
+    }
+  }
+  return result;
 }
 
 /**
diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts
index 46a53bd8f..986ff5221 100644
--- a/src/agents/pi-embedded-runner/run/attempt.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.ts
@@ -42,6 +42,7 @@ import { createMoltbotCodingTools } from "../../pi-tools.js";
 import { resolveSandboxContext } from "../../sandbox.js";
 import { guardSessionManager } from "../../session-tool-result-guard-wrapper.js";
 import { resolveTranscriptPolicy } from "../../transcript-policy.js";
+import { sanitizeToolUseResultPairing } from "../../session-transcript-repair.js";
 import { acquireSessionWriteLock } from "../../session-write-lock.js";
 import {
   applySkillEnvOverrides,
@@ -531,8 +532,11 @@ export async function runEmbeddedAttempt(
         const validated = transcriptPolicy.validateAnthropicTurns
           ? validateAnthropicTurns(validatedGemini)
           : validatedGemini;
+        const repaired = transcriptPolicy.repairToolUseResultPairing
+          ? sanitizeToolUseResultPairing(validated)
+          : validated;
         const limited = limitHistoryTurns(
-          validated,
+          repaired,
           getDmHistoryLimitFromSessionKey(params.sessionKey, params.config),
         );
         cacheTrace?.recordStage("session:limited", { messages: limited });