Skip to content

Commit 58145ba

Browse files
jorgenwhjorgenwh
committed
refactor
1 parent befa392 commit 58145ba

File tree

9 files changed

+89
-106
lines changed

9 files changed

+89
-106
lines changed

‎index.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import type { Plugin } from "@opencode-ai/plugin"
22
import { getConfig } from "./lib/config"
33
import { Logger } from "./lib/logger"
44
import { createSessionState } from "./lib/state"
5-
import { createPruningTool } from "./lib/pruning-tool"
5+
import { createPruningTool } from "./lib/strategies/pruning-tool"
66
import { createChatMessageTransformHandler } from "./lib/hooks"
77

88
const plugin: Plugin = (async (ctx) => {

‎lib/config.ts‎

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import type { PluginInput } from '@opencode-ai/plugin'
66

77
export interface Deduplication {
88
enabled: boolean
9+
protectedTools: string[]
910
}
1011

1112
export interface PruneThinkingBlocks {
@@ -52,6 +53,7 @@ export const VALID_CONFIG_KEYS = new Set([
5253
// strategies.deduplication
5354
'strategies.deduplication',
5455
'strategies.deduplication.enabled',
56+
'strategies.deduplication.protectedTools',
5557
// strategies.pruneThinkingBlocks
5658
'strategies.pruneThinkingBlocks',
5759
'strategies.pruneThinkingBlocks.enabled',
@@ -122,6 +124,9 @@ function validateConfigTypes(config: Record<string, any>): ValidationError[] {
122124
if (strategies.deduplication?.enabled !== undefined && typeof strategies.deduplication.enabled !== 'boolean') {
123125
errors.push({ key: 'strategies.deduplication.enabled', expected: 'boolean', actual: typeof strategies.deduplication.enabled })
124126
}
127+
if (strategies.deduplication?.protectedTools !== undefined && !Array.isArray(strategies.deduplication.protectedTools)) {
128+
errors.push({ key: 'strategies.deduplication.protectedTools', expected: 'string[]', actual: typeof strategies.deduplication.protectedTools })
129+
}
125130

126131
// pruneThinkingBlocks
127132
if (strategies.pruneThinkingBlocks?.enabled !== undefined && typeof strategies.pruneThinkingBlocks.enabled !== 'boolean') {
@@ -217,7 +222,8 @@ const defaultConfig: PluginConfig = {
217222
pruningSummary: 'detailed',
218223
strategies: {
219224
deduplication: {
220-
enabled: true
225+
enabled: true,
226+
protectedTools: [...DEFAULT_PROTECTED_TOOLS]
221227
},
222228
pruneThinkingBlocks: {
223229
enabled: true
@@ -297,7 +303,9 @@ function createDefaultConfig(): void {
297303
"strategies": {
298304
// Remove duplicate tool calls (same tool with same arguments)
299305
"deduplication": {
300-
"enabled": true
306+
"enabled": true,
307+
// Additional tools to protect from pruning
308+
"protectedTools": []
301309
},
302310
// Remove thinking/reasoning LLM blocks
303311
"pruneThinkingBlocks": {
@@ -362,7 +370,13 @@ function mergeStrategies(
362370

363371
return {
364372
deduplication: {
365-
enabled: override.deduplication?.enabled ?? base.deduplication.enabled
373+
enabled: override.deduplication?.enabled ?? base.deduplication.enabled,
374+
protectedTools: [
375+
...new Set([
376+
...base.deduplication.protectedTools,
377+
...(override.deduplication?.protectedTools ?? [])
378+
])
379+
]
366380
},
367381
pruneThinkingBlocks: {
368382
enabled: override.pruneThinkingBlocks?.enabled ?? base.pruneThinkingBlocks.enabled
@@ -396,7 +410,10 @@ function deepCloneConfig(config: PluginConfig): PluginConfig {
396410
return {
397411
...config,
398412
strategies: {
399-
deduplication: { ...config.strategies.deduplication },
413+
deduplication: {
414+
...config.strategies.deduplication,
415+
protectedTools: [...config.strategies.deduplication.protectedTools]
416+
},
400417
pruneThinkingBlocks: { ...config.strategies.pruneThinkingBlocks },
401418
onIdle: {
402419
...config.strategies.onIdle,
@@ -476,4 +493,4 @@ export function getConfig(ctx: PluginInput): PluginConfig {
476493
}
477494

478495
return config
479-
}
496+

‎lib/hooks.ts‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,5 @@ export function createChatMessageTransformHandler(
1818
syncToolCache(state, logger, output.messages);
1919

2020
deduplicate(state, logger, config, output.messages)
21-
pruneTool(state, logger, config, output.messages)
2221
}
2322
}
24-
25-

‎lib/state/persistence.ts‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,12 @@ import * as fs from "fs/promises";
88
import { existsSync } from "fs";
99
import { homedir } from "os";
1010
import { join } from "path";
11-
import type { SessionState, SessionStats } from "./types"
11+
import type { SessionState, SessionStats, Prune } from "./types"
1212
import type { Logger } from "../logger";
1313

1414
export interface PersistedSessionState {
1515
sessionName?: string;
16-
prunedIds: string[];
16+
prune: Prune
1717
stats: SessionStats;
1818
lastUpdated: string;
1919
}
@@ -52,7 +52,7 @@ export async function saveSessionState(
5252

5353
const state: PersistedSessionState = {
5454
sessionName: sessionName,
55-
prunedIds: sessionState.prunedIds,
55+
prune: sessionState.prune,
5656
stats: sessionState.stats,
5757
lastUpdated: new Date().toISOString(),
5858
};
@@ -63,7 +63,6 @@ export async function saveSessionState(
6363

6464
logger.info("persist", "Saved session state to disk", {
6565
sessionId: sessionState.sessionId.slice(0, 8),
66-
prunedIds: state.prunedIds.length,
6766
totalTokensSaved: state.stats.totalTokensSaved,
6867
});
6968
} catch (error: any) {
@@ -88,7 +87,11 @@ export async function loadSessionState(
8887
const content = await fs.readFile(filePath, "utf-8");
8988
const state = JSON.parse(content) as PersistedSessionState;
9089

91-
if (!state || !Array.isArray(state.prunedIds) || !state.stats) {
90+
if (!state ||
91+
!state.prune ||
92+
!Array.isArray(state.prune.toolIds) ||
93+
!state.stats
94+
) {
9295
logger.warn("persist", "Invalid session state file, ignoring", {
9396
sessionId: sessionId.slice(0, 8),
9497
});
@@ -97,7 +100,6 @@ export async function loadSessionState(
97100

98101
logger.info("persist", "Loaded session state from disk", {
99102
sessionId: sessionId.slice(0, 8),
100-
prunedIds: state.prunedIds.length,
101103
totalTokensSaved: state.stats.totalTokensSaved,
102104
});
103105

‎lib/state/state.ts‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,9 @@ import { loadSessionState } from "./persistence"
55
export function createSessionState(): SessionState {
66
return {
77
sessionId: null,
8-
prunedIds: [],
8+
prune: {
9+
toolIds: []
10+
},
911
stats: {
1012
totalToolsPruned: 0,
1113
totalTokensSaved: 0,
@@ -22,7 +24,9 @@ export function createSessionState(): SessionState {
2224

2325
export function resetSessionState(state: SessionState): void {
2426
state.sessionId = null
25-
state.prunedIds = []
27+
state.prune = {
28+
toolIds: []
29+
}
2630
state.stats = {
2731
totalToolsPruned: 0,
2832
totalTokensSaved: 0,
@@ -56,7 +60,9 @@ export async function ensureSessionInitialized(
5660
}
5761

5862
// Populate state with loaded data
59-
state.prunedIds = persisted.prunedIds || []
63+
state.prune = {
64+
toolIds: persisted.prune.toolIds || []
65+
}
6066
state.stats = {
6167
totalToolsPruned: persisted.stats.totalToolsPruned || 0,
6268
totalTokensSaved: persisted.stats.totalTokensSaved || 0,

‎lib/state/types.ts‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,9 @@ export interface ToolParameterEntry {
1414
error?: string
1515
}
1616

17-
export interface GCStats {
18-
tokensCollected: number
19-
toolsDeduped: number
20-
}
21-
2217
export interface SessionStats {
23-
totalToolsPruned: number
24-
totalTokensSaved: number
25-
totalGCTokens: number
26-
totalGCTools: number
18+
pruneTokenCounter: number
19+
totalPruneTokens: number
2720
}
2821

2922
export interface Prune {
@@ -34,6 +27,5 @@ export interface SessionState {
3427
sessionId: string | null
3528
prune: Prune
3629
stats: SessionStats
37-
gcPending: GCStats
3830
toolParameters: Map<string, ToolParameterEntry>
3931
}

lib/pruning-tool.ts renamed to lib/strategies/pruning-tool.ts

Lines changed: 20 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
11
import { tool } from "@opencode-ai/plugin"
2-
import type { SessionState, ToolParameterEntry} from "./state"
3-
import type { PluginConfig } from "./config"
4-
import { findCurrentAgent, buildToolIdList, getPrunedIds } from "./utils"
5-
import { PruneReason, sendUnifiedNotification } from "./ui/notification"
6-
import { formatPruningResultForTool } from "./ui/display-utils"
7-
import { ensureSessionInitialized } from "./state"
8-
import { saveSessionState } from "./state/persistence"
9-
import type { Logger } from "./logger"
10-
import { estimateTokensBatch } from "./tokenizer"
11-
import { loadPrompt } from "./prompt"
2+
import type { SessionState, ToolParameterEntry} from "../state"
3+
import type { PluginConfig } from "../config"
4+
import { findCurrentAgent, buildToolIdList, getPruneToolIds } from "../utils"
5+
import { PruneReason, sendUnifiedNotification } from "../ui/notification"
6+
import { formatPruningResultForTool } from "../ui/display-utils"
7+
import { ensureSessionInitialized } from "../state"
8+
import { saveSessionState } from "../state/persistence"
9+
import type { Logger } from "../logger"
10+
import { estimateTokensBatch } from "../tokenizer"
11+
import { loadPrompt } from "../prompt"
1212

1313
/** Tool description loaded from prompts/tool.txt */
1414
const TOOL_DESCRIPTION = loadPrompt("tool")
@@ -56,8 +56,8 @@ export function createPruningTool(
5656
return "No valid pruning reason found. Use 'completion', 'noise', or 'consolidation' as the first element."
5757
}
5858

59-
const numericIds: number[] = args.ids.slice(1).filter((id): id is number => typeof id === "number")
60-
if (numericIds.length === 0) {
59+
const numericToolIds: number[] = args.ids.slice(1).filter((id): id is number => typeof id === "number")
60+
if (numericToolIds.length === 0) {
6161
return "No numeric IDs provided. Format: [reason, id1, id2, ...] where reason is 'completion', 'noise', or 'consolidation'."
6262
}
6363

@@ -67,22 +67,20 @@ export function createPruningTool(
6767
const messages = await client.session.messages({
6868
path: { id: sessionId }
6969
})
70-
// const messages = messagesResponse.data || messagesResponse // Need this?
7170

7271
const currentAgent: string | undefined = findCurrentAgent(messages)
7372
const toolIdList: string[] = buildToolIdList(messages)
74-
const prunedIds: string[] = getPrunedIds(numericIds, toolIdList)
75-
const tokensSaved = await calculateTokensSavedFromMessages(messages, prunedIds)
73+
const pruneToolIds: string[] = getPruneToolIds(numericToolIds, toolIdList)
74+
const tokensSaved = await calculateTokensSavedFromMessages(messages, pruneToolIds)
7675

77-
state.stats.totalTokensSaved += tokensSaved
78-
state.stats.totalToolsPruned += prunedIds.length
79-
state.prunedIds.push(...prunedIds)
76+
state.stats.pruneTokenCounter += tokensSaved
77+
state.prune.toolIds.push(...pruneToolIds)
8078

8179
saveSessionState(state, logger)
8280
.catch(err => logger.error("prune-tool", "Failed to persist state", { error: err.message }))
8381

8482
const toolMetadata = new Map<string, ToolParameterEntry>()
85-
for (const id of prunedIds) {
83+
for (const id of pruneToolIds) {
8684
const toolParameters = state.toolParameters.get(id)
8785
if (toolParameters) {
8886
toolMetadata.set(id, toolParameters)
@@ -95,20 +93,17 @@ export function createPruningTool(
9593
client,
9694
logger,
9795
config,
96+
state,
9897
sessionId,
99-
prunedIds.length,
100-
tokensSaved,
101-
prunedIds,
98+
pruneToolIds,
10299
toolMetadata,
103-
null,
104-
state.stats,
105100
reason as PruneReason,
106101
currentAgent,
107102
workingDirectory
108103
)
109104

110105
return formatPruningResultForTool(
111-
prunedIds,
106+
pruneToolIds,
112107
toolMetadata,
113108
workingDirectory
114109
)

0 commit comments

Comments
 (0)