From a9fb1b239b77270e113fb3350cac2daf80ee5744 Mon Sep 17 00:00:00 2001 From: jonnyparris <6400000+jonnyparris@users.noreply.github.com> Date: Mon, 25 May 2026 09:26:32 +0100 Subject: [PATCH] feat(own-loop): in-memory tool-result prune harness safety net MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The harness fix for the failure mode that breaks Gemma + Kimi K2.6 as orchestrators on first-turn explorations. ## Problem Think persists the assistant message only AFTER streamText() completes the whole generator. On a long-running first user turn, sessions.getHistory() returns just [user] — realMessages.length === 1. The autocompaction summariser has nothing to summarise. Meanwhile the own-loop's closure variable `messages` accumulates tool-result entries across iterations. Weak orchestrators (Gemma 4 26B, Kimi K2.6) routinely make many tool calls in a single turn whose results compound the local array to 200k+ tokens. The pre-step hard-stop then kicks in before the model can finish, ending the session with "context budget exhausted". ## Fix New pure helper `pruneOversizedToolResults` in src/own-loop-prune.ts: - Walks the own-loop's in-memory messages array - Replaces the largest tool-result payloads with short placeholders - Preserves the tool-call envelope (toolName, toolCallId) so the model's pending tool-call still resolves - Prunes oldest-first, preserves the most recent N (default 2) - Stops as soon as projected tokens are back under target Wired into the pre-step budget check in onChatMessage's own-loop, running AFTER the Think-level compaction attempt and BEFORE the hard-stop. The order is intentional: compaction handles previous turns (cheap when it works), prune handles the current turn (the case compaction can't fix). No LLM calls, no storage writes. The placeholder tells the model the data was pruned and to re-run the tool if it still needs it. ## Why this is the right fix PR #73 made autocompaction reachable on the first turn. PR #74 made it actually call into the summariser with force=true. This PR closes the remaining gap: compaction's input source (Think's persisted history) is structurally empty on turn 1; the own-loop's in-memory history is where the bytes actually live. ## Tests - 7 new unit tests for pruneOversizedToolResults covering: under-budget no-op, oldest-largest first, preservation of recent N, envelope preservation, min-payload-chars floor, early stop, no tool messages - All 21 existing compaction tests still pass - Typecheck clean --- src/coding-agent.ts | 41 ++++++- src/own-loop-prune.ts | 193 +++++++++++++++++++++++++++++ test/own-loop-prune-unit.test.ts | 200 +++++++++++++++++++++++++++++++ 3 files changed, 430 insertions(+), 4 deletions(-) create mode 100644 src/own-loop-prune.ts create mode 100644 test/own-loop-prune-unit.test.ts diff --git a/src/coding-agent.ts b/src/coding-agent.ts index 52da1fd..f1f0c62 100644 --- a/src/coding-agent.ts +++ b/src/coding-agent.ts @@ -55,6 +55,7 @@ import { type WatchdogConfig, } from "./watchdog"; import { shouldCompact, pickCutoff } from "./compaction-policy"; +import { pruneOversizedToolResults } from "./own-loop-prune"; import { nextRetry } from "./overflow-retry"; import { assembleSystemPrompt } from "./prompt-composer"; import { evaluateSession } from "./watchdog-policy"; @@ -1130,16 +1131,48 @@ export class CodingAgent extends Think { } } - // Hard stop if even after compaction we're projected above the budget. + // ─── In-memory tool-result prune (harness safety net) ─── + // Think-level compaction operates on persisted session history. + // On the very first user turn that's just `[user]`, so compaction + // had nothing to summarise and left `finalMessages` untouched. + // For weak orchestrators that fan out many tool calls inside a + // single turn, this is the failure mode that ends sessions + // prematurely with "context budget exhausted". + // + // This is a no-LLM, no-storage fallback: walk `finalMessages` + // and shorten the largest tool-result payloads (oldest first, + // preserving the most recent two) until projected tokens are + // back under the mid-loop threshold. Preserves the tool-call + // envelope (toolName, toolCallId) so the model's chain of + // pending tool calls still resolves. + const projectedAfterCompactionTokens = estimateMessagesTokens(finalMessages); + if (projectedAfterCompactionTokens / tokenBudget >= MID_LOOP_COMPACTION_THRESHOLD) { + const pruneResult = pruneOversizedToolResults(finalMessages, { + targetTokens: Math.floor(tokenBudget * MID_LOOP_COMPACTION_THRESHOLD), + estimate: estimateMessagesTokens, + }); + if (pruneResult.pruned) { + log("info", "own-loop: in-memory tool-result prune", { + sessionId, + step, + partsPruned: pruneResult.partsPruned, + bytesRemoved: pruneResult.bytesRemoved, + tokensBefore: pruneResult.tokensBefore, + tokensAfter: pruneResult.tokensAfter, + }); + } + } + + // Hard stop if even after compaction + prune we're projected above the budget. // Without this, streamText() is guaranteed to throw a context-overflow // error — better to exit cleanly with a wrap-up message. - const projectedAfterCompactionTokens = estimateMessagesTokens(finalMessages); - const projectedAfterUsage = projectedAfterCompactionTokens / tokenBudget; + const projectedAfterPruneTokens = estimateMessagesTokens(finalMessages); + const projectedAfterUsage = projectedAfterPruneTokens / tokenBudget; if (projectedAfterUsage >= HARD_STOP_THRESHOLD) { log("warn", "own-loop: pre-step hard stop — projected tokens exceed budget", { sessionId, step, - projectedInputTokens: projectedAfterCompactionTokens, + projectedInputTokens: projectedAfterPruneTokens, tokenBudget, projectedUsage: `${Math.round(projectedAfterUsage * 100)}%`, }); diff --git a/src/own-loop-prune.ts b/src/own-loop-prune.ts new file mode 100644 index 0000000..2085b65 --- /dev/null +++ b/src/own-loop-prune.ts @@ -0,0 +1,193 @@ +/** + * In-place pruning of the own-loop's in-memory `messages` array. + * + * **Why this exists.** + * The autocompaction safety nets in `coding-agent.ts` operate on Think's + * *persisted* session history. Think persists the assistant message only + * *after* `streamText()` completes — so during a long-running first user + * turn, `sessions.getHistory()` returns just `[user]`. There is nothing + * for the compaction summariser to summarise. + * + * Meanwhile, the own-loop accumulates tool-result messages in its closure + * variable `messages` across iterations. Weak orchestrators (Gemma 4 26B, + * Kimi K2.6) routinely make many tool calls in a single turn whose results + * compound the local array to 200k+ tokens. The pre-step budget check then + * kicks the model into wrap-up before any safety net engages. + * + * This module is the safety net for that case: a pure, no-LLM, no-storage + * function that walks the in-memory messages and shortens the biggest + * tool-result payloads until the projected budget is back under threshold. + * No state, no async, no surprises. + * + * **Design.** + * - Replace large tool-result `value`s with a short placeholder. Keep the + * `tool-result` envelope (toolCallId, toolName) intact so the model's + * tool-call/result pairing isn't broken — that would confuse the AI SDK + * and the model itself. + * - Prune oldest first. The freshest tool result is the one the model is + * about to react to; preserving it preserves the chain of thought. + * - Preserve the user message and the most recent N tool results + * regardless of size — they're load-bearing. + * - Operate on the array directly (mutate) so callers don't have to + * re-assign references. + */ + +import type { ModelMessage } from "ai"; + +/** Identifying shape of an AI SDK tool-result part. */ +interface ToolResultPart { + type: "tool-result"; + toolName?: string; + toolCallId?: string; + output?: { type?: string; value?: unknown }; +} + +export interface PruneOptions { + /** Token budget to stay under. */ + targetTokens: number; + /** + * Estimator used to compute current token count. Injected so the caller + * uses the same estimator as the budget check that triggered the prune. + */ + estimate: (messages: ModelMessage[]) => number; + /** + * How many of the most recent tool messages to leave alone, regardless + * of size. The freshest tool result is what the model is reacting to — + * pruning it would erase the reason it's about to make its next move. + * Defaults to 2. + */ + preserveRecentToolMessages?: number; + /** + * Lower bound on per-tool-result payload size that's eligible for + * pruning. Tiny results (a few hundred bytes) don't move the needle + * and pruning them just adds noise. Defaults to 2000 chars. + */ + minPrunablePayloadChars?: number; +} + +export interface PruneResult { + /** Whether any messages were modified. */ + pruned: boolean; + /** Tokens estimated before pruning. */ + tokensBefore: number; + /** Tokens estimated after pruning. */ + tokensAfter: number; + /** Number of tool-result parts that were shortened. */ + partsPruned: number; + /** Total bytes removed across all shortened parts. */ + bytesRemoved: number; +} + +/** + * Walk `messages` in place and shorten the largest tool-result payloads + * until the projected token count drops below `targetTokens`. Returns a + * summary of what changed; the caller is expected to log it. + * + * The function is deterministic given the same input. + */ +export function pruneOversizedToolResults( + messages: ModelMessage[], + opts: PruneOptions, +): PruneResult { + const preserveRecent = opts.preserveRecentToolMessages ?? 2; + const minPrunable = opts.minPrunablePayloadChars ?? 2_000; + const tokensBefore = opts.estimate(messages); + + if (tokensBefore <= opts.targetTokens) { + return { + pruned: false, + tokensBefore, + tokensAfter: tokensBefore, + partsPruned: 0, + bytesRemoved: 0, + }; + } + + // Index every prunable tool-result part with its serialized size and + // its source message index. We'll process oldest-first, skipping the + // most recent `preserveRecent` tool messages entirely. + const toolMessageIndices: number[] = []; + for (let i = 0; i < messages.length; i++) { + if (messages[i].role === "tool") toolMessageIndices.push(i); + } + // When there are fewer (or equal) tool messages than the preserve + // window, every tool message is protected — cutoffIndex stays at 0 so + // the eligibility loop finds no candidates. Otherwise the cutoff is + // the index of the *first* preserved tool message; anything strictly + // before it is eligible for pruning. + const cutoffIndex = toolMessageIndices.length > preserveRecent + ? toolMessageIndices[toolMessageIndices.length - preserveRecent] + : 0; + + type Candidate = { + msgIdx: number; + partIdx: number; + sizeChars: number; + toolName: string; + }; + const candidates: Candidate[] = []; + + for (let i = 0; i < messages.length; i++) { + if (i >= cutoffIndex) break; + const msg = messages[i]; + if (msg.role !== "tool") continue; + const content = msg.content; + if (!Array.isArray(content)) continue; + for (let p = 0; p < content.length; p++) { + const part = content[p] as unknown as ToolResultPart; + if (part?.type !== "tool-result") continue; + const value = part.output?.value; + if (value === undefined) continue; + const serialized = typeof value === "string" ? value : JSON.stringify(value); + if (serialized.length < minPrunable) continue; + candidates.push({ + msgIdx: i, + partIdx: p, + sizeChars: serialized.length, + toolName: part.toolName ?? "unknown", + }); + } + } + + // Oldest-largest first: oldest messages are dropped before recent ones, + // and within the same age the biggest payload wins. + candidates.sort((a, b) => { + if (a.msgIdx !== b.msgIdx) return a.msgIdx - b.msgIdx; + return b.sizeChars - a.sizeChars; + }); + + let partsPruned = 0; + let bytesRemoved = 0; + + for (const cand of candidates) { + if (opts.estimate(messages) <= opts.targetTokens) break; + + const msg = messages[cand.msgIdx]; + if (msg.role !== "tool" || !Array.isArray(msg.content)) continue; + const part = msg.content[cand.partIdx] as unknown as ToolResultPart; + if (part?.type !== "tool-result") continue; + + const placeholder = + `[Tool result from \`${cand.toolName}\` (${cand.sizeChars} chars) pruned by the harness ` + + `to stay under the context budget. Re-run the tool if you still need this data.]`; + + // Mutate in place. Preserve the envelope (toolCallId, toolName) so the + // assistant's prior tool-call still resolves; only the payload changes. + part.output = { + type: "text", + value: placeholder, + }; + + partsPruned += 1; + bytesRemoved += cand.sizeChars - placeholder.length; + } + + const tokensAfter = opts.estimate(messages); + return { + pruned: partsPruned > 0, + tokensBefore, + tokensAfter, + partsPruned, + bytesRemoved, + }; +} diff --git a/test/own-loop-prune-unit.test.ts b/test/own-loop-prune-unit.test.ts new file mode 100644 index 0000000..c7cb4ed --- /dev/null +++ b/test/own-loop-prune-unit.test.ts @@ -0,0 +1,200 @@ +/** + * Unit tests for `pruneOversizedToolResults` — the harness-level safety + * net that keeps the own-loop's in-memory messages array under the token + * budget on first-turn explorations. + */ +import { describe, expect, it } from "vitest"; +import type { ModelMessage } from "ai"; +import { pruneOversizedToolResults } from "../src/own-loop-prune"; + +// Deterministic estimator for tests: 1 token per 4 chars of JSON, rounded up. +function estimate(messages: ModelMessage[]): number { + let total = 0; + for (const m of messages) total += Math.ceil(JSON.stringify(m).length / 4); + return total; +} + +function userMsg(text: string): ModelMessage { + return { role: "user", content: text }; +} + +function assistantToolCall(toolName: string, toolCallId: string): ModelMessage { + return { + role: "assistant", + content: [ + { + type: "tool-call", + toolCallId, + toolName, + input: {}, + } as unknown as never, + ], + }; +} + +function toolResult(toolName: string, toolCallId: string, payload: string): ModelMessage { + return { + role: "tool", + content: [ + { + type: "tool-result", + toolName, + toolCallId, + output: { type: "text", value: payload }, + } as unknown as never, + ], + }; +} + +describe("pruneOversizedToolResults", () => { + it("returns pruned=false when already under budget", () => { + const messages: ModelMessage[] = [userMsg("hi"), toolResult("read", "t1", "small")]; + const result = pruneOversizedToolResults(messages, { + targetTokens: 1_000, + estimate, + }); + expect(result.pruned).toBe(false); + expect(result.partsPruned).toBe(0); + expect(messages[1].content).toMatchObject([ + { type: "tool-result", output: { value: "small" } }, + ]); + }); + + it("prunes the oldest large tool-result first", () => { + const big = "X".repeat(40_000); // ~10k tokens with the test estimator + const small = "Y".repeat(40); // tiny + const messages: ModelMessage[] = [ + userMsg("explore"), + assistantToolCall("read", "t1"), + toolResult("read", "t1", big), + assistantToolCall("read", "t2"), + toolResult("read", "t2", small), + ]; + const before = estimate(messages); + expect(before).toBeGreaterThan(5_000); + + const result = pruneOversizedToolResults(messages, { + targetTokens: 500, + estimate, + preserveRecentToolMessages: 0, + minPrunablePayloadChars: 1_000, + }); + + expect(result.pruned).toBe(true); + expect(result.partsPruned).toBe(1); + expect(result.tokensAfter).toBeLessThan(result.tokensBefore); + // The big one is now a placeholder; the small one is untouched. + expect((messages[2].content as Array<{ output?: { value: string } }>)[0].output?.value) + .toMatch(/pruned by the harness/); + expect((messages[4].content as Array<{ output?: { value: string } }>)[0].output?.value) + .toBe(small); + }); + + it("preserves the most recent N tool messages regardless of size", () => { + const huge = "Z".repeat(80_000); + const messages: ModelMessage[] = [ + userMsg("explore"), + assistantToolCall("read", "t1"), + toolResult("read", "t1", huge), + assistantToolCall("read", "t2"), + toolResult("read", "t2", huge), + ]; + + const result = pruneOversizedToolResults(messages, { + targetTokens: 1_000, // impossible to hit while preserving both + estimate, + preserveRecentToolMessages: 2, + minPrunablePayloadChars: 1_000, + }); + + // Both recent tool messages are preserved, so nothing is pruned. + expect(result.pruned).toBe(false); + expect(result.partsPruned).toBe(0); + const lastToolValue = (messages[4].content as Array<{ output?: { value: string } }>)[0].output?.value; + expect(lastToolValue).toBe(huge); + }); + + it("preserves tool-call envelope (toolName, toolCallId) when pruning", () => { + const big = "Q".repeat(40_000); + const messages: ModelMessage[] = [ + userMsg("explore"), + assistantToolCall("grep", "abc-123"), + toolResult("grep", "abc-123", big), + ]; + + pruneOversizedToolResults(messages, { + targetTokens: 100, + estimate, + preserveRecentToolMessages: 0, + minPrunablePayloadChars: 1_000, + }); + + const part = (messages[2].content as Array<{ type: string; toolName?: string; toolCallId?: string }>)[0]; + // toolName + toolCallId must survive so the model's prior tool-call still resolves. + expect(part.toolName).toBe("grep"); + expect(part.toolCallId).toBe("abc-123"); + expect(part.type).toBe("tool-result"); + }); + + it("skips payloads smaller than minPrunablePayloadChars", () => { + // Three tool messages, each ~500 chars — below the default 2000 floor. + const messages: ModelMessage[] = [ + userMsg("hi"), + toolResult("read", "t1", "a".repeat(500)), + toolResult("read", "t2", "b".repeat(500)), + toolResult("read", "t3", "c".repeat(500)), + ]; + + const result = pruneOversizedToolResults(messages, { + targetTokens: 1, + estimate, + preserveRecentToolMessages: 0, + }); + + expect(result.partsPruned).toBe(0); + // Originals intact. + expect((messages[1].content as Array<{ output?: { value: string } }>)[0].output?.value) + .toHaveLength(500); + }); + + it("stops pruning once budget is back under target", () => { + const big1 = "1".repeat(20_000); + const big2 = "2".repeat(20_000); + const big3 = "3".repeat(20_000); + const messages: ModelMessage[] = [ + userMsg("explore"), + toolResult("read", "t1", big1), + toolResult("read", "t2", big2), + toolResult("read", "t3", big3), + ]; + + const result = pruneOversizedToolResults(messages, { + // Roughly enough budget that pruning ONE message gets us back under. + targetTokens: Math.ceil(estimate(messages) / 2), + estimate, + preserveRecentToolMessages: 0, + minPrunablePayloadChars: 1_000, + }); + + // Should have pruned at least one but not all three. + expect(result.partsPruned).toBeGreaterThan(0); + expect(result.partsPruned).toBeLessThan(3); + expect(result.tokensAfter).toBeLessThanOrEqual( + Math.ceil(estimate(messages) * 1.05), // tiny slack for rounding + ); + }); + + it("is a no-op when there are no tool messages", () => { + const messages: ModelMessage[] = [ + userMsg("hello"), + { role: "assistant", content: "world" }, + ]; + const before = JSON.stringify(messages); + const result = pruneOversizedToolResults(messages, { + targetTokens: 1, + estimate, + }); + expect(result.pruned).toBe(false); + expect(JSON.stringify(messages)).toBe(before); + }); +});