Skip to content

Commit 7603d49

Browse files
TarquinenTarquinen
authored
Merge pull request #74 from Tarquinen/refactor/consolidate-tool-tracking
refactor: consolidate tool tracking to OpenCode API
2 parents c14ca14 + 585fa58 commit 7603d49

File tree

9 files changed

+23
-138
lines changed

9 files changed

+23
-138
lines changed

‎index.ts‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -43,12 +43,6 @@ const plugin: Plugin = (async (ctx) => {
4343
// Create tool tracker and load prompts for synthetic instruction injection
4444
const toolTracker = createToolTracker()
4545

46-
// Wire up tool name lookup from the cached tool parameters
47-
toolTracker.getToolName = (callId: string) => {
48-
const entry = state.toolParameters.get(callId.toLowerCase())
49-
return entry?.tool
50-
}
51-
5246
const prompts = {
5347
synthInstruction: loadPrompt("synthetic"),
5448
nudgeInstruction: loadPrompt("nudge")

‎lib/fetch-wrapper/formats/bedrock.ts‎

Lines changed: 1 addition & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { FormatDescriptor, ToolOutput, ToolTracker } from "../types"
1+
import type { FormatDescriptor, ToolOutput } from "../types"
22
import type { PluginState } from "../../state"
33

44
function isNudgeMessage(msg: any, nudgeText: string): boolean {
@@ -30,36 +30,6 @@ function injectSynth(messages: any[], instruction: string, nudgeText: string): b
3030
return false
3131
}
3232

33-
function trackNewToolResults(messages: any[], tracker: ToolTracker, protectedTools: Set<string>): number {
34-
let newCount = 0
35-
for (const m of messages) {
36-
if (m.role === 'tool' && m.tool_call_id) {
37-
if (!tracker.seenToolResultIds.has(m.tool_call_id)) {
38-
tracker.seenToolResultIds.add(m.tool_call_id)
39-
const toolName = tracker.getToolName?.(m.tool_call_id)
40-
if (!toolName || !protectedTools.has(toolName)) {
41-
tracker.toolResultCount++
42-
newCount++
43-
}
44-
}
45-
} else if (m.role === 'user' && Array.isArray(m.content)) {
46-
for (const part of m.content) {
47-
if (part.type === 'tool_result' && part.tool_use_id) {
48-
if (!tracker.seenToolResultIds.has(part.tool_use_id)) {
49-
tracker.seenToolResultIds.add(part.tool_use_id)
50-
const toolName = tracker.getToolName?.(part.tool_use_id)
51-
if (!toolName || !protectedTools.has(toolName)) {
52-
tracker.toolResultCount++
53-
newCount++
54-
}
55-
}
56-
}
57-
}
58-
}
59-
}
60-
return newCount
61-
}
62-
6333
function injectPrunableList(messages: any[], injection: string): boolean {
6434
if (!injection) return false
6535
messages.push({ role: 'user', content: injection })
@@ -90,10 +60,6 @@ export const bedrockFormat: FormatDescriptor = {
9060
return injectSynth(data, instruction, nudgeText)
9161
},
9262

93-
trackNewToolResults(data: any[], tracker: ToolTracker, protectedTools: Set<string>): number {
94-
return trackNewToolResults(data, tracker, protectedTools)
95-
},
96-
9763
injectPrunableList(data: any[], injection: string): boolean {
9864
return injectPrunableList(data, injection)
9965
},

‎lib/fetch-wrapper/formats/gemini.ts‎

Lines changed: 1 addition & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { FormatDescriptor, ToolOutput, ToolTracker } from "../types"
1+
import type { FormatDescriptor, ToolOutput } from "../types"
22
import type { PluginState } from "../../state"
33

44
function isNudgeContent(content: any, nudgeText: string): boolean {
@@ -26,29 +26,6 @@ function injectSynth(contents: any[], instruction: string, nudgeText: string): b
2626
return false
2727
}
2828

29-
function trackNewToolResults(contents: any[], tracker: ToolTracker, protectedTools: Set<string>): number {
30-
let newCount = 0
31-
let positionCounter = 0
32-
for (const content of contents) {
33-
if (!Array.isArray(content.parts)) continue
34-
for (const part of content.parts) {
35-
if (part.functionResponse) {
36-
const positionId = `gemini_pos_${positionCounter}`
37-
positionCounter++
38-
if (!tracker.seenToolResultIds.has(positionId)) {
39-
tracker.seenToolResultIds.add(positionId)
40-
const toolName = part.functionResponse.name
41-
if (!toolName || !protectedTools.has(toolName)) {
42-
tracker.toolResultCount++
43-
newCount++
44-
}
45-
}
46-
}
47-
}
48-
}
49-
return newCount
50-
}
51-
5229
function injectPrunableList(contents: any[], injection: string): boolean {
5330
if (!injection) return false
5431
contents.push({ role: 'user', parts: [{ text: injection }] })
@@ -75,10 +52,6 @@ export const geminiFormat: FormatDescriptor = {
7552
return injectSynth(data, instruction, nudgeText)
7653
},
7754

78-
trackNewToolResults(data: any[], tracker: ToolTracker, protectedTools: Set<string>): number {
79-
return trackNewToolResults(data, tracker, protectedTools)
80-
},
81-
8255
injectPrunableList(data: any[], injection: string): boolean {
8356
return injectPrunableList(data, injection)
8457
},

‎lib/fetch-wrapper/formats/openai-chat.ts‎

Lines changed: 1 addition & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { FormatDescriptor, ToolOutput, ToolTracker } from "../types"
1+
import type { FormatDescriptor, ToolOutput } from "../types"
22
import type { PluginState } from "../../state"
33

44
function isNudgeMessage(msg: any, nudgeText: string): boolean {
@@ -30,36 +30,6 @@ function injectSynth(messages: any[], instruction: string, nudgeText: string): b
3030
return false
3131
}
3232

33-
function trackNewToolResults(messages: any[], tracker: ToolTracker, protectedTools: Set<string>): number {
34-
let newCount = 0
35-
for (const m of messages) {
36-
if (m.role === 'tool' && m.tool_call_id) {
37-
if (!tracker.seenToolResultIds.has(m.tool_call_id)) {
38-
tracker.seenToolResultIds.add(m.tool_call_id)
39-
const toolName = tracker.getToolName?.(m.tool_call_id)
40-
if (!toolName || !protectedTools.has(toolName)) {
41-
tracker.toolResultCount++
42-
newCount++
43-
}
44-
}
45-
} else if (m.role === 'user' && Array.isArray(m.content)) {
46-
for (const part of m.content) {
47-
if (part.type === 'tool_result' && part.tool_use_id) {
48-
if (!tracker.seenToolResultIds.has(part.tool_use_id)) {
49-
tracker.seenToolResultIds.add(part.tool_use_id)
50-
const toolName = tracker.getToolName?.(part.tool_use_id)
51-
if (!toolName || !protectedTools.has(toolName)) {
52-
tracker.toolResultCount++
53-
newCount++
54-
}
55-
}
56-
}
57-
}
58-
}
59-
}
60-
return newCount
61-
}
62-
6333
function injectPrunableList(messages: any[], injection: string): boolean {
6434
if (!injection) return false
6535
messages.push({ role: 'user', content: injection })
@@ -81,10 +51,6 @@ export const openaiChatFormat: FormatDescriptor = {
8151
return injectSynth(data, instruction, nudgeText)
8252
},
8353

84-
trackNewToolResults(data: any[], tracker: ToolTracker, protectedTools: Set<string>): number {
85-
return trackNewToolResults(data, tracker, protectedTools)
86-
},
87-
8854
injectPrunableList(data: any[], injection: string): boolean {
8955
return injectPrunableList(data, injection)
9056
},

‎lib/fetch-wrapper/formats/openai-responses.ts‎

Lines changed: 1 addition & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { FormatDescriptor, ToolOutput, ToolTracker } from "../types"
1+
import type { FormatDescriptor, ToolOutput } from "../types"
22
import type { PluginState } from "../../state"
33

44
function isNudgeItem(item: any, nudgeText: string): boolean {
@@ -30,23 +30,6 @@ function injectSynth(input: any[], instruction: string, nudgeText: string): bool
3030
return false
3131
}
3232

33-
function trackNewToolResults(input: any[], tracker: ToolTracker, protectedTools: Set<string>): number {
34-
let newCount = 0
35-
for (const item of input) {
36-
if (item.type === 'function_call_output' && item.call_id) {
37-
if (!tracker.seenToolResultIds.has(item.call_id)) {
38-
tracker.seenToolResultIds.add(item.call_id)
39-
const toolName = tracker.getToolName?.(item.call_id)
40-
if (!toolName || !protectedTools.has(toolName)) {
41-
tracker.toolResultCount++
42-
newCount++
43-
}
44-
}
45-
}
46-
}
47-
return newCount
48-
}
49-
5033
function injectPrunableList(input: any[], injection: string): boolean {
5134
if (!injection) return false
5235
input.push({ type: 'message', role: 'user', content: injection })
@@ -68,10 +51,6 @@ export const openaiResponsesFormat: FormatDescriptor = {
6851
return injectSynth(data, instruction, nudgeText)
6952
},
7053

71-
trackNewToolResults(data: any[], tracker: ToolTracker, protectedTools: Set<string>): number {
72-
return trackNewToolResults(data, tracker, protectedTools)
73-
},
74-
7554
injectPrunableList(data: any[], injection: string): boolean {
7655
return injectPrunableList(data, injection)
7756
},

‎lib/fetch-wrapper/handler.ts‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -67,9 +67,11 @@ export async function handleFormat(
6767
let modified = false
6868

6969
// Sync tool parameters from OpenCode's session API (single source of truth)
70+
// Also tracks new tool results for nudge injection
7071
const sessionId = ctx.state.lastSeenSessionId
72+
const protectedSet = new Set(ctx.config.protectedTools)
7173
if (sessionId) {
72-
await syncToolParametersFromOpenCode(ctx.client, sessionId, ctx.state, ctx.logger)
74+
await syncToolParametersFromOpenCode(ctx.client, sessionId, ctx.state, ctx.toolTracker, protectedSet, ctx.logger)
7375
}
7476

7577
if (ctx.config.strategies.onTool.length > 0) {
@@ -91,8 +93,6 @@ export async function handleFormat(
9193
)
9294

9395
if (prunableList) {
94-
const protectedSet = new Set(ctx.config.protectedTools)
95-
format.trackNewToolResults(data, ctx.toolTracker, protectedSet)
9696
const includeNudge = ctx.config.nudge_freq > 0 && ctx.toolTracker.toolResultCount > ctx.config.nudge_freq
9797

9898
const endInjection = buildEndInjection(prunableList, includeNudge)
@@ -119,14 +119,12 @@ export async function handleFormat(
119119
}
120120

121121
const toolOutputs = format.extractToolOutputs(data, ctx.state)
122-
const protectedToolsLower = new Set(ctx.config.protectedTools.map(t => t.toLowerCase()))
123122
let replacedCount = 0
124123
let prunableCount = 0
125124

126125
for (const output of toolOutputs) {
127-
if (output.toolName && protectedToolsLower.has(output.toolName.toLowerCase())) {
128-
continue
129-
}
126+
// Skip tools not in cache (protected tools are excluded from cache)
127+
if (!output.toolName) continue
130128
prunableCount++
131129

132130
if (allPrunedIds.has(output.id)) {

‎lib/fetch-wrapper/tool-tracker.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ export interface ToolTracker {
22
seenToolResultIds: Set<string>
33
toolResultCount: number // Tools since last prune
44
skipNextIdle: boolean
5-
getToolName?: (callId: string) => string | undefined
65
}
76

87
export function createToolTracker(): ToolTracker {

‎lib/fetch-wrapper/types.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ export interface FormatDescriptor {
1414
detect(body: any): boolean
1515
getDataArray(body: any): any[] | undefined
1616
injectSynth(data: any[], instruction: string, nudgeText: string): boolean
17-
trackNewToolResults(data: any[], tracker: ToolTracker, protectedTools: Set<string>): number
1817
injectPrunableList(data: any[], injection: string): boolean
1918
extractToolOutputs(data: any[], state: PluginState): ToolOutput[]
2019
replaceToolOutput(data: any[], toolId: string, prunedMessage: string, state: PluginState): boolean

‎lib/state/tool-cache.ts‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import type { PluginState, ToolStatus } from "./index"
22
import type { Logger } from "../logger"
3+
import type { ToolTracker } from "../fetch-wrapper/tool-tracker"
34

45
/** Maximum number of entries to keep in the tool parameters cache */
56
const MAX_TOOL_CACHE_SIZE = 500
@@ -13,6 +14,8 @@ export async function syncToolParametersFromOpenCode(
1314
client: any,
1415
sessionId: string,
1516
state: PluginState,
17+
tracker?: ToolTracker,
18+
protectedTools?: Set<string>,
1619
logger?: Logger
1720
): Promise<void> {
1821
try {
@@ -36,8 +39,17 @@ export async function syncToolParametersFromOpenCode(
3639

3740
const id = part.callID.toLowerCase()
3841

39-
// Skip if already cached (optimization)
42+
// Track tool results for nudge injection
43+
if (tracker && !tracker.seenToolResultIds.has(id)) {
44+
tracker.seenToolResultIds.add(id)
45+
// Only count non-protected tools toward nudge threshold
46+
if (!part.tool || !protectedTools?.has(part.tool)) {
47+
tracker.toolResultCount++
48+
}
49+
}
50+
4051
if (state.toolParameters.has(id)) continue
52+
if (part.tool && protectedTools?.has(part.tool)) continue
4153

4254
const status = part.state?.status as ToolStatus | undefined
4355
state.toolParameters.set(id, {
@@ -55,8 +67,7 @@ export async function syncToolParametersFromOpenCode(
5567
if (logger && synced > 0) {
5668
logger.debug("tool-cache", "Synced tool parameters from OpenCode", {
5769
sessionId: sessionId.slice(0, 8),
58-
synced,
59-
totalCached: state.toolParameters.size
70+
synced
6071
})
6172
}
6273
} catch (error) {

0 commit comments

Comments
 (0)