diff --git a/index.ts b/index.ts index 6a97e05..aaf6e64 100644 --- a/index.ts +++ b/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, diff --git a/lib/core/janitor.ts b/lib/core/janitor.ts index 5667c21..50a29bf 100644 --- a/lib/core/janitor.ts +++ b/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 + toolMetadata: Map sessionStats: SessionStats } @@ -263,7 +265,7 @@ async function runLlmAnalysis( messages: any[], unprunedToolCallIds: string[], alreadyPrunedIds: string[], - toolMetadata: Map, + toolMetadata: Map, options: PruningOptions ): Promise { const { client, state, logger, config } = ctx @@ -400,7 +402,7 @@ function replacePrunedToolOutputs(messages: any[], prunedIds: string[]): any[] { interface ParsedMessages { toolCallIds: string[] toolOutputs: Map - toolMetadata: Map + toolMetadata: Map } 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 // ============================================================================ diff --git a/lib/core/strategies/index.ts b/lib/core/strategies/index.ts index c2742e6..b4eb8af 100644 --- a/lib/core/strategies/index.ts +++ b/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[] diff --git a/lib/fetch-wrapper/prunable-list.ts b/lib/fetch-wrapper/prunable-list.ts index e711f2d..4cce826 100644 --- a/lib/fetch-wrapper/prunable-list.ts +++ b/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 = ` 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. diff --git a/lib/fetch-wrapper/types.ts b/lib/fetch-wrapper/types.ts index 7ea1f83..6a74840 100644 --- a/lib/fetch-wrapper/types.ts +++ b/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 diff --git a/lib/hooks.ts b/lib/hooks.ts index 5ee36b7..234acf6 100644 --- a/lib/hooks.ts +++ b/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 +} diff --git a/lib/pruning-tool.ts b/lib/pruning-tool.ts index dfa7fdc..ff3a1ee 100644 --- a/lib/pruning-tool.ts +++ b/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() + const toolMetadata = new Map() 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. diff --git a/lib/ui/notification.ts b/lib/ui/notification.ts index d9aec9c..68fe6c8 100644 --- a/lib/ui/notification.ts +++ b/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 + toolMetadata: Map gcPending: GCStats | null sessionStats: SessionStats | null } @@ -164,7 +165,7 @@ export function formatPruningResultForTool( export function buildToolsSummary( prunedIds: string[], - toolMetadata: Map, + toolMetadata: Map, workingDirectory?: string ): Map { const toolsSummary = new Map()