pi-openai-codex-compat 0.0.2 → 0.0.3

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,27 @@ type OutputSlot =
 
 type ToolCallSlot = Extract<OutputSlot, { type: "toolCall" }>;
 
+type ProcessCodexStreamOptions = {
+  applyServiceTierPricing?(usage: Usage, responseServiceTier: string | undefined): void;
+};
+
+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;
 }
@@ -60,7 +82,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 +150,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 mapStopReason(status: unknown): AssistantMessage["stopReason"] {
-  if (status === "incomplete") return "length";
-  if (status === "failed" || status === "cancelled") return "error";
-  return "stop";
+function normalizeCodexStatus(status: unknown): CodexResponseStatus | undefined {
+  return typeof status === "string" && CODEX_RESPONSE_STATUSES.has(status as CodexResponseStatus)
+    ? (status as CodexResponseStatus)
+    : undefined;
+}
+
+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 +187,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,
@@ -180,6 +230,7 @@ export async function processCodexStream(
       return slot;
     }
     if (item.type === "message") {
+      applyMessagePhaseStopReason(item);
       const block: TextContent = { type: "text", text: "" };
       output.content.push(block);
       const slot = {
@@ -270,12 +321,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 +342,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 +419,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 +436,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;
@@ -440,5 +509,7 @@ export async function processCodexStream(
     }
   }
 
-  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");
+  }
 }