pi-openai-codex-compat 0.0.3 → 0.0.4

package/extensions/openai-codex-compat/codex-stream.ts

@@ -41,6 +41,12 @@ type ToolCallSlot = Extract<OutputSlot, { type: "toolCall" }>;
 
 type ProcessCodexStreamOptions = {
   applyServiceTierPricing?(usage: Usage, responseServiceTier: string | undefined): void;
+  attemptState?: CodexStreamAttemptState;
+};
+
+export type CodexStreamAttemptState = {
+  startedContentIndexes: Set<number>;
+  completedContentIndexes: Set<number>;
 };
 
 type CodexResponseStatus =
@@ -68,6 +74,18 @@ function stringValue(value: unknown): string {
   return typeof value === "string" ? value : "";
 }
 
+function piToolCallName(item: JsonRecord): string {
+  const wireName = stringValue(item.name);
+  const name =
+    item["namespace"] === undefined
+      ? wireName
+      : namespacedToolCallName(item["namespace"], wireName);
+  if (item["namespace"] === undefined && CODEX_NAMESPACED_TOOL_NAMES.has(name)) {
+    throw new Error(`Codex returned namespaced tool "${name}" as a flat function call.`);
+  }
+  return name;
+}
+
 function encodeTextSignature(id: string, phase: unknown): string {
   const payload: TextSignatureV1 = { v: 1, id };
   if (phase === "commentary" || phase === "final_answer") payload.phase = phase;
@@ -216,15 +234,24 @@ export async function processCodexStream(
     });
   };
 
+  const trackStarted = <TSlot extends OutputSlot>(slot: TSlot): TSlot => {
+    options?.attemptState?.startedContentIndexes.add(slot.contentIndex);
+    return slot;
+  };
+
+  const trackCompleted = (slot: OutputSlot): void => {
+    options?.attemptState?.completedContentIndexes.add(slot.contentIndex);
+  };
+
   const createSlot = (index: number, item: JsonRecord): OutputSlot | undefined => {
     if (item.type === "reasoning") {
       const block: ThinkingContent = { type: "thinking", thinking: "" };
       output.content.push(block);
-      const slot = {
+      const slot = trackStarted({
         type: "thinking",
         block,
         contentIndex: output.content.length - 1,
-      } satisfies OutputSlot;
+      } satisfies OutputSlot);
       slots.set(index, slot);
       stream.push({ type: "thinking_start", contentIndex: slot.contentIndex, partial: output });
       return slot;
@@ -233,24 +260,17 @@ export async function processCodexStream(
       applyMessagePhaseStopReason(item);
       const block: TextContent = { type: "text", text: "" };
       output.content.push(block);
-      const slot = {
+      const slot = trackStarted({
         type: "text",
         block,
         contentIndex: output.content.length - 1,
-      } satisfies OutputSlot;
+      } satisfies OutputSlot);
       slots.set(index, slot);
       stream.push({ type: "text_start", contentIndex: slot.contentIndex, partial: output });
       return slot;
     }
     if (item.type === "function_call") {
-      const wireName = stringValue(item.name);
-      const name =
-        item["namespace"] === undefined
-          ? wireName
-          : namespacedToolCallName(item["namespace"], wireName);
-      if (item["namespace"] === undefined && CODEX_NAMESPACED_TOOL_NAMES.has(name)) {
-        throw new Error(`Codex returned namespaced tool "${name}" as a flat function call.`);
-      }
+      const name = piToolCallName(item);
       const block: StreamingToolCall = {
         type: "toolCall",
         id: `${stringValue(item["call_id"])}|${stringValue(item.id)}`,
@@ -259,17 +279,17 @@ export async function processCodexStream(
         partialJson: typeof item.arguments === "string" ? item.arguments : "",
       };
       output.content.push(block);
-      const slot = {
+      const slot = trackStarted({
         type: "toolCall",
         block,
         contentIndex: output.content.length - 1,
-      } satisfies OutputSlot;
+      } satisfies OutputSlot);
       slots.set(index, slot);
       stream.push({ type: "toolcall_start", contentIndex: slot.contentIndex, partial: output });
       return slot;
     }
     if (item.type === "custom_tool_call") {
-      const name = stringValue(item.name);
+      const name = piToolCallName(item);
       const property = grammarToolInputProperties.get(name) ?? "input";
       const input = typeof item["input"] === "string" ? item["input"] : "";
       const block: StreamingToolCall = {
@@ -283,11 +303,11 @@ export async function processCodexStream(
         },
       };
       output.content.push(block);
-      const slot = {
+      const slot = trackStarted({
         type: "toolCall",
         block,
         contentIndex: output.content.length - 1,
-      } satisfies OutputSlot;
+      } satisfies OutputSlot);
       slots.set(index, slot);
       stream.push({ type: "toolcall_start", contentIndex: slot.contentIndex, partial: output });
       return slot;
@@ -448,6 +468,7 @@ export async function processCodexStream(
           content: slot.block.thinking,
           partial: output,
         });
+        trackCompleted(slot);
         slots.delete(index);
       } else if (item.type === "message" && slot?.type === "text") {
         slot.block.text = itemContentText(item);
@@ -460,17 +481,14 @@ export async function processCodexStream(
           content: slot.block.text,
           partial: output,
         });
+        trackCompleted(slot);
         slots.delete(index);
       } else if (
         item.type === "function_call" &&
         slot?.type === "toolCall" &&
         slot.block.partialJson !== undefined
       ) {
-        if (item["namespace"] !== undefined) {
-          slot.block.name = namespacedToolCallName(item["namespace"], item.name);
-        } else if (typeof item.name === "string" && CODEX_NAMESPACED_TOOL_NAMES.has(item.name)) {
-          throw new Error(`Codex returned namespaced tool "${item.name}" as a flat function call.`);
-        }
+        slot.block.name = piToolCallName(item);
         const argumentsJson =
           typeof item.arguments === "string" ? item.arguments : slot.block.partialJson || "{}";
         slot.block.arguments = parseStreamingJson(argumentsJson);
@@ -481,8 +499,10 @@ export async function processCodexStream(
           toolCall: slot.block,
           partial: output,
         });
+        trackCompleted(slot);
         slots.delete(index);
       } else if (item.type === "custom_tool_call" && slot?.type === "toolCall") {
+        slot.block.name = piToolCallName(item);
         const input = typeof item["input"] === "string" ? item["input"] : customInput(slot.block);
         pushToolDelta(slot, appendCustomInput(slot.block, input, true));
         delete slot.block.customInput;
@@ -492,6 +512,7 @@ export async function processCodexStream(
           toolCall: slot.block,
           partial: output,
         });
+        trackCompleted(slot);
         slots.delete(index);
       }
     } else if (
@@ -500,12 +521,17 @@ export async function processCodexStream(
     ) {
       finalize(event.response);
     } else if (event.type === "response.failed") {
-      terminal = true;
       const response = isObject(event.response) ? event.response : undefined;
+      if (response) {
+        finalize(response);
+      } else {
+        terminal = true;
+      }
+      output.stopReason = "error";
+      output.rawStopReason ??= "failed";
       const error = isObject(response?.["error"]) ? response["error"] : undefined;
-      throw new Error(
-        typeof error?.["message"] === "string" ? error["message"] : "Codex response failed",
-      );
+      output.errorMessage =
+        typeof error?.["message"] === "string" ? error["message"] : "Codex response failed";
     }
   }
 

package/extensions/openai-codex-compat/codex-thread-lineage.ts

@@ -0,0 +1,156 @@
+import type { ExtensionAPI, ExtensionContext, SessionEntry } from "@earendil-works/pi-coding-agent";
+import { uuidv7 } from "@earendil-works/pi-ai";
+import { codexCacheKey } from "./codex-cache-key.ts";
+
+export const CODEX_THREAD_MARKER_ENTRY_TYPE = "openai-codex-compat-thread";
+
+export type CodexThreadMarkerData = {
+  version: 1;
+  sessionId: string;
+  threadId: string;
+  forkedFromThreadId: string;
+  branchParentEntryId: string | null;
+};
+
+export type CodexThreadIdentity = {
+  threadId: string;
+  forkedFromThreadId?: string;
+};
+
+type PendingTreeFork = {
+  expectedLeafId: string | null;
+};
+
+function markerData(entry: SessionEntry, sessionId: string): CodexThreadMarkerData | undefined {
+  if (entry.type !== "custom" || entry.customType !== CODEX_THREAD_MARKER_ENTRY_TYPE) {
+    return undefined;
+  }
+  const data = entry.data;
+  if (typeof data !== "object" || data === null || Array.isArray(data)) {
+    throw new Error("The active Pi branch contains an invalid OpenAI Codex thread marker.");
+  }
+  const candidate = data as Record<string, unknown>;
+  if (
+    candidate["version"] !== 1 ||
+    typeof candidate["sessionId"] !== "string" ||
+    candidate["sessionId"].length === 0 ||
+    typeof candidate["threadId"] !== "string" ||
+    candidate["threadId"].length === 0 ||
+    typeof candidate["forkedFromThreadId"] !== "string" ||
+    candidate["forkedFromThreadId"].length === 0 ||
+    (candidate["branchParentEntryId"] !== null &&
+      typeof candidate["branchParentEntryId"] !== "string")
+  ) {
+    throw new Error("The active Pi branch contains an invalid OpenAI Codex thread marker.");
+  }
+  if (candidate["sessionId"] !== sessionId) return undefined;
+  if (entry.parentId !== candidate["branchParentEntryId"]) {
+    throw new Error("The active Pi branch contains a misplaced OpenAI Codex thread marker.");
+  }
+  return candidate as CodexThreadMarkerData;
+}
+
+function latestMarkerIndex(sessionId: string, branch: readonly SessionEntry[]): number {
+  for (let index = branch.length - 1; index >= 0; index -= 1) {
+    if (markerData(branch[index]!, sessionId)) return index;
+  }
+  return -1;
+}
+
+export function resolveCodexThreadIdentity(
+  sessionId: string,
+  branch: readonly SessionEntry[],
+): CodexThreadIdentity {
+  for (let index = branch.length - 1; index >= 0; index -= 1) {
+    const marker = markerData(branch[index]!, sessionId);
+    if (marker) {
+      return {
+        threadId: marker.threadId,
+        forkedFromThreadId: marker.forkedFromThreadId,
+      };
+    }
+  }
+  return { threadId: codexCacheKey(sessionId)! };
+}
+
+function shouldForkOnNextAppend(
+  sessionId: string,
+  branch: readonly SessionEntry[],
+  entries: readonly SessionEntry[],
+): boolean {
+  // Pi does not persist an explicit branch-start flag. Its append order is
+  // stable, so the first child inherits its thread and later children are
+  // forks unless the active path already contains a marker for that fork.
+  const firstChildByParent = new Map<string | null, string>();
+  for (const entry of entries) {
+    if (!firstChildByParent.has(entry.parentId)) {
+      firstChildByParent.set(entry.parentId, entry.id);
+    }
+  }
+
+  const markerIndex = latestMarkerIndex(sessionId, branch);
+  for (let index = markerIndex + 1; index < branch.length; index += 1) {
+    const entry = branch[index]!;
+    if (firstChildByParent.get(entry.parentId) !== entry.id) return true;
+  }
+
+  const leafId = branch.at(-1)?.id ?? null;
+  return entries.some((entry) => entry.parentId === leafId);
+}
+
+function armPendingFork(pending: Map<string, PendingTreeFork>, ctx: ExtensionContext): void {
+  const sessionId = ctx.sessionManager.getSessionId();
+  const branch = ctx.sessionManager.getBranch() as SessionEntry[];
+  const entries = ctx.sessionManager.getEntries() as SessionEntry[];
+  if (!shouldForkOnNextAppend(sessionId, branch, entries)) {
+    pending.delete(sessionId);
+    return;
+  }
+  pending.set(sessionId, { expectedLeafId: ctx.sessionManager.getLeafId() });
+}
+
+function pendingLeafIsActive(pending: PendingTreeFork, branch: readonly SessionEntry[]): boolean {
+  return (
+    pending.expectedLeafId === null || branch.some((entry) => entry.id === pending.expectedLeafId)
+  );
+}
+
+export default function registerCodexThreadLineage(pi: ExtensionAPI): void {
+  const pending = new Map<string, PendingTreeFork>();
+
+  pi.on("session_start", (_event, ctx) => {
+    armPendingFork(pending, ctx);
+  });
+
+  pi.on("session_tree", (_event, ctx) => {
+    armPendingFork(pending, ctx);
+  });
+
+  pi.on("message_end", (event, ctx) => {
+    if (event.message.role !== "user") return;
+    const sessionId = ctx.sessionManager.getSessionId();
+    const candidate = pending.get(sessionId);
+    if (!candidate) return;
+    pending.delete(sessionId);
+
+    const branch = ctx.sessionManager.getBranch() as SessionEntry[];
+    if (!pendingLeafIsActive(candidate, branch)) return;
+
+    const parent = resolveCodexThreadIdentity(sessionId, branch);
+    const branchParentEntryId = ctx.sessionManager.getLeafId();
+    // Pi invokes message_end handlers immediately before persisting the
+    // finalized user message. Advancing the leaf here makes this context-free
+    // marker the user's parent without writing anything during /tree itself.
+    pi.appendEntry<CodexThreadMarkerData>(CODEX_THREAD_MARKER_ENTRY_TYPE, {
+      version: 1,
+      sessionId,
+      threadId: uuidv7(),
+      forkedFromThreadId: parent.threadId,
+      branchParentEntryId,
+    });
+  });
+
+  pi.on("session_shutdown", (_event, ctx) => {
+    pending.delete(ctx.sessionManager.getSessionId());
+  });
+}