From 02291dc153eedb0a7efd458ab5bb0684e9ba4e81 Mon Sep 17 00:00:00 2001 From: Username Date: Sun, 28 Jun 2026 20:11:51 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20add=20tokenBudgetMax=20=E2=80=94=20abso?= =?UTF-8?q?lute,=20window-independent=20injection=20ceiling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Memory injection is sized as tokenBudgetFraction × the model's context window, so a large-window model (e.g. 1M) balloons per-turn injection proportionally — on a 1M-window model a fresh session showed ~170-200k of context coming purely from memory injection, rewritten into the cache every turn. tokenBudgetMax caps injection in absolute tokens, independent of the window: - Stage 1 (base cap): the budget handed to the daemon's assemble is capped to min(window, tokenBudgetMax / tokenBudgetFraction) so the daemon pre-trims its sub-channels. - Stage 2 (enforcer): after all five injection paths land (main assemble, continuity, exact-recall, predictive_context, beforeTurn), the combined systemPromptAddition is truncated to tokenBudgetMax. This is the real ceiling because several paths size against the real window, not the capped budget. Only injection is bounded — the cap is never applied to enforceTokenBudgetInvariant or compaction, so the usable conversation window is untouched. Unset preserves current behavior (no-op). Adds the config type, JSON schema entry, and unit tests for both stages plus the uncapped pass-through. Co-Authored-By: Claude Opus 4.8 --- openclaw.plugin.json | 5 ++ src/context-engine.ts | 63 +++++++++++++++++++++++- src/types.ts | 9 ++++ test/unit/context-engine.test.ts | 83 ++++++++++++++++++++++++++++++++ 4 files changed, 158 insertions(+), 2 deletions(-) diff --git a/openclaw.plugin.json b/openclaw.plugin.json index 0fabd561..c6c45c99 100644 --- a/openclaw.plugin.json +++ b/openclaw.plugin.json @@ -467,6 +467,11 @@ "tokenBudgetFraction": { "type": "number" }, + "tokenBudgetMax": { + "type": "number", + "minimum": 0, + "description": "Absolute ceiling (in tokens) on memory injection per turn, independent of the model's context window. Without it, injection is sized as tokenBudgetFraction × window, so large-window models balloon injection. The daemon budget is pre-capped to min(window, tokenBudgetMax / tokenBudgetFraction), then the combined system prompt addition is truncated to tokenBudgetMax after all injection paths land. Only injection is bounded; the usable conversation window is untouched. Unset disables the cap." + }, "compactThreshold": { "type": "number" }, diff --git a/src/context-engine.ts b/src/context-engine.ts index c216fd77..240057df 100644 --- a/src/context-engine.ts +++ b/src/context-engine.ts @@ -705,6 +705,33 @@ function resolveEffectiveAssembleBudget(tokenBudget: number | undefined): number return Math.max(1, normalized - headroom); } +// tokenBudgetMax — stage 1 (base cap). Daemon-side injection is sized as +// tokenBudgetFraction × tokenBudget, so on a large-window model (e.g. 1M) +// injection balloons proportionally. This caps the budget handed to the daemon +// to min(window, tokenBudgetMax / fraction) so the daemon pre-trims its +// sub-channels. Returns tokenBudget unchanged when tokenBudgetMax is unset. +// NOTE: this only makes the daemon pre-trim — the post-assembly trim +// (truncateSystemPromptAdditionToTokenBudget) is the actual ceiling enforcer, +// because several injection paths size against the real window, not this budget. +function resolveCappedAssembleBudget( + tokenBudget: number | undefined, + tokenBudgetMax: number | undefined, + tokenBudgetFraction: number | undefined, +): number | undefined { + const normalized = normalizeTokenBudget(tokenBudget); + if (normalized == null) return tokenBudget; + if (typeof tokenBudgetMax !== "number" || !Number.isFinite(tokenBudgetMax) || tokenBudgetMax <= 0) { + return tokenBudget; + } + // The daemon's true default injection fraction is unknown; 0.2 is a safe + // fallback. It only affects how aggressively the daemon pre-trims. + const fraction = typeof tokenBudgetFraction === "number" && tokenBudgetFraction > 0 + ? tokenBudgetFraction + : 0.2; + const cap = Math.max(1, Math.floor(tokenBudgetMax / fraction)); + return Math.min(normalized, cap); +} + function normalizeThresholdFraction(fraction: number | undefined): number { if (typeof fraction !== "number" || !Number.isFinite(fraction)) { return DEFAULT_COMPACTION_THRESHOLD_FRACTION; @@ -2563,6 +2590,13 @@ export function buildContextEngineFactory( } const assembleTimeout = cfg.assembleTimeoutMs ?? 30000; + // tokenBudgetMax stage 1: cap the budget the daemon sizes injection + // against (no-op when tokenBudgetMax is unset). + const cappedAssembleBudget = resolveCappedAssembleBudget( + args.tokenBudget, + cfg.tokenBudgetMax, + cfg.tokenBudgetFraction, + ) ?? args.tokenBudget; const resp = await Promise.race([ client.assembleContextInternal({ sessionId, @@ -2570,8 +2604,8 @@ export function buildContextEngineFactory( userId, prompt: retrievalQuery, messages: messages as any, - tokenBudget: args.tokenBudget, - config: buildAssemblyConfig(args.tokenBudget), + tokenBudget: cappedAssembleBudget, + config: buildAssemblyConfig(cappedAssembleBudget), emitDebug: true, }), new Promise((_, reject) => @@ -2688,6 +2722,31 @@ export function buildContextEngineFactory( }); } + // tokenBudgetMax stage 2 (enforcer): truncate the combined injection to + // the configured ceiling after all paths (main assemble, continuity, + // exact-recall, predictive_context, beforeTurn) have landed. Applied to + // systemPromptAddition only — never to the conversation, which + // enforceTokenBudgetInvariant governs against the real window. + if (typeof cfg.tokenBudgetMax === "number" && Number.isFinite(cfg.tokenBudgetMax) && cfg.tokenBudgetMax > 0) { + const injectionBefore = approximateTokenCount(enforced.systemPromptAddition); + if (injectionBefore > cfg.tokenBudgetMax) { + const trimmed = truncateSystemPromptAdditionToTokenBudget( + enforced.systemPromptAddition, + cfg.tokenBudgetMax, + ); + const injectionAfter = approximateTokenCount(trimmed); + enforced = { + ...enforced, + systemPromptAddition: trimmed, + estimatedTokens: Math.max(0, enforced.estimatedTokens - (injectionBefore - injectionAfter)), + }; + logger.info?.( + `LibraVDB tokenBudgetMax trim sessionId=${sessionId} ` + + `injectionBefore=${injectionBefore} injectionAfter=${injectionAfter} ` + + `cap=${cfg.tokenBudgetMax}`, + ); + } + } enforced = enforceTokenBudgetInvariant( enforced, args.tokenBudget, diff --git a/src/types.ts b/src/types.ts index bebfeed9..ab45d0a4 100644 --- a/src/types.ts +++ b/src/types.ts @@ -85,6 +85,15 @@ export interface PluginConfig { recencyLambdaUser?: number; recencyLambdaGlobal?: number; tokenBudgetFraction?: number; + /** Absolute ceiling (in tokens) on memory injection per turn, independent of + * the model's context window. Without it, injection is sized as + * tokenBudgetFraction × window, so a large-window model balloons injection. + * Two stages: the daemon budget is pre-capped to min(window, + * tokenBudgetMax / tokenBudgetFraction), then the combined systemPromptAddition + * is truncated to tokenBudgetMax after all injection paths land. Only the + * injection is bounded — the usable conversation window is untouched. + * Unset disables the cap. */ + tokenBudgetMax?: number; authoredHardBudgetFraction?: number; authoredSoftBudgetFraction?: number; elevatedGuidanceBudgetFraction?: number; diff --git a/test/unit/context-engine.test.ts b/test/unit/context-engine.test.ts index 84b86f22..ee6b78cc 100644 --- a/test/unit/context-engine.test.ts +++ b/test/unit/context-engine.test.ts @@ -2514,3 +2514,86 @@ test("context engine assemble drain handles empty queue gracefully", async () => // consecutive cursor positions, exercising the hasAllToolIdsSeen / recordToolIds // path directly. // --------------------------------------------------------------------------- + +// --------------------------------------------------------------------------- +// tokenBudgetMax — absolute injection ceiling (window-independent) +// --------------------------------------------------------------------------- + +test("tokenBudgetMax caps the daemon budget and truncates the injected system prompt", async () => { + const client = new FakeClient(); + // ~9000 tokens of plain injection (sanitization is a no-op for plain text). + const bigInjection = "alpha ".repeat(6000); + client.assembleResponse = { + messages: [], + estimatedTokens: 9000, + systemPromptAddition: bigInjection, + }; + const engine = buildContextEngineFactory(fakeRuntime(client), { + userId: "fixed-user", + tokenBudgetMax: 1000, + tokenBudgetFraction: 0.2, + crossSessionRecall: false, // skip exact recall + beforeTurnEnabled: false, // skip beforeTurn injection + }); + + const messages = [ + { role: "user", content: "earlier", id: "u0" }, + { role: "assistant", content: "ok", id: "a0" }, + { role: "user", content: "what is the status", id: "u1" }, + ]; + const result = await engine.assemble({ + sessionId: "s-cap", + sessionKey: "agent:main:session:s-cap", + messages, + tokenBudget: 1_000_000, // 1M window + prompt: "what is the status", + }); + + // Stage 1: the daemon received min(1M, 1000 / 0.2) = 5000, not the full window. + const assembleCall = client.calls.find((c) => c.method === "assembleContextInternal"); + assert.ok(assembleCall, "assembleContextInternal should be called"); + assert.equal(assembleCall.params.tokenBudget, 5000); + + // Stage 2: combined injection truncated to tokenBudgetMax (1000 tokens = + // 4000 chars at APPROX_CHARS_PER_TOKEN=4). + assert.ok( + result.systemPromptAddition.length <= 4000, + `injection ${result.systemPromptAddition.length} chars should be <= 4000`, + ); + assert.ok(result.systemPromptAddition.length > 0, "some injection survives the trim"); + assert.ok( + result.systemPromptAddition.length < bigInjection.length, + "injection was actually truncated", + ); +}); + +test("tokenBudgetMax unset: daemon budget passes through uncapped", async () => { + const client = new FakeClient(); + client.assembleResponse = { + messages: [], + estimatedTokens: 100, + systemPromptAddition: "small note", + }; + const engine = buildContextEngineFactory(fakeRuntime(client), { + userId: "fixed-user", + crossSessionRecall: false, + beforeTurnEnabled: false, + }); + + const messages = [ + { role: "user", content: "earlier", id: "u0" }, + { role: "assistant", content: "ok", id: "a0" }, + { role: "user", content: "hi", id: "u1" }, + ]; + await engine.assemble({ + sessionId: "s-uncapped", + sessionKey: "agent:main:session:s-uncapped", + messages, + tokenBudget: 1_000_000, + prompt: "hi", + }); + + const assembleCall = client.calls.find((c) => c.method === "assembleContextInternal"); + assert.ok(assembleCall, "assembleContextInternal should be called"); + assert.equal(assembleCall.params.tokenBudget, 1_000_000, "full window passed when uncapped"); +});