Fix runtimeContext compatibility follow-ups

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,

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<Record<string, unknown>> = [];
+  readonly assembleCalls: Array<Record<string, unknown>> = [];
+
+  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<string, unknown> }): 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<IngestResult> {
+    this.ingestCalls.push({ ...params });
+    this.rejectSessionKey(params);
+    return { ingested: true };
+  }
+
+  async assemble(params: {
+    sessionId: string;
+    sessionKey?: string;
+    messages: AgentMessage[];
+    tokenBudget?: number;
+    runtimeContext?: Record<string, unknown>;
+  }): Promise<AssembleResult> {
+    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<string, unknown>;
+  }): Promise<CompactResult> {
+    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();

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, {