pi-openai-codex-compat 0.0.2 → 0.0.4

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

@@ -8,6 +8,7 @@ import {
   type TextSignatureV1,
   type ThinkingContent,
   type ToolCall,
+  type Usage,
 } from "@earendil-works/pi-ai";
 import { isObject, type JsonRecord } from "./codex-protocol.ts";
 import { CODEX_NAMESPACED_TOOL_NAMES, namespacedToolCallName } from "./namespaced-tools.ts";
@@ -38,6 +39,33 @@ type OutputSlot =
 
 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 =
+  | "completed"
+  | "incomplete"
+  | "failed"
+  | "cancelled"
+  | "queued"
+  | "in_progress";
+
+const CODEX_RESPONSE_STATUSES = new Set<CodexResponseStatus>([
+  "completed",
+  "incomplete",
+  "failed",
+  "cancelled",
+  "queued",
+  "in_progress",
+]);
+
 function outputIndex(event: JsonRecord): number {
   return typeof event["output_index"] === "number" ? event["output_index"] : 0;
 }
@@ -46,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;
@@ -60,7 +100,7 @@ function appendGrammarDelta(
 ): string | undefined {
   if (buffer.closed) {
     if (close && nextInput === buffer.input) return undefined;
-    throw new Error(`grammar tool input for property "${property}" changed after closure`);
+    throw new Error(`grammar tool input for property "${property}" changed after it was closed`);
   }
   if (!nextInput.startsWith(buffer.input)) {
     throw new Error(`grammar tool input for property "${property}" changed non-monotonically`);
@@ -128,13 +168,35 @@ function reasoningText(item: JsonRecord): string {
         .join("\n\n")
     : "";
   if (summary) return summary;
-  return itemContentText(item);
+  return Array.isArray(item.content)
+    ? item.content
+        .filter(isObject)
+        .map((part) => (typeof part.text === "string" ? part.text : ""))
+        .join("\n\n")
+    : "";
+}
+
+function normalizeCodexStatus(status: unknown): CodexResponseStatus | undefined {
+  return typeof status === "string" && CODEX_RESPONSE_STATUSES.has(status as CodexResponseStatus)
+    ? (status as CodexResponseStatus)
+    : undefined;
 }
 
-function mapStopReason(status: unknown): AssistantMessage["stopReason"] {
-  if (status === "incomplete") return "length";
-  if (status === "failed" || status === "cancelled") return "error";
-  return "stop";
+function mapStopReason(
+  status: CodexResponseStatus | undefined,
+  incompleteReason: string | undefined,
+): { stopReason: AssistantMessage["stopReason"]; errorMessage?: string } {
+  if (status === "incomplete") {
+    if (incompleteReason === "max_output_tokens") return { stopReason: "length" };
+    return {
+      stopReason: "error",
+      errorMessage: incompleteReason
+        ? `Response incomplete: ${incompleteReason}`
+        : "Response incomplete without a provider reason",
+    };
+  }
+  if (status === "failed" || status === "cancelled") return { stopReason: "error" };
+  return { stopReason: "stop" };
 }
 
 export async function processCodexStream(
@@ -143,10 +205,16 @@ export async function processCodexStream(
   stream: AssistantMessageEventStream,
   model: Model<any>,
   grammarToolInputProperties: ReadonlyMap<string, string>,
+  options?: ProcessCodexStreamOptions,
 ): Promise<void> {
   let terminal = false;
   const slots = new Map<number, OutputSlot>();
   const reasoningById = new Map<string, ThinkingContent>();
+  const applyMessagePhaseStopReason = (item: JsonRecord): void => {
+    if (item.type === "message" && item["phase"] === "final_answer") {
+      output.stopReason = "stop";
+    }
+  };
 
   const getSlot = <TType extends OutputSlot["type"]>(
     index: number,
@@ -166,40 +234,43 @@ 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;
     }
     if (item.type === "message") {
+      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)}`,
@@ -208,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 = {
@@ -232,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;
@@ -270,12 +341,15 @@ export async function processCodexStream(
           typeof outputDetails?.["reasoning_tokens"] === "number"
             ? outputDetails["reasoning_tokens"]
             : 0,
-        totalTokens:
-          typeof usage.total_tokens === "number" ? usage.total_tokens : input + outputTokens,
+        totalTokens: typeof usage.total_tokens === "number" ? usage.total_tokens || 0 : 0,
         cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
       };
-      calculateCost(model, output.usage);
     }
+    calculateCost(model, output.usage);
+    options?.applyServiceTierPricing?.(
+      output.usage,
+      typeof response.service_tier === "string" ? response.service_tier : undefined,
+    );
     for (const item of responseItems(response["output"])) {
       if (item.type !== "reasoning" || typeof item.id !== "string") continue;
       const block = reasoningById.get(item.id);
@@ -288,8 +362,20 @@ export async function processCodexStream(
         });
       }
     }
-    if (typeof response["status"] === "string") output.rawStopReason = response["status"];
-    output.stopReason = mapStopReason(response["status"]);
+    const status = normalizeCodexStatus(response["status"]);
+    const incompleteDetails = isObject(response["incomplete_details"])
+      ? response["incomplete_details"]
+      : undefined;
+    const incompleteReason =
+      typeof incompleteDetails?.["reason"] === "string" ? incompleteDetails["reason"] : undefined;
+    const rawStopReason =
+      status === "incomplete" && incompleteReason ? `${status}.${incompleteReason}` : status;
+    if (rawStopReason === undefined) delete output.rawStopReason;
+    else output.rawStopReason = rawStopReason;
+    const mappedStop = mapStopReason(status, incompleteReason);
+    output.stopReason = mappedStop.stopReason;
+    if (mappedStop.errorMessage === undefined) delete output.errorMessage;
+    else output.errorMessage = mappedStop.errorMessage;
     if (output.stopReason === "stop" && output.content.some((block) => block.type === "toolCall")) {
       output.stopReason = "toolUse";
     }
@@ -353,8 +439,10 @@ export async function processCodexStream(
       const previous = slot.block.partialJson;
       slot.block.partialJson = event.arguments;
       slot.block.arguments = parseStreamingJson(event.arguments);
-      if (event.arguments.startsWith(previous))
-        pushToolDelta(slot, event.arguments.slice(previous.length));
+      if (event.arguments.startsWith(previous)) {
+        const delta = event.arguments.slice(previous.length);
+        if (delta.length > 0) pushToolDelta(slot, delta);
+      }
     } else if (event.type === "response.custom_tool_call_input.delta") {
       const slot = getSlot(index, "toolCall");
       if (!slot || typeof event["delta"] !== "string") continue;
@@ -368,6 +456,7 @@ export async function processCodexStream(
       pushToolDelta(slot, appendCustomInput(slot.block, event["input"], true));
     } else if (event.type === "response.output_item.done" && isObject(event.item)) {
       const item = event.item;
+      applyMessagePhaseStopReason(item);
       const slot = slotFor(index, item);
       if (item.type === "reasoning" && slot?.type === "thinking") {
         slot.block.thinking = reasoningText(item) || slot.block.thinking;
@@ -379,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);
@@ -391,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);
@@ -412,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;
@@ -423,6 +512,7 @@ export async function processCodexStream(
           toolCall: slot.block,
           partial: output,
         });
+        trackCompleted(slot);
         slots.delete(index);
       }
     } else if (
@@ -431,14 +521,21 @@ 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";
     }
   }
 
-  if (!terminal) throw new Error("Codex stream ended before a terminal response event");
+  if (!terminal) {
+    throw new Error("OpenAI Responses stream ended before a terminal response event");
+  }
 }

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());
+  });
+}