diff --git a/index.ts b/index.ts index 1ccb82e..a97c7b5 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 } @@ -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 0000000..ee580e4 --- /dev/null +++ b/lib/api-formats/prunable-list.ts @@ -0,0 +1,159 @@ +/** + * 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 } 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') +} + +// ============================================================================ +// 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 aa2dd62..7301c17 100644 --- a/lib/api-formats/synth-instruction.ts +++ b/lib/api-formats/synth-instruction.ts @@ -9,83 +9,94 @@ 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 } -/** Adapter interface for format-specific message operations */ -interface MessageFormatAdapter { - countToolResults(messages: any[], tracker: ToolTracker): number - appendNudge(messages: any[], nudgeText: string): void -} - -/** 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 - } - return false -} - -// ============================================================================ -// 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) +/** + * Track new tool results in OpenAI/Anthropic messages. + * 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, 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) + const toolName = tracker.getToolName?.(m.tool_call_id) + if (!toolName || !protectedTools.has(toolName)) { + tracker.toolResultCount++ 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) + } + } 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) + const toolName = tracker.getToolName?.(part.tool_use_id) + if (!toolName || !protectedTools.has(toolName)) { + tracker.toolResultCount++ 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 }) } + return newCount +} + +/** + * 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 trackNewToolResultsGemini(contents: any[], tracker: ToolTracker, protectedTools: Set): 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) { + const positionId = `gemini_pos_${positionCounter}` + positionCounter++ + if (!tracker.seenToolResultIds.has(positionId)) { + tracker.seenToolResultIds.add(positionId) + const toolName = part.functionResponse.name + if (!toolName || !protectedTools.has(toolName)) { + tracker.toolResultCount++ + newCount++ + } + } + } + } + } + return newCount +} + +/** + * 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, 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) + const toolName = tracker.getToolName?.(item.call_id) + if (!toolName || !protectedTools.has(toolName)) { + tracker.toolResultCount++ + newCount++ + } + } + } + } + return newCount } -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') { return msg.content === nudgeText @@ -116,42 +127,6 @@ export function injectSynth(messages: any[], instruction: string, nudgeText: str return false } -// ============================================================================ -// 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) { const part = content.parts[0] @@ -178,39 +153,6 @@ export function injectSynthGemini(contents: any[], instruction: string, nudgeTex return false } -// ============================================================================ -// 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') { return item.content === nudgeText diff --git a/lib/core/janitor.ts b/lib/core/janitor.ts index 70005d8..5667c21 100644 --- a/lib/core/janitor.ts +++ b/lib/core/janitor.ts @@ -84,6 +84,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 +96,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( @@ -172,7 +167,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 +212,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) @@ -369,6 +364,35 @@ async function runLlmAnalysis( 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 // ============================================================================ @@ -379,7 +403,7 @@ interface ParsedMessages { toolMetadata: Map } -function parseMessages( +export function parseMessages( messages: any[], toolParametersCache: Map ): ParsedMessages { @@ -428,40 +452,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/strategies/index.ts b/lib/core/strategies/index.ts index 060bf64..c6c9128 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/formats/gemini.ts b/lib/fetch-wrapper/formats/gemini.ts new file mode 100644 index 0000000..0eee5d6 --- /dev/null +++ b/lib/fetch-wrapper/formats/gemini.ts @@ -0,0 +1,169 @@ +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[] = [] + + let positionMapping: Map | undefined + for (const [_sessionId, mapping] of state.googleToolCallMapping) { + if (mapping && mapping.size > 0) { + positionMapping = mapping + break + } + } + + if (!positionMapping) { + return outputs + } + + 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 { + 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 0000000..0132c87 --- /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 0000000..b481046 --- /dev/null +++ b/lib/fetch-wrapper/formats/openai-chat.ts @@ -0,0 +1,126 @@ +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) { + 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 + }) + } + + 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] + + if (m.role === 'tool' && m.tool_call_id?.toLowerCase() === toolIdLower) { + data[i] = { ...m, content: prunedMessage } + replaced = true + } + + 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 0000000..67ac5b7 --- /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 abc1bd6..0000000 --- a/lib/fetch-wrapper/gemini.ts +++ /dev/null @@ -1,159 +0,0 @@ -import type { FetchHandlerContext, FetchHandlerResult } from "./types" -import { - PRUNED_CONTENT_MESSAGE, - getAllPrunedIds, - fetchSessionMessages -} from "./types" -import { injectNudgeGemini, injectSynthGemini } from "../api-formats/synth-instruction" - -/** - * 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 - - // 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 - } - } - - if (skipIdleBefore && !ctx.toolTracker.skipNextIdle) { - ctx.logger.debug("fetch", "skipNextIdle was reset by new tool results (Gemini)") - } - - if (injectSynthGemini(body.contents, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { - // ctx.logger.info("fetch", "Injected synthetic instruction (Gemini)") - modified = true - } - } - - // 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) - ) - - if (!hasFunctionResponses) { - return { modified, body } - } - - const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger) - - if (allPrunedIds.size === 0) { - 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 - - 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 - - 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) { - // 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) - - if (toolCallId && allPrunedIds.has(toolCallId)) { - 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: { - ...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: totalFunctionResponses - }) - - 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 0000000..004378d --- /dev/null +++ b/lib/fetch-wrapper/handler.ts @@ -0,0 +1,123 @@ +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 + + format.cacheToolParameters(data, ctx.state, ctx.logger) + + 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 + } + } + } + } + + if (!format.hasToolOutputs(data)) { + return { modified, body } + } + + const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger) + + if (allPrunedIds.size === 0) { + return { modified, body } + } + + 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) { + 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 + }) + + 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 868c9b8..abcf5ad 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" @@ -55,29 +54,23 @@ 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 in order - // OpenAI Chat Completions & Anthropic style (body.messages) - if (body.messages && Array.isArray(body.messages)) { - const result = await handleOpenAIChatAndAnthropic(body, ctx, inputUrl) + // Mutually exclusive format handlers + if (openaiResponsesFormat.detect(body)) { + const result = await handleFormat(body, ctx, inputUrl, openaiResponsesFormat) if (result.modified) { modified = true } } - - // Google/Gemini style (body.contents) - if (body.contents && Array.isArray(body.contents)) { - const result = await handleGemini(body, ctx, inputUrl) + else if (openaiChatFormat.detect(body)) { + const result = await handleFormat(body, ctx, inputUrl, openaiChatFormat) 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) + 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 78b522e..0000000 --- a/lib/fetch-wrapper/openai-chat.ts +++ /dev/null @@ -1,132 +0,0 @@ -import type { FetchHandlerContext, FetchHandlerResult } from "./types" -import { - PRUNED_CONTENT_MESSAGE, - getAllPrunedIds, - fetchSessionMessages, - getMostRecentActiveSession -} from "./types" -import { cacheToolParametersFromMessages } from "../state/tool-cache" -import { injectNudge, injectSynth } from "../api-formats/synth-instruction" - -/** - * 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 } - } - - // Cache tool parameters from messages - cacheToolParametersFromMessages(body.messages, ctx.state) - - 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 - } - } - - if (skipIdleBefore && !ctx.toolTracker.skipNextIdle) { - ctx.logger.debug("fetch", "skipNextIdle was reset by new tool results") - } - - if (injectSynth(body.messages, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { - // ctx.logger.info("fetch", "Injected synthetic instruction") - modified = true - } - } - - // Check for tool messages in both formats: - // 1. OpenAI style: role === 'tool' - // 2. Anthropic style: role === 'user' with content containing tool_result - 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 - }) - - 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) => { - // OpenAI style: role === 'tool' with tool_call_id - if (m.role === 'tool' && allPrunedIds.has(m.tool_call_id?.toLowerCase())) { - replacedCount++ - return { - ...m, - content: PRUNED_CONTENT_MESSAGE - } - } - - // 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) => { - 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: toolMessages.length - }) - - 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 b8a1dbd..0000000 --- a/lib/fetch-wrapper/openai-responses.ts +++ /dev/null @@ -1,106 +0,0 @@ -import type { FetchHandlerContext, FetchHandlerResult } from "./types" -import { - PRUNED_CONTENT_MESSAGE, - getAllPrunedIds, - fetchSessionMessages, - getMostRecentActiveSession -} from "./types" -import { cacheToolParametersFromInput } from "../state/tool-cache" -import { injectNudgeResponses, injectSynthResponses } from "../api-formats/synth-instruction" - -/** - * 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 } - } - - // Cache tool parameters from input - cacheToolParametersFromInput(body.input, ctx.state) - - 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 - } - } - - if (skipIdleBefore && !ctx.toolTracker.skipNextIdle) { - ctx.logger.debug("fetch", "skipNextIdle was reset by new tool results (Responses API)") - } - - if (injectSynthResponses(body.input, ctx.prompts.synthInstruction, ctx.prompts.nudgeInstruction)) { - // ctx.logger.info("fetch", "Injected synthetic instruction (Responses API)") - modified = true - } - } - - // Check for function_call_output items - 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 } - } - - 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: functionOutputs.length - }) - - 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 d6cf4ab..b88bf82 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 @@ -48,7 +99,14 @@ 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)) + 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 dac0b54..b2e461e 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 { @@ -31,8 +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 the same work as 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) @@ -73,10 +71,17 @@ export function createChatParamsHandler( providerID = input.message.model.providerID } - // Track the last seen session ID for fetch wrapper correlation + 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) + }) + clearAllMappings() + state.toolParameters.clear() + } + 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) @@ -85,7 +90,7 @@ export function createChatParamsHandler( } } - // Cache model info for the session + // Cache model info for the session (used by janitor for model selection) if (providerID && modelID) { state.model.set(sessionId, { providerID: providerID, @@ -93,8 +98,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({ @@ -104,7 +108,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) { @@ -112,16 +115,24 @@ 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) + + if (!state.toolParameters.has(callId)) { + state.toolParameters.set(callId, { + tool: part.tool, + parameters: part.input ?? {} + }) + } } } } } - // Create position mapping: "toolName:index" -> toolCallId const positionMapping = new Map() for (const [toolName, callIds] of toolCallsByName) { callIds.forEach((callId, index) => { @@ -132,7 +143,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/prompts/tool.txt b/lib/prompts/tool.txt index 579f510..abd2906 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 20db977..b75901d 100644 --- a/lib/pruning-tool.ts +++ b/lib/pruning-tool.ts @@ -1,57 +1,183 @@ 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 + + if (await isSubagentSession(client, sessionId)) { + return "Pruning is unavailable in subagent sessions. Do not call this tool again. Continue with your current task." + } + + if (!args.ids || args.ids.length === 0) { + return "No IDs provided. Check the list for available IDs to prune." + } + + await ensureSessionRestored(state, sessionId, logger) + + const prunedIds = args.ids + .map(numId => getActualId(sessionId, numId)) + .filter((id): id is string => id !== undefined) + + if (prunedIds.length === 0) { + return "None of the provided IDs were valid. Check the list for available IDs." + } + + const tokensSaved = await calculateTokensSaved(client, sessionId, prunedIds) + + 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) + + const alreadyPrunedIds = state.prunedIds.get(sessionId) ?? [] + const allPrunedIds = [...alreadyPrunedIds, ...prunedIds] + state.prunedIds.set(sessionId, allPrunedIds) + + saveSessionState(sessionId, new Set(allPrunedIds), sessionStats, logger) + .catch(err => logger.error("prune-tool", "Failed to persist state", { error: err.message })) + + const toolMetadata = new Map() + for (const id of prunedIds) { + 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(), + hasLower: state.toolParameters.has(id.toLowerCase()) + }) + } } - const result = await runOnTool( - janitorCtx, - ctx.sessionID, - config.strategies.onTool, - args.reason - ) + 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 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 + 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[] +): Promise { + try { + const messagesResponse = await client.session.messages({ + path: { id: sessionId }, + query: { limit: 200 } + }) + const messages = messagesResponse.data || messagesResponse + + 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) + } + 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) + } + } + } + } + + 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 + } + + const tokenCounts = await estimateTokensBatch(contents) + return tokenCounts.reduce((sum, count) => sum + count, 0) + } catch (error: any) { + return prunedIds.length * 500 + } +} diff --git a/lib/state/id-mapping.ts b/lib/state/id-mapping.ts new file mode 100644 index 0000000..fccac6d --- /dev/null +++ b/lib/state/id-mapping.ts @@ -0,0 +1,68 @@ +/** + * 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() + +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 +} + +export function getActualId(sessionId: string, numericId: number): string | undefined { + const mapping = sessionMappings.get(sessionId) + return mapping?.numericToActual.get(numericId) +} + +export function clearSessionMapping(sessionId: string): void { + sessionMappings.delete(sessionId) +} + +export function clearAllMappings(): void { + sessionMappings.clear() +} diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index 27549ea..f0ae3c6 100644 --- a/lib/state/tool-cache.ts +++ b/lib/state/tool-cache.ts @@ -1,36 +1,69 @@ 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 - } + 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, { + try { + const params = typeof toolCall.function.arguments === 'string' + ? JSON.parse(toolCall.function.arguments) + : toolCall.function.arguments + state.toolParameters.set(toolCall.id.toLowerCase(), { tool: toolCall.function.name, parameters: params }) - } catch (error) { - // Silently ignore parse errors + openaiCached++ + } catch (error) { + } + } + } + + 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.toLowerCase(), { + tool: part.name, + parameters: part.input ?? {} + }) + anthropicCached++ } } } + + 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 +72,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 @@ -50,14 +86,21 @@ 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 }) + 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 */