Skip to content

Commit 3143b87

Browse files
jorgenwhjorgenwh
committed
deduplication and prune tool works. MUCH nicer code
1 parent 1a2fc6c commit 3143b87

File tree

14 files changed

+166
-119
lines changed

14 files changed

+166
-119
lines changed

‎index.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ const plugin: Plugin = (async (ctx) => {
2222
const state = createSessionState()
2323

2424
// Log initialization
25-
logger.info("plugin", "DCP initialized", {
25+
logger.info("DCP initialized", {
2626
strategies: config.strategies,
2727
})
2828

‎lib/hooks.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import type { PluginConfig } from "./config"
44
import { syncToolCache } from "./state/tool-cache"
55
import { deduplicate } from "./strategies"
66
import { prune, insertPruneToolContext } from "./messages"
7+
import { checkSession } from "./state"
78

89

910
export function createChatMessageTransformHandler(
@@ -16,6 +17,7 @@ export function createChatMessageTransformHandler(
1617
input: {},
1718
output: { messages: WithParts[] }
1819
) => {
20+
checkSession(state, logger, output.messages);
1921
syncToolCache(state, logger, output.messages);
2022

2123
deduplicate(state, logger, config, output.messages)

‎lib/logger.ts‎

Lines changed: 31 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,29 @@ export class Logger {
4545
return parts.join(" ")
4646
}
4747

48+
private getCallerFile(skipFrames: number = 3): string {
49+
const originalPrepareStackTrace = Error.prepareStackTrace
50+
try {
51+
const err = new Error()
52+
Error.prepareStackTrace = (_, stack) => stack
53+
const stack = err.stack as unknown as NodeJS.CallSite[]
54+
Error.prepareStackTrace = originalPrepareStackTrace
55+
56+
// Skip specified number of frames to get to actual caller
57+
for (let i = skipFrames; i < stack.length; i++) {
58+
const filename = stack[i]?.getFileName()
59+
if (filename && !filename.includes('/logger.')) {
60+
// Extract just the filename without path and extension
61+
const match = filename.match(/([^/\\]+)\.[tj]s$/)
62+
return match ? match[1] : filename
63+
}
64+
}
65+
return 'unknown'
66+
} catch {
67+
return 'unknown'
68+
}
69+
}
70+
4871
private async write(level: string, component: string, message: string, data?: any) {
4972
if (!this.enabled) return
5073

@@ -67,19 +90,23 @@ export class Logger {
6790
}
6891
}
6992

70-
info(component: string, message: string, data?: any) {
93+
info(message: string, data?: any) {
94+
const component = this.getCallerFile(2)
7195
return this.write("INFO", component, message, data)
7296
}
7397

74-
debug(component: string, message: string, data?: any) {
98+
debug(message: string, data?: any) {
99+
const component = this.getCallerFile(2)
75100
return this.write("DEBUG", component, message, data)
76101
}
77102

78-
warn(component: string, message: string, data?: any) {
103+
warn(message: string, data?: any) {
104+
const component = this.getCallerFile(2)
79105
return this.write("WARN", component, message, data)
80106
}
81107

82-
error(component: string, message: string, data?: any) {
108+
error(message: string, data?: any) {
109+
const component = this.getCallerFile(2)
83110
return this.write("ERROR", component, message, data)
84111
}
85112

‎lib/messages/prune.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ const buildPrunableToolsList = (
2626
const paramKey = extractParameterKey(toolParameterEntry.tool, toolParameterEntry.parameters)
2727
const description = paramKey ? `${toolParameterEntry.tool}, ${paramKey}` : toolParameterEntry.tool
2828
lines.push(`${numericId}: ${description}`)
29+
logger.debug(`Prunable tool found - ID: ${numericId}, Tool: ${toolParameterEntry.tool}, Call ID: ${toolCallId}`)
2930
})
3031

3132
return `<prunable-tools>\nThe following tools have been invoked and are available for pruning. This list does not mandate immediate action. Consider your current goals and the resources you need before discarding valuable tool outputs. Keep the context free of noise.\n${lines.join('\n')}\n</prunable-tools>`
@@ -89,7 +90,7 @@ const pruneToolOutputs = (
8990
if (part.type !== 'tool') {
9091
continue
9192
}
92-
if (!state.prune.toolIds.includes(part.id)) {
93+
if (!state.prune.toolIds.includes(part.callID)) {
9394
continue
9495
}
9596
if (part.state.status === 'completed') {

‎lib/model-selector.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ async function importOpencodeAI(logger?: Logger, maxRetries: number = 3, delayMs
5555
lastError = error;
5656

5757
if (error.message?.includes('before initialization')) {
58-
logger?.debug('model-selector', `Import attempt ${attempt}/${maxRetries} failed, will retry`, {
58+
logger?.debug(`Import attempt ${attempt}/${maxRetries} failed, will retry`, {
5959
error: error.message
6060
});
6161

@@ -85,7 +85,7 @@ export async function selectModel(
8585
if (configModel) {
8686
const parts = configModel.split('/');
8787
if (parts.length !== 2) {
88-
logger?.warn('model-selector', 'Invalid config model format', { configModel });
88+
logger?.warn('Invalid config model format', { configModel });
8989
} else {
9090
const [providerID, modelID] = parts;
9191

@@ -98,7 +98,7 @@ export async function selectModel(
9898
reason: 'Using model specified in dcp.jsonc config'
9999
};
100100
} catch (error: any) {
101-
logger?.warn('model-selector', `Config model failed: ${providerID}/${modelID}`, {
101+
logger?.warn(`Config model failed: ${providerID}/${modelID}`, {
102102
error: error.message
103103
});
104104
failedModelInfo = { providerID, modelID };

‎lib/state/persistence.ts‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -61,13 +61,13 @@ export async function saveSessionState(
6161
const content = JSON.stringify(state, null, 2);
6262
await fs.writeFile(filePath, content, "utf-8");
6363

64-
logger.info("persist", "Saved session state to disk", {
65-
sessionId: sessionState.sessionId.slice(0, 8),
64+
logger.info("Saved session state to disk", {
65+
sessionId: sessionState.sessionId,
6666
totalTokensSaved: state.stats.totalPruneTokens
6767
});
6868
} catch (error: any) {
69-
logger.error("persist", "Failed to save session state", {
70-
sessionId: sessionState.sessionId?.slice(0, 8),
69+
logger.error("Failed to save session state", {
70+
sessionId: sessionState.sessionId,
7171
error: error?.message,
7272
});
7373
}
@@ -92,21 +92,21 @@ export async function loadSessionState(
9292
!Array.isArray(state.prune.toolIds) ||
9393
!state.stats
9494
) {
95-
logger.warn("persist", "Invalid session state file, ignoring", {
96-
sessionId: sessionId.slice(0, 8),
95+
logger.warn("Invalid session state file, ignoring", {
96+
sessionId: sessionId,
9797
});
9898
return null;
9999
}
100100

101-
logger.info("persist", "Loaded session state from disk", {
102-
sessionId: sessionId.slice(0, 8),
101+
logger.info("Loaded session state from disk", {
102+
sessionId: sessionId,
103103
totalTokensSaved: state.stats.totalPruneTokens
104104
});
105105

106106
return state;
107107
} catch (error: any) {
108-
logger.warn("persist", "Failed to load session state", {
109-
sessionId: sessionId.slice(0, 8),
108+
logger.warn("Failed to load session state", {
109+
sessionId: sessionId,
110110
error: error?.message,
111111
});
112112
return null;

‎lib/state/state.ts‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,32 @@
1-
import type { SessionState, ToolParameterEntry } from "./types"
1+
import type { SessionState, ToolParameterEntry, WithParts } from "./types"
22
import type { Logger } from "../logger"
33
import { loadSessionState } from "./persistence"
4+
import { getLastUserMessage } from "../messages/utils"
5+
6+
export const checkSession = (
7+
state: SessionState,
8+
logger: Logger,
9+
messages: WithParts[]
10+
) => {
11+
12+
const lastUserMessage = getLastUserMessage(messages)
13+
if (!lastUserMessage) {
14+
return
15+
}
16+
17+
const lastSessionId = lastUserMessage.info.sessionID
18+
19+
if (state.sessionId === null || state.sessionId !== lastSessionId) {
20+
logger.info(`Session changed: ${state.sessionId} -> ${lastSessionId}`)
21+
ensureSessionInitialized(
22+
state,
23+
lastSessionId,
24+
logger
25+
).catch((err) => {
26+
logger.error("Failed to initialize session state", { error: err.message })
27+
} )
28+
}
29+
}
430

531
export function createSessionState(): SessionState {
632
return {
@@ -37,6 +63,9 @@ export async function ensureSessionInitialized(
3763
return;
3864
}
3965

66+
logger.info("session ID = " + sessionId)
67+
logger.info("Initializing session state", { sessionId: sessionId })
68+
4069
// Clear previous session data
4170
resetSessionState(state)
4271
state.sessionId = sessionId

‎lib/state/tool-cache.ts‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,8 @@ export async function syncToolCache(
1414
messages: WithParts[],
1515
): Promise<void> {
1616
try {
17+
logger.info("Syncing tool parameters from OpenCode messages")
1718
for (const msg of messages) {
18-
if (!msg.parts) {
19-
continue
20-
}
21-
2219
for (const part of msg.parts) {
2320
if (part.type !== "tool" || !part.callID || state.toolParameters.has(part.callID)) {
2421
continue
@@ -38,7 +35,7 @@ export async function syncToolCache(
3835

3936
trimToolParametersCache(state)
4037
} catch (error) {
41-
logger.warn("tool-cache", "Failed to sync tool parameters from OpenCode", {
38+
logger.warn("Failed to sync tool parameters from OpenCode", {
4239
error: error instanceof Error ? error.message : String(error)
4340
})
4441
}

‎lib/strategies/deduplication.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { PluginConfig } from "../config"
22
import { Logger } from "../logger"
33
import type { SessionState, WithParts } from "../state"
4+
import { calculateTokensSaved } from "../utils"
45

56
/**
67
* Deduplication strategy - prunes older tool calls that have identical
@@ -39,7 +40,7 @@ export const deduplicate = (
3940
for (const id of unprunedIds) {
4041
const metadata = state.toolParameters.get(id)
4142
if (!metadata) {
42-
logger.warn("deduplication", `Missing metadata for tool call ID: ${id}`)
43+
logger.warn(`Missing metadata for tool call ID: ${id}`)
4344
continue
4445
}
4546

@@ -66,9 +67,11 @@ export const deduplicate = (
6667
}
6768
}
6869

70+
state.stats.totalPruneTokens += calculateTokensSaved(messages, newPruneIds)
71+
6972
if (newPruneIds.length > 0) {
7073
state.prune.toolIds.push(...newPruneIds)
71-
logger.debug("deduplication", `Marked ${newPruneIds.length} duplicate tool calls for pruning`)
74+
logger.debug(`Marked ${newPruneIds.length} duplicate tool calls for pruning`)
7275
}
7376
}
7477

‎lib/strategies/prune-tool.ts‎

Lines changed: 8 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,12 @@
11
import { tool } from "@opencode-ai/plugin"
22
import type { SessionState, ToolParameterEntry, WithParts } from "../state"
33
import type { PluginConfig } from "../config"
4-
import { findCurrentAgent, buildToolIdList, getPruneToolIds } from "../utils"
4+
import { findCurrentAgent, buildToolIdList, getPruneToolIds, calculateTokensSaved } from "../utils"
55
import { PruneReason, sendUnifiedNotification } from "../ui/notification"
66
import { formatPruningResultForTool } from "../ui/display-utils"
77
import { ensureSessionInitialized } from "../state"
88
import { saveSessionState } from "../state/persistence"
99
import type { Logger } from "../logger"
10-
import { estimateTokensBatch } from "../tokenizer"
1110
import { loadPrompt } from "../prompt"
1211

1312
/** Tool description loaded from prompts/tool.txt */
@@ -67,29 +66,28 @@ export function createPruneTool(
6766
const messagesResponse = await client.session.messages({
6867
path: { id: sessionId }
6968
})
70-
const messages = messagesResponse.data || messagesResponse
69+
const messages: WithParts[] = messagesResponse.data || messagesResponse
7170

7271
const currentAgent: string | undefined = findCurrentAgent(messages)
7372
const toolIdList: string[] = buildToolIdList(messages)
7473
const pruneToolIds: string[] = getPruneToolIds(numericToolIds, toolIdList)
75-
const tokensSaved = await calculateTokensSavedFromMessages(messages, pruneToolIds)
76-
77-
state.stats.pruneTokenCounter += tokensSaved
7874
state.prune.toolIds.push(...pruneToolIds)
7975

8076
saveSessionState(state, logger)
81-
.catch(err => logger.error("prune-tool", "Failed to persist state", { error: err.message }))
77+
.catch(err => logger.error("Failed to persist state", { error: err.message }))
8278

8379
const toolMetadata = new Map<string, ToolParameterEntry>()
8480
for (const id of pruneToolIds) {
8581
const toolParameters = state.toolParameters.get(id)
8682
if (toolParameters) {
8783
toolMetadata.set(id, toolParameters)
8884
} else {
89-
logger.debug("prune-tool", "No metadata found for ID", { id })
85+
logger.debug("No metadata found for ID", { id })
9086
}
9187
}
9288

89+
state.stats.pruneTokenCounter += calculateTokensSaved(messages, pruneToolIds)
90+
9391
await sendUnifiedNotification(
9492
client,
9593
logger,
@@ -102,6 +100,8 @@ export function createPruneTool(
102100
currentAgent,
103101
workingDirectory
104102
)
103+
state.stats.totalPruneTokens += state.stats.pruneTokenCounter
104+
state.stats.pruneTokenCounter = 0
105105

106106
return formatPruningResultForTool(
107107
pruneToolIds,
@@ -112,50 +112,3 @@ export function createPruneTool(
112112
})
113113
}
114114

115-
/**
116-
* Calculates approximate tokens saved by pruning the given tool call IDs.
117-
* Uses pre-fetched messages to avoid duplicate API calls.
118-
*/
119-
async function calculateTokensSavedFromMessages(
120-
messages: any[],
121-
prunedIds: string[]
122-
): Promise<number> {
123-
try {
124-
const toolOutputs = new Map<string, string>()
125-
for (const msg of messages) {
126-
if (msg.role === 'tool' && msg.tool_call_id) {
127-
const content = typeof msg.content === 'string'
128-
? msg.content
129-
: JSON.stringify(msg.content)
130-
toolOutputs.set(msg.tool_call_id.toLowerCase(), content)
131-
}
132-
if (msg.role === 'user' && Array.isArray(msg.content)) {
133-
for (const part of msg.content) {
134-
if (part.type === 'tool_result' && part.tool_use_id) {
135-
const content = typeof part.content === 'string'
136-
? part.content
137-
: JSON.stringify(part.content)
138-
toolOutputs.set(part.tool_use_id.toLowerCase(), content)
139-
}
140-
}
141-
}
142-
}
143-
144-
const contents: string[] = []
145-
for (const id of prunedIds) {
146-
const content = toolOutputs.get(id.toLowerCase())
147-
if (content) {
148-
contents.push(content)
149-
}
150-
}
151-
152-
if (contents.length === 0) {
153-
return prunedIds.length * 500
154-
}
155-
156-
const tokenCounts = await estimateTokensBatch(contents)
157-
return tokenCounts.reduce((sum, count) => sum + count, 0)
158-
} catch (error: any) {
159-
return prunedIds.length * 500
160-
}
161-
}

0 commit comments

Comments
 (0)