codex-grok-bridge 1.7.0 → 1.7.2

package/src/imagine.mjs

@@ -0,0 +1,150 @@
+// Codex posts OpenAI Images requests on the provider base
+// (`POST /v1/images/generations`, `POST /v1/images/edits`). grok-build sends
+// the same two paths to the Imagine API with the login bearer, not an
+// XAI_API_KEY. The bodies are not the same: Codex hardcodes `gpt-image-2`
+// plus `size` / `quality` / `background`, and edits use `{image_url}` or
+// `{file_id}`. Imagine wants `grok-imagine-image-quality`, `response_format:
+// b64_json`, and `{url}` references.
+
+export const DEFAULT_IMAGINE_API_BASE = "https://api.x.ai/v1";
+export const IMAGINE_MODEL = "grok-imagine-image-quality";
+
+const IMAGINE_TIMEOUT_MS = 300_000;
+
+function promptOf(input) {
+  const prompt = typeof input?.prompt === "string" ? input.prompt.trim() : "";
+  if (!prompt) {
+    const error = new Error("Image generation requires a prompt.");
+    error.status = 400;
+    throw error;
+  }
+  return prompt;
+}
+
+function basePayload(prompt) {
+  return {
+    model: IMAGINE_MODEL,
+    prompt,
+    n: 1,
+    resolution: "1k",
+    response_format: "b64_json",
+  };
+}
+
+export function imagineGenerationBody(input) {
+  return basePayload(promptOf(input));
+}
+
+export function imagineEditBody(input) {
+  const prompt = promptOf(input);
+  const images = Array.isArray(input?.images) ? input.images : [];
+  const urls = [];
+  for (const image of images) {
+    if (!image || typeof image !== "object") continue;
+    if (typeof image.image_url === "string" && image.image_url.trim()) {
+      urls.push(image.image_url.trim());
+      continue;
+    }
+    if (typeof image.file_id === "string" && image.file_id.trim()) {
+      const error = new Error("OpenAI file ids cannot be edited with Grok.");
+      error.status = 400;
+      throw error;
+    }
+  }
+  if (!urls.length) {
+    const error = new Error("Image edit requires a reference image.");
+    error.status = 400;
+    throw error;
+  }
+  const payload = basePayload(prompt);
+  if (urls.length === 1) payload.image = { url: urls[0] };
+  else {
+    payload.images = urls.map((url) => ({ url }));
+    payload.aspect_ratio = "auto";
+  }
+  return payload;
+}
+
+export function codexImageResponse(payload) {
+  const data = [];
+  if (Array.isArray(payload?.data)) {
+    for (const item of payload.data) {
+      if (item && typeof item.b64_json === "string" && item.b64_json.length)
+        data.push({ b64_json: item.b64_json });
+    }
+  }
+  if (!data.length) {
+    const error = new Error("Image generation returned no image.");
+    error.status = 502;
+    throw error;
+  }
+  const created = Number.isInteger(payload?.created)
+    ? payload.created
+    : Math.floor(Date.now() / 1000);
+  return { created, data };
+}
+
+function scrub(text, token) {
+  let out = String(text ?? "");
+  out = out.replace(/Bearer\s+\S+/gi, "Bearer [redacted]");
+  if (token) out = out.split(token).join("[redacted]");
+  return out;
+}
+
+async function imagineHttp(fetchImpl, url, init) {
+  const response = await fetchImpl(url, { ...init, redirect: "manual" });
+  const text = await response.text().catch(() => "");
+  if (response.status >= 300 && response.status < 400) {
+    const error = new Error("Image generation failed (redirect).");
+    error.status = 502;
+    throw error;
+  }
+  let body = null;
+  if (text) {
+    try {
+      body = JSON.parse(text);
+    } catch {
+      body = null;
+    }
+  }
+  return { status: response.status, body };
+}
+
+export async function forwardImagine(options) {
+  const token = options.token;
+  const kind = options.kind === "edits" ? "edits" : "generations";
+  const payload =
+    kind === "edits" ? imagineEditBody(options.body) : imagineGenerationBody(options.body);
+  const fetchImpl = options.fetchImpl ?? fetch;
+  const base = (options.baseUrl ?? DEFAULT_IMAGINE_API_BASE).replace(/\/$/, "");
+  const timeout = AbortSignal.timeout(IMAGINE_TIMEOUT_MS);
+  const signal = options.signal ? AbortSignal.any([options.signal, timeout]) : timeout;
+  let result;
+  try {
+    result = await imagineHttp(fetchImpl, `${base}/images/${kind}`, {
+      method: "POST",
+      headers: {
+        authorization: `Bearer ${token}`,
+        "content-type": "application/json",
+      },
+      body: JSON.stringify(payload),
+      signal,
+    });
+  } catch (error) {
+    if (error?.status) throw error;
+    const wrapped = new Error(scrub(error?.message || "Image generation failed.", token));
+    wrapped.status = 502;
+    throw wrapped;
+  }
+  if (result.status === 401 || result.status === 403) {
+    const error = new Error("Grok login expired. Run grok login.");
+    error.status = 401;
+    throw error;
+  }
+  if (result.status < 200 || result.status >= 300) {
+    const error = new Error(`Image generation failed (${result.status}).`);
+    error.status = result.status >= 400 && result.status < 500 ? result.status : 502;
+    throw error;
+  }
+  return codexImageResponse(result.body);
+}

package/src/proxy.mjs

@@ -8,8 +8,13 @@ import { BRIDGE_ERROR, classifyBridgeError } from "./errors.mjs";
 export const DEFAULT_PROXY_BASE = "https://cli-chat-proxy.grok.com/v1";
 export const DEFAULT_CLIENT_IDENTIFIER = "grok-shell";
 
+const UNKNOWN_CLIENT_VERSION = "unknown";
 let cachedClientVersion;
 
+export function clearClientVersionCache() {
+  cachedClientVersion = undefined;
+}
+
 export function detectGrokClientVersion(home = homedir()) {
   if (cachedClientVersion) return cachedClientVersion;
   try {
@@ -20,9 +25,9 @@ export function detectGrokClientVersion(home = homedir()) {
       timeout: 5000,
     });
     const match = output.match(/grok\s+(\S+(?:\s+\([^)]+\))?)/i);
-    cachedClientVersion = match ? match[1].trim() : "1.0.24";
+    cachedClientVersion = match ? match[1].trim() : UNKNOWN_CLIENT_VERSION;
   } catch {
-    cachedClientVersion = "1.0.24";
+    cachedClientVersion = UNKNOWN_CLIENT_VERSION;
   }
   return cachedClientVersion;
 }
@@ -74,9 +79,9 @@ export async function openProxyStream(options) {
   return response;
 }
 
-// Resolve when the consumer has drained, reject if it goes away first. Without
-// this, a slow client makes the bridge buffer the whole Grok response in memory,
-// and a client that vanishes mid-stream is never noticed.
+// The upstream reply is already buffered until it finishes. This wait applies
+// only while that buffer is flushed: a slow Codex exerts backpressure, and a
+// client that vanishes is noticed instead of being written forever.
 function drained(output) {
   return new Promise((resolve, reject) => {
     const settle = (fn, value) => {
@@ -100,24 +105,68 @@ function drained(output) {
   });
 }
 
-async function writeBlock(output, block) {
-  if (output.destroyed || output.writableEnded)
-    throw Object.assign(new Error("The client closed the stream"), {
-      code: "ERR_STREAM_PREMATURE_CLOSE",
-    });
-  if (!output.write(block)) await drained(output);
-}
+// First try, then two more. That is the budget Codex used to spend on
+// request_max_retries / stream_max_retries. A retry sends the prompt again, so
+// the input tokens for that attempt can be billed again.
+export const PROXY_ATTEMPTS = 3;
 
-// Nothing has reached Codex until the first SSE block is written, so a stream
-// that dies before it opens is safe to send again. After that it never is:
-// re-sending would duplicate items Codex has already recorded.
+// DNS and connect failures are included. A socket reset (ECONNRESET, EPIPE,
+// UND_ERR_SOCKET) or a body that ends before response.completed is included
+// only while Codex has not been sent a byte. User aborts and deterministic
+// 422s are not in this set.
 const RETRYABLE = new Set([
   BRIDGE_ERROR.DNS,
   BRIDGE_ERROR.CONNECT,
   BRIDGE_ERROR.UPSTREAM_CLOSED,
 ]);
 
-export async function openProxyStreamWithRetry(options, attempts = 2) {
+// response.incomplete is the other usage event this code already treats as the
+// end of a response. response.failed here is a finished upstream event, not a
+// dropped socket, so it is forwarded instead of retried.
+const TERMINAL_EVENTS = new Set([
+  "response.completed",
+  "response.incomplete",
+  "response.failed",
+]);
+
+export function isTerminalSseBlock(part) {
+  for (const line of part.split("\n")) {
+    if (line.startsWith("event:") && TERMINAL_EVENTS.has(line.slice(6).trim()))
+      return true;
+    if (!line.startsWith("data:")) continue;
+    const payload = line.slice(5).trim();
+    if (!payload || payload === "[DONE]") continue;
+    try {
+      const value = JSON.parse(payload);
+      if (TERMINAL_EVENTS.has(value?.type)) return true;
+    } catch {
+      // A non-JSON data line cannot be the terminal event.
+    }
+  }
+  return false;
+}
+
+function prematureUpstreamClose() {
+  return Object.assign(new Error("Grok closed the connection before finishing"), {
+    code: "ERR_UPSTREAM_PREMATURE_CLOSE",
+  });
+}
+
+function clientClosed() {
+  return Object.assign(new Error("The client closed the stream"), {
+    code: "ERR_STREAM_PREMATURE_CLOSE",
+  });
+}
+
+// Returns the retry kind, or null when this failure must surface as-is.
+function retryKind(error, options, attempt, rounds, clientBytes) {
+  if (clientBytes > 0 || options.signal?.aborted || attempt === rounds)
+    return null;
+  const kind = classifyBridgeError(error);
+  return RETRYABLE.has(kind) ? kind : null;
+}
+
+export async function openProxyStreamWithRetry(options, attempts = PROXY_ATTEMPTS) {
   const rounds = Math.max(1, attempts);
   let lastError;
   for (let attempt = 1; attempt <= rounds; attempt += 1) {
@@ -125,46 +174,196 @@ export async function openProxyStreamWithRetry(options, attempts = 2) {
       return await openProxyStream(options);
     } catch (error) {
       lastError = error;
-      const kind = classifyBridgeError(error);
-      if (
-        !RETRYABLE.has(kind) ||
-        options.signal?.aborted ||
-        attempt === rounds
-      )
-        break;
+      const kind = retryKind(error, options, attempt, rounds, 0);
+      if (!kind) break;
       options.onRetry?.({ attempt, kind });
     }
   }
   throw lastError;
 }
 
-export async function pipeProxySse(stream, output, map, usageBox) {
+function completeBlocks(buffer) {
+  const parts = buffer.split("\n\n");
+  const rest = parts.pop() ?? "";
+  return { parts: parts.filter((part) => part.trim()), rest };
+}
+
+// Read the upstream body to memory. Write to Codex only after a terminal
+// event, so a reset before that has not shown Codex a partial reply.
+export async function pipeProxySse(stream, output, map, usageBox, onClientByte) {
   const reader = stream.getReader();
   const decoder = new TextDecoder();
   let buffer = "";
   const rewrite = createSseRewriter(map);
+  const pending = [];
+  let clientBytes = 0;
+  let finished = false;
+  let failed = false;
+
+  const flush = async (parts) => {
+    for (const part of parts) {
+      const rewritten = rewrite(part);
+      if (rewritten === null) continue;
+      const block = `${rewritten}\n\n`;
+      if (output.destroyed || output.writableEnded) throw clientClosed();
+      // Count the byte as soon as the client stream accepts it. A later drain
+      // error must not look like "Codex saw nothing" and resubmit the prompt.
+      const accepted = output.write(block);
+      clientBytes += Buffer.byteLength(block);
+      onClientByte?.();
+      if (!accepted) await drained(output);
+    }
+  };
+
+  const accept = async (parts) => {
+    if (!parts.length) return;
+    if (finished) {
+      await flush(parts);
+      return;
+    }
+    const terminalAt = parts.findIndex((part) => isTerminalSseBlock(part));
+    if (terminalAt === -1) {
+      pending.push(...parts);
+      return;
+    }
+    pending.push(...parts.slice(0, terminalAt + 1));
+    await flush(pending);
+    pending.length = 0;
+    finished = true;
+    await flush(parts.slice(terminalAt + 1));
+  };
+
   try {
     while (true) {
+      if (output.destroyed || output.writableEnded) throw clientClosed();
       const { done, value } = await reader.read();
       if (done) break;
       buffer += decoder.decode(value, { stream: true });
-      const parts = buffer.split("\n\n");
-      buffer = parts.pop();
-      for (const part of parts) {
-        if (!part.trim()) continue;
-        const rewritten = rewrite(part);
-        if (rewritten !== null) await writeBlock(output, rewritten + "\n\n");
-      }
+      const split = completeBlocks(buffer);
+      buffer = split.rest;
+      await accept(split.parts);
     }
     buffer += decoder.decode();
-    if (buffer.trim()) {
-      const rewritten = rewrite(buffer);
-      if (rewritten !== null) await writeBlock(output, rewritten + "\n\n");
-    }
+    if (buffer.trim()) await accept([buffer]);
+    if (!finished) throw prematureUpstreamClose();
   } catch (error) {
-    // Stop pulling from Grok the moment the client is gone; leaving the body
+    failed = true;
+    // Stop pulling from Grok the moment this attempt is over; leaving the body
     // unread holds the upstream socket open for the rest of the response.
     await reader.cancel(error).catch(() => {});
+    if (error && typeof error === "object") error.clientBytes = clientBytes;
+    throw error;
+  } finally {
+    if (!failed) reader.releaseLock();
+    if (usageBox) usageBox.cacheUsage = rewrite.cacheUsage ?? null;
+  }
+}
+
+// Open and read one attempt. If the socket dies before Codex has a byte, send
+// the same request again. Once a byte has been written, never resubmit.
+async function relayAttempts(options, consume, attempts = PROXY_ATTEMPTS) {
+  const rounds = Math.max(1, attempts);
+  let lastError;
+  for (let attempt = 1; attempt <= rounds; attempt += 1) {
+    try {
+      const response = await openProxyStream(options);
+      await consume(response);
+      return response;
+    } catch (error) {
+      lastError = error;
+      const kind = retryKind(
+        error,
+        options,
+        attempt,
+        rounds,
+        error?.clientBytes ?? 0,
+      );
+      if (!kind) break;
+      options.onRetry?.({ attempt, kind });
+    }
+  }
+  throw lastError;
+}
+
+export async function relayProxySse(options, output, map, usageBox, attempts = PROXY_ATTEMPTS) {
+  return relayAttempts(
+    options,
+    (response) =>
+      pipeProxySse(response.body, output, map, usageBox, options.onClientByte),
+    attempts,
+  );
+}
+
+// Read until response.completed (or the other terminal events). Do not write
+// those bytes yet: a video tool call has to be answered before Codex sees the
+// turn, and a reset before that write is still safe to retry.
+export async function readUntilTerminal(stream, output) {
+  const reader = stream.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+  const collected = [];
+  let released = false;
+  const finish = async (error) => {
+    if (released) return;
+    released = true;
+    await reader.cancel(error).catch(() => {});
+  };
+  try {
+    while (true) {
+      if (output?.destroyed || output?.writableEnded) throw clientClosed();
+      const { done, value } = await reader.read();
+      if (done) break;
+      buffer += decoder.decode(value, { stream: true });
+      const split = completeBlocks(buffer);
+      buffer = split.rest;
+      collected.push(...split.parts);
+      if (collected.some((part) => isTerminalSseBlock(part))) {
+        await finish();
+        return collected;
+      }
+    }
+    buffer += decoder.decode();
+    if (buffer.trim()) collected.push(buffer);
+    if (!collected.some((part) => isTerminalSseBlock(part)))
+      throw prematureUpstreamClose();
+    await finish();
+    return collected;
+  } catch (error) {
+    await finish(error);
+    if (error && typeof error === "object" && error.clientBytes == null)
+      error.clientBytes = 0;
+    throw error;
+  }
+}
+
+export async function readRelayProxySse(options, output, attempts = PROXY_ATTEMPTS) {
+  let parts;
+  await relayAttempts(
+    options,
+    async (response) => {
+      parts = await readUntilTerminal(response.body, output);
+    },
+    attempts,
+  );
+  return parts;
+}
+
+export async function emitProxySse(parts, output, map, usageBox, onClientByte) {
+  const rewrite = createSseRewriter(map);
+  let clientBytes = 0;
+  try {
+    for (const part of parts) {
+      const rewritten = rewrite(part);
+      if (rewritten === null) continue;
+      const block = `${rewritten}\n\n`;
+      if (output.destroyed || output.writableEnded) throw clientClosed();
+      const accepted = output.write(block);
+      clientBytes += Buffer.byteLength(block);
+      onClientByte?.();
+      if (!accepted) await drained(output);
+    }
+  } catch (error) {
+    if (error && typeof error === "object") error.clientBytes = clientBytes;
     throw error;
   } finally {
     if (usageBox) usageBox.cacheUsage = rewrite.cacheUsage ?? null;

package/src/runtime.mjs

@@ -29,14 +29,11 @@ export async function startRuntime(options = {}) {
     env_key: "CODEX_GROK_BRIDGE_TOKEN",
     wire_api: "responses",
     requires_openai_auth: false,
-    // The upstream resets a response mid-stream now and then — measured three
-    // times as ECONNRESET at 25s, 27s and 253s, on requests from 726 KB to
-    // 22 MB. The bridge cannot retry that: Codex has already received part of
-    // the answer, so re-sending would duplicate it. Codex can, because it owns
-    // the conversation and re-sends the same request. These were both 0, which
-    // turned every such reset into a dead turn.
-    request_max_retries: 2,
-    stream_max_retries: 2,
+    // The bridge buffers the upstream reply and retries a reset itself, up to
+    // two more attempts, before Codex has seen any SSE byte. Codex must not
+    // also resend the full prompt.
+    request_max_retries: 0,
+    stream_max_retries: 0,
     // Do not inherit whatever Codex's default happens to be across upgrades.
     stream_idle_timeout_ms: 300000,
   };