Skip to content

Commit 1a2fc6c

Browse files
jorgenwhjorgenwh
committed
wip
1 parent 8c48761 commit 1a2fc6c

File tree

10 files changed

+204
-47
lines changed

10 files changed

+204
-47
lines changed

‎lib/constants.ts‎

Lines changed: 0 additions & 1 deletion
This file was deleted.

‎lib/hooks.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import type { Logger } from "./logger"
33
import type { PluginConfig } from "./config"
44
import { syncToolCache } from "./state/tool-cache"
55
import { deduplicate } from "./strategies"
6-
import { prune } from "./prune"
6+
import { prune, insertPruneToolContext } from "./messages"
77

88

99
export function createChatMessageTransformHandler(
@@ -21,5 +21,7 @@ export function createChatMessageTransformHandler(
2121
deduplicate(state, logger, config, output.messages)
2222

2323
prune(state, logger, config, output.messages)
24+
25+
insertPruneToolContext(state, config, logger, output.messages)
2426
}
2527
}

‎lib/messages/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
export { prune, insertPruneToolContext } from "./prune"

‎lib/messages/prune.ts‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
import type { SessionState, WithParts } from "../state"
2+
import type { Logger } from "../logger"
3+
import type { PluginConfig } from "../config"
4+
import { buildToolIdList } from "../utils"
5+
import { getLastUserMessage, extractParameterKey } from "./utils"
6+
7+
const PRUNED_TOOL_OUTPUT_REPLACEMENT = '[Output removed to save context - information superseded or no longer needed]'
8+
9+
const buildPrunableToolsList = (
10+
state: SessionState,
11+
config: PluginConfig,
12+
logger: Logger,
13+
messages: WithParts[],
14+
): string => {
15+
const lines: string[] = []
16+
const toolIdList: string[] = buildToolIdList(messages)
17+
18+
state.toolParameters.forEach((toolParameterEntry, toolCallId) => {
19+
if (state.prune.toolIds.includes(toolCallId)) {
20+
return
21+
}
22+
if (config.strategies.pruneTool.protectedTools.includes(toolParameterEntry.tool)) {
23+
return
24+
}
25+
const numericId = toolIdList.indexOf(toolCallId)
26+
const paramKey = extractParameterKey(toolParameterEntry.tool, toolParameterEntry.parameters)
27+
const description = paramKey ? `${toolParameterEntry.tool}, ${paramKey}` : toolParameterEntry.tool
28+
lines.push(`${numericId}: ${description}`)
29+
})
30+
31+
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>`
32+
}
33+
34+
export const insertPruneToolContext = (
35+
state: SessionState,
36+
config: PluginConfig,
37+
logger: Logger,
38+
messages: WithParts[]
39+
): void => {
40+
const lastUserMessage = getLastUserMessage(messages)
41+
if (!lastUserMessage || lastUserMessage.info.role !== 'user') {
42+
return
43+
}
44+
45+
const prunableToolsList = buildPrunableToolsList(state, config, logger, messages)
46+
47+
const userMessage: WithParts = {
48+
info: {
49+
id: "msg_01234567890123456789012345",
50+
sessionID: lastUserMessage.info.sessionID,
51+
role: "user",
52+
time: { created: Date.now() },
53+
agent: lastUserMessage.info.agent || "build",
54+
model: {
55+
providerID: lastUserMessage.info.model.providerID,
56+
modelID: lastUserMessage.info.model.modelID
57+
}
58+
},
59+
parts: [
60+
{
61+
id: "prt_01234567890123456789012345",
62+
sessionID: lastUserMessage.info.sessionID,
63+
messageID: "msg_01234567890123456789012345",
64+
type: "text",
65+
text: prunableToolsList,
66+
}
67+
]
68+
}
69+
70+
messages.push(userMessage)
71+
}
72+
73+
export const prune = (
74+
state: SessionState,
75+
logger: Logger,
76+
config: PluginConfig,
77+
messages: WithParts[]
78+
): void => {
79+
pruneToolOutputs(state, logger, messages)
80+
}
81+
82+
const pruneToolOutputs = (
83+
state: SessionState,
84+
logger: Logger,
85+
messages: WithParts[]
86+
): void => {
87+
for (const msg of messages) {
88+
for (const part of msg.parts) {
89+
if (part.type !== 'tool') {
90+
continue
91+
}
92+
if (!state.prune.toolIds.includes(part.id)) {
93+
continue
94+
}
95+
if (part.state.status === 'completed') {
96+
part.state.output = PRUNED_TOOL_OUTPUT_REPLACEMENT
97+
}
98+
// if (part.state.status === 'error') {
99+
// part.state.error = PRUNED_TOOL_OUTPUT_REPLACEMENT
100+
// }
101+
}
102+
}
103+
}

‎lib/messages/utils.ts‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import type { WithParts } from "../state"
2+
3+
/**
4+
* Extracts a human-readable key from tool metadata for display purposes.
5+
* Used by both deduplication and AI analysis to show what was pruned.
6+
*/
7+
export const extractParameterKey = (tool: string, parameters: any): string => {
8+
if (tool === "read" && parameters.filePath) {
9+
return parameters.filePath
10+
}
11+
if (tool === "write" && parameters.filePath) {
12+
return parameters.filePath
13+
}
14+
if (tool === "edit" && parameters.filePath) {
15+
return parameters.filePath
16+
}
17+
18+
if (tool === "list") {
19+
return parameters.path || '(current directory)'
20+
}
21+
if (tool === "glob") {
22+
if (parameters.pattern) {
23+
const pathInfo = parameters.path ? ` in ${parameters.path}` : ""
24+
return `"${parameters.pattern}"${pathInfo}`
25+
}
26+
return '(unknown pattern)'
27+
}
28+
if (tool === "grep") {
29+
if (parameters.pattern) {
30+
const pathInfo = parameters.path ? ` in ${parameters.path}` : ""
31+
return `"${parameters.pattern}"${pathInfo}`
32+
}
33+
return '(unknown pattern)'
34+
}
35+
36+
if (tool === "bash") {
37+
if (parameters.description) return parameters.description
38+
if (parameters.command) {
39+
return parameters.command.length > 50
40+
? parameters.command.substring(0, 50) + "..."
41+
: parameters.command
42+
}
43+
}
44+
45+
if (tool === "webfetch" && parameters.url) {
46+
return parameters.url
47+
}
48+
if (tool === "websearch" && parameters.query) {
49+
return `"${parameters.query}"`
50+
}
51+
if (tool === "codesearch" && parameters.query) {
52+
return `"${parameters.query}"`
53+
}
54+
55+
if (tool === "todowrite") {
56+
return `${parameters.todos?.length || 0} todos`
57+
}
58+
if (tool === "todoread") {
59+
return "read todo list"
60+
}
61+
62+
if (tool === "task" && parameters.description) {
63+
return parameters.description
64+
}
65+
66+
const paramStr = JSON.stringify(parameters)
67+
if (paramStr === '{}' || paramStr === '[]' || paramStr === 'null') {
68+
return ''
69+
}
70+
return paramStr.substring(0, 50)
71+
}
72+
73+
export const getLastUserMessage = (
74+
messages: WithParts[]
75+
): WithParts | null => {
76+
for (let i = messages.length - 1; i >= 0; i--) {
77+
const msg = messages[i]
78+
if (msg.info.role === 'user') {
79+
return msg
80+
}
81+
}
82+
return null
83+
}

‎lib/prune.ts‎

Lines changed: 0 additions & 21 deletions
This file was deleted.

‎lib/state/state.ts‎

Lines changed: 6 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,8 @@ export function createSessionState(): SessionState {
99
toolIds: []
1010
},
1111
stats: {
12-
totalToolsPruned: 0,
13-
totalTokensSaved: 0,
14-
totalGCTokens: 0,
15-
totalGCTools: 0
16-
},
17-
gcPending: {
18-
tokensCollected: 0,
19-
toolsDeduped: 0
12+
pruneTokenCounter: 0,
13+
totalPruneTokens: 0,
2014
},
2115
toolParameters: new Map<string, ToolParameterEntry>()
2216
}
@@ -28,14 +22,8 @@ export function resetSessionState(state: SessionState): void {
2822
toolIds: []
2923
}
3024
state.stats = {
31-
totalToolsPruned: 0,
32-
totalTokensSaved: 0,
33-
totalGCTokens: 0,
34-
totalGCTools: 0
35-
}
36-
state.gcPending = {
37-
tokensCollected: 0,
38-
toolsDeduped: 0
25+
pruneTokenCounter: 0,
26+
totalPruneTokens: 0,
3927
}
4028
state.toolParameters.clear()
4129
}
@@ -64,9 +52,7 @@ export async function ensureSessionInitialized(
6452
toolIds: persisted.prune.toolIds || []
6553
}
6654
state.stats = {
67-
totalToolsPruned: persisted.stats.totalToolsPruned || 0,
68-
totalTokensSaved: persisted.stats.totalTokensSaved || 0,
69-
totalGCTokens: persisted.stats.totalGCTokens || 0,
70-
totalGCTools: persisted.stats.totalGCTools || 0
55+
pruneTokenCounter: persisted.stats?.pruneTokenCounter || 0,
56+
totalPruneTokens: persisted.stats?.totalPruneTokens || 0,
7157
}
7258
}

‎lib/strategies/deduplication.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ export const deduplicate = (
6868

6969
if (newPruneIds.length > 0) {
7070
state.prune.toolIds.push(...newPruneIds)
71-
logger.debug("dedulication", `Marked ${newPruneIds.length} duplicate tool calls for pruning`)
71+
logger.debug("deduplication", `Marked ${newPruneIds.length} duplicate tool calls for pruning`)
7272
}
7373
}
7474

‎lib/strategies/prune-tool.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { tool } from "@opencode-ai/plugin"
2-
import type { SessionState, ToolParameterEntry } from "../state"
2+
import type { SessionState, ToolParameterEntry, WithParts } from "../state"
33
import type { PluginConfig } from "../config"
44
import { findCurrentAgent, buildToolIdList, getPruneToolIds } from "../utils"
55
import { PruneReason, sendUnifiedNotification } from "../ui/notification"
@@ -64,9 +64,10 @@ export function createPruneTool(
6464
await ensureSessionInitialized(state, sessionId, logger)
6565

6666
// Fetch messages to calculate tokens and find current agent
67-
const messages = await client.session.messages({
67+
const messagesResponse = await client.session.messages({
6868
path: { id: sessionId }
6969
})
70+
const messages = messagesResponse.data || messagesResponse
7071

7172
const currentAgent: string | undefined = findCurrentAgent(messages)
7273
const toolIdList: string[] = buildToolIdList(messages)

‎lib/utils.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
import { Logger } from "./logger"
2+
import { WithParts } from "./state"
3+
14
/**
25
* Checks if a session is a subagent session by looking for a parentID.
36
*/
@@ -27,7 +30,7 @@ export function findCurrentAgent(messages: any[]): string | undefined {
2730
/**
2831
* Builds a list of tool call IDs from messages.
2932
*/
30-
export function buildToolIdList(messages: any[]): string[] {
33+
export function buildToolIdList(messages: WithParts[]): string[] {
3134
const toolIds: string[] = []
3235
for (const msg of messages) {
3336
if (msg.parts) {

0 commit comments

Comments
 (0)