pi-openai-codex-compat 0.0.2 → 0.0.3

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

@@ -34,8 +34,15 @@ import {
   type JsonRecord,
   type ResponsesItem,
 } from "./codex-protocol.ts";
+import { codexCacheKey } from "./codex-cache-key.ts";
 import { processCodexStream } from "./codex-stream.ts";
-import { CodexTransport } from "./codex-transport.ts";
+import {
+  CodexTransport,
+  validateCodexAuthentication,
+  type CodexContinuationHandle,
+  type CodexTransportDiagnostic,
+  type CodexWebSocketResponseHandle,
+} from "./codex-transport.ts";
 import type { CodexCompatConfig, ImageDetail } from "./config.ts";
 import { nativeResponseData, NATIVE_RESPONSE_ENTRY_TYPE } from "./native-history.ts";
 import {
@@ -49,10 +56,10 @@ import {
   createGrammarToolInputProperties,
   type ResponsesItem as SerializedResponsesItem,
 } from "./vendor/pi-ai/openai-responses-serialization.ts";
+import { formatProviderError } from "./provider-error.ts";
 
 const CODEX_PROVIDER = "openai-codex";
 const CODEX_API = "openai-codex-responses";
-const CHECKPOINT_STATUS_ID = "openai-codex-compat-compaction";
 const CODEX_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode"]);
 
 type ConfigResolver = (ctx: ExtensionContext) => CodexCompatConfig;
@@ -78,7 +85,6 @@ type RuntimeScope = {
   config: CodexCompatConfig;
   hasUI: boolean;
   notify(message: string, level: "info" | "warning" | "error"): void;
-  setStatus(message: string | undefined): void;
 };
 
 type RequestTemplate = {
@@ -94,11 +100,6 @@ type CodexCompat = {
   supportsOpenAIGrammarTools?: boolean;
 };
 
-function clampPromptCacheKey(key: string | undefined): string | undefined {
-  if (key === undefined) return undefined;
-  return Array.from(key).slice(0, 64).join("");
-}
-
 function markerSummary(): string {
   return `OpenAI Codex remote compaction checkpoint (${randomUUID()}).`;
 }
@@ -153,7 +154,6 @@ function nativeOverrideRequired(
 function captureRawEvents(
   events: AsyncIterable<JsonRecord>,
   items: ResponsesItem[],
-  metadata: { serviceTier?: string },
 ): AsyncIterable<JsonRecord> {
   return {
     async *[Symbol.asyncIterator]() {
@@ -171,12 +171,23 @@ function captureRawEvents(
             items.splice(0, items.length, ...terminalItems.map((item) => structuredClone(item)));
           }
         }
-        if (
-          (event.type === "response.completed" || event.type === "response.incomplete") &&
-          isObject(event.response) &&
-          typeof event.response.service_tier === "string"
-        ) {
-          metadata.serviceTier = event.response.service_tier;
+        yield event;
+      }
+    },
+  };
+}
+
+function startOnFirstEvent(
+  events: AsyncIterable<JsonRecord>,
+  onStart: () => void,
+): AsyncIterable<JsonRecord> {
+  return {
+    async *[Symbol.asyncIterator]() {
+      let started = false;
+      for await (const event of events) {
+        if (!started) {
+          started = true;
+          onStart();
         }
         yield event;
       }
@@ -184,6 +195,13 @@ function captureRawEvents(
   };
 }
 
+function clearStreamingScratchState(message: AssistantMessage): void {
+  for (const block of message.content) {
+    delete (block as { partialJson?: string }).partialJson;
+    delete (block as { customInput?: unknown }).customInput;
+  }
+}
+
 function updateInput(payload: JsonRecord, input: readonly ResponsesItem[]): JsonRecord {
   const result: JsonRecord = {
     ...payload,
@@ -194,14 +212,15 @@ function updateInput(payload: JsonRecord, input: readonly ResponsesItem[]): Json
   return result;
 }
 
-function successfulStopReason(
+function assertSuccessfulOutput(
   message: AssistantMessage,
-): message is AssistantMessage & { stopReason: "stop" | "length" | "toolUse" } {
-  return (
-    message.stopReason === "stop" ||
-    message.stopReason === "length" ||
-    message.stopReason === "toolUse"
-  );
+): asserts message is AssistantMessage & { stopReason: "stop" | "length" | "toolUse" } {
+  if (message.stopReason === "pending") {
+    throw new Error("Codex stream ended without a stop reason");
+  }
+  if (message.stopReason === "error" || message.stopReason === "aborted") {
+    throw new Error(message.errorMessage || "An unknown error occurred");
+  }
 }
 
 function applyServiceTierPricing(
@@ -308,7 +327,6 @@ export class CodexProviderRuntime {
       config: this.resolveConfig(ctx),
       hasUI: ctx.hasUI,
       notify: (message, level) => ctx.ui.notify(message, level),
-      setStatus: (message) => ctx.ui.setStatus(CHECKPOINT_STATUS_ID, message),
     });
   }
 
@@ -319,7 +337,10 @@ export class CodexProviderRuntime {
     this.transport.close(sessionId);
   }
 
-  private async acquireRequest(sessionId: string | undefined): Promise<() => void> {
+  private async acquireRequest(
+    sessionId: string | undefined,
+    signal?: AbortSignal,
+  ): Promise<() => void> {
     if (!sessionId) return () => {};
     const previous = this.requestTails.get(sessionId) ?? Promise.resolve();
     let releaseCurrent!: () => void;
@@ -327,11 +348,44 @@ export class CodexProviderRuntime {
       releaseCurrent = resolve;
     });
     this.requestTails.set(sessionId, current);
-    await previous;
-    return () => {
+    let released = false;
+    const release = () => {
+      if (released) return;
+      released = true;
       releaseCurrent();
       if (this.requestTails.get(sessionId) === current) this.requestTails.delete(sessionId);
     };
+
+    if (signal?.aborted) {
+      void previous.then(release);
+      throw new Error("Request was aborted");
+    }
+
+    let onAbort: (() => void) | undefined;
+    try {
+      await Promise.race([
+        previous,
+        ...(signal
+          ? [
+              new Promise<never>((_resolve, reject) => {
+                onAbort = () => reject(new Error("Request was aborted"));
+                signal.addEventListener("abort", onAbort, { once: true });
+              }),
+            ]
+          : []),
+      ]);
+    } catch (error) {
+      void previous.then(release);
+      throw error;
+    } finally {
+      if (onAbort) signal?.removeEventListener("abort", onAbort);
+    }
+
+    if (signal?.aborted) {
+      release();
+      throw new Error("Request was aborted");
+    }
+    return release;
   }
 
   createProvider(base: Provider): Provider {
@@ -381,7 +435,8 @@ export class CodexProviderRuntime {
     model: Model<any>,
     context: Context,
     options: OpenAICodexResponsesOptions,
-    sessionId: string | undefined,
+    runtimeSessionId: string | undefined,
+    cacheSessionId: string | undefined,
     grammarToolInputProperties: GrammarToolInputProperties,
   ): JsonRecord {
     const compat = model.compat as CodexCompat | undefined;
@@ -393,12 +448,12 @@ export class CodexProviderRuntime {
       instructions: context.systemPrompt || "You are a helpful assistant.",
       input: this.wireHistory(
         model,
-        Object.assign({}, context, { sessionId }),
+        Object.assign({}, context, { sessionId: runtimeSessionId }),
         grammarToolInputProperties,
       ),
       text: { verbosity: options.textVerbosity ?? "low" },
       include: ["reasoning.encrypted_content"],
-      prompt_cache_key: sessionId,
+      prompt_cache_key: cacheSessionId,
       tool_choice: options.toolChoice ?? "auto",
       parallel_tool_calls: true,
     };
@@ -439,21 +494,38 @@ export class CodexProviderRuntime {
   }): Promise<{ checkpoint: CheckpointData; usage?: Usage }> {
     const sessionId = options.requestOptions.sessionId;
     if (!sessionId) throw new Error("Codex compaction requires a Pi session id.");
+    const accountId = validateCodexAuthentication(options.model, options.requestOptions.apiKey);
     const payload = remoteCompactionPayload({
       template: options.template,
       modelId: options.model.id,
       history: options.history,
       instructions: options.instructions,
-      sessionId,
+      sessionId:
+        options.requestOptions.cacheRetention === "none" ? undefined : codexCacheKey(sessionId),
       priority: options.priority,
     });
     const transformed = await options.requestOptions.onPayload?.(payload, options.model);
-    const request = isObject(transformed) ? transformed : payload;
-    const compacted = await collectRemoteCompaction(
-      this.transport.request(options.model, request, options.requestOptions),
-      options.model,
-      options.priority,
-    );
+    const request = transformed === undefined ? payload : (transformed as JsonRecord);
+    let webSocketResponseHandle: CodexWebSocketResponseHandle | undefined;
+    let compacted: Awaited<ReturnType<typeof collectRemoteCompaction>>;
+    try {
+      compacted = await collectRemoteCompaction(
+        this.transport.request(options.model, request, {
+          ...options.requestOptions,
+          accountId,
+          onWebSocketResponseHandle(handle) {
+            webSocketResponseHandle = handle;
+          },
+        }),
+        options.model,
+        options.priority,
+      );
+    } catch (error) {
+      if (!options.requestOptions.signal?.aborted) {
+        webSocketResponseHandle?.failParsing(error);
+      }
+      throw error;
+    }
     return {
       checkpoint: checkpointData(
         options.model.id,
@@ -515,44 +587,39 @@ export class CodexProviderRuntime {
       return body;
     }
 
-    scope.setStatus("Codex compacting…");
-    try {
-      const compacted = await this.performCompaction({
-        model,
-        requestOptions: options,
-        history: split.history,
-        postCompactionTail: split.tail,
-        template: withoutConversationInput(body),
-        instructions:
-          typeof body.instructions === "string"
-            ? body.instructions
-            : context.systemPrompt || "You are a helpful assistant.",
-        grammarToolInputProperties,
-        priority: scope.config.fastMode,
-      });
-      const firstKeptEntryId = userEntryAfterLastSampled(branch)?.id ?? scope.manager.getLeafId();
-      if (!firstKeptEntryId || typeof scope.manager.appendCompaction !== "function") {
-        throw new Error("Pi's mutable SessionManager is unavailable for percentage compaction.");
-      }
-      if (scope.manager.getLeafId() !== scope.leafId) {
-        throw new Error("Pi's active session branch changed while Codex was compacting.");
-      }
-      scope.manager.appendCompaction(
-        markerSummary(),
-        firstKeptEntryId,
-        scope.contextTokens ?? 0,
-        compacted.checkpoint,
-        true,
-        compacted.usage,
-      );
-      scope.notify(
-        `OpenAI Codex context compacted at ${scope.contextPercent.toFixed(1)}% and will continue.`,
-        "info",
-      );
-      return updateInput(body, compacted.checkpoint.history);
-    } finally {
-      scope.setStatus(undefined);
+    const compacted = await this.performCompaction({
+      model,
+      requestOptions: options,
+      history: split.history,
+      postCompactionTail: split.tail,
+      template: withoutConversationInput(body),
+      instructions:
+        typeof body.instructions === "string"
+          ? body.instructions
+          : context.systemPrompt || "You are a helpful assistant.",
+      grammarToolInputProperties,
+      priority: scope.config.fastMode,
+    });
+    const firstKeptEntryId = userEntryAfterLastSampled(branch)?.id ?? scope.manager.getLeafId();
+    if (!firstKeptEntryId || typeof scope.manager.appendCompaction !== "function") {
+      throw new Error("Pi's mutable SessionManager is unavailable for percentage compaction.");
     }
+    if (scope.manager.getLeafId() !== scope.leafId) {
+      throw new Error("Pi's active session branch changed while Codex was compacting.");
+    }
+    scope.manager.appendCompaction(
+      markerSummary(),
+      firstKeptEntryId,
+      scope.contextTokens ?? 0,
+      compacted.checkpoint,
+      true,
+      compacted.usage,
+    );
+    scope.notify(
+      `OpenAI Codex context compacted at ${scope.contextPercent.toFixed(1)}% and will continue.`,
+      "info",
+    );
+    return updateInput(body, compacted.checkpoint.history);
   }
 
   stream(
@@ -581,9 +648,12 @@ export class CodexProviderRuntime {
         timestamp: Date.now(),
       };
       const runtimeSessionId = requestOptions.sessionId;
-      const releaseRequest = await this.acquireRequest(runtimeSessionId);
+      let releaseRequest = () => {};
       try {
-        const cacheSessionId = clampPromptCacheKey(runtimeSessionId);
+        const accountId = validateCodexAuthentication(model, requestOptions.apiKey);
+        releaseRequest = await this.acquireRequest(runtimeSessionId, requestOptions.signal);
+        const cacheSessionId =
+          requestOptions.cacheRetention === "none" ? undefined : codexCacheKey(runtimeSessionId);
         const grammarToolInputProperties = createGrammarToolInputProperties(
           context.tools,
           (model.compat as CodexCompat | undefined)?.supportsOpenAIGrammarTools ?? false,
@@ -592,11 +662,12 @@ export class CodexProviderRuntime {
           model,
           context,
           requestOptions,
+          runtimeSessionId,
           cacheSessionId,
           grammarToolInputProperties,
         );
         const transformed = await requestOptions.onPayload?.(body, model);
-        if (isObject(transformed)) body = transformed;
+        if (transformed !== undefined) body = transformed as JsonRecord;
         if (runtimeSessionId) {
           this.templates.set(runtimeSessionId, {
             modelId: model.id,
@@ -614,28 +685,60 @@ export class CodexProviderRuntime {
         );
 
         const rawItems: ResponsesItem[] = [];
-        const responseMetadata: { serviceTier?: string } = {};
-        stream.push({ type: "start", partial: output });
-        await processCodexStream(
-          captureRawEvents(
-            this.transport.request(model, body, requestOptions),
-            rawItems,
-            responseMetadata,
-          ),
-          output,
-          stream,
-          model,
-          grammarToolInputProperties,
-        );
-        if (!successfulStopReason(output)) {
-          throw new Error(output.errorMessage || "Codex stream ended without a successful stop.");
+        let continuationHandle: CodexContinuationHandle | undefined;
+        let webSocketResponseHandle: CodexWebSocketResponseHandle | undefined;
+        let startEmitted = false;
+        const emitStart = () => {
+          if (startEmitted) return;
+          startEmitted = true;
+          stream.push({ type: "start", partial: output });
+        };
+        const transportRequestOptions = {
+          ...requestOptions,
+          accountId,
+          onContinuationReady(handle: CodexContinuationHandle) {
+            continuationHandle = handle;
+          },
+          onWebSocketResponseHandle(handle: CodexWebSocketResponseHandle) {
+            webSocketResponseHandle = handle;
+          },
+          onTransportStart: emitStart,
+          onTransportDiagnostic(diagnostic: CodexTransportDiagnostic) {
+            output.diagnostics = [...(output.diagnostics ?? []), diagnostic];
+          },
+        };
+        try {
+          await processCodexStream(
+            startOnFirstEvent(
+              captureRawEvents(
+                this.transport.request(model, body, transportRequestOptions),
+                rawItems,
+              ),
+              emitStart,
+            ),
+            output,
+            stream,
+            model,
+            grammarToolInputProperties,
+            {
+              applyServiceTierPricing(usage, responseServiceTier) {
+                applyServiceTierPricing(usage, model, body.service_tier, responseServiceTier);
+              },
+            },
+          );
+        } catch (error) {
+          if (!requestOptions.signal?.aborted) {
+            webSocketResponseHandle?.failParsing(error);
+          }
+          throw error;
+        }
+        if (requestOptions.signal?.aborted) throw new Error("Request was aborted");
+        try {
+          assertSuccessfulOutput(output);
+        } catch (error) {
+          webSocketResponseHandle?.discard();
+          throw error;
         }
-        applyServiceTierPricing(
-          output.usage,
-          model,
-          body.service_tier,
-          responseMetadata.serviceTier,
-        );
 
         const compat = model.compat as CodexCompat | undefined;
         const canonicalContext: Context = {
@@ -667,19 +770,26 @@ export class CodexProviderRuntime {
           (item) =>
             item["type"] !== "function_call_output" && item["type"] !== "custom_tool_call_output",
         );
-        if (rawItems.length > 0 && nativeOverrideRequired(rawItems, canonicalItems)) {
+        const persistNativeItems =
+          rawItems.length > 0 && nativeOverrideRequired(rawItems, canonicalItems);
+        if (persistNativeItems) {
           if (!output.responseId) throw new Error("Codex response is missing a response id.");
           this.pi.appendEntry(
             NATIVE_RESPONSE_ENTRY_TYPE,
             nativeResponseData(model.id, output.responseId, rawItems),
           );
         }
+        const readyContinuation = continuationHandle;
+        if (readyContinuation && readyContinuation.responseId === output.responseId) {
+          readyContinuation.replaceResponseItems(persistNativeItems ? rawItems : canonicalItems);
+        }
 
         stream.push({ type: "done", reason: output.stopReason, message: output });
         stream.end();
       } catch (error) {
+        clearStreamingScratchState(output);
         output.stopReason = requestOptions.signal?.aborted ? "aborted" : "error";
-        output.errorMessage = error instanceof Error ? error.message : String(error);
+        output.errorMessage = formatProviderError(error);
         stream.push({ type: "error", reason: output.stopReason, error: output });
         stream.end();
       } finally {
@@ -694,6 +804,7 @@ export class CodexProviderRuntime {
     context: Context,
     options?: SimpleStreamOptions,
   ): AssistantMessageEventStream {
+    if (!options?.apiKey) throw new Error(`No API key for provider: ${model.provider}`);
     const effort = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
     return this.stream(model, context, {
       ...options,
@@ -705,7 +816,7 @@ export class CodexProviderRuntime {
     return this.templates.get(sessionId);
   }
 
-  compact(options: {
+  async compact(options: {
     model: Model<any>;
     requestOptions: OpenAICodexResponsesOptions;
     history: ResponsesItem[];
@@ -714,13 +825,16 @@ export class CodexProviderRuntime {
     template: JsonRecord;
     priority: boolean;
   }): Promise<{ checkpoint: CheckpointData; usage?: Usage }> {
-    return this.acquireRequest(options.requestOptions.sessionId).then(async (release) => {
-      try {
-        return await this.performCompaction(options);
-      } finally {
-        release();
-      }
-    });
+    validateCodexAuthentication(options.model, options.requestOptions.apiKey);
+    const release = await this.acquireRequest(
+      options.requestOptions.sessionId,
+      options.requestOptions.signal,
+    );
+    try {
+      return await this.performCompaction(options);
+    } finally {
+      release();
+    }
   }
 }