From 203b798c619415dcb3b7dc5d79154f20aa061afc Mon Sep 17 00:00:00 2001 From: Sawyer Cutler Date: Sat, 17 Oct 2026 19:21:52 -0700 Subject: [PATCH] fix(agent): compact before the completion budget overflows the window --- src/agent/compaction.test.ts | 93 +++++++++++++++++- src/agent/compaction.ts | 85 ++++++++++++++-- src/agent/director.ts | 2 + src/config/index.ts | 12 +++ .../context-window-thresholds.test.ts | 14 +++ src/provider/context-window.ts | 98 ++++++++++++++++++- src/subagent/nudge-director.ts | 2 + 7 files changed, 292 insertions(+), 14 deletions(-) diff --git a/src/agent/compaction.test.ts b/src/agent/compaction.test.ts index 92d27d3b8..16c64a55b 100644 --- a/src/agent/compaction.test.ts +++ b/src/agent/compaction.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, test } from "bun:test"; +import { afterEach, describe, expect, test } from "bun:test"; import type { ConversationTurn, ReactorAction, @@ -18,8 +18,12 @@ import { createChatDirector, } from "./director.js"; import { - compactionWideResumeDeltaFor, + clearSessionContextWindowPins, compactionThresholdFor, + compactionWideResumeDeltaFor, + contextWindowFor, + setModelContextWindows, + setReservedCompletionTokens, } from "../provider/context-window.js"; import { COMPACT_SPACER_TEXT, @@ -138,6 +142,30 @@ function emptyMessage(): ReactorInboundEvent { } as ReactorInboundEvent; } +function userMessage(content: string): ReactorInboundEvent { + return { + type: "message.received", + message: { content }, + } as ReactorInboundEvent; +} + +const VLLM_OVERFLOW = + "Requested token count exceeds the model's maximum context length of 262144 tokens. You requested a total of 269714 tokens: 236946 tokens from the input messages and 32768 tokens for the completion."; + +function mislabeledOverflow(category: string): ReactorInboundEvent { + return { + type: "inference.error", + error: { category, message: VLLM_OVERFLOW }, + partial: { text: "" }, + } as ReactorInboundEvent; +} + +afterEach(() => { + setReservedCompletionTokens(undefined); + setModelContextWindows(undefined); + clearSessionContextWindowPins(); +}); + function overflowError(): ReactorInboundEvent { return { type: "inference.error", @@ -152,6 +180,67 @@ const inferAction: ReactorAction[] = [{ type: "infer" }]; const tenTurns = turnsOfLength(10, 1); describe("compaction governor", () => { + test("compacts a follow-up when the live estimate plus maxTokens crosses the window", () => { + setReservedCompletionTokens({ m: 40_000 }); + const grown = turnsOfLength(10, 36_000); + const arm = (governor: ReturnType) => { + governor.noteInferenceDone(inferenceDone(20_000), tenTurns); + governor.syncFromTurns(grown); + }; + + const toolGovernor = createCompactionGovernor(() => undefined); + arm(toolGovernor); + const toolActions = toolGovernor.interceptActions( + toolDone(), + inferAction, + capabilities, + ); + expect(toolActions?.some((a) => a.type === "compact")).toBe(true); + expect(toolActions?.some((a) => a.type === "infer")).toBe(false); + + const userGovernor = createCompactionGovernor(() => undefined); + arm(userGovernor); + const userActions = userGovernor.interceptActions( + userMessage("continue"), + inferAction, + capabilities, + ); + expect(userActions?.some((a) => a.type === "compact")).toBe(true); + expect(userActions?.some((a) => a.type === "infer")).toBe(false); + }); + + test("a non-context_overflow body that names a context overflow compact-and-retries twice", () => { + const governor = createCompactionGovernor(undefined); + const event = mislabeledOverflow("fatal"); + expect( + governor + .interceptOverflow(event, capabilities) + ?.some((a) => a.type === "compact" && a.reason === "context-overflow"), + ).toBe(true); + expect( + governor + .interceptOverflow(event, capabilities) + ?.some((a) => a.type === "compact"), + ).toBe(true); + expect(governor.interceptOverflow(event, capabilities)).toBeNull(); + }); + + test("a server-named context limit pins the window used by the next threshold", () => { + setModelContextWindows({ m: 1_000_000 }); + const reserved = 32_768; + const governor = createCompactionGovernor(undefined); + governor.noteInferenceDone(inferenceDone(1000), tenTurns); + expect(contextWindowFor("m")).toBe(1_000_000); + expect(compactionThresholdFor("m", reserved)).toBe( + Math.floor((1_000_000 - reserved) * 0.6), + ); + governor.interceptOverflow(mislabeledOverflow("retryable"), capabilities); + expect(contextWindowFor("m")).toBe(262_144); + expect(compactionThresholdFor("m", reserved)).toBe( + Math.floor((262_144 - reserved) * 0.6), + ); + }); + test("stays inert below the threshold or with few turns", () => { const governor = createCompactionGovernor(() => undefined); governor.noteInferenceDone(inferenceDone(1000), tenTurns); diff --git a/src/agent/compaction.ts b/src/agent/compaction.ts index 14867d73e..0238fe0e5 100644 --- a/src/agent/compaction.ts +++ b/src/agent/compaction.ts @@ -8,10 +8,14 @@ import type { } from "@intx/types/runtime"; import { compactionThresholdFor, + completionReservationExceedsWindow, contextTokensFromUsage, hasWideResumeGap, isAtOrUnderCompactThreshold, + pinSessionContextWindow, + reservedMaxTokensFor, } from "../provider/context-window.js"; +import { looksLikeContextOverflow } from "../inference-error-message.js"; import { assistantTextIsCompactSpacerEcho, COMPACT_NOOP_FLOOR, @@ -30,6 +34,9 @@ const MIN_TURNS_TO_COMPACT = COMPACT_NOOP_FLOOR; // Cap overflow retries so an unshrinkable history cannot loop forever; // cleared once usage drops back under the watermark. const MAX_OVERFLOW_RECOVERIES = 2; +// vLLM / OpenAI-compatible bodies name the enforced window even when the +// HTTP classifier did not tag context_overflow (non-400). +const SERVER_CONTEXT_LIMIT = /maximum context length of\s+(\d+)/i; // Consecutive threshold compacts with no under-watermark relief; tool-call // occupancy must not reset it, or tool traffic re-enables the loop. const MAX_CONSECUTIVE_THRESHOLD_COMPACTS = 2; @@ -151,6 +158,7 @@ export function createCompactionGovernor( // True when arming fell back to the local estimate (usage omitted or // zero), so a meter can flag the number as approximate. let usingEstimate = false; + let lastReportedTokens = 0; let lastModel: string | undefined; let turnCount = 0; // Growth latch: snapshot the first inference.done after a compact and @@ -177,9 +185,13 @@ export function createCompactionGovernor( return estimate.syncFromTurns(turns); } + function reserved(): number { + return reservedMaxTokensFor(lastModel); + } + function isOverThreshold(contextTokens: number): boolean { if (turnCount <= MIN_TURNS_TO_COMPACT) return false; - const high = compactionThresholdFor(lastModel); + const high = compactionThresholdFor(lastModel, reserved()); if (contextTokens <= high) return false; if (tokensAtLastCompact !== undefined) { return hasWideResumeGap(tokensAtLastCompact, contextTokens, lastModel); @@ -187,6 +199,24 @@ export function createCompactionGovernor( return true; } + // Reported usage wins for the 60% latch, but a follow-up can grow the live + // estimate past that sample. The hard check takes the larger of the two so + // a prompt that fits the last sample still compacts when adding maxTokens + // would cross the window. + function occupancyTokens(): number { + if (usingEstimate || lastReportedTokens <= 0) return estimate.tokens; + return Math.max(lastReportedTokens, estimate.tokens); + } + + function wouldExceedWindow(): boolean { + if (turnCount <= MIN_TURNS_TO_COMPACT) return false; + return completionReservationExceedsWindow( + occupancyTokens(), + lastModel, + reserved(), + ); + } + function noteCompactIssued(): void { awaitingPostCompactMeasurement = true; } @@ -244,13 +274,14 @@ export function createCompactionGovernor( lastModel = event.source?.model; const reportedTokens = contextTokensFromUsage(event.usage); usingEstimate = reportedTokens <= 0; + lastReportedTokens = reportedTokens; const contextTokens = usingEstimate ? estimate.tokens : reportedTokens; // Snapshot on the first inference.done after a compact (the post-compact // infer); intercept has no fresh usage. if (awaitingPostCompactMeasurement) { tokensAtLastCompact = contextTokens; awaitingPostCompactMeasurement = false; - if (!isAtOrUnderCompactThreshold(contextTokens, lastModel)) { + if (!isAtOrUnderCompactThreshold(contextTokens, lastModel, reserved())) { foldNonConverged = true; consecutiveThresholdCompacts = MAX_CONSECUTIVE_THRESHOLD_COMPACTS; } @@ -259,7 +290,7 @@ export function createCompactionGovernor( // compacts, overflow recoveries) and clears foldNonConverged; the growth // latch stays. Nothing else resets them, or compact→infer→compact loops // forever. - if (isAtOrUnderCompactThreshold(contextTokens, lastModel)) { + if (isAtOrUnderCompactThreshold(contextTokens, lastModel, reserved())) { consecutiveThresholdCompacts = 0; overflowRecoveries = 0; foldNonConverged = false; @@ -269,22 +300,43 @@ export function createCompactionGovernor( pending = isOverThreshold(contextTokens); } + function issueOverflowPreflight( + capabilities: ReactorCapabilities, + ): ReactorAction[] | null { + if (manualPending || atThresholdCompactCap()) return null; + if (turnCount <= MIN_TURNS_TO_COMPACT) return null; + if (!wouldExceedWindow()) return null; + clearManualArming(); + postCompactInfer = true; + issueThresholdCompact(); + return [ + capabilities.compact(COMPACTOR_NAME, THRESHOLD_COMPACT_REASON), + ...continuationActions(capabilities), + ]; + } + // Compact at the natural pause between a tool batch and its follow-up - // infer: drop the infer, compact, and re-enter via the continuation. + // infer, and before a user message's infer: drop the infer, compact, and + // re-enter via the continuation. // `pending` is the last inference.done snapshot — authoritative with real // usage; with omitted usage it came from the estimate, so re-derive // against the live estimate instead of trusting a `pending` stale by one - // tool batch. + // tool batch. A reported sample can still be under the watermark while the + // live estimate plus maxTokens already fills the window; that preflight + // does not wait for the next inference.done. function interceptActions( event: ReactorInboundEvent, actions: ReactorAction[], capabilities: ReactorCapabilities, ): ReactorAction[] | null { - if (event.type !== "tool.done") return null; + if (event.type !== "tool.done" && event.type !== "message.received") + return null; const operator = manualPending; + const hard = !operator && wouldExceedWindow(); if ( !operator && !pending && + !hard && !(usingEstimate && isOverThreshold(estimate.tokens)) ) return null; @@ -390,17 +442,22 @@ export function createCompactionGovernor( } // A context-overflow error would terminate the loop; compact and retry, - // bounded so an unshrinkable history does not loop forever. + // bounded so an unshrinkable history does not loop forever. Message text + // is enough: the vendor classifier only tags HTTP 400, and vLLM can phrase + // the same overflow on another status. function interceptOverflow( event: ReactorInboundEvent, capabilities: ReactorCapabilities, ): ReactorAction[] | null { + if (event.type !== "inference.error") return null; + const message = event.error.message ?? ""; if ( - event.type !== "inference.error" || - event.error.category !== "context_overflow" + event.error.category !== "context_overflow" && + !looksLikeContextOverflow(message) ) { return null; } + pinNamedServerLimit(message); if (overflowRecoveries >= MAX_OVERFLOW_RECOVERIES) return null; overflowRecoveries++; // Overflow compact spends any operator arming so a queued handoff @@ -414,6 +471,15 @@ export function createCompactionGovernor( ]; } + function pinNamedServerLimit(message: string): void { + if (lastModel === undefined) return; + const match = SERVER_CONTEXT_LIMIT.exec(message); + if (match === null) return; + const limit = Number(match[1]); + if (!Number.isSafeInteger(limit) || limit <= 0) return; + pinSessionContextWindow(lastModel, limit); + } + // A keep-prior compact (verify abort, summarizer failure) must not resume // into inference: the continuation is already queued, and inferring on the // full transcript is what 401/429s the handoff. @@ -564,6 +630,7 @@ export function createCompactionGovernor( notePostCompact, noteIdleTurn, interceptActions, + preflightInfer: issueOverflowPreflight, interceptIdleContinuation, interceptOverflow, resumeAfterCompact, diff --git a/src/agent/director.ts b/src/agent/director.ts index 49fa51946..0a14c34ff 100644 --- a/src/agent/director.ts +++ b/src/agent/director.ts @@ -873,6 +873,8 @@ class ChatDirectorImpl extends DefaultDirector { this.compaction.notePostCompact(state.turns ?? []); // The meter-only re-entry stays idle — no unprompted inference. if (afterCompact === "meter") return capabilities.wait(); + const preflight = this.compaction.preflightInfer(capabilities); + if (preflight !== null) return preflight; return capabilities.infer(); } const idleCompact = this.compaction.interceptIdleContinuation( diff --git a/src/config/index.ts b/src/config/index.ts index f85eff33c..f02cb4159 100644 --- a/src/config/index.ts +++ b/src/config/index.ts @@ -20,7 +20,9 @@ import { } from "../provider/reasoning-effort.js"; import { buildProviderContextWindowOverrides, + buildReservedCompletionTokens, setProviderContextWindowOverrides, + setReservedCompletionTokens, } from "../provider/context-window.js"; import { bootstrapPricingMetadata } from "../cost/pricing-metadata.js"; import { @@ -1490,11 +1492,13 @@ export function refreshProviderContextWindows( entry.codexProfile === undefined && entry.xaiProfile === undefined ? settings?.providers[entry.name]?.contextWindow : undefined; + const maxTokens = settings?.providers[entry.name]?.maxTokens; return [ entry.name, { models: entry.models, ...(window !== undefined ? { contextWindow: window } : {}), + ...(maxTokens !== undefined ? { maxTokens } : {}), }, ]; }), @@ -1502,6 +1506,14 @@ export function refreshProviderContextWindows( setProviderContextWindowOverrides( buildProviderContextWindowOverrides(providers, activeProvider, activeModel), ); + setReservedCompletionTokens( + buildReservedCompletionTokens( + providers, + activeProvider, + activeModel, + SOURCE_MAX_TOKENS, + ), + ); } export function catalogEntryAsProviderSettings( diff --git a/src/provider/context-window-thresholds.test.ts b/src/provider/context-window-thresholds.test.ts index 8b3fd4d47..fc7dd21d2 100644 --- a/src/provider/context-window-thresholds.test.ts +++ b/src/provider/context-window-thresholds.test.ts @@ -12,6 +12,7 @@ import { CONTEXT_METER_DANGER_FRACTION, setModelContextWindows, setProviderContextWindowOverrides, + clearSessionContextWindowPins, } from "./context-window.js"; function usage(overrides: Partial): TokenUsage { @@ -28,6 +29,7 @@ function usage(overrides: Partial): TokenUsage { afterEach(() => { setModelContextWindows(undefined); setProviderContextWindowOverrides(undefined); + clearSessionContextWindowPins(); }); describe("contextWindowFor", () => { @@ -63,6 +65,18 @@ describe("compactionThresholdFor", () => { test("falls back to the default window when the model is unknown", () => { expect(compactionThresholdFor(undefined)).toBe(76_800); }); + + test("is 60 percent of the window minus the reserved completion budget", () => { + setModelContextWindows({ "qwen-test": 262_144 }); + const reserved = 32_768; + const usable = 262_144 - reserved; + const threshold = compactionThresholdFor("qwen-test", reserved); + expect(threshold).toBe(Math.floor(usable * COMPACTION_WINDOW_FRACTION)); + expect(threshold).toBeLessThan(usable); + expect(compactionThresholdFor("qwen-test", 0)).toBe( + Math.floor(262_144 * COMPACTION_WINDOW_FRACTION), + ); + }); }); describe("compactionWideResumeDeltaFor", () => { diff --git a/src/provider/context-window.ts b/src/provider/context-window.ts index 3c4592f8f..7ceab59ab 100644 --- a/src/provider/context-window.ts +++ b/src/provider/context-window.ts @@ -26,6 +26,15 @@ let contextWindowRegistry: Record = {}; // later models.dev refresh because it lives beside the registry, not in it. let contextWindowOverrides: Record = {}; +// Completion budget the active source will send (`maxTokens`). Separate from +// the window tables so a reservation cannot be mistaken for a context limit. +let reservedCompletionTokens: Record = {}; + +// Server-enforced limits learned from an overflow body this process. They win +// over settings and models.dev, which can advertise a larger window than the +// server actually accepts, and they are never written back to user settings. +let sessionContextWindowPins: Record = {}; + // Both tables are wholesale-replaced on refresh; one factory keeps the two // trivial setters from drifting. `undefined` (no cache yet) means empty. function createRegistrySetter( @@ -44,6 +53,21 @@ export const setProviderContextWindowOverrides = createRegistrySetter( }, ); +export const setReservedCompletionTokens = createRegistrySetter((windows) => { + reservedCompletionTokens = windows; +}); + +/** Drop server limits learned from overflow bodies. Test isolation only. */ +export function clearSessionContextWindowPins(): void { + sessionContextWindowPins = {}; +} + +/** Remember a server-named context limit for this process. Not persisted. */ +export function pinSessionContextWindow(model: string, window: number): void { + if (model.length === 0 || !isPositiveWindow(window)) return; + sessionContextWindowPins = { ...sessionContextWindowPins, [model]: window }; +} + export type ProviderContextWindowSource = { models: readonly string[]; contextWindow?: number; @@ -81,6 +105,39 @@ export function buildProviderContextWindowOverrides( return overrides; } +// Same keying as context-window overrides. Unset provider maxTokens falls +// back to the source default the client actually sends; 0 means unknown. +export function buildReservedCompletionTokens( + providers: Record< + string, + ProviderContextWindowSource & { maxTokens?: number } + >, + resolvedProviderName: string, + resolvedModel: string, + fallbackMaxTokens: number, +): Record { + const reserved: Record = {}; + const fallback = isPositiveWindow(fallbackMaxTokens) ? fallbackMaxTokens : 0; + for (const [name, provider] of Object.entries(providers)) { + const budget = isPositiveWindow(provider.maxTokens) + ? provider.maxTokens + : fallback; + if (budget === 0) continue; + const models = new Set(provider.models); + if (name === resolvedProviderName && resolvedModel.length > 0) { + models.add(resolvedModel); + } + for (const model of models) { + if (model.length === 0) continue; + reserved[`${name}:${model}`] = budget; + if (name === resolvedProviderName) { + reserved[model] = budget; + } + } + } + return reserved; +} + function heuristicWindow(model: string): number { const m = model.toLowerCase(); if (m.includes("gpt-6")) return 1_000_000; @@ -142,12 +199,20 @@ export function hasContextWindowFor(model: string): boolean { export function contextWindowFor(model: string): number { return ( + lookupWindow(sessionContextWindowPins, model) ?? lookupWindow(contextWindowOverrides, model) ?? lookupWindow(contextWindowRegistry, model) ?? heuristicWindow(model) ); } +/** maxTokens the source will send for `model`, or 0 when unset. */ +export function reservedMaxTokensFor(model: string | undefined): number { + if (model === undefined) return 0; + const reserved = lookupWindow(reservedCompletionTokens, model); + return isPositiveWindow(reserved) ? reserved : 0; +} + // Fraction of the window at which proactive compaction fires. Kept well below // the hard limit so summarization happens while the model still reasons well // and before any provider rejects the request. Also the status-bar meter's @@ -182,10 +247,21 @@ export function contextMeterBand(percentUsed: number): ContextMeterBand { // Token threshold at which the director should compact, sized to the model's // real window. `model` may be undefined early in a session (no cycle yet); // fall back to the default window. -export function compactionThresholdFor(model: string | undefined): number { +export function compactionThresholdFor( + model: string | undefined, + reservedMaxTokens = 0, +): number { const window = model !== undefined ? contextWindowFor(model) : DEFAULT_CONTEXT_WINDOW; - return Math.floor(window * COMPACTION_WINDOW_FRACTION); + // Reserve the completion budget the client adds on top of the prompt. + // 60% of what remains sits strictly below `window - maxTokens` (the fraction + // is under 1), so a steady-state prompt cannot be legal locally and illegal + // once that reservation is sent. + const reserved = + Number.isFinite(reservedMaxTokens) && reservedMaxTokens > 0 + ? Math.min(reservedMaxTokens, window) + : 0; + return Math.floor((window - reserved) * COMPACTION_WINDOW_FRACTION); } /** Tokens of growth past the last post-compact measurement before the @@ -205,8 +281,24 @@ export function compactionWideResumeDeltaFor( export function isAtOrUnderCompactThreshold( contextTokens: number, model: string | undefined, + reservedMaxTokens = 0, +): boolean { + return contextTokens <= compactionThresholdFor(model, reservedMaxTokens); +} + +/** Prompt plus the completion reservation would be rejected by the window. */ +export function completionReservationExceedsWindow( + contextTokens: number, + model: string | undefined, + reservedMaxTokens: number, ): boolean { - return contextTokens <= compactionThresholdFor(model); + const window = + model !== undefined ? contextWindowFor(model) : DEFAULT_CONTEXT_WINDOW; + const reserved = + Number.isFinite(reservedMaxTokens) && reservedMaxTokens > 0 + ? reservedMaxTokens + : 0; + return contextTokens + reserved >= window; } /** Shared re-arm rule for a session latched after a compact (automatic, diff --git a/src/subagent/nudge-director.ts b/src/subagent/nudge-director.ts index a90505e28..199a76c21 100644 --- a/src/subagent/nudge-director.ts +++ b/src/subagent/nudge-director.ts @@ -262,6 +262,8 @@ export class SubAgentDirector extends DefaultDirector { // post-compact resume once this leaf already replied. if (afterCompact === "meter" || this.reportReplied) return capabilities.wait(); + const preflight = this.compaction.preflightInfer(capabilities); + if (preflight !== null) return preflight; return this.applyPendingNudge([capabilities.infer()], capabilities); } const idleCompact = this.compaction.interceptIdleContinuation(