@shanepadgett/tau-agent 0.33.0 → 0.34.0

package/src/image-generation/index.ts

@@ -11,6 +11,7 @@ export type { ImageProvider } from "./constants.ts";
 
 const MAX_INPUT_BYTES = 50 * 1024 * 1024;
 const MAX_ERROR_MESSAGE_LENGTH = 2000;
+const OUTPUT_EXTENSIONS = new Set([".jpg", ".jpeg", ".png", ".webp"]);
 
 export type ImageGenerationContext = Pick<ExtensionContext, "cwd" | "model" | "modelRegistry">;
 
@@ -33,12 +34,21 @@ export interface GeneratedImageResult {
 	height: number;
 }
 
+type ReferencedImage = {
+	mimeType: "image/png" | "image/jpeg" | "image/webp";
+	data: string;
+};
+
 function outputExtension(mimeType: GeneratedImageResult["mimeType"]): string {
 	if (mimeType === "image/png") return ".png";
 	if (mimeType === "image/webp") return ".webp";
 	return ".jpg";
 }
 
+function providerLabel(provider: ImageProvider): string {
+	return provider === "openai" ? "OpenAI Codex" : "xAI";
+}
+
 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";
@@ -56,162 +66,195 @@ function sanitizedError(error: unknown, signal: AbortSignal | undefined, secrets
 	return new Error(message || "Image generation failed");
 }
 
-export async function generateImage(
+function parsePrompt(request: GenerateImageRequest): string {
+	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");
+	}
+	return prompt;
+}
+
+function resolveRequestedAbsolutePath(cwd: string, path: string | undefined): string | undefined {
+	const requestedPath = path?.startsWith("@") ? path.slice(1) : path;
+	if (requestedPath !== undefined && !requestedPath.trim()) throw new Error("Image path cannot be empty");
+	if (!requestedPath) return undefined;
+	const absolutePath = isAbsolute(requestedPath) ? requestedPath : resolve(cwd, requestedPath);
+	if (!OUTPUT_EXTENSIONS.has(extname(absolutePath).toLowerCase())) {
+		throw new Error("Image path must end in .jpg, .jpeg, .png, or .webp");
+	}
+	return absolutePath;
+}
+
+function providerCandidates(
+	requestProvider: ImageProvider | undefined,
+	parentUsesXai: boolean,
+): readonly ImageProvider[] {
+	const preferred: ImageProvider = requestProvider ?? (parentUsesXai ? "xai" : "openai");
+	if (requestProvider) return [requestProvider];
+	return preferred === "xai" ? ["xai", "openai"] : ["openai", "xai"];
+}
+
+async function resolveProviderAuth(
 	ctx: ImageGenerationContext,
-	request: GenerateImageRequest,
-): Promise<GeneratedImageResult> {
-	const signal = request.signal;
-	const secrets: string[] = [];
-	try {
+	requestProvider: ImageProvider | undefined,
+	signal: AbortSignal | undefined,
+	secrets: string[],
+): Promise<{ provider: ImageProvider; token: string }> {
+	const parentUsesXai =
+		ctx.model?.provider.toLowerCase() === XAI_PROVIDER || ctx.model?.id.toLowerCase().includes("grok") === true;
+	for (const candidate of providerCandidates(requestProvider, parentUsesXai)) {
 		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 candidateToken = await ctx.modelRegistry.getApiKeyForProvider(
+			candidate === "openai" ? OPENAI_PROVIDER : XAI_PROVIDER,
+		);
+		signal?.throwIfAborted();
+		if (!candidateToken) continue;
+		secrets.push(candidateToken);
+		return { provider: candidate, token: candidateToken };
+	}
+	if (requestProvider === "openai") {
+		throw new Error("OpenAI Codex authentication is unavailable. Run /login openai-codex.");
+	}
+	if (requestProvider === "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 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");
-		}
+async function loadOneReferencedImage(
+	cwd: string,
+	path: string,
+	signal: AbortSignal | undefined,
+): Promise<ReferencedImage> {
+	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(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}`);
+	}
+	return { mimeType, data: bytes.toString("base64") };
+}
 
-		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,
+async function loadReferencedImages(
+	cwd: string,
+	paths: readonly string[] | undefined,
+	signal: AbortSignal | undefined,
+): Promise<ReferencedImage[]> {
+	const images: ReferencedImage[] = [];
+	for (const path of paths ?? []) images.push(await loadOneReferencedImage(cwd, path, signal));
+	return images;
+}
+
+function readImageDimensions(bytes: Buffer, provider: ImageProvider): { width: number; height: number } {
+	let dimensions: ReturnType<typeof imageSize>;
+	try {
+		dimensions = imageSize(bytes);
+	} catch {
+		throw new Error(`${providerLabel(provider)} returned invalid image dimensions`);
+	}
+	if (
+		!Number.isInteger(dimensions.width) ||
+		dimensions.width <= 0 ||
+		!Number.isInteger(dimensions.height) ||
+		dimensions.height <= 0
+	) {
+		throw new Error(`${providerLabel(provider)} returned invalid image dimensions`);
+	}
+	return { width: dimensions.width, height: dimensions.height };
+}
+
+function assertDestinationMatchesMime(
+	requestedAbsolutePath: string | undefined,
+	mimeType: GeneratedImageResult["mimeType"],
+	provider: ImageProvider,
+): string {
+	const generatedExtension = outputExtension(mimeType);
+	if (requestedAbsolutePath) {
+		const requestedExtension = extname(requestedAbsolutePath).toLowerCase();
+		const matches =
+			requestedExtension === generatedExtension || (generatedExtension === ".jpg" && requestedExtension === ".jpeg");
+		if (!matches) {
+			throw new Error(
+				`${providerLabel(provider)} returned ${mimeType}; destination must end in ${generatedExtension}`,
 			);
-			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.");
 		}
+	}
+	return generatedExtension;
+}
 
-		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);
+async function writeImageAtomically(
+	absolutePath: string,
+	bytes: Buffer,
+	generatedExtension: string,
+	signal: AbortSignal | undefined,
+): Promise<void> {
+	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, bytes, { flag: "wx", signal });
 			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") });
+			await link(temporaryPath, absolutePath);
+		} catch (error) {
+			await rm(temporaryPath, { force: true }).catch(() => undefined);
+			throw error;
 		}
+		await rm(temporaryPath, { force: true }).catch(() => undefined);
+	});
+}
 
+export async function generateImage(
+	ctx: ImageGenerationContext,
+	request: GenerateImageRequest,
+): Promise<GeneratedImageResult> {
+	const signal = request.signal;
+	const secrets: string[] = [];
+	try {
+		signal?.throwIfAborted();
+		const prompt = parsePrompt(request);
+		const requestedAbsolutePath = resolveRequestedAbsolutePath(ctx.cwd, request.path);
+		const { provider, token } = await resolveProviderAuth(ctx, request.provider, signal, secrets);
+		const images = await loadReferencedImages(ctx.cwd, request.referencedImagePaths, signal);
 		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);
+				? await requestGeneratedImage(provider, prompt, token, signal)
+				: await editImage(provider, prompt, images, token, 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 generatedExtension = assertDestinationMatchesMime(requestedAbsolutePath, generated.mimeType, provider);
+		const dimensions = readImageDimensions(generated.bytes, provider);
 		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);
-		});
-
+		await writeImageAtomically(absolutePath, generated.bytes, generatedExtension, signal);
 		return {
 			bytes: generated.bytes,
 			path: absolutePath,
-			provider: selectedProvider,
+			provider,
 			model,
 			operation,
 			mimeType: generated.mimeType,

package/src/index.ts

@@ -5,3 +5,5 @@ export {
 	type ImageGenerationContext,
 	type ImageProvider,
 } from "./image-generation/index.ts";
+export { FILE_INJECTION_TYPE, prepareFileInjection } from "./file-injection/index.ts";
+export { registerDeferredToolGroup, type DeferredToolGroup } from "./tool-loading/index.ts";

package/src/tool-loading/index.ts

@@ -0,0 +1,58 @@
+import type { ExtensionAPI, ToolDefinition } from "@earendil-works/pi-coding-agent";
+import { emitTauEvent, onTauEventImmediately } from "../../shared/events.ts";
+
+export interface DeferredToolGroup {
+	id: string;
+	description: string;
+	tools: readonly ToolDefinition[];
+}
+
+export type DeferredToolHost = Pick<ExtensionAPI, "events" | "on" | "registerTool">;
+
+export interface DeferredToolGroupInfo {
+	id: string;
+	description: string;
+	toolNames: readonly string[];
+}
+
+// Pi evaluates each extension in an isolated module graph, so module-local registries are not shared.
+const DEFERRED_TOOL_GROUP_REQUEST_EVENT = "tau:deferred-tool-group.request";
+
+export function registerDeferredToolGroup(pi: DeferredToolHost, group: DeferredToolGroup): void {
+	const id = group.id.trim();
+	if (id.length === 0 || id !== group.id) throw new Error("Deferred tool group ID must be non-empty and trimmed");
+	if (group.description.trim().length === 0) throw new Error(`Deferred tool group ${id} needs a description`);
+	if (group.tools.length === 0) throw new Error(`Deferred tool group ${id} needs at least one tool`);
+
+	const toolNames = group.tools.map((tool) => tool.name);
+	if (new Set(toolNames).size !== toolNames.length) {
+		throw new Error(`Deferred tool group ${id} contains duplicate tool names`);
+	}
+	if (toolNames.includes("load_tools")) throw new Error("load_tools cannot be part of a deferred tool group");
+
+	const groups = registeredDeferredToolGroups(pi);
+	if (groups.some((existing) => existing.id === id)) {
+		throw new Error(`Deferred tool group is already registered: ${id}`);
+	}
+	for (const existing of groups) {
+		if (toolNames.some((name) => existing.toolNames.includes(name))) {
+			throw new Error(`Deferred tool group ${id} overlaps an existing tool group`);
+		}
+	}
+
+	for (const tool of group.tools) pi.registerTool(tool);
+	const info = { id, description: group.description, toolNames } satisfies DeferredToolGroupInfo;
+	onTauEventImmediately(pi, `deferred-tool-group.${id}`, DEFERRED_TOOL_GROUP_REQUEST_EVENT, ({ accept }) => {
+		accept(info);
+	});
+}
+
+export function registeredDeferredToolGroups(pi: Pick<ExtensionAPI, "events">): readonly DeferredToolGroupInfo[] {
+	const groups: DeferredToolGroupInfo[] = [];
+	emitTauEvent(pi, DEFERRED_TOOL_GROUP_REQUEST_EVENT, {
+		accept(group: DeferredToolGroupInfo) {
+			groups.push(group);
+		},
+	});
+	return groups;
+}