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); + }); +});