diff --git a/lib/commands/sweep.ts b/lib/commands/sweep.ts index 394c91b7..4e0ce331 100644 --- a/lib/commands/sweep.ts +++ b/lib/commands/sweep.ts @@ -15,7 +15,7 @@ import { formatPrunedItemsList } from "../ui/utils" import { getCurrentParams, getTotalToolTokens } from "../strategies/utils" import { buildToolIdList, isIgnoredUserMessage } from "../messages/utils" import { saveSessionState } from "../state/persistence" -import { isMessageCompacted } from "../shared-utils" +import { getLastUserMessage, isMessageCompacted } from "../shared-utils" import { getFilePathsFromParameters, isProtected } from "../protected-file-patterns" import { syncToolCache } from "../state/tool-cache" @@ -215,11 +215,21 @@ export async function handleSweepCommand(ctx: SweepCommandContext): Promise { + if (!state.prune.origins?.size) { + return + } + + const messageIds = new Set(messages.map((msg) => msg.info.id)) + let removedToolCount = 0 + let removedOriginCount = 0 + + for (const [toolId, origin] of state.prune.origins.entries()) { + if (!state.prune.tools.has(toolId)) { + state.prune.origins.delete(toolId) + removedOriginCount++ + continue + } + + if (!messageIds.has(origin.originMessageId)) { + state.prune.origins.delete(toolId) + removedOriginCount++ + if (state.prune.tools.delete(toolId)) { + removedToolCount++ + } + } + } + + if (removedToolCount > 0 || removedOriginCount > 0) { + logger.info("Synced prune origins", { + removedToolCount, + removedOriginCount, + }) + } +} diff --git a/lib/messages/utils.ts b/lib/messages/utils.ts index 1b9d31c7..59994f98 100644 --- a/lib/messages/utils.ts +++ b/lib/messages/utils.ts @@ -84,10 +84,15 @@ export const createSyntheticToolPart = ( const partId = generateStableId("prt_dcp_tool", deterministicSeed) const callId = generateStableId("call_dcp_tool", deterministicSeed) - // Gemini requires thoughtSignature bypass to accept synthetic tool parts + // Gemini requires a thought signature on synthetic function calls. + // Keep this metadata both on the part and on state so whichever + // conversion path is used can forward it to providerOptions. const toolPartMetadata = isGeminiModel(modelID) - ? { google: { thoughtSignature: "skip_thought_signature_validator" } } - : {} + ? { + google: { thoughtSignature: "skip_thought_signature_validator" }, + vertex: { thoughtSignature: "skip_thought_signature_validator" }, + } + : undefined return { id: partId, @@ -96,14 +101,15 @@ export const createSyntheticToolPart = ( type: "tool" as const, callID: callId, tool: "context_info", + ...(toolPartMetadata ? { metadata: toolPartMetadata } : {}), state: { status: "completed" as const, input: {}, output: content, title: "Context Info", - metadata: toolPartMetadata, + ...(toolPartMetadata ? { metadata: toolPartMetadata } : {}), time: { start: now, end: now }, - }, + } as any, } } diff --git a/lib/state/persistence.ts b/lib/state/persistence.ts index 4652ce18..a57e1be3 100644 --- a/lib/state/persistence.ts +++ b/lib/state/persistence.ts @@ -8,7 +8,7 @@ import * as fs from "fs/promises" import { existsSync } from "fs" import { homedir } from "os" import { join } from "path" -import type { SessionState, SessionStats, CompressSummary } from "./types" +import type { SessionState, SessionStats, CompressSummary, PruneOrigin } from "./types" import type { Logger } from "../logger" /** Prune state as stored on disk */ @@ -16,6 +16,7 @@ export interface PersistedPrune { // New format: tool/message IDs with token counts tools?: Record messages?: Record + origins?: Record // Legacy format: plain ID arrays (backward compatibility) toolIds?: string[] messageIds?: string[] @@ -64,6 +65,7 @@ export async function saveSessionState( prune: { tools: Object.fromEntries(sessionState.prune.tools), messages: Object.fromEntries(sessionState.prune.messages), + origins: Object.fromEntries(sessionState.prune.origins), }, compressSummaries: sessionState.compressSummaries, stats: sessionState.stats, diff --git a/lib/state/state.ts b/lib/state/state.ts index a9f34ddf..d59fa8f3 100644 --- a/lib/state/state.ts +++ b/lib/state/state.ts @@ -7,6 +7,7 @@ import { countTurns, resetOnCompaction, loadPruneMap, + loadPruneOriginMap, } from "./utils" import { getLastUserMessage } from "../shared-utils" @@ -67,6 +68,7 @@ export function createSessionState(): SessionState { prune: { tools: new Map(), messages: new Map(), + origins: new Map(), }, compressSummaries: [], stats: { @@ -97,6 +99,7 @@ export function resetSessionState(state: SessionState): void { state.prune = { tools: new Map(), messages: new Map(), + origins: new Map(), } state.compressSummaries = [] state.stats = { @@ -151,6 +154,7 @@ export async function ensureSessionInitialized( state.prune.tools = loadPruneMap(persisted.prune.tools, persisted.prune.toolIds) state.prune.messages = loadPruneMap(persisted.prune.messages, persisted.prune.messageIds) + state.prune.origins = loadPruneOriginMap(persisted.prune.origins) state.compressSummaries = persisted.compressSummaries || [] state.stats = { pruneTokenCounter: persisted.stats?.pruneTokenCounter || 0, diff --git a/lib/state/types.ts b/lib/state/types.ts index 218756cc..683f1ced 100644 --- a/lib/state/types.ts +++ b/lib/state/types.ts @@ -27,9 +27,23 @@ export interface CompressSummary { summary: string } +export type PruneOriginSource = + | "prune" + | "distill" + | "sweep" + | "deduplication" + | "supersedeWrites" + | "purgeErrors" + +export interface PruneOrigin { + source: PruneOriginSource + originMessageId: string +} + export interface Prune { tools: Map messages: Map + origins: Map } export interface PendingManualTrigger { diff --git a/lib/state/utils.ts b/lib/state/utils.ts index f5e0918b..163ccf18 100644 --- a/lib/state/utils.ts +++ b/lib/state/utils.ts @@ -1,4 +1,4 @@ -import type { SessionState, WithParts } from "./types" +import type { PruneOrigin, SessionState, WithParts } from "./types" import { isMessageCompacted } from "../shared-utils" export async function isSubAgentSession(client: any, sessionID: string): Promise { @@ -45,10 +45,32 @@ export function loadPruneMap( return new Map() } +export function loadPruneOriginMap(obj?: Record): Map { + if (!obj || typeof obj !== "object") { + return new Map() + } + + const entries: [string, PruneOrigin][] = [] + for (const [toolId, origin] of Object.entries(obj)) { + if ( + origin && + typeof origin === "object" && + typeof origin.source === "string" && + typeof origin.originMessageId === "string" && + origin.originMessageId.length > 0 + ) { + entries.push([toolId, origin]) + } + } + + return new Map(entries) +} + export function resetOnCompaction(state: SessionState): void { state.toolParameters.clear() state.prune.tools = new Map() state.prune.messages = new Map() + state.prune.origins = new Map() state.compressSummaries = [] state.messageIds = { byRawId: new Map(), diff --git a/lib/strategies/deduplication.ts b/lib/strategies/deduplication.ts index aff2d019..5c6f5748 100644 --- a/lib/strategies/deduplication.ts +++ b/lib/strategies/deduplication.ts @@ -2,6 +2,7 @@ import { PluginConfig } from "../config" import { Logger } from "../logger" import type { SessionState, WithParts } from "../state" import { getFilePathsFromParameters, isProtected } from "../protected-file-patterns" +import { getLastUserMessage } from "../shared-utils" import { getTotalToolTokens } from "./utils" /** @@ -79,11 +80,21 @@ export const deduplicate = ( } state.stats.totalPruneTokens += getTotalToolTokens(state, newPruneIds) + const decisionMessageId = getLastUserMessage(messages)?.info.id || "" if (newPruneIds.length > 0) { + if (!decisionMessageId) { + logger.warn("Deduplication prune origin unavailable - missing user message") + } for (const id of newPruneIds) { const entry = state.toolParameters.get(id) state.prune.tools.set(id, entry?.tokenCount ?? 0) + if (decisionMessageId) { + state.prune.origins.set(id, { + source: "deduplication", + originMessageId: decisionMessageId, + }) + } } logger.debug(`Marked ${newPruneIds.length} duplicate tool calls for pruning`) } diff --git a/lib/strategies/purge-errors.ts b/lib/strategies/purge-errors.ts index 25939eda..b579cafe 100644 --- a/lib/strategies/purge-errors.ts +++ b/lib/strategies/purge-errors.ts @@ -2,6 +2,7 @@ import { PluginConfig } from "../config" import { Logger } from "../logger" import type { SessionState, WithParts } from "../state" import { getFilePathsFromParameters, isProtected } from "../protected-file-patterns" +import { getLastUserMessage } from "../shared-utils" import { getTotalToolTokens } from "./utils" /** @@ -72,10 +73,20 @@ export const purgeErrors = ( } if (newPruneIds.length > 0) { + const decisionMessageId = getLastUserMessage(messages)?.info.id || "" + if (!decisionMessageId) { + logger.warn("Purge errors prune origin unavailable - missing user message") + } state.stats.totalPruneTokens += getTotalToolTokens(state, newPruneIds) for (const id of newPruneIds) { const entry = state.toolParameters.get(id) state.prune.tools.set(id, entry?.tokenCount ?? 0) + if (decisionMessageId) { + state.prune.origins.set(id, { + source: "purgeErrors", + originMessageId: decisionMessageId, + }) + } } logger.debug( `Marked ${newPruneIds.length} error tool calls for pruning (older than ${turnThreshold} turns)`, diff --git a/lib/strategies/supersede-writes.ts b/lib/strategies/supersede-writes.ts index e61b1d0a..ea8b142f 100644 --- a/lib/strategies/supersede-writes.ts +++ b/lib/strategies/supersede-writes.ts @@ -2,6 +2,7 @@ import { PluginConfig } from "../config" import { Logger } from "../logger" import type { SessionState, WithParts } from "../state" import { getFilePathsFromParameters, isProtected } from "../protected-file-patterns" +import { getLastUserMessage } from "../shared-utils" import { getTotalToolTokens } from "./utils" /** @@ -105,10 +106,20 @@ export const supersedeWrites = ( } if (newPruneIds.length > 0) { + const decisionMessageId = getLastUserMessage(messages)?.info.id || "" + if (!decisionMessageId) { + logger.warn("Supersede writes prune origin unavailable - missing user message") + } state.stats.totalPruneTokens += getTotalToolTokens(state, newPruneIds) for (const id of newPruneIds) { const entry = state.toolParameters.get(id) state.prune.tools.set(id, entry?.tokenCount ?? 0) + if (decisionMessageId) { + state.prune.origins.set(id, { + source: "supersedeWrites", + originMessageId: decisionMessageId, + }) + } } logger.debug(`Marked ${newPruneIds.length} superseded write tool calls for pruning`) } diff --git a/lib/tools/distill.ts b/lib/tools/distill.ts index b67d56b4..6cbd56d6 100644 --- a/lib/tools/distill.ts +++ b/lib/tools/distill.ts @@ -53,6 +53,7 @@ export function createDistillTool(ctx: PruneToolContext): ReturnType { const { client, state, logger, config, workingDirectory } = ctx @@ -114,6 +115,12 @@ export async function executePruneOperation( continue } + if (state.prune.tools.has(id)) { + logger.debug("Rejecting prune request - already pruned", { index, id }) + skippedIds.push(index.toString()) + continue + } + validNumericIds.push(index) } @@ -126,9 +133,21 @@ export async function executePruneOperation( } const pruneToolIds: string[] = validNumericIds.map((index) => toolIdList[index]) + const originMessageId = + typeof toolCtx.messageID === "string" && toolCtx.messageID.length > 0 + ? toolCtx.messageID + : "" + + if (!originMessageId) { + logger.warn(`Missing tool message ID for ${toolName} prune origin tracking`) + } + for (const id of pruneToolIds) { const entry = state.toolParameters.get(id) state.prune.tools.set(id, entry?.tokenCount ?? 0) + if (originMessageId) { + state.prune.origins.set(id, { source, originMessageId }) + } } const toolMetadata = new Map() @@ -167,7 +186,7 @@ export async function executePruneOperation( let result = formatPruningResultForTool(pruneToolIds, toolMetadata, workingDirectory) if (skippedIds.length > 0) { - result += `\n\nNote: ${skippedIds.length} IDs were skipped (invalid, protected, or missing metadata): ${skippedIds.join(", ")}` + result += `\n\nNote: ${skippedIds.length} IDs were skipped (invalid, protected, already pruned, or missing metadata): ${skippedIds.join(", ")}` } return result } diff --git a/lib/tools/prune.ts b/lib/tools/prune.ts index 17065aa9..1cd19a51 100644 --- a/lib/tools/prune.ts +++ b/lib/tools/prune.ts @@ -30,7 +30,7 @@ export function createPruneTool(ctx: PruneToolContext): ReturnType const numericIds = args.ids const reason = "noise" - return executePruneOperation(ctx, toolCtx, numericIds, reason, "Prune") + return executePruneOperation(ctx, toolCtx, numericIds, reason, "Prune", "prune") }, }) }