From b8458f0de69677d3b26cd9de4f340e91c3e11e14 Mon Sep 17 00:00:00 2001 From: Codex Date: Sat, 21 Mar 2026 02:18:55 +0800 Subject: [PATCH] Pass runtime context into assemble --- docs/concepts/context-engine.md | 7 +- .../assemble-runtime-context.ts | 51 +++++++ src/agents/pi-embedded-runner/run/attempt.ts | 23 +++ src/context-engine/context-engine.test.ts | 88 ++++++++++++ src/context-engine/legacy.ts | 1 + src/context-engine/registry.ts | 131 ++++++++++++------ src/context-engine/types.ts | 25 +++- 7 files changed, 282 insertions(+), 44 deletions(-) create mode 100644 src/agents/pi-embedded-runner/assemble-runtime-context.ts diff --git a/docs/concepts/context-engine.md b/docs/concepts/context-engine.md index 0b2ec1cd78b..79689a4b84c 100644 --- a/docs/concepts/context-engine.md +++ b/docs/concepts/context-engine.md @@ -128,7 +128,7 @@ export default function register(api) { return { ingested: true }; }, - async assemble({ sessionId, messages, tokenBudget }) { + async assemble({ sessionId, messages, tokenBudget, runtimeContext }) { // Return messages that fit the budget return { messages: buildContext(messages, tokenBudget), @@ -181,6 +181,11 @@ Required members: decisions and diagnostic reporting. - `systemPromptAddition` (optional, `string`) — prepended to the system prompt. +`assemble(params)` also receives an optional `runtimeContext` object. OpenClaw +uses it to pass caller-owned budget signals such as system-prompt size, +tool-schema size, and current prompt size so plugin engines can make more +accurate token-budget decisions without having to guess from inside the plugin. + Optional members: | Member | Kind | Purpose | diff --git a/src/agents/pi-embedded-runner/assemble-runtime-context.ts b/src/agents/pi-embedded-runner/assemble-runtime-context.ts new file mode 100644 index 00000000000..7f7a8973fca --- /dev/null +++ b/src/agents/pi-embedded-runner/assemble-runtime-context.ts @@ -0,0 +1,51 @@ +import type { SessionSystemPromptReport } from "../../config/sessions/types.js"; + +export type EmbeddedAssembleRuntimeContext = { + systemPromptChars?: number; + systemPromptTokensEstimate?: number; + skillsPromptChars?: number; + toolListChars?: number; + toolSchemaChars?: number; + currentPromptChars?: number; + currentPromptTokensEstimate?: number; + reservedContextCharsEstimate?: number; + reservedContextTokensEstimate?: number; +}; + +function normalizeChars(value: unknown): number { + if (!Number.isFinite(value) || Number(value) <= 0) { + return 0; + } + return Math.max(0, Math.floor(Number(value))); +} + +function estimateTokensFromChars(chars: number): number { + if (chars <= 0) { + return 0; + } + return Math.max(1, Math.ceil(chars / 4)); +} + +export function buildEmbeddedAssembleRuntimeContext(params: { + systemPromptText?: string | null; + prompt?: string | null; + systemPromptReport?: Pick | null; +}): EmbeddedAssembleRuntimeContext { + const report = params.systemPromptReport ?? undefined; + const systemPromptChars = normalizeChars( + report?.systemPrompt?.chars ?? params.systemPromptText?.length ?? 0, + ); + const currentPromptChars = normalizeChars(params.prompt?.length ?? 0); + const reservedContextCharsEstimate = Math.max(0, systemPromptChars + currentPromptChars); + return { + systemPromptChars, + systemPromptTokensEstimate: estimateTokensFromChars(systemPromptChars), + skillsPromptChars: normalizeChars(report?.skills?.promptChars ?? 0), + toolListChars: normalizeChars(report?.tools?.listChars ?? 0), + toolSchemaChars: normalizeChars(report?.tools?.schemaChars ?? 0), + currentPromptChars, + currentPromptTokensEstimate: estimateTokensFromChars(currentPromptChars), + reservedContextCharsEstimate, + reservedContextTokensEstimate: estimateTokensFromChars(reservedContextCharsEstimate), + }; +} diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index f89759606de..daeb946f2f4 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -100,6 +100,7 @@ import { normalizeToolName } from "../../tool-policy.js"; import { resolveTranscriptPolicy } from "../../transcript-policy.js"; import { DEFAULT_BOOTSTRAP_FILENAME } from "../../workspace.js"; import { isRunnerAbortError } from "../abort.js"; +import { buildEmbeddedAssembleRuntimeContext } from "../assemble-runtime-context.js"; import { appendCacheTtlTimestamp, isCacheTtlEligibleProvider } from "../cache-ttl.js"; import type { CompactEmbeddedPiSessionParams } from "../compact.js"; import { buildEmbeddedCompactionRuntimeContext } from "../compaction-runtime-context.js"; @@ -1327,6 +1328,23 @@ export function buildAfterTurnRuntimeContext(params: { }); } +/** Build runtime context passed into context-engine assemble hooks. */ +export function buildAssembleRuntimeContext(params: { + prompt?: string | null; + systemPromptText?: string | null; + systemPromptReport?: { + systemPrompt?: { chars?: number }; + skills?: { promptChars?: number }; + tools?: { listChars?: number; schemaChars?: number }; + } | null; +}) { + return buildEmbeddedAssembleRuntimeContext({ + prompt: params.prompt, + systemPromptText: params.systemPromptText, + systemPromptReport: params.systemPromptReport, + }); +} + function summarizeMessagePayload(msg: AgentMessage): { textChars: number; imageBlocks: number } { const content = (msg as { content?: unknown }).content; if (typeof content === "string") { @@ -2167,6 +2185,11 @@ export async function runEmbeddedAttempt( sessionKey: params.sessionKey, messages: activeSession.messages, tokenBudget: params.contextTokenBudget, + runtimeContext: buildAssembleRuntimeContext({ + prompt: params.prompt, + systemPromptText, + systemPromptReport, + }), }); if (assembled.messages !== activeSession.messages) { activeSession.agent.replaceMessages(assembled.messages); diff --git a/src/context-engine/context-engine.test.ts b/src/context-engine/context-engine.test.ts index cf24bfd7a07..382929c326c 100644 --- a/src/context-engine/context-engine.test.ts +++ b/src/context-engine/context-engine.test.ts @@ -76,6 +76,7 @@ class MockContextEngine implements ContextEngine { sessionKey?: string; messages: AgentMessage[]; tokenBudget?: number; + runtimeContext?: Record; }): Promise { return { messages: params.messages, @@ -143,6 +144,7 @@ class LegacySessionKeyStrictEngine implements ContextEngine { sessionKey?: string; messages: AgentMessage[]; tokenBudget?: number; + runtimeContext?: Record; }): Promise { this.assembleCalls.push({ ...params }); this.rejectSessionKey(params); @@ -174,6 +176,59 @@ class LegacySessionKeyStrictEngine implements ContextEngine { } } +class LegacyRuntimeContextStrictEngine implements ContextEngine { + readonly info: ContextEngineInfo = { + id: "legacy-runtimecontext-strict", + name: "Legacy RuntimeContext Strict Engine", + }; + readonly assembleCalls: Array> = []; + + private rejectRuntimeContext(params: { runtimeContext?: Record }): void { + if (Object.prototype.hasOwnProperty.call(params, "runtimeContext")) { + throw new Error("Unrecognized key(s) in object: 'runtimeContext'"); + } + } + + async ingest(_params: { + sessionId: string; + sessionKey?: string; + message: AgentMessage; + isHeartbeat?: boolean; + }): Promise { + return { ingested: true }; + } + + async assemble(params: { + sessionId: string; + sessionKey?: string; + messages: AgentMessage[]; + tokenBudget?: number; + runtimeContext?: Record; + }): Promise { + this.assembleCalls.push({ ...params }); + this.rejectRuntimeContext(params); + return { + messages: params.messages, + estimatedTokens: 9, + }; + } + + async compact(_params: { + sessionId: string; + sessionKey?: string; + sessionFile: string; + tokenBudget?: number; + compactionTarget?: "budget" | "threshold"; + customInstructions?: string; + runtimeContext?: Record; + }): Promise { + return { + ok: true, + compacted: false, + }; + } +} + class SessionKeyRuntimeErrorEngine implements ContextEngine { readonly info: ContextEngineInfo = { id: "sessionkey-runtime-error", @@ -196,6 +251,7 @@ class SessionKeyRuntimeErrorEngine implements ContextEngine { sessionKey?: string; messages: AgentMessage[]; tokenBudget?: number; + runtimeContext?: Record; }): Promise { this.assembleCalls += 1; throw new Error(this.errorMessage); @@ -463,6 +519,38 @@ describe("Legacy sessionKey compatibility", () => { expect(strictEngine.ingestedMessages).toEqual([firstMessage, secondMessage]); }); + it("retries strict assemble once when runtimeContext is rejected and memoizes that field", async () => { + const engineId = `legacy-runtimecontext-${Date.now().toString(36)}`; + const strictEngine = new LegacyRuntimeContextStrictEngine(); + registerContextEngine(engineId, () => strictEngine); + + const engine = await resolveContextEngine(configWithSlot(engineId)); + const runtimeContext = { reservedContextTokensEstimate: 321 }; + + const firstAssembled = await engine.assemble({ + sessionId: "s1", + sessionKey: "agent:main:test", + messages: [makeMockMessage()], + runtimeContext, + }); + const secondAssembled = await engine.assemble({ + sessionId: "s1", + sessionKey: "agent:main:test", + messages: [makeMockMessage("assistant", "second")], + runtimeContext, + }); + + expect(firstAssembled.estimatedTokens).toBe(9); + expect(secondAssembled.estimatedTokens).toBe(9); + expect(strictEngine.assembleCalls).toHaveLength(3); + expect(strictEngine.assembleCalls[0]).toHaveProperty("runtimeContext", runtimeContext); + expect(strictEngine.assembleCalls[0]).toHaveProperty("sessionKey", "agent:main:test"); + expect(strictEngine.assembleCalls[1]).not.toHaveProperty("runtimeContext"); + expect(strictEngine.assembleCalls[1]).toHaveProperty("sessionKey", "agent:main:test"); + expect(strictEngine.assembleCalls[2]).not.toHaveProperty("runtimeContext"); + expect(strictEngine.assembleCalls[2]).toHaveProperty("sessionKey", "agent:main:test"); + }); + it("does not retry non-compat runtime errors", async () => { const engineId = `sessionkey-runtime-${Date.now().toString(36)}`; const runtimeErrorEngine = new SessionKeyRuntimeErrorEngine(); diff --git a/src/context-engine/legacy.ts b/src/context-engine/legacy.ts index 09659c968fb..38fac11f323 100644 --- a/src/context-engine/legacy.ts +++ b/src/context-engine/legacy.ts @@ -40,6 +40,7 @@ export class LegacyContextEngine implements ContextEngine { sessionKey?: string; messages: AgentMessage[]; tokenBudget?: number; + runtimeContext?: ContextEngineRuntimeContext; }): Promise { // Pass-through: the existing sanitize -> validate -> limit -> repair pipeline // in attempt.ts handles context assembly for the legacy engine. diff --git a/src/context-engine/registry.ts b/src/context-engine/registry.ts index 2c5cac439c0..9101eb5856c 100644 --- a/src/context-engine/registry.ts +++ b/src/context-engine/registry.ts @@ -1,6 +1,6 @@ import type { OpenClawConfig } from "../config/config.js"; import { defaultSlotIdForKey } from "../plugins/slots.js"; -import type { ContextEngine } from "./types.js"; +import type { ContextEngine, ContextEngineRuntimeContext } from "./types.js"; /** * A factory that creates a ContextEngine instance. @@ -24,8 +24,11 @@ const SESSION_KEY_COMPAT_METHODS = [ ] as const; type SessionKeyCompatMethodName = (typeof SESSION_KEY_COMPAT_METHODS)[number]; +const LEGACY_COMPAT_FIELDS = ["sessionKey", "runtimeContext"] as const; +type LegacyCompatFieldName = (typeof LEGACY_COMPAT_FIELDS)[number]; type SessionKeyCompatParams = { sessionKey?: string; + runtimeContext?: ContextEngineRuntimeContext; }; function isSessionKeyCompatMethodName(value: PropertyKey): value is SessionKeyCompatMethodName { @@ -34,21 +37,29 @@ function isSessionKeyCompatMethodName(value: PropertyKey): value is SessionKeyCo ); } -function hasOwnSessionKey(params: unknown): params is SessionKeyCompatParams { - return ( - params !== null && - typeof params === "object" && - Object.prototype.hasOwnProperty.call(params, "sessionKey") - ); +function getOwnLegacyCompatFields(params: unknown): LegacyCompatFieldName[] { + if (!params || typeof params !== "object") { + return []; + } + return LEGACY_COMPAT_FIELDS.filter((field) => Object.prototype.hasOwnProperty.call(params, field)); } -function withoutSessionKey(params: T): T { +function hasOwnLegacyCompatFields(params: unknown): params is SessionKeyCompatParams { + return getOwnLegacyCompatFields(params).length > 0; +} + +function withoutLegacyCompatFields( + params: T, + fields: readonly LegacyCompatFieldName[], +): T { const legacyParams = { ...params }; - delete legacyParams.sessionKey; + for (const field of fields) { + delete legacyParams[field]; + } return legacyParams; } -function issueRejectsSessionKeyStrictly(issue: unknown): boolean { +function issueRejectsLegacyFieldStrictly(issue: unknown, field: LegacyCompatFieldName): boolean { if (!issue || typeof issue !== "object") { return false; } @@ -61,12 +72,12 @@ function issueRejectsSessionKeyStrictly(issue: unknown): boolean { if ( issueRecord.code === "unrecognized_keys" && Array.isArray(issueRecord.keys) && - issueRecord.keys.some((key) => key === "sessionKey") + issueRecord.keys.some((key) => key === field) ) { return true; } - return isSessionKeyCompatibilityError(issueRecord.message); + return isLegacyCompatFieldCompatibilityError(issueRecord.message, field); } function* iterateErrorChain(error: unknown) { @@ -82,31 +93,50 @@ function* iterateErrorChain(error: unknown) { } } -const SESSION_KEY_UNKNOWN_FIELD_PATTERNS = [ - /\bunrecognized key(?:\(s\)|s)? in object:.*['"`]sessionKey['"`]/i, - /\badditional propert(?:y|ies)\b.*['"`]sessionKey['"`]/i, - /\bmust not have additional propert(?:y|ies)\b.*['"`]sessionKey['"`]/i, - /\b(?:unexpected|extraneous)\s+(?:property|properties|field|fields|key|keys)\b.*['"`]sessionKey['"`]/i, - /\b(?:unknown|invalid)\s+(?:property|properties|field|fields|key|keys)\b.*['"`]sessionKey['"`]/i, - /['"`]sessionKey['"`].*\b(?:was|is)\s+not allowed\b/i, - /"code"\s*:\s*"unrecognized_keys"[^]*"sessionKey"/i, -] as const; - -function isSessionKeyUnknownFieldValidationMessage(message: string): boolean { - return SESSION_KEY_UNKNOWN_FIELD_PATTERNS.some((pattern) => pattern.test(message)); +function buildUnknownFieldPatterns(field: LegacyCompatFieldName): readonly RegExp[] { + return [ + new RegExp(`\\bunrecognized key(?:\\(s\\)|s)? in object:.*['"\`]${field}['"\`]`, "i"), + new RegExp(`\\badditional propert(?:y|ies)\\b.*['"\`]${field}['"\`]`, "i"), + new RegExp(`\\bmust not have additional propert(?:y|ies)\\b.*['"\`]${field}['"\`]`, "i"), + new RegExp( + `\\b(?:unexpected|extraneous)\\s+(?:property|properties|field|fields|key|keys)\\b.*['"\`]${field}['"\`]`, + "i", + ), + new RegExp( + `\\b(?:unknown|invalid)\\s+(?:property|properties|field|fields|key|keys)\\b.*['"\`]${field}['"\`]`, + "i", + ), + new RegExp(`['"\`]${field}['"\`].*\\b(?:was|is)\\s+not allowed\\b`, "i"), + new RegExp(`"code"\\s*:\\s*"unrecognized_keys"[^]*"${field}"`, "i"), + ] as const; } -function isSessionKeyCompatibilityError(error: unknown): boolean { +const LEGACY_UNKNOWN_FIELD_PATTERNS = { + sessionKey: buildUnknownFieldPatterns("sessionKey"), + runtimeContext: buildUnknownFieldPatterns("runtimeContext"), +} as const satisfies Record; + +function isLegacyCompatFieldUnknownFieldValidationMessage( + message: string, + field: LegacyCompatFieldName, +): boolean { + return LEGACY_UNKNOWN_FIELD_PATTERNS[field].some((pattern) => pattern.test(message)); +} + +function isLegacyCompatFieldCompatibilityError( + error: unknown, + field: LegacyCompatFieldName, +): boolean { for (const candidate of iterateErrorChain(error)) { if (Array.isArray(candidate)) { - if (candidate.some((entry) => issueRejectsSessionKeyStrictly(entry))) { + if (candidate.some((entry) => issueRejectsLegacyFieldStrictly(entry, field))) { return true; } continue; } if (typeof candidate === "string") { - if (isSessionKeyUnknownFieldValidationMessage(candidate)) { + if (isLegacyCompatFieldUnknownFieldValidationMessage(candidate, field)) { return true; } continue; @@ -124,21 +154,21 @@ function isSessionKeyCompatibilityError(error: unknown): boolean { if ( Array.isArray(issueContainer.issues) && - issueContainer.issues.some((issue) => issueRejectsSessionKeyStrictly(issue)) + issueContainer.issues.some((issue) => issueRejectsLegacyFieldStrictly(issue, field)) ) { return true; } if ( Array.isArray(issueContainer.errors) && - issueContainer.errors.some((issue) => issueRejectsSessionKeyStrictly(issue)) + issueContainer.errors.some((issue) => issueRejectsLegacyFieldStrictly(issue, field)) ) { return true; } if ( typeof issueContainer.message === "string" && - isSessionKeyUnknownFieldValidationMessage(issueContainer.message) + isLegacyCompatFieldUnknownFieldValidationMessage(issueContainer.message, field) ) { return true; } @@ -147,25 +177,35 @@ function isSessionKeyCompatibilityError(error: unknown): boolean { return false; } +function collectLegacyCompatRejectedFields(error: unknown): LegacyCompatFieldName[] { + return LEGACY_COMPAT_FIELDS.filter((field) => isLegacyCompatFieldCompatibilityError(error, field)); +} + async function invokeWithLegacySessionKeyCompat( method: (params: TParams) => Promise | TResult, params: TParams, opts?: { - onLegacyModeDetected?: () => void; + onLegacyModeDetected?: (fields: LegacyCompatFieldName[]) => void; }, ): Promise { - if (!hasOwnSessionKey(params)) { + if (!hasOwnLegacyCompatFields(params)) { return await method(params); } - try { - return await method(params); - } catch (error) { - if (!isSessionKeyCompatibilityError(error)) { - throw error; + let currentParams = params; + while (true) { + try { + return await method(currentParams); + } catch (error) { + const presentRejectedFields = collectLegacyCompatRejectedFields(error).filter((field) => + getOwnLegacyCompatFields(currentParams).includes(field), + ); + if (presentRejectedFields.length === 0) { + throw error; + } + opts?.onLegacyModeDetected?.(presentRejectedFields); + currentParams = withoutLegacyCompatFields(currentParams, presentRejectedFields); } - opts?.onLegacyModeDetected?.(); - return await method(withoutSessionKey(params)); } } @@ -178,6 +218,7 @@ function wrapContextEngineWithSessionKeyCompat(engine: ContextEngine): ContextEn } let isLegacy = false; + const legacyFields = new Set(); const proxy: ContextEngine = new Proxy(engine, { get(target, property, receiver) { if (property === LEGACY_SESSION_KEY_COMPAT) { @@ -195,12 +236,18 @@ function wrapContextEngineWithSessionKeyCompat(engine: ContextEngine): ContextEn return (params: SessionKeyCompatParams) => { const method = value.bind(target) as (params: SessionKeyCompatParams) => unknown; - if (isLegacy && hasOwnSessionKey(params)) { - return method(withoutSessionKey(params)); + const knownLegacyFields = getOwnLegacyCompatFields(params).filter((field) => + legacyFields.has(field), + ); + if (isLegacy && knownLegacyFields.length > 0) { + return method(withoutLegacyCompatFields(params, knownLegacyFields)); } return invokeWithLegacySessionKeyCompat(method, params, { - onLegacyModeDetected: () => { + onLegacyModeDetected: (detectedFields) => { isLegacy = true; + for (const field of detectedFields) { + legacyFields.add(field); + } }, }); }; diff --git a/src/context-engine/types.ts b/src/context-engine/types.ts index 7ddd695b5b6..f6ca84e0a3f 100644 --- a/src/context-engine/types.ts +++ b/src/context-engine/types.ts @@ -57,7 +57,28 @@ export type SubagentSpawnPreparation = { }; export type SubagentEndReason = "deleted" | "completed" | "swept" | "released"; -export type ContextEngineRuntimeContext = Record; +export type ContextEngineRuntimeContext = Record & { + /** + * Approximate size of the host-owned system prompt before context-engine + * additions are prepended. + */ + systemPromptChars?: number; + systemPromptTokensEstimate?: number; + /** Skill blocks are already included in systemPromptChars; these are breakdowns. */ + skillsPromptChars?: number; + /** Tool list/schema chars are already included in systemPromptChars; these are breakdowns. */ + toolListChars?: number; + toolSchemaChars?: number; + /** Approximate size of the current user prompt for the pending run. */ + currentPromptChars?: number; + currentPromptTokensEstimate?: number; + /** + * Host-owned context reserved outside the engine-controlled history slice. + * This usually includes the base system prompt and current user prompt. + */ + reservedContextCharsEstimate?: number; + reservedContextTokensEstimate?: number; +}; /** * ContextEngine defines the pluggable contract for context management. @@ -131,6 +152,8 @@ export interface ContextEngine { sessionKey?: string; messages: AgentMessage[]; tokenBudget?: number; + /** Optional runtime-owned context for engines that need caller state. */ + runtimeContext?: ContextEngineRuntimeContext; }): Promise; /**