pi-openai-codex-compat 0.0.2 → 0.0.4

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

@@ -17,7 +17,12 @@ import {
   type SimpleStreamOptions,
   type Tool,
   type Usage,
+  uuidv7,
 } from "@earendil-works/pi-ai";
+import {
+  codexCacheDiagnosticContext,
+  type CodexCacheDiagnosticContext,
+} from "./codex-cache-diagnostics.ts";
 import {
   checkpointData,
   providerHistory,
@@ -34,9 +39,26 @@ import {
   type JsonRecord,
   type ResponsesItem,
 } from "./codex-protocol.ts";
-import { processCodexStream } from "./codex-stream.ts";
-import { CodexTransport } from "./codex-transport.ts";
-import type { CodexCompatConfig, ImageDetail } from "./config.ts";
+import { codexCacheKey } from "./codex-cache-key.ts";
+import { resolveCodexInstallationId } from "./codex-installation.ts";
+import {
+  type CodexCompactionMetadata,
+  type CodexMetadataIdentity,
+  responsesCompactionV2Metadata,
+  withCodexRequestMetadata,
+} from "./codex-metadata.ts";
+import { resolveCodexThreadIdentity, type CodexThreadIdentity } from "./codex-thread-lineage.ts";
+import { applyResponsesLite } from "./responses-lite.ts";
+import { processCodexStream, type CodexStreamAttemptState } from "./codex-stream.ts";
+import {
+  CodexTransport,
+  CodexTurnState,
+  validateCodexAuthentication,
+  type CodexContinuationHandle,
+  type CodexTransportDiagnostic,
+  type CodexWebSocketResponseHandle,
+} from "./codex-transport.ts";
+import { DEFAULT_CONFIG, type CodexCompatConfig, type ImageDetail } from "./config.ts";
 import { nativeResponseData, NATIVE_RESPONSE_ENTRY_TYPE } from "./native-history.ts";
 import {
   CODEX_NAMESPACED_TOOL_NAMES,
@@ -49,10 +71,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 +100,6 @@ type RuntimeScope = {
   config: CodexCompatConfig;
   hasUI: boolean;
   notify(message: string, level: "info" | "warning" | "error"): void;
-  setStatus(message: string | undefined): void;
 };
 
 type RequestTemplate = {
@@ -88,16 +109,32 @@ type RequestTemplate = {
   requestOptions: OpenAICodexResponsesOptions;
 };
 
+type ActiveAgentTurn = {
+  turnId: string;
+  startedAtUnixMs: number;
+  turnState: CodexTurnState;
+};
+
 type CodexCompat = {
   supportsToolSearch?: boolean;
   supportsStrictMode?: boolean;
   supportsOpenAIGrammarTools?: boolean;
 };
 
-function clampPromptCacheKey(key: string | undefined): string | undefined {
-  if (key === undefined) return undefined;
-  return Array.from(key).slice(0, 64).join("");
-}
+type CodexTerminalState = {
+  type?: "response.completed" | "response.incomplete" | "response.failed";
+  response?: JsonRecord;
+};
+
+export type CodexResponseRetryPolicy = {
+  maxRetries: number;
+  baseDelayMs: number;
+};
+
+const DEFAULT_RESPONSE_RETRY_POLICY: CodexResponseRetryPolicy = {
+  maxRetries: 5,
+  baseDelayMs: 200,
+};
 
 function markerSummary(): string {
   return `OpenAI Codex remote compaction checkpoint (${randomUUID()}).`;
@@ -153,7 +190,7 @@ function nativeOverrideRequired(
 function captureRawEvents(
   events: AsyncIterable<JsonRecord>,
   items: ResponsesItem[],
-  metadata: { serviceTier?: string },
+  terminalState?: CodexTerminalState,
 ): AsyncIterable<JsonRecord> {
   return {
     async *[Symbol.asyncIterator]() {
@@ -162,7 +199,9 @@ function captureRawEvents(
           items.push(structuredClone(event.item));
         }
         if (
-          (event.type === "response.completed" || event.type === "response.incomplete") &&
+          (event.type === "response.completed" ||
+            event.type === "response.incomplete" ||
+            event.type === "response.failed") &&
           isObject(event.response) &&
           Array.isArray(event.response["output"])
         ) {
@@ -172,11 +211,35 @@ function captureRawEvents(
           }
         }
         if (
-          (event.type === "response.completed" || event.type === "response.incomplete") &&
-          isObject(event.response) &&
-          typeof event.response.service_tier === "string"
+          terminalState &&
+          (event.type === "response.completed" ||
+            event.type === "response.incomplete" ||
+            event.type === "response.failed")
         ) {
-          metadata.serviceTier = event.response.service_tier;
+          terminalState.type = event.type;
+          if (isObject(event.response)) {
+            terminalState.response = structuredClone(event.response);
+          } else {
+            delete terminalState.response;
+          }
+        }
+        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 +247,99 @@ 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 emptyUsage(): Usage {
+  return {
+    input: 0,
+    output: 0,
+    cacheRead: 0,
+    cacheWrite: 0,
+    totalTokens: 0,
+    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+  };
+}
+
+function discardIncompleteAttemptContent(
+  message: AssistantMessage,
+  attempt: CodexStreamAttemptState,
+): void {
+  const incomplete = [...attempt.startedContentIndexes]
+    .filter((index) => !attempt.completedContentIndexes.has(index))
+    .sort((left, right) => right - left);
+  for (const index of incomplete) message.content.splice(index, 1);
+}
+
+function accumulateUsage(previous: Usage, current: Usage): Usage {
+  const previousReasoning = previous.reasoning;
+  const currentReasoning = current.reasoning;
+  return {
+    input: previous.input + current.input,
+    output: previous.output + current.output,
+    cacheRead: previous.cacheRead + current.cacheRead,
+    cacheWrite: previous.cacheWrite + current.cacheWrite,
+    ...(previousReasoning === undefined && currentReasoning === undefined
+      ? {}
+      : { reasoning: (previousReasoning ?? 0) + (currentReasoning ?? 0) }),
+    totalTokens: previous.totalTokens + current.totalTokens,
+    cost: {
+      input: previous.cost.input + current.cost.input,
+      output: previous.cost.output + current.cost.output,
+      cacheRead: previous.cost.cacheRead + current.cost.cacheRead,
+      cacheWrite: previous.cost.cacheWrite + current.cost.cacheWrite,
+      total: previous.cost.total + current.cost.total,
+    },
+  };
+}
+
+function retryableResponseFailure(response: JsonRecord | undefined): boolean {
+  const error = isObject(response?.["error"]) ? response["error"] : undefined;
+  const code = typeof error?.["code"] === "string" ? error["code"].toLowerCase() : "";
+  return !(
+    code === "context_length_exceeded" ||
+    code === "insufficient_quota" ||
+    code === "usage_not_included" ||
+    code === "cyber_policy" ||
+    code === "invalid_prompt" ||
+    code === "bio_policy"
+  );
+}
+
+function responseRetryDelayMs(baseDelayMs: number, attempt: number): number {
+  if (baseDelayMs <= 0) return 0;
+  const exponential = baseDelayMs * 2 ** Math.max(0, attempt - 1);
+  return Math.floor(exponential * (0.9 + Math.random() * 0.2));
+}
+
+function waitForResponseRetry(milliseconds: number, signal?: AbortSignal): Promise<void> {
+  if (signal?.aborted) return Promise.reject(new Error("Request was aborted"));
+  if (milliseconds <= 0) return Promise.resolve();
+  return new Promise((resolve, reject) => {
+    const onAbort = () => {
+      clearTimeout(timer);
+      reject(new Error("Request was aborted"));
+    };
+    const timer = setTimeout(() => {
+      signal?.removeEventListener("abort", onAbort);
+      resolve();
+    }, milliseconds);
+    signal?.addEventListener("abort", onAbort, { once: true });
+  });
+}
+
+function continueResponseBody(
+  body: JsonRecord,
+  responseItems: readonly ResponsesItem[],
+): JsonRecord | undefined {
+  if (!Array.isArray(body.input) || !body.input.every(isResponsesItem)) return undefined;
+  return updateInput(body, [...body.input, ...responseItems]);
+}
+
 function updateInput(payload: JsonRecord, input: readonly ResponsesItem[]): JsonRecord {
   const result: JsonRecord = {
     ...payload,
@@ -194,14 +350,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(
@@ -286,13 +443,38 @@ export class CodexProviderRuntime {
   readonly transport = new CodexTransport();
   private readonly scopes = new Map<string, RuntimeScope>();
   private readonly templates = new Map<string, RequestTemplate>();
+  private readonly prewarmedTemplates = new Set<string>();
   private readonly requestTails = new Map<string, Promise<void>>();
+  private readonly activeAgentTurns = new Map<string, ActiveAgentTurn>();
+  private readonly windowNumbers = new Map<string, number>();
+  private readonly activeThreadIds = new Map<string, string>();
   private readonly pi: ExtensionAPI;
   private readonly resolveConfig: ConfigResolver;
+  private readonly installationId: string;
+  private readonly responseRetryPolicy: CodexResponseRetryPolicy;
 
-  constructor(pi: ExtensionAPI, resolveConfig: ConfigResolver) {
+  constructor(
+    pi: ExtensionAPI,
+    resolveConfig: ConfigResolver,
+    installationId: string = randomUUID(),
+    responseRetryPolicy: Partial<CodexResponseRetryPolicy> = {},
+  ) {
     this.pi = pi;
     this.resolveConfig = resolveConfig;
+    this.installationId = installationId;
+    const maxRetries = responseRetryPolicy.maxRetries ?? DEFAULT_RESPONSE_RETRY_POLICY.maxRetries;
+    const baseDelayMs =
+      responseRetryPolicy.baseDelayMs ?? DEFAULT_RESPONSE_RETRY_POLICY.baseDelayMs;
+    if (!Number.isFinite(maxRetries) || maxRetries < 0) {
+      throw new Error(`Invalid Codex response maxRetries: ${String(maxRetries)}`);
+    }
+    if (!Number.isFinite(baseDelayMs) || baseDelayMs < 0) {
+      throw new Error(`Invalid Codex response baseDelayMs: ${String(baseDelayMs)}`);
+    }
+    this.responseRetryPolicy = {
+      maxRetries: Math.floor(maxRetries),
+      baseDelayMs: Math.floor(baseDelayMs),
+    };
   }
 
   captureScope(ctx: ExtensionContext): void {
@@ -308,18 +490,161 @@ 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),
     });
   }
 
+  beginAgentTurn(ctx: ExtensionContext): void {
+    this.activeAgentTurns.set(ctx.sessionManager.getSessionId(), {
+      turnId: uuidv7(),
+      startedAtUnixMs: Date.now(),
+      turnState: new CodexTurnState(),
+    });
+  }
+
+  endAgentTurn(ctx: ExtensionContext): void {
+    this.activeAgentTurns.delete(ctx.sessionManager.getSessionId());
+  }
+
+  updateSessionConfig(sessionId: string, config: CodexCompatConfig): void {
+    const scope = this.scopes.get(sessionId);
+    if (scope) scope.config = config;
+  }
+
+  private responsesLiteEnabled(sessionId: string | undefined): boolean {
+    return (
+      (sessionId ? this.scopes.get(sessionId)?.config.responsesLite : undefined) ??
+      DEFAULT_CONFIG.responsesLite
+    );
+  }
+
+  private agentTurn(sessionId: string | undefined): ActiveAgentTurn {
+    return (
+      (sessionId ? this.activeAgentTurns.get(sessionId) : undefined) ?? {
+        turnId: uuidv7(),
+        startedAtUnixMs: Date.now(),
+        turnState: new CodexTurnState(),
+      }
+    );
+  }
+
+  private metadataIdentity(
+    metadataSessionId: string | undefined,
+    turn?: ActiveAgentTurn,
+    runtimeSessionId = metadataSessionId,
+  ): CodexMetadataIdentity {
+    const thread: Partial<CodexThreadIdentity> = runtimeSessionId
+      ? this.threadIdentity(runtimeSessionId)
+      : metadataSessionId
+        ? { threadId: metadataSessionId }
+        : {};
+    const windowKey =
+      runtimeSessionId && thread.threadId ? `${runtimeSessionId}\0${thread.threadId}` : undefined;
+    return {
+      installationId: this.installationId,
+      ...(thread.threadId ? { threadId: thread.threadId } : {}),
+      ...(thread.forkedFromThreadId ? { forkedFromThreadId: thread.forkedFromThreadId } : {}),
+      windowNumber: windowKey ? (this.windowNumbers.get(windowKey) ?? 0) : 0,
+      ...(turn ? { turnStartedAtUnixMs: turn.startedAtUnixMs } : {}),
+      threadSource: "user",
+      // Pi extensions execute without Codex's platform sandbox.
+      sandbox: "none",
+    };
+  }
+
+  private threadIdentity(sessionId: string): CodexThreadIdentity {
+    const branch = this.scopes.get(sessionId)?.manager.getBranch() as SessionEntry[] | undefined;
+    return resolveCodexThreadIdentity(sessionId, branch ?? []);
+  }
+
+  private clearPrewarmState(sessionId: string): void {
+    for (const key of this.prewarmedTemplates) {
+      if (key.startsWith(`${sessionId}\0`)) this.prewarmedTemplates.delete(key);
+    }
+  }
+
+  private activateThread(sessionId: string | undefined): void {
+    if (!sessionId) return;
+    const threadId = this.threadIdentity(sessionId).threadId;
+    const previous = this.activeThreadIds.get(sessionId);
+    if (previous && previous !== threadId) {
+      this.transport.close(sessionId);
+      this.clearPrewarmState(sessionId);
+    }
+    this.activeThreadIds.set(sessionId, threadId);
+  }
+
+  private advanceWindow(sessionId: string): void {
+    const threadId = this.threadIdentity(sessionId).threadId;
+    const key = `${sessionId}\0${threadId}`;
+    this.windowNumbers.set(key, (this.windowNumbers.get(key) ?? 0) + 1);
+  }
+
   clearSession(sessionId: string): void {
     this.scopes.delete(sessionId);
     this.templates.delete(sessionId);
+    this.clearPrewarmState(sessionId);
     this.requestTails.delete(sessionId);
+    this.activeAgentTurns.delete(sessionId);
+    for (const key of this.windowNumbers.keys()) {
+      if (key.startsWith(`${sessionId}\0`)) this.windowNumbers.delete(key);
+    }
+    this.activeThreadIds.delete(sessionId);
     this.transport.close(sessionId);
   }
 
-  private async acquireRequest(sessionId: string | undefined): Promise<() => void> {
+  private async maybePrewarm(options: {
+    model: Model<any>;
+    body: JsonRecord;
+    fullBody: JsonRecord;
+    requestOptions: OpenAICodexResponsesOptions;
+    accountId: string;
+    diagnostics: CodexTransportDiagnostic[];
+    turnState: CodexTurnState;
+    cacheDiagnostics: CodexCacheDiagnosticContext;
+  }): Promise<void> {
+    const sessionId = options.requestOptions.sessionId;
+    if (
+      !sessionId ||
+      options.requestOptions.cacheRetention === "none" ||
+      options.requestOptions.transport === "sse" ||
+      !Array.isArray(options.body.input) ||
+      !Array.isArray(options.fullBody.input)
+    ) {
+      return;
+    }
+    const key = `${sessionId}\0${options.model.id}\0${options.cacheDiagnostics.envelope}`;
+    if (this.prewarmedTemplates.has(key)) return;
+    this.prewarmedTemplates.add(key);
+
+    const cacheSessionId = codexCacheKey(sessionId);
+    const prewarmBody = withCodexRequestMetadata(
+      options.body,
+      cacheSessionId,
+      { kind: "prewarm" },
+      "",
+      this.metadataIdentity(cacheSessionId, undefined, sessionId),
+    );
+    try {
+      await this.transport.prewarm(options.model, prewarmBody, {
+        ...options.requestOptions,
+        accountId: options.accountId,
+        turnState: options.turnState,
+        cacheDiagnostics: options.cacheDiagnostics,
+        requestKind: "prewarm",
+        onTransportDiagnostic(diagnostic) {
+          options.diagnostics.push(diagnostic);
+        },
+      });
+    } catch {
+      // Warmup is best-effort. The transport has already activated sticky SSE
+      // after exhausting its WebSocket retry budget.
+    }
+  }
+
+  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 +652,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 {
@@ -358,7 +716,7 @@ export class CodexProviderRuntime {
         grammarToolInputProperties,
         deferredTools: splitDeferredTools(context, Boolean(compat?.supportsToolSearch)).deferred,
         toolOptions: {
-          strict: null,
+          strict: false,
           supportsStrictMode: compat?.supportsStrictMode ?? true,
           supportsOpenAIGrammarTools: compat?.supportsOpenAIGrammarTools ?? false,
         },
@@ -381,32 +739,45 @@ export class CodexProviderRuntime {
     model: Model<any>,
     context: Context,
     options: OpenAICodexResponsesOptions,
-    sessionId: string | undefined,
+    runtimeSessionId: string | undefined,
+    cacheSessionId: string | undefined,
     grammarToolInputProperties: GrammarToolInputProperties,
+    turnId: string,
   ): JsonRecord {
     const compat = model.compat as CodexCompat | undefined;
     const toolPlacement = splitDeferredTools(context, Boolean(compat?.supportsToolSearch));
-    const body: JsonRecord = {
+    let body: JsonRecord = {
       model: model.id,
       store: false,
       stream: true,
       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,
+      tools: [],
     };
-    if (options.temperature !== undefined) body["temperature"] = options.temperature;
+    body = withCodexRequestMetadata(
+      body,
+      cacheSessionId,
+      { kind: "turn" },
+      turnId,
+      this.metadataIdentity(
+        cacheSessionId,
+        this.activeAgentTurns.get(runtimeSessionId ?? ""),
+        runtimeSessionId,
+      ),
+    );
     if (options.serviceTier !== undefined) body.service_tier = options.serviceTier;
     if (toolPlacement.immediate.length > 0) {
       body.tools = convertResponsesTools(toolPlacement.immediate, {
-        strict: null,
+        strict: false,
         supportsStrictMode: compat?.supportsStrictMode ?? true,
         supportsOpenAIGrammarTools: compat?.supportsOpenAIGrammarTools ?? false,
         namespacedToolNames: CODEX_NAMESPACED_TOOL_NAMES,
@@ -436,24 +807,74 @@ export class CodexProviderRuntime {
     instructions: string;
     grammarToolInputProperties: GrammarToolInputProperties;
     priority: boolean;
+    compactionMetadata: CodexCompactionMetadata;
+    agentTurn?: ActiveAgentTurn;
+    responsesLiteEnabled?: boolean;
   }): Promise<{ checkpoint: CheckpointData; usage?: Usage }> {
     const sessionId = options.requestOptions.sessionId;
     if (!sessionId) throw new Error("Codex compaction requires a Pi session id.");
-    const payload = remoteCompactionPayload({
-      template: options.template,
-      modelId: options.model.id,
-      history: options.history,
-      instructions: options.instructions,
-      sessionId,
-      priority: options.priority,
-    });
+    this.activateThread(sessionId);
+    const agentTurn = options.agentTurn ?? this.agentTurn(sessionId);
+    const responsesLiteEnabled =
+      options.responsesLiteEnabled ?? this.responsesLiteEnabled(sessionId);
+    const accountId = validateCodexAuthentication(options.model, options.requestOptions.apiKey);
+    const payload = withCodexRequestMetadata(
+      remoteCompactionPayload({
+        template: options.template,
+        modelId: options.model.id,
+        history: options.history,
+        instructions: options.instructions,
+        sessionId:
+          options.requestOptions.cacheRetention === "none" ? undefined : codexCacheKey(sessionId),
+        priority: options.priority,
+      }),
+      options.requestOptions.cacheRetention === "none" ? undefined : codexCacheKey(sessionId),
+      { kind: "compaction", compaction: options.compactionMetadata },
+      agentTurn.turnId,
+      this.metadataIdentity(
+        options.requestOptions.cacheRetention === "none" ? undefined : codexCacheKey(sessionId),
+        agentTurn,
+        sessionId,
+      ),
+    );
     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 ordinaryRequest = transformed === undefined ? payload : (transformed as JsonRecord);
+    const request = applyResponsesLite(ordinaryRequest, options.model.id, responsesLiteEnabled);
+    const staticRequest = applyResponsesLite(
+      updateInput(ordinaryRequest, []),
+      options.model.id,
+      responsesLiteEnabled,
+    );
+    const cacheDiagnostics = codexCacheDiagnosticContext(
+      ordinaryRequest,
+      request,
+      staticRequest,
+      options.model.id,
+      responsesLiteEnabled,
     );
+    let webSocketResponseHandle: CodexWebSocketResponseHandle | undefined;
+    let compacted: Awaited<ReturnType<typeof collectRemoteCompaction>>;
+    try {
+      compacted = await collectRemoteCompaction(
+        this.transport.request(options.model, request, {
+          ...options.requestOptions,
+          accountId,
+          requestKind: "compaction",
+          turnState: agentTurn.turnState,
+          cacheDiagnostics,
+          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,
@@ -471,12 +892,16 @@ export class CodexProviderRuntime {
     options: OpenAICodexResponsesOptions,
     body: JsonRecord,
     grammarToolInputProperties: GrammarToolInputProperties,
+    agentTurn: ActiveAgentTurn,
+    responsesLiteEnabled: boolean,
   ): Promise<JsonRecord> {
     const sessionId = options.sessionId;
     const scope = sessionId ? this.scopes.get(sessionId) : undefined;
     const threshold = scope?.config.autoCompactAtPercent;
+    if (!sessionId || !scope) {
+      return body;
+    }
     if (
-      !scope ||
       threshold === undefined ||
       scope.contextPercent === null ||
       scope.contextPercent < threshold ||
@@ -515,44 +940,50 @@ 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,
+      compactionMetadata: responsesCompactionV2Metadata("auto", "context_limit", "pre_turn"),
+      agentTurn,
+      responsesLiteEnabled,
+    });
+    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",
+    );
+    this.advanceWindow(sessionId);
+    const cacheSessionId = options.cacheRetention === "none" ? undefined : codexCacheKey(sessionId);
+    return withCodexRequestMetadata(
+      updateInput(body, compacted.checkpoint.history),
+      cacheSessionId,
+      { kind: "turn" },
+      agentTurn.turnId,
+      this.metadataIdentity(cacheSessionId, agentTurn, sessionId),
+    );
   }
 
   stream(
@@ -569,21 +1000,20 @@ export class CodexProviderRuntime {
         api: CODEX_API,
         provider: model.provider,
         model: model.id,
-        usage: {
-          input: 0,
-          output: 0,
-          cacheRead: 0,
-          cacheWrite: 0,
-          totalTokens: 0,
-          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-        },
+        usage: emptyUsage(),
         stopReason: "pending",
         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);
+        this.activateThread(runtimeSessionId);
+        const cacheSessionId =
+          requestOptions.cacheRetention === "none" ? undefined : codexCacheKey(runtimeSessionId);
+        const agentTurn = this.agentTurn(runtimeSessionId);
+        const responsesLiteEnabled = this.responsesLiteEnabled(runtimeSessionId);
         const grammarToolInputProperties = createGrammarToolInputProperties(
           context.tools,
           (model.compat as CodexCompat | undefined)?.supportsOpenAIGrammarTools ?? false,
@@ -592,11 +1022,13 @@ export class CodexProviderRuntime {
           model,
           context,
           requestOptions,
+          runtimeSessionId,
           cacheSessionId,
           grammarToolInputProperties,
+          agentTurn.turnId,
         );
         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,
@@ -611,31 +1043,151 @@ export class CodexProviderRuntime {
           requestOptions,
           body,
           grammarToolInputProperties,
+          agentTurn,
+          responsesLiteEnabled,
+        );
+        const ordinaryBody = body;
+        const staticBody = applyResponsesLite(
+          updateInput(ordinaryBody, []),
+          model.id,
+          responsesLiteEnabled,
+        );
+        body = applyResponsesLite(ordinaryBody, model.id, responsesLiteEnabled);
+        const cacheDiagnostics = codexCacheDiagnosticContext(
+          ordinaryBody,
+          body,
+          staticBody,
+          model.id,
+          responsesLiteEnabled,
         );
 
         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,
+        const prewarmDiagnostics: CodexTransportDiagnostic[] = [];
+        await this.maybePrewarm({
           model,
-          grammarToolInputProperties,
-        );
-        if (!successfulStopReason(output)) {
-          throw new Error(output.errorMessage || "Codex stream ended without a successful stop.");
+          body: staticBody,
+          fullBody: body,
+          requestOptions,
+          accountId,
+          diagnostics: prewarmDiagnostics,
+          turnState: agentTurn.turnState,
+          cacheDiagnostics,
+        });
+        if (prewarmDiagnostics.length > 0) output.diagnostics = prewarmDiagnostics;
+        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,
+          requestKind: "turn" as const,
+          turnState: agentTurn.turnState,
+          cacheDiagnostics,
+          onContinuationReady(handle: CodexContinuationHandle) {
+            continuationHandle = handle;
+          },
+          onWebSocketResponseHandle(handle: CodexWebSocketResponseHandle) {
+            webSocketResponseHandle = handle;
+          },
+          onTransportStart: emitStart,
+          onTransportDiagnostic(diagnostic: CodexTransportDiagnostic) {
+            output.diagnostics = [...(output.diagnostics ?? []), diagnostic];
+          },
+        };
+        let responseRequests = 0;
+        let responseRetries = 0;
+        while (true) {
+          responseRequests += 1;
+          const attemptItems: ResponsesItem[] = [];
+          const terminalState: CodexTerminalState = {};
+          const attemptState: CodexStreamAttemptState = {
+            startedContentIndexes: new Set(),
+            completedContentIndexes: new Set(),
+          };
+          const usageBeforeAttempt = structuredClone(output.usage);
+          output.usage = emptyUsage();
+          try {
+            await processCodexStream(
+              startOnFirstEvent(
+                captureRawEvents(
+                  this.transport.request(model, body, transportRequestOptions),
+                  attemptItems,
+                  terminalState,
+                ),
+                emitStart,
+              ),
+              output,
+              stream,
+              model,
+              grammarToolInputProperties,
+              {
+                attemptState,
+                applyServiceTierPricing(usage, responseServiceTier) {
+                  applyServiceTierPricing(usage, model, body.service_tier, responseServiceTier);
+                },
+              },
+            );
+          } catch (error) {
+            output.usage = usageBeforeAttempt;
+            if (!requestOptions.signal?.aborted) {
+              webSocketResponseHandle?.failParsing(error);
+            }
+            throw error;
+          }
+          output.usage = accumulateUsage(usageBeforeAttempt, output.usage);
+          rawItems.push(...attemptItems.map((item) => structuredClone(item)));
+
+          const nextBody = continueResponseBody(body, attemptItems);
+          const attemptHasToolCall = [...attemptState.completedContentIndexes].some(
+            (index) => output.content[index]?.type === "toolCall",
+          );
+          discardIncompleteAttemptContent(output, attemptState);
+          const retryableTerminal =
+            terminalState.type === "response.incomplete" ||
+            (terminalState.type === "response.failed" &&
+              retryableResponseFailure(terminalState.response));
+          if (
+            retryableTerminal &&
+            nextBody &&
+            responseRetries < this.responseRetryPolicy.maxRetries
+          ) {
+            responseRetries += 1;
+            body = nextBody;
+            output.stopReason = "pending";
+            delete output.errorMessage;
+            delete output.rawStopReason;
+            await waitForResponseRetry(
+              responseRetryDelayMs(this.responseRetryPolicy.baseDelayMs, responseRetries),
+              requestOptions.signal,
+            );
+            continue;
+          }
+
+          if (
+            terminalState.type === "response.completed" &&
+            terminalState.response?.["end_turn"] === false &&
+            !attemptHasToolCall
+          ) {
+            if (!nextBody) {
+              throw new Error(
+                "Codex requested a follow-up response, but its completed output could not be appended to request history.",
+              );
+            }
+            responseRetries = 0;
+            body = nextBody;
+            output.stopReason = "pending";
+            delete output.errorMessage;
+            delete output.rawStopReason;
+            continue;
+          }
+          break;
         }
-        applyServiceTierPricing(
-          output.usage,
-          model,
-          body.service_tier,
-          responseMetadata.serviceTier,
-        );
+        if (requestOptions.signal?.aborted) throw new Error("Request was aborted");
 
         const compat = model.compat as CodexCompat | undefined;
         const canonicalContext: Context = {
@@ -652,7 +1204,7 @@ export class CodexProviderRuntime {
             deferredTools: splitDeferredTools(context, Boolean(compat?.supportsToolSearch))
               .deferred,
             toolOptions: {
-              strict: null,
+              strict: false,
               supportsStrictMode: compat?.supportsStrictMode ?? true,
               supportsOpenAIGrammarTools: compat?.supportsOpenAIGrammarTools ?? false,
             },
@@ -667,19 +1219,37 @@ 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 &&
+          (responseRequests > 1 || 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),
           );
         }
+        try {
+          assertSuccessfulOutput(output);
+        } catch (error) {
+          webSocketResponseHandle?.discard();
+          throw error;
+        }
+        const readyContinuation = continuationHandle;
+        if (
+          responseRequests === 1 &&
+          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 +1264,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 +1276,7 @@ export class CodexProviderRuntime {
     return this.templates.get(sessionId);
   }
 
-  compact(options: {
+  async compact(options: {
     model: Model<any>;
     requestOptions: OpenAICodexResponsesOptions;
     history: ResponsesItem[];
@@ -713,14 +1284,21 @@ export class CodexProviderRuntime {
     grammarToolInputProperties: GrammarToolInputProperties;
     template: JsonRecord;
     priority: boolean;
+    compactionMetadata: CodexCompactionMetadata;
   }): 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 {
+      const compacted = await this.performCompaction(options);
+      const sessionId = options.requestOptions.sessionId;
+      if (sessionId) this.advanceWindow(sessionId);
+      return compacted;
+    } finally {
+      release();
+    }
   }
 }
 
@@ -728,7 +1306,7 @@ export function registerCodexProvider(
   pi: ExtensionAPI,
   resolveConfig: ConfigResolver,
 ): CodexProviderRuntime {
-  const runtime = new CodexProviderRuntime(pi, resolveConfig);
+  const runtime = new CodexProviderRuntime(pi, resolveConfig, resolveCodexInstallationId());
   pi.on("session_start", (_event, ctx) => {
     const base =
       ctx.modelRegistry.getRegisteredNativeProvider(CODEX_PROVIDER) ??
@@ -736,5 +1314,7 @@ export function registerCodexProvider(
     if (!base) throw new Error("Pi's built-in OpenAI Codex provider is unavailable.");
     pi.registerProvider(runtime.createProvider(base));
   });
+  pi.on("agent_start", (_event, ctx) => runtime.beginAgentTurn(ctx));
+  pi.on("agent_end", (_event, ctx) => runtime.endAgentTurn(ctx));
   return runtime;
 }