Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 78 additions & 2 deletions src/agent/compaction.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,25 @@ import type {
} from "@intx/types/runtime";
import { createCompactionGovernor } from "./compaction.js";
import { compactionThresholdFor } from "../provider/context-window.js";
import { COMPACTOR_KEEP_RECENT_TURNS, compactorNoOpFloor } from "../session/compactor.js";

const capabilities = {
infer: (options?: unknown) => ({ type: "infer", ...(options !== undefined ? { options } : {}) }),
compact: (compactor: string, reason: string) => ({ type: "compact", compactor, reason }),
} as unknown as ReactorCapabilities;

// Distinct, non-zero cacheRead/cacheWrite so a test asserting on the total
// would fail if compaction.ts ever stopped routing through the shared
// contextTokensFromUsage and summed only `input` again.
function usage(input: number): TokenUsage {
return { input, output: 0, cacheRead: 0, cacheWrite: 0, thinking: 0 };
return { input, output: 0, cacheRead: 3, cacheWrite: 5, thinking: 0 };
}

// A provider that truly omits usage reports every field as zero, not just
// `input` — distinct from usage(0), which still carries the fixture's
// non-zero cache values above.
function zeroUsage(): TokenUsage {
return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, thinking: 0 };
}

function turnsOfLength(count: number, textLength: number): ConversationTurn[] {
Expand All @@ -39,7 +50,7 @@ function inferenceDoneWithoutUsage(): Extract<ReactorInboundEvent, { type: "infe
return {
type: "inference.done",
turn: { role: "assistant", content: [{ type: "text", text: "ok" }] },
usage: usage(0),
usage: zeroUsage(),
source: { sourceId: "s", provider: "p", model: "m" },
} as unknown as Extract<ReactorInboundEvent, { type: "inference.done" }>;
}
Expand Down Expand Up @@ -255,4 +266,69 @@ describe("compaction governor", () => {
expect(governor.interceptActions(inferenceDone(overThreshold), inferAction, capabilities)).toBeNull();
expect(governor.interceptActions(toolDone(), [{ type: "reply", content: "x" }], capabilities)).toBeNull();
});

test("stays inert below the minimum-turn floor no matter how far over threshold", () => {
// Two turns is well under MIN_TURNS_TO_COMPACT. createPruningCompactor
// no-ops at the same floor (see session/compactor.ts), so arming here
// would spend a reactor cycle that cannot shrink anything.
const governor = createCompactionGovernor(() => {});
governor.noteInferenceDone(inferenceDone(overThreshold * 10), turnsOfLength(2, 1));
expect(governor.interceptActions(toolDone(), inferAction, capabilities)).toBeNull();
});

test("arms on tool.done from a live estimate even when the last snapshot was under threshold", () => {
// Usage is omitted (pending is derived from the local estimate, which
// starts small and stays false), but the tool result that follows is
// itself large enough to cross the ordinary threshold before the next
// inference.done ever runs.
const governor = createCompactionGovernor(() => {});
governor.noteInferenceDone(inferenceDoneWithoutUsage(), tenTurns);
expect(governor.interceptActions(toolDone(), inferAction, capabilities)).toBeNull();

const overThresholdChars = (compactionThresholdFor("m") + 1) * 4;
governor.syncFromTurns(turnsOfLength(10, Math.ceil(overThresholdChars / 10)));

const actions = governor.interceptActions(toolDone(), inferAction, capabilities);
expect(actions).not.toBeNull();
expect(actions?.some((a) => a.type === "compact")).toBe(true);
});

test("never arms at the exact turn count createPruningCompactor no-ops on", () => {
// createPruningCompactor's own no-op floor (session/compactor.ts) is
// compactorNoOpFloor(COMPACTOR_KEEP_RECENT_TURNS). Arming at or below it
// would spend a reactor cycle that is guaranteed to shrink nothing.
const floor = compactorNoOpFloor(COMPACTOR_KEEP_RECENT_TURNS);
const governor = createCompactionGovernor(() => {});
governor.noteInferenceDone(inferenceDone(overThreshold), turnsOfLength(floor, 1));
expect(governor.interceptActions(toolDone(), inferAction, capabilities)).toBeNull();
});

test("arms one turn past the floor createPruningCompactor no-ops on", () => {
const floor = compactorNoOpFloor(COMPACTOR_KEEP_RECENT_TURNS);
const governor = createCompactionGovernor(() => {});
governor.noteInferenceDone(inferenceDone(overThreshold), turnsOfLength(floor + 1, 1));
const actions = governor.interceptActions(toolDone(), inferAction, capabilities);
expect(actions).not.toBeNull();
expect(actions?.some((a) => a.type === "compact")).toBe(true);
});

test("does not catch a huge tool result mid-cycle when the provider reported real usage", () => {
// Disclosed, accepted gap: the live tool.done re-check only re-derives
// arming from the local estimate when the last inference.done snapshot
// came from that same estimate (usingEstimate). When the provider
// reported real usage under threshold, that snapshot is trusted as
// authoritative until the next inference.done — a huge tool result
// arriving in between is not caught until then, unlike the
// usage-omitted case covered above.
const governor = createCompactionGovernor(() => {});
governor.noteInferenceDone(inferenceDone(1000), tenTurns);
expect(governor.interceptActions(toolDone(), inferAction, capabilities)).toBeNull();

const overThresholdChars = (compactionThresholdFor("m") + 1) * 4;
governor.syncFromTurns(turnsOfLength(10, Math.ceil(overThresholdChars / 10)));

// Still null: the live estimate is now over threshold, but the last
// arming decision trusted reported usage, so it is not re-checked here.
expect(governor.interceptActions(toolDone(), inferAction, capabilities)).toBeNull();
});
});
80 changes: 62 additions & 18 deletions src/agent/compaction.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,19 @@ import type {
ReactorAction,
ReactorCapabilities,
ReactorInboundEvent,
ToolDefinition,
} from "@intx/types/runtime";
import { compactionThresholdFor } from "../provider/context-window.js";
import { createContextEstimate } from "./context-estimate.js";
import { compactionThresholdFor, contextTokensFromUsage } from "../provider/context-window.js";
import { COMPACTOR_KEEP_RECENT_TURNS, compactorNoOpFloor } from "../session/compactor.js";
import { createContextEstimate, estimateOverheadTokens } from "./context-estimate.js";

const COMPACTOR_NAME = "pruning-compactor";
const MIN_TURNS_TO_COMPACT = 6;
// The exact turn count `createPruningCompactor` (session/compactor.ts) is
// guaranteed to no-op on. Derived from the same keepRecentTurns both real
// registrations (session, sub-agent) use, so this floor cannot silently
// drift from what the compactor will actually do — arming at or below it
// would spend a reactor cycle that shrinks nothing.
const MIN_TURNS_TO_COMPACT = compactorNoOpFloor(COMPACTOR_KEEP_RECENT_TURNS);
const MAX_OVERFLOW_RECOVERIES = 2;

// A compact action runs in its own reactor cycle, after which the reactor
Expand All @@ -20,51 +27,82 @@ const MAX_OVERFLOW_RECOVERIES = 2;
// would be worse than growing the context.
export type CompactionGovernor = ReturnType<typeof createCompactionGovernor>;

export function createCompactionGovernor(requestContinuation?: () => void) {
export function createCompactionGovernor(
requestContinuation?: () => void,
systemPrompt = "",
toolDefinitions: readonly ToolDefinition[] = [],
) {
let pending = false;
let idlePending = false;
let postCompactInfer = false;
let overflowRecoveries = 0;
// Set whenever the arming decision fell back to the local estimate because
// the provider omitted usage or reported zero, so callers rendering a meter
// can flag the number as approximate instead of implying provider-grade
// precision.
let usingEstimate = false;
// Model of the last inference.done turn, kept for live re-checks between
// inference cycles (see interceptActions) where the event carries no model.
let lastModel: string | undefined;
let turnCount = 0;

// Running local estimate of the turns we send. Providers that omit usage or
// report zero leave the proactive path blind; the estimate fills that gap.
// When the provider reports real usage we prefer it so a coarse local count
// cannot thrash against a trustworthy signal.
const estimate = createContextEstimate();
// Running local estimate of the turns we send, plus the fixed system-prompt
// and tool-schema overhead every request carries. Providers that omit usage
// or report zero leave the proactive path blind; the estimate fills that
// gap. When the provider reports real usage we prefer it so a coarse local
// count cannot thrash against a trustworthy signal.
const estimate = createContextEstimate(estimateOverheadTokens(systemPrompt, toolDefinitions));

// Re-sync after turn appends, tool results, and compaction rewrites. Callers
// pass the full turn list so the estimate stays accurate without incremental
// add/subtract bookkeeping.
function syncFromTurns(turns: readonly ConversationTurn[]): number {
turnCount = turns.length;
return estimate.syncFromTurns(turns);
}

function isOverThreshold(contextTokens: number): boolean {
return contextTokens > compactionThresholdFor(lastModel) && turnCount > MIN_TURNS_TO_COMPACT;
}

function noteInferenceDone(
event: Extract<ReactorInboundEvent, { type: "inference.done" }>,
turns: readonly ConversationTurn[],
): void {
overflowRecoveries = 0;
if (requestContinuation === undefined) return;
syncFromTurns(turns);
const reportedTokens = event.usage?.input ?? 0;
const contextTokens = reportedTokens > 0 ? reportedTokens : estimate.tokens;
// Assign, don't OR: an under-threshold follow-up must disarm a sticky pending
// left from an earlier over-threshold turn (e.g. after the provider reports
// real usage that lands below the threshold).
pending =
contextTokens > compactionThresholdFor(event.source?.model) &&
turns.length > MIN_TURNS_TO_COMPACT;
lastModel = event.source?.model;
const reportedTokens = contextTokensFromUsage(event.usage);
usingEstimate = reportedTokens <= 0;
const contextTokens = usingEstimate ? estimate.tokens : reportedTokens;
// Assign, don't OR: an under-threshold follow-up must disarm a sticky
// pending left from an earlier over-threshold turn (e.g. after the
// provider reports real usage that lands below the threshold).
pending = isOverThreshold(contextTokens);
}

// Compaction waits for the natural pause between a tool batch finishing and
// the follow-up infer: the infer is dropped from the action set, the compact
// cycle runs, and the continuation message re-enters inference.
//
// `pending` reflects the snapshot as of the last inference.done, which
// predates any tool result produced by that turn's own tool batch. When the
// provider is reporting real usage, that snapshot is authoritative and
// `pending` alone is trusted (there is no fresher provider number to check
// against until the next inference.done). But when usage was omitted or
// zero, `pending` was itself derived from the local estimate — in that case
// a large tool result can push the estimate over threshold before the next
// inference.done ever runs, so this re-derives the same arming rule against
// the live estimate (already re-synced this cycle by the director) instead
// of trusting a `pending` that can be stale by exactly one tool batch.
function interceptActions(
event: ReactorInboundEvent,
actions: ReactorAction[],
capabilities: ReactorCapabilities,
): ReactorAction[] | null {
if (!pending || event.type !== "tool.done") return null;
if (event.type !== "tool.done") return null;
if (!pending && !(usingEstimate && isOverThreshold(estimate.tokens))) return null;
if (!actions.some((a) => a.type === "infer")) return null;
pending = false;
postCompactInfer = true;
Expand Down Expand Up @@ -140,6 +178,12 @@ export function createCompactionGovernor(requestContinuation?: () => void) {
get estimatedTokens(): number {
return estimate.tokens;
},
// True once the provider has omitted or zeroed usage on the current
// turn, so a status-bar meter reading this can mark itself approximate
// rather than silently understating a real number.
get usingEstimate(): boolean {
return usingEstimate;
},
syncFromTurns,
noteInferenceDone,
noteIdleTurn,
Expand Down
26 changes: 25 additions & 1 deletion src/agent/context-estimate.test.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import { describe, expect, test } from "bun:test";
import type { ContentBlock, ConversationTurn, MediaSource } from "@intx/types/runtime";
import type { ContentBlock, ConversationTurn, MediaSource, ToolDefinition } from "@intx/types/runtime";
import {
createContextEstimate,
estimateContentBlockTokens,
estimateContextTokens,
estimateMediaSourceTokens,
estimateOverheadTokens,
estimateTokensFromChars,
} from "./context-estimate.js";

Expand Down Expand Up @@ -85,7 +86,30 @@ describe("estimateContextTokens", () => {
});
});

describe("estimateOverheadTokens", () => {
test("counts the system prompt and every tool's name, description, and schema", () => {
const systemPrompt = "x".repeat(40);
const tools: ToolDefinition[] = [
{ name: "run_shell", description: "y".repeat(20), inputSchema: { command: "string" } },
];
const expectedChars =
40 + "run_shell".length + 20 + JSON.stringify({ command: "string" }).length;
expect(estimateOverheadTokens(systemPrompt, tools)).toBe(estimateTokensFromChars(expectedChars));
});

test("is zero for an empty prompt and no tools", () => {
expect(estimateOverheadTokens("", [])).toBe(0);
});
});

describe("createContextEstimate", () => {
test("folds a fixed overhead into every sync", () => {
const estimate = createContextEstimate(100);
expect(estimate.tokens).toBe(100);
expect(estimate.syncFromTurns([textTurn("xxxx")])).toBe(101);
expect(estimate.tokens).toBe(101);
});

test("re-syncs from the full turn list after each append", () => {
const estimate = createContextEstimate();
expect(estimate.tokens).toBe(0);
Expand Down
33 changes: 28 additions & 5 deletions src/agent/context-estimate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,12 @@
// tool payloads, images) so proactive compaction still has a signal. This is
// a lower bound: system prompt, tool schemas, and framing are not counted.

import type { ContentBlock, ConversationTurn, MediaSource } from "@intx/types/runtime";
import type {
ContentBlock,
ConversationTurn,
MediaSource,
ToolDefinition,
} from "@intx/types/runtime";

const CHARS_PER_TOKEN = 4;

Expand Down Expand Up @@ -77,17 +82,35 @@ export function estimateContextTokens(turns: readonly ConversationTurn[]): numbe
return total;
}

// The system prompt and tool schemas ride on every request the same way turns
// do, but they never appear in `turns` — they're framing the harness supplies
// out of band. Without this, the estimate undercounts by whatever AGENTS.md
// and the active tool roster cost, which is often tens of thousands of tokens
// before a single turn is sent.
export function estimateOverheadTokens(
systemPrompt: string,
toolDefinitions: readonly ToolDefinition[],
): number {
let chars = systemPrompt.length;
for (const tool of toolDefinitions) {
chars += tool.name.length + tool.description.length + JSON.stringify(tool.inputSchema).length;
}
return estimateTokensFromChars(chars);
}

// Mutable running estimate. Callers re-sync from the full turn list after each
// append so compaction rewrites and tool results stay accurate without
// incremental add/subtract bookkeeping.
// incremental add/subtract bookkeeping. `overheadTokens` is fixed per session
// (system prompt + tool schemas do not change turn to turn) and is folded into
// every sync so the total tracks what actually goes out on the wire.
export type ContextEstimate = ReturnType<typeof createContextEstimate>;

export function createContextEstimate() {
let tokens = 0;
export function createContextEstimate(overheadTokens = 0) {
let tokens = overheadTokens;
let turnCount = 0;

function syncFromTurns(turns: readonly ConversationTurn[]): number {
tokens = estimateContextTokens(turns);
tokens = overheadTokens + estimateContextTokens(turns);
turnCount = turns.length;
return tokens;
}
Expand Down
10 changes: 9 additions & 1 deletion src/agent/director.ts
Original file line number Diff line number Diff line change
Expand Up @@ -361,7 +361,7 @@ class ChatDirectorImpl extends DefaultDirector {
this.onActivateTools = onActivateTools;
this.workflowCoordinator = workflowCoordinator;
this.onTasksChange = onTasksChange;
this.compaction = createCompactionGovernor(requestContinuation);
this.compaction = createCompactionGovernor(requestContinuation, systemPrompt, toolDefinitions);
this.modelFamilyPolicy = modelFamilyPolicy ?? resolveModelFamilyPolicy({ providerName: "" });
}

Expand All @@ -385,6 +385,13 @@ class ChatDirectorImpl extends DefaultDirector {
return [...this.tasks];
}

// The status bar's context meter falls back to this when a provider omits
// or zeroes usage on the latest turn — a local lower-then-corrected bound
// beats displaying a number the provider never actually reported.
getContextEstimate(): { tokens: number; isEstimate: boolean } {
return { tokens: this.compaction.estimatedTokens, isEstimate: this.compaction.usingEstimate };
}

private openTaskIds(): string[] {
return this.tasks
.filter((t) => t.status === "todo" || t.status === "doing")
Expand Down Expand Up @@ -796,4 +803,5 @@ export interface ChatDirector extends ReactorDirector {
setGoalGovernor(goal: GoalGovernor | undefined): void;
getGoalGovernor(): GoalGovernor | undefined;
getTasks(): Task[];
getContextEstimate(): { tokens: number; isEstimate: boolean };
}
Loading
Loading