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 <clio-agent@sisyphuslabs.ai>
This commit is contained in:
co-authored by
Sisyphus
parent
9e03512cb0
commit
83f4399e64
@@ -141,3 +141,57 @@ export function prunePreservingItems(options: PruneOptions): ModelMessage[] {
|
|||||||
const pruned = pruneMessages(options);
|
const pruned = pruneMessages(options);
|
||||||
return dropOrphanedResults(detachOrphanedItems(options.messages, pruned));
|
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 };
|
||||||
|
|||||||
+41
-7
@@ -15,7 +15,7 @@ import { Notebook, type NotebookState } from './notebook';
|
|||||||
import { Permissions, type PermissionConfig } from './permission';
|
import { Permissions, type PermissionConfig } from './permission';
|
||||||
import type { PluginHost } from './plugins';
|
import type { PluginHost } from './plugins';
|
||||||
import { systemPrompt } from './prompt';
|
import { systemPrompt } from './prompt';
|
||||||
import { prunePreservingItems } from './prune';
|
import { pruneToFit } from './prune';
|
||||||
import { createSkillTool, renderSkills, type Skill } from './skills';
|
import { createSkillTool, renderSkills, type Skill } from './skills';
|
||||||
import { disabledToolNames, onBashOutput, tools as builtinTools, type ToolSetName } from './tools';
|
import { disabledToolNames, onBashOutput, tools as builtinTools, type ToolSetName } from './tools';
|
||||||
|
|
||||||
@@ -29,6 +29,8 @@ export type ApprovalRequest = {
|
|||||||
suggestedPattern: string;
|
suggestedPattern: string;
|
||||||
/** Set when the call is being asked about because it repeated, not because of a rule. */
|
/** Set when the call is being asked about because it repeated, not because of a rule. */
|
||||||
repeated?: boolean;
|
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. */
|
/** '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<boolean> {
|
||||||
|
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 {
|
setModel(model: LanguageModel): void {
|
||||||
this.model = model;
|
this.model = model;
|
||||||
}
|
}
|
||||||
@@ -318,6 +353,7 @@ export class Session {
|
|||||||
): AsyncGenerator<AgentEvent> {
|
): AsyncGenerator<AgentEvent> {
|
||||||
// Each iteration is one model run. A run ends either finished, or suspended
|
// 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.
|
// on tool approvals, in which case we collect decisions and run again.
|
||||||
|
let compactionReported = false;
|
||||||
while (true) {
|
while (true) {
|
||||||
const pending: ApprovalRequest[] = [];
|
const pending: ApprovalRequest[] = [];
|
||||||
const compactions: Extract<AgentEvent, { type: 'compacted' }>[] = [];
|
const compactions: Extract<AgentEvent, { type: 'compacted' }>[] = [];
|
||||||
@@ -341,14 +377,12 @@ export class Session {
|
|||||||
// visible to the steps that follow it, not only to the next turn.
|
// visible to the steps that follow it, not only to the next turn.
|
||||||
const instructions = this.systemFor();
|
const instructions = this.systemFor();
|
||||||
if (estimateTokens(messages) <= threshold) return { instructions };
|
if (estimateTokens(messages) <= threshold) return { instructions };
|
||||||
const pruned = prunePreservingItems({
|
const pruned = pruneToFit({ messages, threshold, estimate: estimateTokens });
|
||||||
messages,
|
|
||||||
reasoning: 'all',
|
|
||||||
toolCalls: 'before-last-3-messages',
|
|
||||||
emptyMessages: 'remove',
|
|
||||||
});
|
|
||||||
// prepareStep cannot yield, so queue the notice and drain it in the loop.
|
// prepareStep cannot yield, so queue the notice and drain it in the loop.
|
||||||
|
if (!compactionReported) {
|
||||||
compactions.push({ type: 'compacted', before: messages.length, after: pruned.length });
|
compactions.push({ type: 'compacted', before: messages.length, after: pruned.length });
|
||||||
|
compactionReported = true;
|
||||||
|
}
|
||||||
return { instructions, messages: pruned };
|
return { instructions, messages: pruned };
|
||||||
},
|
},
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -252,6 +252,30 @@ test('a compacted turn still shows the model the tool calls it already made', as
|
|||||||
}
|
}
|
||||||
}), 20_000);
|
}), 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 () =>
|
test('a compacted turn sends no assistant item reference whose reasoning was pruned', async () =>
|
||||||
inTempDir(async () => {
|
inTempDir(async () => {
|
||||||
await Bun.write(join(process.cwd(), 'big.txt'), 'lorem ipsum dolor sit amet\n'.repeat(1500));
|
await Bun.write(join(process.cwd(), 'big.txt'), 'lorem ipsum dolor sit amet\n'.repeat(1500));
|
||||||
|
|||||||
+107
-1
@@ -1,6 +1,6 @@
|
|||||||
import { expect, test } from 'bun:test';
|
import { expect, test } from 'bun:test';
|
||||||
import type { ModelMessage } from 'ai';
|
import type { ModelMessage } from 'ai';
|
||||||
import { detachOrphanedItems, dropOrphanedResults, prunePreservingItems } from '../src/prune';
|
import { detachOrphanedItems, dropOrphanedResults, pruneToFit, prunePreservingItems } from '../src/prune';
|
||||||
|
|
||||||
const kinds = (messages: ModelMessage[]) =>
|
const kinds = (messages: ModelMessage[]) =>
|
||||||
messages.map((m) => (Array.isArray(m.content) ? `${m.role}:${m.content.map((p) => p.type).join('+')}` : m.role));
|
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<string>();
|
||||||
|
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');
|
||||||
|
});
|
||||||
|
|||||||
Reference in New Issue
Block a user