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
This commit is contained in:
Kastrah 2026-01-28 03:24:01 +01:00
parent 711e087ba5
commit 8bdaba0ad8
2 changed files with 79 additions and 4 deletions

View File

@ -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 * Limits conversation history to the last N user turns (and their associated
* assistant responses). This reduces token usage for long-running DM sessions. * assistant responses). This reduces token usage for long-running DM sessions.
* Tool result messages are not counted as new user turns. * 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( export function limitHistoryTurns(
messages: AgentMessage[], messages: AgentMessage[],
@ -43,18 +76,56 @@ export function limitHistoryTurns(
for (let i = messages.length - 1; i >= 0; i--) { for (let i = messages.length - 1; i >= 0; i--) {
const msg = messages[i]; const msg = messages[i];
// Only count genuine user messages, not tool results
if (msg.role === "user" && !isToolResultMessage(msg)) { if (msg.role === "user" && !isToolResultMessage(msg)) {
userCount++; userCount++;
if (userCount > limit) { if (userCount > limit) {
return messages.slice(lastUserIndex); break;
} }
lastUserIndex = i; lastUserIndex = i;
} }
} }
if (lastUserIndex === 0 || lastUserIndex === messages.length) {
return messages; 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;
}
/** /**
* Extract provider + user ID from a session key and look up dmHistoryLimit. * Extract provider + user ID from a session key and look up dmHistoryLimit.
* Supports per-DM overrides and provider defaults. * Supports per-DM overrides and provider defaults.

View File

@ -42,6 +42,7 @@ import { createMoltbotCodingTools } from "../../pi-tools.js";
import { resolveSandboxContext } from "../../sandbox.js"; import { resolveSandboxContext } from "../../sandbox.js";
import { guardSessionManager } from "../../session-tool-result-guard-wrapper.js"; import { guardSessionManager } from "../../session-tool-result-guard-wrapper.js";
import { resolveTranscriptPolicy } from "../../transcript-policy.js"; import { resolveTranscriptPolicy } from "../../transcript-policy.js";
import { sanitizeToolUseResultPairing } from "../../session-transcript-repair.js";
import { acquireSessionWriteLock } from "../../session-write-lock.js"; import { acquireSessionWriteLock } from "../../session-write-lock.js";
import { import {
applySkillEnvOverrides, applySkillEnvOverrides,
@ -531,8 +532,11 @@ export async function runEmbeddedAttempt(
const validated = transcriptPolicy.validateAnthropicTurns const validated = transcriptPolicy.validateAnthropicTurns
? validateAnthropicTurns(validatedGemini) ? validateAnthropicTurns(validatedGemini)
: validatedGemini; : validatedGemini;
const repaired = transcriptPolicy.repairToolUseResultPairing
? sanitizeToolUseResultPairing(validated)
: validated;
const limited = limitHistoryTurns( const limited = limitHistoryTurns(
validated, repaired,
getDmHistoryLimitFromSessionKey(params.sessionKey, params.config), getDmHistoryLimitFromSessionKey(params.sessionKey, params.config),
); );
cacheTrace?.recordStage("session:limited", { messages: limited }); cacheTrace?.recordStage("session:limited", { messages: limited });