From 04f2e405aa212ff0bd05d6f81f361477ef5db5e2 Mon Sep 17 00:00:00 2001 From: Rodrigo Uroz Date: Fri, 27 Feb 2026 18:46:56 +0000 Subject: [PATCH] Compaction/Safeguard: preserve tool results in recent tail --- .../compaction-safeguard.test.ts | 37 ++++++++++++ .../pi-extensions/compaction-safeguard.ts | 58 ++++++++++++++++++- 2 files changed, 92 insertions(+), 3 deletions(-) diff --git a/src/agents/pi-extensions/compaction-safeguard.test.ts b/src/agents/pi-extensions/compaction-safeguard.test.ts index 8e0f0453196..750eb2b4f38 100644 --- a/src/agents/pi-extensions/compaction-safeguard.test.ts +++ b/src/agents/pi-extensions/compaction-safeguard.test.ts @@ -460,6 +460,7 @@ describe("compaction-safeguard recent-turn preservation", () => { expect(split.preservedMessages.map((msg) => msg.role)).toEqual([ "user", "assistant", + "toolResult", "assistant", ]); expect( @@ -475,6 +476,42 @@ describe("compaction-safeguard recent-turn preservation", () => { expect(summarizableToolResultIds).not.toContain("call_recent"); }); + it("includes preserved tool results in the preserved-turns section", () => { + const split = splitPreservedRecentTurns({ + messages: [ + { role: "user", content: "older ask", timestamp: 1 }, + { + role: "assistant", + content: [{ type: "text", text: "older answer" }], + timestamp: 2, + } as unknown as AgentMessage, + { role: "user", content: "recent ask", timestamp: 3 }, + { + role: "assistant", + content: [{ type: "toolCall", id: "call_recent", name: "read", arguments: {} }], + timestamp: 4, + } as unknown as AgentMessage, + { + role: "toolResult", + toolCallId: "call_recent", + toolName: "read", + content: [{ type: "text", text: "recent raw output" }], + timestamp: 5, + } as unknown as AgentMessage, + { + role: "assistant", + content: [{ type: "text", text: "recent final answer" }], + timestamp: 6, + } as unknown as AgentMessage, + ], + recentTurnsPreserve: 1, + }); + + const section = formatPreservedTurnsSection(split.preservedMessages); + expect(section).toContain("- Tool result (read): recent raw output"); + expect(section).toContain("- User: recent ask"); + }); + it("formats preserved non-text messages with placeholders", () => { const section = formatPreservedTurnsSection([ { diff --git a/src/agents/pi-extensions/compaction-safeguard.ts b/src/agents/pi-extensions/compaction-safeguard.ts index 68ab218f588..df3889f5704 100644 --- a/src/agents/pi-extensions/compaction-safeguard.ts +++ b/src/agents/pi-extensions/compaction-safeguard.ts @@ -19,6 +19,7 @@ import { } from "../compaction.js"; import { collectTextContentBlocks } from "../content-blocks.js"; import { repairToolUseResultPairing } from "../session-transcript-repair.js"; +import { extractToolCallsFromAssistant, extractToolResultId } from "../tool-call-id.js"; import { getCompactionSafeguardRuntime } from "./compaction-safeguard-runtime.js"; const log = createSubsystemLogger("compaction-safeguard"); @@ -279,6 +280,46 @@ function splitPreservedRecentTurns(params: { if (preservedIndexSet.size === 0) { return { summarizableMessages: params.messages, preservedMessages: [] }; } + const preservedToolCallIds = new Set(); + for (let i = 0; i < params.messages.length; i += 1) { + if (!preservedIndexSet.has(i)) { + continue; + } + const message = params.messages[i]; + const role = (message as { role?: unknown }).role; + if (role !== "assistant") { + continue; + } + const toolCalls = extractToolCallsFromAssistant( + message as Extract, + ); + for (const toolCall of toolCalls) { + preservedToolCallIds.add(toolCall.id); + } + } + if (preservedToolCallIds.size > 0) { + let preservedStartIndex = -1; + for (let i = 0; i < params.messages.length; i += 1) { + if (preservedIndexSet.has(i)) { + preservedStartIndex = i; + break; + } + } + if (preservedStartIndex >= 0) { + for (let i = preservedStartIndex; i < params.messages.length; i += 1) { + const message = params.messages[i]; + if ((message as { role?: unknown }).role !== "toolResult") { + continue; + } + const toolResultId = extractToolResultId( + message as Extract, + ); + if (toolResultId && preservedToolCallIds.has(toolResultId)) { + preservedIndexSet.add(i); + } + } + } + } const summarizableMessages = params.messages.filter((_, idx) => !preservedIndexSet.has(idx)); // Preserving recent assistant turns can orphan downstream toolResult messages. // Repair pairings here so compaction summarization doesn't trip strict providers. @@ -287,7 +328,7 @@ function splitPreservedRecentTurns(params: { .filter((_, idx) => preservedIndexSet.has(idx)) .filter((msg) => { const role = (msg as { role?: unknown }).role; - return role === "user" || role === "assistant"; + return role === "user" || role === "assistant" || role === "toolResult"; }); return { summarizableMessages: repairedSummarizableMessages, preservedMessages }; } @@ -298,7 +339,18 @@ function formatPreservedTurnsSection(messages: AgentMessage[]): string { } const lines = messages .map((message) => { - const role = message.role === "assistant" ? "Assistant" : "User"; + let roleLabel: string; + if (message.role === "assistant") { + roleLabel = "Assistant"; + } else if (message.role === "user") { + roleLabel = "User"; + } else if (message.role === "toolResult") { + const toolName = (message as { toolName?: unknown }).toolName; + const safeToolName = typeof toolName === "string" && toolName.trim() ? toolName : "tool"; + roleLabel = `Tool result (${safeToolName})`; + } else { + return null; + } const text = extractMessageText(message); const nonTextPlaceholder = formatNonTextPlaceholder( (message as { content?: unknown }).content, @@ -311,7 +363,7 @@ function formatPreservedTurnsSection(messages: AgentMessage[]): string { rendered.length > MAX_RECENT_TURN_TEXT_CHARS ? `${rendered.slice(0, MAX_RECENT_TURN_TEXT_CHARS)}...` : rendered; - return `- ${role}: ${trimmed}`; + return `- ${roleLabel}: ${trimmed}`; }) .filter((line): line is string => Boolean(line)); if (lines.length === 0) {