From c30177a375e31a3901aab130e8928ec4ee5b0f79 Mon Sep 17 00:00:00 2001
From: Codex <codex@openai.invalid>
Date: Sat, 21 Mar 2026 04:54:25 +0800
Subject: [PATCH] Fix assemble runtimeContext CI follow-ups

---
 .../assemble-runtime-context.ts                   | 15 ++++++++++-----
 src/context-engine/context-engine.test.ts         |  8 +++-----
 src/context-engine/registry.ts                    |  6 +++++-
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/src/agents/pi-embedded-runner/assemble-runtime-context.ts b/src/agents/pi-embedded-runner/assemble-runtime-context.ts
index 2c3b3e11b27..2c8bf04df40 100644
--- a/src/agents/pi-embedded-runner/assemble-runtime-context.ts
+++ b/src/agents/pi-embedded-runner/assemble-runtime-context.ts
@@ -1,5 +1,3 @@
-import type { SessionSystemPromptReport } from "../../config/sessions/types.js";
-
 export type EmbeddedAssembleRuntimeContext = {
   systemPromptChars?: number;
   systemPromptTokensEstimate?: number;
@@ -13,9 +11,16 @@ export type EmbeddedAssembleRuntimeContext = {
 };
 
 type AssembleRuntimePromptReport = {
-  systemPrompt?: Pick<SessionSystemPromptReport["systemPrompt"], "chars"> | null;
-  skills?: Pick<SessionSystemPromptReport["skills"], "promptChars"> | null;
-  tools?: Pick<SessionSystemPromptReport["tools"], "listChars" | "schemaChars"> | null;
+  systemPrompt?: {
+    chars?: number;
+  } | null;
+  skills?: {
+    promptChars?: number;
+  } | null;
+  tools?: {
+    listChars?: number;
+    schemaChars?: number;
+  } | null;
 };
 
 function normalizeChars(value: unknown): number {
diff --git a/src/context-engine/context-engine.test.ts b/src/context-engine/context-engine.test.ts
index 5236e3a4575..441c187b454 100644
--- a/src/context-engine/context-engine.test.ts
+++ b/src/context-engine/context-engine.test.ts
@@ -638,13 +638,11 @@ describe("Legacy sessionKey compatibility", () => {
     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).toHaveLength(2);
+    expect(strictEngine.assembleCalls[0]).not.toHaveProperty("sessionKey");
     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");
+    expect(strictEngine.assembleCalls[1]).not.toHaveProperty("runtimeContext");
   });
 
   it("does not retry non-compat runtime errors", async () => {
diff --git a/src/context-engine/registry.ts b/src/context-engine/registry.ts
index f861d292da5..1c77cf14581 100644
--- a/src/context-engine/registry.ts
+++ b/src/context-engine/registry.ts
@@ -244,7 +244,11 @@ function wrapContextEngineWithSessionKeyCompat(engine: ContextEngine): ContextEn
         if (isLegacy && knownLegacyFields.length > 0 && !hasUntestedCompatFields) {
           return method(withoutLegacyCompatFields(params, knownLegacyFields));
         }
-        return invokeWithLegacySessionKeyCompat(method, params, {
+        const retryParams =
+          isLegacy && knownLegacyFields.length > 0
+            ? withoutLegacyCompatFields(params, knownLegacyFields)
+            : params;
+        return invokeWithLegacySessionKeyCompat(method, retryParams, {
           onLegacyModeDetected: (detectedFields) => {
             isLegacy = true;
             for (const field of detectedFields) {