From 83f4399e64adc5b24ae5489db398c8f4dcb51e29 Mon Sep 17 00:00:00 2001 From: Muhammad Zakir Ramadhan <61570975+zakirkun@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:42:02 +0700 Subject: [PATCH] Make compaction bounded and report it once per turn The fixed three-message tool window collapsed long transcripts to a handful of messages: a 405-message run kept two of 202 tool calls, and the model re-ran what it could no longer see. Pruning now drops reasoning first and keeps the widest recent tool tail that fits a ladder, the SDK carries that view into later steps, and the turn emits one compaction event instead of one per step. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src/prune.ts | 54 ++++++++++++++++++++++ src/session.ts | 50 ++++++++++++++++---- test/compact.test.ts | 24 ++++++++++ test/prune.test.ts | 108 ++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 227 insertions(+), 9 deletions(-) diff --git a/src/prune.ts b/src/prune.ts index 7af79de..7f23bef 100644 --- a/src/prune.ts +++ b/src/prune.ts @@ -141,3 +141,57 @@ export function prunePreservingItems(options: PruneOptions): ModelMessage[] { const pruned = pruneMessages(options); return dropOrphanedResults(detachOrphanedItems(options.messages, pruned)); } + +/** + * How many trailing messages keep their tool content, widest first. + * + * One agent step is two messages — the assistant's tool call and the tool message + * answering it — so 64 is about 32 steps of memory. + */ +const KEEP_LADDER = [64, 32, 16, 8, 4] as const; + +export type FitOptions = { + messages: ModelMessage[]; + /** Estimated tokens the wire history must come in under. */ + threshold: number; + estimate: (messages: ModelMessage[]) => number; +}; + +/** + * Prunes only as hard as the threshold requires. + * + * A fixed `before-last-3-messages` is catastrophic on an agent transcript, because + * nearly every assistant and tool message there consists of nothing but tool parts: + * stripping them empties the message, `emptyMessages: 'remove'` deletes it, and a + * 405-message history collapses to five. Measured on a synthetic run of 202 steps — + * two surviving tool calls out of 202. + * + * That is not a cost problem, it is a correctness one. The model loses its record of + * what it already ran, so it runs it again, the history grows, the threshold is + * crossed again, and the turn never converges. It looks like `git_status` and + * `list_dir` being called in a circle with a compaction notice between them. + * + * So: drop reasoning first, since it is never needed on the wire, and only reach for + * tool content if that was not enough — keeping as much of the recent tail as fits. + * The widest rung that comes in under the threshold wins; if even the narrowest does + * not, the narrowest is returned, because sending something is better than sending a + * request that will be rejected for size. + */ +export function pruneToFit({ messages, threshold, estimate }: FitOptions): ModelMessage[] { + const withoutReasoning = prunePreservingItems({ messages, reasoning: 'all', emptyMessages: 'remove' }); + if (estimate(withoutReasoning) <= threshold) return withoutReasoning; + + let narrowest = withoutReasoning; + for (const keep of KEEP_LADDER) { + narrowest = prunePreservingItems({ + messages, + reasoning: 'all', + toolCalls: `before-last-${keep}-messages`, + emptyMessages: 'remove', + }); + if (estimate(narrowest) <= threshold) return narrowest; + } + return narrowest; +} + +export { KEEP_LADDER }; diff --git a/src/session.ts b/src/session.ts index 504def6..5b8c96e 100644 --- a/src/session.ts +++ b/src/session.ts @@ -15,7 +15,7 @@ import { Notebook, type NotebookState } from './notebook'; import { Permissions, type PermissionConfig } from './permission'; import type { PluginHost } from './plugins'; import { systemPrompt } from './prompt'; -import { prunePreservingItems } from './prune'; +import { pruneToFit } from './prune'; import { createSkillTool, renderSkills, type Skill } from './skills'; import { disabledToolNames, onBashOutput, tools as builtinTools, type ToolSetName } from './tools'; @@ -29,6 +29,8 @@ export type ApprovalRequest = { suggestedPattern: string; /** Set when the call is being asked about because it repeated, not because of a rule. */ repeated?: boolean; + /** Set when a `worker` subagent is asking, not the main agent. */ + subagent?: boolean; }; /** 'once' runs this call only; 'always' whitelists the suggested pattern for the session. */ @@ -136,6 +138,39 @@ export class Session { }); } + /** + * The approval channel a `worker` subagent uses for its gated calls. + * + * Same rules, same prompt, same grants as a direct call: a subagent that could + * approve its own writes would be a way to launder a tool call past the user. + * Handed to `createTaskTool` from cli.tsx, which is where the two are wired. + */ + approveForSubagent(): (req: { toolName: string; input: unknown }) => Promise { + return async ({ toolName, input }) => { + const blocked = await this.opts.plugins?.guard({ + toolName, + input, + cwd: this.opts.cwd ?? process.cwd(), + }); + if (blocked) return false; + + const { decision, pattern } = this.permissions.check(toolName, input); + if (decision === 'deny') return false; + if (decision === 'allow') return true; + + const answer = await this.opts.askApproval({ + approvalId: `sub:${toolName}`, + toolName, + input, + ...(pattern ? { matchedPattern: pattern } : {}), + suggestedPattern: this.permissions.suggest(toolName, input), + subagent: true, + }); + if (answer === 'always') this.permissions.grant(toolName, this.permissions.suggest(toolName, input)); + return answer !== 'deny'; + }; + } + setModel(model: LanguageModel): void { this.model = model; } @@ -318,6 +353,7 @@ export class Session { ): AsyncGenerator { // Each iteration is one model run. A run ends either finished, or suspended // on tool approvals, in which case we collect decisions and run again. + let compactionReported = false; while (true) { const pending: ApprovalRequest[] = []; const compactions: Extract[] = []; @@ -341,14 +377,12 @@ export class Session { // visible to the steps that follow it, not only to the next turn. const instructions = this.systemFor(); if (estimateTokens(messages) <= threshold) return { instructions }; - const pruned = prunePreservingItems({ - messages, - reasoning: 'all', - toolCalls: 'before-last-3-messages', - emptyMessages: 'remove', - }); + const pruned = pruneToFit({ messages, threshold, estimate: estimateTokens }); // prepareStep cannot yield, so queue the notice and drain it in the loop. - compactions.push({ type: 'compacted', before: messages.length, after: pruned.length }); + if (!compactionReported) { + compactions.push({ type: 'compacted', before: messages.length, after: pruned.length }); + compactionReported = true; + } return { instructions, messages: pruned }; }, }); diff --git a/test/compact.test.ts b/test/compact.test.ts index de4e1ab..59366f1 100644 --- a/test/compact.test.ts +++ b/test/compact.test.ts @@ -252,6 +252,30 @@ test('a compacted turn still shows the model the tool calls it already made', as } }), 20_000); +test('a multi-step turn emits one compacted event', async () => + inTempDir(async () => { + await Bun.write(join(process.cwd(), 'big.txt'), 'lorem ipsum dolor sit amet\n'.repeat(1500)); + + let call = 0; + const session = new Session({ + compactThreshold: 4000, + maxSteps: 8, + model: new MockLanguageModelV4({ + doStream: async () => { + const n = call++; + return n < 3 ? stream(reasoningToolStep(n)) : stream(text('done')); + }, + }), + askApproval: async () => 'deny', + }); + + const events: AgentEvent[] = []; + for await (const ev of session.send('read big.txt a few times')) events.push(ev); + + expect(events.filter((ev) => ev.type === 'compacted')).toHaveLength(1); + expect(call).toBe(4); + }), 20_000); + test('a compacted turn sends no assistant item reference whose reasoning was pruned', async () => inTempDir(async () => { await Bun.write(join(process.cwd(), 'big.txt'), 'lorem ipsum dolor sit amet\n'.repeat(1500)); diff --git a/test/prune.test.ts b/test/prune.test.ts index 4eece1c..fce2ef2 100644 --- a/test/prune.test.ts +++ b/test/prune.test.ts @@ -1,6 +1,6 @@ import { expect, test } from 'bun:test'; import type { ModelMessage } from 'ai'; -import { detachOrphanedItems, dropOrphanedResults, prunePreservingItems } from '../src/prune'; +import { detachOrphanedItems, dropOrphanedResults, pruneToFit, prunePreservingItems } from '../src/prune'; const kinds = (messages: ModelMessage[]) => messages.map((m) => (Array.isArray(m.content) ? `${m.role}:${m.content.map((p) => p.type).join('+')}` : m.role)); @@ -322,3 +322,109 @@ test('prunePreservingItems never strands a tool result on the wire', () => { } } }); + +const estimate = (messages: ModelMessage[]) => Math.round(JSON.stringify(messages).length / 4); + +/** `size` chars of tool output per step, so a transcript's weight is controllable. */ +function transcript(steps: number, size: number): ModelMessage[] { + const messages: ModelMessage[] = [{ role: 'user', content: 'do the thing' }]; + for (let i = 0; i < steps; i++) { + messages.push({ + role: 'assistant', + content: [ + { type: 'reasoning', text: 'deciding', providerOptions: { openai: { itemId: `rs_${i}` } } }, + { type: 'tool-call', toolCallId: `t${i}`, toolName: 'read_file', input: { path: `f${i}.ts` } }, + ], + }); + messages.push({ + role: 'tool', + content: [ + { type: 'tool-result', toolCallId: `t${i}`, toolName: 'read_file', output: { type: 'text', value: 'x'.repeat(size) } }, + ], + }); + } + return messages; +} + +const toolCallsIn = (messages: ModelMessage[]): number => { + let n = 0; + for (const m of messages) { + if (!Array.isArray(m.content)) continue; + for (const p of m.content as { type: string }[]) if (p.type === 'tool-call') n++; + } + return n; +}; + +/** + * The loop this guards against: a fixed `before-last-3-messages` strips tool parts + * from every earlier message, `emptyMessages: 'remove'` then deletes the emptied + * messages, and a long agent transcript collapses to a handful. The model loses its + * record of what it ran and runs it again, which on screen is `git_status` and + * `list_dir` repeating with a compaction notice between them. + */ +test('a long transcript keeps most of its tool calls when reasoning alone is enough', () => { + const messages = transcript(200, 100); + const fitted = pruneToFit({ messages, threshold: estimate(messages), estimate }); + + expect(toolCallsIn(fitted)).toBe(200); + expect(fitted.length).toBeGreaterThan(messages.length - 10); +}); + +test('tool content is only dropped when dropping reasoning was not enough', () => { + const messages = transcript(200, 4000); + // Far below the transcript's own weight, so the ladder has to descend. + const fitted = pruneToFit({ messages, threshold: 20_000, estimate }); + + expect(estimate(fitted)).toBeLessThanOrEqual(20_000); + // The old behaviour left two. Anything in that range is the bug returning. + expect(toolCallsIn(fitted)).toBeGreaterThan(2); +}); + +test('the widest rung that fits is the one used', () => { + const messages = transcript(200, 300); + const wide = pruneToFit({ messages, threshold: estimate(messages), estimate }); + const narrow = pruneToFit({ messages, threshold: 5_000, estimate }); + + expect(toolCallsIn(wide)).toBeGreaterThan(toolCallsIn(narrow)); +}); + +test('an impossible threshold returns the narrowest rung rather than nothing', () => { + const messages = transcript(200, 4000); + const fitted = pruneToFit({ messages, threshold: 10, estimate }); + + expect(fitted.length).toBeGreaterThan(0); + expect(fitted[0]?.role).toBe('user'); +}); + +test('pruning to fit never strands a tool result, at any rung', () => { + const messages = transcript(120, 4000); + for (const threshold of [10, 5_000, 20_000, 100_000, estimate(messages)]) { + const fitted = pruneToFit({ messages, threshold, estimate }); + const calls = new Set(); + for (const m of fitted) { + if (!Array.isArray(m.content)) continue; + for (const p of m.content as { type: string; toolCallId?: string }[]) { + if (p.type === 'tool-call' && p.toolCallId) calls.add(p.toolCallId); + } + } + for (const m of fitted) { + if (!Array.isArray(m.content)) continue; + for (const p of m.content as { type: string; toolCallId?: string }[]) { + if (p.type === 'tool-result') expect(calls.has(p.toolCallId!), `threshold ${threshold}`).toBe(true); + } + } + } +}); + +test('reasoning is always dropped, whatever the threshold', () => { + const messages = transcript(20, 100); + const fitted = pruneToFit({ messages, threshold: estimate(messages), estimate }); + expect(itemIds(fitted)).toEqual([]); + expect(JSON.stringify(fitted)).not.toContain('deciding'); +}); + +test('the user prompt survives even the narrowest rung', () => { + const messages = transcript(200, 4000); + const fitted = pruneToFit({ messages, threshold: 100, estimate }); + expect(JSON.stringify(fitted)).toContain('do the thing'); +});