Merge pull request #80 from Tarquinen/refactor/consolidate-duplicates

Consolidate ToolMetadata type and findCurrentAgent function

Author: Tarquinen

index.ts

@@ -82,7 +82,7 @@ const plugin: Plugin = (async (ctx) => {
 
     return {
         event: createEventHandler(ctx.client, janitorCtx, logger, config, toolTracker),
-        "chat.params": createChatParamsHandler(ctx.client, state, logger),
+        "chat.params": createChatParamsHandler(ctx.client, state, logger, toolTracker),
         tool: config.strategies.onTool.length > 0 ? {
             prune: createPruningTool({
                 client: ctx.client,

lib/core/janitor.ts

@@ -2,6 +2,8 @@ import { z } from "zod"
 import type { Logger } from "../logger"
 import type { PruningStrategy } from "../config"
 import type { PluginState } from "../state"
+import type { ToolMetadata } from "../fetch-wrapper/types"
+import { findCurrentAgent } from "../hooks"
 import { buildAnalysisPrompt } from "./prompt"
 import { selectModel, extractModelFromSession } from "../model-selector"
 import { estimateTokensBatch, formatTokenCount } from "../tokenizer"
@@ -28,7 +30,7 @@ export interface PruningResult {
     prunedCount: number
     tokensSaved: number
     llmPrunedIds: string[]
-    toolMetadata: Map<string, { tool: string, parameters?: any }>
+    toolMetadata: Map<string, ToolMetadata>
     sessionStats: SessionStats
 }
 
@@ -263,7 +265,7 @@ async function runLlmAnalysis(
     messages: any[],
     unprunedToolCallIds: string[],
     alreadyPrunedIds: string[],
-    toolMetadata: Map<string, { tool: string, parameters?: any }>,
+    toolMetadata: Map<string, ToolMetadata>,
     options: PruningOptions
 ): Promise<string[]> {
     const { client, state, logger, config } = ctx
@@ -400,7 +402,7 @@ function replacePrunedToolOutputs(messages: any[], prunedIds: string[]): any[] {
 interface ParsedMessages {
     toolCallIds: string[]
     toolOutputs: Map<string, string>
-    toolMetadata: Map<string, { tool: string, parameters?: any }>
+    toolMetadata: Map<string, ToolMetadata>
 }
 
 export function parseMessages(
@@ -437,17 +439,6 @@ export function parseMessages(
     return { toolCallIds, toolOutputs, toolMetadata }
 }
 
-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
-}
-
 // ============================================================================
 // Helpers
 // ============================================================================

lib/core/strategies/index.ts

@@ -3,17 +3,15 @@
  */
 
 import { deduplicationStrategy } from "./deduplication"
+import type { ToolMetadata } from "../../fetch-wrapper/types"
+
+export type { ToolMetadata }
 
 /**
  * Common interface for rule-based pruning strategies.
  * Each strategy analyzes tool metadata and returns IDs that should be pruned.
  */
 
-export interface ToolMetadata {
-    tool: string
-    parameters?: any
-}
-
 export interface StrategyResult {
     /** Tool call IDs that should be pruned */
     prunedIds: string[]

lib/fetch-wrapper/prunable-list.ts

@@ -1,10 +1,6 @@
 import { extractParameterKey } from '../ui/display-utils'
 import { getOrCreateNumericId } from '../state/id-mapping'
-
-export interface ToolMetadata {
-    tool: string
-    parameters?: any
-}
+import type { ToolMetadata } from './types'
 
 const SYSTEM_REMINDER = `<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.

lib/fetch-wrapper/types.ts

@@ -9,6 +9,11 @@ export interface ToolOutput {
     toolName?: string
 }
 
+export interface ToolMetadata {
+    tool: string
+    parameters?: any
+}
+
 export interface FormatDescriptor {
     name: string
     detect(body: any): boolean

lib/hooks.ts

@@ -60,7 +60,8 @@ export function createEventHandler(
 export function createChatParamsHandler(
     client: any,
     state: PluginState,
-    logger: Logger
+    logger: Logger,
+    toolTracker?: ToolTracker
 ) {
     return async (input: any, _output: any) => {
         const sessionId = input.sessionID
@@ -73,11 +74,14 @@ export function createChatParamsHandler(
 
         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)
+                from: state.lastSeenSessionId,
+                to: sessionId
             })
             clearAllMappings()
             state.toolParameters.clear()
+            if (toolTracker) {
+                resetToolTrackerCount(toolTracker)
+            }
         }
 
         state.lastSeenSessionId = sessionId
@@ -148,3 +152,17 @@ export function createChatParamsHandler(
         }
     }
 }
+
+/**
+ * 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
+}

lib/pruning-tool.ts

@@ -2,8 +2,9 @@ import { tool } from "@opencode-ai/plugin"
 import type { PluginState } from "./state"
 import type { PluginConfig } from "./config"
 import type { ToolTracker } from "./fetch-wrapper/tool-tracker"
+import type { ToolMetadata } from "./fetch-wrapper/types"
 import { resetToolTrackerCount } from "./fetch-wrapper/tool-tracker"
-import { isSubagentSession } from "./hooks"
+import { isSubagentSession, findCurrentAgent } from "./hooks"
 import { getActualId } from "./state/id-mapping"
 import { sendUnifiedNotification, type NotificationContext } from "./ui/notification"
 import { ensureSessionRestored } from "./state"
@@ -92,7 +93,7 @@ export function createPruningTool(
             saveSessionState(sessionId, new Set(allPrunedIds), sessionStats, logger)
                 .catch(err => logger.error("prune-tool", "Failed to persist state", { error: err.message }))
 
-            const toolMetadata = new Map<string, { tool: string, parameters?: any }>()
+            const toolMetadata = new Map<string, ToolMetadata>()
             for (const id of prunedIds) {
                 const meta = state.toolParameters.get(id.toLowerCase())
                 if (meta) {
@@ -127,20 +128,6 @@ export function createPruningTool(
     })
 }
 
-/**
- * Finds the current agent from messages (same logic as janitor.ts).
- */
-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
-}
-
 /**
  * Calculates approximate tokens saved by pruning the given tool call IDs.
  * Uses pre-fetched messages to avoid duplicate API calls.

lib/ui/notification.ts

@@ -1,5 +1,6 @@
 import type { Logger } from "../logger"
 import type { SessionStats, GCStats, PruningResult } from "../core/janitor"
+import type { ToolMetadata } from "../fetch-wrapper/types"
 import { formatTokenCount } from "../tokenizer"
 import { extractParameterKey } from "./display-utils"
 
@@ -20,7 +21,7 @@ export interface NotificationData {
     aiPrunedCount: number
     aiTokensSaved: number
     aiPrunedIds: string[]
-    toolMetadata: Map<string, { tool: string, parameters?: any }>
+    toolMetadata: Map<string, ToolMetadata>
     gcPending: GCStats | null
     sessionStats: SessionStats | null
 }
@@ -164,7 +165,7 @@ export function formatPruningResultForTool(
 
 export function buildToolsSummary(
     prunedIds: string[],
-    toolMetadata: Map<string, { tool: string, parameters?: any }>,
+    toolMetadata: Map<string, ToolMetadata>,
     workingDirectory?: string
 ): Map<string, string[]> {
     const toolsSummary = new Map<string, string[]>()