@shanepadgett/tau-agent 0.5.1 → 0.7.0

package/extensions/explore/read.ts

@@ -2,17 +2,21 @@ import {
 	createReadToolDefinition,
 	DEFAULT_MAX_BYTES,
 	formatSize,
+	generateUnifiedPatch,
 	truncateHead,
 	type ReadToolDetails,
 	type ToolDefinition,
 } from "@earendil-works/pi-coding-agent";
-import { Text } from "@earendil-works/pi-tui";
+import { Container, Text } from "@earendil-works/pi-tui";
+import { createHash } from "node:crypto";
 import { readFile } from "node:fs/promises";
 import { isAbsolute, resolve } from "node:path";
 import { Type, type Static } from "typebox";
 import { formatToolRowTitle, type ToolRowStateStore } from "../../shared/tool-row-state.js";
 import { normalizeCountLimit } from "./limits.ts";
 import { stripLeadingAt } from "./path-display.ts";
+import { createReadCacheStore, type ReadCacheMetaV1, type ReadCacheStore } from "./read-cache.ts";
+import { createReadSnapshotStore, type ReadSnapshotStore } from "./read-snapshots.ts";
 
 const readSchema = Type.Object({
 	path: Type.String({ description: "Path to the file to read (relative or absolute)" }),
@@ -23,11 +27,26 @@ const readSchema = Type.Object({
 
 type ReadToolInput = Static<typeof readSchema>;
 type BaseReadDefinition = ReturnType<typeof createReadToolDefinition>;
-type ReadDefinition = ToolDefinition<typeof readSchema, ReadToolDetails | undefined>;
+interface ExploreReadDetails extends ReadToolDetails {
+	readCache?: ReadCacheMetaV1;
+}
+type ReadDefinition = ToolDefinition<typeof readSchema, ExploreReadDetails | undefined>;
 type ReadExecute = ReadDefinition["execute"];
 type ReadRenderCall = NonNullable<ReadDefinition["renderCall"]>;
 type ReadRenderResult = NonNullable<ReadDefinition["renderResult"]>;
 
+interface BaselineTextResult {
+	text: string;
+	details: ReadToolDetails | undefined;
+	totalLines: number;
+	startLine: number;
+	endLine: number;
+	completeFile: boolean;
+	scopeKey: string;
+	summary: string;
+	cacheable: boolean;
+}
+
 const readDefinitionByCwd = new Map<string, BaseReadDefinition>();
 
 function readDefinitionForCwd(cwd: string): BaseReadDefinition {
@@ -67,7 +86,114 @@ function renderCallSummary(args: ReadToolInput | undefined): string {
 	return `${path}:${start}${end === "" ? "" : `-${end}`}`;
 }
 
-export function createExploreReadTool(rowState: ToolRowStateStore): ReadDefinition {
+function estimateTokens(text: string): number {
+	return Math.ceil(text.length / 4);
+}
+
+function baselineText(text: string, params: ReadToolInput): BaselineTextResult {
+	const allLines = text.split("\n");
+	const startIndex = params.offset ? Math.max(0, params.offset - 1) : 0;
+	const startLine = startIndex + 1;
+	if (startIndex >= allLines.length) {
+		throw new Error(`Offset ${params.offset} is beyond end of file (${allLines.length} lines total)`);
+	}
+
+	const selectedEnd =
+		params.limit === undefined ? allLines.length : Math.min(startIndex + params.limit, allLines.length);
+	const selectedLines = allLines.slice(startIndex, selectedEnd);
+	const selectedContent = selectedLines
+		.map((line, index) => (params.lineNumbers ? `${startLine + index}: ${line}` : line))
+		.join("\n");
+	const truncation = truncateHead(selectedContent);
+	let outputText: string;
+	let details: ReadToolDetails | undefined;
+	let outputLines = selectedLines.length;
+	let cacheable = true;
+
+	if (truncation.firstLineExceedsLimit) {
+		const firstLineSize = formatSize(Buffer.byteLength(selectedLines[0] ?? "", "utf-8"));
+		outputText = `[Line ${startLine} is ${firstLineSize}, exceeds ${formatSize(DEFAULT_MAX_BYTES)} limit. Use bash: sed -n '${startLine}p' ${params.path} | head -c ${DEFAULT_MAX_BYTES}]`;
+		details = { truncation };
+		outputLines = 0;
+		cacheable = false;
+	} else if (truncation.truncated) {
+		outputLines = truncation.outputLines;
+		const endLineDisplay = startLine + outputLines - 1;
+		const nextOffset = endLineDisplay + 1;
+		outputText = truncation.content;
+		outputText +=
+			truncation.truncatedBy === "lines"
+				? `\n\n[Showing lines ${startLine}-${endLineDisplay} of ${allLines.length}. Use offset=${nextOffset} to continue.]`
+				: `\n\n[Showing lines ${startLine}-${endLineDisplay} of ${allLines.length} (${formatSize(DEFAULT_MAX_BYTES)} limit). Use offset=${nextOffset} to continue.]`;
+		details = { truncation };
+	} else if (selectedEnd < allLines.length) {
+		outputText = `${truncation.content}\n\n[${allLines.length - selectedEnd} more lines in file. Use offset=${selectedEnd + 1} to continue.]`;
+	} else {
+		outputText = truncation.content;
+	}
+
+	const endLine = outputLines === 0 ? startLine : startLine + outputLines - 1;
+	const completeFile = startIndex === 0 && selectedEnd === allLines.length && !truncation.truncated;
+	const scopeKey = `${completeFile ? "full" : `r:${startLine}:${endLine}`}:n${params.lineNumbers ? 1 : 0}`;
+	const summary = completeFile ? `${allLines.length} lines` : `${outputLines} lines`;
+	return {
+		text: outputText,
+		details,
+		totalLines: allLines.length,
+		startLine,
+		endLine,
+		completeFile,
+		scopeKey,
+		summary,
+		cacheable,
+	};
+}
+
+function withMeta(baseline: BaselineTextResult, meta: ReadCacheMetaV1, text = baseline.text) {
+	return {
+		content: [{ type: "text" as const, text }],
+		details: { ...baseline.details, readCache: meta },
+	};
+}
+
+function createMeta(
+	baseline: BaselineTextResult,
+	pathKey: string,
+	hash: string,
+	mode: ReadCacheMetaV1["mode"],
+	returnedText: string,
+	baseHash?: string,
+	summary = baseline.summary,
+): ReadCacheMetaV1 {
+	return {
+		v: 1,
+		pathKey,
+		scopeKey: baseline.scopeKey,
+		servedHash: hash,
+		baseHash,
+		mode,
+		baselineTokens: estimateTokens(baseline.text),
+		returnedTokens: estimateTokens(returnedText),
+		totalLines: baseline.totalLines,
+		summary,
+	};
+}
+
+function countDiffLines(patch: string): { added: number; removed: number } {
+	let added = 0;
+	let removed = 0;
+	for (const line of patch.split("\n")) {
+		if (line.startsWith("+") && !line.startsWith("+++")) added += 1;
+		else if (line.startsWith("-") && !line.startsWith("---")) removed += 1;
+	}
+	return { added, removed };
+}
+
+export function createExploreReadTool(
+	rowState: ToolRowStateStore,
+	cache: ReadCacheStore = createReadCacheStore(),
+	snapshots: ReadSnapshotStore = createReadSnapshotStore(),
+): ReadDefinition {
 	const baseDefinition = readDefinitionForCwd(process.cwd());
 	return {
 		...baseDefinition,
@@ -81,7 +207,7 @@ export function createExploreReadTool(rowState: ToolRowStateStore): ReadDefiniti
 		) {
 			const definition = readDefinitionForCwd(ctx.cwd);
 			const normalized = normalizeReadParams(params);
-			const path = isAbsolute(normalized.path) ? normalized.path : resolve(ctx.cwd, normalized.path);
+			const path = isAbsolute(normalized.path) ? resolve(normalized.path) : resolve(ctx.cwd, normalized.path);
 			const buffer = await readFile(path);
 			if (isSupportedImage(buffer)) {
 				return definition.execute(
@@ -94,42 +220,49 @@ export function createExploreReadTool(rowState: ToolRowStateStore): ReadDefiniti
 			}
 
 			if (signal?.aborted) throw new Error("Operation aborted");
-			const allLines = buffer.toString("utf-8").split("\n");
-			const startLine = normalized.offset ? Math.max(0, normalized.offset - 1) : 0;
-			const startLineDisplay = startLine + 1;
-			if (startLine >= allLines.length) {
-				throw new Error(`Offset ${normalized.offset} is beyond end of file (${allLines.length} lines total)`);
+			let text: string;
+			try {
+				text = new TextDecoder("utf-8", { fatal: true }).decode(buffer);
+			} catch {
+				return definition.execute(toolCallId, normalized, signal, onUpdate, ctx);
+			}
+
+			const baseline = baselineText(text, normalized);
+			if (!baseline.cacheable) {
+				return { content: [{ type: "text", text: baseline.text }], details: baseline.details };
 			}
+			const hash = createHash("sha256").update(buffer).digest("hex");
+			const decision = cache.decision(ctx, path, baseline.scopeKey);
+			let output = baseline.text;
+			let mode: ReadCacheMetaV1["mode"] = decision.recovery ? "recovery" : "baseline";
+			let summary = baseline.summary;
 
-			const endLine =
-				normalized.limit === undefined ? allLines.length : Math.min(startLine + normalized.limit, allLines.length);
-			const selectedLines = allLines.slice(startLine, endLine);
-			const selectedContent = selectedLines
-				.map((line, index) => (normalized.lineNumbers ? `${startLineDisplay + index}: ${line}` : line))
-				.join("\n");
-			const truncation = truncateHead(selectedContent);
-			let outputText: string;
-			let details: ReadToolDetails | undefined;
-			if (truncation.firstLineExceedsLimit) {
-				const firstLineSize = formatSize(Buffer.byteLength(selectedLines[0] ?? "", "utf-8"));
-				outputText = `[Line ${startLineDisplay} is ${firstLineSize}, exceeds ${formatSize(DEFAULT_MAX_BYTES)} limit. Use bash: sed -n '${startLineDisplay}p' ${normalized.path} | head -c ${DEFAULT_MAX_BYTES}]`;
-				details = { truncation };
-			} else if (truncation.truncated) {
-				const endLineDisplay = startLineDisplay + truncation.outputLines - 1;
-				const nextOffset = endLineDisplay + 1;
-				outputText = truncation.content;
-				outputText +=
-					truncation.truncatedBy === "lines"
-						? `\n\n[Showing lines ${startLineDisplay}-${endLineDisplay} of ${allLines.length}. Use offset=${nextOffset} to continue.]`
-						: `\n\n[Showing lines ${startLineDisplay}-${endLineDisplay} of ${allLines.length} (${formatSize(DEFAULT_MAX_BYTES)} limit). Use offset=${nextOffset} to continue.]`;
-				details = { truncation };
-			} else if (endLine < allLines.length) {
-				outputText = `${truncation.content}\n\n[${allLines.length - endLine} more lines in file. Use offset=${endLine + 1} to continue.]`;
-			} else {
-				outputText = truncation.content;
+			if (!decision.recovery && decision.baseHash === hash) {
+				output = baseline.completeFile
+					? `unchanged, ${baseline.totalLines} lines`
+					: `unchanged, lines ${baseline.startLine}-${baseline.endLine} of ${baseline.totalLines}`;
+				mode = "unchanged";
+				summary = output;
+			} else if (!decision.recovery && decision.baseHash && baseline.completeFile && !normalized.lineNumbers) {
+				const baseText = snapshots.get(decision.baseHash);
+				if (baseText !== undefined) {
+					const patch = generateUnifiedPatch(normalized.path, baseText, text, 3);
+					const counts = countDiffLines(patch);
+					const candidate = `[read: ${counts.added} lines added, ${counts.removed} removed of ${baseline.totalLines}]\n${patch}`;
+					const candidateTruncation = truncateHead(candidate);
+					if (!candidateTruncation.truncated && estimateTokens(candidate) < estimateTokens(baseline.text)) {
+						output = candidate;
+						mode = "diff";
+						summary = `+${counts.added} -${counts.removed}`;
+					}
+				}
 			}
 
-			return { content: [{ type: "text", text: outputText }], details };
+			if (signal?.aborted) throw new Error("Operation aborted");
+			snapshots.set(hash, text, buffer.byteLength);
+			const meta = createMeta(baseline, path, hash, mode, output, decision.baseHash, summary);
+			cache.record(ctx, meta);
+			return withMeta(baseline, meta, output);
 		},
 		renderCall(
 			args: Parameters<ReadRenderCall>[0],
@@ -138,6 +271,10 @@ export function createExploreReadTool(rowState: ToolRowStateStore): ReadDefiniti
 		) {
 			rowState.watch(context.toolCallId, context.invalidate);
 			const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
+			if (context.executionStarted) {
+				text.setText("");
+				return text;
+			}
 			const title = formatToolRowTitle(rowState, context.toolCallId, "read", theme);
 			text.setText(`${title} ${theme.fg("muted", renderCallSummary(args))}`);
 			return text;
@@ -148,13 +285,32 @@ export function createExploreReadTool(rowState: ToolRowStateStore): ReadDefiniti
 			theme: Parameters<ReadRenderResult>[2],
 			context: Parameters<ReadRenderResult>[3],
 		) {
-			if (!context.expanded) {
-				const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
-				text.setText("");
-				return text;
+			rowState.watch(context.toolCallId, context.invalidate);
+			const container = new Container();
+			const title = formatToolRowTitle(rowState, context.toolCallId, "read", theme);
+			const args = context.args as ReadToolInput | undefined;
+			const details = result.details as ExploreReadDetails | undefined;
+			const summary = details?.readCache?.summary;
+			const summaryText = context.isError
+				? theme.fg("error", "error")
+				: summary
+					? theme.fg(
+							details?.readCache?.mode === "unchanged"
+								? "success"
+								: details?.readCache?.mode === "diff"
+									? "accent"
+									: "muted",
+							summary,
+						)
+					: "";
+			const header = `${title} ${theme.fg("muted", renderCallSummary(args))}${summaryText ? `  ${summaryText}` : ""}`;
+			container.addChild(new Text(header, 0, 0));
+			if (options.expanded) {
+				const definition = readDefinitionForCwd(context.cwd);
+				const body = definition.renderResult?.(result, { ...options, expanded: true }, theme, context);
+				if (body) container.addChild(body);
 			}
-			const definition = readDefinitionForCwd(context.cwd);
-			return definition.renderResult?.(result, { ...options, expanded: true }, theme, context) ?? new Text("", 0, 0);
+			return container;
 		},
 	};
 }

package/extensions/image-gen/README.md

@@ -1,11 +1,11 @@
 # Image Generation
 
-`image_gen` generates raster images and edits up to five local raster images with OpenAI Codex. It uses the current `gpt-image-2` model and saves results under `~/.local/share/tau-agent/images/` by default. Pass an explicit PNG path when the image should be saved in the current repository or another chosen location.
+`image_gen` generates raster images and edits up to three local raster images with Grok Imagine. It uses `grok-imagine-image-quality` and saves results under `~/.local/share/tau-agent/images/` by default. Pass an explicit path with the expected image extension when the image should be saved in the current repository or another chosen location.
 
-Use `/login` and select OpenAI Codex before invoking the tool. No `OPENAI_API_KEY` is needed.
+Use `/login` and select xAI (Grok subscription OAuth) before invoking the tool. No xAI API key is used.
 
 Run `/reload` after installing or changing the extension.
 
-The model invokes `image_gen` with a prompt. Size, quality, and background use the Codex automatic settings. For edits, it also supplies one to five local PNG, JPEG, or WebP paths. Successful PNGs up to 12 MiB are returned inline for inspection; larger results remain available at the saved path.
+The model invokes `image_gen` with a prompt. For edits, it also supplies one to three local PNG, JPEG, or WebP paths. Successful images up to 12 MiB are returned inline for inspection; larger results remain available at the saved path.
 
-This extension calls a private Codex backend. OpenAI may change its availability or protocol without notice.
+This extension uses xAI's undocumented subscription OAuth access. xAI may change its availability, entitlement rules, or protocol without notice.

package/extensions/image-gen/client.ts

@@ -1,60 +1,34 @@
-const CODEX_IMAGES_URL = "https://chatgpt.com/backend-api/codex/images";
-const MODEL = "gpt-image-2";
-const PNG_SIGNATURE = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]);
+import { XAI_API_BASE_URL, XAI_IMAGE_MODEL } from "../xai/constants.ts";
+
 const MAX_ERROR_BODY_BYTES = 8192;
 const MAX_ERROR_MESSAGE_LENGTH = 2000;
-
-interface FetchResponse {
-	ok: boolean;
-	status: number;
-	body: ReadableStream<Uint8Array> | null;
-	json(): Promise<unknown>;
-}
-
-export interface CodexAuth {
-	token: string;
-	accountId: string;
-}
+const REQUEST_TIMEOUT_MS = 60_000;
+const MAX_ATTEMPTS = 3;
+const RETRY_BASE_DELAY_MS = 500;
 
 export interface EditImage {
 	mimeType: "image/png" | "image/jpeg" | "image/webp";
 	data: string;
 }
 
-interface GeneratedImage {
+export interface GeneratedImage {
 	bytes: Buffer;
 	base64: string;
-	mimeType: "image/png";
+	mimeType: EditImage["mimeType"];
 }
 
-function isRecord(value: unknown): value is Record<string, unknown> {
-	return typeof value === "object" && value !== null && !Array.isArray(value);
+interface HttpResponse {
+	ok: boolean;
+	status: number;
+	body: ReadableStream<Uint8Array> | null;
+	json(): Promise<unknown>;
 }
 
-export function resolveCodexAuth(token: string): CodexAuth {
-	const invalidCredential = () =>
-		new Error("The OpenAI Codex credential does not contain a usable ChatGPT account ID. Run /login again.");
-	const segments = token.split(".");
-	if (segments.length !== 3 || !segments[1]) throw invalidCredential();
-	if (!/^[A-Za-z0-9_-]+$/.test(segments[1]) || segments[1].length % 4 === 1) throw invalidCredential();
-
-	let payload: unknown;
-	try {
-		const decoded = Buffer.from(segments[1], "base64url");
-		if (decoded.toString("base64url") !== segments[1]) throw invalidCredential();
-		payload = JSON.parse(decoded.toString("utf8"));
-	} catch {
-		throw invalidCredential();
-	}
-	if (!isRecord(payload)) throw invalidCredential();
-	const authClaim = payload["https://api.openai.com/auth"];
-	if (!isRecord(authClaim)) throw invalidCredential();
-	const accountId = authClaim.chatgpt_account_id;
-	if (typeof accountId !== "string" || !accountId.trim()) throw invalidCredential();
-	return { token, accountId: accountId.trim() };
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
 }
 
-async function readBoundedError(response: FetchResponse): Promise<string> {
+async function boundedError(response: HttpResponse): Promise<string> {
 	if (!response.body) return "";
 	const reader = response.body.getReader();
 	const chunks: Uint8Array[] = [];
@@ -63,8 +37,7 @@ async function readBoundedError(response: FetchResponse): Promise<string> {
 		while (length < MAX_ERROR_BODY_BYTES) {
 			const result = await reader.read();
 			if (result.done) break;
-			const remaining = MAX_ERROR_BODY_BYTES - length;
-			const chunk = result.value.subarray(0, remaining);
+			const chunk = result.value.subarray(0, MAX_ERROR_BODY_BYTES - length);
 			chunks.push(chunk);
 			length += chunk.length;
 			if (chunk.length < result.value.length) break;
@@ -72,23 +45,21 @@ async function readBoundedError(response: FetchResponse): Promise<string> {
 	} finally {
 		await reader.cancel().catch(() => undefined);
 	}
-	const bytes = new Uint8Array(length);
-	let offset = 0;
-	for (const chunk of chunks) {
-		bytes.set(chunk, offset);
-		offset += chunk.length;
-	}
-	return new TextDecoder().decode(bytes);
+	return Buffer.concat(
+		chunks.map((chunk) => Buffer.from(chunk)),
+		length,
+	)
+		.toString()
+		.trim();
 }
 
 function serverErrorMessage(body: string, token: string): string {
-	let message = body.trim();
+	let message = body;
 	try {
-		const parsed: unknown = JSON.parse(message);
-		if (isRecord(parsed)) {
-			const error = parsed.error;
-			if (isRecord(error) && typeof error.message === "string") message = error.message;
-			else if (typeof parsed.message === "string") message = parsed.message;
+		const value: unknown = JSON.parse(body);
+		if (isRecord(value)) {
+			if (isRecord(value.error) && typeof value.error.message === "string") message = value.error.message;
+			else if (typeof value.message === "string") message = value.message;
 		}
 	} catch {
 		// Plain-text error body.
@@ -96,70 +67,117 @@ function serverErrorMessage(body: string, token: string): string {
 	return message.replaceAll(token, "[redacted]").slice(0, MAX_ERROR_MESSAGE_LENGTH).trim();
 }
 
+export function detectImageMimeType(bytes: Buffer): GeneratedImage["mimeType"] | undefined {
+	if (
+		bytes.length >= 8 &&
+		bytes.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))
+	) {
+		return "image/png";
+	}
+	if (bytes.length >= 3 && bytes[0] === 0xff && bytes[1] === 0xd8 && bytes[2] === 0xff) return "image/jpeg";
+	if (
+		bytes.length >= 12 &&
+		bytes.subarray(0, 4).toString("ascii") === "RIFF" &&
+		bytes.subarray(8, 12).toString("ascii") === "WEBP"
+	) {
+		return "image/webp";
+	}
+	return undefined;
+}
+
 function decodeImageResponse(value: unknown): GeneratedImage {
-	if (!isRecord(value) || !Array.isArray(value.data) || value.data.length === 0 || !isRecord(value.data[0])) {
-		throw new Error("OpenAI Codex returned an invalid image response");
+	if (!isRecord(value) || !Array.isArray(value.data) || !isRecord(value.data[0])) {
+		throw new Error("xAI returned an invalid image response");
 	}
 	const encoded = value.data[0].b64_json;
-	if (typeof encoded !== "string" || !encoded.trim()) {
-		throw new Error("OpenAI Codex returned an invalid image response");
-	}
+	if (typeof encoded !== "string" || !encoded.trim()) throw new Error("xAI returned an invalid image response");
 	const base64 = encoded.trim();
 	if (!/^[A-Za-z0-9+/]*={0,2}$/.test(base64) || base64.length % 4 === 1) {
-		throw new Error("OpenAI Codex returned invalid base64 image data");
+		throw new Error("xAI returned invalid base64 image data");
 	}
 	const bytes = Buffer.from(base64, "base64");
 	if (bytes.toString("base64").replace(/=+$/, "") !== base64.replace(/=+$/, "")) {
-		throw new Error("OpenAI Codex returned invalid base64 image data");
-	}
-	if (bytes.length < PNG_SIGNATURE.length || !bytes.subarray(0, PNG_SIGNATURE.length).equals(PNG_SIGNATURE)) {
-		throw new Error("OpenAI Codex returned image data that is not PNG");
+		throw new Error("xAI returned invalid base64 image data");
 	}
-	return { bytes, base64: bytes.toString("base64"), mimeType: "image/png" };
+	const mimeType = detectImageMimeType(bytes);
+	if (!mimeType) throw new Error("xAI returned unsupported image data");
+	return { bytes, base64: bytes.toString("base64"), mimeType };
+}
+
+function retryable(status: number): boolean {
+	return status === 408 || status === 409 || status === 425 || status === 429 || status >= 500;
+}
+
+async function delay(milliseconds: number, signal?: AbortSignal): Promise<void> {
+	await new Promise<void>((resolve, reject) => {
+		const timeout = setTimeout(() => {
+			signal?.removeEventListener("abort", abort);
+			resolve();
+		}, milliseconds);
+		const abort = () => {
+			clearTimeout(timeout);
+			signal?.removeEventListener("abort", abort);
+			reject(signal?.reason ?? new DOMException("Aborted", "AbortError"));
+		};
+		if (signal?.aborted) abort();
+		else signal?.addEventListener("abort", abort, { once: true });
+	});
 }
 
 async function requestImage(
 	operation: "generation" | "edit",
 	body: Record<string, unknown>,
-	auth: CodexAuth,
-	signal: AbortSignal | undefined,
+	token: string,
+	signal?: AbortSignal,
 ): Promise<GeneratedImage> {
 	const route = operation === "generation" ? "generations" : "edits";
-	const response = (await fetch(`${CODEX_IMAGES_URL}/${route}`, {
-		method: "POST",
-		headers: {
-			Authorization: `Bearer ${auth.token}`,
-			"chatgpt-account-id": auth.accountId,
-			originator: "pi",
-			Accept: "application/json",
-			"Content-Type": "application/json",
-		},
-		body: JSON.stringify(body),
-		signal,
-	})) as FetchResponse;
-	if (!response.ok) {
-		const message = serverErrorMessage(await readBoundedError(response), auth.token);
-		throw new Error(`Image ${operation} failed with status ${response.status}${message ? `: ${message}` : ""}`);
-	}
-
-	let value: unknown;
-	try {
-		value = await response.json();
-	} catch {
-		throw new Error("OpenAI Codex returned a non-JSON image response");
+	for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
+		const requestSignal = signal
+			? AbortSignal.any([signal, AbortSignal.timeout(REQUEST_TIMEOUT_MS)])
+			: AbortSignal.timeout(REQUEST_TIMEOUT_MS);
+		let response: HttpResponse;
+		try {
+			response = (await fetch(`${XAI_API_BASE_URL}/images/${route}`, {
+				method: "POST",
+				headers: {
+					Accept: "application/json",
+					Authorization: `Bearer ${token}`,
+					"Content-Type": "application/json",
+				},
+				body: JSON.stringify(body),
+				signal: requestSignal,
+			})) as HttpResponse;
+		} catch (error) {
+			if (signal?.aborted) throw signal.reason;
+			if (attempt === MAX_ATTEMPTS) throw error;
+			await delay(RETRY_BASE_DELAY_MS * 2 ** (attempt - 1), signal);
+			continue;
+		}
+		if (response.ok) {
+			let value: unknown;
+			try {
+				value = await response.json();
+			} catch {
+				throw new Error("xAI returned a non-JSON image response");
+			}
+			return decodeImageResponse(value);
+		}
+		const message = serverErrorMessage(await boundedError(response), token);
+		if (!retryable(response.status) || attempt === MAX_ATTEMPTS) {
+			throw new Error(
+				`xAI image ${operation} failed with status ${response.status}${message ? `: ${message}` : ""}`,
+			);
+		}
+		await delay(RETRY_BASE_DELAY_MS * 2 ** (attempt - 1), signal);
 	}
-	return decodeImageResponse(value);
+	throw new Error(`xAI image ${operation} failed`);
 }
 
-export function generateImage(
-	prompt: string,
-	auth: CodexAuth,
-	signal: AbortSignal | undefined,
-): Promise<GeneratedImage> {
+export function generateImage(prompt: string, token: string, signal?: AbortSignal): Promise<GeneratedImage> {
 	return requestImage(
 		"generation",
-		{ prompt, model: MODEL, background: "auto", quality: "auto", size: "auto" },
-		auth,
+		{ model: XAI_IMAGE_MODEL, prompt, n: 1, resolution: "1k", response_format: "b64_json" },
+		token,
 		signal,
 	);
 }
@@ -167,20 +185,21 @@ export function generateImage(
 export function editImage(
 	prompt: string,
 	images: readonly EditImage[],
-	auth: CodexAuth,
-	signal: AbortSignal | undefined,
+	token: string,
+	signal?: AbortSignal,
 ): Promise<GeneratedImage> {
+	const references = images.map((image) => ({ url: `data:${image.mimeType};base64,${image.data}` }));
 	return requestImage(
 		"edit",
 		{
-			images: images.map((image) => ({ image_url: `data:${image.mimeType};base64,${image.data}` })),
+			model: XAI_IMAGE_MODEL,
 			prompt,
-			model: MODEL,
-			background: "auto",
-			quality: "auto",
-			size: "auto",
+			n: 1,
+			resolution: "1k",
+			response_format: "b64_json",
+			...(references.length === 1 ? { image: references[0] } : { images: references }),
 		},
-		auth,
+		token,
 		signal,
 	);
 }