Skip to content

Commit 94d9056

Browse files
jalehmanRomneyDa
jalehman
authored and
RomneyDa
committed
fix: preserve preflight overflow token counts into recovery budgeting
Budget context engine assembly against the reserve and rendered prompt pressure, and carry the preflight estimated prompt tokens, prompt budget, and overflow tokens into the outer overflow recovery loop so compaction engines compact against the prompt OpenClaw actually rendered instead of a minimally over-budget guess.
1 parent c7295e4 commit 94d9056

6 files changed

Lines changed: 158 additions & 8 deletions

File tree

‎src/agents/embedded-agent-runner/run.overflow-compaction.test.ts‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2294,6 +2294,41 @@ describe("runEmbeddedAgent overflow compaction trigger routing", () => {
22942294
expect(result.meta.error).toBeUndefined();
22952295
});
22962296

2297+
it("passes preflight prompt estimates into synthetic overflow compaction", async () => {
2298+
mockedExtractObservedOverflowTokenCount.mockReturnValueOnce(undefined);
2299+
mockedRunEmbeddedAttempt
2300+
.mockResolvedValueOnce(
2301+
makeAttemptResult({
2302+
promptError: makeOverflowError(),
2303+
promptErrorSource: "precheck",
2304+
preflightRecovery: {
2305+
route: "compact_then_truncate",
2306+
estimatedPromptTokens: 268138,
2307+
promptBudgetBeforeReserve: 241616,
2308+
overflowTokens: 26522,
2309+
},
2310+
}),
2311+
)
2312+
.mockResolvedValueOnce(makeAttemptResult({ promptError: null }));
2313+
mockedCompactDirect.mockResolvedValueOnce(
2314+
makeCompactionSuccess({
2315+
summary: "Compacted session",
2316+
firstKeptEntryId: "entry-preflight",
2317+
tokensBefore: 268138,
2318+
}),
2319+
);
2320+
2321+
const result = await runEmbeddedAgent(overflowBaseRunParams);
2322+
2323+
expectMockCallFields(mockedCompactDirect, {
2324+
currentTokenCount: 268138,
2325+
});
2326+
expectRecordFields(expectMockCallFields(mockedCompactDirect, {}).runtimeContext, {
2327+
currentTokenCount: 268138,
2328+
});
2329+
expect(result.meta.error).toBeUndefined();
2330+
});
2331+
22972332
it("passes minimally over-budget count when overflow text is confirmed but unparseable", async () => {
22982333
mockedExtractObservedOverflowTokenCount.mockReturnValueOnce(undefined);
22992334
mockedRunEmbeddedAttempt

‎src/agents/embedded-agent-runner/run.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2674,8 +2674,15 @@ async function runEmbeddedAgentInternal(
26742674
const errorText = contextOverflowError.text;
26752675
const msgCount = attempt.messagesSnapshot?.length ?? 0;
26762676
const observedOverflowTokens = extractObservedOverflowTokenCount(errorText);
2677+
const preflightEstimatedPromptTokens =
2678+
typeof preflightRecovery?.estimatedPromptTokens === "number" &&
2679+
Number.isFinite(preflightRecovery.estimatedPromptTokens) &&
2680+
preflightRecovery.estimatedPromptTokens > 0
2681+
? Math.ceil(preflightRecovery.estimatedPromptTokens)
2682+
: undefined;
26772683
const overflowTokenCountForCompaction =
26782684
observedOverflowTokens ??
2685+
preflightEstimatedPromptTokens ??
26792686
(ctxInfo.tokens > 0
26802687
? // Confirmed overflow with an unparseable provider message still carries a
26812688
// minimally over-budget count for compaction engines and diagnostics.
@@ -2687,6 +2694,7 @@ async function runEmbeddedAgentInternal(
26872694
`messages=${msgCount} sessionFile=${activeSessionFile} ` +
26882695
`diagId=${overflowDiagId} compactionAttempts=${overflowCompactionAttempts} ` +
26892696
`observedTokens=${observedOverflowTokens ?? "unknown"} ` +
2697+
`preflightEstimatedTokens=${preflightEstimatedPromptTokens ?? "unknown"} ` +
26902698
`compactionTokens=${overflowTokenCountForCompaction ?? "unknown"} ` +
26912699
`error=${errorText.slice(0, 200)}`,
26922700
);

‎src/agents/embedded-agent-runner/run/attempt.spawn-workspace.context-engine.test.ts‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3023,7 +3023,13 @@ describe("runEmbeddedAttempt context engine mid-turn precheck integration", () =
30233023
});
30243024

30253025
expect(result.promptErrorSource).toBe("precheck");
3026-
expect(result.preflightRecovery).toEqual({ route: "compact_only", source: "mid-turn" });
3026+
expect(result.preflightRecovery).toEqual({
3027+
route: "compact_only",
3028+
source: "mid-turn",
3029+
estimatedPromptTokens: 9000,
3030+
promptBudgetBeforeReserve: 7000,
3031+
overflowTokens: 2000,
3032+
});
30273033
expect(result.messagesSnapshot).toEqual([seedMessage]);
30283034
});
30293035
});
@@ -3066,6 +3072,34 @@ describe("runEmbeddedAttempt tool-result guard budget wiring", () => {
30663072
).toBe(1_000_000);
30673073
});
30683074

3075+
it("passes context engines the message budget after reserve and rendered prompt pressure", async () => {
3076+
const contextEngine = createContextEngineBootstrapAndAssemble();
3077+
hoisted.compactionReserveTokens = 20_000;
3078+
3079+
await createContextEngineAttemptRunner({
3080+
contextEngine,
3081+
sessionKey,
3082+
tempPaths,
3083+
attemptOverrides: {
3084+
contextTokenBudget: 100_000,
3085+
prompt: "current prompt",
3086+
transcriptPrompt: "current prompt",
3087+
},
3088+
});
3089+
3090+
const assembleParams = mockParams(
3091+
contextEngine.assemble as MockCallSource,
3092+
0,
3093+
"assemble params",
3094+
);
3095+
expect(assembleParams.tokenBudget).toBeLessThan(80_000);
3096+
expect(assembleParams.runtimeSettings).toMatchObject({
3097+
limits: {
3098+
maxOutputTokens: 20_000,
3099+
},
3100+
});
3101+
});
3102+
30693103
it("preserves the cacheable prefix while bounding current prompt results", async () => {
30703104
const toolText = "process output ".repeat(70);
30713105
const sessionMessages: AgentMessage[] = [{ role: "user", content: "seed", timestamp: 1 }];

‎src/agents/embedded-agent-runner/run/attempt.spawn-workspace.test-support.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ type AttemptSpawnWorkspaceHoisted = {
9696
>;
9797
limitHistoryTurnsMock: Mock<<T>(messages: T, limit: number | undefined) => T>;
9898
preemptiveCompactionCalls: Parameters<ShouldPreemptivelyCompactBeforePromptFn>[0][];
99+
compactionReserveTokens: number;
99100
systemPromptTexts: string[];
100101
embeddedSystemPromptInputs: unknown[];
101102
sessionManager: SessionManagerMocks;
@@ -201,6 +202,7 @@ const hoisted = vi.hoisted((): AttemptSpawnWorkspaceHoisted => {
201202
(messages) => messages,
202203
);
203204
const preemptiveCompactionCalls: Parameters<ShouldPreemptivelyCompactBeforePromptFn>[0][] = [];
205+
const compactionReserveTokens = 0;
204206
const systemPromptTexts: string[] = [];
205207
const embeddedSystemPromptInputs: unknown[] = [];
206208
const sessionManager = {
@@ -245,6 +247,7 @@ const hoisted = vi.hoisted((): AttemptSpawnWorkspaceHoisted => {
245247
getHistoryLimitFromSessionKeyMock,
246248
limitHistoryTurnsMock,
247249
preemptiveCompactionCalls,
250+
compactionReserveTokens,
248251
systemPromptTexts,
249252
embeddedSystemPromptInputs,
250253
sessionManager,
@@ -437,7 +440,7 @@ vi.mock("../../docs-path.js", () => ({
437440
vi.mock("../../agent-project-settings.js", () => ({
438441
createPreparedEmbeddedAgentSettingsManager: () => ({
439442
reload: async () => {},
440-
getCompactionReserveTokens: () => 0,
443+
getCompactionReserveTokens: () => hoisted.compactionReserveTokens,
441444
getCompactionKeepRecentTokens: () => 40_000,
442445
getDefaultProvider: () => undefined,
443446
getDefaultModel: () => undefined,
@@ -1033,6 +1036,7 @@ export function resetEmbeddedAttemptHarness(
10331036
hoisted.getHistoryLimitFromSessionKeyMock.mockReset().mockReturnValue(undefined);
10341037
hoisted.limitHistoryTurnsMock.mockReset().mockImplementation((messages) => messages);
10351038
hoisted.preemptiveCompactionCalls.length = 0;
1039+
hoisted.compactionReserveTokens = 0;
10361040
hoisted.systemPromptTexts.length = 0;
10371041
hoisted.embeddedSystemPromptInputs.length = 0;
10381042
hoisted.sessionManager.getLeafEntry.mockReset().mockReturnValue(null);

‎src/agents/embedded-agent-runner/run/attempt.ts‎

Lines changed: 69 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,7 @@ import {
500500
PREEMPTIVE_OVERFLOW_ERROR_TEXT,
501501
buildPrePromptContextBudgetStatus,
502502
estimateLlmBoundaryTokenPressure,
503+
estimateRenderedLlmBoundaryTokenPressure,
503504
formatPrePromptPrecheckLog,
504505
shouldPreemptivelyCompactBeforePrompt,
505506
} from "./preemptive-compaction.js";
@@ -510,6 +511,22 @@ import {
510511
} from "./runtime-context-prompt.js";
511512
import type { EmbeddedRunAttemptParams, EmbeddedRunAttemptResult } from "./types.js";
512513

514+
type PreflightRecoveryBudgetSnapshot = Pick<
515+
MidTurnPrecheckRequest,
516+
"estimatedPromptTokens" | "promptBudgetBeforeReserve" | "overflowTokens"
517+
>;
518+
519+
// Carries the measured prompt budget into the outer recovery loop. The synthetic
520+
// precheck error is only a routing signal, so compaction engines need these
521+
// fields to compact against the prompt OpenClaw actually rendered.
522+
function buildPreflightRecoveryBudgetSnapshot(snapshot: PreflightRecoveryBudgetSnapshot) {
523+
return {
524+
estimatedPromptTokens: snapshot.estimatedPromptTokens,
525+
promptBudgetBeforeReserve: snapshot.promptBudgetBeforeReserve,
526+
overflowTokens: snapshot.overflowTokens,
527+
};
528+
}
529+
513530
export {
514531
appendAttemptCacheTtlIfNeeded,
515532
composeSystemPromptWithHookContext,
@@ -3334,15 +3351,42 @@ export async function runEmbeddedAttempt(
33343351
// history in place would otherwise leave the precheck reading
33353352
// already-windowed messages instead of the true pre-assembly state.
33363353
const preassemblyContextEngineMessagesForPrecheck = activeSession.messages.slice();
3354+
const contextEngineAssembleReserveTokens = Math.max(
3355+
0,
3356+
Math.floor(settingsManager.getCompactionReserveTokens()),
3357+
);
3358+
const contextEngineAssembleContextTokenBudget = Math.max(
3359+
1,
3360+
Math.floor(
3361+
params.contextTokenBudget ??
3362+
params.model.contextWindow ??
3363+
params.model.maxTokens ??
3364+
DEFAULT_CONTEXT_TOKENS,
3365+
),
3366+
);
3367+
const contextEngineAssemblePromptBudget = Math.max(
3368+
1,
3369+
contextEngineAssembleContextTokenBudget - contextEngineAssembleReserveTokens,
3370+
);
3371+
const contextEngineAssembleRenderedPromptTokens =
3372+
estimateRenderedLlmBoundaryTokenPressure({
3373+
systemPrompt: systemPromptText,
3374+
prompt: params.prompt ?? "",
3375+
});
3376+
const contextEngineAssembleMessageBudget = Math.max(
3377+
1,
3378+
contextEngineAssemblePromptBudget - contextEngineAssembleRenderedPromptTokens,
3379+
);
33373380
const assembled = await assembleAttemptContextEngine({
33383381
contextEngine: activeContextEngine,
33393382
sessionId: params.sessionId,
33403383
sessionKey: params.sessionKey,
33413384
messages: activeSession.messages,
3342-
tokenBudget: params.contextTokenBudget,
3385+
tokenBudget: contextEngineAssembleMessageBudget,
33433386
availableTools: new Set(capabilityToolNames),
33443387
citationsMode: params.config?.memory?.citations,
33453388
modelId: params.modelId,
3389+
maxOutputTokens: contextEngineAssembleReserveTokens,
33463390
contextEngineHostSupport: OPENCLAW_EMBEDDED_CONTEXT_ENGINE_HOST,
33473391
providerId: params.provider,
33483392
requestedModelId: params.requestedModelId,
@@ -3911,6 +3955,7 @@ export async function runEmbeddedAttempt(
39113955
preflightRecovery = {
39123956
route: "truncate_tool_results_only",
39133957
source: "mid-turn",
3958+
...buildPreflightRecoveryBudgetSnapshot(request),
39143959
handled: true,
39153960
truncatedCount: truncationResult.truncatedCount,
39163961
};
@@ -3921,7 +3966,11 @@ export async function runEmbeddedAttempt(
39213966
`handled=true truncatedCount=${truncationResult.truncatedCount}`,
39223967
);
39233968
} else {
3924-
preflightRecovery = { route: "compact_only", source: "mid-turn" };
3969+
preflightRecovery = {
3970+
route: "compact_only",
3971+
source: "mid-turn",
3972+
...buildPreflightRecoveryBudgetSnapshot(request),
3973+
};
39253974
promptError = new Error(PREEMPTIVE_OVERFLOW_ERROR_TEXT);
39263975
promptErrorSource = "precheck";
39273976
logMidTurnPrecheck(
@@ -3930,7 +3979,11 @@ export async function runEmbeddedAttempt(
39303979
);
39313980
}
39323981
} else {
3933-
preflightRecovery = { route: request.route, source: "mid-turn" };
3982+
preflightRecovery = {
3983+
route: request.route,
3984+
source: "mid-turn",
3985+
...buildPreflightRecoveryBudgetSnapshot(request),
3986+
};
39343987
promptError = new Error(PREEMPTIVE_OVERFLOW_ERROR_TEXT);
39353988
promptErrorSource = "precheck";
39363989
logMidTurnPrecheck(request.route);
@@ -4761,6 +4814,7 @@ export async function runEmbeddedAttempt(
47614814
if (truncationResult.truncated) {
47624815
preflightRecovery = {
47634816
route: "truncate_tool_results_only",
4817+
...buildPreflightRecoveryBudgetSnapshot(preemptiveCompaction),
47644818
handled: true,
47654819
truncatedCount: truncationResult.truncatedCount,
47664820
};
@@ -4783,7 +4837,10 @@ export async function runEmbeddedAttempt(
47834837
`${params.provider}/${params.modelId}; falling back to compaction ` +
47844838
`reason=${truncationResult.reason ?? "unknown"} sessionFile=${params.sessionFile}`,
47854839
);
4786-
preflightRecovery = { route: "compact_only" };
4840+
preflightRecovery = {
4841+
route: "compact_only",
4842+
...buildPreflightRecoveryBudgetSnapshot(preemptiveCompaction),
4843+
};
47874844
promptError = new Error(PREEMPTIVE_OVERFLOW_ERROR_TEXT);
47884845
promptErrorSource = "precheck";
47894846
skipPromptSubmission = true;
@@ -4792,8 +4849,14 @@ export async function runEmbeddedAttempt(
47924849
if (preemptiveCompaction?.shouldCompact) {
47934850
preflightRecovery =
47944851
preemptiveCompaction.route === "compact_then_truncate"
4795-
? { route: "compact_then_truncate" }
4796-
: { route: "compact_only" };
4852+
? {
4853+
route: "compact_then_truncate",
4854+
...buildPreflightRecoveryBudgetSnapshot(preemptiveCompaction),
4855+
}
4856+
: {
4857+
route: "compact_only",
4858+
...buildPreflightRecoveryBudgetSnapshot(preemptiveCompaction),
4859+
};
47974860
promptError = new Error(PREEMPTIVE_OVERFLOW_ERROR_TEXT);
47984861
promptErrorSource = "precheck";
47994862
log.warn(

‎src/agents/embedded-agent-runner/run/types.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,12 +142,18 @@ export type EmbeddedRunAttemptResult = {
142142
| {
143143
route: Exclude<PreemptiveCompactionRoute, "fits">;
144144
source?: "mid-turn";
145+
estimatedPromptTokens?: number;
146+
promptBudgetBeforeReserve?: number;
147+
overflowTokens?: number;
145148
handled: true;
146149
truncatedCount?: number;
147150
}
148151
| {
149152
route: Exclude<PreemptiveCompactionRoute, "fits">;
150153
source?: "mid-turn";
154+
estimatedPromptTokens?: number;
155+
promptBudgetBeforeReserve?: number;
156+
overflowTokens?: number;
151157
handled?: false;
152158
};
153159
sessionIdUsed: string;

0 commit comments

Comments
 (0)