@shanepadgett/tau-agent 0.21.0 → 0.22.0

package/README.md

@@ -2,6 +2,8 @@
 
 Tau is a custom agentic harness built with Pi extensions: tools, commands, prompts, skills, and themes.
 
+The package also exports programmatic Tau capabilities for trusted Pi extensions. The built-in extensions and public API use the same underlying operations.
+
 ## Install
 
 ```bash
@@ -13,6 +15,22 @@ pi install ./path/to/tau-agent
 pi install ./path/to/tau-agent/packages/agent
 ```
 
+## Programmatic use
+
+Install the package in the extension's project, then import from its root:
+
+```ts
+import { generateImage } from "@shanepadgett/tau-agent";
+
+const image = await generateImage(ctx, {
+  prompt: "A quiet mountain lake",
+  path: "assets/lake.png",
+  signal,
+});
+```
+
+See [Extending Tau Agent](./docs/extending-tau-agent.md) for the supported API.
+
 ## Development
 
 From the monorepo root:

package/docs/extending-tau-agent.md

@@ -1,8 +1,40 @@
 # Extending Tau Agent
 
-Tau Agent is a Pi extension harness. External integration uses Pi's native `pi.events` bus.
+Tau Agent is a Pi extension harness and a programmatic API for trusted Pi extensions.
 
-The caller and Tau Agent must be loaded in the same Pi runtime. External callers use string channel names and documented payloads. They do not import Tau Agent internals.
+Install `@shanepadgett/tau-agent` as a dependency when calling its API. Import only from the package root. Files under `extensions/` and other package paths are private.
+
+## Image generation
+
+`generateImage()` runs the same image-generation pipeline as Tau's `image_gen` tool without opening UI, injecting messages, or starting an agent turn.
+
+```ts
+import { generateImage } from "@shanepadgett/tau-agent";
+
+const result = await generateImage(ctx, {
+  prompt: "A stone archway at sunrise",
+  provider: "openai",
+  path: candidatePath,
+  referencedImagePaths: [perspectivePath, palettePath],
+  signal,
+});
+```
+
+The context must provide `cwd`, `model`, and `modelRegistry`; Pi's `ExtensionContext` and `ExtensionCommandContext` already do. Relative destination and reference paths resolve from `ctx.cwd`.
+
+Fields:
+
+- `prompt`: complete non-empty image prompt.
+- `provider`: optional `"openai"` or `"xai"` override. Without one, Tau follows the active model and falls back when authentication is unavailable.
+- `path`: optional destination. Tau uses its external image store when omitted and never overwrites an explicit destination.
+- `referencedImagePaths`: optional list of up to three local PNG, JPEG, or WebP files. Supplying references selects image editing or composition.
+- `signal`: optional cancellation signal.
+
+The promise resolves after publication with the generated bytes, absolute path, provider, model, operation, MIME type, width, and height. The bytes are the same bytes written to `result.path`, so callers can attach the image without reopening the file. Failures reject with a bounded, sanitized error. Cancellation before publication rejects with the signal's abort reason and removes staged files.
+
+## Events
+
+External event integration uses Pi's native `pi.events` bus. The caller and Tau Agent must be loaded in the same Pi runtime. Event callers use string channel names and documented payloads. They do not import Tau Agent internals.
 
 Only events documented in this file are public. Extensions run trusted in-process; event emitters can ask Tau Agent to do work.
 

package/extensions/context-pruning/index.ts

@@ -235,9 +235,7 @@ export default function contextPruningExtension(pi: ExtensionAPI): void {
 		const boundary = baseline + reminder * nudgeEveryPercent;
 		if (boundary <= nudgeState.highestBoundary) return undefined;
 		const tierCount = nudgeInstructions.length;
-		const tierFloor = nudgeState.terminalTierReached
-			? tierCount
-			: Math.min(nudgeState.highestTier, tierCount);
+		const tierFloor = nudgeState.terminalTierReached ? tierCount : Math.min(nudgeState.highestTier, tierCount);
 		const tier = Math.max(Math.min(reminder, tierCount), tierFloor);
 		const instruction = nudgeInstructions[tier - 1] ?? nudgeInstructions[0];
 		const details: ContextPruningNudgeDetailsV2 = {
@@ -328,7 +326,7 @@ function reconstructNudgeState(branch: readonly SessionEntry[], anchorToolCallId
 			(growthBaselinePercent !== undefined && details.growthBaselinePercent !== growthBaselinePercent)
 		)
 			continue;
-		const expectedTierFloor = terminalTierReached
+		const expectedTierFloor: number = terminalTierReached
 			? details.tierCount
 			: Math.min(highestTier, details.tierCount);
 		if (details.boundary <= highestBoundary || details.tierFloor !== expectedTierFloor) continue;

package/extensions/context-pruning/projection.ts

@@ -15,9 +15,7 @@ export function projectContext(
 			message?.role === "assistant" &&
 			message.content.some(
 				(block) =>
-					block.type === "toolCall" &&
-					block.id === state.latestAnchorToolCallId &&
-					block.name === "context_prune",
+					block.type === "toolCall" && block.id === state.latestAnchorToolCallId && block.name === "context_prune",
 			)
 		) {
 			anchorIndex = index;

package/extensions/context-pruning/prune.ts

@@ -1,13 +1,7 @@
 import { resolve } from "node:path";
-import {
-	type ExtensionContext,
-	type SessionEntry,
-} from "@earendil-works/pi-coding-agent";
+import { type ExtensionContext, type SessionEntry } from "@earendil-works/pi-coding-agent";
 import { type Static, Type } from "typebox";
-import type {
-	ContextPruneDeferredFileV2,
-	ContextPruneDetailsV2,
-} from "../../shared/context-pruning-state.ts";
+import type { ContextPruneDeferredFileV2, ContextPruneDetailsV2 } from "../../shared/context-pruning-state.ts";
 import { prepareAutoreadMessage, type PreparedAutoreadMessage } from "../explore/autoread.ts";
 import { MAX_COMPLETE_FILE_SNAPSHOT_BYTES } from "../explore/full-file-knowledge.ts";
 
@@ -119,9 +113,7 @@ export async function executeContextPrune(options: ContextPruneExecutionOptions)
 		content: [
 			{ type: "text", text: status },
 			...preparedSnapshots.map((snapshot) => ({ type: "text" as const, text: snapshot.content })),
-			...(deferredFiles.length === 0
-				? []
-				: [{ type: "text" as const, text: deferredFileText(deferredFiles) }]),
+			...(deferredFiles.length === 0 ? [] : [{ type: "text" as const, text: deferredFileText(deferredFiles) }]),
 		],
 		details,
 	};
@@ -169,7 +161,8 @@ function collectPrunedRows(
 		if (entry.message.toolName !== "context_prune" || !isRecord(entry.message.details)) continue;
 		const files = entry.message.details.refreshedFiles;
 		if (!Array.isArray(files)) continue;
-		for (const file of files) if (isRecord(file) && typeof file.rowId === "string") prunedAutoreadRowIds.add(file.rowId);
+		for (const file of files)
+			if (isRecord(file) && typeof file.rowId === "string") prunedAutoreadRowIds.add(file.rowId);
 	}
 	return { prunedToolCallIds: [...prunedToolCallIds], prunedAutoreadRowIds: [...prunedAutoreadRowIds] };
 }

package/extensions/context-pruning/render.ts

@@ -128,11 +128,7 @@ export function renderContextPruningNudge(details: unknown, theme: Theme): Marke
 				? ["Prune requested."]
 				: [
 						`${parsed.percent}%`,
-						...(parsed.tier === parsed.tierCount
-							? ["Prune now."]
-							: parsed.tier > 1
-								? ["Prune soon."]
-								: []),
+						...(parsed.tier === parsed.tierCount ? ["Prune now."] : parsed.tier > 1 ? ["Prune soon."] : []),
 					],
 	});
 }

package/extensions/handoff/index.ts

@@ -35,10 +35,7 @@ export default function handoffExtension(pi: ExtensionAPI): void {
 				ctx.ui.notify("The current session must be persisted before handoff.", "error");
 				return;
 			}
-			const messages = buildSessionContext(
-				ctx.sessionManager.getEntries(),
-				ctx.sessionManager.getLeafId(),
-			).messages;
+			const messages = buildSessionContext(ctx.sessionManager.getEntries(), ctx.sessionManager.getLeafId()).messages;
 			if (messages.length === 0) {
 				ctx.ui.notify("No conversation context to hand off.", "error");
 				return;
@@ -53,7 +50,13 @@ export default function handoffExtension(pi: ExtensionAPI): void {
 						const thinking = pi.getThinkingLevel();
 						const candidates = await resolveCandidates(
 							ctx,
-							[{ provider: model.provider, model: model.id, reasoning: thinking === "off" ? undefined : thinking }],
+							[
+								{
+									provider: model.provider,
+									model: model.id,
+									reasoning: thinking === "off" ? undefined : thinking,
+								},
+							],
 							false,
 						);
 						const conversation = serializeConversation(convertToLlm(messages));
@@ -82,7 +85,10 @@ export default function handoffExtension(pi: ExtensionAPI): void {
 			});
 
 			if (!draft) {
-				ctx.ui.notify(generationError ? `Handoff failed: ${generationError}` : "Handoff cancelled.", generationError ? "error" : "info");
+				ctx.ui.notify(
+					generationError ? `Handoff failed: ${generationError}` : "Handoff cancelled.",
+					generationError ? "error" : "info",
+				);
 				return;
 			}
 

package/extensions/image-gen/index.ts

@@ -1,20 +1,20 @@
-import { defineTool, withFileMutationQueue, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import { randomUUID } from "node:crypto";
-import { link, mkdir, readFile, rm, stat, writeFile } from "node:fs/promises";
-import { homedir } from "node:os";
-import { basename, dirname, extname, isAbsolute, join, resolve } from "node:path";
+import { defineTool, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import {
+	generateImage,
+	type GeneratedImageResult,
+	type GenerateImageRequest,
+	type ImageGenerationContext,
+	type ImageProvider,
+} from "@shanepadgett/tau-agent";
 import { type Static, Type } from "typebox";
-import { detectImageMimeType, editImage, generateImage, type EditImage, type GeneratedImage } from "./client.ts";
-import { OPENAI_IMAGE_MODEL, OPENAI_PROVIDER, XAI_IMAGE_MODEL, XAI_PROVIDER, type ImageProvider } from "./constants.ts";
 
-const MAX_INPUT_BYTES = 50 * 1024 * 1024;
 const MAX_INLINE_BYTES = 12 * 1024 * 1024;
 
 const imageGenSchema = Type.Object(
 	{
 		prompt: Type.String({ minLength: 1 }),
 		provider: Type.Optional(
-			Type.Union([Type.Literal("openai"), Type.Literal("xai")], {
+			Type.Union([Type.Literal("openai" satisfies ImageProvider), Type.Literal("xai" satisfies ImageProvider)], {
 				description:
 					"Image provider override. Omit to follow the parent model, preferring OpenAI for GPT and xAI for Grok.",
 			}),
@@ -28,19 +28,7 @@ const imageGenSchema = Type.Object(
 );
 
 type ImageGenParams = Static<typeof imageGenSchema>;
-
-interface ImageGenDetails {
-	path: string;
-	provider: ImageProvider;
-	model: typeof OPENAI_IMAGE_MODEL | typeof XAI_IMAGE_MODEL;
-	operation: "generate" | "edit";
-}
-
-function outputExtension(image: GeneratedImage): string {
-	if (image.mimeType === "image/png") return ".png";
-	if (image.mimeType === "image/webp") return ".webp";
-	return ".jpg";
-}
+type ImageGenDetails = Omit<GeneratedImageResult, "bytes">;
 
 export default function imageGenExtension(pi: ExtensionAPI): void {
 	pi.registerTool(
@@ -51,127 +39,30 @@ export default function imageGenExtension(pi: ExtensionAPI): void {
 				"Generate a requested raster image or AI-edit existing images with OpenAI GPT Image or xAI Grok Imagine. Omit provider to follow the parent model; set it to openai or xai to override. Omit referenced_image_paths to generate; pass one to three local paths to edit or compose. Omit path to use Tau's external image store; pass path only when the user explicitly requests a repository file or other destination. Returns the image for inspection.",
 			parameters: imageGenSchema,
 			async execute(_toolCallId, params: ImageGenParams, signal, onUpdate, ctx) {
-				signal?.throwIfAborted();
-				const prompt = params.prompt.trim();
-				if (!prompt) throw new Error("Image prompt cannot be empty");
-				const requestedPath = params.path?.startsWith("@") ? params.path.slice(1) : params.path;
-				if (requestedPath !== undefined && !requestedPath.trim()) throw new Error("Image path cannot be empty");
-				const requestedAbsolutePath = requestedPath
-					? isAbsolute(requestedPath)
-						? requestedPath
-						: resolve(ctx.cwd, requestedPath)
-					: undefined;
-				if (
-					requestedAbsolutePath &&
-					![".jpg", ".jpeg", ".png", ".webp"].includes(extname(requestedAbsolutePath).toLowerCase())
-				) {
-					throw new Error("Image path must end in .jpg, .jpeg, .png, or .webp");
-				}
-
-				const parentUsesXai =
-					ctx.model?.provider.toLowerCase() === XAI_PROVIDER ||
-					ctx.model?.id.toLowerCase().includes("grok") === true;
-				const preferredProvider: ImageProvider = params.provider ?? (parentUsesXai ? "xai" : "openai");
-				const providers: readonly ImageProvider[] = params.provider
-					? [params.provider]
-					: preferredProvider === "xai"
-						? ["xai", "openai"]
-						: ["openai", "xai"];
-				let provider: ImageProvider | undefined;
-				let token: string | undefined;
-				for (const candidate of providers) {
-					const candidateToken = await ctx.modelRegistry.getApiKeyForProvider(
-						candidate === "openai" ? OPENAI_PROVIDER : XAI_PROVIDER,
-					);
-					if (candidateToken) {
-						provider = candidate;
-						token = candidateToken;
-						break;
-					}
-				}
-				if (!provider || !token) {
-					if (params.provider === "openai") {
-						throw new Error("OpenAI Codex authentication is unavailable. Run /login openai-codex.");
-					}
-					if (params.provider === "xai") {
-						throw new Error("xAI authentication is unavailable. Run /login xai and choose a login method.");
-					}
-					throw new Error("Image generation authentication is unavailable. Run /login for OpenAI Codex or xAI.");
-				}
-				const model = provider === "openai" ? OPENAI_IMAGE_MODEL : XAI_IMAGE_MODEL;
-				const images: EditImage[] = [];
-				for (const path of params.referenced_image_paths ?? []) {
-					const rawPath = path.startsWith("@") ? path.slice(1) : path;
-					if (!rawPath.trim()) throw new Error("Referenced image path cannot be empty");
-					const absolutePath = isAbsolute(rawPath) ? rawPath : resolve(ctx.cwd, rawPath);
-					const metadata = await stat(absolutePath);
-					if (!metadata.isFile()) throw new Error(`Referenced image is not a file: ${absolutePath}`);
-					if (metadata.size > MAX_INPUT_BYTES) throw new Error(`Referenced image exceeds 50 MiB: ${absolutePath}`);
-					const bytes = await readFile(absolutePath);
-					if (bytes.length > MAX_INPUT_BYTES) throw new Error(`Referenced image exceeds 50 MiB: ${absolutePath}`);
-					const mimeType = detectImageMimeType(bytes);
-					if (!mimeType)
-						throw new Error(`Referenced image is not a supported PNG, JPEG, or WebP file: ${absolutePath}`);
-					images.push({ mimeType, data: bytes.toString("base64") });
-				}
-
-				const operation = images.length === 0 ? "generate" : "edit";
 				await onUpdate?.({
 					content: [
 						{
 							type: "text",
-							text:
-								operation === "generate"
-									? `Generating image with ${model}...`
-									: `Editing image with ${model}...`,
+							text: "Processing image request...",
 						},
 					],
 					details: undefined,
 				});
-				const generated =
-					operation === "generate"
-						? await generateImage(provider, prompt, token, signal)
-						: await editImage(provider, prompt, images, token, signal);
-				signal?.throwIfAborted();
-				const generatedExtension = outputExtension(generated);
-				if (requestedAbsolutePath) {
-					const requestedExtension = extname(requestedAbsolutePath).toLowerCase();
-					const matches =
-						requestedExtension === generatedExtension ||
-						(generatedExtension === ".jpg" && requestedExtension === ".jpeg");
-					if (!matches)
-						throw new Error(
-							`${provider === "openai" ? "OpenAI Codex" : "xAI"} returned ${generated.mimeType}; destination must end in ${generatedExtension}`,
-						);
-				}
-				const absolutePath =
-					requestedAbsolutePath ??
-					join(homedir(), ".local", "share", "tau-agent", "images", `image-${randomUUID()}${generatedExtension}`);
-
-				const outputDirectory = dirname(absolutePath);
-				await withFileMutationQueue(absolutePath, async () => {
-					await mkdir(outputDirectory, { recursive: true });
-					const temporaryPath = join(
-						outputDirectory,
-						`.${basename(absolutePath)}.${randomUUID()}.tmp${generatedExtension}`,
-					);
-					try {
-						await writeFile(temporaryPath, generated.bytes, { flag: "wx" });
-						signal?.throwIfAborted();
-						await link(temporaryPath, absolutePath);
-					} finally {
-						await rm(temporaryPath, { force: true });
-					}
-				});
-
-				const verb = operation === "generate" ? "Generated" : "Edited";
-				const details: ImageGenDetails = { path: absolutePath, provider, model, operation };
-				if (generated.bytes.length > MAX_INLINE_BYTES) {
+				const request = {
+					prompt: params.prompt,
+					provider: params.provider,
+					path: params.path,
+					referencedImagePaths: params.referenced_image_paths,
+					signal,
+				} satisfies GenerateImageRequest;
+				const { bytes, ...details } = await generateImage(ctx satisfies ImageGenerationContext, request);
+				const verb = details.operation === "generate" ? "Generated" : "Edited";
+				if (bytes.length > MAX_INLINE_BYTES) {
 					return {
 						content: [
 							{
 								type: "text",
-								text: `${verb} image saved to ${absolutePath}. The image exceeds the 12 MiB attachment limit, so it was not added to model context.`,
+								text: `${verb} image saved to ${details.path}. The image exceeds the 12 MiB attachment limit, so it was not added to model context.`,
 							},
 						],
 						details,
@@ -179,8 +70,8 @@ export default function imageGenExtension(pi: ExtensionAPI): void {
 				}
 				return {
 					content: [
-						{ type: "text", text: `${verb} image saved to ${absolutePath}` },
-						{ type: "image", data: generated.base64, mimeType: generated.mimeType },
+						{ type: "text", text: `${verb} image saved to ${details.path}` },
+						{ type: "image", data: bytes.toString("base64"), mimeType: details.mimeType },
 					],
 					details,
 				};

package/package.json

@@ -1,8 +1,13 @@
 {
 	"name": "@shanepadgett/tau-agent",
-	"version": "0.21.0",
+	"version": "0.22.0",
 	"description": "Tau is a custom agentic harness built with pi extensions",
 	"type": "module",
+	"main": "./src/index.ts",
+	"types": "./src/index.ts",
+	"exports": {
+		".": "./src/index.ts"
+	},
 	"license": "MIT",
 	"keywords": [
 		"pi-package",
@@ -18,6 +23,7 @@
 		"access": "public"
 	},
 	"files": [
+		"src/**/*",
 		"extensions/**/*",
 		"shared/**/*",
 		"prompts/**/*",
@@ -28,8 +34,9 @@
 		"README.md"
 	],
 	"dependencies": {
-		"@shanepadgett/tau-tui": "0.21.0",
+		"@shanepadgett/tau-tui": "0.22.0",
 		"@toon-format/toon": "2.3.0",
+		"image-size": "2.0.2",
 		"smol-toml": "1.7.0"
 	},
 	"peerDependencies": {

package/shared/model-fallback/index.ts

@@ -24,7 +24,7 @@ interface ModelFallbackOptions {
 
 export async function resolveCandidates(
 	ctx: Pick<ExtensionContext, "modelRegistry" | "model" | "cwd" | "isProjectTrusted">,
-	preferredModels: ReadonlyArray<{ provider: string; model: string; reasoning: ThinkingLevel }>,
+	preferredModels: ReadonlyArray<{ provider: string; model: string; reasoning: ThinkingLevel | undefined }>,
 	includeParentModel: boolean,
 ): Promise<ModelCandidate[]> {
 	const settings = await loadTauExtensionSettings(ctx, modelFallbackSettings);

package/{extensions/image-gen → src/image-generation}/client.ts

@@ -24,7 +24,6 @@ export interface EditImage {
 
 export interface GeneratedImage {
 	bytes: Buffer;
-	base64: string;
 	mimeType: EditImage["mimeType"];
 }
 
@@ -58,7 +57,7 @@ export function resolveCodexAuth(token: string): CodexAuth {
 	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();
+	if (typeof accountId !== "string" || !/^[\x21-\x7e]+$/.test(accountId.trim())) throw invalidCredential();
 	return { token, accountId: accountId.trim() };
 }
 
@@ -87,7 +86,7 @@ async function boundedError(response: HttpResponse): Promise<string> {
 		.trim();
 }
 
-function serverErrorMessage(body: string, token: string): string {
+function serverErrorMessage(body: string, secrets: readonly string[]): string {
 	let message = body;
 	try {
 		const value: unknown = JSON.parse(body);
@@ -98,7 +97,10 @@ function serverErrorMessage(body: string, token: string): string {
 	} catch {
 		// Plain-text error body.
 	}
-	return message.replaceAll(token, "[redacted]").slice(0, MAX_ERROR_MESSAGE_LENGTH).trim();
+	for (const secret of secrets) {
+		if (secret) message = message.replaceAll(secret, "[redacted]");
+	}
+	return message.slice(0, MAX_ERROR_MESSAGE_LENGTH).trim();
 }
 
 export function detectImageMimeType(bytes: Buffer): GeneratedImage["mimeType"] | undefined {
@@ -137,7 +139,7 @@ function decodeImageResponse(value: unknown, service: "OpenAI Codex" | "xAI", pn
 	if (!mimeType || (pngOnly && mimeType !== "image/png")) {
 		throw new Error(`${service} returned unsupported image data`);
 	}
-	return { bytes, base64: bytes.toString("base64"), mimeType };
+	return { bytes, mimeType };
 }
 
 async function requestOpenAIImage(
@@ -147,20 +149,28 @@ async function requestOpenAIImage(
 	signal?: AbortSignal,
 ): Promise<GeneratedImage> {
 	const route = operation === "generation" ? "generations" : "edits";
-	const response = (await fetch(`${OPENAI_IMAGE_API_BASE_URL}/${route}`, {
-		method: "POST",
-		headers: {
-			Accept: "application/json",
-			Authorization: `Bearer ${auth.token}`,
-			"chatgpt-account-id": auth.accountId,
-			"Content-Type": "application/json",
-			originator: "pi",
-		},
-		body: JSON.stringify(body),
-		signal,
-	})) as HttpResponse;
+	let response: HttpResponse;
+	try {
+		response = (await fetch(`${OPENAI_IMAGE_API_BASE_URL}/${route}`, {
+			method: "POST",
+			headers: {
+				Accept: "application/json",
+				Authorization: `Bearer ${auth.token}`,
+				"chatgpt-account-id": auth.accountId,
+				"Content-Type": "application/json",
+				originator: "pi",
+			},
+			body: JSON.stringify(body),
+			signal,
+		})) as HttpResponse;
+	} catch (error) {
+		if (signal?.aborted) throw signal.reason;
+		let message = error instanceof Error ? error.message : "OpenAI Codex image request failed";
+		for (const secret of [auth.token, auth.accountId]) message = message.replaceAll(secret, "[redacted]");
+		throw new Error(message.slice(0, MAX_ERROR_MESSAGE_LENGTH).trim() || "OpenAI Codex image request failed");
+	}
 	if (!response.ok) {
-		const message = serverErrorMessage(await boundedError(response), auth.token);
+		const message = serverErrorMessage(await boundedError(response), [auth.token, auth.accountId]);
 		throw new Error(
 			`OpenAI Codex image ${operation} failed with status ${response.status}${message ? `: ${message}` : ""}`,
 		);
@@ -232,7 +242,7 @@ async function requestXaiImage(
 			}
 			return decodeImageResponse(value, "xAI", false);
 		}
-		const message = serverErrorMessage(await boundedError(response), token);
+		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}` : ""}`,
@@ -243,7 +253,7 @@ async function requestXaiImage(
 	throw new Error(`xAI image ${operation} failed`);
 }
 
-export function generateImage(
+export function requestGeneratedImage(
 	provider: ImageProvider,
 	prompt: string,
 	token: string,

package/src/image-generation/index.ts

@@ -0,0 +1,224 @@
+import { withFileMutationQueue, type ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { imageSize } from "image-size";
+import { randomUUID } from "node:crypto";
+import { link, mkdir, readFile, rm, stat, writeFile } from "node:fs/promises";
+import { homedir } from "node:os";
+import { basename, dirname, extname, isAbsolute, join, resolve } from "node:path";
+import { detectImageMimeType, editImage, requestGeneratedImage } from "./client.ts";
+import { OPENAI_IMAGE_MODEL, OPENAI_PROVIDER, XAI_IMAGE_MODEL, XAI_PROVIDER, type ImageProvider } from "./constants.ts";
+
+export type { ImageProvider } from "./constants.ts";
+
+const MAX_INPUT_BYTES = 50 * 1024 * 1024;
+const MAX_ERROR_MESSAGE_LENGTH = 2000;
+
+export type ImageGenerationContext = Pick<ExtensionContext, "cwd" | "model" | "modelRegistry">;
+
+export interface GenerateImageRequest {
+	prompt: string;
+	provider?: ImageProvider;
+	path?: string;
+	referencedImagePaths?: readonly string[];
+	signal?: AbortSignal;
+}
+
+export interface GeneratedImageResult {
+	bytes: Buffer;
+	path: string;
+	provider: ImageProvider;
+	model: typeof OPENAI_IMAGE_MODEL | typeof XAI_IMAGE_MODEL;
+	operation: "generate" | "edit";
+	mimeType: "image/png" | "image/jpeg" | "image/webp";
+	width: number;
+	height: number;
+}
+
+function outputExtension(mimeType: GeneratedImageResult["mimeType"]): string {
+	if (mimeType === "image/png") return ".png";
+	if (mimeType === "image/webp") return ".webp";
+	return ".jpg";
+}
+
+function sanitizedError(error: unknown, signal: AbortSignal | undefined, secrets: readonly string[]): unknown {
+	if (signal?.aborted) return signal.reason;
+	let message = error instanceof Error ? error.message : "Image generation failed";
+	for (const secret of secrets) {
+		if (secret) message = message.replaceAll(secret, "[redacted]");
+	}
+	message = [...message]
+		.filter((character) => {
+			const code = character.charCodeAt(0);
+			return code === 9 || code === 10 || code === 13 || (code >= 32 && code !== 127);
+		})
+		.join("")
+		.slice(0, MAX_ERROR_MESSAGE_LENGTH)
+		.trim();
+	return new Error(message || "Image generation failed");
+}
+
+export async function generateImage(
+	ctx: ImageGenerationContext,
+	request: GenerateImageRequest,
+): Promise<GeneratedImageResult> {
+	const signal = request.signal;
+	const secrets: string[] = [];
+	try {
+		signal?.throwIfAborted();
+		if (typeof request.prompt !== "string") throw new Error("Image prompt must be a string");
+		if (request.provider !== undefined && request.provider !== "openai" && request.provider !== "xai") {
+			throw new Error("Image provider must be openai or xai");
+		}
+		if (request.path !== undefined && typeof request.path !== "string")
+			throw new Error("Image path must be a string");
+		if (request.referencedImagePaths !== undefined && !Array.isArray(request.referencedImagePaths)) {
+			throw new Error("Referenced image paths must be an array");
+		}
+		const prompt = request.prompt.trim();
+		if (!prompt) throw new Error("Image prompt cannot be empty");
+		if (request.referencedImagePaths && request.referencedImagePaths.length > 3) {
+			throw new Error("Image generation accepts at most three referenced images");
+		}
+
+		const requestedPath = request.path?.startsWith("@") ? request.path.slice(1) : request.path;
+		if (requestedPath !== undefined && !requestedPath.trim()) throw new Error("Image path cannot be empty");
+		const requestedAbsolutePath = requestedPath
+			? isAbsolute(requestedPath)
+				? requestedPath
+				: resolve(ctx.cwd, requestedPath)
+			: undefined;
+		if (
+			requestedAbsolutePath &&
+			![".jpg", ".jpeg", ".png", ".webp"].includes(extname(requestedAbsolutePath).toLowerCase())
+		) {
+			throw new Error("Image path must end in .jpg, .jpeg, .png, or .webp");
+		}
+
+		const parentUsesXai =
+			ctx.model?.provider.toLowerCase() === XAI_PROVIDER || ctx.model?.id.toLowerCase().includes("grok") === true;
+		const preferredProvider: ImageProvider = request.provider ?? (parentUsesXai ? "xai" : "openai");
+		const providers: readonly ImageProvider[] = request.provider
+			? [request.provider]
+			: preferredProvider === "xai"
+				? ["xai", "openai"]
+				: ["openai", "xai"];
+		let provider: ImageProvider | undefined;
+		let token: string | undefined;
+		for (const candidate of providers) {
+			signal?.throwIfAborted();
+			const candidateToken = await ctx.modelRegistry.getApiKeyForProvider(
+				candidate === "openai" ? OPENAI_PROVIDER : XAI_PROVIDER,
+			);
+			signal?.throwIfAborted();
+			if (candidateToken) {
+				provider = candidate;
+				token = candidateToken;
+				secrets.push(candidateToken);
+				break;
+			}
+		}
+		if (!provider || !token) {
+			if (request.provider === "openai") {
+				throw new Error("OpenAI Codex authentication is unavailable. Run /login openai-codex.");
+			}
+			if (request.provider === "xai") {
+				throw new Error("xAI authentication is unavailable. Run /login xai and choose a login method.");
+			}
+			throw new Error("Image generation authentication is unavailable. Run /login for OpenAI Codex or xAI.");
+		}
+
+		const images: Array<{
+			mimeType: "image/png" | "image/jpeg" | "image/webp";
+			data: string;
+		}> = [];
+		for (const path of request.referencedImagePaths ?? []) {
+			signal?.throwIfAborted();
+			if (typeof path !== "string") throw new Error("Referenced image path must be a string");
+			const rawPath = path.startsWith("@") ? path.slice(1) : path;
+			if (!rawPath.trim()) throw new Error("Referenced image path cannot be empty");
+			const absolutePath = isAbsolute(rawPath) ? rawPath : resolve(ctx.cwd, rawPath);
+			const metadata = await stat(absolutePath);
+			signal?.throwIfAborted();
+			if (!metadata.isFile()) throw new Error(`Referenced image is not a file: ${absolutePath}`);
+			if (metadata.size > MAX_INPUT_BYTES) throw new Error(`Referenced image exceeds 50 MiB: ${absolutePath}`);
+			const bytes = await readFile(absolutePath, { signal });
+			if (bytes.length > MAX_INPUT_BYTES) throw new Error(`Referenced image exceeds 50 MiB: ${absolutePath}`);
+			const mimeType = detectImageMimeType(bytes);
+			if (!mimeType) {
+				throw new Error(`Referenced image is not a supported PNG, JPEG, or WebP file: ${absolutePath}`);
+			}
+			images.push({ mimeType, data: bytes.toString("base64") });
+		}
+
+		const operation = images.length === 0 ? "generate" : "edit";
+		const model = provider === "openai" ? OPENAI_IMAGE_MODEL : XAI_IMAGE_MODEL;
+		const selectedProvider = provider;
+		const selectedToken = token;
+		signal?.throwIfAborted();
+		const generated =
+			operation === "generate"
+				? await requestGeneratedImage(selectedProvider, prompt, selectedToken, signal)
+				: await editImage(selectedProvider, prompt, images, selectedToken, signal);
+		signal?.throwIfAborted();
+		const generatedExtension = outputExtension(generated.mimeType);
+		if (requestedAbsolutePath) {
+			const requestedExtension = extname(requestedAbsolutePath).toLowerCase();
+			const matches =
+				requestedExtension === generatedExtension ||
+				(generatedExtension === ".jpg" && requestedExtension === ".jpeg");
+			if (!matches) {
+				throw new Error(
+					`${selectedProvider === "openai" ? "OpenAI Codex" : "xAI"} returned ${generated.mimeType}; destination must end in ${generatedExtension}`,
+				);
+			}
+		}
+
+		let dimensions: ReturnType<typeof imageSize>;
+		try {
+			dimensions = imageSize(generated.bytes);
+		} catch {
+			throw new Error(`${selectedProvider === "openai" ? "OpenAI Codex" : "xAI"} returned invalid image dimensions`);
+		}
+		if (
+			!Number.isInteger(dimensions.width) ||
+			dimensions.width <= 0 ||
+			!Number.isInteger(dimensions.height) ||
+			dimensions.height <= 0
+		) {
+			throw new Error(`${selectedProvider === "openai" ? "OpenAI Codex" : "xAI"} returned invalid image dimensions`);
+		}
+
+		const absolutePath =
+			requestedAbsolutePath ??
+			join(homedir(), ".local", "share", "tau-agent", "images", `image-${randomUUID()}${generatedExtension}`);
+		const outputDirectory = dirname(absolutePath);
+		await withFileMutationQueue(absolutePath, async () => {
+			await mkdir(outputDirectory, { recursive: true });
+			const temporaryPath = join(
+				outputDirectory,
+				`.${basename(absolutePath)}.${randomUUID()}.tmp${generatedExtension}`,
+			);
+			try {
+				await writeFile(temporaryPath, generated.bytes, { flag: "wx", signal });
+				signal?.throwIfAborted();
+				await link(temporaryPath, absolutePath);
+			} catch (error) {
+				await rm(temporaryPath, { force: true }).catch(() => undefined);
+				throw error;
+			}
+			await rm(temporaryPath, { force: true }).catch(() => undefined);
+		});
+
+		return {
+			bytes: generated.bytes,
+			path: absolutePath,
+			provider: selectedProvider,
+			model,
+			operation,
+			mimeType: generated.mimeType,
+			width: dimensions.width,
+			height: dimensions.height,
+		};
+	} catch (error) {
+		throw sanitizedError(error, signal, secrets);
+	}
+}

package/src/index.ts

@@ -0,0 +1,7 @@
+export {
+	generateImage,
+	type GeneratedImageResult,
+	type GenerateImageRequest,
+	type ImageGenerationContext,
+	type ImageProvider,
+} from "./image-generation/index.ts";