diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index daeb946f2f4..9f01e6772ce 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -1721,6 +1721,13 @@ export async function runEmbeddedAttempt( const heartbeatPrompt = isDefaultAgent ? resolveHeartbeatPrompt(params.config?.agents?.defaults?.heartbeat?.prompt) : undefined; + const promptWithBootstrapWarning = prependBootstrapPromptWarning( + params.prompt, + bootstrapPromptWarning.lines, + { + preserveExactPrompt: heartbeatPrompt, + }, + ); const appendPrompt = buildEmbeddedSystemPrompt({ workspaceDir: effectiveWorkspace, @@ -2186,7 +2193,7 @@ export async function runEmbeddedAttempt( messages: activeSession.messages, tokenBudget: params.contextTokenBudget, runtimeContext: buildAssembleRuntimeContext({ - prompt: params.prompt, + prompt: promptWithBootstrapWarning, systemPromptText, systemPromptReport, }), @@ -2437,13 +2444,7 @@ export async function runEmbeddedAttempt( // Run before_prompt_build hooks to allow plugins to inject prompt context. // Legacy compatibility: before_agent_start is also checked for context fields. - let effectivePrompt = prependBootstrapPromptWarning( - params.prompt, - bootstrapPromptWarning.lines, - { - preserveExactPrompt: heartbeatPrompt, - }, - ); + let effectivePrompt = promptWithBootstrapWarning; const hookCtx = { agentId: hookAgentId, sessionKey: params.sessionKey, diff --git a/src/context-engine/context-engine.test.ts b/src/context-engine/context-engine.test.ts index 382929c326c..5236e3a4575 100644 --- a/src/context-engine/context-engine.test.ts +++ b/src/context-engine/context-engine.test.ts @@ -229,6 +229,69 @@ class LegacyRuntimeContextStrictEngine implements ContextEngine { } } +class LegacySessionKeyAndRuntimeContextStrictEngine implements ContextEngine { + readonly info: ContextEngineInfo = { + id: "legacy-sessionkey-runtimecontext-strict", + name: "Legacy SessionKey + RuntimeContext Strict Engine", + }; + readonly ingestCalls: Array> = []; + readonly assembleCalls: Array> = []; + + private rejectSessionKey(params: { sessionKey?: string }): void { + if (Object.prototype.hasOwnProperty.call(params, "sessionKey")) { + throw new Error("Unrecognized key(s) in object: 'sessionKey'"); + } + } + + 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 { + this.ingestCalls.push({ ...params }); + this.rejectSessionKey(params); + return { ingested: true }; + } + + async assemble(params: { + sessionId: string; + sessionKey?: string; + messages: AgentMessage[]; + tokenBudget?: number; + runtimeContext?: Record; + }): Promise { + this.assembleCalls.push({ ...params }); + this.rejectSessionKey(params); + this.rejectRuntimeContext(params); + return { + messages: params.messages, + estimatedTokens: 11, + }; + } + + 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", @@ -551,6 +614,39 @@ describe("Legacy sessionKey compatibility", () => { expect(strictEngine.assembleCalls[2]).toHaveProperty("sessionKey", "agent:main:test"); }); + it("still discovers runtimeContext after sessionKey legacy mode was learned earlier", async () => { + const engineId = `legacy-sessionkey-runtimecontext-${Date.now().toString(36)}`; + const strictEngine = new LegacySessionKeyAndRuntimeContextStrictEngine(); + registerContextEngine(engineId, () => strictEngine); + + const engine = await resolveContextEngine(configWithSlot(engineId)); + await engine.ingest({ + sessionId: "s1", + sessionKey: "agent:main:test", + message: makeMockMessage("user", "first"), + }); + + const runtimeContext = { reservedContextTokensEstimate: 321 }; + const assembled = await engine.assemble({ + sessionId: "s1", + sessionKey: "agent:main:test", + messages: [makeMockMessage("assistant", "second")], + runtimeContext, + }); + + expect(assembled.estimatedTokens).toBe(11); + expect(strictEngine.ingestCalls).toHaveLength(2); + expect(strictEngine.ingestCalls[0]).toHaveProperty("sessionKey", "agent:main:test"); + expect(strictEngine.ingestCalls[1]).not.toHaveProperty("sessionKey"); + expect(strictEngine.assembleCalls).toHaveLength(3); + expect(strictEngine.assembleCalls[0]).toHaveProperty("sessionKey", "agent:main:test"); + expect(strictEngine.assembleCalls[0]).toHaveProperty("runtimeContext", runtimeContext); + expect(strictEngine.assembleCalls[1]).not.toHaveProperty("sessionKey"); + expect(strictEngine.assembleCalls[1]).toHaveProperty("runtimeContext", runtimeContext); + expect(strictEngine.assembleCalls[2]).not.toHaveProperty("sessionKey"); + expect(strictEngine.assembleCalls[2]).not.toHaveProperty("runtimeContext"); + }); + 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/registry.ts b/src/context-engine/registry.ts index 9101eb5856c..f861d292da5 100644 --- a/src/context-engine/registry.ts +++ b/src/context-engine/registry.ts @@ -236,10 +236,12 @@ function wrapContextEngineWithSessionKeyCompat(engine: ContextEngine): ContextEn return (params: SessionKeyCompatParams) => { const method = value.bind(target) as (params: SessionKeyCompatParams) => unknown; - const knownLegacyFields = getOwnLegacyCompatFields(params).filter((field) => - legacyFields.has(field), + const compatFieldsInParams = getOwnLegacyCompatFields(params); + const knownLegacyFields = compatFieldsInParams.filter((field) => legacyFields.has(field)); + const hasUntestedCompatFields = compatFieldsInParams.some( + (field) => !legacyFields.has(field), ); - if (isLegacy && knownLegacyFields.length > 0) { + if (isLegacy && knownLegacyFields.length > 0 && !hasUntestedCompatFields) { return method(withoutLegacyCompatFields(params, knownLegacyFields)); } return invokeWithLegacySessionKeyCompat(method, params, {