Merge context meter accuracy and compaction arming

Author: TheGreatAxios

src/agent/compaction.test.ts

@@ -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[] {
@@ -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" }>;
 }
@@ -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();
+  });
 });

src/agent/compaction.ts

@@ -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
@@ -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;
@@ -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,

src/agent/context-estimate.test.ts

@@ -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";
 
@@ -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);

src/agent/context-estimate.ts

@@ -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;
 
@@ -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;
   }

src/agent/director.ts

@@ -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: "" });
   }
 
@@ -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")
@@ -796,4 +803,5 @@ export interface ChatDirector extends ReactorDirector {
   setGoalGovernor(goal: GoalGovernor | undefined): void;
   getGoalGovernor(): GoalGovernor | undefined;
   getTasks(): Task[];
+  getContextEstimate(): { tokens: number; isEstimate: boolean };
 }