From ff9a6826ddd50e53cc832a4699eececcd7f8a1bb Mon Sep 17 00:00:00 2001 From: can1357 Date: Sun, 31 May 2026 07:12:27 +0200 Subject: [PATCH] fix(compaction): made `read` tool-results prunable except for `skill://` paths - Replaced flat `protectedTools: string[]` with `ProtectedToolMatcher[]` supporting predicate functions. - Regular file/URL `read` calls are now eligible for pruning and shake compaction. - `read` calls whose `path` starts with `skill://` remain protected like native `skill` results. - Added `collectToolCallsById` to correlate tool results with their originating call arguments. --- packages/agent/CHANGELOG.md | 3 + packages/agent/src/compaction/pruning.ts | 15 ++- .../agent/src/compaction/tool-protection.ts | 46 +++++++++ packages/agent/test/tool-protection.test.ts | 93 +++++++++++++++++++ 4 files changed, 153 insertions(+), 4 deletions(-) create mode 100644 packages/agent/src/compaction/tool-protection.ts create mode 100644 packages/agent/test/tool-protection.test.ts diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 953b9132e..b8ea8d0f9 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -1,6 +1,9 @@ # Changelog ## [Unreleased] +### Fixed + +- Fixed tool-output pruning and shake protection for `read`: ordinary file/URL reads are now eligible for compaction, while `read` calls whose `path` starts with `skill://` remain protected like native `skill` results. ## [15.5.15] - 2026-05-30 ### Added diff --git a/packages/agent/src/compaction/pruning.ts b/packages/agent/src/compaction/pruning.ts index 7e99cb87e..9aee37131 100644 --- a/packages/agent/src/compaction/pruning.ts +++ b/packages/agent/src/compaction/pruning.ts @@ -6,20 +6,26 @@ import type { ToolResultMessage } from "@oh-my-pi/pi-ai"; import type { AgentMessage } from "../types"; import { estimateTokens } from "./compaction"; import type { SessionEntry, SessionMessageEntry } from "./entries"; +import { + collectToolCallsById, + isProtectedToolResult, + isSkillReadToolResult, + type ProtectedToolMatcher, +} from "./tool-protection"; export interface PruneConfig { /** Keep the most recent tool output tokens intact. */ protectTokens: number; /** Only prune if total savings meets this threshold. */ minimumSavings: number; - /** Tool names that should never be pruned. */ - protectedTools: string[]; + /** Tool-result protection matchers. String entries protect every result from that tool; predicates may inspect the paired tool call. */ + protectedTools: ProtectedToolMatcher[]; } export const DEFAULT_PRUNE_CONFIG: PruneConfig = { protectTokens: 40_000, minimumSavings: 20_000, - protectedTools: ["skill", "read"], + protectedTools: ["skill", isSkillReadToolResult], }; export interface PruneResult { @@ -49,6 +55,7 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig = let prunedCount = 0; const candidates: Array<{ entry: SessionMessageEntry; tokens: number }> = []; + const toolCallsById = collectToolCallsById(entries); for (let i = entries.length - 1; i >= 0; i--) { const entry = entries[i]; @@ -56,7 +63,7 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig = if (!message) continue; const tokens = estimateTokens(message as AgentMessage); - const isProtected = config.protectedTools.includes(message.toolName); + const isProtected = isProtectedToolResult(message, toolCallsById.get(message.toolCallId), config.protectedTools); if (message.prunedAt !== undefined) { accumulatedTokens += tokens; diff --git a/packages/agent/src/compaction/tool-protection.ts b/packages/agent/src/compaction/tool-protection.ts new file mode 100644 index 000000000..e2aec93aa --- /dev/null +++ b/packages/agent/src/compaction/tool-protection.ts @@ -0,0 +1,46 @@ +import type { ToolResultMessage } from "@oh-my-pi/pi-ai"; +import type { AgentToolCall } from "../types"; +import type { SessionEntry } from "./entries"; + +export interface ProtectedToolContext { + readonly toolResult: ToolResultMessage; + readonly toolCall: AgentToolCall | undefined; +} + +export type ProtectedToolMatcher = string | ((context: ProtectedToolContext) => boolean); + +const SKILL_INTERNAL_URL_PREFIX = "skill://"; + +export function collectToolCallsById(entries: readonly SessionEntry[]): Map { + const toolCalls = new Map(); + for (const entry of entries) { + if (entry.type !== "message") continue; + const message = entry.message; + if (message.role !== "assistant") continue; + for (const block of message.content) { + if (block.type === "toolCall") toolCalls.set(block.id, block); + } + } + return toolCalls; +} + +export function isSkillReadToolResult({ toolResult, toolCall }: ProtectedToolContext): boolean { + if (toolResult.toolName !== "read" || toolCall?.name !== "read") return false; + const path = (toolCall.arguments as Record).path; + return typeof path === "string" && path.startsWith(SKILL_INTERNAL_URL_PREFIX); +} + +export function isProtectedToolResult( + toolResult: ToolResultMessage, + toolCall: AgentToolCall | undefined, + matchers: readonly ProtectedToolMatcher[], +): boolean { + for (const matcher of matchers) { + if (typeof matcher === "string") { + if (toolResult.toolName === matcher) return true; + continue; + } + if (matcher({ toolResult, toolCall })) return true; + } + return false; +} diff --git a/packages/agent/test/tool-protection.test.ts b/packages/agent/test/tool-protection.test.ts new file mode 100644 index 000000000..59d427fe2 --- /dev/null +++ b/packages/agent/test/tool-protection.test.ts @@ -0,0 +1,93 @@ +import { describe, expect, it } from "bun:test"; +import type { SessionMessageEntry } from "@oh-my-pi/pi-agent-core/compaction/entries"; +import { DEFAULT_PRUNE_CONFIG, pruneToolOutputs } from "@oh-my-pi/pi-agent-core/compaction/pruning"; +import { AGGRESSIVE_SHAKE_CONFIG, collectShakeRegions } from "@oh-my-pi/pi-agent-core/compaction/shake"; +import type { AssistantMessage, TextContent, ToolResultMessage, Usage } from "@oh-my-pi/pi-ai"; + +function usage(): Usage { + return { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function messageEntry(id: string, message: AssistantMessage | ToolResultMessage): SessionMessageEntry { + return { + type: "message", + id, + parentId: null, + timestamp: "2026-05-31T00:00:00.000Z", + message, + }; +} + +function assistantReadCall(toolCallId: string, path: string): SessionMessageEntry { + return messageEntry(`assistant-${toolCallId}`, { + role: "assistant", + content: [{ type: "toolCall", id: toolCallId, name: "read", arguments: { path } }], + api: "mock", + provider: "mock", + model: "mock-model", + usage: usage(), + stopReason: "toolUse", + timestamp: 0, + }); +} + +function readResult(toolCallId: string, text: string): SessionMessageEntry { + const content: TextContent[] = [{ type: "text", text }]; + return messageEntry(`result-${toolCallId}`, { + role: "toolResult", + toolCallId, + toolName: "read", + content, + isError: false, + timestamp: 0, + }); +} + +describe("conditional tool-result protection", () => { + it("prunes regular read results but keeps skill:// reads", () => { + const skillResult = readResult("skill-read", "skill read output that must remain intact"); + const fileResult = readResult("file-read", "file read output that can be pruned"); + const entries = [ + assistantReadCall("skill-read", "skill://session-memory"), + skillResult, + assistantReadCall("file-read", "packages/agent/src/index.ts"), + fileResult, + ]; + + const result = pruneToolOutputs(entries, { ...DEFAULT_PRUNE_CONFIG, protectTokens: 0, minimumSavings: 0 }); + + expect(result.prunedCount).toBe(1); + expect((skillResult.message as ToolResultMessage).prunedAt).toBeUndefined(); + expect((skillResult.message as ToolResultMessage).content).toEqual([ + { type: "text", text: "skill read output that must remain intact" }, + ]); + expect(typeof (fileResult.message as ToolResultMessage).prunedAt).toBe("number"); + expect(((fileResult.message as ToolResultMessage).content[0] as TextContent).text).toStartWith( + "[Output truncated - ", + ); + }); + + it("shakes regular read results but excludes skill:// reads", () => { + const skillResult = readResult("skill-read", "skill read output that must not be shaken"); + const fileResult = readResult("file-read", "file read output that is eligible for shake"); + const entries = [ + assistantReadCall("skill-read", "skill://session-memory"), + skillResult, + assistantReadCall("file-read", "src/index.ts"), + fileResult, + ]; + + const regions = collectShakeRegions(entries, AGGRESSIVE_SHAKE_CONFIG); + + expect(regions).toHaveLength(1); + expect(regions[0]?.kind).toBe("toolResult"); + expect(regions[0]?.entry).toBe(fileResult); + }); +});