fix: uniform assemble tokenBudget and messages-only currentTokenCount

jalehman · jalehman · commit 65deda82e458 · 2026-07-06T16:21:27.000-07:00
tokenBudget passed to ContextEngine.assemble now means the same thing on
every harness path: the budget available for assembled messages, computed
by the shared computeContextEngineMessageBudget helper (context window
minus compaction reserve minus rendered system/user prompt pressure).
currentTokenCount is now a messages-only estimate
(estimateTranscriptTokenPressure), so engines size systemPromptAddition
as tokenBudget - currentTokenCount with no double-counting.

- PI main assemble: budget block now uses the shared helper; count is
  messages-only
- PI tool loop: new assembleTokenBudget hook param carries the loop
  budget; afterTurn keeps the raw window per its contract; the hook
  resolves getRuntimeContext once per iteration instead of twice
- Codex harness: derives the budget with reserve 0 (native compaction
  owns the reserve concept there)
- plugin SDK re-exports both helpers instead of the removed
  estimatePrePromptTokens alias; docs and interface comments updated
diff --git a/docs/concepts/context-engine.md b/docs/concepts/context-engine.md
@@ -74,7 +74,7 @@ Every time OpenClaw runs a model prompt, the context engine participates at four
     Called when a new message is added to the session. The engine can store or index the message in its own data store.
   </Accordion>
   <Accordion title="2. Assemble">
-    Called before each model run. The engine returns an ordered set of messages (and an optional `systemPromptAddition`) that fit within the token budget. OpenClaw also passes an optional `currentTokenCount` — a best-effort estimate of tokens already consumed by `messages + system prompt + incoming prompt` — so engines can bound any `systemPromptAddition` against the remaining headroom.
+    Called before each model run. The engine returns an ordered set of messages (and an optional `systemPromptAddition`) that fit within the token budget. OpenClaw also passes an optional `currentTokenCount` — a best-effort estimate of the tokens the pre-assembly `messages` already consume — so engines can bound any `systemPromptAddition` against the remaining headroom.
   </Accordion>
   <Accordion title="3. Compact">
     Called when the context window is full, or when the user runs `/compact`. The engine summarizes older history to free space.
@@ -108,19 +108,18 @@ The `assemble` method can return a `systemPromptAddition` string. OpenClaw prepe
 `assemble` receives two values that together let engines bound their injection:
 
 <ParamField path="tokenBudget" type="number">
-  The model's full input context window for this run (e.g. `200_000`, `1_048_576`). Not the headroom remaining — engines must subtract what is already consumed.
+  The budget available for the assembled messages. The runtime computes this as the model's context window minus its compaction reserve (when the harness holds one) and the rendered system/user prompt pressure — engines do not subtract those themselves.
 </ParamField>
 <ParamField path="currentTokenCount" type="number">
-  Best-effort pre-assembly estimate of tokens already consumed by `messages + system prompt + incoming prompt`. Computed by the runtime via the same estimator the preemptive overflow precheck uses, so engines see the same number the runtime would. Optional — may be `undefined` on older OpenClaw runtimes; treat absence as "no headroom info, inject conservatively or skip the guard."
+  Best-effort estimate of the tokens the pre-assembly `messages` already consume, computed with the same transcript estimator the runtime uses (`estimateTranscriptTokenPressure` on `openclaw/plugin-sdk/agent-harness-runtime`), so engines see the same number the runtime would. Optional — `undefined` on older OpenClaw runtimes; treat absence as "no headroom info, inject conservatively or skip the guard."
 </ParamField>
 
-A typical guard looks like:
+Both values live on the same scale, so the remaining headroom for a `systemPromptAddition` is their difference:
 
 ```ts
-const responseReserve = 8_192; // assistant response budget; not in currentTokenCount
 const remaining =
   tokenBudget !== undefined && currentTokenCount !== undefined
-    ? Math.max(0, tokenBudget - currentTokenCount - responseReserve)
+    ? Math.max(0, tokenBudget - currentTokenCount)
     : undefined;
 
 if (remaining !== undefined && estimateTokens(additionContent) > remaining) {
diff --git a/extensions/codex/src/app-server/run-attempt.context-engine.test.ts b/extensions/codex/src/app-server/run-attempt.context-engine.test.ts
@@ -408,7 +408,12 @@ describe("runCodexAppServerAttempt context-engine lifecycle", () => {
     >[0];
     expect(assembleParams.sessionId).toBe("session-1");
     expect(assembleParams.sessionKey).toBe("agent:main:session-1");
-    expect(assembleParams.tokenBudget).toBe(321);
+    // The harness passes the message budget: contextTokenBudget (321) minus the
+    // rendered developer-instructions/prompt pressure, with no reserve on the
+    // codex path (native compaction).
+    expect(typeof assembleParams.tokenBudget).toBe("number");
+    expect(assembleParams.tokenBudget as number).toBeGreaterThan(0);
+    expect(assembleParams.tokenBudget as number).toBeLessThan(321);
     expect(assembleParams.citationsMode).toBe("on");
     expect(assembleParams.model).toBe("gpt-5.4-codex");
     expect(assembleParams.runtimeSettings).toMatchObject({
diff --git a/extensions/codex/src/app-server/run-attempt.ts b/extensions/codex/src/app-server/run-attempt.ts
@@ -11,7 +11,8 @@ import {
   clearActiveEmbeddedRun,
   embeddedAgentLog,
   emitAgentEvent as emitGlobalAgentEvent,
-  estimateLlmBoundaryTokenPressure,
+  computeContextEngineMessageBudget,
+  estimateTranscriptTokenPressure,
   finalizeHarnessContextEngineTurn,
   FAST_MODE_AUTO_PROGRESS_KIND,
   formatFastModeAutoProgressText,
@@ -1025,18 +1026,27 @@ export async function runCodexAppServerAttempt(
     if (!activeContextEngine) {
       return;
     }
-    // Pre-assembly token estimate so engines can bound any systemPromptAddition against tokenBudget.
-    const preassemblyCurrentTokenCount = estimateLlmBoundaryTokenPressure({
-      messages: historyMessages,
-      systemPrompt: developerInstructions,
-      prompt: params.prompt ?? "",
-    });
+    // Codex app-server owns compaction natively, so no OpenClaw compaction
+    // reserve is held back on this path — the budget only sets aside the
+    // rendered developer-instructions/prompt pressure.
+    const assembleTokenBudget =
+      typeof params.contextTokenBudget === "number"
+        ? computeContextEngineMessageBudget({
+            contextWindowTokens: params.contextTokenBudget,
+            compactionReserveTokens: 0,
+            systemPrompt: developerInstructions,
+            prompt: params.prompt ?? "",
+          })
+        : undefined;
+    // Messages-only estimate: engines size any systemPromptAddition as
+    // tokenBudget - currentTokenCount.
+    const preassemblyCurrentTokenCount = estimateTranscriptTokenPressure(historyMessages);
     const assembled = await assembleHarnessContextEngine({
       contextEngine: activeContextEngine,
       sessionId: activeSessionId,
       sessionKey: contextSessionKey,
       messages: historyMessages,
-      tokenBudget: params.contextTokenBudget,
+      tokenBudget: assembleTokenBudget,
       currentTokenCount: preassemblyCurrentTokenCount,
       availableTools: new Set(
         flattenCodexDynamicToolFunctions(toolBridge.availableSpecs)
diff --git a/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.context-engine.test.ts b/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.context-engine.test.ts
@@ -3013,6 +3013,11 @@ describe("runEmbeddedAttempt context engine mid-turn precheck integration", () =
     });
     expect(typeof runtimeContext?.currentTokenCount).toBe("number");
     expect(runtimeContext?.currentTokenCount as number).toBeGreaterThan(0);
+
+    // The install site also derives a loop assemble budget (window minus
+    // reserve and rendered system-prompt pressure) for the hook.
+    expect(typeof loopHookParams.assembleTokenBudget).toBe("number");
+    expect(loopHookParams.assembleTokenBudget as number).toBeGreaterThan(0);
   });
 
   it("recovers when the runtime persists the mid-turn precheck as an assistant error", async () => {
diff --git a/src/agents/embedded-agent-runner/run/attempt.ts b/src/agents/embedded-agent-runner/run/attempt.ts
@@ -499,8 +499,9 @@ import {
 import {
   PREEMPTIVE_OVERFLOW_ERROR_TEXT,
   buildPrePromptContextBudgetStatus,
+  computeContextEngineMessageBudget,
   estimateLlmBoundaryTokenPressure,
-  estimateRenderedLlmBoundaryTokenPressure,
+  estimateTranscriptTokenPressure,
   formatPrePromptPrecheckLog,
   shouldPreemptivelyCompactBeforePrompt,
 } from "./preemptive-compaction.js";
@@ -2671,13 +2672,23 @@ export async function runEmbeddedAttempt(
           fallbackReason: params.fallbackReason,
           degradedReason: params.degradedReason,
         });
+        // Mid-tool-loop assemble calls carry no new user prompt — the original
+        // prompt is already part of the transcript — so the loop budget only
+        // sets aside the system prompt and the compaction reserve.
+        const loopAssembleTokenBudget = computeContextEngineMessageBudget({
+          contextWindowTokens: contextTokenBudgetForGuard,
+          compactionReserveTokens: settingsManager.getCompactionReserveTokens(),
+          systemPrompt: systemPromptText,
+          prompt: "",
+        });
         const removeContextEngineLoopHook = installContextEngineLoopHook({
           agent: activeSession.agent,
           contextEngine: activeContextEngine,
           sessionId: params.sessionId,
           sessionKey: params.sessionKey,
           sessionFile: params.sessionFile,
           tokenBudget: params.contextTokenBudget,
+          assembleTokenBudget: loopAssembleTokenBudget,
           modelId: params.modelId,
           ...(transcriptPolicy.repairToolUseResultPairing
             ? {
@@ -2696,14 +2707,9 @@ export async function runEmbeddedAttempt(
               cwd: effectiveCwd,
               agentDir,
               tokenBudget: params.contextTokenBudget,
-              // Mid-tool-loop the next model call has no new user prompt — the
-              // continuation is driven by tool results that are already in
-              // `messages`. Estimate against messages + system prompt only.
-              currentTokenCount: estimateLlmBoundaryTokenPressure({
-                messages,
-                systemPrompt: systemPromptText,
-                prompt: "",
-              }),
+              // Messages-only estimate to match the assemble contract: the loop
+              // assemble budget already sets aside the system prompt and reserve.
+              currentTokenCount: estimateTranscriptTokenPressure(messages),
               promptCache:
                 promptCache ??
                 buildLoopPromptCacheInfo({
@@ -3372,25 +3378,18 @@ export async function runEmbeddedAttempt(
                   DEFAULT_CONTEXT_TOKENS,
               ),
             );
-            const contextEngineAssemblePromptBudget = Math.max(
-              1,
-              contextEngineAssembleContextTokenBudget - contextEngineAssembleReserveTokens,
-            );
-            const contextEngineAssembleRenderedPromptTokens =
-              estimateRenderedLlmBoundaryTokenPressure({
-                systemPrompt: systemPromptText,
-                prompt: params.prompt ?? "",
-              });
-            const contextEngineAssembleMessageBudget = Math.max(
-              1,
-              contextEngineAssemblePromptBudget - contextEngineAssembleRenderedPromptTokens,
-            );
-            // Pre-assembly token estimate so engines can bound any systemPromptAddition against tokenBudget.
-            const preassemblyCurrentTokenCount = estimateLlmBoundaryTokenPressure({
-              messages: activeSession.messages,
+            const contextEngineAssembleMessageBudget = computeContextEngineMessageBudget({
+              contextWindowTokens: contextEngineAssembleContextTokenBudget,
+              compactionReserveTokens: contextEngineAssembleReserveTokens,
               systemPrompt: systemPromptText,
               prompt: params.prompt ?? "",
             });
+            // Messages-only estimate: tokenBudget already accounts for the rendered
+            // system/user prompt and the compaction reserve, so engines size any
+            // systemPromptAddition as tokenBudget - currentTokenCount.
+            const preassemblyCurrentTokenCount = estimateTranscriptTokenPressure(
+              activeSession.messages,
+            );
             const assembled = await assembleAttemptContextEngine({
               contextEngine: activeContextEngine,
               sessionId: params.sessionId,
diff --git a/src/agents/embedded-agent-runner/run/preemptive-compaction.ts b/src/agents/embedded-agent-runner/run/preemptive-compaction.ts
@@ -287,6 +287,41 @@ export function estimateRenderedLlmBoundaryTokenPressure(params: {
   return Math.max(0, Math.ceil((systemTokens + promptTokens) * SAFETY_MARGIN));
 }
 
+/**
+ * Estimates the token pressure of transcript messages alone — no system prompt or
+ * incoming prompt. This is the estimator behind `ContextEngine.assemble`'s
+ * `currentTokenCount`: it pairs with `computeContextEngineMessageBudget` so engines
+ * can compute assemble headroom as `tokenBudget - currentTokenCount`.
+ */
+export function estimateTranscriptTokenPressure(messages: AgentMessage[]): number {
+  const historyTokens = messages.reduce(
+    (sum, message) => sum + estimateMessageTokenPressure(message),
+    0,
+  );
+  return Math.max(0, Math.ceil(historyTokens * SAFETY_MARGIN));
+}
+
+/**
+ * Budget available to a context engine's assembled messages: the resolved context
+ * window minus the runtime's compaction reserve and the rendered system/user prompt
+ * pressure. Every `ContextEngine.assemble` call site derives `tokenBudget` through
+ * this helper so the field means the same thing on every harness path.
+ */
+export function computeContextEngineMessageBudget(params: {
+  contextWindowTokens: number;
+  compactionReserveTokens: number;
+  systemPrompt?: string;
+  prompt: string;
+}): number {
+  const reserveTokens = Math.max(0, Math.floor(params.compactionReserveTokens));
+  const promptBudget = Math.max(1, Math.floor(params.contextWindowTokens) - reserveTokens);
+  const renderedPromptTokens = estimateRenderedLlmBoundaryTokenPressure({
+    systemPrompt: params.systemPrompt,
+    prompt: params.prompt,
+  });
+  return Math.max(1, promptBudget - renderedPromptTokens);
+}
+
 function normalizeLlmBoundaryTokenPressure(
   pressure: LlmBoundaryTokenPressure | undefined,
 ): LlmBoundaryTokenPressure | undefined {
diff --git a/src/agents/embedded-agent-runner/tool-result-context-guard.test.ts b/src/agents/embedded-agent-runner/tool-result-context-guard.test.ts
@@ -737,6 +737,29 @@ describe("installContextEngineLoopHook", () => {
     expect(assembleParams).not.toHaveProperty("currentTokenCount");
   });
 
+  it("prefers assembleTokenBudget for assemble while afterTurn keeps tokenBudget", async () => {
+    const agent = makeGuardableAgent();
+    const engine = makeMockEngine();
+    installContextEngineLoopHook({
+      agent,
+      contextEngine: engine,
+      sessionId,
+      sessionKey,
+      sessionFile,
+      tokenBudget,
+      assembleTokenBudget: 2048,
+      modelId,
+      getPrePromptMessageCount: () => 1,
+    });
+
+    const messages = [makeUser("first"), makeToolResult("call_1", "result")];
+    await callTransform(agent, messages);
+
+    expect(recordMockArg(engine.afterTurn).tokenBudget).toBe(tokenBudget);
+    const assembleParams = engine.assemble.mock.calls.at(0)?.[0];
+    expect(assembleParams?.tokenBudget).toBe(2048);
+  });
+
   it("passes loop messages and the prompt fence into the runtimeContext callback", async () => {
     const agent = makeGuardableAgent();
     const engine = makeMockEngine();
diff --git a/src/agents/embedded-agent-runner/tool-result-context-guard.ts b/src/agents/embedded-agent-runner/tool-result-context-guard.ts
@@ -330,6 +330,8 @@ export function installContextEngineLoopHook(params: {
   sessionKey?: string;
   sessionFile: string;
   tokenBudget?: number;
+  /** Budget for assemble calls: context window minus compaction reserve and rendered system-prompt pressure. Falls back to `tokenBudget` when absent. */
+  assembleTokenBudget?: number;
   modelId: string;
   repairAssembledMessages?: (messages: AgentMessage[]) => AgentMessage[];
   getPrePromptMessageCount?: () => number;
@@ -393,6 +395,12 @@ export function installContextEngineLoopHook(params: {
       return lastAssembledView ?? providerMessages;
     }
     try {
+      // Resolved once per iteration: afterTurn and assemble share it, so the
+      // currentTokenCount estimate inside the callback is computed once.
+      const loopRuntimeContext = params.getRuntimeContext?.({
+        messages: transcriptMessages,
+        prePromptMessageCount,
+      });
       if (typeof contextEngine.afterTurn === "function") {
         await contextEngine.afterTurn({
           sessionId,
@@ -401,10 +409,7 @@ export function installContextEngineLoopHook(params: {
           messages: transcriptMessages,
           prePromptMessageCount,
           tokenBudget,
-          runtimeContext: params.getRuntimeContext?.({
-            messages: transcriptMessages,
-            prePromptMessageCount,
-          }),
+          runtimeContext: loopRuntimeContext,
           runtimeSettings: params.runtimeSettings,
           isHeartbeat: params.isHeartbeat,
         });
@@ -433,15 +438,12 @@ export function installContextEngineLoopHook(params: {
       lastSeenLength = transcriptMessages.length;
       params.onAfterTurnCheckpoint?.(lastSeenLength);
       lastSourceMessages = transcriptMessages;
-      const loopCurrentTokenCount = params.getRuntimeContext?.({
-        messages: sourceMessages,
-        prePromptMessageCount,
-      })?.currentTokenCount;
+      const loopCurrentTokenCount = loopRuntimeContext?.currentTokenCount;
       const assembled = await contextEngine.assemble({
         sessionId,
         sessionKey,
         messages: providerMessages,
-        tokenBudget,
+        tokenBudget: params.assembleTokenBudget ?? tokenBudget,
         ...(typeof loopCurrentTokenCount === "number"
           ? { currentTokenCount: loopCurrentTokenCount }
           : {}),
diff --git a/src/context-engine/types.ts b/src/context-engine/types.ts
@@ -381,8 +381,9 @@ export interface ContextEngine {
     sessionId: string;
     sessionKey?: string;
     messages: AgentMessage[];
+    /** Budget available for the assembled messages. The runtime has already subtracted its compaction reserve (when it holds one) and the rendered system/user prompt pressure from the context window. */
     tokenBudget?: number;
-    /** Best-effort pre-assembly prompt token estimate so engines can bound any `systemPromptAddition` against `tokenBudget`. */
+    /** Best-effort token estimate of `messages` alone, from the runtime's transcript estimator. Remaining headroom for any `systemPromptAddition` is `tokenBudget - currentTokenCount`. Omitted on runtimes that predate the field. */
     currentTokenCount?: number;
     /** Tool names available for this run so engines can align prompt guidance with runtime tool access. */
     availableTools?: Set<string>;
diff --git a/src/plugin-sdk/agent-harness-runtime.ts b/src/plugin-sdk/agent-harness-runtime.ts
@@ -357,8 +357,9 @@ export {
   resolveCompactionTimeoutMs,
 } from "../agents/embedded-agent-runner/compaction-safety-timeout.js";
 export {
-  estimateLlmBoundaryTokenPressure,
+  computeContextEngineMessageBudget,
   estimateRenderedLlmBoundaryTokenPressure,
+  estimateTranscriptTokenPressure,
   formatPrePromptPrecheckLog,
   PREEMPTIVE_OVERFLOW_ERROR_TEXT,
   shouldPreemptivelyCompactBeforePrompt,
