Skip to content

Commit eb4b8af

Browse files
author
Gio Della-Libera
committed
feat(agents): persist estimated context budget status
1 parent d003bc2 commit eb4b8af

14 files changed

Lines changed: 398 additions & 3 deletions

File tree

‎src/agents/command/session-store.test.ts‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -623,6 +623,75 @@ describe("updateSessionStoreAfterAgentRun", () => {
623623
});
624624
});
625625

626+
it("persists estimated context budget status without marking stale usage fresh", async () => {
627+
await withTempSessionStore(async ({ storePath }) => {
628+
const cfg = {} as OpenClawConfig;
629+
const sessionKey = "agent:main:explicit:test-context-budget-status";
630+
const sessionId = "test-context-budget-status-session";
631+
const sessionStore: Record<string, SessionEntry> = {
632+
[sessionKey]: {
633+
sessionId,
634+
updatedAt: 1,
635+
totalTokens: 21225,
636+
totalTokensFresh: true,
637+
},
638+
};
639+
await fs.writeFile(storePath, JSON.stringify(sessionStore, null, 2));
640+
641+
const result: EmbeddedPiRunResult = {
642+
meta: {
643+
durationMs: 500,
644+
agentMeta: {
645+
sessionId,
646+
provider: "minimax",
647+
model: "MiniMax-M2.7",
648+
contextBudgetStatus: {
649+
schemaVersion: 1,
650+
source: "pre-prompt-estimate",
651+
updatedAt: 123,
652+
provider: "minimax",
653+
model: "MiniMax-M2.7",
654+
route: "fits",
655+
shouldCompact: false,
656+
estimatedPromptTokens: 18_000,
657+
contextTokenBudget: 32_000,
658+
promptBudgetBeforeReserve: 28_000,
659+
reserveTokens: 4_000,
660+
effectiveReserveTokens: 4_000,
661+
remainingPromptBudgetTokens: 10_000,
662+
overflowTokens: 0,
663+
toolResultReducibleChars: 0,
664+
messageCount: 4,
665+
unwindowedMessageCount: 4,
666+
},
667+
},
668+
},
669+
};
670+
671+
await updateSessionStoreAfterAgentRun({
672+
cfg,
673+
sessionId,
674+
sessionKey,
675+
storePath,
676+
sessionStore,
677+
defaultProvider: "minimax",
678+
defaultModel: "MiniMax-M2.7",
679+
result,
680+
});
681+
682+
expect(sessionStore[sessionKey]?.totalTokens).toBe(21225);
683+
expect(sessionStore[sessionKey]?.totalTokensFresh).toBe(false);
684+
expect(sessionStore[sessionKey]?.contextBudgetStatus).toMatchObject({
685+
source: "pre-prompt-estimate",
686+
estimatedPromptTokens: 18_000,
687+
contextTokenBudget: 32_000,
688+
});
689+
690+
const persisted = loadSessionStore(storePath);
691+
expect(persisted[sessionKey]?.contextBudgetStatus?.estimatedPromptTokens).toBe(18_000);
692+
});
693+
});
694+
626695
it("does not treat CLI cumulative usage as a fresh context snapshot", async () => {
627696
await withTempSessionStore(async ({ storePath }) => {
628697
const cfg = {

‎src/agents/command/session-store.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ export async function updateSessionStoreAfterAgentRun(params: {
9595
const providerUsed = result.meta.agentMeta?.provider ?? fallbackProvider ?? defaultProvider;
9696
const agentHarnessId = normalizeOptionalString(result.meta.agentMeta?.agentHarnessId);
9797
const runtimeContextTokens = resolvePositiveInteger(result.meta.agentMeta?.contextTokens);
98+
const contextBudgetStatus = result.meta.agentMeta?.contextBudgetStatus;
9899
const contextTokens =
99100
runtimeContextTokens !== undefined
100101
? runtimeContextTokens
@@ -174,6 +175,9 @@ export async function updateSessionStoreAfterAgentRun(params: {
174175
if (result.meta.systemPromptReport) {
175176
next.systemPromptReport = result.meta.systemPromptReport;
176177
}
178+
if (contextBudgetStatus) {
179+
next.contextBudgetStatus = contextBudgetStatus;
180+
}
177181
if (hasNonzeroUsage(usage)) {
178182
const { estimateUsageCost, resolveModelCostConfig } = await getUsageFormatModule();
179183
const input = usage.input ?? 0;

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1023,6 +1023,7 @@ export async function runEmbeddedPiAgent(
10231023
let lastRunPromptUsage: ReturnType<typeof normalizeUsage> | undefined;
10241024
let autoCompactionCount = 0;
10251025
let lastCompactionTokensAfter: number | undefined;
1026+
let lastContextBudgetStatus: EmbeddedPiAgentMeta["contextBudgetStatus"];
10261027
let runLoopIterations = 0;
10271028
let overloadProfileRotations = 0;
10281029
let planningOnlyRetryAttempts = 0;
@@ -1623,6 +1624,9 @@ export async function runEmbeddedPiAgent(
16231624
) {
16241625
lastCompactionTokensAfter = Math.floor(attempt.compactionTokensAfter);
16251626
}
1627+
if (attempt.contextBudgetStatus) {
1628+
lastContextBudgetStatus = attempt.contextBudgetStatus;
1629+
}
16261630
const activeErrorContext = resolveActiveErrorContext({
16271631
provider,
16281632
model: modelId,
@@ -2627,6 +2631,7 @@ export async function runEmbeddedPiAgent(
26272631
usage: usageMeta.usage,
26282632
lastCallUsage: usageMeta.lastCallUsage,
26292633
promptTokens: usageMeta.promptTokens,
2634+
...(lastContextBudgetStatus ? { contextBudgetStatus: lastContextBudgetStatus } : {}),
26302635
compactionCount: autoCompactionCount > 0 ? autoCompactionCount : undefined,
26312636
compactionTokensAfter: lastCompactionTokensAfter,
26322637
};

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,7 @@ import {
389389
} from "./midturn-precheck.js";
390390
import {
391391
PREEMPTIVE_OVERFLOW_ERROR_TEXT,
392+
buildPrePromptContextBudgetStatus,
392393
formatPrePromptPrecheckLog,
393394
shouldPreemptivelyCompactBeforePrompt,
394395
} from "./preemptive-compaction.js";
@@ -3320,6 +3321,7 @@ export async function runEmbeddedAttempt(
33203321
let cacheBreak: PromptCacheBreak | null = null;
33213322
let promptCache: EmbeddedRunAttemptResult["promptCache"];
33223323
let lastCallUsage: NormalizedUsage | undefined;
3324+
let contextBudgetStatus: EmbeddedRunAttemptResult["contextBudgetStatus"];
33233325
let compactionOccurredThisAttempt = false;
33243326
let finalPromptText: string | undefined;
33253327
if (params.replyOperation) {
@@ -4012,6 +4014,20 @@ export async function runEmbeddedAttempt(
40124014
}),
40134015
});
40144016
if (preemptiveCompaction) {
4017+
contextBudgetStatus = buildPrePromptContextBudgetStatus({
4018+
result: preemptiveCompaction,
4019+
provider: params.provider,
4020+
modelId: params.modelId,
4021+
messageCount: activeSession.messages.length,
4022+
contextTokenBudget,
4023+
reserveTokens,
4024+
...(params.sessionId ? { sessionId: params.sessionId } : {}),
4025+
...(contextEnginePromptAuthority === "preassembly_may_overflow" &&
4026+
unwindowedContextEngineMessagesForPrecheck
4027+
? { unwindowedMessageCount: unwindowedContextEngineMessagesForPrecheck.length }
4028+
: {}),
4029+
...(params.sessionFile ? { sessionFile: params.sessionFile } : {}),
4030+
});
40154031
log.debug(
40164032
formatPrePromptPrecheckLog({
40174033
result: preemptiveCompaction,
@@ -4818,6 +4834,7 @@ export async function runEmbeddedAttempt(
48184834
),
48194835
attemptUsage,
48204836
promptCache,
4837+
contextBudgetStatus,
48214838
compactionCount: getCompactionCount(),
48224839
compactionTokensAfter: getLastCompactionTokensAfter(),
48234840
// Client tool calls detected (OpenResponses hosted tools).

‎src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import "../../test-helpers/pi-coding-agent-token-mock.js";
44
import { estimateToolResultReductionPotential } from "../tool-result-truncation.js";
55

66
let PREEMPTIVE_OVERFLOW_ERROR_TEXT: typeof import("./preemptive-compaction.js").PREEMPTIVE_OVERFLOW_ERROR_TEXT;
7+
let buildPrePromptContextBudgetStatus: typeof import("./preemptive-compaction.js").buildPrePromptContextBudgetStatus;
78
let estimatePrePromptTokens: typeof import("./preemptive-compaction.js").estimatePrePromptTokens;
89
let formatPrePromptPrecheckLog: typeof import("./preemptive-compaction.js").formatPrePromptPrecheckLog;
910
let shouldPreemptivelyCompactBeforePrompt: typeof import("./preemptive-compaction.js").shouldPreemptivelyCompactBeforePrompt;
@@ -12,6 +13,7 @@ beforeAll(async () => {
1213
vi.resetModules();
1314
({
1415
PREEMPTIVE_OVERFLOW_ERROR_TEXT,
16+
buildPrePromptContextBudgetStatus,
1517
estimatePrePromptTokens,
1618
formatPrePromptPrecheckLog,
1719
shouldPreemptivelyCompactBeforePrompt,
@@ -131,6 +133,51 @@ describe("preemptive-compaction", () => {
131133
expect(line).toContain("unwindowedMessages=3");
132134
});
133135

136+
it("builds a durable estimated context budget status snapshot", () => {
137+
const result = shouldPreemptivelyCompactBeforePrompt({
138+
messages: [makeAssistantHistory("short history")],
139+
systemPrompt: "sys",
140+
prompt: "hello",
141+
contextTokenBudget: 10_000,
142+
reserveTokens: 1_000,
143+
});
144+
145+
const status = buildPrePromptContextBudgetStatus({
146+
result,
147+
provider: "anthropic",
148+
modelId: "claude-opus-4-6",
149+
messageCount: 1,
150+
unwindowedMessageCount: 3,
151+
contextTokenBudget: 10_000,
152+
reserveTokens: 1_000,
153+
sessionId: "session-1",
154+
sessionFile: "sessions/session-1.json",
155+
now: 123,
156+
});
157+
158+
expect(status).toMatchObject({
159+
schemaVersion: 1,
160+
source: "pre-prompt-estimate",
161+
updatedAt: 123,
162+
provider: "anthropic",
163+
model: "claude-opus-4-6",
164+
route: "fits",
165+
shouldCompact: false,
166+
contextTokenBudget: 10_000,
167+
promptBudgetBeforeReserve: result.promptBudgetBeforeReserve,
168+
reserveTokens: 1_000,
169+
effectiveReserveTokens: result.effectiveReserveTokens,
170+
overflowTokens: 0,
171+
messageCount: 1,
172+
unwindowedMessageCount: 3,
173+
sessionId: "session-1",
174+
sessionFile: "sessions/session-1.json",
175+
});
176+
expect(status.remainingPromptBudgetTokens).toBe(
177+
result.promptBudgetBeforeReserve - result.estimatedPromptTokens,
178+
);
179+
});
180+
134181
it("uses the larger unwindowed message estimate when explicitly provided", () => {
135182
const result = shouldPreemptivelyCompactBeforePrompt({
136183
messages: [makeAssistantHistory("small assembled window")],

‎src/agents/pi-embedded-runner/run/preemptive-compaction.ts‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import type { AgentMessage } from "@earendil-works/pi-agent-core";
22
import { estimateTokens } from "@earendil-works/pi-coding-agent";
3+
import type { SessionContextBudgetStatus } from "../../../config/sessions.js";
34
import { SAFETY_MARGIN, estimateMessagesTokens } from "../../compaction.js";
45
import {
56
MIN_PROMPT_BUDGET_RATIO,
@@ -150,3 +151,46 @@ export function formatPrePromptPrecheckLog(params: {
150151
`sessionFile=${params.sessionFile}`
151152
);
152153
}
154+
155+
export function buildPrePromptContextBudgetStatus(params: {
156+
result: PreemptiveCompactionDecision;
157+
provider: string;
158+
modelId: string;
159+
messageCount: number;
160+
unwindowedMessageCount?: number;
161+
contextTokenBudget: number;
162+
reserveTokens: number;
163+
sessionId?: string;
164+
sessionFile?: string;
165+
now?: number;
166+
}): SessionContextBudgetStatus {
167+
const { result } = params;
168+
const remainingPromptBudgetTokens = Math.max(
169+
0,
170+
result.promptBudgetBeforeReserve - result.estimatedPromptTokens,
171+
);
172+
return {
173+
schemaVersion: 1,
174+
source: "pre-prompt-estimate",
175+
updatedAt: params.now ?? Date.now(),
176+
provider: params.provider,
177+
model: params.modelId,
178+
route: result.route,
179+
shouldCompact: result.shouldCompact,
180+
estimatedPromptTokens: result.estimatedPromptTokens,
181+
contextTokenBudget: Math.max(1, Math.floor(params.contextTokenBudget)),
182+
promptBudgetBeforeReserve: result.promptBudgetBeforeReserve,
183+
reserveTokens: Math.max(0, Math.floor(params.reserveTokens)),
184+
effectiveReserveTokens: result.effectiveReserveTokens,
185+
remainingPromptBudgetTokens,
186+
overflowTokens: result.overflowTokens,
187+
toolResultReducibleChars: result.toolResultReducibleChars,
188+
messageCount: Math.max(0, Math.floor(params.messageCount)),
189+
unwindowedMessageCount: Math.max(
190+
0,
191+
Math.floor(params.unwindowedMessageCount ?? params.messageCount),
192+
),
193+
...(params.sessionId ? { sessionId: params.sessionId } : {}),
194+
...(params.sessionFile ? { sessionFile: params.sessionFile } : {}),
195+
};
196+
}

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,10 @@ import type { Api, AssistantMessage, Model } from "@earendil-works/pi-ai";
33
import type { AuthStorage, ModelRegistry } from "@earendil-works/pi-coding-agent";
44
import type { HeartbeatToolResponse } from "../../../auto-reply/heartbeat-tool-response.js";
55
import type { ThinkLevel } from "../../../auto-reply/thinking.js";
6-
import type { SessionSystemPromptReport } from "../../../config/sessions/types.js";
6+
import type {
7+
SessionContextBudgetStatus,
8+
SessionSystemPromptReport,
9+
} from "../../../config/sessions/types.js";
710
import type { ContextEngine, ContextEnginePromptCacheInfo } from "../../../context-engine/types.js";
811
import type { DiagnosticTraceContext } from "../../../infra/diagnostic-trace-context.js";
912
import type { PluginHookBeforeAgentStartResult } from "../../../plugins/hook-before-agent-start.types.js";
@@ -132,6 +135,7 @@ export type EmbeddedRunAttemptResult = {
132135
cloudCodeAssistFormatError: boolean;
133136
attemptUsage?: NormalizedUsage;
134137
promptCache?: ContextEnginePromptCacheInfo;
138+
contextBudgetStatus?: SessionContextBudgetStatus;
135139
compactionCount?: number;
136140
compactionTokensAfter?: number;
137141
/**

‎src/agents/pi-embedded-runner/types.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
import type { HeartbeatToolResponse } from "../../auto-reply/heartbeat-tool-response.js";
2-
import type { CliSessionBinding, SessionSystemPromptReport } from "../../config/sessions/types.js";
2+
import type {
3+
CliSessionBinding,
4+
SessionContextBudgetStatus,
5+
SessionSystemPromptReport,
6+
} from "../../config/sessions/types.js";
37
import type { DiagnosticTraceContext } from "../../infra/diagnostic-trace-context.js";
48
import type { FallbackAttempt } from "../model-fallback.types.js";
59
import type {
@@ -50,6 +54,7 @@ export type EmbeddedPiAgentMeta = {
5054
cacheWrite?: number;
5155
total?: number;
5256
};
57+
contextBudgetStatus?: SessionContextBudgetStatus;
5358
};
5459

5560
export type TraceAttempt = {

0 commit comments

Comments
 (0)