@shanepadgett/tau-agent 0.1.0

package/extensions/image-gen/client.ts

@@ -0,0 +1,186 @@
+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]);
+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;
+}
+
+export interface EditImage {
+	mimeType: "image/png" | "image/jpeg" | "image/webp";
+	data: string;
+}
+
+interface GeneratedImage {
+	bytes: Buffer;
+	base64: string;
+	mimeType: "image/png";
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+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() };
+}
+
+async function readBoundedError(response: FetchResponse): Promise<string> {
+	if (!response.body) return "";
+	const reader = response.body.getReader();
+	const chunks: Uint8Array[] = [];
+	let length = 0;
+	try {
+		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);
+			chunks.push(chunk);
+			length += chunk.length;
+			if (chunk.length < result.value.length) break;
+		}
+	} 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);
+}
+
+function serverErrorMessage(body: string, token: string): string {
+	let message = body.trim();
+	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;
+		}
+	} catch {
+		// Plain-text error body.
+	}
+	return message.replaceAll(token, "[redacted]").slice(0, MAX_ERROR_MESSAGE_LENGTH).trim();
+}
+
+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");
+	}
+	const encoded = value.data[0].b64_json;
+	if (typeof encoded !== "string" || !encoded.trim()) {
+		throw new Error("OpenAI Codex 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");
+	}
+	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");
+	}
+	return { bytes, base64: bytes.toString("base64"), mimeType: "image/png" };
+}
+
+async function requestImage(
+	operation: "generation" | "edit",
+	body: Record<string, unknown>,
+	auth: CodexAuth,
+	signal: AbortSignal | undefined,
+): 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");
+	}
+	return decodeImageResponse(value);
+}
+
+export function generateImage(
+	prompt: string,
+	auth: CodexAuth,
+	signal: AbortSignal | undefined,
+): Promise<GeneratedImage> {
+	return requestImage(
+		"generation",
+		{ prompt, model: MODEL, background: "auto", quality: "auto", size: "auto" },
+		auth,
+		signal,
+	);
+}
+
+export function editImage(
+	prompt: string,
+	images: readonly EditImage[],
+	auth: CodexAuth,
+	signal: AbortSignal | undefined,
+): Promise<GeneratedImage> {
+	return requestImage(
+		"edit",
+		{
+			images: images.map((image) => ({ image_url: `data:${image.mimeType};base64,${image.data}` })),
+			prompt,
+			model: MODEL,
+			background: "auto",
+			quality: "auto",
+			size: "auto",
+		},
+		auth,
+		signal,
+	);
+}

package/extensions/image-gen/index.ts

@@ -0,0 +1,152 @@
+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, isAbsolute, join, resolve } from "node:path";
+import { type Static, Type } from "typebox";
+import { editImage, generateImage, resolveCodexAuth, type EditImage } from "./client.ts";
+
+const MODEL = "gpt-image-2";
+const MAX_INPUT_BYTES = 50 * 1024 * 1024;
+const MAX_INLINE_BYTES = 12 * 1024 * 1024;
+const PNG_SIGNATURE = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]);
+
+const imageGenSchema = Type.Object(
+	{
+		prompt: Type.String({ minLength: 1 }),
+		path: Type.Optional(
+			Type.String({ description: "Explicit PNG destination path; defaults to Tau's external image store" }),
+		),
+		referenced_image_paths: Type.Optional(Type.Array(Type.String({ minLength: 1 }), { minItems: 1, maxItems: 5 })),
+	},
+	{ additionalProperties: false },
+);
+
+type ImageGenParams = Static<typeof imageGenSchema>;
+
+interface ImageGenDetails {
+	path: string;
+	model: typeof MODEL;
+	operation: "generate" | "edit";
+}
+
+function detectImageMimeType(bytes: Buffer): EditImage["mimeType"] | undefined {
+	if (bytes.length >= PNG_SIGNATURE.length && bytes.subarray(0, PNG_SIGNATURE.length).equals(PNG_SIGNATURE)) {
+		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;
+}
+
+export default function imageGenExtension(pi: ExtensionAPI): void {
+	pi.registerTool(
+		defineTool<typeof imageGenSchema, ImageGenDetails | undefined>({
+			name: "image_gen",
+			label: "Image Generation",
+			description:
+				"Generate a new raster image from a prompt, or edit one to five local raster images. Uses the existing OpenAI Codex OAuth login, saves a PNG under Tau's external image store unless an explicit path is provided, and returns the image for inspection.",
+			promptSnippet: "Generate or edit raster images with OpenAI Codex",
+			promptGuidelines: [
+				"Use image_gen when the user asks for a generated raster image or an AI edit of local raster images.",
+				"Omit referenced_image_paths when image_gen should create a new image.",
+				"Pass one to five local paths in referenced_image_paths when image_gen should edit or compose existing images.",
+				"Omit path for temporary external storage. Pass path only when the user wants the generated image saved in their repository or another explicit location.",
+			],
+			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 absolutePath = requestedPath
+					? isAbsolute(requestedPath)
+						? requestedPath
+						: resolve(ctx.cwd, requestedPath)
+					: join(homedir(), ".local", "share", "tau-agent", "images", `image-${randomUUID()}.png`);
+				if (!absolutePath.toLowerCase().endsWith(".png")) throw new Error("Image path must end in .png");
+
+				const token = await ctx.modelRegistry.getApiKeyForProvider("openai-codex");
+				if (!token) {
+					throw new Error("OpenAI Codex OAuth is unavailable. Run /login and select OpenAI Codex.");
+				}
+				const auth = resolveCodexAuth(token);
+				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}...`,
+						},
+					],
+					details: undefined,
+				});
+				const generated =
+					operation === "generate"
+						? await generateImage(prompt, auth, signal)
+						: await editImage(prompt, images, auth, signal);
+				signal?.throwIfAborted();
+
+				const outputDirectory = dirname(absolutePath);
+				await withFileMutationQueue(absolutePath, async () => {
+					await mkdir(outputDirectory, { recursive: true });
+					const temporaryPath = join(outputDirectory, `.${basename(absolutePath)}.${randomUUID()}.tmp.png`);
+					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, model: MODEL, operation };
+				if (generated.bytes.length > MAX_INLINE_BYTES) {
+					return {
+						content: [
+							{
+								type: "text",
+								text: `${verb} image saved to ${absolutePath}. The PNG exceeds the 12 MiB attachment limit, so it was not added to model context.`,
+							},
+						],
+						details,
+					};
+				}
+				return {
+					content: [
+						{ type: "text", text: `${verb} image saved to ${absolutePath}` },
+						{ type: "image", data: generated.base64, mimeType: generated.mimeType },
+					],
+					details,
+				};
+			},
+		}),
+	);
+}

package/extensions/manage-sessions/README.md

@@ -0,0 +1,12 @@
+# Manage sessions
+
+Clean up saved Pi sessions from the TUI.
+
+## Commands
+
+- `/manage-sessions` opens a session manager with active and archived sessions.
+- `/sweep` starts a new blank session, then archives or deletes the session you were just using.
+
+The manager lets you archive, unarchive, and delete saved sessions in bulk. It can show sessions for the current folder or all projects.
+
+Archived sessions stay out of the active session list and can be restored later.

package/extensions/manage-sessions/index.ts

@@ -0,0 +1,58 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { showSessionManager } from "./manager-ui.ts";
+import { archiveSession, deleteSessionFile } from "./sessions.ts";
+
+export default function manageSessionsExtension(pi: ExtensionAPI): void {
+	pi.registerCommand("manage-sessions", {
+		description: "Manage saved sessions in the TUI",
+		handler: async (_args, ctx) => {
+			await ctx.waitForIdle();
+
+			if (ctx.mode !== "tui") {
+				ctx.ui.notify("Session manager requires TUI mode.", "error");
+				return;
+			}
+
+			await showSessionManager(ctx);
+		},
+	});
+
+	pi.registerCommand("sweep", {
+		description: "Archive or delete the current session after starting a new one",
+		handler: async (_args, ctx) => {
+			await ctx.waitForIdle();
+
+			if (!ctx.hasUI) {
+				ctx.ui.notify("Sweep requires interactive UI.", "error");
+				return;
+			}
+
+			const sessionFile = ctx.sessionManager.getSessionFile();
+			if (!sessionFile) {
+				ctx.ui.notify("Current session is not persisted.", "error");
+				return;
+			}
+
+			const choice = await ctx.ui.select("Sweep current session:", ["Archive", "Delete"]);
+			if (choice !== "Archive" && choice !== "Delete") return;
+
+			const confirmed = await ctx.ui.confirm(`${choice} this session?`, "Will launch a new session.");
+			if (!confirmed) return;
+
+			await ctx.newSession({
+				parentSession: sessionFile,
+				withSession: async (newCtx) => {
+					try {
+						if (choice === "Archive") await archiveSession(sessionFile);
+						else await deleteSessionFile(sessionFile);
+
+						newCtx.ui.notify(`Session ${choice.toLowerCase()}d.`, "info");
+					} catch (error) {
+						const message = error instanceof Error ? error.message : String(error);
+						newCtx.ui.notify(`Sweep failed: ${message}`, "error");
+					}
+				},
+			});
+		},
+	});
+}