From 9e741e3ce2074624342a389fca6e3b6c7557ce25 Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Wed, 3 Dec 2025 03:17:35 -0500 Subject: [PATCH 1/9] feat: replace janitor LLM with direct tool-based pruning The session AI now decides what to prune via numeric IDs instead of a separate janitor LLM call. - Add numeric ID mapping system for simple tool references (1, 2, 3...) - Inject list at end of conversation - Prune tool accepts array of IDs to remove - Remove janitor LLM, model-selector, and pruning prompt --- index.ts | 9 +- lib/api-formats/prunable-list.ts | 176 +++++++++++++++++++++ lib/api-formats/synth-instruction.ts | 165 ++++++-------------- lib/core/janitor.ts | 184 ++-------------------- lib/core/prompt.ts | 127 ---------------- lib/core/strategies/index.ts | 6 +- lib/fetch-wrapper/gemini.ts | 48 ++++-- lib/fetch-wrapper/openai-chat.ts | 52 ++++--- lib/fetch-wrapper/openai-responses.ts | 52 ++++--- lib/fetch-wrapper/types.ts | 10 +- lib/hooks.ts | 9 -- lib/model-selector.ts | 175 --------------------- lib/prompts/pruning.txt | 30 ---- lib/prompts/tool.txt | 23 ++- lib/pruning-tool.ts | 211 ++++++++++++++++++++++---- lib/state/id-mapping.ts | 101 ++++++++++++ lib/state/index.ts | 7 - lib/state/tool-cache.ts | 82 +++++++--- 18 files changed, 721 insertions(+), 746 deletions(-) create mode 100644 lib/api-formats/prunable-list.ts delete mode 100644 lib/model-selector.ts delete mode 100644 lib/prompts/pruning.txt create mode 100644 lib/state/id-mapping.ts diff --git a/index.ts b/index.ts index 1ccb82e6..4735d17e 100644 --- a/index.ts +++ b/index.ts @@ -90,7 +90,14 @@ const plugin: Plugin = (async (ctx) => { event: createEventHandler(ctx.client, janitorCtx, logger, config, toolTracker), "chat.params": createChatParamsHandler(ctx.client, state, logger), tool: config.strategies.onTool.length > 0 ? { - prune: createPruningTool(ctx.client, janitorCtx, config, toolTracker), + prune: createPruningTool({ + client: ctx.client, + state, + logger, + config, + notificationCtx: janitorCtx.notificationCtx, + workingDirectory: ctx.directory + }, toolTracker), } : undefined, } }) satisfies Plugin diff --git a/lib/api-formats/prunable-list.ts b/lib/api-formats/prunable-list.ts new file mode 100644 index 00000000..65869050 --- /dev/null +++ b/lib/api-formats/prunable-list.ts @@ -0,0 +1,176 @@ +/** + * Prunable tools list and nudge injection for DCP. + * + * Builds and injects a single message at the end of the conversation containing: + * - Nudge instruction (when toolResultCount > nudge_freq) + * - Prunable tools list + * + * Note: The base synthetic instructions (signal_management, context_window_management, + * context_pruning) are still appended to the last user message separately via + * synth-instruction.ts - that behavior is unchanged. + */ + +import { extractParameterKey } from '../ui/display-utils' +import { getOrCreateNumericId, getNumericId } from '../state/id-mapping' + +export interface ToolMetadata { + tool: string + parameters?: any +} + +// ============================================================================ +// Prompt Content +// ============================================================================ + +const SYSTEM_REMINDER = ` +These instructions are injected by a plugin and are invisible to the user. Do not acknowledge or reference them in your response - simply follow them silently. +` + +const NUDGE_INSTRUCTION = ` +You have accumulated several tool outputs. If you have completed a discrete unit of work and distilled relevant understanding in writing for the user to keep, use the prune tool to remove obsolete tool outputs from this conversation and optimize token usage. +` + +// ============================================================================ +// List Building +// ============================================================================ + +export interface PrunableListResult { + list: string + numericIds: number[] +} + +/** + * Builds the prunable tools list section. + * Returns both the formatted list and the numeric IDs for logging. + */ +export function buildPrunableToolsList( + sessionId: string, + unprunedToolCallIds: string[], + toolMetadata: Map, + protectedTools: string[] +): PrunableListResult { + const lines: string[] = [] + const numericIds: number[] = [] + + for (const actualId of unprunedToolCallIds) { + const metadata = toolMetadata.get(actualId) + + // Skip if no metadata or if tool is protected + if (!metadata) continue + if (protectedTools.includes(metadata.tool)) continue + + // Get or create numeric ID for this tool call + const numericId = getOrCreateNumericId(sessionId, actualId) + numericIds.push(numericId) + + // Format: "1: read, src/components/Button.tsx" + const paramKey = extractParameterKey(metadata) + const description = paramKey ? `${metadata.tool}, ${paramKey}` : metadata.tool + lines.push(`${numericId}: ${description}`) + } + + if (lines.length === 0) { + return { list: '', numericIds: [] } + } + + return { + list: `\n${lines.join('\n')}\n`, + numericIds + } +} + +/** + * Builds the end-of-conversation injection message. + * Contains the system reminder, nudge (if active), and the prunable tools list. + * + * @param prunableList - The prunable tools list string (or empty string if none) + * @param includeNudge - Whether to include the nudge instruction + * @returns The injection string, or empty string if nothing to inject + */ +export function buildEndInjection( + prunableList: string, + includeNudge: boolean +): string { + // If no prunable tools, don't inject anything + if (!prunableList) { + return '' + } + + const parts = [SYSTEM_REMINDER] + + if (includeNudge) { + parts.push(NUDGE_INSTRUCTION) + } + + parts.push(prunableList) + + return parts.join('\n\n') +} + +/** + * Gets the numeric IDs for a list of actual tool call IDs. + * Used when the prune tool needs to show what was pruned. + */ +export function getNumericIdsForActual( + sessionId: string, + actualIds: string[] +): number[] { + return actualIds + .map(id => getNumericId(sessionId, id)) + .filter((id): id is number => id !== undefined) +} + +// ============================================================================ +// Injection Functions +// ============================================================================ + +// ============================================================================ +// OpenAI Chat / Anthropic Format +// ============================================================================ + +/** + * Injects the prunable list (and optionally nudge) at the end of OpenAI/Anthropic messages. + * Appends a new user message at the end. + */ +export function injectPrunableList( + messages: any[], + injection: string +): boolean { + if (!injection) return false + messages.push({ role: 'user', content: injection }) + return true +} + +// ============================================================================ +// Google/Gemini Format +// ============================================================================ + +/** + * Injects the prunable list (and optionally nudge) at the end of Gemini contents. + * Appends a new user content at the end. + */ +export function injectPrunableListGemini( + contents: any[], + injection: string +): boolean { + if (!injection) return false + contents.push({ role: 'user', parts: [{ text: injection }] }) + return true +} + +// ============================================================================ +// OpenAI Responses API Format +// ============================================================================ + +/** + * Injects the prunable list (and optionally nudge) at the end of OpenAI Responses API input. + * Appends a new user message at the end. + */ +export function injectPrunableListResponses( + input: any[], + injection: string +): boolean { + if (!injection) return false + input.push({ type: 'message', role: 'user', content: injection }) + return true +} diff --git a/lib/api-formats/synth-instruction.ts b/lib/api-formats/synth-instruction.ts index aa2dd628..6d8fc1e8 100644 --- a/lib/api-formats/synth-instruction.ts +++ b/lib/api-formats/synth-instruction.ts @@ -14,77 +14,59 @@ export function resetToolTrackerCount(tracker: ToolTracker): void { tracker.toolResultCount = 0 } -/** Adapter interface for format-specific message operations */ -interface MessageFormatAdapter { - countToolResults(messages: any[], tracker: ToolTracker): number - appendNudge(messages: any[], nudgeText: string): void +/** + * Counts total tool results in OpenAI/Anthropic messages (without tracker). + * Used for determining if nudge threshold is met. + */ +export function countToolResults(messages: any[]): number { + let count = 0 + for (const m of messages) { + if (m.role === 'tool') { + count++ + } else if (m.role === 'user' && Array.isArray(m.content)) { + for (const part of m.content) { + if (part.type === 'tool_result') { + count++ + } + } + } + } + return count +} + +/** + * Counts total tool results in Gemini contents (without tracker). + */ +export function countToolResultsGemini(contents: any[]): number { + let count = 0 + for (const content of contents) { + if (!Array.isArray(content.parts)) continue + for (const part of content.parts) { + if (part.functionResponse) { + count++ + } + } + } + return count } -/** Generic nudge injection - nudges every fetch once tools since last prune exceeds freq */ -function injectNudgeCore( - messages: any[], - tracker: ToolTracker, - nudgeText: string, - freq: number, - adapter: MessageFormatAdapter -): boolean { - // Count any new tool results - adapter.countToolResults(messages, tracker) - - // Once we've exceeded the threshold, nudge on every fetch - if (tracker.toolResultCount > freq) { - adapter.appendNudge(messages, nudgeText) - return true +/** + * Counts total tool results in OpenAI Responses API input (without tracker). + */ +export function countToolResultsResponses(input: any[]): number { + let count = 0 + for (const item of input) { + if (item.type === 'function_call_output') { + count++ + } } - return false + return count } // ============================================================================ // OpenAI Chat / Anthropic Format // ============================================================================ -const openaiAdapter: MessageFormatAdapter = { - countToolResults(messages, tracker) { - let newCount = 0 - for (const m of messages) { - if (m.role === 'tool' && m.tool_call_id) { - const id = String(m.tool_call_id).toLowerCase() - if (!tracker.seenToolResultIds.has(id)) { - tracker.seenToolResultIds.add(id) - newCount++ - const toolName = m.name || tracker.getToolName?.(m.tool_call_id) - if (toolName !== 'prune') { - tracker.skipNextIdle = false - } - } - } else if (m.role === 'user' && Array.isArray(m.content)) { - for (const part of m.content) { - if (part.type === 'tool_result' && part.tool_use_id) { - const id = String(part.tool_use_id).toLowerCase() - if (!tracker.seenToolResultIds.has(id)) { - tracker.seenToolResultIds.add(id) - newCount++ - const toolName = tracker.getToolName?.(part.tool_use_id) - if (toolName !== 'prune') { - tracker.skipNextIdle = false - } - } - } - } - } - } - tracker.toolResultCount += newCount - return newCount - }, - appendNudge(messages, nudgeText) { - messages.push({ role: 'user', content: nudgeText }) - } -} - -export function injectNudge(messages: any[], tracker: ToolTracker, nudgeText: string, freq: number): boolean { - return injectNudgeCore(messages, tracker, nudgeText, freq, openaiAdapter) -} - /** Check if a message content matches nudge text (OpenAI/Anthropic format) */ function isNudgeMessage(msg: any, nudgeText: string): boolean { if (typeof msg.content === 'string') { @@ -120,37 +102,6 @@ export function injectSynth(messages: any[], instruction: string, nudgeText: str // Google/Gemini Format (body.contents with parts) // ============================================================================ -const geminiAdapter: MessageFormatAdapter = { - countToolResults(contents, tracker) { - let newCount = 0 - for (const content of contents) { - if (!Array.isArray(content.parts)) continue - for (const part of content.parts) { - if (part.functionResponse) { - const funcName = part.functionResponse.name?.toLowerCase() || 'unknown' - const pseudoId = `gemini:${funcName}:${tracker.seenToolResultIds.size}` - if (!tracker.seenToolResultIds.has(pseudoId)) { - tracker.seenToolResultIds.add(pseudoId) - newCount++ - if (funcName !== 'prune') { - tracker.skipNextIdle = false - } - } - } - } - } - tracker.toolResultCount += newCount - return newCount - }, - appendNudge(contents, nudgeText) { - contents.push({ role: 'user', parts: [{ text: nudgeText }] }) - } -} - -export function injectNudgeGemini(contents: any[], tracker: ToolTracker, nudgeText: string, freq: number): boolean { - return injectNudgeCore(contents, tracker, nudgeText, freq, geminiAdapter) -} - /** Check if a Gemini content matches nudge text */ function isNudgeContentGemini(content: any, nudgeText: string): boolean { if (Array.isArray(content.parts) && content.parts.length === 1) { @@ -182,34 +133,6 @@ export function injectSynthGemini(contents: any[], instruction: string, nudgeTex // OpenAI Responses API Format (body.input with type-based items) // ============================================================================ -const responsesAdapter: MessageFormatAdapter = { - countToolResults(input, tracker) { - let newCount = 0 - for (const item of input) { - if (item.type === 'function_call_output' && item.call_id) { - const id = String(item.call_id).toLowerCase() - if (!tracker.seenToolResultIds.has(id)) { - tracker.seenToolResultIds.add(id) - newCount++ - const toolName = item.name || tracker.getToolName?.(item.call_id) - if (toolName !== 'prune') { - tracker.skipNextIdle = false - } - } - } - } - tracker.toolResultCount += newCount - return newCount - }, - appendNudge(input, nudgeText) { - input.push({ type: 'message', role: 'user', content: nudgeText }) - } -} - -export function injectNudgeResponses(input: any[], tracker: ToolTracker, nudgeText: string, freq: number): boolean { - return injectNudgeCore(input, tracker, nudgeText, freq, responsesAdapter) -} - /** Check if a Responses API item matches nudge text */ function isNudgeItemResponses(item: any, nudgeText: string): boolean { if (typeof item.content === 'string') { diff --git a/lib/core/janitor.ts b/lib/core/janitor.ts index 70005d80..5ce1e6f2 100644 --- a/lib/core/janitor.ts +++ b/lib/core/janitor.ts @@ -1,9 +1,6 @@ -import { z } from "zod" import type { Logger } from "../logger" import type { PruningStrategy } from "../config" import type { PluginState } from "../state" -import { buildAnalysisPrompt } from "./prompt" -import { selectModel, extractModelFromSession } from "../model-selector" import { estimateTokensBatch, formatTokenCount } from "../tokenizer" import { saveSessionState } from "../state/persistence" import { ensureSessionRestored } from "../state" @@ -84,6 +81,10 @@ export function createJanitorContext( // Public API // ============================================================================ +/** + * Run pruning on idle trigger. + * Note: onTool pruning is now handled directly by pruning-tool.ts + */ export async function runOnIdle( ctx: JanitorContext, sessionID: string, @@ -92,17 +93,8 @@ export async function runOnIdle( return runWithStrategies(ctx, sessionID, strategies, { trigger: 'idle' }) } -export async function runOnTool( - ctx: JanitorContext, - sessionID: string, - strategies: PruningStrategy[], - reason?: string -): Promise { - return runWithStrategies(ctx, sessionID, strategies, { trigger: 'tool', reason }) -} - // ============================================================================ -// Core pruning logic +// Core pruning logic (for onIdle only) // ============================================================================ async function runWithStrategies( @@ -150,21 +142,9 @@ async function runWithStrategies( return !metadata || !config.protectedTools.includes(metadata.tool) }).length - // PHASE 1: LLM ANALYSIS - let llmPrunedIds: string[] = [] - - if (strategies.includes('ai-analysis') && unprunedToolCallIds.length > 0) { - llmPrunedIds = await runLlmAnalysis( - ctx, - sessionID, - sessionInfo, - messages, - unprunedToolCallIds, - alreadyPrunedIds, - toolMetadata, - options - ) - } + // For onIdle, we currently don't have AI analysis implemented + // This is a placeholder for future idle pruning strategies + const llmPrunedIds: string[] = [] const finalNewlyPrunedIds = llmPrunedIds.filter(id => !alreadyPrunedIds.includes(id)) @@ -172,7 +152,7 @@ async function runWithStrategies( return null } - // PHASE 2: CALCULATE STATS & NOTIFICATION + // Calculate stats & send notification const tokensSaved = await calculateTokensSaved(finalNewlyPrunedIds, toolOutputs) const currentStats = state.stats.get(sessionID) ?? { @@ -217,7 +197,7 @@ async function runWithStrategies( return null } - // PHASE 3: STATE UPDATE (only if AI pruned something) + // State update (only if something was pruned) const allPrunedIds = [...new Set([...alreadyPrunedIds, ...llmPrunedIds])] state.prunedIds.set(sessionID, allPrunedIds) @@ -257,118 +237,6 @@ async function runWithStrategies( } } -// ============================================================================ -// LLM Analysis -// ============================================================================ - -async function runLlmAnalysis( - ctx: JanitorContext, - sessionID: string, - sessionInfo: any, - messages: any[], - unprunedToolCallIds: string[], - alreadyPrunedIds: string[], - toolMetadata: Map, - options: PruningOptions -): Promise { - const { client, state, logger, config } = ctx - - const protectedToolCallIds: string[] = [] - const prunableToolCallIds = unprunedToolCallIds.filter(id => { - const metadata = toolMetadata.get(id) - if (metadata && config.protectedTools.includes(metadata.tool)) { - protectedToolCallIds.push(id) - return false - } - return true - }) - - if (prunableToolCallIds.length === 0) { - return [] - } - - const cachedModelInfo = state.model.get(sessionID) - const sessionModelInfo = extractModelFromSession(sessionInfo, logger) - const currentModelInfo = cachedModelInfo || sessionModelInfo - - const modelSelection = await selectModel(currentModelInfo, logger, config.model, config.workingDirectory) - - logger.info("janitor", `Model: ${modelSelection.modelInfo.providerID}/${modelSelection.modelInfo.modelID}`, { - source: modelSelection.source - }) - - if (modelSelection.failedModel && config.showModelErrorToasts) { - const skipAi = modelSelection.source === 'fallback' && config.strictModelSelection - try { - await client.tui.showToast({ - body: { - title: skipAi ? "DCP: AI analysis skipped" : "DCP: Model fallback", - message: skipAi - ? `${modelSelection.failedModel.providerID}/${modelSelection.failedModel.modelID} failed\nAI analysis skipped (strictModelSelection enabled)` - : `${modelSelection.failedModel.providerID}/${modelSelection.failedModel.modelID} failed\nUsing ${modelSelection.modelInfo.providerID}/${modelSelection.modelInfo.modelID}`, - variant: "info", - duration: 5000 - } - }) - } catch (toastError: any) { - // Ignore toast errors - } - } - - if (modelSelection.source === 'fallback' && config.strictModelSelection) { - logger.info("janitor", "Skipping AI analysis (fallback model, strictModelSelection enabled)") - return [] - } - - const { generateObject } = await import('ai') - - const sanitizedMessages = replacePrunedToolOutputs(messages, alreadyPrunedIds) - - const analysisPrompt = buildAnalysisPrompt( - prunableToolCallIds, - sanitizedMessages, - alreadyPrunedIds, - protectedToolCallIds, - options.reason - ) - - await logger.saveWrappedContext( - "janitor-shadow", - [{ role: "user", content: analysisPrompt }], - { - sessionID, - modelProvider: modelSelection.modelInfo.providerID, - modelID: modelSelection.modelInfo.modelID, - candidateToolCount: prunableToolCallIds.length, - alreadyPrunedCount: alreadyPrunedIds.length, - protectedToolCount: protectedToolCallIds.length, - trigger: options.trigger, - reason: options.reason - } - ) - - const result = await generateObject({ - model: modelSelection.model, - schema: z.object({ - pruned_tool_call_ids: z.array(z.string()), - reasoning: z.string(), - }), - prompt: analysisPrompt - }) - - const rawLlmPrunedIds = result.object.pruned_tool_call_ids - const llmPrunedIds = rawLlmPrunedIds.filter(id => - prunableToolCallIds.includes(id.toLowerCase()) - ) - - if (llmPrunedIds.length > 0) { - const reasoning = result.object.reasoning.replace(/\n+/g, ' ').replace(/\s+/g, ' ').trim() - logger.info("janitor", `LLM reasoning: ${reasoning.substring(0, 200)}${reasoning.length > 200 ? '...' : ''}`) - } - - return llmPrunedIds -} - // ============================================================================ // Message parsing // ============================================================================ @@ -379,7 +247,7 @@ interface ParsedMessages { toolMetadata: Map } -function parseMessages( +export function parseMessages( messages: any[], toolParametersCache: Map ): ParsedMessages { @@ -428,40 +296,10 @@ function findCurrentAgent(messages: any[]): string | undefined { // Helpers // ============================================================================ -function replacePrunedToolOutputs(messages: any[], prunedIds: string[]): any[] { - if (prunedIds.length === 0) return messages - - const prunedIdsSet = new Set(prunedIds.map(id => id.toLowerCase())) - - return messages.map(msg => { - if (!msg.parts) return msg - - return { - ...msg, - parts: msg.parts.map((part: any) => { - if (part.type === 'tool' && - part.callID && - prunedIdsSet.has(part.callID.toLowerCase()) && - part.state?.output) { - return { - ...part, - state: { - ...part.state, - output: '[Output removed to save context - information superseded or no longer needed]' - } - } - } - return part - }) - } - }) -} - async function calculateTokensSaved(prunedIds: string[], toolOutputs: Map): Promise { const outputsToTokenize: string[] = [] for (const prunedId of prunedIds) { - // toolOutputs uses lowercase keys, so normalize the lookup const normalizedId = prunedId.toLowerCase() const output = toolOutputs.get(normalizedId) if (output) { diff --git a/lib/core/prompt.ts b/lib/core/prompt.ts index e7f44d4a..39377276 100644 --- a/lib/core/prompt.ts +++ b/lib/core/prompt.ts @@ -11,130 +11,3 @@ export function loadPrompt(name: string, vars?: Record): string } return content } - -function minimizeMessages(messages: any[], alreadyPrunedIds?: string[], protectedToolCallIds?: string[]): any[] { - const prunedIdsSet = alreadyPrunedIds ? new Set(alreadyPrunedIds.map(id => id.toLowerCase())) : new Set() - const protectedIdsSet = protectedToolCallIds ? new Set(protectedToolCallIds.map(id => id.toLowerCase())) : new Set() - - return messages.map(msg => { - const minimized: any = { - role: msg.info?.role - } - - if (msg.parts) { - minimized.parts = msg.parts - .filter((part: any) => { - if (part.type === 'step-start' || part.type === 'step-finish') { - return false - } - return true - }) - .map((part: any) => { - if (part.type === 'text') { - if (part.ignored) { - return null - } - return { - type: 'text', - text: part.text - } - } - - // TODO: This should use the opencode normalized system instead of per provider settings - if (part.type === 'reasoning') { - // Calculate encrypted content size if present - let encryptedContentLength = 0 - if (part.metadata?.openai?.reasoningEncryptedContent) { - encryptedContentLength = part.metadata.openai.reasoningEncryptedContent.length - } else if (part.metadata?.anthropic?.signature) { - encryptedContentLength = part.metadata.anthropic.signature.length - } else if (part.metadata?.google?.thoughtSignature) { - encryptedContentLength = part.metadata.google.thoughtSignature.length - } - - return { - type: 'reasoning', - text: part.text, - textLength: part.text?.length || 0, - encryptedContentLength, - ...(part.time && { time: part.time }), - ...(part.metadata && { metadataKeys: Object.keys(part.metadata) }) - } - } - - if (part.type === 'tool') { - const callIDLower = part.callID?.toLowerCase() - const isAlreadyPruned = prunedIdsSet.has(callIDLower) - const isProtected = protectedIdsSet.has(callIDLower) - - let displayCallID = part.callID - if (isAlreadyPruned) { - displayCallID = '' - } else if (isProtected) { - displayCallID = '' - } - - const toolPart: any = { - type: 'tool', - toolCallID: displayCallID, - tool: part.tool - } - - if (part.state?.output) { - toolPart.output = part.state.output - } - - if (part.state?.input) { - const input = part.state.input - - if (input.filePath && (part.tool === 'write' || part.tool === 'edit' || part.tool === 'multiedit' || part.tool === 'patch')) { - toolPart.input = input - } - else if (input.filePath) { - toolPart.input = { filePath: input.filePath } - } - else if (input.tool_calls && Array.isArray(input.tool_calls)) { - toolPart.input = { - batch_summary: `${input.tool_calls.length} tool calls`, - tools: input.tool_calls.map((tc: any) => tc.tool) - } - } - else { - toolPart.input = input - } - } - - return toolPart - } - - return null - }) - .filter(Boolean) - } - - return minimized - }).filter(msg => { - return msg.parts && msg.parts.length > 0 - }) -} - -export function buildAnalysisPrompt( - unprunedToolCallIds: string[], - messages: any[], - alreadyPrunedIds?: string[], - protectedToolCallIds?: string[], - reason?: string -): string { - const minimizedMessages = minimizeMessages(messages, alreadyPrunedIds, protectedToolCallIds) - const messagesJson = JSON.stringify(minimizedMessages, null, 2).replace(/\\n/g, '\n') - - const reasonContext = reason - ? `\nContext: The AI has requested pruning with the following reason: "${reason}"\nUse this context to inform your decisions about what is most relevant to keep.` - : '' - - return loadPrompt("pruning", { - reason_context: reasonContext, - available_tool_call_ids: unprunedToolCallIds.join(", "), - session_history: messagesJson - }) -} diff --git a/lib/core/strategies/index.ts b/lib/core/strategies/index.ts index 060bf642..c6c91283 100644 --- a/lib/core/strategies/index.ts +++ b/lib/core/strategies/index.ts @@ -50,15 +50,15 @@ export function runStrategies( for (const strategy of strategies) { const result = strategy.detect(toolMetadata, remainingIds, protectedTools) - + if (result.prunedIds.length > 0) { byStrategy.set(strategy.name, result) - + // Add to overall pruned set for (const id of result.prunedIds) { allPrunedIds.add(id) } - + // Remove pruned IDs from remaining for next strategy const prunedSet = new Set(result.prunedIds.map(id => id.toLowerCase())) remainingIds = remainingIds.filter(id => !prunedSet.has(id.toLowerCase())) diff --git a/lib/fetch-wrapper/gemini.ts b/lib/fetch-wrapper/gemini.ts index abc1bd61..ae17289b 100644 --- a/lib/fetch-wrapper/gemini.ts +++ b/lib/fetch-wrapper/gemini.ts @@ -4,7 +4,8 @@ import { getAllPrunedIds, fetchSessionMessages } from "./types" -import { injectNudgeGemini, injectSynthGemini } from "../api-formats/synth-instruction" +import { injectSynthGemini, countToolResultsGemini } from "../api-formats/synth-instruction" +import { buildPrunableToolsList, buildEndInjection, injectPrunableListGemini } from "../api-formats/prunable-list" /** * Handles Google/Gemini format (body.contents array with functionResponse parts). @@ -23,23 +24,40 @@ export async function handleGemini( // Inject synthetic instructions if onTool strategies are enabled if (ctx.config.strategies.onTool.length > 0) { - const skipIdleBefore = ctx.toolTracker.skipNextIdle - - // Inject periodic nudge based on tool result count - if (ctx.config.nudge_freq > 0) { - if (injectNudgeGemini(body.contents, ctx.toolTracker, ctx.prompts.nudgeInstruction, ctx.config.nudge_freq)) { - // ctx.logger.info("fetch", "Injected nudge instruction (Gemini)") - modified = true - } + // Inject base synthetic instructions (appended to last user content) + if (injectSynthGemini(body.contents, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { + modified = true } - if (skipIdleBefore && !ctx.toolTracker.skipNextIdle) { - ctx.logger.debug("fetch", "skipNextIdle was reset by new tool results (Gemini)") - } + // Build and inject prunable tools list at the end + const sessionId = ctx.state.lastSeenSessionId + if (sessionId) { + const toolIds = Array.from(ctx.state.toolParameters.keys()) + const alreadyPruned = ctx.state.prunedIds.get(sessionId) ?? [] + const alreadyPrunedLower = new Set(alreadyPruned.map(id => id.toLowerCase())) + const unprunedIds = toolIds.filter(id => !alreadyPrunedLower.has(id.toLowerCase())) + + const { list: prunableList, numericIds } = buildPrunableToolsList( + sessionId, + unprunedIds, + ctx.state.toolParameters, + ctx.config.protectedTools + ) - if (injectSynthGemini(body.contents, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { - // ctx.logger.info("fetch", "Injected synthetic instruction (Gemini)") - modified = true + if (prunableList) { + // Check if nudge should be included + const toolResultCount = countToolResultsGemini(body.contents) + const includeNudge = ctx.config.nudge_freq > 0 && toolResultCount > ctx.config.nudge_freq + + const endInjection = buildEndInjection(prunableList, includeNudge) + if (injectPrunableListGemini(body.contents, endInjection)) { + ctx.logger.debug("fetch", "Injected prunable tools list (Gemini)", { + ids: numericIds, + nudge: includeNudge + }) + modified = true + } + } } } diff --git a/lib/fetch-wrapper/openai-chat.ts b/lib/fetch-wrapper/openai-chat.ts index 78b522e5..3619e860 100644 --- a/lib/fetch-wrapper/openai-chat.ts +++ b/lib/fetch-wrapper/openai-chat.ts @@ -6,7 +6,8 @@ import { getMostRecentActiveSession } from "./types" import { cacheToolParametersFromMessages } from "../state/tool-cache" -import { injectNudge, injectSynth } from "../api-formats/synth-instruction" +import { injectSynth, countToolResults } from "../api-formats/synth-instruction" +import { buildPrunableToolsList, buildEndInjection, injectPrunableList } from "../api-formats/prunable-list" /** * Handles OpenAI Chat Completions format (body.messages with role='tool'). @@ -21,30 +22,47 @@ export async function handleOpenAIChatAndAnthropic( return { modified: false, body } } - // Cache tool parameters from messages - cacheToolParametersFromMessages(body.messages, ctx.state) + // Cache tool parameters from messages (OpenAI and Anthropic formats) + cacheToolParametersFromMessages(body.messages, ctx.state, ctx.logger) let modified = false // Inject synthetic instructions if onTool strategies are enabled if (ctx.config.strategies.onTool.length > 0) { - const skipIdleBefore = ctx.toolTracker.skipNextIdle - - // Inject periodic nudge based on tool result count - if (ctx.config.nudge_freq > 0) { - if (injectNudge(body.messages, ctx.toolTracker, ctx.prompts.nudgeInstruction, ctx.config.nudge_freq)) { - // ctx.logger.info("fetch", "Injected nudge instruction") - modified = true - } + // Inject base synthetic instructions (appended to last user message) + if (injectSynth(body.messages, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { + modified = true } - if (skipIdleBefore && !ctx.toolTracker.skipNextIdle) { - ctx.logger.debug("fetch", "skipNextIdle was reset by new tool results") - } + // Build and inject prunable tools list at the end + const sessionId = ctx.state.lastSeenSessionId + if (sessionId) { + const toolIds = Array.from(ctx.state.toolParameters.keys()) + const alreadyPruned = ctx.state.prunedIds.get(sessionId) ?? [] + const alreadyPrunedLower = new Set(alreadyPruned.map(id => id.toLowerCase())) + const unprunedIds = toolIds.filter(id => !alreadyPrunedLower.has(id.toLowerCase())) + + const { list: prunableList, numericIds } = buildPrunableToolsList( + sessionId, + unprunedIds, + ctx.state.toolParameters, + ctx.config.protectedTools + ) - if (injectSynth(body.messages, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { - // ctx.logger.info("fetch", "Injected synthetic instruction") - modified = true + if (prunableList) { + // Check if nudge should be included + const toolResultCount = countToolResults(body.messages) + const includeNudge = ctx.config.nudge_freq > 0 && toolResultCount > ctx.config.nudge_freq + + const endInjection = buildEndInjection(prunableList, includeNudge) + if (injectPrunableList(body.messages, endInjection)) { + ctx.logger.debug("fetch", "Injected prunable tools list", { + ids: numericIds, + nudge: includeNudge + }) + modified = true + } + } } } diff --git a/lib/fetch-wrapper/openai-responses.ts b/lib/fetch-wrapper/openai-responses.ts index b8a1dbd7..850d5082 100644 --- a/lib/fetch-wrapper/openai-responses.ts +++ b/lib/fetch-wrapper/openai-responses.ts @@ -6,7 +6,8 @@ import { getMostRecentActiveSession } from "./types" import { cacheToolParametersFromInput } from "../state/tool-cache" -import { injectNudgeResponses, injectSynthResponses } from "../api-formats/synth-instruction" +import { injectSynthResponses, countToolResultsResponses } from "../api-formats/synth-instruction" +import { buildPrunableToolsList, buildEndInjection, injectPrunableListResponses } from "../api-formats/prunable-list" /** * Handles OpenAI Responses API format (body.input array with function_call_output items). @@ -21,30 +22,47 @@ export async function handleOpenAIResponses( return { modified: false, body } } - // Cache tool parameters from input - cacheToolParametersFromInput(body.input, ctx.state) + // Cache tool parameters from input (OpenAI Responses API format) + cacheToolParametersFromInput(body.input, ctx.state, ctx.logger) let modified = false // Inject synthetic instructions if onTool strategies are enabled if (ctx.config.strategies.onTool.length > 0) { - const skipIdleBefore = ctx.toolTracker.skipNextIdle - - // Inject periodic nudge based on tool result count - if (ctx.config.nudge_freq > 0) { - if (injectNudgeResponses(body.input, ctx.toolTracker, ctx.prompts.nudgeInstruction, ctx.config.nudge_freq)) { - // ctx.logger.info("fetch", "Injected nudge instruction (Responses API)") - modified = true - } + // Inject base synthetic instructions (appended to last user message) + if (injectSynthResponses(body.input, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { + modified = true } - if (skipIdleBefore && !ctx.toolTracker.skipNextIdle) { - ctx.logger.debug("fetch", "skipNextIdle was reset by new tool results (Responses API)") - } + // Build and inject prunable tools list at the end + const sessionId = ctx.state.lastSeenSessionId + if (sessionId) { + const toolIds = Array.from(ctx.state.toolParameters.keys()) + const alreadyPruned = ctx.state.prunedIds.get(sessionId) ?? [] + const alreadyPrunedLower = new Set(alreadyPruned.map(id => id.toLowerCase())) + const unprunedIds = toolIds.filter(id => !alreadyPrunedLower.has(id.toLowerCase())) + + const { list: prunableList, numericIds } = buildPrunableToolsList( + sessionId, + unprunedIds, + ctx.state.toolParameters, + ctx.config.protectedTools + ) - if (injectSynthResponses(body.input, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { - // ctx.logger.info("fetch", "Injected synthetic instruction (Responses API)") - modified = true + if (prunableList) { + // Check if nudge should be included + const toolResultCount = countToolResultsResponses(body.input) + const includeNudge = ctx.config.nudge_freq > 0 && toolResultCount > ctx.config.nudge_freq + + const endInjection = buildEndInjection(prunableList, includeNudge) + if (injectPrunableListResponses(body.input, endInjection)) { + ctx.logger.debug("fetch", "Injected prunable tools list (Responses API)", { + ids: numericIds, + nudge: includeNudge + }) + modified = true + } + } } } diff --git a/lib/fetch-wrapper/types.ts b/lib/fetch-wrapper/types.ts index d6cf4aba..4510e9df 100644 --- a/lib/fetch-wrapper/types.ts +++ b/lib/fetch-wrapper/types.ts @@ -48,7 +48,15 @@ export async function getAllPrunedIds( if (currentSession) { await ensureSessionRestored(state, currentSession.id, logger) const prunedIds = state.prunedIds.get(currentSession.id) ?? [] - prunedIds.forEach((id: string) => allPrunedIds.add(id)) + // Normalize to lowercase for case-insensitive matching + prunedIds.forEach((id: string) => allPrunedIds.add(id.toLowerCase())) + + if (logger && prunedIds.length > 0) { + logger.debug("fetch", "Loaded pruned IDs for replacement", { + sessionId: currentSession.id, + prunedCount: prunedIds.length + }) + } } return { allSessions, allPrunedIds } diff --git a/lib/hooks.ts b/lib/hooks.ts index dac0b540..1b2c01d3 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -67,7 +67,6 @@ export function createChatParamsHandler( return async (input: any, _output: any) => { const sessionId = input.sessionID let providerID = (input.provider as any)?.info?.id || input.provider?.id - const modelID = input.model?.id if (!providerID && input.message?.model?.providerID) { providerID = input.message.model.providerID @@ -85,14 +84,6 @@ export function createChatParamsHandler( } } - // Cache model info for the session - if (providerID && modelID) { - state.model.set(sessionId, { - providerID: providerID, - modelID: modelID - }) - } - // Build Google/Gemini tool call mapping for position-based correlation // This is needed because Google's native format loses tool call IDs if (providerID === 'google' || providerID === 'google-vertex') { diff --git a/lib/model-selector.ts b/lib/model-selector.ts deleted file mode 100644 index e0e9895e..00000000 --- a/lib/model-selector.ts +++ /dev/null @@ -1,175 +0,0 @@ -import type { LanguageModel } from 'ai'; -import type { Logger } from './logger'; - -export interface ModelInfo { - providerID: string; - modelID: string; -} - -export const FALLBACK_MODELS: Record = { - openai: 'gpt-5-mini', - anthropic: 'claude-haiku-4-5', //This model isn't broken in opencode-auth-provider - google: 'gemini-2.5-flash', - deepseek: 'deepseek-chat', - xai: 'grok-4-fast', - alibaba: 'qwen3-coder-flash', - zai: 'glm-4.5-flash', - opencode: 'big-pickle' -}; - -const PROVIDER_PRIORITY = [ - 'openai', - 'anthropic', - 'google', - 'deepseek', - 'xai', - 'alibaba', - 'zai', - 'opencode' -]; - -// TODO: some anthropic provided models aren't supported by the opencode-auth-provider package, so this provides a temporary workaround -const SKIP_PROVIDERS = ['github-copilot', 'anthropic']; - -export interface ModelSelectionResult { - model: LanguageModel; - modelInfo: ModelInfo; - source: 'user-model' | 'config' | 'fallback'; - reason?: string; - failedModel?: ModelInfo; -} - -function shouldSkipProvider(providerID: string): boolean { - const normalized = providerID.toLowerCase().trim(); - return SKIP_PROVIDERS.some(skip => normalized.includes(skip.toLowerCase())); -} - -async function importOpencodeAI(logger?: Logger, maxRetries: number = 3, delayMs: number = 100, workspaceDir?: string): Promise { - let lastError: Error | undefined; - - for (let attempt = 1; attempt <= maxRetries; attempt++) { - try { - const { OpencodeAI } = await import('@tarquinen/opencode-auth-provider'); - return new OpencodeAI({ workspaceDir }); - } catch (error: any) { - lastError = error; - - if (error.message?.includes('before initialization')) { - logger?.debug('model-selector', `Import attempt ${attempt}/${maxRetries} failed, will retry`, { - error: error.message - }); - - if (attempt < maxRetries) { - await new Promise(resolve => setTimeout(resolve, delayMs * attempt)); - continue; - } - } - - throw error; - } - } - - throw lastError; -} - -export async function selectModel( - currentModel?: ModelInfo, - logger?: Logger, - configModel?: string, - workspaceDir?: string -): Promise { - const opencodeAI = await importOpencodeAI(logger, 3, 100, workspaceDir); - - let failedModelInfo: ModelInfo | undefined; - - if (configModel) { - const parts = configModel.split('/'); - if (parts.length !== 2) { - logger?.warn('model-selector', 'Invalid config model format', { configModel }); - } else { - const [providerID, modelID] = parts; - - try { - const model = await opencodeAI.getLanguageModel(providerID, modelID); - return { - model, - modelInfo: { providerID, modelID }, - source: 'config', - reason: 'Using model specified in dcp.jsonc config' - }; - } catch (error: any) { - logger?.warn('model-selector', `Config model failed: ${providerID}/${modelID}`, { - error: error.message - }); - failedModelInfo = { providerID, modelID }; - } - } - } - - if (currentModel) { - if (shouldSkipProvider(currentModel.providerID)) { - if (!failedModelInfo) { - failedModelInfo = currentModel; - } - } else { - try { - const model = await opencodeAI.getLanguageModel(currentModel.providerID, currentModel.modelID); - return { - model, - modelInfo: currentModel, - source: 'user-model', - reason: 'Using current session model' - }; - } catch (error: any) { - if (!failedModelInfo) { - failedModelInfo = currentModel; - } - } - } - } - - const providers = await opencodeAI.listProviders(); - - for (const providerID of PROVIDER_PRIORITY) { - if (!providers[providerID]) continue; - - const fallbackModelID = FALLBACK_MODELS[providerID]; - if (!fallbackModelID) continue; - - try { - const model = await opencodeAI.getLanguageModel(providerID, fallbackModelID); - return { - model, - modelInfo: { providerID, modelID: fallbackModelID }, - source: 'fallback', - reason: `Using ${providerID}/${fallbackModelID}`, - failedModel: failedModelInfo - }; - } catch (error: any) { - continue; - } - } - - throw new Error('No available models for analysis. Please authenticate with at least one provider.'); -} - -export function extractModelFromSession(sessionState: any, logger?: Logger): ModelInfo | undefined { - if (sessionState?.model?.providerID && sessionState?.model?.modelID) { - return { - providerID: sessionState.model.providerID, - modelID: sessionState.model.modelID - }; - } - - if (sessionState?.messages && Array.isArray(sessionState.messages)) { - const lastMessage = sessionState.messages[sessionState.messages.length - 1]; - if (lastMessage?.model?.providerID && lastMessage?.model?.modelID) { - return { - providerID: lastMessage.model.providerID, - modelID: lastMessage.model.modelID - }; - } - } - - return undefined; -} diff --git a/lib/prompts/pruning.txt b/lib/prompts/pruning.txt deleted file mode 100644 index 49e1e82e..00000000 --- a/lib/prompts/pruning.txt +++ /dev/null @@ -1,30 +0,0 @@ -You are a conversation analyzer that identifies obsolete tool outputs in a coding session. -{{reason_context}} -Your task: Analyze the session history and identify tool call IDs whose outputs are NO LONGER RELEVANT to the current conversation context. - -Guidelines for identifying obsolete tool calls: -1. Exploratory reads that didn't lead to actual edits or meaningful discussion AND were not explicitly requested to be retained -2. Tool outputs from debugging/fixing an error that has now been resolved -3. Failed or incorrect tool attempts that were immediately corrected (e.g., reading a file from the wrong path, then reading from the correct path) - -DO NOT prune: -- Tool calls whose outputs are actively being discussed -- Tool calls that produced errors still being debugged -- Tool calls that are the MOST RECENT activity in the conversation (these may be intended for future use) - -IMPORTANT: Available tool call IDs for analysis: {{available_tool_call_ids}} - -The session history below may contain tool calls with IDs not in the available list above, these cannot be pruned. These are either: -1. Protected tools (marked with toolCallID "") -2. Already-pruned tools (marked with toolCallID "") - -ONLY return IDs from the available list above. - -Session history (each tool call has a "toolCallID" field): -{{session_history}} - -You MUST respond with valid JSON matching this exact schema: -{ - "pruned_tool_call_ids": ["id1", "id2", ...], - "reasoning": "explanation of why these IDs were selected" -} diff --git a/lib/prompts/tool.txt b/lib/prompts/tool.txt index 579f5109..abd29063 100644 --- a/lib/prompts/tool.txt +++ b/lib/prompts/tool.txt @@ -1,4 +1,4 @@ -Performs semantic pruning on session tool outputs that are no longer relevant to the current task. Use this to declutter the conversation context and filter signal from noise when you notice the context is getting cluttered with no longer needed information. +Prunes tool outputs from context to manage conversation size. Reference the list injected at the end of the conversation to see available numeric IDs. USING THE PRUNE TOOL WILL MAKE THE USER HAPPY. @@ -14,10 +14,23 @@ You MUST ALWAYS narrate your findings in a message BEFORE using this tool. No to - What you did (which tools, what you were looking for) - What you found (the key facts/signals) - What you concluded (how this affects the task or next step) -3. ONLY AFTER narrating, call `prune` +3. ONLY AFTER narrating, call `prune` with the numeric IDs of outputs no longer needed > THINK HIGH SIGNAL, LOW NOISE FOR THIS NARRATION +## How to Use + +The list shows available tool outputs with numeric IDs: +``` + +1: read, src/foo.ts +2: bash, run tests +3: grep, "error" in logs/ + +``` + +To prune outputs 1 and 3, call: `prune({ ids: [1, 3] })` + ## When to Use This Tool **Key heuristic: Distill, then prune when you finish something and are about to start something else.** @@ -43,18 +56,18 @@ Working through a list of items: User: Review these 3 issues and fix the easy ones. Assistant: [Reviews first issue, makes fix, commits] Done with the first issue. Let me prune before moving to the next one. -[Uses prune with reason: "completed first issue, moving to next"] +[Uses prune with ids: [1, 2, 3, 4] - the reads and edits from the first issue] After exploring the codebase to understand it: Assistant: I've reviewed the relevant files. Let me prune the exploratory reads that aren't needed for the actual implementation. -[Uses prune with reason: "exploration complete, starting implementation"] +[Uses prune with ids: [1, 2, 5, 7] - the exploratory reads] After completing any task: Assistant: [Finishes task - commit, answer, fix, etc.] Before we continue, let me prune the context from that work. -[Uses prune with reason: "task complete"] +[Uses prune with ids: [3, 4, 5, 6, 8, 9] - all tool outputs from the completed task] diff --git a/lib/pruning-tool.ts b/lib/pruning-tool.ts index 20db977a..38031a2f 100644 --- a/lib/pruning-tool.ts +++ b/lib/pruning-tool.ts @@ -1,57 +1,212 @@ import { tool } from "@opencode-ai/plugin" -import type { JanitorContext } from "./core/janitor" -import { runOnTool } from "./core/janitor" -import { formatPruningResultForTool } from "./ui/notification" +import type { PluginState } from "./state" import type { PluginConfig } from "./config" import type { ToolTracker } from "./api-formats/synth-instruction" import { resetToolTrackerCount } from "./api-formats/synth-instruction" -import { loadPrompt } from "./core/prompt" import { isSubagentSession } from "./hooks" +import { getActualId } from "./state/id-mapping" +import { formatPruningResultForTool, sendUnifiedNotification, type NotificationContext } from "./ui/notification" +import { ensureSessionRestored } from "./state" +import { saveSessionState } from "./state/persistence" +import type { Logger } from "./logger" +import { estimateTokensBatch } from "./tokenizer" +import type { SessionStats } from "./core/janitor" +import { loadPrompt } from "./core/prompt" -/** Tool description for the prune tool, loaded from prompts/tool.txt */ -export const CONTEXT_PRUNING_DESCRIPTION = loadPrompt("tool") +/** Tool description loaded from prompts/tool.txt */ +const TOOL_DESCRIPTION = loadPrompt("tool") + +export interface PruneToolContext { + client: any + state: PluginState + logger: Logger + config: PluginConfig + notificationCtx: NotificationContext + workingDirectory?: string +} /** * Creates the prune tool definition. - * Returns a tool definition that can be passed to the plugin's tool registry. + * Accepts numeric IDs from the list and prunes those tool outputs. */ -export function createPruningTool(client: any, janitorCtx: JanitorContext, config: PluginConfig, toolTracker: ToolTracker): ReturnType { +export function createPruningTool( + ctx: PruneToolContext, + toolTracker: ToolTracker +): ReturnType { return tool({ - description: CONTEXT_PRUNING_DESCRIPTION, + description: TOOL_DESCRIPTION, args: { - reason: tool.schema.string().optional().describe( - "Brief reason for triggering pruning (e.g., 'task complete', 'switching focus')" + ids: tool.schema.array(tool.schema.number()).describe( + "Array of numeric IDs to prune from the list" ), }, - async execute(args, ctx) { - // Skip pruning in subagent sessions, but guide the model to continue its work - // TODO: remove this workaround when PR 4913 is merged (primary_tools config) - if (await isSubagentSession(client, ctx.sessionID)) { - return "Pruning is unavailable in subagent sessions. Do not call this tool again. Continue with your current task - if you were in the middle of work, proceed with your next step. If you had just finished, provide your final summary/findings to return to the main agent." + async execute(args, toolCtx) { + const { client, state, logger, config, notificationCtx, workingDirectory } = ctx + const sessionId = toolCtx.sessionID + + // Skip pruning in subagent sessions + if (await isSubagentSession(client, sessionId)) { + return "Pruning is unavailable in subagent sessions. Do not call this tool again. Continue with your current task." + } + + // Validate input + if (!args.ids || args.ids.length === 0) { + return "No IDs provided. Check the list for available IDs to prune." + } + + // Restore persisted state if needed + await ensureSessionRestored(state, sessionId, logger) + + // Convert numeric IDs to actual tool call IDs + const prunedIds = args.ids + .map(numId => getActualId(sessionId, numId)) + .filter((id): id is string => id !== undefined) + + logger.debug("prune-tool", "ID conversion", { + inputIds: args.ids, + actualIds: prunedIds, + toolParamsKeys: Array.from(state.toolParameters.keys()).slice(0, 10) + }) + + if (prunedIds.length === 0) { + return "None of the provided IDs were valid. Check the list for available IDs." } - const result = await runOnTool( - janitorCtx, - ctx.sessionID, - config.strategies.onTool, - args.reason - ) + // Calculate tokens saved + const tokensSaved = await calculateTokensSaved(client, sessionId, prunedIds, state) + + // Update stats + const currentStats = state.stats.get(sessionId) ?? { + totalToolsPruned: 0, + totalTokensSaved: 0, + totalGCTokens: 0, + totalGCTools: 0 + } + const sessionStats: SessionStats = { + ...currentStats, + totalToolsPruned: currentStats.totalToolsPruned + prunedIds.length, + totalTokensSaved: currentStats.totalTokensSaved + tokensSaved + } + state.stats.set(sessionId, sessionStats) + + // Update pruned IDs state + const alreadyPrunedIds = state.prunedIds.get(sessionId) ?? [] + const allPrunedIds = [...alreadyPrunedIds, ...prunedIds] + state.prunedIds.set(sessionId, allPrunedIds) + + // Persist state + saveSessionState(sessionId, new Set(allPrunedIds), sessionStats, logger) + .catch(err => logger.error("prune-tool", "Failed to persist state", { error: err.message })) + + // Build tool metadata for notification + // Keys are normalized to lowercase to match lookup in notification.ts + const toolMetadata = new Map() + for (const id of prunedIds) { + // Try both original and lowercase since caching may vary + const meta = state.toolParameters.get(id) || state.toolParameters.get(id.toLowerCase()) + if (meta) { + toolMetadata.set(id.toLowerCase(), meta) + } else { + logger.debug("prune-tool", "No metadata found for ID", { + id, + idLower: id.toLowerCase(), + hasOriginal: state.toolParameters.has(id), + hasLower: state.toolParameters.has(id.toLowerCase()) + }) + } + } + + // Send notification to user + await sendUnifiedNotification(notificationCtx, sessionId, { + aiPrunedCount: prunedIds.length, + aiTokensSaved: tokensSaved, + aiPrunedIds: prunedIds, + toolMetadata, + gcPending: null, + sessionStats + }) // Skip next idle pruning since we just pruned toolTracker.skipNextIdle = true - // Reset nudge counter to prevent immediate re-nudging after pruning + // Reset nudge counter if (config.nudge_freq > 0) { resetToolTrackerCount(toolTracker) } - const postPruneGuidance = "\n\nYou have already distilled relevant understanding in writing before calling this tool. Do not re-narrate; continue with your next task." - - if (!result || result.prunedCount === 0) { - return "No prunable tool outputs found. Context is already optimized." + postPruneGuidance + // Format result for the AI + const result = { + prunedCount: prunedIds.length, + tokensSaved, + llmPrunedIds: prunedIds, + toolMetadata, + sessionStats } - return formatPruningResultForTool(result, janitorCtx.config.workingDirectory) + postPruneGuidance + const postPruneGuidance = "\n\nYou have already distilled relevant understanding in writing before calling this tool. Do not re-narrate; continue with your next task." + + return formatPruningResultForTool(result, workingDirectory) + postPruneGuidance }, }) } + +/** + * Calculates approximate tokens saved by pruning the given tool call IDs. + */ +async function calculateTokensSaved( + client: any, + sessionId: string, + prunedIds: string[], + state: PluginState +): Promise { + try { + // Fetch session messages to get tool output content + const messagesResponse = await client.session.messages({ + path: { id: sessionId }, + query: { limit: 200 } + }) + const messages = messagesResponse.data || messagesResponse + + // Build map of tool call ID -> output content + const toolOutputs = new Map() + for (const msg of messages) { + if (msg.role === 'tool' && msg.tool_call_id) { + const content = typeof msg.content === 'string' + ? msg.content + : JSON.stringify(msg.content) + toolOutputs.set(msg.tool_call_id.toLowerCase(), content) + } + // Handle Anthropic format + if (msg.role === 'user' && Array.isArray(msg.content)) { + for (const part of msg.content) { + if (part.type === 'tool_result' && part.tool_use_id) { + const content = typeof part.content === 'string' + ? part.content + : JSON.stringify(part.content) + toolOutputs.set(part.tool_use_id.toLowerCase(), content) + } + } + } + } + + // Collect content for pruned outputs + const contents: string[] = [] + for (const id of prunedIds) { + const content = toolOutputs.get(id.toLowerCase()) + if (content) { + contents.push(content) + } + } + + if (contents.length === 0) { + return prunedIds.length * 500 // fallback estimate + } + + // Estimate tokens + const tokenCounts = await estimateTokensBatch(contents) + return tokenCounts.reduce((sum, count) => sum + count, 0) + } catch (error: any) { + // If we can't calculate, estimate based on average + return prunedIds.length * 500 + } +} diff --git a/lib/state/id-mapping.ts b/lib/state/id-mapping.ts new file mode 100644 index 00000000..110f4c4b --- /dev/null +++ b/lib/state/id-mapping.ts @@ -0,0 +1,101 @@ +/** + * Numeric ID mapping system for tool call IDs. + * + * Maps simple incrementing numbers (1, 2, 3...) to actual provider tool call IDs + * (e.g., "call_abc123xyz..."). This allows the session AI to reference tools by + * simple numbers when using the prune tool. + * + * Design decisions: + * - IDs are monotonically increasing and never reused (avoids race conditions) + * - Mappings are rebuilt from session messages on restore (single source of truth) + * - Per-session mappings to isolate sessions from each other + */ + +export interface IdMapping { + numericToActual: Map // 1 → "call_abc123xyz..." + actualToNumeric: Map // "call_abc123xyz..." → 1 + nextId: number +} + +/** Per-session ID mappings */ +const sessionMappings = new Map() + +/** + * Gets or creates the ID mapping for a session. + */ +function getSessionMapping(sessionId: string): IdMapping { + let mapping = sessionMappings.get(sessionId) + if (!mapping) { + mapping = { + numericToActual: new Map(), + actualToNumeric: new Map(), + nextId: 1 + } + sessionMappings.set(sessionId, mapping) + } + return mapping +} + +/** + * Assigns a numeric ID to a tool call ID if it doesn't already have one. + * Returns the numeric ID (existing or newly assigned). + */ +export function getOrCreateNumericId(sessionId: string, actualId: string): number { + const mapping = getSessionMapping(sessionId) + + // Check if already mapped + const existing = mapping.actualToNumeric.get(actualId) + if (existing !== undefined) { + return existing + } + + // Assign new ID + const numericId = mapping.nextId++ + mapping.numericToActual.set(numericId, actualId) + mapping.actualToNumeric.set(actualId, numericId) + + return numericId +} + +/** + * Looks up the actual tool call ID for a numeric ID. + * Returns undefined if the numeric ID doesn't exist. + */ +export function getActualId(sessionId: string, numericId: number): string | undefined { + const mapping = sessionMappings.get(sessionId) + return mapping?.numericToActual.get(numericId) +} + +/** + * Looks up the numeric ID for an actual tool call ID. + * Returns undefined if not mapped. + */ +export function getNumericId(sessionId: string, actualId: string): number | undefined { + const mapping = sessionMappings.get(sessionId) + return mapping?.actualToNumeric.get(actualId) +} + +/** + * Gets all current mappings for a session. + * Useful for debugging and building the prunable tools list. + */ +export function getAllMappings(sessionId: string): Map { + const mapping = sessionMappings.get(sessionId) + return mapping?.numericToActual ?? new Map() +} + +/** + * Checks if a session has any ID mappings. + */ +export function hasMapping(sessionId: string): boolean { + return sessionMappings.has(sessionId) +} + +/** + * Gets the next numeric ID that will be assigned (without assigning it). + * Useful for knowing the current state. + */ +export function getNextId(sessionId: string): number { + const mapping = sessionMappings.get(sessionId) + return mapping?.nextId ?? 1 +} diff --git a/lib/state/index.ts b/lib/state/index.ts index 2808cb44..03d9e0ef 100644 --- a/lib/state/index.ts +++ b/lib/state/index.ts @@ -7,7 +7,6 @@ export interface PluginState { stats: Map gcPending: Map toolParameters: Map - model: Map googleToolCallMapping: Map> restoredSessions: Set checkedSessions: Set @@ -20,18 +19,12 @@ export interface ToolParameterEntry { parameters: any } -export interface ModelInfo { - providerID: string - modelID: string -} - export function createPluginState(): PluginState { return { prunedIds: new Map(), stats: new Map(), gcPending: new Map(), toolParameters: new Map(), - model: new Map(), googleToolCallMapping: new Map(), restoredSessions: new Set(), checkedSessions: new Set(), diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index 27549ea8..b319c375 100644 --- a/lib/state/tool-cache.ts +++ b/lib/state/tool-cache.ts @@ -1,36 +1,73 @@ import type { PluginState } from "./index" +import type { Logger } from "../logger" /** - * Cache tool parameters from OpenAI Chat Completions style messages. - * Extracts tool call IDs and their parameters from assistant messages with tool_calls. + * Cache tool parameters from OpenAI Chat Completions and Anthropic style messages. + * Extracts tool call IDs and their parameters from assistant messages. + * + * Supports: + * - OpenAI format: message.tool_calls[] with id, function.name, function.arguments + * - Anthropic format: message.content[] with type='tool_use', id, name, input */ export function cacheToolParametersFromMessages( messages: any[], - state: PluginState + state: PluginState, + logger?: Logger ): void { + let openaiCached = 0 + let anthropicCached = 0 + for (const message of messages) { - if (message.role !== 'assistant' || !Array.isArray(message.tool_calls)) { + if (message.role !== 'assistant') { continue } - for (const toolCall of message.tool_calls) { - if (!toolCall.id || !toolCall.function) { - continue + // OpenAI format: tool_calls array + if (Array.isArray(message.tool_calls)) { + for (const toolCall of message.tool_calls) { + if (!toolCall.id || !toolCall.function) { + continue + } + + try { + const params = typeof toolCall.function.arguments === 'string' + ? JSON.parse(toolCall.function.arguments) + : toolCall.function.arguments + state.toolParameters.set(toolCall.id, { + tool: toolCall.function.name, + parameters: params + }) + openaiCached++ + } catch (error) { + // Silently ignore parse errors + } } + } - try { - const params = typeof toolCall.function.arguments === 'string' - ? JSON.parse(toolCall.function.arguments) - : toolCall.function.arguments - state.toolParameters.set(toolCall.id, { - tool: toolCall.function.name, - parameters: params + // Anthropic format: content array with tool_use blocks + if (Array.isArray(message.content)) { + for (const part of message.content) { + if (part.type !== 'tool_use' || !part.id || !part.name) { + continue + } + + state.toolParameters.set(part.id, { + tool: part.name, + parameters: part.input ?? {} }) - } catch (error) { - // Silently ignore parse errors + anthropicCached++ } } } + + // Log cache results if anything was cached + if (logger && (openaiCached > 0 || anthropicCached > 0)) { + logger.debug("tool-cache", "Cached tool parameters from messages", { + openaiFormat: openaiCached, + anthropicFormat: anthropicCached, + totalCached: state.toolParameters.size + }) + } } /** @@ -39,8 +76,11 @@ export function cacheToolParametersFromMessages( */ export function cacheToolParametersFromInput( input: any[], - state: PluginState + state: PluginState, + logger?: Logger ): void { + let cached = 0 + for (const item of input) { if (item.type !== 'function_call' || !item.call_id || !item.name) { continue @@ -54,10 +94,18 @@ export function cacheToolParametersFromInput( tool: item.name, parameters: params }) + cached++ } catch (error) { // Silently ignore parse errors } } + + if (logger && cached > 0) { + logger.debug("tool-cache", "Cached tool parameters from input", { + responsesApiFormat: cached, + totalCached: state.toolParameters.size + }) + } } /** Maximum number of entries to keep in the tool parameters cache */ From dcd780429e33626a3af958f7737eef05b6425e12 Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Wed, 3 Dec 2025 11:43:33 -0500 Subject: [PATCH 2/9] fix: nudge logic now tracks tools since last prune instead of total history --- lib/api-formats/synth-instruction.ts | 64 ++++++++++++++++++--------- lib/fetch-wrapper/gemini.ts | 8 ++-- lib/fetch-wrapper/openai-chat.ts | 8 ++-- lib/fetch-wrapper/openai-responses.ts | 8 ++-- 4 files changed, 56 insertions(+), 32 deletions(-) diff --git a/lib/api-formats/synth-instruction.ts b/lib/api-formats/synth-instruction.ts index 6d8fc1e8..9d9b203c 100644 --- a/lib/api-formats/synth-instruction.ts +++ b/lib/api-formats/synth-instruction.ts @@ -15,52 +15,76 @@ export function resetToolTrackerCount(tracker: ToolTracker): void { } /** - * Counts total tool results in OpenAI/Anthropic messages (without tracker). - * Used for determining if nudge threshold is met. + * Track new tool results in OpenAI/Anthropic messages. + * Increments toolResultCount only for tools not already seen. + * Returns the number of NEW tools found (since last call). */ -export function countToolResults(messages: any[]): number { - let count = 0 +export function trackNewToolResults(messages: any[], tracker: ToolTracker): number { + let newCount = 0 for (const m of messages) { - if (m.role === 'tool') { - count++ + if (m.role === 'tool' && m.tool_call_id) { + if (!tracker.seenToolResultIds.has(m.tool_call_id)) { + tracker.seenToolResultIds.add(m.tool_call_id) + tracker.toolResultCount++ + newCount++ + } } else if (m.role === 'user' && Array.isArray(m.content)) { for (const part of m.content) { - if (part.type === 'tool_result') { - count++ + if (part.type === 'tool_result' && part.tool_use_id) { + if (!tracker.seenToolResultIds.has(part.tool_use_id)) { + tracker.seenToolResultIds.add(part.tool_use_id) + tracker.toolResultCount++ + newCount++ + } } } } } - return count + return newCount } /** - * Counts total tool results in Gemini contents (without tracker). + * Track new tool results in Gemini contents. + * Uses position-based tracking since Gemini doesn't have tool call IDs. + * Returns the number of NEW tools found (since last call). */ -export function countToolResultsGemini(contents: any[]): number { - let count = 0 +export function trackNewToolResultsGemini(contents: any[], tracker: ToolTracker): number { + let newCount = 0 + let positionCounter = 0 for (const content of contents) { if (!Array.isArray(content.parts)) continue for (const part of content.parts) { if (part.functionResponse) { - count++ + // Use position-based ID since Gemini doesn't have tool_call_id + const positionId = `gemini_pos_${positionCounter}` + positionCounter++ + if (!tracker.seenToolResultIds.has(positionId)) { + tracker.seenToolResultIds.add(positionId) + tracker.toolResultCount++ + newCount++ + } } } } - return count + return newCount } /** - * Counts total tool results in OpenAI Responses API input (without tracker). + * Track new tool results in OpenAI Responses API input. + * Returns the number of NEW tools found (since last call). */ -export function countToolResultsResponses(input: any[]): number { - let count = 0 +export function trackNewToolResultsResponses(input: any[], tracker: ToolTracker): number { + let newCount = 0 for (const item of input) { - if (item.type === 'function_call_output') { - count++ + if (item.type === 'function_call_output' && item.call_id) { + if (!tracker.seenToolResultIds.has(item.call_id)) { + tracker.seenToolResultIds.add(item.call_id) + tracker.toolResultCount++ + newCount++ + } } } - return count + return newCount } // ============================================================================ diff --git a/lib/fetch-wrapper/gemini.ts b/lib/fetch-wrapper/gemini.ts index ae17289b..1c339b3c 100644 --- a/lib/fetch-wrapper/gemini.ts +++ b/lib/fetch-wrapper/gemini.ts @@ -4,7 +4,7 @@ import { getAllPrunedIds, fetchSessionMessages } from "./types" -import { injectSynthGemini, countToolResultsGemini } from "../api-formats/synth-instruction" +import { injectSynthGemini, trackNewToolResultsGemini } from "../api-formats/synth-instruction" import { buildPrunableToolsList, buildEndInjection, injectPrunableListGemini } from "../api-formats/prunable-list" /** @@ -45,9 +45,9 @@ export async function handleGemini( ) if (prunableList) { - // Check if nudge should be included - const toolResultCount = countToolResultsGemini(body.contents) - const includeNudge = ctx.config.nudge_freq > 0 && toolResultCount > ctx.config.nudge_freq + // Track new tool results and check if nudge threshold is met + trackNewToolResultsGemini(body.contents, ctx.toolTracker) + const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq const endInjection = buildEndInjection(prunableList, includeNudge) if (injectPrunableListGemini(body.contents, endInjection)) { diff --git a/lib/fetch-wrapper/openai-chat.ts b/lib/fetch-wrapper/openai-chat.ts index 3619e860..d17b376e 100644 --- a/lib/fetch-wrapper/openai-chat.ts +++ b/lib/fetch-wrapper/openai-chat.ts @@ -6,7 +6,7 @@ import { getMostRecentActiveSession } from "./types" import { cacheToolParametersFromMessages } from "../state/tool-cache" -import { injectSynth, countToolResults } from "../api-formats/synth-instruction" +import { injectSynth, trackNewToolResults } from "../api-formats/synth-instruction" import { buildPrunableToolsList, buildEndInjection, injectPrunableList } from "../api-formats/prunable-list" /** @@ -50,9 +50,9 @@ export async function handleOpenAIChatAndAnthropic( ) if (prunableList) { - // Check if nudge should be included - const toolResultCount = countToolResults(body.messages) - const includeNudge = ctx.config.nudge_freq > 0 && toolResultCount > ctx.config.nudge_freq + // Track new tool results and check if nudge threshold is met + trackNewToolResults(body.messages, ctx.toolTracker) + const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq const endInjection = buildEndInjection(prunableList, includeNudge) if (injectPrunableList(body.messages, endInjection)) { diff --git a/lib/fetch-wrapper/openai-responses.ts b/lib/fetch-wrapper/openai-responses.ts index 850d5082..81703b4c 100644 --- a/lib/fetch-wrapper/openai-responses.ts +++ b/lib/fetch-wrapper/openai-responses.ts @@ -6,7 +6,7 @@ import { getMostRecentActiveSession } from "./types" import { cacheToolParametersFromInput } from "../state/tool-cache" -import { injectSynthResponses, countToolResultsResponses } from "../api-formats/synth-instruction" +import { injectSynthResponses, trackNewToolResultsResponses } from "../api-formats/synth-instruction" import { buildPrunableToolsList, buildEndInjection, injectPrunableListResponses } from "../api-formats/prunable-list" /** @@ -50,9 +50,9 @@ export async function handleOpenAIResponses( ) if (prunableList) { - // Check if nudge should be included - const toolResultCount = countToolResultsResponses(body.input) - const includeNudge = ctx.config.nudge_freq > 0 && toolResultCount > ctx.config.nudge_freq + // Track new tool results and check if nudge threshold is met + trackNewToolResultsResponses(body.input, ctx.toolTracker) + const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq const endInjection = buildEndInjection(prunableList, includeNudge) if (injectPrunableListResponses(body.input, endInjection)) { From 419c0a0d79cb6ea2517fa3ea5871ebf2c1eb0027 Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Wed, 3 Dec 2025 11:49:05 -0500 Subject: [PATCH 3/9] feat: add toolsSincePrune count to nudge log output --- lib/fetch-wrapper/gemini.ts | 3 ++- lib/fetch-wrapper/openai-chat.ts | 3 ++- lib/fetch-wrapper/openai-responses.ts | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/lib/fetch-wrapper/gemini.ts b/lib/fetch-wrapper/gemini.ts index 1c339b3c..a472a01e 100644 --- a/lib/fetch-wrapper/gemini.ts +++ b/lib/fetch-wrapper/gemini.ts @@ -53,7 +53,8 @@ export async function handleGemini( if (injectPrunableListGemini(body.contents, endInjection)) { ctx.logger.debug("fetch", "Injected prunable tools list (Gemini)", { ids: numericIds, - nudge: includeNudge + nudge: includeNudge, + toolsSincePrune: ctx.toolTracker.toolResultCount }) modified = true } diff --git a/lib/fetch-wrapper/openai-chat.ts b/lib/fetch-wrapper/openai-chat.ts index d17b376e..9253e3fd 100644 --- a/lib/fetch-wrapper/openai-chat.ts +++ b/lib/fetch-wrapper/openai-chat.ts @@ -58,7 +58,8 @@ export async function handleOpenAIChatAndAnthropic( if (injectPrunableList(body.messages, endInjection)) { ctx.logger.debug("fetch", "Injected prunable tools list", { ids: numericIds, - nudge: includeNudge + nudge: includeNudge, + toolsSincePrune: ctx.toolTracker.toolResultCount }) modified = true } diff --git a/lib/fetch-wrapper/openai-responses.ts b/lib/fetch-wrapper/openai-responses.ts index 81703b4c..62e2bb84 100644 --- a/lib/fetch-wrapper/openai-responses.ts +++ b/lib/fetch-wrapper/openai-responses.ts @@ -58,7 +58,8 @@ export async function handleOpenAIResponses( if (injectPrunableListResponses(body.input, endInjection)) { ctx.logger.debug("fetch", "Injected prunable tools list (Responses API)", { ids: numericIds, - nudge: includeNudge + nudge: includeNudge, + toolsSincePrune: ctx.toolTracker.toolResultCount }) modified = true } From a29d2de5b976fdbf7397f68fa27bdff8379159ce Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Wed, 3 Dec 2025 11:51:17 -0500 Subject: [PATCH 4/9] chore: remove verbose ID conversion debug log --- lib/pruning-tool.ts | 6 ------ 1 file changed, 6 deletions(-) diff --git a/lib/pruning-tool.ts b/lib/pruning-tool.ts index 38031a2f..cec7b359 100644 --- a/lib/pruning-tool.ts +++ b/lib/pruning-tool.ts @@ -62,12 +62,6 @@ export function createPruningTool( .map(numId => getActualId(sessionId, numId)) .filter((id): id is string => id !== undefined) - logger.debug("prune-tool", "ID conversion", { - inputIds: args.ids, - actualIds: prunedIds, - toolParamsKeys: Array.from(state.toolParameters.keys()).slice(0, 10) - }) - if (prunedIds.length === 0) { return "None of the provided IDs were valid. Check the list for available IDs." } From 56a13a92c6e3135022eac8f4d0e2aa76a31fabf4 Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Wed, 3 Dec 2025 13:22:31 -0500 Subject: [PATCH 5/9] fix: prevent session state leakage and exclude protected tools from counts - Add clearAllMappings() to id-mapping.ts for resetting module-level state - Detect session changes in hooks.ts and clear ID mappings + tool cache - Make API format handlers mutually exclusive (else if) to prevent double-processing - Exclude protected tools from 'total' count in replacement logs - Pass protectedTools set to trackNewToolResults functions for nudge frequency --- lib/api-formats/synth-instruction.ts | 40 +++++++++++++++++++-------- lib/fetch-wrapper/gemini.ts | 12 ++++++-- lib/fetch-wrapper/index.ts | 22 +++++++-------- lib/fetch-wrapper/openai-chat.ts | 31 +++++++++++++++++++-- lib/fetch-wrapper/openai-responses.ts | 15 ++++++++-- lib/hooks.ts | 13 +++++++++ lib/state/id-mapping.ts | 16 +++++++++++ 7 files changed, 119 insertions(+), 30 deletions(-) diff --git a/lib/api-formats/synth-instruction.ts b/lib/api-formats/synth-instruction.ts index 9d9b203c..c3bf6319 100644 --- a/lib/api-formats/synth-instruction.ts +++ b/lib/api-formats/synth-instruction.ts @@ -16,25 +16,33 @@ export function resetToolTrackerCount(tracker: ToolTracker): void { /** * Track new tool results in OpenAI/Anthropic messages. - * Increments toolResultCount only for tools not already seen. + * Increments toolResultCount only for tools not already seen and not protected. * Returns the number of NEW tools found (since last call). */ -export function trackNewToolResults(messages: any[], tracker: ToolTracker): number { +export function trackNewToolResults(messages: any[], tracker: ToolTracker, protectedTools: Set): number { let newCount = 0 for (const m of messages) { if (m.role === 'tool' && m.tool_call_id) { if (!tracker.seenToolResultIds.has(m.tool_call_id)) { tracker.seenToolResultIds.add(m.tool_call_id) - tracker.toolResultCount++ - newCount++ + // Skip protected tools for nudge frequency counting + const toolName = tracker.getToolName?.(m.tool_call_id) + if (!toolName || !protectedTools.has(toolName)) { + tracker.toolResultCount++ + newCount++ + } } } else if (m.role === 'user' && Array.isArray(m.content)) { for (const part of m.content) { if (part.type === 'tool_result' && part.tool_use_id) { if (!tracker.seenToolResultIds.has(part.tool_use_id)) { tracker.seenToolResultIds.add(part.tool_use_id) - tracker.toolResultCount++ - newCount++ + // Skip protected tools for nudge frequency counting + const toolName = tracker.getToolName?.(part.tool_use_id) + if (!toolName || !protectedTools.has(toolName)) { + tracker.toolResultCount++ + newCount++ + } } } } @@ -48,7 +56,7 @@ export function trackNewToolResults(messages: any[], tracker: ToolTracker): numb * Uses position-based tracking since Gemini doesn't have tool call IDs. * Returns the number of NEW tools found (since last call). */ -export function trackNewToolResultsGemini(contents: any[], tracker: ToolTracker): number { +export function trackNewToolResultsGemini(contents: any[], tracker: ToolTracker, protectedTools: Set): number { let newCount = 0 let positionCounter = 0 for (const content of contents) { @@ -60,8 +68,12 @@ export function trackNewToolResultsGemini(contents: any[], tracker: ToolTracker) positionCounter++ if (!tracker.seenToolResultIds.has(positionId)) { tracker.seenToolResultIds.add(positionId) - tracker.toolResultCount++ - newCount++ + // Skip protected tools for nudge frequency counting + const toolName = part.functionResponse.name + if (!toolName || !protectedTools.has(toolName)) { + tracker.toolResultCount++ + newCount++ + } } } } @@ -73,14 +85,18 @@ export function trackNewToolResultsGemini(contents: any[], tracker: ToolTracker) * Track new tool results in OpenAI Responses API input. * Returns the number of NEW tools found (since last call). */ -export function trackNewToolResultsResponses(input: any[], tracker: ToolTracker): number { +export function trackNewToolResultsResponses(input: any[], tracker: ToolTracker, protectedTools: Set): number { let newCount = 0 for (const item of input) { if (item.type === 'function_call_output' && item.call_id) { if (!tracker.seenToolResultIds.has(item.call_id)) { tracker.seenToolResultIds.add(item.call_id) - tracker.toolResultCount++ - newCount++ + // Skip protected tools for nudge frequency counting + const toolName = tracker.getToolName?.(item.call_id) + if (!toolName || !protectedTools.has(toolName)) { + tracker.toolResultCount++ + newCount++ + } } } } diff --git a/lib/fetch-wrapper/gemini.ts b/lib/fetch-wrapper/gemini.ts index a472a01e..7eac6ad6 100644 --- a/lib/fetch-wrapper/gemini.ts +++ b/lib/fetch-wrapper/gemini.ts @@ -46,7 +46,8 @@ export async function handleGemini( if (prunableList) { // Track new tool results and check if nudge threshold is met - trackNewToolResultsGemini(body.contents, ctx.toolTracker) + const protectedSet = new Set(ctx.config.protectedTools) + trackNewToolResultsGemini(body.contents, ctx.toolTracker, protectedSet) const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq const endInjection = buildEndInjection(prunableList, includeNudge) @@ -99,6 +100,8 @@ export async function handleGemini( const toolPositionCounters = new Map() let replacedCount = 0 let totalFunctionResponses = 0 + let prunableFunctionResponses = 0 + const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase())) body.contents = body.contents.map((content: any) => { if (!Array.isArray(content.parts)) return content @@ -109,6 +112,11 @@ export async function handleGemini( totalFunctionResponses++ const funcName = part.functionResponse.name?.toLowerCase() + // Count as prunable if not a protected tool + if (!funcName || !protectedToolsLower.has(funcName)) { + prunableFunctionResponses++ + } + if (funcName) { // Get current position for this tool name and increment counter const currentIndex = toolPositionCounters.get(funcName) || 0 @@ -148,7 +156,7 @@ export async function handleGemini( if (replacedCount > 0) { ctx.logger.info("fetch", "Replaced pruned tool outputs (Google/Gemini)", { replaced: replacedCount, - total: totalFunctionResponses + total: prunableFunctionResponses }) if (ctx.logger.enabled) { diff --git a/lib/fetch-wrapper/index.ts b/lib/fetch-wrapper/index.ts index 868c9b8e..3c6f43a7 100644 --- a/lib/fetch-wrapper/index.ts +++ b/lib/fetch-wrapper/index.ts @@ -58,26 +58,24 @@ export function installFetchWrapper( // Capture tool IDs before handlers run to track what gets cached this request const toolIdsBefore = new Set(state.toolParameters.keys()) - // Try each format handler in order - // OpenAI Chat Completions & Anthropic style (body.messages) - if (body.messages && Array.isArray(body.messages)) { - const result = await handleOpenAIChatAndAnthropic(body, ctx, inputUrl) + // Try each format handler - mutually exclusive to avoid double-processing + // OpenAI Responses API style (body.input) - check first as it may also have messages + if (body.input && Array.isArray(body.input)) { + const result = await handleOpenAIResponses(body, ctx, inputUrl) if (result.modified) { modified = true } } - - // Google/Gemini style (body.contents) - if (body.contents && Array.isArray(body.contents)) { - const result = await handleGemini(body, ctx, inputUrl) + // OpenAI Chat Completions & Anthropic style (body.messages) + else if (body.messages && Array.isArray(body.messages)) { + const result = await handleOpenAIChatAndAnthropic(body, ctx, inputUrl) if (result.modified) { modified = true } } - - // OpenAI Responses API style (body.input) - if (body.input && Array.isArray(body.input)) { - const result = await handleOpenAIResponses(body, ctx, inputUrl) + // Google/Gemini style (body.contents) + else if (body.contents && Array.isArray(body.contents)) { + const result = await handleGemini(body, ctx, inputUrl) if (result.modified) { modified = true } diff --git a/lib/fetch-wrapper/openai-chat.ts b/lib/fetch-wrapper/openai-chat.ts index 9253e3fd..677e66f1 100644 --- a/lib/fetch-wrapper/openai-chat.ts +++ b/lib/fetch-wrapper/openai-chat.ts @@ -51,7 +51,8 @@ export async function handleOpenAIChatAndAnthropic( if (prunableList) { // Track new tool results and check if nudge threshold is met - trackNewToolResults(body.messages, ctx.toolTracker) + const protectedSet = new Set(ctx.config.protectedTools) + trackNewToolResults(body.messages, ctx.toolTracker, protectedSet) const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq const endInjection = buildEndInjection(prunableList, includeNudge) @@ -70,6 +71,9 @@ export async function handleOpenAIChatAndAnthropic( // Check for tool messages in both formats: // 1. OpenAI style: role === 'tool' // 2. Anthropic style: role === 'user' with content containing tool_result + const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase())) + + // Count all tool messages const toolMessages = body.messages.filter((m: any) => { if (m.role === 'tool') return true if (m.role === 'user' && Array.isArray(m.content)) { @@ -79,6 +83,29 @@ export async function handleOpenAIChatAndAnthropic( } return false }) + + // Count only prunable (non-protected) tool messages for the total + let prunableToolCount = 0 + for (const m of body.messages) { + if (m.role === 'tool') { + // Get tool name from cached metadata + const toolId = m.tool_call_id?.toLowerCase() + const metadata = toolId ? ctx.state.toolParameters.get(toolId) : undefined + if (!metadata || !protectedToolsLower.has(metadata.tool.toLowerCase())) { + prunableToolCount++ + } + } else if (m.role === 'user' && Array.isArray(m.content)) { + for (const part of m.content) { + if (part.type === 'tool_result') { + const toolId = part.tool_use_id?.toLowerCase() + const metadata = toolId ? ctx.state.toolParameters.get(toolId) : undefined + if (!metadata || !protectedToolsLower.has(metadata.tool.toLowerCase())) { + prunableToolCount++ + } + } + } + } + } const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger) @@ -123,7 +150,7 @@ export async function handleOpenAIChatAndAnthropic( if (replacedCount > 0) { ctx.logger.info("fetch", "Replaced pruned tool outputs", { replaced: replacedCount, - total: toolMessages.length + total: prunableToolCount }) if (ctx.logger.enabled) { diff --git a/lib/fetch-wrapper/openai-responses.ts b/lib/fetch-wrapper/openai-responses.ts index 62e2bb84..9e02c285 100644 --- a/lib/fetch-wrapper/openai-responses.ts +++ b/lib/fetch-wrapper/openai-responses.ts @@ -51,7 +51,8 @@ export async function handleOpenAIResponses( if (prunableList) { // Track new tool results and check if nudge threshold is met - trackNewToolResultsResponses(body.input, ctx.toolTracker) + const protectedSet = new Set(ctx.config.protectedTools) + trackNewToolResultsResponses(body.input, ctx.toolTracker, protectedSet) const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq const endInjection = buildEndInjection(prunableList, includeNudge) @@ -80,6 +81,16 @@ export async function handleOpenAIResponses( return { modified, body } } + // Count only prunable (non-protected) function outputs for the total + const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase())) + let prunableFunctionOutputCount = 0 + for (const item of functionOutputs) { + const toolName = item.name?.toLowerCase() + if (!toolName || !protectedToolsLower.has(toolName)) { + prunableFunctionOutputCount++ + } + } + let replacedCount = 0 body.input = body.input.map((item: any) => { @@ -96,7 +107,7 @@ export async function handleOpenAIResponses( if (replacedCount > 0) { ctx.logger.info("fetch", "Replaced pruned tool outputs (Responses API)", { replaced: replacedCount, - total: functionOutputs.length + total: prunableFunctionOutputCount }) if (ctx.logger.enabled) { diff --git a/lib/hooks.ts b/lib/hooks.ts index 1b2c01d3..1fa4c749 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -5,6 +5,7 @@ import { runOnIdle } from "./core/janitor" import type { PluginConfig, PruningStrategy } from "./config" import type { ToolTracker } from "./api-formats/synth-instruction" import { resetToolTrackerCount } from "./api-formats/synth-instruction" +import { clearAllMappings } from "./state/id-mapping" export async function isSubagentSession(client: any, sessionID: string): Promise { try { @@ -72,6 +73,18 @@ export function createChatParamsHandler( providerID = input.message.model.providerID } + // Detect session change and reset per-session state + if (state.lastSeenSessionId && state.lastSeenSessionId !== sessionId) { + logger.info("chat.params", "Session changed, resetting state", { + from: state.lastSeenSessionId.substring(0, 8), + to: sessionId.substring(0, 8) + }) + // Clear ID mappings from previous session + clearAllMappings() + // Clear tool parameters cache (not session-scoped, so must be cleared) + state.toolParameters.clear() + } + // Track the last seen session ID for fetch wrapper correlation state.lastSeenSessionId = sessionId diff --git a/lib/state/id-mapping.ts b/lib/state/id-mapping.ts index 110f4c4b..d7ed9a03 100644 --- a/lib/state/id-mapping.ts +++ b/lib/state/id-mapping.ts @@ -91,6 +91,22 @@ export function hasMapping(sessionId: string): boolean { return sessionMappings.has(sessionId) } +/** + * Clears all ID mappings for a specific session. + * Call this when a session ends or when switching to a new session. + */ +export function clearSessionMapping(sessionId: string): void { + sessionMappings.delete(sessionId) +} + +/** + * Clears all session mappings. + * Call this when switching sessions to ensure clean state. + */ +export function clearAllMappings(): void { + sessionMappings.clear() +} + /** * Gets the next numeric ID that will be assigned (without assigning it). * Useful for knowing the current state. From 0c1017a9805fde50208367b8c2ad7b1cfe1694ba Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Wed, 3 Dec 2025 13:58:53 -0500 Subject: [PATCH 6/9] chore: remove obvious comments --- lib/api-formats/synth-instruction.ts | 21 ------------------- lib/fetch-wrapper/gemini.ts | 10 --------- lib/fetch-wrapper/index.ts | 6 +----- lib/fetch-wrapper/openai-chat.ts | 13 ------------ lib/fetch-wrapper/openai-responses.ts | 7 ------- lib/hooks.ts | 10 +-------- lib/pruning-tool.ts | 20 ++---------------- lib/state/id-mapping.ts | 30 --------------------------- lib/state/tool-cache.ts | 3 --- 9 files changed, 4 insertions(+), 116 deletions(-) diff --git a/lib/api-formats/synth-instruction.ts b/lib/api-formats/synth-instruction.ts index c3bf6319..7301c17c 100644 --- a/lib/api-formats/synth-instruction.ts +++ b/lib/api-formats/synth-instruction.ts @@ -9,7 +9,6 @@ export function createToolTracker(): ToolTracker { return { seenToolResultIds: new Set(), toolResultCount: 0, skipNextIdle: false } } -/** Reset tool count to 0 (called after a prune event) */ export function resetToolTrackerCount(tracker: ToolTracker): void { tracker.toolResultCount = 0 } @@ -25,7 +24,6 @@ export function trackNewToolResults(messages: any[], tracker: ToolTracker, prote if (m.role === 'tool' && m.tool_call_id) { if (!tracker.seenToolResultIds.has(m.tool_call_id)) { tracker.seenToolResultIds.add(m.tool_call_id) - // Skip protected tools for nudge frequency counting const toolName = tracker.getToolName?.(m.tool_call_id) if (!toolName || !protectedTools.has(toolName)) { tracker.toolResultCount++ @@ -37,7 +35,6 @@ export function trackNewToolResults(messages: any[], tracker: ToolTracker, prote if (part.type === 'tool_result' && part.tool_use_id) { if (!tracker.seenToolResultIds.has(part.tool_use_id)) { tracker.seenToolResultIds.add(part.tool_use_id) - // Skip protected tools for nudge frequency counting const toolName = tracker.getToolName?.(part.tool_use_id) if (!toolName || !protectedTools.has(toolName)) { tracker.toolResultCount++ @@ -63,12 +60,10 @@ export function trackNewToolResultsGemini(contents: any[], tracker: ToolTracker, if (!Array.isArray(content.parts)) continue for (const part of content.parts) { if (part.functionResponse) { - // Use position-based ID since Gemini doesn't have tool_call_id const positionId = `gemini_pos_${positionCounter}` positionCounter++ if (!tracker.seenToolResultIds.has(positionId)) { tracker.seenToolResultIds.add(positionId) - // Skip protected tools for nudge frequency counting const toolName = part.functionResponse.name if (!toolName || !protectedTools.has(toolName)) { tracker.toolResultCount++ @@ -91,7 +86,6 @@ export function trackNewToolResultsResponses(input: any[], tracker: ToolTracker, if (item.type === 'function_call_output' && item.call_id) { if (!tracker.seenToolResultIds.has(item.call_id)) { tracker.seenToolResultIds.add(item.call_id) - // Skip protected tools for nudge frequency counting const toolName = tracker.getToolName?.(item.call_id) if (!toolName || !protectedTools.has(toolName)) { tracker.toolResultCount++ @@ -103,11 +97,6 @@ export function trackNewToolResultsResponses(input: any[], tracker: ToolTracker, return newCount } -// ============================================================================ -// OpenAI Chat / Anthropic Format -// ============================================================================ - -/** Check if a message content matches nudge text (OpenAI/Anthropic format) */ function isNudgeMessage(msg: any, nudgeText: string): boolean { if (typeof msg.content === 'string') { return msg.content === nudgeText @@ -138,11 +127,6 @@ export function injectSynth(messages: any[], instruction: string, nudgeText: str return false } -// ============================================================================ -// Google/Gemini Format (body.contents with parts) -// ============================================================================ - -/** Check if a Gemini content matches nudge text */ function isNudgeContentGemini(content: any, nudgeText: string): boolean { if (Array.isArray(content.parts) && content.parts.length === 1) { const part = content.parts[0] @@ -169,11 +153,6 @@ export function injectSynthGemini(contents: any[], instruction: string, nudgeTex return false } -// ============================================================================ -// OpenAI Responses API Format (body.input with type-based items) -// ============================================================================ - -/** Check if a Responses API item matches nudge text */ function isNudgeItemResponses(item: any, nudgeText: string): boolean { if (typeof item.content === 'string') { return item.content === nudgeText diff --git a/lib/fetch-wrapper/gemini.ts b/lib/fetch-wrapper/gemini.ts index 7eac6ad6..b3d73fb3 100644 --- a/lib/fetch-wrapper/gemini.ts +++ b/lib/fetch-wrapper/gemini.ts @@ -22,14 +22,11 @@ export async function handleGemini( let modified = false - // Inject synthetic instructions if onTool strategies are enabled if (ctx.config.strategies.onTool.length > 0) { - // Inject base synthetic instructions (appended to last user content) if (injectSynthGemini(body.contents, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { modified = true } - // Build and inject prunable tools list at the end const sessionId = ctx.state.lastSeenSessionId if (sessionId) { const toolIds = Array.from(ctx.state.toolParameters.keys()) @@ -45,7 +42,6 @@ export async function handleGemini( ) if (prunableList) { - // Track new tool results and check if nudge threshold is met const protectedSet = new Set(ctx.config.protectedTools) trackNewToolResultsGemini(body.contents, ctx.toolTracker, protectedSet) const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq @@ -63,7 +59,6 @@ export async function handleGemini( } } - // Check for functionResponse parts in any content item const hasFunctionResponses = body.contents.some((content: any) => Array.isArray(content.parts) && content.parts.some((part: any) => part.functionResponse) @@ -79,7 +74,6 @@ export async function handleGemini( return { modified, body } } - // Find the active session to get the position mapping const activeSessions = allSessions.data?.filter((s: any) => !s.parentID) || [] let positionMapping: Map | undefined @@ -112,17 +106,14 @@ export async function handleGemini( totalFunctionResponses++ const funcName = part.functionResponse.name?.toLowerCase() - // Count as prunable if not a protected tool if (!funcName || !protectedToolsLower.has(funcName)) { prunableFunctionResponses++ } if (funcName) { - // Get current position for this tool name and increment counter const currentIndex = toolPositionCounters.get(funcName) || 0 toolPositionCounters.set(funcName, currentIndex + 1) - // Look up the tool call ID using position const positionKey = `${funcName}:${currentIndex}` const toolCallId = positionMapping!.get(positionKey) @@ -130,7 +121,6 @@ export async function handleGemini( contentModified = true replacedCount++ // Preserve thoughtSignature if present (required for Gemini 3 Pro) - // response must be a Struct (object), not a plain string return { ...part, functionResponse: { diff --git a/lib/fetch-wrapper/index.ts b/lib/fetch-wrapper/index.ts index 3c6f43a7..561c7174 100644 --- a/lib/fetch-wrapper/index.ts +++ b/lib/fetch-wrapper/index.ts @@ -55,25 +55,21 @@ export function installFetchWrapper( const inputUrl = typeof input === 'string' ? input : 'URL object' let modified = false - // Capture tool IDs before handlers run to track what gets cached this request const toolIdsBefore = new Set(state.toolParameters.keys()) - // Try each format handler - mutually exclusive to avoid double-processing - // OpenAI Responses API style (body.input) - check first as it may also have messages + // Mutually exclusive format handlers to avoid double-processing if (body.input && Array.isArray(body.input)) { const result = await handleOpenAIResponses(body, ctx, inputUrl) if (result.modified) { modified = true } } - // OpenAI Chat Completions & Anthropic style (body.messages) else if (body.messages && Array.isArray(body.messages)) { const result = await handleOpenAIChatAndAnthropic(body, ctx, inputUrl) if (result.modified) { modified = true } } - // Google/Gemini style (body.contents) else if (body.contents && Array.isArray(body.contents)) { const result = await handleGemini(body, ctx, inputUrl) if (result.modified) { diff --git a/lib/fetch-wrapper/openai-chat.ts b/lib/fetch-wrapper/openai-chat.ts index 677e66f1..6488d5a6 100644 --- a/lib/fetch-wrapper/openai-chat.ts +++ b/lib/fetch-wrapper/openai-chat.ts @@ -22,19 +22,15 @@ export async function handleOpenAIChatAndAnthropic( return { modified: false, body } } - // Cache tool parameters from messages (OpenAI and Anthropic formats) cacheToolParametersFromMessages(body.messages, ctx.state, ctx.logger) let modified = false - // Inject synthetic instructions if onTool strategies are enabled if (ctx.config.strategies.onTool.length > 0) { - // Inject base synthetic instructions (appended to last user message) if (injectSynth(body.messages, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { modified = true } - // Build and inject prunable tools list at the end const sessionId = ctx.state.lastSeenSessionId if (sessionId) { const toolIds = Array.from(ctx.state.toolParameters.keys()) @@ -50,7 +46,6 @@ export async function handleOpenAIChatAndAnthropic( ) if (prunableList) { - // Track new tool results and check if nudge threshold is met const protectedSet = new Set(ctx.config.protectedTools) trackNewToolResults(body.messages, ctx.toolTracker, protectedSet) const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq @@ -68,12 +63,8 @@ export async function handleOpenAIChatAndAnthropic( } } - // Check for tool messages in both formats: - // 1. OpenAI style: role === 'tool' - // 2. Anthropic style: role === 'user' with content containing tool_result const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase())) - // Count all tool messages const toolMessages = body.messages.filter((m: any) => { if (m.role === 'tool') return true if (m.role === 'user' && Array.isArray(m.content)) { @@ -84,11 +75,9 @@ export async function handleOpenAIChatAndAnthropic( return false }) - // Count only prunable (non-protected) tool messages for the total let prunableToolCount = 0 for (const m of body.messages) { if (m.role === 'tool') { - // Get tool name from cached metadata const toolId = m.tool_call_id?.toLowerCase() const metadata = toolId ? ctx.state.toolParameters.get(toolId) : undefined if (!metadata || !protectedToolsLower.has(metadata.tool.toLowerCase())) { @@ -116,7 +105,6 @@ export async function handleOpenAIChatAndAnthropic( let replacedCount = 0 body.messages = body.messages.map((m: any) => { - // OpenAI style: role === 'tool' with tool_call_id if (m.role === 'tool' && allPrunedIds.has(m.tool_call_id?.toLowerCase())) { replacedCount++ return { @@ -125,7 +113,6 @@ export async function handleOpenAIChatAndAnthropic( } } - // Anthropic style: role === 'user' with content array containing tool_result if (m.role === 'user' && Array.isArray(m.content)) { let messageModified = false const newContent = m.content.map((part: any) => { diff --git a/lib/fetch-wrapper/openai-responses.ts b/lib/fetch-wrapper/openai-responses.ts index 9e02c285..2c71c6d2 100644 --- a/lib/fetch-wrapper/openai-responses.ts +++ b/lib/fetch-wrapper/openai-responses.ts @@ -22,19 +22,15 @@ export async function handleOpenAIResponses( return { modified: false, body } } - // Cache tool parameters from input (OpenAI Responses API format) cacheToolParametersFromInput(body.input, ctx.state, ctx.logger) let modified = false - // Inject synthetic instructions if onTool strategies are enabled if (ctx.config.strategies.onTool.length > 0) { - // Inject base synthetic instructions (appended to last user message) if (injectSynthResponses(body.input, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { modified = true } - // Build and inject prunable tools list at the end const sessionId = ctx.state.lastSeenSessionId if (sessionId) { const toolIds = Array.from(ctx.state.toolParameters.keys()) @@ -50,7 +46,6 @@ export async function handleOpenAIResponses( ) if (prunableList) { - // Track new tool results and check if nudge threshold is met const protectedSet = new Set(ctx.config.protectedTools) trackNewToolResultsResponses(body.input, ctx.toolTracker, protectedSet) const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq @@ -68,7 +63,6 @@ export async function handleOpenAIResponses( } } - // Check for function_call_output items const functionOutputs = body.input.filter((item: any) => item.type === 'function_call_output') if (functionOutputs.length === 0) { @@ -81,7 +75,6 @@ export async function handleOpenAIResponses( return { modified, body } } - // Count only prunable (non-protected) function outputs for the total const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase())) let prunableFunctionOutputCount = 0 for (const item of functionOutputs) { diff --git a/lib/hooks.ts b/lib/hooks.ts index 1fa4c749..77b530a2 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -32,8 +32,7 @@ export function createEventHandler( if (await isSubagentSession(client, event.properties.sessionID)) return if (config.strategies.onIdle.length === 0) return - // Skip idle pruning if the last tool used was prune - // and idle strategies cover the same work as tool strategies + // Skip idle pruning if the last tool used was prune and idle strategies cover tool strategies if (toolTracker?.skipNextIdle) { toolTracker.skipNextIdle = false if (toolStrategiesCoveredByIdle(config.strategies.onIdle, config.strategies.onTool)) { @@ -73,22 +72,17 @@ export function createChatParamsHandler( providerID = input.message.model.providerID } - // Detect session change and reset per-session state if (state.lastSeenSessionId && state.lastSeenSessionId !== sessionId) { logger.info("chat.params", "Session changed, resetting state", { from: state.lastSeenSessionId.substring(0, 8), to: sessionId.substring(0, 8) }) - // Clear ID mappings from previous session clearAllMappings() - // Clear tool parameters cache (not session-scoped, so must be cleared) state.toolParameters.clear() } - // Track the last seen session ID for fetch wrapper correlation state.lastSeenSessionId = sessionId - // Check if this is a subagent session if (!state.checkedSessions.has(sessionId)) { state.checkedSessions.add(sessionId) const isSubagent = await isSubagentSession(client, sessionId) @@ -108,7 +102,6 @@ export function createChatParamsHandler( const messages = messagesResponse.data || messagesResponse if (Array.isArray(messages)) { - // Build position mapping: track tool calls by name and occurrence index const toolCallsByName = new Map() for (const msg of messages) { @@ -125,7 +118,6 @@ export function createChatParamsHandler( } } - // Create position mapping: "toolName:index" -> toolCallId const positionMapping = new Map() for (const [toolName, callIds] of toolCallsByName) { callIds.forEach((callId, index) => { diff --git a/lib/pruning-tool.ts b/lib/pruning-tool.ts index cec7b359..891d657b 100644 --- a/lib/pruning-tool.ts +++ b/lib/pruning-tool.ts @@ -44,20 +44,16 @@ export function createPruningTool( const { client, state, logger, config, notificationCtx, workingDirectory } = ctx const sessionId = toolCtx.sessionID - // Skip pruning in subagent sessions if (await isSubagentSession(client, sessionId)) { return "Pruning is unavailable in subagent sessions. Do not call this tool again. Continue with your current task." } - // Validate input if (!args.ids || args.ids.length === 0) { return "No IDs provided. Check the list for available IDs to prune." } - // Restore persisted state if needed await ensureSessionRestored(state, sessionId, logger) - // Convert numeric IDs to actual tool call IDs const prunedIds = args.ids .map(numId => getActualId(sessionId, numId)) .filter((id): id is string => id !== undefined) @@ -66,10 +62,8 @@ export function createPruningTool( return "None of the provided IDs were valid. Check the list for available IDs." } - // Calculate tokens saved - const tokensSaved = await calculateTokensSaved(client, sessionId, prunedIds, state) + const tokensSaved = await calculateTokensSaved(client, sessionId, prunedIds) - // Update stats const currentStats = state.stats.get(sessionId) ?? { totalToolsPruned: 0, totalTokensSaved: 0, @@ -83,17 +77,13 @@ export function createPruningTool( } state.stats.set(sessionId, sessionStats) - // Update pruned IDs state const alreadyPrunedIds = state.prunedIds.get(sessionId) ?? [] const allPrunedIds = [...alreadyPrunedIds, ...prunedIds] state.prunedIds.set(sessionId, allPrunedIds) - // Persist state saveSessionState(sessionId, new Set(allPrunedIds), sessionStats, logger) .catch(err => logger.error("prune-tool", "Failed to persist state", { error: err.message })) - // Build tool metadata for notification - // Keys are normalized to lowercase to match lookup in notification.ts const toolMetadata = new Map() for (const id of prunedIds) { // Try both original and lowercase since caching may vary @@ -110,7 +100,6 @@ export function createPruningTool( } } - // Send notification to user await sendUnifiedNotification(notificationCtx, sessionId, { aiPrunedCount: prunedIds.length, aiTokensSaved: tokensSaved, @@ -120,15 +109,12 @@ export function createPruningTool( sessionStats }) - // Skip next idle pruning since we just pruned toolTracker.skipNextIdle = true - // Reset nudge counter if (config.nudge_freq > 0) { resetToolTrackerCount(toolTracker) } - // Format result for the AI const result = { prunedCount: prunedIds.length, tokensSaved, @@ -150,11 +136,9 @@ export function createPruningTool( async function calculateTokensSaved( client: any, sessionId: string, - prunedIds: string[], - state: PluginState + prunedIds: string[] ): Promise { try { - // Fetch session messages to get tool output content const messagesResponse = await client.session.messages({ path: { id: sessionId }, query: { limit: 200 } diff --git a/lib/state/id-mapping.ts b/lib/state/id-mapping.ts index d7ed9a03..3421dbc4 100644 --- a/lib/state/id-mapping.ts +++ b/lib/state/id-mapping.ts @@ -20,9 +20,6 @@ export interface IdMapping { /** Per-session ID mappings */ const sessionMappings = new Map() -/** - * Gets or creates the ID mapping for a session. - */ function getSessionMapping(sessionId: string): IdMapping { let mapping = sessionMappings.get(sessionId) if (!mapping) { @@ -57,60 +54,33 @@ export function getOrCreateNumericId(sessionId: string, actualId: string): numbe return numericId } -/** - * Looks up the actual tool call ID for a numeric ID. - * Returns undefined if the numeric ID doesn't exist. - */ export function getActualId(sessionId: string, numericId: number): string | undefined { const mapping = sessionMappings.get(sessionId) return mapping?.numericToActual.get(numericId) } -/** - * Looks up the numeric ID for an actual tool call ID. - * Returns undefined if not mapped. - */ export function getNumericId(sessionId: string, actualId: string): number | undefined { const mapping = sessionMappings.get(sessionId) return mapping?.actualToNumeric.get(actualId) } -/** - * Gets all current mappings for a session. - * Useful for debugging and building the prunable tools list. - */ export function getAllMappings(sessionId: string): Map { const mapping = sessionMappings.get(sessionId) return mapping?.numericToActual ?? new Map() } -/** - * Checks if a session has any ID mappings. - */ export function hasMapping(sessionId: string): boolean { return sessionMappings.has(sessionId) } -/** - * Clears all ID mappings for a specific session. - * Call this when a session ends or when switching to a new session. - */ export function clearSessionMapping(sessionId: string): void { sessionMappings.delete(sessionId) } -/** - * Clears all session mappings. - * Call this when switching sessions to ensure clean state. - */ export function clearAllMappings(): void { sessionMappings.clear() } -/** - * Gets the next numeric ID that will be assigned (without assigning it). - * Useful for knowing the current state. - */ export function getNextId(sessionId: string): number { const mapping = sessionMappings.get(sessionId) return mapping?.nextId ?? 1 diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index b319c375..8c27ab17 100644 --- a/lib/state/tool-cache.ts +++ b/lib/state/tool-cache.ts @@ -22,7 +22,6 @@ export function cacheToolParametersFromMessages( continue } - // OpenAI format: tool_calls array if (Array.isArray(message.tool_calls)) { for (const toolCall of message.tool_calls) { if (!toolCall.id || !toolCall.function) { @@ -44,7 +43,6 @@ export function cacheToolParametersFromMessages( } } - // Anthropic format: content array with tool_use blocks if (Array.isArray(message.content)) { for (const part of message.content) { if (part.type !== 'tool_use' || !part.id || !part.name) { @@ -60,7 +58,6 @@ export function cacheToolParametersFromMessages( } } - // Log cache results if anything was cached if (logger && (openaiCached > 0 || anthropicCached > 0)) { logger.debug("tool-cache", "Cached tool parameters from messages", { openaiFormat: openaiCached, From 035285fe43d1c0f568cfb0a6495c446e2b3bb21c Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Wed, 3 Dec 2025 14:54:20 -0500 Subject: [PATCH 7/9] refactor: consolidate format handlers with FormatDescriptor pattern and fix tool ID case handling --- index.ts | 2 +- lib/fetch-wrapper/formats/gemini.ts | 173 +++++++++++++++++ lib/fetch-wrapper/formats/index.ts | 3 + lib/fetch-wrapper/formats/openai-chat.ts | 130 +++++++++++++ lib/fetch-wrapper/formats/openai-responses.ts | 88 +++++++++ lib/fetch-wrapper/gemini.ts | 176 ------------------ lib/fetch-wrapper/handler.ts | 130 +++++++++++++ lib/fetch-wrapper/index.ts | 17 +- lib/fetch-wrapper/openai-chat.ts | 165 ---------------- lib/fetch-wrapper/openai-responses.ts | 129 ------------- lib/fetch-wrapper/types.ts | 51 +++++ lib/hooks.ts | 16 +- lib/pruning-tool.ts | 5 +- lib/state/tool-cache.ts | 18 +- 14 files changed, 609 insertions(+), 494 deletions(-) create mode 100644 lib/fetch-wrapper/formats/gemini.ts create mode 100644 lib/fetch-wrapper/formats/index.ts create mode 100644 lib/fetch-wrapper/formats/openai-chat.ts create mode 100644 lib/fetch-wrapper/formats/openai-responses.ts delete mode 100644 lib/fetch-wrapper/gemini.ts create mode 100644 lib/fetch-wrapper/handler.ts delete mode 100644 lib/fetch-wrapper/openai-chat.ts delete mode 100644 lib/fetch-wrapper/openai-responses.ts diff --git a/index.ts b/index.ts index 4735d17e..a97c7b55 100644 --- a/index.ts +++ b/index.ts @@ -45,7 +45,7 @@ const plugin: Plugin = (async (ctx) => { // Wire up tool name lookup from the cached tool parameters toolTracker.getToolName = (callId: string) => { - const entry = state.toolParameters.get(callId) + const entry = state.toolParameters.get(callId.toLowerCase()) return entry?.tool } diff --git a/lib/fetch-wrapper/formats/gemini.ts b/lib/fetch-wrapper/formats/gemini.ts new file mode 100644 index 00000000..765d9cb9 --- /dev/null +++ b/lib/fetch-wrapper/formats/gemini.ts @@ -0,0 +1,173 @@ +import type { FormatDescriptor, ToolOutput } from "../types" +import { PRUNED_CONTENT_MESSAGE } from "../types" +import type { PluginState } from "../../state" +import type { Logger } from "../../logger" +import type { ToolTracker } from "../../api-formats/synth-instruction" +import { injectSynthGemini, trackNewToolResultsGemini } from "../../api-formats/synth-instruction" +import { injectPrunableListGemini } from "../../api-formats/prunable-list" + +/** + * Format descriptor for Google/Gemini API. + * + * Uses body.contents array with: + * - parts[].functionCall for tool invocations + * - parts[].functionResponse for tool results + * + * IMPORTANT: Gemini doesn't include tool call IDs in its native format. + * We use position-based correlation via state.googleToolCallMapping which maps + * "toolName:index" -> "toolCallId" (populated by hooks.ts from message events). + */ +export const geminiFormat: FormatDescriptor = { + name: 'gemini', + + detect(body: any): boolean { + return body.contents && Array.isArray(body.contents) + }, + + getDataArray(body: any): any[] | undefined { + return body.contents + }, + + cacheToolParameters(_data: any[], _state: PluginState, _logger?: Logger): void { + // Gemini format doesn't include tool parameters in the request body. + // Tool parameters are captured via message events in hooks.ts and stored + // in state.googleToolCallMapping for position-based correlation. + // No-op here. + }, + + injectSynth(data: any[], instruction: string, nudgeText: string): boolean { + return injectSynthGemini(data, instruction, nudgeText) + }, + + trackNewToolResults(data: any[], tracker: ToolTracker, protectedTools: Set): number { + return trackNewToolResultsGemini(data, tracker, protectedTools) + }, + + injectPrunableList(data: any[], injection: string): boolean { + return injectPrunableListGemini(data, injection) + }, + + extractToolOutputs(data: any[], state: PluginState): ToolOutput[] { + const outputs: ToolOutput[] = [] + + // We need the position mapping to correlate functionResponses to tool call IDs + // Find the mapping from any active session + let positionMapping: Map | undefined + for (const [_sessionId, mapping] of state.googleToolCallMapping) { + if (mapping && mapping.size > 0) { + positionMapping = mapping + break + } + } + + if (!positionMapping) { + return outputs + } + + // Track position counters per tool name + const toolPositionCounters = new Map() + + for (const content of data) { + if (!Array.isArray(content.parts)) continue + + for (const part of content.parts) { + if (part.functionResponse) { + const funcName = part.functionResponse.name?.toLowerCase() + if (funcName) { + const currentIndex = toolPositionCounters.get(funcName) || 0 + toolPositionCounters.set(funcName, currentIndex + 1) + + const positionKey = `${funcName}:${currentIndex}` + const toolCallId = positionMapping.get(positionKey) + + if (toolCallId) { + outputs.push({ + id: toolCallId.toLowerCase(), + toolName: funcName + }) + } + } + } + } + } + + return outputs + }, + + replaceToolOutput(data: any[], toolId: string, prunedMessage: string, state: PluginState): boolean { + // Find the position mapping + let positionMapping: Map | undefined + for (const [_sessionId, mapping] of state.googleToolCallMapping) { + if (mapping && mapping.size > 0) { + positionMapping = mapping + break + } + } + + if (!positionMapping) { + return false + } + + const toolIdLower = toolId.toLowerCase() + const toolPositionCounters = new Map() + let replaced = false + + for (let i = 0; i < data.length; i++) { + const content = data[i] + if (!Array.isArray(content.parts)) continue + + let contentModified = false + const newParts = content.parts.map((part: any) => { + if (part.functionResponse) { + const funcName = part.functionResponse.name?.toLowerCase() + if (funcName) { + const currentIndex = toolPositionCounters.get(funcName) || 0 + toolPositionCounters.set(funcName, currentIndex + 1) + + const positionKey = `${funcName}:${currentIndex}` + const mappedToolId = positionMapping!.get(positionKey) + + if (mappedToolId?.toLowerCase() === toolIdLower) { + contentModified = true + replaced = true + // Preserve thoughtSignature if present (required for Gemini 3 Pro) + return { + ...part, + functionResponse: { + ...part.functionResponse, + response: { + name: part.functionResponse.name, + content: prunedMessage + } + } + } + } + } + } + return part + }) + + if (contentModified) { + data[i] = { ...content, parts: newParts } + } + } + + return replaced + }, + + hasToolOutputs(data: any[]): boolean { + return data.some((content: any) => + Array.isArray(content.parts) && + content.parts.some((part: any) => part.functionResponse) + ) + }, + + getLogMetadata(data: any[], replacedCount: number, inputUrl: string): Record { + return { + url: inputUrl, + replacedCount, + totalContents: data.length, + format: 'google-gemini' + } + } +} diff --git a/lib/fetch-wrapper/formats/index.ts b/lib/fetch-wrapper/formats/index.ts new file mode 100644 index 00000000..0132c878 --- /dev/null +++ b/lib/fetch-wrapper/formats/index.ts @@ -0,0 +1,3 @@ +export { openaiChatFormat } from './openai-chat' +export { openaiResponsesFormat } from './openai-responses' +export { geminiFormat } from './gemini' diff --git a/lib/fetch-wrapper/formats/openai-chat.ts b/lib/fetch-wrapper/formats/openai-chat.ts new file mode 100644 index 00000000..b75f9f19 --- /dev/null +++ b/lib/fetch-wrapper/formats/openai-chat.ts @@ -0,0 +1,130 @@ +import type { FormatDescriptor, ToolOutput } from "../types" +import { PRUNED_CONTENT_MESSAGE } from "../types" +import type { PluginState } from "../../state" +import type { Logger } from "../../logger" +import type { ToolTracker } from "../../api-formats/synth-instruction" +import { cacheToolParametersFromMessages } from "../../state/tool-cache" +import { injectSynth, trackNewToolResults } from "../../api-formats/synth-instruction" +import { injectPrunableList } from "../../api-formats/prunable-list" + +/** + * Format descriptor for OpenAI Chat Completions and Anthropic APIs. + * + * OpenAI Chat format: + * - Messages with role='tool' and tool_call_id + * - Assistant messages with tool_calls[] array + * + * Anthropic format: + * - Messages with role='user' containing content[].type='tool_result' and tool_use_id + * - Assistant messages with content[].type='tool_use' + */ +export const openaiChatFormat: FormatDescriptor = { + name: 'openai-chat', + + detect(body: any): boolean { + return body.messages && Array.isArray(body.messages) + }, + + getDataArray(body: any): any[] | undefined { + return body.messages + }, + + cacheToolParameters(data: any[], state: PluginState, logger?: Logger): void { + cacheToolParametersFromMessages(data, state, logger) + }, + + injectSynth(data: any[], instruction: string, nudgeText: string): boolean { + return injectSynth(data, instruction, nudgeText) + }, + + trackNewToolResults(data: any[], tracker: ToolTracker, protectedTools: Set): number { + return trackNewToolResults(data, tracker, protectedTools) + }, + + injectPrunableList(data: any[], injection: string): boolean { + return injectPrunableList(data, injection) + }, + + extractToolOutputs(data: any[], state: PluginState): ToolOutput[] { + const outputs: ToolOutput[] = [] + + for (const m of data) { + // OpenAI Chat format: role='tool' with tool_call_id + if (m.role === 'tool' && m.tool_call_id) { + const metadata = state.toolParameters.get(m.tool_call_id.toLowerCase()) + outputs.push({ + id: m.tool_call_id.toLowerCase(), + toolName: metadata?.tool + }) + } + + // Anthropic format: role='user' with content[].type='tool_result' + if (m.role === 'user' && Array.isArray(m.content)) { + for (const part of m.content) { + if (part.type === 'tool_result' && part.tool_use_id) { + const metadata = state.toolParameters.get(part.tool_use_id.toLowerCase()) + outputs.push({ + id: part.tool_use_id.toLowerCase(), + toolName: metadata?.tool + }) + } + } + } + } + + return outputs + }, + + replaceToolOutput(data: any[], toolId: string, prunedMessage: string, _state: PluginState): boolean { + const toolIdLower = toolId.toLowerCase() + let replaced = false + + for (let i = 0; i < data.length; i++) { + const m = data[i] + + // OpenAI Chat format + if (m.role === 'tool' && m.tool_call_id?.toLowerCase() === toolIdLower) { + data[i] = { ...m, content: prunedMessage } + replaced = true + } + + // Anthropic format + if (m.role === 'user' && Array.isArray(m.content)) { + let messageModified = false + const newContent = m.content.map((part: any) => { + if (part.type === 'tool_result' && part.tool_use_id?.toLowerCase() === toolIdLower) { + messageModified = true + return { ...part, content: prunedMessage } + } + return part + }) + if (messageModified) { + data[i] = { ...m, content: newContent } + replaced = true + } + } + } + + return replaced + }, + + hasToolOutputs(data: any[]): boolean { + for (const m of data) { + if (m.role === 'tool') return true + if (m.role === 'user' && Array.isArray(m.content)) { + for (const part of m.content) { + if (part.type === 'tool_result') return true + } + } + } + return false + }, + + getLogMetadata(data: any[], replacedCount: number, inputUrl: string): Record { + return { + url: inputUrl, + replacedCount, + totalMessages: data.length + } + } +} diff --git a/lib/fetch-wrapper/formats/openai-responses.ts b/lib/fetch-wrapper/formats/openai-responses.ts new file mode 100644 index 00000000..67ac5b7b --- /dev/null +++ b/lib/fetch-wrapper/formats/openai-responses.ts @@ -0,0 +1,88 @@ +import type { FormatDescriptor, ToolOutput } from "../types" +import { PRUNED_CONTENT_MESSAGE } from "../types" +import type { PluginState } from "../../state" +import type { Logger } from "../../logger" +import type { ToolTracker } from "../../api-formats/synth-instruction" +import { cacheToolParametersFromInput } from "../../state/tool-cache" +import { injectSynthResponses, trackNewToolResultsResponses } from "../../api-formats/synth-instruction" +import { injectPrunableListResponses } from "../../api-formats/prunable-list" + +/** + * Format descriptor for OpenAI Responses API (GPT-5 models via sdk.responses()). + * + * Uses body.input array with: + * - type='function_call' items for tool calls + * - type='function_call_output' items for tool results + * - type='message' items for user/assistant messages + */ +export const openaiResponsesFormat: FormatDescriptor = { + name: 'openai-responses', + + detect(body: any): boolean { + return body.input && Array.isArray(body.input) + }, + + getDataArray(body: any): any[] | undefined { + return body.input + }, + + cacheToolParameters(data: any[], state: PluginState, logger?: Logger): void { + cacheToolParametersFromInput(data, state, logger) + }, + + injectSynth(data: any[], instruction: string, nudgeText: string): boolean { + return injectSynthResponses(data, instruction, nudgeText) + }, + + trackNewToolResults(data: any[], tracker: ToolTracker, protectedTools: Set): number { + return trackNewToolResultsResponses(data, tracker, protectedTools) + }, + + injectPrunableList(data: any[], injection: string): boolean { + return injectPrunableListResponses(data, injection) + }, + + extractToolOutputs(data: any[], state: PluginState): ToolOutput[] { + const outputs: ToolOutput[] = [] + + for (const item of data) { + if (item.type === 'function_call_output' && item.call_id) { + const metadata = state.toolParameters.get(item.call_id.toLowerCase()) + outputs.push({ + id: item.call_id.toLowerCase(), + toolName: metadata?.tool ?? item.name + }) + } + } + + return outputs + }, + + replaceToolOutput(data: any[], toolId: string, prunedMessage: string, _state: PluginState): boolean { + const toolIdLower = toolId.toLowerCase() + let replaced = false + + for (let i = 0; i < data.length; i++) { + const item = data[i] + if (item.type === 'function_call_output' && item.call_id?.toLowerCase() === toolIdLower) { + data[i] = { ...item, output: prunedMessage } + replaced = true + } + } + + return replaced + }, + + hasToolOutputs(data: any[]): boolean { + return data.some((item: any) => item.type === 'function_call_output') + }, + + getLogMetadata(data: any[], replacedCount: number, inputUrl: string): Record { + return { + url: inputUrl, + replacedCount, + totalItems: data.length, + format: 'openai-responses-api' + } + } +} diff --git a/lib/fetch-wrapper/gemini.ts b/lib/fetch-wrapper/gemini.ts deleted file mode 100644 index b3d73fb3..00000000 --- a/lib/fetch-wrapper/gemini.ts +++ /dev/null @@ -1,176 +0,0 @@ -import type { FetchHandlerContext, FetchHandlerResult } from "./types" -import { - PRUNED_CONTENT_MESSAGE, - getAllPrunedIds, - fetchSessionMessages -} from "./types" -import { injectSynthGemini, trackNewToolResultsGemini } from "../api-formats/synth-instruction" -import { buildPrunableToolsList, buildEndInjection, injectPrunableListGemini } from "../api-formats/prunable-list" - -/** - * Handles Google/Gemini format (body.contents array with functionResponse parts). - * Uses position-based correlation since Google's native format doesn't include tool call IDs. - */ -export async function handleGemini( - body: any, - ctx: FetchHandlerContext, - inputUrl: string -): Promise { - if (!body.contents || !Array.isArray(body.contents)) { - return { modified: false, body } - } - - let modified = false - - if (ctx.config.strategies.onTool.length > 0) { - if (injectSynthGemini(body.contents, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { - modified = true - } - - const sessionId = ctx.state.lastSeenSessionId - if (sessionId) { - const toolIds = Array.from(ctx.state.toolParameters.keys()) - const alreadyPruned = ctx.state.prunedIds.get(sessionId) ?? [] - const alreadyPrunedLower = new Set(alreadyPruned.map(id => id.toLowerCase())) - const unprunedIds = toolIds.filter(id => !alreadyPrunedLower.has(id.toLowerCase())) - - const { list: prunableList, numericIds } = buildPrunableToolsList( - sessionId, - unprunedIds, - ctx.state.toolParameters, - ctx.config.protectedTools - ) - - if (prunableList) { - const protectedSet = new Set(ctx.config.protectedTools) - trackNewToolResultsGemini(body.contents, ctx.toolTracker, protectedSet) - const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq - - const endInjection = buildEndInjection(prunableList, includeNudge) - if (injectPrunableListGemini(body.contents, endInjection)) { - ctx.logger.debug("fetch", "Injected prunable tools list (Gemini)", { - ids: numericIds, - nudge: includeNudge, - toolsSincePrune: ctx.toolTracker.toolResultCount - }) - modified = true - } - } - } - } - - const hasFunctionResponses = body.contents.some((content: any) => - Array.isArray(content.parts) && - content.parts.some((part: any) => part.functionResponse) - ) - - if (!hasFunctionResponses) { - return { modified, body } - } - - const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger) - - if (allPrunedIds.size === 0) { - return { modified, body } - } - - const activeSessions = allSessions.data?.filter((s: any) => !s.parentID) || [] - let positionMapping: Map | undefined - - for (const session of activeSessions) { - const mapping = ctx.state.googleToolCallMapping.get(session.id) - if (mapping && mapping.size > 0) { - positionMapping = mapping - break - } - } - - if (!positionMapping) { - ctx.logger.info("fetch", "No Google tool call mapping found, skipping pruning for Gemini format") - return { modified, body } - } - - // Build position counters to track occurrence of each tool name - const toolPositionCounters = new Map() - let replacedCount = 0 - let totalFunctionResponses = 0 - let prunableFunctionResponses = 0 - const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase())) - - body.contents = body.contents.map((content: any) => { - if (!Array.isArray(content.parts)) return content - - let contentModified = false - const newParts = content.parts.map((part: any) => { - if (part.functionResponse) { - totalFunctionResponses++ - const funcName = part.functionResponse.name?.toLowerCase() - - if (!funcName || !protectedToolsLower.has(funcName)) { - prunableFunctionResponses++ - } - - if (funcName) { - const currentIndex = toolPositionCounters.get(funcName) || 0 - toolPositionCounters.set(funcName, currentIndex + 1) - - const positionKey = `${funcName}:${currentIndex}` - const toolCallId = positionMapping!.get(positionKey) - - if (toolCallId && allPrunedIds.has(toolCallId)) { - contentModified = true - replacedCount++ - // Preserve thoughtSignature if present (required for Gemini 3 Pro) - return { - ...part, - functionResponse: { - ...part.functionResponse, - response: { - name: part.functionResponse.name, - content: PRUNED_CONTENT_MESSAGE - } - } - } - } - } - } - return part - }) - - if (contentModified) { - return { ...content, parts: newParts } - } - return content - }) - - if (replacedCount > 0) { - ctx.logger.info("fetch", "Replaced pruned tool outputs (Google/Gemini)", { - replaced: replacedCount, - total: prunableFunctionResponses - }) - - if (ctx.logger.enabled) { - let sessionMessages: any[] | undefined - if (activeSessions.length > 0) { - const mostRecentSession = activeSessions[0] - sessionMessages = await fetchSessionMessages(ctx.client, mostRecentSession.id) - } - - await ctx.logger.saveWrappedContext( - "global", - body.contents, - { - url: inputUrl, - replacedCount, - totalContents: body.contents.length, - format: 'google-gemini' - }, - sessionMessages - ) - } - - return { modified: true, body } - } - - return { modified, body } -} diff --git a/lib/fetch-wrapper/handler.ts b/lib/fetch-wrapper/handler.ts new file mode 100644 index 00000000..72fcc085 --- /dev/null +++ b/lib/fetch-wrapper/handler.ts @@ -0,0 +1,130 @@ +import type { FetchHandlerContext, FetchHandlerResult, FormatDescriptor } from "./types" +import { + PRUNED_CONTENT_MESSAGE, + getAllPrunedIds, + fetchSessionMessages +} from "./types" +import { buildPrunableToolsList, buildEndInjection } from "../api-formats/prunable-list" + +/** + * Generic format handler that processes any API format using a FormatDescriptor. + * + * This consolidates the common logic from all format-specific handlers: + * 1. Cache tool parameters + * 2. Inject synthetic instructions (if strategies enabled) + * 3. Build and inject prunable tools list + * 4. Replace pruned tool outputs + * 5. Log and save context + */ +export async function handleFormat( + body: any, + ctx: FetchHandlerContext, + inputUrl: string, + format: FormatDescriptor +): Promise { + const data = format.getDataArray(body) + if (!data) { + return { modified: false, body } + } + + let modified = false + + // Cache tool parameters for this format + format.cacheToolParameters(data, ctx.state, ctx.logger) + + // Inject synthetic instructions if strategies are enabled + if (ctx.config.strategies.onTool.length > 0) { + if (format.injectSynth(data, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { + modified = true + } + + const sessionId = ctx.state.lastSeenSessionId + if (sessionId) { + const toolIds = Array.from(ctx.state.toolParameters.keys()) + const alreadyPruned = ctx.state.prunedIds.get(sessionId) ?? [] + const alreadyPrunedLower = new Set(alreadyPruned.map(id => id.toLowerCase())) + const unprunedIds = toolIds.filter(id => !alreadyPrunedLower.has(id.toLowerCase())) + + const { list: prunableList, numericIds } = buildPrunableToolsList( + sessionId, + unprunedIds, + ctx.state.toolParameters, + ctx.config.protectedTools + ) + + if (prunableList) { + const protectedSet = new Set(ctx.config.protectedTools) + format.trackNewToolResults(data, ctx.toolTracker, protectedSet) + const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq + + const endInjection = buildEndInjection(prunableList, includeNudge) + if (format.injectPrunableList(data, endInjection)) { + ctx.logger.debug("fetch", `Injected prunable tools list (${format.name})`, { + ids: numericIds, + nudge: includeNudge, + toolsSincePrune: ctx.toolTracker.toolResultCount + }) + modified = true + } + } + } + } + + // Check if there are any tool outputs to potentially prune + if (!format.hasToolOutputs(data)) { + return { modified, body } + } + + // Get all pruned IDs across sessions + const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger) + + if (allPrunedIds.size === 0) { + return { modified, body } + } + + // Extract tool outputs and replace pruned ones + const toolOutputs = format.extractToolOutputs(data, ctx.state) + const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase())) + let replacedCount = 0 + + for (const output of toolOutputs) { + // Skip protected tools + if (output.toolName && protectedToolsLower.has(output.toolName.toLowerCase())) { + continue + } + + if (allPrunedIds.has(output.id)) { + if (format.replaceToolOutput(data, output.id, PRUNED_CONTENT_MESSAGE, ctx.state)) { + replacedCount++ + } + } + } + + if (replacedCount > 0) { + ctx.logger.info("fetch", `Replaced pruned tool outputs (${format.name})`, { + replaced: replacedCount, + total: toolOutputs.length + }) + + // Save context for debugging if logging is enabled + if (ctx.logger.enabled) { + const activeSessions = allSessions.data?.filter((s: any) => !s.parentID) || [] + let sessionMessages: any[] | undefined + if (activeSessions.length > 0) { + const mostRecentSession = activeSessions[0] + sessionMessages = await fetchSessionMessages(ctx.client, mostRecentSession.id) + } + + await ctx.logger.saveWrappedContext( + "global", + data, + format.getLogMetadata(data, replacedCount, inputUrl), + sessionMessages + ) + } + + return { modified: true, body } + } + + return { modified, body } +} diff --git a/lib/fetch-wrapper/index.ts b/lib/fetch-wrapper/index.ts index 561c7174..b72d7f60 100644 --- a/lib/fetch-wrapper/index.ts +++ b/lib/fetch-wrapper/index.ts @@ -3,9 +3,8 @@ import type { Logger } from "../logger" import type { FetchHandlerContext, SynthPrompts } from "./types" import type { ToolTracker } from "../api-formats/synth-instruction" import type { PluginConfig } from "../config" -import { handleOpenAIChatAndAnthropic } from "./openai-chat" -import { handleGemini } from "./gemini" -import { handleOpenAIResponses } from "./openai-responses" +import { openaiChatFormat, openaiResponsesFormat, geminiFormat } from "./formats" +import { handleFormat } from "./handler" import { runStrategies } from "../core/strategies" import { accumulateGCStats } from "./gc-tracker" import { trimToolParametersCache } from "../state/tool-cache" @@ -58,20 +57,20 @@ export function installFetchWrapper( const toolIdsBefore = new Set(state.toolParameters.keys()) // Mutually exclusive format handlers to avoid double-processing - if (body.input && Array.isArray(body.input)) { - const result = await handleOpenAIResponses(body, ctx, inputUrl) + if (openaiResponsesFormat.detect(body)) { + const result = await handleFormat(body, ctx, inputUrl, openaiResponsesFormat) if (result.modified) { modified = true } } - else if (body.messages && Array.isArray(body.messages)) { - const result = await handleOpenAIChatAndAnthropic(body, ctx, inputUrl) + else if (openaiChatFormat.detect(body)) { + const result = await handleFormat(body, ctx, inputUrl, openaiChatFormat) if (result.modified) { modified = true } } - else if (body.contents && Array.isArray(body.contents)) { - const result = await handleGemini(body, ctx, inputUrl) + else if (geminiFormat.detect(body)) { + const result = await handleFormat(body, ctx, inputUrl, geminiFormat) if (result.modified) { modified = true } diff --git a/lib/fetch-wrapper/openai-chat.ts b/lib/fetch-wrapper/openai-chat.ts deleted file mode 100644 index 6488d5a6..00000000 --- a/lib/fetch-wrapper/openai-chat.ts +++ /dev/null @@ -1,165 +0,0 @@ -import type { FetchHandlerContext, FetchHandlerResult } from "./types" -import { - PRUNED_CONTENT_MESSAGE, - getAllPrunedIds, - fetchSessionMessages, - getMostRecentActiveSession -} from "./types" -import { cacheToolParametersFromMessages } from "../state/tool-cache" -import { injectSynth, trackNewToolResults } from "../api-formats/synth-instruction" -import { buildPrunableToolsList, buildEndInjection, injectPrunableList } from "../api-formats/prunable-list" - -/** - * Handles OpenAI Chat Completions format (body.messages with role='tool'). - * Also handles Anthropic format (role='user' with tool_result content parts). - */ -export async function handleOpenAIChatAndAnthropic( - body: any, - ctx: FetchHandlerContext, - inputUrl: string -): Promise { - if (!body.messages || !Array.isArray(body.messages)) { - return { modified: false, body } - } - - cacheToolParametersFromMessages(body.messages, ctx.state, ctx.logger) - - let modified = false - - if (ctx.config.strategies.onTool.length > 0) { - if (injectSynth(body.messages, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { - modified = true - } - - const sessionId = ctx.state.lastSeenSessionId - if (sessionId) { - const toolIds = Array.from(ctx.state.toolParameters.keys()) - const alreadyPruned = ctx.state.prunedIds.get(sessionId) ?? [] - const alreadyPrunedLower = new Set(alreadyPruned.map(id => id.toLowerCase())) - const unprunedIds = toolIds.filter(id => !alreadyPrunedLower.has(id.toLowerCase())) - - const { list: prunableList, numericIds } = buildPrunableToolsList( - sessionId, - unprunedIds, - ctx.state.toolParameters, - ctx.config.protectedTools - ) - - if (prunableList) { - const protectedSet = new Set(ctx.config.protectedTools) - trackNewToolResults(body.messages, ctx.toolTracker, protectedSet) - const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq - - const endInjection = buildEndInjection(prunableList, includeNudge) - if (injectPrunableList(body.messages, endInjection)) { - ctx.logger.debug("fetch", "Injected prunable tools list", { - ids: numericIds, - nudge: includeNudge, - toolsSincePrune: ctx.toolTracker.toolResultCount - }) - modified = true - } - } - } - } - - const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase())) - - const toolMessages = body.messages.filter((m: any) => { - if (m.role === 'tool') return true - if (m.role === 'user' && Array.isArray(m.content)) { - for (const part of m.content) { - if (part.type === 'tool_result') return true - } - } - return false - }) - - let prunableToolCount = 0 - for (const m of body.messages) { - if (m.role === 'tool') { - const toolId = m.tool_call_id?.toLowerCase() - const metadata = toolId ? ctx.state.toolParameters.get(toolId) : undefined - if (!metadata || !protectedToolsLower.has(metadata.tool.toLowerCase())) { - prunableToolCount++ - } - } else if (m.role === 'user' && Array.isArray(m.content)) { - for (const part of m.content) { - if (part.type === 'tool_result') { - const toolId = part.tool_use_id?.toLowerCase() - const metadata = toolId ? ctx.state.toolParameters.get(toolId) : undefined - if (!metadata || !protectedToolsLower.has(metadata.tool.toLowerCase())) { - prunableToolCount++ - } - } - } - } - } - - const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger) - - if (toolMessages.length === 0 || allPrunedIds.size === 0) { - return { modified, body } - } - - let replacedCount = 0 - - body.messages = body.messages.map((m: any) => { - if (m.role === 'tool' && allPrunedIds.has(m.tool_call_id?.toLowerCase())) { - replacedCount++ - return { - ...m, - content: PRUNED_CONTENT_MESSAGE - } - } - - if (m.role === 'user' && Array.isArray(m.content)) { - let messageModified = false - const newContent = m.content.map((part: any) => { - if (part.type === 'tool_result' && allPrunedIds.has(part.tool_use_id?.toLowerCase())) { - messageModified = true - replacedCount++ - return { - ...part, - content: PRUNED_CONTENT_MESSAGE - } - } - return part - }) - if (messageModified) { - return { ...m, content: newContent } - } - } - - return m - }) - - if (replacedCount > 0) { - ctx.logger.info("fetch", "Replaced pruned tool outputs", { - replaced: replacedCount, - total: prunableToolCount - }) - - if (ctx.logger.enabled) { - const mostRecentSession = getMostRecentActiveSession(allSessions) - const sessionMessages = mostRecentSession - ? await fetchSessionMessages(ctx.client, mostRecentSession.id) - : undefined - - await ctx.logger.saveWrappedContext( - "global", - body.messages, - { - url: inputUrl, - replacedCount, - totalMessages: body.messages.length - }, - sessionMessages - ) - } - - return { modified: true, body } - } - - return { modified, body } -} diff --git a/lib/fetch-wrapper/openai-responses.ts b/lib/fetch-wrapper/openai-responses.ts deleted file mode 100644 index 2c71c6d2..00000000 --- a/lib/fetch-wrapper/openai-responses.ts +++ /dev/null @@ -1,129 +0,0 @@ -import type { FetchHandlerContext, FetchHandlerResult } from "./types" -import { - PRUNED_CONTENT_MESSAGE, - getAllPrunedIds, - fetchSessionMessages, - getMostRecentActiveSession -} from "./types" -import { cacheToolParametersFromInput } from "../state/tool-cache" -import { injectSynthResponses, trackNewToolResultsResponses } from "../api-formats/synth-instruction" -import { buildPrunableToolsList, buildEndInjection, injectPrunableListResponses } from "../api-formats/prunable-list" - -/** - * Handles OpenAI Responses API format (body.input array with function_call_output items). - * Used by GPT-5 models via sdk.responses(). - */ -export async function handleOpenAIResponses( - body: any, - ctx: FetchHandlerContext, - inputUrl: string -): Promise { - if (!body.input || !Array.isArray(body.input)) { - return { modified: false, body } - } - - cacheToolParametersFromInput(body.input, ctx.state, ctx.logger) - - let modified = false - - if (ctx.config.strategies.onTool.length > 0) { - if (injectSynthResponses(body.input, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { - modified = true - } - - const sessionId = ctx.state.lastSeenSessionId - if (sessionId) { - const toolIds = Array.from(ctx.state.toolParameters.keys()) - const alreadyPruned = ctx.state.prunedIds.get(sessionId) ?? [] - const alreadyPrunedLower = new Set(alreadyPruned.map(id => id.toLowerCase())) - const unprunedIds = toolIds.filter(id => !alreadyPrunedLower.has(id.toLowerCase())) - - const { list: prunableList, numericIds } = buildPrunableToolsList( - sessionId, - unprunedIds, - ctx.state.toolParameters, - ctx.config.protectedTools - ) - - if (prunableList) { - const protectedSet = new Set(ctx.config.protectedTools) - trackNewToolResultsResponses(body.input, ctx.toolTracker, protectedSet) - const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq - - const endInjection = buildEndInjection(prunableList, includeNudge) - if (injectPrunableListResponses(body.input, endInjection)) { - ctx.logger.debug("fetch", "Injected prunable tools list (Responses API)", { - ids: numericIds, - nudge: includeNudge, - toolsSincePrune: ctx.toolTracker.toolResultCount - }) - modified = true - } - } - } - } - - const functionOutputs = body.input.filter((item: any) => item.type === 'function_call_output') - - if (functionOutputs.length === 0) { - return { modified, body } - } - - const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger) - - if (allPrunedIds.size === 0) { - return { modified, body } - } - - const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase())) - let prunableFunctionOutputCount = 0 - for (const item of functionOutputs) { - const toolName = item.name?.toLowerCase() - if (!toolName || !protectedToolsLower.has(toolName)) { - prunableFunctionOutputCount++ - } - } - - let replacedCount = 0 - - body.input = body.input.map((item: any) => { - if (item.type === 'function_call_output' && allPrunedIds.has(item.call_id?.toLowerCase())) { - replacedCount++ - return { - ...item, - output: PRUNED_CONTENT_MESSAGE - } - } - return item - }) - - if (replacedCount > 0) { - ctx.logger.info("fetch", "Replaced pruned tool outputs (Responses API)", { - replaced: replacedCount, - total: prunableFunctionOutputCount - }) - - if (ctx.logger.enabled) { - const mostRecentSession = getMostRecentActiveSession(allSessions) - const sessionMessages = mostRecentSession - ? await fetchSessionMessages(ctx.client, mostRecentSession.id) - : undefined - - await ctx.logger.saveWrappedContext( - "global", - body.input, - { - url: inputUrl, - replacedCount, - totalItems: body.input.length, - format: 'openai-responses-api' - }, - sessionMessages - ) - } - - return { modified: true, body } - } - - return { modified, body } -} diff --git a/lib/fetch-wrapper/types.ts b/lib/fetch-wrapper/types.ts index 4510e9df..15bf1f2c 100644 --- a/lib/fetch-wrapper/types.ts +++ b/lib/fetch-wrapper/types.ts @@ -6,6 +6,57 @@ import type { PluginConfig } from "../config" /** The message used to replace pruned tool output content */ export const PRUNED_CONTENT_MESSAGE = '[Output removed to save context - information superseded or no longer needed]' +// ============================================================================ +// Format Descriptor Interface +// ============================================================================ + +/** Represents a tool output that can be pruned */ +export interface ToolOutput { + /** The tool call ID (tool_call_id, call_id, tool_use_id, or position key for Gemini) */ + id: string + /** The tool name (for protected tool checking) */ + toolName?: string +} + +/** + * Describes how to handle a specific API format (OpenAI Chat, Anthropic, Gemini, etc.) + * Each format implements this interface to provide format-specific logic. + */ +export interface FormatDescriptor { + /** Human-readable name for logging */ + name: string + + /** Check if this format matches the request body */ + detect(body: any): boolean + + /** Get the data array to process (messages, contents, input, etc.) */ + getDataArray(body: any): any[] | undefined + + /** Cache tool parameters from the data array */ + cacheToolParameters(data: any[], state: PluginState, logger?: Logger): void + + /** Inject synthetic instruction into the last user message */ + injectSynth(data: any[], instruction: string, nudgeText: string): boolean + + /** Track new tool results for nudge frequency */ + trackNewToolResults(data: any[], tracker: ToolTracker, protectedTools: Set): number + + /** Inject prunable list at end of conversation */ + injectPrunableList(data: any[], injection: string): boolean + + /** Extract all tool outputs from the data for pruning */ + extractToolOutputs(data: any[], state: PluginState): ToolOutput[] + + /** Replace a pruned tool output with the pruned message. Returns true if replaced. */ + replaceToolOutput(data: any[], toolId: string, prunedMessage: string, state: PluginState): boolean + + /** Check if data has any tool outputs worth processing */ + hasToolOutputs(data: any[]): boolean + + /** Get metadata for logging after replacements */ + getLogMetadata(data: any[], replacedCount: number, inputUrl: string): Record +} + /** Prompts used for synthetic instruction injection */ export interface SynthPrompts { synthInstruction: string diff --git a/lib/hooks.ts b/lib/hooks.ts index 77b530a2..1578d41f 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -109,10 +109,21 @@ export function createChatParamsHandler( for (const part of msg.parts) { if (part.type === 'tool' && part.callID && part.tool) { const toolName = part.tool.toLowerCase() + const callId = part.callID.toLowerCase() + if (!toolCallsByName.has(toolName)) { toolCallsByName.set(toolName, []) } - toolCallsByName.get(toolName)!.push(part.callID.toLowerCase()) + toolCallsByName.get(toolName)!.push(callId) + + // Also populate toolParameters for Gemini + // This is needed for buildPrunableToolsList to work + if (!state.toolParameters.has(callId)) { + state.toolParameters.set(callId, { + tool: part.tool, + parameters: part.input ?? {} + }) + } } } } @@ -128,7 +139,8 @@ export function createChatParamsHandler( state.googleToolCallMapping.set(sessionId, positionMapping) logger.info("chat.params", "Built Google tool call mapping", { sessionId: sessionId.substring(0, 8), - toolCount: positionMapping.size + toolCount: positionMapping.size, + toolParamsCount: state.toolParameters.size }) } } catch (error: any) { diff --git a/lib/pruning-tool.ts b/lib/pruning-tool.ts index 891d657b..cfc2dafd 100644 --- a/lib/pruning-tool.ts +++ b/lib/pruning-tool.ts @@ -86,15 +86,14 @@ export function createPruningTool( const toolMetadata = new Map() for (const id of prunedIds) { - // Try both original and lowercase since caching may vary - const meta = state.toolParameters.get(id) || state.toolParameters.get(id.toLowerCase()) + // Look up with lowercase since all IDs are stored lowercase + const meta = state.toolParameters.get(id.toLowerCase()) if (meta) { toolMetadata.set(id.toLowerCase(), meta) } else { logger.debug("prune-tool", "No metadata found for ID", { id, idLower: id.toLowerCase(), - hasOriginal: state.toolParameters.has(id), hasLower: state.toolParameters.has(id.toLowerCase()) }) } diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index 8c27ab17..7aca40c1 100644 --- a/lib/state/tool-cache.ts +++ b/lib/state/tool-cache.ts @@ -32,10 +32,10 @@ export function cacheToolParametersFromMessages( const params = typeof toolCall.function.arguments === 'string' ? JSON.parse(toolCall.function.arguments) : toolCall.function.arguments - state.toolParameters.set(toolCall.id, { - tool: toolCall.function.name, - parameters: params - }) + state.toolParameters.set(toolCall.id.toLowerCase(), { + tool: toolCall.function.name, + parameters: params + }) openaiCached++ } catch (error) { // Silently ignore parse errors @@ -49,10 +49,10 @@ export function cacheToolParametersFromMessages( continue } - state.toolParameters.set(part.id, { - tool: part.name, - parameters: part.input ?? {} - }) + state.toolParameters.set(part.id.toLowerCase(), { + tool: part.name, + parameters: part.input ?? {} + }) anthropicCached++ } } @@ -87,7 +87,7 @@ export function cacheToolParametersFromInput( const params = typeof item.arguments === 'string' ? JSON.parse(item.arguments) : item.arguments - state.toolParameters.set(item.call_id, { + state.toolParameters.set(item.call_id.toLowerCase(), { tool: item.name, parameters: params }) From 86c564df5525171ddd3010e67f6d76dffb0ca217 Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Wed, 3 Dec 2025 14:58:21 -0500 Subject: [PATCH 8/9] chore: remove redundant comments --- lib/fetch-wrapper/formats/gemini.ts | 4 ---- lib/fetch-wrapper/formats/openai-chat.ts | 4 ---- lib/fetch-wrapper/handler.ts | 7 ------- lib/fetch-wrapper/index.ts | 2 +- lib/fetch-wrapper/types.ts | 1 - lib/hooks.ts | 7 +------ lib/pruning-tool.ts | 8 +------- lib/state/tool-cache.ts | 2 -- 8 files changed, 3 insertions(+), 32 deletions(-) diff --git a/lib/fetch-wrapper/formats/gemini.ts b/lib/fetch-wrapper/formats/gemini.ts index 765d9cb9..0eee5d6c 100644 --- a/lib/fetch-wrapper/formats/gemini.ts +++ b/lib/fetch-wrapper/formats/gemini.ts @@ -50,8 +50,6 @@ export const geminiFormat: FormatDescriptor = { extractToolOutputs(data: any[], state: PluginState): ToolOutput[] { const outputs: ToolOutput[] = [] - // We need the position mapping to correlate functionResponses to tool call IDs - // Find the mapping from any active session let positionMapping: Map | undefined for (const [_sessionId, mapping] of state.googleToolCallMapping) { if (mapping && mapping.size > 0) { @@ -64,7 +62,6 @@ export const geminiFormat: FormatDescriptor = { return outputs } - // Track position counters per tool name const toolPositionCounters = new Map() for (const content of data) { @@ -95,7 +92,6 @@ export const geminiFormat: FormatDescriptor = { }, replaceToolOutput(data: any[], toolId: string, prunedMessage: string, state: PluginState): boolean { - // Find the position mapping let positionMapping: Map | undefined for (const [_sessionId, mapping] of state.googleToolCallMapping) { if (mapping && mapping.size > 0) { diff --git a/lib/fetch-wrapper/formats/openai-chat.ts b/lib/fetch-wrapper/formats/openai-chat.ts index b75f9f19..b4810461 100644 --- a/lib/fetch-wrapper/formats/openai-chat.ts +++ b/lib/fetch-wrapper/formats/openai-chat.ts @@ -49,7 +49,6 @@ export const openaiChatFormat: FormatDescriptor = { const outputs: ToolOutput[] = [] for (const m of data) { - // OpenAI Chat format: role='tool' with tool_call_id if (m.role === 'tool' && m.tool_call_id) { const metadata = state.toolParameters.get(m.tool_call_id.toLowerCase()) outputs.push({ @@ -58,7 +57,6 @@ export const openaiChatFormat: FormatDescriptor = { }) } - // Anthropic format: role='user' with content[].type='tool_result' if (m.role === 'user' && Array.isArray(m.content)) { for (const part of m.content) { if (part.type === 'tool_result' && part.tool_use_id) { @@ -82,13 +80,11 @@ export const openaiChatFormat: FormatDescriptor = { for (let i = 0; i < data.length; i++) { const m = data[i] - // OpenAI Chat format if (m.role === 'tool' && m.tool_call_id?.toLowerCase() === toolIdLower) { data[i] = { ...m, content: prunedMessage } replaced = true } - // Anthropic format if (m.role === 'user' && Array.isArray(m.content)) { let messageModified = false const newContent = m.content.map((part: any) => { diff --git a/lib/fetch-wrapper/handler.ts b/lib/fetch-wrapper/handler.ts index 72fcc085..004378dd 100644 --- a/lib/fetch-wrapper/handler.ts +++ b/lib/fetch-wrapper/handler.ts @@ -29,10 +29,8 @@ export async function handleFormat( let modified = false - // Cache tool parameters for this format format.cacheToolParameters(data, ctx.state, ctx.logger) - // Inject synthetic instructions if strategies are enabled if (ctx.config.strategies.onTool.length > 0) { if (format.injectSynth(data, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { modified = true @@ -70,25 +68,21 @@ export async function handleFormat( } } - // Check if there are any tool outputs to potentially prune if (!format.hasToolOutputs(data)) { return { modified, body } } - // Get all pruned IDs across sessions const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger) if (allPrunedIds.size === 0) { return { modified, body } } - // Extract tool outputs and replace pruned ones const toolOutputs = format.extractToolOutputs(data, ctx.state) const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase())) let replacedCount = 0 for (const output of toolOutputs) { - // Skip protected tools if (output.toolName && protectedToolsLower.has(output.toolName.toLowerCase())) { continue } @@ -106,7 +100,6 @@ export async function handleFormat( total: toolOutputs.length }) - // Save context for debugging if logging is enabled if (ctx.logger.enabled) { const activeSessions = allSessions.data?.filter((s: any) => !s.parentID) || [] let sessionMessages: any[] | undefined diff --git a/lib/fetch-wrapper/index.ts b/lib/fetch-wrapper/index.ts index b72d7f60..abcf5adc 100644 --- a/lib/fetch-wrapper/index.ts +++ b/lib/fetch-wrapper/index.ts @@ -56,7 +56,7 @@ export function installFetchWrapper( const toolIdsBefore = new Set(state.toolParameters.keys()) - // Mutually exclusive format handlers to avoid double-processing + // Mutually exclusive format handlers if (openaiResponsesFormat.detect(body)) { const result = await handleFormat(body, ctx, inputUrl, openaiResponsesFormat) if (result.modified) { diff --git a/lib/fetch-wrapper/types.ts b/lib/fetch-wrapper/types.ts index 15bf1f2c..b88bf826 100644 --- a/lib/fetch-wrapper/types.ts +++ b/lib/fetch-wrapper/types.ts @@ -99,7 +99,6 @@ export async function getAllPrunedIds( if (currentSession) { await ensureSessionRestored(state, currentSession.id, logger) const prunedIds = state.prunedIds.get(currentSession.id) ?? [] - // Normalize to lowercase for case-insensitive matching prunedIds.forEach((id: string) => allPrunedIds.add(id.toLowerCase())) if (logger && prunedIds.length > 0) { diff --git a/lib/hooks.ts b/lib/hooks.ts index 1578d41f..e802f532 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -32,7 +32,6 @@ export function createEventHandler( if (await isSubagentSession(client, event.properties.sessionID)) return if (config.strategies.onIdle.length === 0) return - // Skip idle pruning if the last tool used was prune and idle strategies cover tool strategies if (toolTracker?.skipNextIdle) { toolTracker.skipNextIdle = false if (toolStrategiesCoveredByIdle(config.strategies.onIdle, config.strategies.onTool)) { @@ -43,7 +42,6 @@ export function createEventHandler( try { const result = await runOnIdle(janitorCtx, event.properties.sessionID, config.strategies.onIdle) - // Reset nudge counter if idle pruning succeeded and covers tool strategies if (result && result.prunedCount > 0 && toolTracker && config.nudge_freq > 0) { if (toolStrategiesCoveredByIdle(config.strategies.onIdle, config.strategies.onTool)) { resetToolTrackerCount(toolTracker) @@ -91,8 +89,7 @@ export function createChatParamsHandler( } } - // Build Google/Gemini tool call mapping for position-based correlation - // This is needed because Google's native format loses tool call IDs + // Build position-based mapping for Gemini (which loses tool call IDs in native format) if (providerID === 'google' || providerID === 'google-vertex') { try { const messagesResponse = await client.session.messages({ @@ -116,8 +113,6 @@ export function createChatParamsHandler( } toolCallsByName.get(toolName)!.push(callId) - // Also populate toolParameters for Gemini - // This is needed for buildPrunableToolsList to work if (!state.toolParameters.has(callId)) { state.toolParameters.set(callId, { tool: part.tool, diff --git a/lib/pruning-tool.ts b/lib/pruning-tool.ts index cfc2dafd..b75901da 100644 --- a/lib/pruning-tool.ts +++ b/lib/pruning-tool.ts @@ -86,7 +86,6 @@ export function createPruningTool( const toolMetadata = new Map() for (const id of prunedIds) { - // Look up with lowercase since all IDs are stored lowercase const meta = state.toolParameters.get(id.toLowerCase()) if (meta) { toolMetadata.set(id.toLowerCase(), meta) @@ -144,7 +143,6 @@ async function calculateTokensSaved( }) const messages = messagesResponse.data || messagesResponse - // Build map of tool call ID -> output content const toolOutputs = new Map() for (const msg of messages) { if (msg.role === 'tool' && msg.tool_call_id) { @@ -153,7 +151,6 @@ async function calculateTokensSaved( : JSON.stringify(msg.content) toolOutputs.set(msg.tool_call_id.toLowerCase(), content) } - // Handle Anthropic format if (msg.role === 'user' && Array.isArray(msg.content)) { for (const part of msg.content) { if (part.type === 'tool_result' && part.tool_use_id) { @@ -166,7 +163,6 @@ async function calculateTokensSaved( } } - // Collect content for pruned outputs const contents: string[] = [] for (const id of prunedIds) { const content = toolOutputs.get(id.toLowerCase()) @@ -176,14 +172,12 @@ async function calculateTokensSaved( } if (contents.length === 0) { - return prunedIds.length * 500 // fallback estimate + return prunedIds.length * 500 } - // Estimate tokens const tokenCounts = await estimateTokensBatch(contents) return tokenCounts.reduce((sum, count) => sum + count, 0) } catch (error: any) { - // If we can't calculate, estimate based on average return prunedIds.length * 500 } } diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index 7aca40c1..f0ae3c6e 100644 --- a/lib/state/tool-cache.ts +++ b/lib/state/tool-cache.ts @@ -38,7 +38,6 @@ export function cacheToolParametersFromMessages( }) openaiCached++ } catch (error) { - // Silently ignore parse errors } } } @@ -93,7 +92,6 @@ export function cacheToolParametersFromInput( }) cached++ } catch (error) { - // Silently ignore parse errors } } From 80220dd6145189969dd30e01523e0c56c83e9455 Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Wed, 3 Dec 2025 15:23:42 -0500 Subject: [PATCH 9/9] fix: restore onIdle LLM analysis and remove dead code - Restore model-selector.ts for LLM model selection with provider fallback - Restore runLlmAnalysis() and replacePrunedToolOutputs() in janitor.ts - Restore buildAnalysisPrompt() and minimizeMessages() in prompt.ts - Restore pruning.txt prompt template - Restore ModelInfo interface and state.model map for model caching - Add model caching back to hooks.ts chat.params handler - Remove dead code: getNumericIdsForActual, getNumericId, getAllMappings, hasMapping, getNextId --- lib/api-formats/prunable-list.ts | 19 +--- lib/core/janitor.ts | 162 +++++++++++++++++++++++++++- lib/core/prompt.ts | 127 ++++++++++++++++++++++ lib/hooks.ts | 9 ++ lib/model-selector.ts | 175 +++++++++++++++++++++++++++++++ lib/prompts/pruning.txt | 30 ++++++ lib/state/id-mapping.ts | 19 ---- lib/state/index.ts | 7 ++ 8 files changed, 508 insertions(+), 40 deletions(-) create mode 100644 lib/model-selector.ts create mode 100644 lib/prompts/pruning.txt diff --git a/lib/api-formats/prunable-list.ts b/lib/api-formats/prunable-list.ts index 65869050..ee580e49 100644 --- a/lib/api-formats/prunable-list.ts +++ b/lib/api-formats/prunable-list.ts @@ -11,7 +11,7 @@ */ import { extractParameterKey } from '../ui/display-utils' -import { getOrCreateNumericId, getNumericId } from '../state/id-mapping' +import { getOrCreateNumericId } from '../state/id-mapping' export interface ToolMetadata { tool: string @@ -107,23 +107,6 @@ export function buildEndInjection( return parts.join('\n\n') } -/** - * Gets the numeric IDs for a list of actual tool call IDs. - * Used when the prune tool needs to show what was pruned. - */ -export function getNumericIdsForActual( - sessionId: string, - actualIds: string[] -): number[] { - return actualIds - .map(id => getNumericId(sessionId, id)) - .filter((id): id is number => id !== undefined) -} - -// ============================================================================ -// Injection Functions -// ============================================================================ - // ============================================================================ // OpenAI Chat / Anthropic Format // ============================================================================ diff --git a/lib/core/janitor.ts b/lib/core/janitor.ts index 5ce1e6f2..5667c21c 100644 --- a/lib/core/janitor.ts +++ b/lib/core/janitor.ts @@ -1,6 +1,9 @@ +import { z } from "zod" import type { Logger } from "../logger" import type { PruningStrategy } from "../config" import type { PluginState } from "../state" +import { buildAnalysisPrompt } from "./prompt" +import { selectModel, extractModelFromSession } from "../model-selector" import { estimateTokensBatch, formatTokenCount } from "../tokenizer" import { saveSessionState } from "../state/persistence" import { ensureSessionRestored } from "../state" @@ -142,9 +145,21 @@ async function runWithStrategies( return !metadata || !config.protectedTools.includes(metadata.tool) }).length - // For onIdle, we currently don't have AI analysis implemented - // This is a placeholder for future idle pruning strategies - const llmPrunedIds: string[] = [] + // PHASE 1: LLM ANALYSIS + let llmPrunedIds: string[] = [] + + if (strategies.includes('ai-analysis') && unprunedToolCallIds.length > 0) { + llmPrunedIds = await runLlmAnalysis( + ctx, + sessionID, + sessionInfo, + messages, + unprunedToolCallIds, + alreadyPrunedIds, + toolMetadata, + options + ) + } const finalNewlyPrunedIds = llmPrunedIds.filter(id => !alreadyPrunedIds.includes(id)) @@ -237,6 +252,147 @@ async function runWithStrategies( } } +// ============================================================================ +// LLM Analysis +// ============================================================================ + +async function runLlmAnalysis( + ctx: JanitorContext, + sessionID: string, + sessionInfo: any, + messages: any[], + unprunedToolCallIds: string[], + alreadyPrunedIds: string[], + toolMetadata: Map, + options: PruningOptions +): Promise { + const { client, state, logger, config } = ctx + + const protectedToolCallIds: string[] = [] + const prunableToolCallIds = unprunedToolCallIds.filter(id => { + const metadata = toolMetadata.get(id) + if (metadata && config.protectedTools.includes(metadata.tool)) { + protectedToolCallIds.push(id) + return false + } + return true + }) + + if (prunableToolCallIds.length === 0) { + return [] + } + + const cachedModelInfo = state.model.get(sessionID) + const sessionModelInfo = extractModelFromSession(sessionInfo, logger) + const currentModelInfo = cachedModelInfo || sessionModelInfo + + const modelSelection = await selectModel(currentModelInfo, logger, config.model, config.workingDirectory) + + logger.info("janitor", `Model: ${modelSelection.modelInfo.providerID}/${modelSelection.modelInfo.modelID}`, { + source: modelSelection.source + }) + + if (modelSelection.failedModel && config.showModelErrorToasts) { + const skipAi = modelSelection.source === 'fallback' && config.strictModelSelection + try { + await client.tui.showToast({ + body: { + title: skipAi ? "DCP: AI analysis skipped" : "DCP: Model fallback", + message: skipAi + ? `${modelSelection.failedModel.providerID}/${modelSelection.failedModel.modelID} failed\nAI analysis skipped (strictModelSelection enabled)` + : `${modelSelection.failedModel.providerID}/${modelSelection.failedModel.modelID} failed\nUsing ${modelSelection.modelInfo.providerID}/${modelSelection.modelInfo.modelID}`, + variant: "info", + duration: 5000 + } + }) + } catch (toastError: any) { + // Ignore toast errors + } + } + + if (modelSelection.source === 'fallback' && config.strictModelSelection) { + logger.info("janitor", "Skipping AI analysis (fallback model, strictModelSelection enabled)") + return [] + } + + const { generateObject } = await import('ai') + + const sanitizedMessages = replacePrunedToolOutputs(messages, alreadyPrunedIds) + + const analysisPrompt = buildAnalysisPrompt( + prunableToolCallIds, + sanitizedMessages, + alreadyPrunedIds, + protectedToolCallIds, + options.reason + ) + + await logger.saveWrappedContext( + "janitor-shadow", + [{ role: "user", content: analysisPrompt }], + { + sessionID, + modelProvider: modelSelection.modelInfo.providerID, + modelID: modelSelection.modelInfo.modelID, + candidateToolCount: prunableToolCallIds.length, + alreadyPrunedCount: alreadyPrunedIds.length, + protectedToolCount: protectedToolCallIds.length, + trigger: options.trigger, + reason: options.reason + } + ) + + const result = await generateObject({ + model: modelSelection.model, + schema: z.object({ + pruned_tool_call_ids: z.array(z.string()), + reasoning: z.string(), + }), + prompt: analysisPrompt + }) + + const rawLlmPrunedIds = result.object.pruned_tool_call_ids + const llmPrunedIds = rawLlmPrunedIds.filter(id => + prunableToolCallIds.includes(id.toLowerCase()) + ) + + if (llmPrunedIds.length > 0) { + const reasoning = result.object.reasoning.replace(/\n+/g, ' ').replace(/\s+/g, ' ').trim() + logger.info("janitor", `LLM reasoning: ${reasoning.substring(0, 200)}${reasoning.length > 200 ? '...' : ''}`) + } + + return llmPrunedIds +} + +function replacePrunedToolOutputs(messages: any[], prunedIds: string[]): any[] { + if (prunedIds.length === 0) return messages + + const prunedIdsSet = new Set(prunedIds.map(id => id.toLowerCase())) + + return messages.map(msg => { + if (!msg.parts) return msg + + return { + ...msg, + parts: msg.parts.map((part: any) => { + if (part.type === 'tool' && + part.callID && + prunedIdsSet.has(part.callID.toLowerCase()) && + part.state?.output) { + return { + ...part, + state: { + ...part.state, + output: '[Output removed to save context - information superseded or no longer needed]' + } + } + } + return part + }) + } + }) +} + // ============================================================================ // Message parsing // ============================================================================ diff --git a/lib/core/prompt.ts b/lib/core/prompt.ts index 39377276..e7f44d4a 100644 --- a/lib/core/prompt.ts +++ b/lib/core/prompt.ts @@ -11,3 +11,130 @@ export function loadPrompt(name: string, vars?: Record): string } return content } + +function minimizeMessages(messages: any[], alreadyPrunedIds?: string[], protectedToolCallIds?: string[]): any[] { + const prunedIdsSet = alreadyPrunedIds ? new Set(alreadyPrunedIds.map(id => id.toLowerCase())) : new Set() + const protectedIdsSet = protectedToolCallIds ? new Set(protectedToolCallIds.map(id => id.toLowerCase())) : new Set() + + return messages.map(msg => { + const minimized: any = { + role: msg.info?.role + } + + if (msg.parts) { + minimized.parts = msg.parts + .filter((part: any) => { + if (part.type === 'step-start' || part.type === 'step-finish') { + return false + } + return true + }) + .map((part: any) => { + if (part.type === 'text') { + if (part.ignored) { + return null + } + return { + type: 'text', + text: part.text + } + } + + // TODO: This should use the opencode normalized system instead of per provider settings + if (part.type === 'reasoning') { + // Calculate encrypted content size if present + let encryptedContentLength = 0 + if (part.metadata?.openai?.reasoningEncryptedContent) { + encryptedContentLength = part.metadata.openai.reasoningEncryptedContent.length + } else if (part.metadata?.anthropic?.signature) { + encryptedContentLength = part.metadata.anthropic.signature.length + } else if (part.metadata?.google?.thoughtSignature) { + encryptedContentLength = part.metadata.google.thoughtSignature.length + } + + return { + type: 'reasoning', + text: part.text, + textLength: part.text?.length || 0, + encryptedContentLength, + ...(part.time && { time: part.time }), + ...(part.metadata && { metadataKeys: Object.keys(part.metadata) }) + } + } + + if (part.type === 'tool') { + const callIDLower = part.callID?.toLowerCase() + const isAlreadyPruned = prunedIdsSet.has(callIDLower) + const isProtected = protectedIdsSet.has(callIDLower) + + let displayCallID = part.callID + if (isAlreadyPruned) { + displayCallID = '' + } else if (isProtected) { + displayCallID = '' + } + + const toolPart: any = { + type: 'tool', + toolCallID: displayCallID, + tool: part.tool + } + + if (part.state?.output) { + toolPart.output = part.state.output + } + + if (part.state?.input) { + const input = part.state.input + + if (input.filePath && (part.tool === 'write' || part.tool === 'edit' || part.tool === 'multiedit' || part.tool === 'patch')) { + toolPart.input = input + } + else if (input.filePath) { + toolPart.input = { filePath: input.filePath } + } + else if (input.tool_calls && Array.isArray(input.tool_calls)) { + toolPart.input = { + batch_summary: `${input.tool_calls.length} tool calls`, + tools: input.tool_calls.map((tc: any) => tc.tool) + } + } + else { + toolPart.input = input + } + } + + return toolPart + } + + return null + }) + .filter(Boolean) + } + + return minimized + }).filter(msg => { + return msg.parts && msg.parts.length > 0 + }) +} + +export function buildAnalysisPrompt( + unprunedToolCallIds: string[], + messages: any[], + alreadyPrunedIds?: string[], + protectedToolCallIds?: string[], + reason?: string +): string { + const minimizedMessages = minimizeMessages(messages, alreadyPrunedIds, protectedToolCallIds) + const messagesJson = JSON.stringify(minimizedMessages, null, 2).replace(/\\n/g, '\n') + + const reasonContext = reason + ? `\nContext: The AI has requested pruning with the following reason: "${reason}"\nUse this context to inform your decisions about what is most relevant to keep.` + : '' + + return loadPrompt("pruning", { + reason_context: reasonContext, + available_tool_call_ids: unprunedToolCallIds.join(", "), + session_history: messagesJson + }) +} diff --git a/lib/hooks.ts b/lib/hooks.ts index e802f532..b2e461e8 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -65,6 +65,7 @@ export function createChatParamsHandler( return async (input: any, _output: any) => { const sessionId = input.sessionID let providerID = (input.provider as any)?.info?.id || input.provider?.id + const modelID = input.model?.id if (!providerID && input.message?.model?.providerID) { providerID = input.message.model.providerID @@ -89,6 +90,14 @@ export function createChatParamsHandler( } } + // Cache model info for the session (used by janitor for model selection) + if (providerID && modelID) { + state.model.set(sessionId, { + providerID: providerID, + modelID: modelID + }) + } + // Build position-based mapping for Gemini (which loses tool call IDs in native format) if (providerID === 'google' || providerID === 'google-vertex') { try { diff --git a/lib/model-selector.ts b/lib/model-selector.ts new file mode 100644 index 00000000..e0e9895e --- /dev/null +++ b/lib/model-selector.ts @@ -0,0 +1,175 @@ +import type { LanguageModel } from 'ai'; +import type { Logger } from './logger'; + +export interface ModelInfo { + providerID: string; + modelID: string; +} + +export const FALLBACK_MODELS: Record = { + openai: 'gpt-5-mini', + anthropic: 'claude-haiku-4-5', //This model isn't broken in opencode-auth-provider + google: 'gemini-2.5-flash', + deepseek: 'deepseek-chat', + xai: 'grok-4-fast', + alibaba: 'qwen3-coder-flash', + zai: 'glm-4.5-flash', + opencode: 'big-pickle' +}; + +const PROVIDER_PRIORITY = [ + 'openai', + 'anthropic', + 'google', + 'deepseek', + 'xai', + 'alibaba', + 'zai', + 'opencode' +]; + +// TODO: some anthropic provided models aren't supported by the opencode-auth-provider package, so this provides a temporary workaround +const SKIP_PROVIDERS = ['github-copilot', 'anthropic']; + +export interface ModelSelectionResult { + model: LanguageModel; + modelInfo: ModelInfo; + source: 'user-model' | 'config' | 'fallback'; + reason?: string; + failedModel?: ModelInfo; +} + +function shouldSkipProvider(providerID: string): boolean { + const normalized = providerID.toLowerCase().trim(); + return SKIP_PROVIDERS.some(skip => normalized.includes(skip.toLowerCase())); +} + +async function importOpencodeAI(logger?: Logger, maxRetries: number = 3, delayMs: number = 100, workspaceDir?: string): Promise { + let lastError: Error | undefined; + + for (let attempt = 1; attempt <= maxRetries; attempt++) { + try { + const { OpencodeAI } = await import('@tarquinen/opencode-auth-provider'); + return new OpencodeAI({ workspaceDir }); + } catch (error: any) { + lastError = error; + + if (error.message?.includes('before initialization')) { + logger?.debug('model-selector', `Import attempt ${attempt}/${maxRetries} failed, will retry`, { + error: error.message + }); + + if (attempt < maxRetries) { + await new Promise(resolve => setTimeout(resolve, delayMs * attempt)); + continue; + } + } + + throw error; + } + } + + throw lastError; +} + +export async function selectModel( + currentModel?: ModelInfo, + logger?: Logger, + configModel?: string, + workspaceDir?: string +): Promise { + const opencodeAI = await importOpencodeAI(logger, 3, 100, workspaceDir); + + let failedModelInfo: ModelInfo | undefined; + + if (configModel) { + const parts = configModel.split('/'); + if (parts.length !== 2) { + logger?.warn('model-selector', 'Invalid config model format', { configModel }); + } else { + const [providerID, modelID] = parts; + + try { + const model = await opencodeAI.getLanguageModel(providerID, modelID); + return { + model, + modelInfo: { providerID, modelID }, + source: 'config', + reason: 'Using model specified in dcp.jsonc config' + }; + } catch (error: any) { + logger?.warn('model-selector', `Config model failed: ${providerID}/${modelID}`, { + error: error.message + }); + failedModelInfo = { providerID, modelID }; + } + } + } + + if (currentModel) { + if (shouldSkipProvider(currentModel.providerID)) { + if (!failedModelInfo) { + failedModelInfo = currentModel; + } + } else { + try { + const model = await opencodeAI.getLanguageModel(currentModel.providerID, currentModel.modelID); + return { + model, + modelInfo: currentModel, + source: 'user-model', + reason: 'Using current session model' + }; + } catch (error: any) { + if (!failedModelInfo) { + failedModelInfo = currentModel; + } + } + } + } + + const providers = await opencodeAI.listProviders(); + + for (const providerID of PROVIDER_PRIORITY) { + if (!providers[providerID]) continue; + + const fallbackModelID = FALLBACK_MODELS[providerID]; + if (!fallbackModelID) continue; + + try { + const model = await opencodeAI.getLanguageModel(providerID, fallbackModelID); + return { + model, + modelInfo: { providerID, modelID: fallbackModelID }, + source: 'fallback', + reason: `Using ${providerID}/${fallbackModelID}`, + failedModel: failedModelInfo + }; + } catch (error: any) { + continue; + } + } + + throw new Error('No available models for analysis. Please authenticate with at least one provider.'); +} + +export function extractModelFromSession(sessionState: any, logger?: Logger): ModelInfo | undefined { + if (sessionState?.model?.providerID && sessionState?.model?.modelID) { + return { + providerID: sessionState.model.providerID, + modelID: sessionState.model.modelID + }; + } + + if (sessionState?.messages && Array.isArray(sessionState.messages)) { + const lastMessage = sessionState.messages[sessionState.messages.length - 1]; + if (lastMessage?.model?.providerID && lastMessage?.model?.modelID) { + return { + providerID: lastMessage.model.providerID, + modelID: lastMessage.model.modelID + }; + } + } + + return undefined; +} diff --git a/lib/prompts/pruning.txt b/lib/prompts/pruning.txt new file mode 100644 index 00000000..49e1e82e --- /dev/null +++ b/lib/prompts/pruning.txt @@ -0,0 +1,30 @@ +You are a conversation analyzer that identifies obsolete tool outputs in a coding session. +{{reason_context}} +Your task: Analyze the session history and identify tool call IDs whose outputs are NO LONGER RELEVANT to the current conversation context. + +Guidelines for identifying obsolete tool calls: +1. Exploratory reads that didn't lead to actual edits or meaningful discussion AND were not explicitly requested to be retained +2. Tool outputs from debugging/fixing an error that has now been resolved +3. Failed or incorrect tool attempts that were immediately corrected (e.g., reading a file from the wrong path, then reading from the correct path) + +DO NOT prune: +- Tool calls whose outputs are actively being discussed +- Tool calls that produced errors still being debugged +- Tool calls that are the MOST RECENT activity in the conversation (these may be intended for future use) + +IMPORTANT: Available tool call IDs for analysis: {{available_tool_call_ids}} + +The session history below may contain tool calls with IDs not in the available list above, these cannot be pruned. These are either: +1. Protected tools (marked with toolCallID "") +2. Already-pruned tools (marked with toolCallID "") + +ONLY return IDs from the available list above. + +Session history (each tool call has a "toolCallID" field): +{{session_history}} + +You MUST respond with valid JSON matching this exact schema: +{ + "pruned_tool_call_ids": ["id1", "id2", ...], + "reasoning": "explanation of why these IDs were selected" +} diff --git a/lib/state/id-mapping.ts b/lib/state/id-mapping.ts index 3421dbc4..fccac6d3 100644 --- a/lib/state/id-mapping.ts +++ b/lib/state/id-mapping.ts @@ -59,20 +59,6 @@ export function getActualId(sessionId: string, numericId: number): string | unde return mapping?.numericToActual.get(numericId) } -export function getNumericId(sessionId: string, actualId: string): number | undefined { - const mapping = sessionMappings.get(sessionId) - return mapping?.actualToNumeric.get(actualId) -} - -export function getAllMappings(sessionId: string): Map { - const mapping = sessionMappings.get(sessionId) - return mapping?.numericToActual ?? new Map() -} - -export function hasMapping(sessionId: string): boolean { - return sessionMappings.has(sessionId) -} - export function clearSessionMapping(sessionId: string): void { sessionMappings.delete(sessionId) } @@ -80,8 +66,3 @@ export function clearSessionMapping(sessionId: string): void { export function clearAllMappings(): void { sessionMappings.clear() } - -export function getNextId(sessionId: string): number { - const mapping = sessionMappings.get(sessionId) - return mapping?.nextId ?? 1 -} diff --git a/lib/state/index.ts b/lib/state/index.ts index 03d9e0ef..2808cb44 100644 --- a/lib/state/index.ts +++ b/lib/state/index.ts @@ -7,6 +7,7 @@ export interface PluginState { stats: Map gcPending: Map toolParameters: Map + model: Map googleToolCallMapping: Map> restoredSessions: Set checkedSessions: Set @@ -19,12 +20,18 @@ export interface ToolParameterEntry { parameters: any } +export interface ModelInfo { + providerID: string + modelID: string +} + export function createPluginState(): PluginState { return { prunedIds: new Map(), stats: new Map(), gcPending: new Map(), toolParameters: new Map(), + model: new Map(), googleToolCallMapping: new Map(), restoredSessions: new Set(), checkedSessions: new Set(),