diff --git a/lib/messages/prune.ts b/lib/messages/prune.ts index fc97aee..e0cec3b 100644 --- a/lib/messages/prune.ts +++ b/lib/messages/prune.ts @@ -1,8 +1,7 @@ import type { SessionState, WithParts } from "../state" import type { Logger } from "../logger" import type { PluginConfig } from "../config" -import { buildToolIdList } from "../utils" -import { getLastUserMessage, extractParameterKey } from "./utils" +import { getLastUserMessage, extractParameterKey, buildToolIdList } from "./utils" import { loadPrompt } from "../prompt" const PRUNED_TOOL_OUTPUT_REPLACEMENT = '[Output removed to save context - information superseded or no longer needed]' diff --git a/lib/messages/utils.ts b/lib/messages/utils.ts index 1a9dbf6..54abbd8 100644 --- a/lib/messages/utils.ts +++ b/lib/messages/utils.ts @@ -2,7 +2,6 @@ import type { WithParts } from "../state" /** * Extracts a human-readable key from tool metadata for display purposes. - * Used by both deduplication and AI analysis to show what was pruned. */ export const extractParameterKey = (tool: string, parameters: any): string => { if (!parameters) return '' @@ -83,3 +82,33 @@ export const getLastUserMessage = ( } return null } + +export function findCurrentAgent(messages: WithParts[]): string | undefined { + const userMsg = getLastUserMessage(messages) + if (!userMsg) return undefined + return (userMsg.info as any).agent || 'build' +} + +export function buildToolIdList(messages: WithParts[]): string[] { + const toolIds: string[] = [] + for (const msg of messages) { + if (msg.parts) { + for (const part of msg.parts) { + if (part.type === 'tool' && part.callID && part.tool) { + toolIds.push(part.callID) + } + } + } + } + return toolIds +} + +export function getPruneToolIds(numericToolIds: number[], toolIdList: string[]): string[] { + const pruneToolIds: string[] = [] + for (const index of numericToolIds) { + if (!isNaN(index) && index >= 0 && index < toolIdList.length) { + pruneToolIds.push(toolIdList[index]) + } + } + return pruneToolIds +} diff --git a/lib/strategies/deduplication.ts b/lib/strategies/deduplication.ts index 50b30fa..61cc484 100644 --- a/lib/strategies/deduplication.ts +++ b/lib/strategies/deduplication.ts @@ -1,7 +1,8 @@ import { PluginConfig } from "../config" import { Logger } from "../logger" import type { SessionState, WithParts } from "../state" -import { buildToolIdList, calculateTokensSaved } from "../utils" +import { calculateTokensSaved } from "../utils" +import { buildToolIdList } from "../messages/utils" /** * Deduplication strategy - prunes older tool calls that have identical diff --git a/lib/strategies/on-idle.ts b/lib/strategies/on-idle.ts index a0f07d0..49887d3 100644 --- a/lib/strategies/on-idle.ts +++ b/lib/strategies/on-idle.ts @@ -4,7 +4,8 @@ import type { Logger } from "../logger" import type { PluginConfig } from "../config" import { buildAnalysisPrompt } from "../prompt" import { selectModel, extractModelFromSession, ModelInfo } from "../model-selector" -import { calculateTokensSaved, findCurrentAgent } from "../utils" +import { calculateTokensSaved } from "../utils" +import { findCurrentAgent } from "../messages/utils" import { saveSessionState } from "../state/persistence" import { sendUnifiedNotification } from "../ui/notification" diff --git a/lib/strategies/prune-tool.ts b/lib/strategies/prune-tool.ts index 1fccb44..d0a07ab 100644 --- a/lib/strategies/prune-tool.ts +++ b/lib/strategies/prune-tool.ts @@ -1,7 +1,8 @@ import { tool } from "@opencode-ai/plugin" import type { SessionState, ToolParameterEntry, WithParts } from "../state" import type { PluginConfig } from "../config" -import { findCurrentAgent, buildToolIdList, getPruneToolIds, calculateTokensSaved } from "../utils" +import { findCurrentAgent, buildToolIdList, getPruneToolIds } from "../messages/utils" +import { calculateTokensSaved } from "../utils" import { PruneReason, sendUnifiedNotification } from "../ui/notification" import { formatPruningResultForTool } from "../ui/display-utils" import { ensureSessionInitialized } from "../state" diff --git a/lib/utils.ts b/lib/utils.ts index 471729c..842b964 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -3,7 +3,6 @@ import { encode } from 'gpt-tokenizer' /** * Estimates token counts for a batch of texts using gpt-tokenizer. - * TODO: ensure we aren't falling back to catch branch */ function estimateTokensBatch(texts: string[]): number[] { try { @@ -15,7 +14,6 @@ function estimateTokensBatch(texts: string[]): number[] { /** * Calculates approximate tokens saved by pruning the given tool call IDs. - * Uses pre-fetched messages to avoid duplicate API calls. * TODO: Make it count message content that are not tool outputs. Currently it ONLY covers tool outputs and errors */ export const calculateTokensSaved = ( @@ -57,9 +55,6 @@ export function formatTokenCount(tokens: number): string { return tokens.toString() + ' tokens' } -/** - * Checks if a session is a subagent session by looking for a parentID. - */ export async function isSubAgentSession(client: any, sessionID: string): Promise { try { const result = await client.session.get({ path: { id: sessionID } }) @@ -68,47 +63,3 @@ export async function isSubAgentSession(client: any, sessionID: string): Promise return false } } - -/** - * Finds the current agent from messages by scanning backward for user messages. - */ -export function findCurrentAgent(messages: any[]): string | undefined { - for (let i = messages.length - 1; i >= 0; i--) { - const msg = messages[i] - const info = msg.info - if (info?.role === 'user') { - return info.agent || 'build' - } - } - return undefined -} - -/** - * Builds a list of tool call IDs from messages. - */ -export function buildToolIdList(messages: WithParts[]): string[] { - const toolIds: string[] = [] - for (const msg of messages) { - if (msg.parts) { - for (const part of msg.parts) { - if (part.type === 'tool' && part.callID && part.tool) { - toolIds.push(part.callID) - } - } - } - } - return toolIds -} - -/** - * Prunes numeric tool IDs to valid tool call IDs based on the provided tool ID list. - */ -export function getPruneToolIds(numericToolIds: number[], toolIdList: string[]): string[] { - const pruneToolIds: string[] = [] - for (const index of numericToolIds) { - if (!isNaN(index) && index >= 0 && index < toolIdList.length) { - pruneToolIds.push(toolIdList[index]) - } - } - return pruneToolIds -}