Skip to content

Commit 9e741e3

Browse files
TarquinenTarquinen
committed
feat: replace janitor LLM with direct tool-based pruning
The session AI now decides what to prune via numeric IDs instead of a separate janitor LLM call. - Add numeric ID mapping system for simple tool references (1, 2, 3...) - Inject <prunable-tools> list at end of conversation - Prune tool accepts array of IDs to remove - Remove janitor LLM, model-selector, and pruning prompt
1 parent c10328e commit 9e741e3

File tree

18 files changed

+721
-746
lines changed

18 files changed

+721
-746
lines changed

‎index.ts‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,14 @@ const plugin: Plugin = (async (ctx) => {
9090
event: createEventHandler(ctx.client, janitorCtx, logger, config, toolTracker),
9191
"chat.params": createChatParamsHandler(ctx.client, state, logger),
9292
tool: config.strategies.onTool.length > 0 ? {
93-
prune: createPruningTool(ctx.client, janitorCtx, config, toolTracker),
93+
prune: createPruningTool({
94+
client: ctx.client,
95+
state,
96+
logger,
97+
config,
98+
notificationCtx: janitorCtx.notificationCtx,
99+
workingDirectory: ctx.directory
100+
}, toolTracker),
94101
} : undefined,
95102
}
96103
}) satisfies Plugin

‎lib/api-formats/prunable-list.ts‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
/**
2+
* Prunable tools list and nudge injection for DCP.
3+
*
4+
* Builds and injects a single message at the end of the conversation containing:
5+
* - Nudge instruction (when toolResultCount > nudge_freq)
6+
* - Prunable tools list
7+
*
8+
* Note: The base synthetic instructions (signal_management, context_window_management,
9+
* context_pruning) are still appended to the last user message separately via
10+
* synth-instruction.ts - that behavior is unchanged.
11+
*/
12+
13+
import { extractParameterKey } from '../ui/display-utils'
14+
import { getOrCreateNumericId, getNumericId } from '../state/id-mapping'
15+
16+
export interface ToolMetadata {
17+
tool: string
18+
parameters?: any
19+
}
20+
21+
// ============================================================================
22+
// Prompt Content
23+
// ============================================================================
24+
25+
const SYSTEM_REMINDER = `<system-reminder>
26+
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.
27+
</system-reminder>`
28+
29+
const NUDGE_INSTRUCTION = `<instruction name=agent_nudge>
30+
You have accumulated several tool outputs. If you have completed a discrete unit of work and distilled relevant understanding in writing for the user to keep, use the prune tool to remove obsolete tool outputs from this conversation and optimize token usage.
31+
</instruction>`
32+
33+
// ============================================================================
34+
// List Building
35+
// ============================================================================
36+
37+
export interface PrunableListResult {
38+
list: string
39+
numericIds: number[]
40+
}
41+
42+
/**
43+
* Builds the prunable tools list section.
44+
* Returns both the formatted list and the numeric IDs for logging.
45+
*/
46+
export function buildPrunableToolsList(
47+
sessionId: string,
48+
unprunedToolCallIds: string[],
49+
toolMetadata: Map<string, ToolMetadata>,
50+
protectedTools: string[]
51+
): PrunableListResult {
52+
const lines: string[] = []
53+
const numericIds: number[] = []
54+
55+
for (const actualId of unprunedToolCallIds) {
56+
const metadata = toolMetadata.get(actualId)
57+
58+
// Skip if no metadata or if tool is protected
59+
if (!metadata) continue
60+
if (protectedTools.includes(metadata.tool)) continue
61+
62+
// Get or create numeric ID for this tool call
63+
const numericId = getOrCreateNumericId(sessionId, actualId)
64+
numericIds.push(numericId)
65+
66+
// Format: "1: read, src/components/Button.tsx"
67+
const paramKey = extractParameterKey(metadata)
68+
const description = paramKey ? `${metadata.tool}, ${paramKey}` : metadata.tool
69+
lines.push(`${numericId}: ${description}`)
70+
}
71+
72+
if (lines.length === 0) {
73+
return { list: '', numericIds: [] }
74+
}
75+
76+
return {
77+
list: `<prunable-tools>\n${lines.join('\n')}\n</prunable-tools>`,
78+
numericIds
79+
}
80+
}
81+
82+
/**
83+
* Builds the end-of-conversation injection message.
84+
* Contains the system reminder, nudge (if active), and the prunable tools list.
85+
*
86+
* @param prunableList - The prunable tools list string (or empty string if none)
87+
* @param includeNudge - Whether to include the nudge instruction
88+
* @returns The injection string, or empty string if nothing to inject
89+
*/
90+
export function buildEndInjection(
91+
prunableList: string,
92+
includeNudge: boolean
93+
): string {
94+
// If no prunable tools, don't inject anything
95+
if (!prunableList) {
96+
return ''
97+
}
98+
99+
const parts = [SYSTEM_REMINDER]
100+
101+
if (includeNudge) {
102+
parts.push(NUDGE_INSTRUCTION)
103+
}
104+
105+
parts.push(prunableList)
106+
107+
return parts.join('\n\n')
108+
}
109+
110+
/**
111+
* Gets the numeric IDs for a list of actual tool call IDs.
112+
* Used when the prune tool needs to show what was pruned.
113+
*/
114+
export function getNumericIdsForActual(
115+
sessionId: string,
116+
actualIds: string[]
117+
): number[] {
118+
return actualIds
119+
.map(id => getNumericId(sessionId, id))
120+
.filter((id): id is number => id !== undefined)
121+
}
122+
123+
// ============================================================================
124+
// Injection Functions
125+
// ============================================================================
126+
127+
// ============================================================================
128+
// OpenAI Chat / Anthropic Format
129+
// ============================================================================
130+
131+
/**
132+
* Injects the prunable list (and optionally nudge) at the end of OpenAI/Anthropic messages.
133+
* Appends a new user message at the end.
134+
*/
135+
export function injectPrunableList(
136+
messages: any[],
137+
injection: string
138+
): boolean {
139+
if (!injection) return false
140+
messages.push({ role: 'user', content: injection })
141+
return true
142+
}
143+
144+
// ============================================================================
145+
// Google/Gemini Format
146+
// ============================================================================
147+
148+
/**
149+
* Injects the prunable list (and optionally nudge) at the end of Gemini contents.
150+
* Appends a new user content at the end.
151+
*/
152+
export function injectPrunableListGemini(
153+
contents: any[],
154+
injection: string
155+
): boolean {
156+
if (!injection) return false
157+
contents.push({ role: 'user', parts: [{ text: injection }] })
158+
return true
159+
}
160+
161+
// ============================================================================
162+
// OpenAI Responses API Format
163+
// ============================================================================
164+
165+
/**
166+
* Injects the prunable list (and optionally nudge) at the end of OpenAI Responses API input.
167+
* Appends a new user message at the end.
168+
*/
169+
export function injectPrunableListResponses(
170+
input: any[],
171+
injection: string
172+
): boolean {
173+
if (!injection) return false
174+
input.push({ type: 'message', role: 'user', content: injection })
175+
return true
176+
}

‎lib/api-formats/synth-instruction.ts‎

Lines changed: 44 additions & 121 deletions
Original file line numberDiff line numberDiff line change
@@ -14,77 +14,59 @@ export function resetToolTrackerCount(tracker: ToolTracker): void {
1414
tracker.toolResultCount = 0
1515
}
1616

17-
/** Adapter interface for format-specific message operations */
18-
interface MessageFormatAdapter {
19-
countToolResults(messages: any[], tracker: ToolTracker): number
20-
appendNudge(messages: any[], nudgeText: string): void
17+
/**
18+
* Counts total tool results in OpenAI/Anthropic messages (without tracker).
19+
* Used for determining if nudge threshold is met.
20+
*/
21+
export function countToolResults(messages: any[]): number {
22+
let count = 0
23+
for (const m of messages) {
24+
if (m.role === 'tool') {
25+
count++
26+
} else if (m.role === 'user' && Array.isArray(m.content)) {
27+
for (const part of m.content) {
28+
if (part.type === 'tool_result') {
29+
count++
30+
}
31+
}
32+
}
33+
}
34+
return count
35+
}
36+
37+
/**
38+
* Counts total tool results in Gemini contents (without tracker).
39+
*/
40+
export function countToolResultsGemini(contents: any[]): number {
41+
let count = 0
42+
for (const content of contents) {
43+
if (!Array.isArray(content.parts)) continue
44+
for (const part of content.parts) {
45+
if (part.functionResponse) {
46+
count++
47+
}
48+
}
49+
}
50+
return count
2151
}
2252

23-
/** Generic nudge injection - nudges every fetch once tools since last prune exceeds freq */
24-
function injectNudgeCore(
25-
messages: any[],
26-
tracker: ToolTracker,
27-
nudgeText: string,
28-
freq: number,
29-
adapter: MessageFormatAdapter
30-
): boolean {
31-
// Count any new tool results
32-
adapter.countToolResults(messages, tracker)
33-
34-
// Once we've exceeded the threshold, nudge on every fetch
35-
if (tracker.toolResultCount > freq) {
36-
adapter.appendNudge(messages, nudgeText)
37-
return true
53+
/**
54+
* Counts total tool results in OpenAI Responses API input (without tracker).
55+
*/
56+
export function countToolResultsResponses(input: any[]): number {
57+
let count = 0
58+
for (const item of input) {
59+
if (item.type === 'function_call_output') {
60+
count++
61+
}
3862
}
39-
return false
63+
return count
4064
}
4165

4266
// ============================================================================
4367
// OpenAI Chat / Anthropic Format
4468
// ============================================================================
4569

46-
const openaiAdapter: MessageFormatAdapter = {
47-
countToolResults(messages, tracker) {
48-
let newCount = 0
49-
for (const m of messages) {
50-
if (m.role === 'tool' && m.tool_call_id) {
51-
const id = String(m.tool_call_id).toLowerCase()
52-
if (!tracker.seenToolResultIds.has(id)) {
53-
tracker.seenToolResultIds.add(id)
54-
newCount++
55-
const toolName = m.name || tracker.getToolName?.(m.tool_call_id)
56-
if (toolName !== 'prune') {
57-
tracker.skipNextIdle = false
58-
}
59-
}
60-
} else if (m.role === 'user' && Array.isArray(m.content)) {
61-
for (const part of m.content) {
62-
if (part.type === 'tool_result' && part.tool_use_id) {
63-
const id = String(part.tool_use_id).toLowerCase()
64-
if (!tracker.seenToolResultIds.has(id)) {
65-
tracker.seenToolResultIds.add(id)
66-
newCount++
67-
const toolName = tracker.getToolName?.(part.tool_use_id)
68-
if (toolName !== 'prune') {
69-
tracker.skipNextIdle = false
70-
}
71-
}
72-
}
73-
}
74-
}
75-
}
76-
tracker.toolResultCount += newCount
77-
return newCount
78-
},
79-
appendNudge(messages, nudgeText) {
80-
messages.push({ role: 'user', content: nudgeText })
81-
}
82-
}
83-
84-
export function injectNudge(messages: any[], tracker: ToolTracker, nudgeText: string, freq: number): boolean {
85-
return injectNudgeCore(messages, tracker, nudgeText, freq, openaiAdapter)
86-
}
87-
8870
/** Check if a message content matches nudge text (OpenAI/Anthropic format) */
8971
function isNudgeMessage(msg: any, nudgeText: string): boolean {
9072
if (typeof msg.content === 'string') {
@@ -120,37 +102,6 @@ export function injectSynth(messages: any[], instruction: string, nudgeText: str
120102
// Google/Gemini Format (body.contents with parts)
121103
// ============================================================================
122104

123-
const geminiAdapter: MessageFormatAdapter = {
124-
countToolResults(contents, tracker) {
125-
let newCount = 0
126-
for (const content of contents) {
127-
if (!Array.isArray(content.parts)) continue
128-
for (const part of content.parts) {
129-
if (part.functionResponse) {
130-
const funcName = part.functionResponse.name?.toLowerCase() || 'unknown'
131-
const pseudoId = `gemini:${funcName}:${tracker.seenToolResultIds.size}`
132-
if (!tracker.seenToolResultIds.has(pseudoId)) {
133-
tracker.seenToolResultIds.add(pseudoId)
134-
newCount++
135-
if (funcName !== 'prune') {
136-
tracker.skipNextIdle = false
137-
}
138-
}
139-
}
140-
}
141-
}
142-
tracker.toolResultCount += newCount
143-
return newCount
144-
},
145-
appendNudge(contents, nudgeText) {
146-
contents.push({ role: 'user', parts: [{ text: nudgeText }] })
147-
}
148-
}
149-
150-
export function injectNudgeGemini(contents: any[], tracker: ToolTracker, nudgeText: string, freq: number): boolean {
151-
return injectNudgeCore(contents, tracker, nudgeText, freq, geminiAdapter)
152-
}
153-
154105
/** Check if a Gemini content matches nudge text */
155106
function isNudgeContentGemini(content: any, nudgeText: string): boolean {
156107
if (Array.isArray(content.parts) && content.parts.length === 1) {
@@ -182,34 +133,6 @@ export function injectSynthGemini(contents: any[], instruction: string, nudgeTex
182133
// OpenAI Responses API Format (body.input with type-based items)
183134
// ============================================================================
184135

185-
const responsesAdapter: MessageFormatAdapter = {
186-
countToolResults(input, tracker) {
187-
let newCount = 0
188-
for (const item of input) {
189-
if (item.type === 'function_call_output' && item.call_id) {
190-
const id = String(item.call_id).toLowerCase()
191-
if (!tracker.seenToolResultIds.has(id)) {
192-
tracker.seenToolResultIds.add(id)
193-
newCount++
194-
const toolName = item.name || tracker.getToolName?.(item.call_id)
195-
if (toolName !== 'prune') {
196-
tracker.skipNextIdle = false
197-
}
198-
}
199-
}
200-
}
201-
tracker.toolResultCount += newCount
202-
return newCount
203-
},
204-
appendNudge(input, nudgeText) {
205-
input.push({ type: 'message', role: 'user', content: nudgeText })
206-
}
207-
}
208-
209-
export function injectNudgeResponses(input: any[], tracker: ToolTracker, nudgeText: string, freq: number): boolean {
210-
return injectNudgeCore(input, tracker, nudgeText, freq, responsesAdapter)
211-
}
212-
213136
/** Check if a Responses API item matches nudge text */
214137
function isNudgeItemResponses(item: any, nudgeText: string): boolean {
215138
if (typeof item.content === 'string') {

0 commit comments

Comments
 (0)