Skip to content

Commit 7f8c30d

Browse files
author
Gio Della-Libera
committed
feat(agents): classify context budget pressure
1 parent eb4b8af commit 7f8c30d

8 files changed

Lines changed: 223 additions & 3 deletions

File tree

‎src/auto-reply/status.test.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ describe("buildStatusMessage", () => {
174174
});
175175
const normalized = normalizeTestText(text);
176176

177-
expect(normalized).toContain("Context: ~640k/1.0m (64% est)");
177+
expect(normalized).toContain("Context: ~640k/1.0m (64% est) · Budget: watch");
178178
expect(normalized).not.toContain("Context: ?/1.0m");
179179
expect(normalized).not.toContain("Context: 3.8m/1.0m");
180180
});
@@ -220,6 +220,7 @@ describe("buildStatusMessage", () => {
220220
const normalized = normalizeTestText(text);
221221

222222
expect(normalized).toContain("Context: 36k/1.0m (4%)");
223+
expect(normalized).toContain("Budget: watch");
223224
expect(normalized).not.toContain("~640k");
224225
});
225226

@@ -261,7 +262,7 @@ describe("buildStatusMessage", () => {
261262
});
262263
const normalized = normalizeTestText(text);
263264

264-
expect(normalized).toContain("Context: ~125k/1.0m (13% est)");
265+
expect(normalized).toContain("Context: ~125k/1.0m (13% est) · Budget: safe");
265266
expect(normalized).not.toContain("Context: 0/1.0m");
266267
});
267268

‎src/config/sessions.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ export * from "./sessions/paths.js";
99
export * from "./sessions/reset.js";
1010
export * from "./sessions/session-key.js";
1111
export * from "./sessions/store.js";
12+
export * from "./sessions/context-budget-policy.js";
1213
export * from "./sessions/types.js";
1314
export * from "./sessions/transcript.js";
1415
export * from "./sessions/session-file.js";

‎src/config/sessions/context-budget-policy.test.ts‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
import { describe, expect, it } from "vitest";
2+
import { resolveSessionContextBudgetPolicy } from "./context-budget-policy.js";
3+
import type { SessionContextBudgetStatus } from "./types.js";
4+
5+
function makeStatus(patch: Partial<SessionContextBudgetStatus> = {}): SessionContextBudgetStatus {
6+
return {
7+
schemaVersion: 1,
8+
source: "pre-prompt-estimate",
9+
updatedAt: 1,
10+
provider: "anthropic",
11+
model: "claude-opus-4-6",
12+
route: "fits",
13+
shouldCompact: false,
14+
estimatedPromptTokens: 100,
15+
contextTokenBudget: 1_000,
16+
promptBudgetBeforeReserve: 900,
17+
reserveTokens: 100,
18+
effectiveReserveTokens: 100,
19+
remainingPromptBudgetTokens: 800,
20+
overflowTokens: 0,
21+
toolResultReducibleChars: 0,
22+
messageCount: 1,
23+
unwindowedMessageCount: 1,
24+
...patch,
25+
};
26+
}
27+
28+
describe("resolveSessionContextBudgetPolicy", () => {
29+
it("classifies low estimated prompt usage as safe", () => {
30+
expect(
31+
resolveSessionContextBudgetPolicy(
32+
makeStatus({
33+
estimatedPromptTokens: 125_000,
34+
contextTokenBudget: 1_000_000,
35+
promptBudgetBeforeReserve: 900_000,
36+
remainingPromptBudgetTokens: 775_000,
37+
}),
38+
),
39+
).toMatchObject({
40+
pressure: "safe",
41+
contextBudgetPct: 13,
42+
promptBudgetPct: 14,
43+
remainingPromptBudgetTokens: 775_000,
44+
});
45+
});
46+
47+
it("classifies reserve-budget pressure before overflow", () => {
48+
expect(
49+
resolveSessionContextBudgetPolicy(
50+
makeStatus({
51+
estimatedPromptTokens: 640_000,
52+
contextTokenBudget: 1_000_000,
53+
promptBudgetBeforeReserve: 900_000,
54+
remainingPromptBudgetTokens: 260_000,
55+
}),
56+
),
57+
).toMatchObject({
58+
pressure: "watch",
59+
contextBudgetPct: 64,
60+
promptBudgetPct: 71,
61+
});
62+
63+
expect(
64+
resolveSessionContextBudgetPolicy(
65+
makeStatus({
66+
estimatedPromptTokens: 780_000,
67+
contextTokenBudget: 1_000_000,
68+
promptBudgetBeforeReserve: 900_000,
69+
remainingPromptBudgetTokens: 120_000,
70+
}),
71+
)?.pressure,
72+
).toBe("pressure");
73+
});
74+
75+
it("classifies non-fitting precheck routes as overflow risk", () => {
76+
expect(
77+
resolveSessionContextBudgetPolicy(
78+
makeStatus({
79+
route: "compact_then_truncate",
80+
shouldCompact: true,
81+
estimatedPromptTokens: 920_000,
82+
contextTokenBudget: 1_000_000,
83+
promptBudgetBeforeReserve: 900_000,
84+
remainingPromptBudgetTokens: 0,
85+
overflowTokens: 20_000,
86+
}),
87+
)?.pressure,
88+
).toBe("overflow-risk");
89+
});
90+
});

‎src/config/sessions/context-budget-policy.ts‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
import type { SessionContextBudgetStatus } from "./types.js";
2+
3+
export type SessionContextBudgetPressure = "safe" | "watch" | "pressure" | "overflow-risk";
4+
5+
export type SessionContextBudgetPolicy = {
6+
pressure: SessionContextBudgetPressure;
7+
estimatedPromptTokens: number;
8+
contextBudgetPct?: number;
9+
promptBudgetPct?: number;
10+
remainingPromptBudgetTokens: number;
11+
overflowTokens: number;
12+
route: SessionContextBudgetStatus["route"];
13+
};
14+
15+
const WATCH_PROMPT_BUDGET_PCT = 65;
16+
const PRESSURE_PROMPT_BUDGET_PCT = 85;
17+
18+
function resolveNonNegativeInteger(value: number | undefined): number | undefined {
19+
if (typeof value !== "number" || !Number.isFinite(value) || value < 0) {
20+
return undefined;
21+
}
22+
return Math.floor(value);
23+
}
24+
25+
function resolvePositiveInteger(value: number | undefined): number | undefined {
26+
if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) {
27+
return undefined;
28+
}
29+
return Math.floor(value);
30+
}
31+
32+
function pct(numerator: number, denominator: number | undefined): number | undefined {
33+
if (denominator === undefined) {
34+
return undefined;
35+
}
36+
return Math.min(999, Math.max(0, Math.round((numerator / denominator) * 100)));
37+
}
38+
39+
export function resolveSessionContextBudgetPolicy(
40+
status: SessionContextBudgetStatus | undefined,
41+
): SessionContextBudgetPolicy | undefined {
42+
if (!status || status.source !== "pre-prompt-estimate") {
43+
return undefined;
44+
}
45+
const estimatedPromptTokens = resolveNonNegativeInteger(status.estimatedPromptTokens);
46+
if (estimatedPromptTokens === undefined) {
47+
return undefined;
48+
}
49+
const contextTokenBudget = resolvePositiveInteger(status.contextTokenBudget);
50+
const promptBudgetBeforeReserve = resolvePositiveInteger(status.promptBudgetBeforeReserve);
51+
const overflowTokens = resolveNonNegativeInteger(status.overflowTokens) ?? 0;
52+
const remainingPromptBudgetTokens =
53+
resolveNonNegativeInteger(status.remainingPromptBudgetTokens) ??
54+
Math.max(0, (promptBudgetBeforeReserve ?? 0) - estimatedPromptTokens);
55+
const promptBudgetPct = pct(estimatedPromptTokens, promptBudgetBeforeReserve);
56+
const contextBudgetPct = pct(estimatedPromptTokens, contextTokenBudget);
57+
const pressure: SessionContextBudgetPressure =
58+
overflowTokens > 0 || status.route !== "fits"
59+
? "overflow-risk"
60+
: promptBudgetPct !== undefined && promptBudgetPct >= PRESSURE_PROMPT_BUDGET_PCT
61+
? "pressure"
62+
: promptBudgetPct !== undefined && promptBudgetPct >= WATCH_PROMPT_BUDGET_PCT
63+
? "watch"
64+
: "safe";
65+
return {
66+
pressure,
67+
estimatedPromptTokens,
68+
...(contextBudgetPct !== undefined ? { contextBudgetPct } : {}),
69+
...(promptBudgetPct !== undefined ? { promptBudgetPct } : {}),
70+
remainingPromptBudgetTokens,
71+
overflowTokens,
72+
route: status.route,
73+
};
74+
}

‎src/gateway/session-utils.test.ts‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,41 @@ describe("gateway session utils", () => {
335335
});
336336
});
337337

338+
test("session rows expose derived context budget pressure", () => {
339+
const row = buildGatewaySessionRow({
340+
cfg: createModelDefaultsConfig({ primary: "anthropic/claude-opus-4-6" }),
341+
storePath: "",
342+
store: {},
343+
key: "agent:main:main",
344+
entry: {
345+
sessionId: "session-1",
346+
updatedAt: 1,
347+
contextBudgetStatus: {
348+
schemaVersion: 1,
349+
source: "pre-prompt-estimate",
350+
updatedAt: 1,
351+
provider: "anthropic",
352+
model: "claude-opus-4-6",
353+
route: "fits",
354+
shouldCompact: false,
355+
estimatedPromptTokens: 640_000,
356+
contextTokenBudget: 1_000_000,
357+
promptBudgetBeforeReserve: 900_000,
358+
reserveTokens: 100_000,
359+
effectiveReserveTokens: 100_000,
360+
remainingPromptBudgetTokens: 260_000,
361+
overflowTokens: 0,
362+
toolResultReducibleChars: 0,
363+
messageCount: 2,
364+
unwindowedMessageCount: 2,
365+
},
366+
},
367+
});
368+
369+
expect(row.contextBudgetPressure).toBe("watch");
370+
expect(row.contextBudgetStatus?.estimatedPromptTokens).toBe(640_000);
371+
});
372+
338373
test("async session list reuses thinking metadata for lightweight rows", async () => {
339374
const resolveThinkingProfile = vi.fn(() => ({
340375
levels: [{ id: "off" as const }, { id: "medium" as const }],

‎src/gateway/session-utils.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ import {
5050
resolveAllAgentSessionStoreTargetsSync,
5151
resolveAgentMainSessionKey,
5252
resolveFreshSessionTotalTokens,
53+
resolveSessionContextBudgetPolicy,
5354
resolveStorePath,
5455
type SessionEntry,
5556
type SessionStoreTarget,
@@ -1861,6 +1862,7 @@ export function buildGatewaySessionRow(params: {
18611862
allowAsyncLoad: false,
18621863
}),
18631864
));
1865+
const contextBudgetPolicy = resolveSessionContextBudgetPolicy(entry?.contextBudgetStatus);
18641866

18651867
let derivedTitle: string | undefined;
18661868
let lastMessagePreview: string | undefined;
@@ -1946,6 +1948,7 @@ export function buildGatewaySessionRow(params: {
19461948
agentRuntime,
19471949
contextTokens,
19481950
contextBudgetStatus: entry?.contextBudgetStatus,
1951+
contextBudgetPressure: contextBudgetPolicy?.pressure,
19491952
deliveryContext: deliveryFields.deliveryContext,
19501953
lastChannel: deliveryFields.lastChannel ?? entry?.lastChannel,
19511954
lastTo: deliveryFields.lastTo ?? entry?.lastTo,

‎src/gateway/session-utils.types.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
import type { ChatType } from "../channels/chat-type.js";
2-
import type { SessionCompactionCheckpoint, SessionEntry } from "../config/sessions/types.js";
2+
import type {
3+
SessionCompactionCheckpoint,
4+
SessionContextBudgetPressure,
5+
SessionEntry,
6+
} from "../config/sessions.js";
37
import type { PluginSessionExtensionProjection } from "../plugins/host-hooks.js";
48
import type {
59
GatewayAgentRuntime,
@@ -85,6 +89,7 @@ export type GatewaySessionRow = {
8589
agentRuntime?: GatewayAgentRuntime;
8690
contextTokens?: number;
8791
contextBudgetStatus?: SessionEntry["contextBudgetStatus"];
92+
contextBudgetPressure?: SessionContextBudgetPressure;
8893
deliveryContext?: DeliveryContext;
8994
lastChannel?: SessionEntry["lastChannel"];
9095
lastTo?: string;

‎src/status/status-message.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,9 @@ import {
2929
resolveSessionPluginStatusLines,
3030
resolveSessionPluginTraceLines,
3131
resolveFreshSessionTotalTokens,
32+
resolveSessionContextBudgetPolicy,
3233
type SessionEntry,
34+
type SessionContextBudgetPressure,
3335
type SessionScope,
3436
} from "../config/sessions.js";
3537
import { hasSessionAutoModelFallbackProvenance } from "../config/sessions/model-override-provenance.js";
@@ -243,6 +245,13 @@ const formatEstimatedContextBudgetTokens = (
243245
return `~${totalLabel}/${ctxLabel}${pct !== null ? ` (${pct}% est)` : " (est)"}`;
244246
};
245247

248+
const formatContextBudgetPressure = (pressure: SessionContextBudgetPressure | undefined) => {
249+
if (!pressure) {
250+
return null;
251+
}
252+
return `Budget: ${pressure}`;
253+
};
254+
246255
export const formatContextUsageShort = (
247256
total: number | null | undefined,
248257
contextTokens: number | null | undefined,
@@ -857,8 +866,10 @@ export function buildStatusMessage(args: StatusArgs): string {
857866
? (formatEstimatedContextBudgetTokens(entry?.contextBudgetStatus, contextTokens) ??
858867
formatTokens(totalTokens, contextTokens ?? null))
859868
: formatTokens(totalTokens, contextTokens ?? null);
869+
const contextBudgetPolicy = resolveSessionContextBudgetPolicy(entry?.contextBudgetStatus);
860870
const contextLine = [
861871
`Context: ${contextUsageLabel}`,
872+
formatContextBudgetPressure(contextBudgetPolicy?.pressure),
862873
`🧹 Compactions: ${entry?.compactionCount ?? 0}`,
863874
]
864875
.filter(Boolean)

0 commit comments

Comments
 (0)