pi-openai-codex-compat 0.0.1-alpha.0

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

@@ -0,0 +1,444 @@
+import {
+  calculateCost,
+  parseStreamingJson,
+  type AssistantMessage,
+  type AssistantMessageEventStream,
+  type Model,
+  type TextContent,
+  type TextSignatureV1,
+  type ThinkingContent,
+  type ToolCall,
+} from "@earendil-works/pi-ai";
+import { isObject, type JsonRecord } from "./codex-protocol.ts";
+import { CODEX_NAMESPACED_TOOL_NAMES, namespacedToolCallName } from "./namespaced-tools.ts";
+
+/**
+ * Focused adaptation of @earendil-works/pi-ai@0.83.0
+ * src/api/openai-responses-shared.ts stream processing.
+ */
+
+type GrammarJsonBuffer = {
+  input: string;
+  started: boolean;
+  closed: boolean;
+};
+
+type StreamingToolCall = ToolCall & {
+  partialJson?: string;
+  customInput?: {
+    property: string;
+    jsonBuffer: GrammarJsonBuffer;
+  };
+};
+
+type OutputSlot =
+  | { type: "thinking"; block: ThinkingContent; contentIndex: number }
+  | { type: "text"; block: TextContent; contentIndex: number }
+  | { type: "toolCall"; block: StreamingToolCall; contentIndex: number };
+
+type ToolCallSlot = Extract<OutputSlot, { type: "toolCall" }>;
+
+function outputIndex(event: JsonRecord): number {
+  return typeof event["output_index"] === "number" ? event["output_index"] : 0;
+}
+
+function stringValue(value: unknown): string {
+  return typeof value === "string" ? value : "";
+}
+
+function encodeTextSignature(id: string, phase: unknown): string {
+  const payload: TextSignatureV1 = { v: 1, id };
+  if (phase === "commentary" || phase === "final_answer") payload.phase = phase;
+  return JSON.stringify(payload);
+}
+
+function appendGrammarDelta(
+  buffer: GrammarJsonBuffer,
+  property: string,
+  nextInput: string,
+  close: boolean,
+): string | undefined {
+  if (buffer.closed) {
+    if (close && nextInput === buffer.input) return undefined;
+    throw new Error(`grammar tool input for property "${property}" changed after closure`);
+  }
+  if (!nextInput.startsWith(buffer.input)) {
+    throw new Error(`grammar tool input for property "${property}" changed non-monotonically`);
+  }
+  const inputDelta = nextInput.slice(buffer.input.length);
+  if (!close && inputDelta.length === 0) return undefined;
+
+  let delta = "";
+  if (!buffer.started) {
+    delta = `{${JSON.stringify(property)}:"`;
+    buffer.started = true;
+  }
+  delta += JSON.stringify(inputDelta).slice(1, -1);
+  buffer.input = nextInput;
+  if (close) {
+    delta += '"}';
+    buffer.closed = true;
+  }
+  return delta;
+}
+
+function customInput(block: StreamingToolCall): string {
+  const property = block.customInput?.property;
+  if (!property) return "";
+  const value = block.arguments[property];
+  return typeof value === "string" ? value : "";
+}
+
+function appendCustomInput(
+  block: StreamingToolCall,
+  nextInput: string,
+  close: boolean,
+): string | undefined {
+  const state = block.customInput;
+  if (!state) return undefined;
+  const delta = appendGrammarDelta(state.jsonBuffer, state.property, nextInput, close);
+  block.arguments = { [state.property]: nextInput };
+  return delta;
+}
+
+function responseItems(value: unknown): JsonRecord[] {
+  if (!Array.isArray(value)) return [];
+  return value.filter(isObject);
+}
+
+function itemContentText(item: JsonRecord): string {
+  if (!Array.isArray(item.content)) return "";
+  return item.content
+    .filter(isObject)
+    .map((content) =>
+      typeof content.text === "string"
+        ? content.text
+        : typeof content["refusal"] === "string"
+          ? content["refusal"]
+          : "",
+    )
+    .join("");
+}
+
+function reasoningText(item: JsonRecord): string {
+  const summary = Array.isArray(item["summary"])
+    ? item["summary"]
+        .filter(isObject)
+        .map((part) => (typeof part.text === "string" ? part.text : ""))
+        .join("\n\n")
+    : "";
+  if (summary) return summary;
+  return itemContentText(item);
+}
+
+function mapStopReason(status: unknown): AssistantMessage["stopReason"] {
+  if (status === "incomplete") return "length";
+  if (status === "failed" || status === "cancelled") return "error";
+  return "stop";
+}
+
+export async function processCodexStream(
+  events: AsyncIterable<JsonRecord>,
+  output: AssistantMessage,
+  stream: AssistantMessageEventStream,
+  model: Model<any>,
+  grammarToolInputProperties: ReadonlyMap<string, string>,
+): Promise<void> {
+  let terminal = false;
+  const slots = new Map<number, OutputSlot>();
+  const reasoningById = new Map<string, ThinkingContent>();
+
+  const getSlot = <TType extends OutputSlot["type"]>(
+    index: number,
+    type: TType,
+  ): Extract<OutputSlot, { type: TType }> | undefined => {
+    const slot = slots.get(index);
+    return slot?.type === type ? (slot as Extract<OutputSlot, { type: TType }>) : undefined;
+  };
+
+  const pushToolDelta = (slot: ToolCallSlot, delta: string | undefined): void => {
+    if (delta === undefined) return;
+    stream.push({
+      type: "toolcall_delta",
+      contentIndex: slot.contentIndex,
+      delta,
+      partial: output,
+    });
+  };
+
+  const createSlot = (index: number, item: JsonRecord): OutputSlot | undefined => {
+    if (item.type === "reasoning") {
+      const block: ThinkingContent = { type: "thinking", thinking: "" };
+      output.content.push(block);
+      const slot = {
+        type: "thinking",
+        block,
+        contentIndex: output.content.length - 1,
+      } satisfies OutputSlot;
+      slots.set(index, slot);
+      stream.push({ type: "thinking_start", contentIndex: slot.contentIndex, partial: output });
+      return slot;
+    }
+    if (item.type === "message") {
+      const block: TextContent = { type: "text", text: "" };
+      output.content.push(block);
+      const slot = {
+        type: "text",
+        block,
+        contentIndex: output.content.length - 1,
+      } 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 block: StreamingToolCall = {
+        type: "toolCall",
+        id: `${stringValue(item["call_id"])}|${stringValue(item.id)}`,
+        name,
+        arguments: {},
+        partialJson: typeof item.arguments === "string" ? item.arguments : "",
+      };
+      output.content.push(block);
+      const slot = {
+        type: "toolCall",
+        block,
+        contentIndex: output.content.length - 1,
+      } 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 property = grammarToolInputProperties.get(name) ?? "input";
+      const input = typeof item["input"] === "string" ? item["input"] : "";
+      const block: StreamingToolCall = {
+        type: "toolCall",
+        id: `${stringValue(item["call_id"])}|${stringValue(item.id)}`,
+        name,
+        arguments: { [property]: input },
+        customInput: {
+          property,
+          jsonBuffer: { input: "", started: false, closed: false },
+        },
+      };
+      output.content.push(block);
+      const slot = {
+        type: "toolCall",
+        block,
+        contentIndex: output.content.length - 1,
+      } satisfies OutputSlot;
+      slots.set(index, slot);
+      stream.push({ type: "toolcall_start", contentIndex: slot.contentIndex, partial: output });
+      return slot;
+    }
+    return undefined;
+  };
+
+  const slotFor = (index: number, item: JsonRecord): OutputSlot | undefined =>
+    slots.get(index) ?? createSlot(index, item);
+
+  const finalize = (response: JsonRecord): void => {
+    terminal = true;
+    if (typeof response.id === "string") output.responseId = response.id;
+    const usage = isObject(response.usage) ? response.usage : undefined;
+    if (usage) {
+      const details = isObject(usage.input_tokens_details) ? usage.input_tokens_details : undefined;
+      const cached = typeof details?.cached_tokens === "number" ? details.cached_tokens : 0;
+      const cacheWrite =
+        typeof details?.cache_write_tokens === "number" ? details.cache_write_tokens : 0;
+      const input = typeof usage.input_tokens === "number" ? usage.input_tokens : 0;
+      const outputTokens = typeof usage.output_tokens === "number" ? usage.output_tokens : 0;
+      const outputDetails = isObject(usage["output_tokens_details"])
+        ? usage["output_tokens_details"]
+        : undefined;
+      output.usage = {
+        input: Math.max(0, input - cached - cacheWrite),
+        output: outputTokens,
+        cacheRead: cached,
+        cacheWrite,
+        reasoning:
+          typeof outputDetails?.["reasoning_tokens"] === "number"
+            ? outputDetails["reasoning_tokens"]
+            : 0,
+        totalTokens:
+          typeof usage.total_tokens === "number" ? usage.total_tokens : input + outputTokens,
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+      };
+      calculateCost(model, output.usage);
+    }
+    for (const item of responseItems(response["output"])) {
+      if (item.type !== "reasoning" || typeof item.id !== "string") continue;
+      const block = reasoningById.get(item.id);
+      if (!block?.thinkingSignature || typeof item.encrypted_content !== "string") continue;
+      const stored = JSON.parse(block.thinkingSignature) as JsonRecord;
+      if (typeof stored.encrypted_content !== "string") {
+        block.thinkingSignature = JSON.stringify({
+          ...stored,
+          encrypted_content: item.encrypted_content,
+        });
+      }
+    }
+    if (typeof response["status"] === "string") output.rawStopReason = response["status"];
+    output.stopReason = mapStopReason(response["status"]);
+    if (output.stopReason === "stop" && output.content.some((block) => block.type === "toolCall")) {
+      output.stopReason = "toolUse";
+    }
+  };
+
+  for await (const event of events) {
+    const index = outputIndex(event);
+    if (event.type === "response.created" && isObject(event.response)) {
+      if (typeof event.response.id === "string") output.responseId = event.response.id;
+    } else if (event.type === "response.output_item.added" && isObject(event.item)) {
+      createSlot(index, event.item);
+    } else if (
+      event.type === "response.reasoning_summary_text.delta" ||
+      event.type === "response.reasoning_text.delta"
+    ) {
+      const slot = getSlot(index, "thinking");
+      if (!slot || typeof event["delta"] !== "string") continue;
+      slot.block.thinking += event["delta"];
+      stream.push({
+        type: "thinking_delta",
+        contentIndex: slot.contentIndex,
+        delta: event["delta"],
+        partial: output,
+      });
+    } else if (event.type === "response.reasoning_summary_part.done") {
+      const slot = getSlot(index, "thinking");
+      if (!slot) continue;
+      slot.block.thinking += "\n\n";
+      stream.push({
+        type: "thinking_delta",
+        contentIndex: slot.contentIndex,
+        delta: "\n\n",
+        partial: output,
+      });
+    } else if (
+      event.type === "response.output_text.delta" ||
+      event.type === "response.refusal.delta"
+    ) {
+      const slot = getSlot(index, "text");
+      if (!slot || typeof event["delta"] !== "string") continue;
+      slot.block.text += event["delta"];
+      stream.push({
+        type: "text_delta",
+        contentIndex: slot.contentIndex,
+        delta: event["delta"],
+        partial: output,
+      });
+    } else if (event.type === "response.function_call_arguments.delta") {
+      const slot = getSlot(index, "toolCall");
+      if (!slot || slot.block.partialJson === undefined || typeof event["delta"] !== "string") {
+        continue;
+      }
+      slot.block.partialJson += event["delta"];
+      slot.block.arguments = parseStreamingJson(slot.block.partialJson);
+      pushToolDelta(slot, event["delta"]);
+    } else if (event.type === "response.function_call_arguments.done") {
+      const slot = getSlot(index, "toolCall");
+      if (!slot || slot.block.partialJson === undefined || typeof event.arguments !== "string") {
+        continue;
+      }
+      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));
+    } else if (event.type === "response.custom_tool_call_input.delta") {
+      const slot = getSlot(index, "toolCall");
+      if (!slot || typeof event["delta"] !== "string") continue;
+      pushToolDelta(
+        slot,
+        appendCustomInput(slot.block, customInput(slot.block) + event["delta"], false),
+      );
+    } else if (event.type === "response.custom_tool_call_input.done") {
+      const slot = getSlot(index, "toolCall");
+      if (!slot || typeof event["input"] !== "string") continue;
+      pushToolDelta(slot, appendCustomInput(slot.block, event["input"], true));
+    } else if (event.type === "response.output_item.done" && isObject(event.item)) {
+      const item = event.item;
+      const slot = slotFor(index, item);
+      if (item.type === "reasoning" && slot?.type === "thinking") {
+        slot.block.thinking = reasoningText(item) || slot.block.thinking;
+        slot.block.thinkingSignature = JSON.stringify(item);
+        if (typeof item.id === "string") reasoningById.set(item.id, slot.block);
+        stream.push({
+          type: "thinking_end",
+          contentIndex: slot.contentIndex,
+          content: slot.block.thinking,
+          partial: output,
+        });
+        slots.delete(index);
+      } else if (item.type === "message" && slot?.type === "text") {
+        slot.block.text = itemContentText(item);
+        if (typeof item.id === "string") {
+          slot.block.textSignature = encodeTextSignature(item.id, item["phase"]);
+        }
+        stream.push({
+          type: "text_end",
+          contentIndex: slot.contentIndex,
+          content: slot.block.text,
+          partial: output,
+        });
+        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.`);
+        }
+        const argumentsJson =
+          typeof item.arguments === "string" ? item.arguments : slot.block.partialJson || "{}";
+        slot.block.arguments = parseStreamingJson(argumentsJson);
+        delete slot.block.partialJson;
+        stream.push({
+          type: "toolcall_end",
+          contentIndex: slot.contentIndex,
+          toolCall: slot.block,
+          partial: output,
+        });
+        slots.delete(index);
+      } else if (item.type === "custom_tool_call" && slot?.type === "toolCall") {
+        const input = typeof item["input"] === "string" ? item["input"] : customInput(slot.block);
+        pushToolDelta(slot, appendCustomInput(slot.block, input, true));
+        delete slot.block.customInput;
+        stream.push({
+          type: "toolcall_end",
+          contentIndex: slot.contentIndex,
+          toolCall: slot.block,
+          partial: output,
+        });
+        slots.delete(index);
+      }
+    } else if (
+      (event.type === "response.completed" || event.type === "response.incomplete") &&
+      isObject(event.response)
+    ) {
+      finalize(event.response);
+    } else if (event.type === "response.failed") {
+      terminal = true;
+      const response = isObject(event.response) ? event.response : undefined;
+      const error = isObject(response?.["error"]) ? response["error"] : undefined;
+      throw new Error(
+        typeof error?.["message"] === "string" ? error["message"] : "Codex response failed",
+      );
+    }
+  }
+
+  if (!terminal) throw new Error("Codex stream ended before a terminal response event");
+}

package/extensions/openai-codex-compat/codex-tool-surface.ts

@@ -0,0 +1,186 @@
+import type { Theme } from "@earendil-works/pi-coding-agent";
+import { type Component, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
+import type { CodexToolBackground } from "./config.ts";
+
+export type CodexToolSurfaceStatus = "pending" | "success" | "error";
+export type CodexToolBackgroundResolver = () => CodexToolBackground;
+type ThemeBg =
+  | "selectedBg"
+  | "userMessageBg"
+  | "customMessageBg"
+  | "toolPendingBg"
+  | "toolSuccessBg"
+  | "toolErrorBg";
+
+const ANSI_RESET_BACKGROUND = "\u001b[49m";
+const DARK_256_SURFACE_BG = "\u001b[48;5;234m";
+const LIGHT_256_SURFACE_BG = "\u001b[48;5;255m";
+const TRUECOLOR_BACKGROUND_PATTERN = new RegExp(String.raw`\u001b\[48;2;(\d+);(\d+);(\d+)m`);
+const INDEXED_BACKGROUND_PATTERN = new RegExp(String.raw`\u001b\[48;5;(\d+)m`);
+const BACKGROUND_RESET_PATTERN = new RegExp(String.raw`\u001b\[(?:0|49)m`, "g");
+
+function xtermChannel(index: number): number {
+  return index === 0 ? 0 : 55 + index * 40;
+}
+
+function xterm256ToRgb(index: number): [number, number, number] {
+  if (index < 16) {
+    const standard: Array<[number, number, number]> = [
+      [0, 0, 0],
+      [128, 0, 0],
+      [0, 128, 0],
+      [128, 128, 0],
+      [0, 0, 128],
+      [128, 0, 128],
+      [0, 128, 128],
+      [192, 192, 192],
+      [128, 128, 128],
+      [255, 0, 0],
+      [0, 255, 0],
+      [255, 255, 0],
+      [0, 0, 255],
+      [255, 0, 255],
+      [0, 255, 255],
+      [255, 255, 255],
+    ];
+    return standard[index] ?? [0, 0, 0];
+  }
+  if (index < 232) {
+    const offset = index - 16;
+    return [
+      xtermChannel(Math.floor(offset / 36)),
+      xtermChannel(Math.floor((offset % 36) / 6)),
+      xtermChannel(offset % 6),
+    ];
+  }
+  const gray = 8 + (index - 232) * 10;
+  return [gray, gray, gray];
+}
+
+function backgroundRgb(theme: Theme, background: ThemeBg): [number, number, number] | undefined {
+  const ansi = theme.getBgAnsi?.(background);
+  if (!ansi) return undefined;
+  const truecolor = ansi.match(TRUECOLOR_BACKGROUND_PATTERN);
+  if (truecolor) {
+    return [Number(truecolor[1]), Number(truecolor[2]), Number(truecolor[3])];
+  }
+  const indexed = ansi.match(INDEXED_BACKGROUND_PATTERN);
+  return indexed ? xterm256ToRgb(Number(indexed[1])) : undefined;
+}
+
+export function usesLightToolPalette(theme: Theme): boolean {
+  const rgb = backgroundRgb(theme, "toolSuccessBg");
+  if (rgb) {
+    const red = rgb[0] / 255;
+    const green = rgb[1] / 255;
+    const blue = rgb[2] / 255;
+    const luminance = 0.2126 * red + 0.7152 * green + 0.0722 * blue;
+    return luminance >= 0.6;
+  }
+  return theme.name?.toLowerCase().includes("light") ?? false;
+}
+
+function blendRgb(
+  overlay: [number, number, number],
+  background: [number, number, number],
+  alpha: number,
+): [number, number, number] {
+  return [
+    Math.round(overlay[0] * alpha + background[0] * (1 - alpha)),
+    Math.round(overlay[1] * alpha + background[1] * (1 - alpha)),
+    Math.round(overlay[2] * alpha + background[2] * (1 - alpha)),
+  ];
+}
+
+function subtleBackground(theme: Theme): string {
+  const light = usesLightToolPalette(theme);
+  if (theme.getColorMode?.() === "256color") {
+    return light ? LIGHT_256_SURFACE_BG : DARK_256_SURFACE_BG;
+  }
+
+  const base =
+    backgroundRgb(theme, "toolPendingBg") ??
+    (light ? ([232, 232, 240] as const) : ([40, 40, 50] as const));
+  const overlay: [number, number, number] = light ? [255, 255, 255] : [0, 0, 0];
+  const alpha = light ? 0.55 : 0.35;
+  const [red, green, blue] = blendRgb(overlay, base, alpha);
+  return `\u001b[48;2;${red};${green};${blue}m`;
+}
+
+function statusBackground(theme: Theme, status: CodexToolSurfaceStatus): string {
+  const token: ThemeBg =
+    status === "pending" ? "toolPendingBg" : status === "error" ? "toolErrorBg" : "toolSuccessBg";
+  return theme.getBgAnsi?.(token) ?? "";
+}
+
+function surfaceBackground(
+  theme: Theme,
+  style: CodexToolBackground,
+  status: CodexToolSurfaceStatus,
+): string {
+  if (style === "none") return "";
+  return style === "status" ? statusBackground(theme, status) : subtleBackground(theme);
+}
+
+function restoreBackgroundAfterResets(text: string, background: string): string {
+  if (!background) return text;
+  return text.replace(BACKGROUND_RESET_PATTERN, (reset) => `${reset}${background}`);
+}
+
+function withBackground(text: string, background: string): string {
+  if (!background) return text;
+  return `${background}${restoreBackgroundAfterResets(text, background)}${ANSI_RESET_BACKGROUND}`;
+}
+
+function fillLine(line: string, width: number, background: string): string {
+  const truncated = truncateToWidth(line, width, "");
+  const padding = " ".repeat(Math.max(0, width - visibleWidth(truncated)));
+  return withBackground(`${truncated}${padding}`, background);
+}
+
+export class CodexToolSurfaceComponent implements Component {
+  private readonly component: Component;
+  private readonly theme: Theme;
+  private readonly resolveBackground: CodexToolBackgroundResolver;
+  private readonly status: CodexToolSurfaceStatus;
+  private readonly topPadding: boolean;
+  private readonly bottomPadding: boolean;
+
+  constructor(
+    component: Component,
+    theme: Theme,
+    options: {
+      background: CodexToolBackgroundResolver;
+      status: CodexToolSurfaceStatus;
+      top: boolean;
+      bottom: boolean;
+    },
+  ) {
+    this.component = component;
+    this.theme = theme;
+    this.resolveBackground = options.background;
+    this.status = options.status;
+    this.topPadding = options.top;
+    this.bottomPadding = options.bottom;
+  }
+
+  render(width: number): string[] {
+    const effectiveWidth = Math.max(1, width);
+    const background = surfaceBackground(this.theme, this.resolveBackground(), this.status);
+    const horizontalPadding = effectiveWidth > 2 ? 1 : 0;
+    const contentWidth = Math.max(1, effectiveWidth - horizontalPadding * 2);
+    const lines = this.component
+      .render(contentWidth)
+      .map((line) =>
+        fillLine(`${" ".repeat(horizontalPadding)}${line}`, effectiveWidth, background),
+      );
+    const blankLine = fillLine("", effectiveWidth, background);
+    if (this.topPadding) lines.unshift(blankLine);
+    if (this.bottomPadding) lines.push(blankLine);
+    return lines;
+  }
+
+  invalidate(): void {
+    this.component.invalidate();
+  }
+}