Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
19 changes: 5 additions & 14 deletions lib/core/janitor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
}

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
// ============================================================================
Expand Down
8 changes: 3 additions & 5 deletions lib/core/strategies/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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[]
Expand Down
6 changes: 1 addition & 5 deletions lib/fetch-wrapper/prunable-list.ts
Original file line number Diff line number Diff line change
@@ -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.
Expand Down
5 changes: 5 additions & 0 deletions lib/fetch-wrapper/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
24 changes: 21 additions & 3 deletions lib/hooks.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
}
19 changes: 3 additions & 16 deletions lib/pruning-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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.
Expand Down
5 changes: 3 additions & 2 deletions lib/ui/notification.ts
Original file line number Diff line number Diff line change
@@ -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"

Expand All @@ -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
}
Expand Down Expand Up @@ -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[]>()
Expand Down