refactor: consolidate message utilities and eliminate duplicate iteration

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]'

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
+}

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

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"
 

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"

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<boolean> {
     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
-}