@prometheus-ai/agent 0.5.4 → 0.5.8

package/src/tools/image-gen.ts

@@ -1,16 +1,16 @@
 import * as os from "node:os";
 import * as path from "node:path";
-import { getAntigravityUserAgent, getEnvApiKey, type Model } from "@prometheus-ai/ai";
+import { type ApiKey, type FetchImpl, getEnvApiKey, type Model, ProviderHttpError, withAuth } from "@prometheus-ai/ai";
 import {
 	CODEX_BASE_URL,
 	getCodexAccountId,
 	OPENAI_HEADER_VALUES,
 	OPENAI_HEADERS,
 	URL_PATHS,
-} from "@prometheus-ai/ai/providers/openai-codex/constants";
+} from "@prometheus-ai/catalog/wire/codex";
+import { getAntigravityUserAgent } from "@prometheus-ai/catalog/wire/gemini-headers";
 import {
 	$env,
-	APP_DISPLAY_NAME,
 	isEnoent,
 	parseImageMetadata,
 	prompt,
@@ -19,11 +19,12 @@ import {
 	Snowflake,
 	untilAborted,
 } from "@prometheus-ai/utils";
-import * as z from "zod/v4";
+import { z } from "zod/v4";
 import packageJson from "../../package.json" with { type: "json" };
+
 import { isAuthenticated, type ModelRegistry } from "../config/model-registry";
 import type { CustomTool } from "../extensibility/custom-tools/types";
-import { prometheusXAIUserAgent, resolveXAIHttpCredentials } from "../lib/xai-http";
+import { getPrometheusXAIUserAgent, resolveXAIHttpCredentials } from "../lib/xai-http";
 import imageGenDescription from "../prompts/tools/image-gen.md" with { type: "text" };
 import { resolveReadPath } from "./path-utils";
 
@@ -46,7 +47,7 @@ export type ImageProviderPreference = Exclude<ImageProvider, "openai-codex"> | "
 
 interface ImageApiKey {
 	provider: ImageProvider;
-	apiKey: string;
+	apiKey: ApiKey;
 	projectId?: string;
 	model?: Model;
 }
@@ -366,7 +367,11 @@ function toDataUrl(image: InlineImageData): string {
 	return `data:${image.mimeType};base64,${image.data}`;
 }
 
-async function loadImageFromUrl(imageUrl: string, signal?: AbortSignal): Promise<InlineImageData> {
+async function loadImageFromUrl(
+	imageUrl: string,
+	fetchImpl: FetchImpl,
+	signal?: AbortSignal,
+): Promise<InlineImageData> {
 	if (imageUrl.startsWith("data:")) {
 		const normalized = normalizeDataUrl(imageUrl.trim());
 		if (!normalized.mimeType) {
@@ -378,7 +383,7 @@ async function loadImageFromUrl(imageUrl: string, signal?: AbortSignal): Promise
 		return { data: normalized.data, mimeType: normalized.mimeType };
 	}
 
-	const response = await fetch(imageUrl, { signal });
+	const response = await fetchImpl(imageUrl, { signal });
 	if (!response.ok) {
 		const rawText = await response.text();
 		throw new Error(`Image download failed (${response.status}): ${rawText}`);
@@ -467,8 +472,13 @@ function parseAntigravityCredentials(raw: string): ParsedAntigravityCredentials
 	return null;
 }
 
-async function findAntigravityCredentials(modelRegistry: ModelRegistry): Promise<ImageApiKey | null> {
-	const apiKey = await modelRegistry.getApiKeyForProvider("google-antigravity");
+async function findAntigravityCredentials(
+	modelRegistry: ModelRegistry,
+	sessionId?: string,
+): Promise<ImageApiKey | null> {
+	const apiKey = await modelRegistry.getApiKeyForProvider("google-antigravity", sessionId, {
+		modelId: DEFAULT_ANTIGRAVITY_MODEL,
+	});
 	if (!apiKey) return null;
 
 	const parsed = parseAntigravityCredentials(apiKey);
@@ -492,6 +502,39 @@ async function findXAIImageCredentials(modelRegistry?: ModelRegistry): Promise<I
 	return null;
 }
 
+async function findOpenRouterImageCredentials(
+	modelRegistry?: ModelRegistry,
+	sessionId?: string,
+): Promise<ImageApiKey | null> {
+	if (modelRegistry) {
+		// AuthStorage.getApiKey already falls back to env keys, so this covers OPENROUTER_API_KEY too.
+		const apiKey = await modelRegistry.getApiKeyForProvider("openrouter", sessionId);
+		if (apiKey) return { provider: "openrouter", apiKey: modelRegistry.resolver("openrouter", { sessionId }) };
+		return null;
+	}
+	const apiKey = getEnvApiKey("openrouter");
+	if (apiKey) return { provider: "openrouter", apiKey };
+	return null;
+}
+
+async function findGeminiImageCredentials(
+	modelRegistry?: ModelRegistry,
+	sessionId?: string,
+): Promise<ImageApiKey | null> {
+	if (modelRegistry) {
+		// AuthStorage.getApiKey already falls back to env keys (GEMINI_API_KEY), so only
+		// GOOGLE_API_KEY needs the explicit check below.
+		const apiKey = await modelRegistry.getApiKeyForProvider("google", sessionId);
+		if (apiKey) return { provider: "gemini", apiKey: modelRegistry.resolver("google", { sessionId }) };
+	} else {
+		const envKey = getEnvApiKey("google");
+		if (envKey) return { provider: "gemini", apiKey: envKey };
+	}
+	const googleKey = $env.GOOGLE_API_KEY;
+	if (googleKey) return { provider: "gemini", apiKey: googleKey };
+	return null;
+}
+
 async function findOpenAIHostedImageCredentials(
 	modelRegistry: ModelRegistry | undefined,
 	activeModel: Model | undefined,
@@ -518,18 +561,16 @@ async function findImageApiKey(
 		if (openAI) return openAI;
 		// Fall through to auto-detect if preferred provider key not found.
 	} else if (preferredImageProvider === "antigravity" && modelRegistry) {
-		const antigravity = await findAntigravityCredentials(modelRegistry);
+		const antigravity = await findAntigravityCredentials(modelRegistry, sessionId);
 		if (antigravity) return antigravity;
 		// Fall through to auto-detect if preferred provider key not found.
 	} else if (preferredImageProvider === "gemini") {
-		const geminiKey = getEnvApiKey("google");
-		if (geminiKey) return { provider: "gemini", apiKey: geminiKey };
-		const googleKey = $env.GOOGLE_API_KEY;
-		if (googleKey) return { provider: "gemini", apiKey: googleKey };
+		const gemini = await findGeminiImageCredentials(modelRegistry, sessionId);
+		if (gemini) return gemini;
 		// Fall through to auto-detect if preferred provider key not found.
 	} else if (preferredImageProvider === "openrouter") {
-		const openRouterKey = getEnvApiKey("openrouter");
-		if (openRouterKey) return { provider: "openrouter", apiKey: openRouterKey };
+		const openRouter = await findOpenRouterImageCredentials(modelRegistry, sessionId);
+		if (openRouter) return openRouter;
 		// Fall through to auto-detect if preferred provider key not found.
 	} else if (preferredImageProvider === "xai") {
 		const xai = await findXAIImageCredentials(modelRegistry);
@@ -542,21 +583,18 @@ async function findImageApiKey(
 	if (openAI) return openAI;
 
 	if (modelRegistry) {
-		const antigravity = await findAntigravityCredentials(modelRegistry);
+		const antigravity = await findAntigravityCredentials(modelRegistry, sessionId);
 		if (antigravity) return antigravity;
 	}
 
 	const xai = await findXAIImageCredentials(modelRegistry);
 	if (xai) return xai;
 
-	const openRouterKey = getEnvApiKey("openrouter");
-	if (openRouterKey) return { provider: "openrouter", apiKey: openRouterKey };
-
-	const geminiKey = getEnvApiKey("google");
-	if (geminiKey) return { provider: "gemini", apiKey: geminiKey };
+	const openRouter = await findOpenRouterImageCredentials(modelRegistry, sessionId);
+	if (openRouter) return openRouter;
 
-	const googleKey = $env.GOOGLE_API_KEY;
-	if (googleKey) return { provider: "gemini", apiKey: googleKey };
+	const gemini = await findGeminiImageCredentials(modelRegistry, sessionId);
+	if (gemini) return gemini;
 
 	return null;
 }
@@ -803,7 +841,7 @@ function buildOpenAIImageHeaders(model: Model, apiKey: string, sessionId: string
 		headers.set(OPENAI_HEADERS.ACCOUNT_ID, accountId);
 		headers.set(OPENAI_HEADERS.BETA, OPENAI_HEADER_VALUES.BETA_RESPONSES);
 		headers.set(OPENAI_HEADERS.ORIGINATOR, OPENAI_HEADER_VALUES.ORIGINATOR_CODEX);
-		headers.set("User-Agent", `prometheus/${packageJson.version} (${os.platform()} ${os.release()}; ${os.arch()})`);
+		headers.set("User-Agent", `pi/${packageJson.version} (${os.platform()} ${os.release()}; ${os.arch()})`);
 		if (sessionId) {
 			headers.set(OPENAI_HEADERS.CONVERSATION_ID, sessionId);
 			headers.set(OPENAI_HEADERS.SESSION_ID, sessionId);
@@ -850,13 +888,14 @@ async function generateOpenAIHostedImage(
 	model: Model,
 	params: ImageGenParams,
 	inputImages: InlineImageData[],
+	fetchImpl: FetchImpl,
 	signal: AbortSignal | undefined,
 	sessionId: string | undefined,
 ): Promise<OpenAIHostedImageResult> {
 	const promptText = assemblePrompt(params);
 	const stream = model.api === "openai-codex-responses" || model.provider === "openai-codex";
 	const requestBody = buildOpenAIHostedImageRequest(model, promptText, params, inputImages, stream);
-	const response = await fetch(getOpenAIResponsesUrl(model), {
+	const response = await fetchImpl(getOpenAIResponsesUrl(model), {
 		method: "POST",
 		headers: buildOpenAIImageHeaders(model, apiKey, sessionId),
 		body: JSON.stringify(requestBody),
@@ -865,7 +904,10 @@ async function generateOpenAIHostedImage(
 
 	if (!response.ok) {
 		const errorText = await response.text();
-		throw new Error(`OpenAI image request failed (${response.status}): ${getOpenAIResponseErrorMessage(errorText)}`);
+		throw Object.assign(
+			new Error(`OpenAI image request failed (${response.status}): ${getOpenAIResponseErrorMessage(errorText)}`),
+			{ status: response.status },
+		);
 	}
 
 	const contentType = response.headers.get("content-type") ?? "";
@@ -1032,19 +1074,29 @@ export const imageGenTool: CustomTool<typeof imageGenSchema, ImageGenToolDetails
 			}
 
 			const requestSignal = ptree.combineSignals(signal, IMAGE_TIMEOUT);
+			const fetchImpl = ctx.fetch ?? fetch;
 
 			if (provider === "openai" || provider === "openai-codex") {
 				if (!apiKey.model) {
 					throw new Error("Missing active GPT model for OpenAI image generation");
 				}
 
-				const parsed = await generateOpenAIHostedImage(
-					apiKey.apiKey,
-					apiKey.model,
-					params,
-					resolvedImages,
-					requestSignal,
-					sessionId,
+				const hostedModel = apiKey.model;
+				const hostedKey: ApiKey = ctx.modelRegistry.resolver(hostedModel, sessionId);
+
+				const parsed = await withAuth(
+					hostedKey,
+					key =>
+						generateOpenAIHostedImage(
+							key,
+							hostedModel,
+							params,
+							resolvedImages,
+							fetchImpl,
+							requestSignal,
+							sessionId,
+						),
+					{ signal: requestSignal },
 				);
 
 				if (parsed.images.length === 0) {
@@ -1089,38 +1141,60 @@ export const imageGenTool: CustomTool<typeof imageGenSchema, ImageGenToolDetails
 				}
 
 				const prompt = assemblePrompt(params);
-				const requestBody = buildAntigravityRequest(
-					prompt,
-					model,
-					apiKey.projectId,
-					params.aspect_ratio,
-					params.image_size,
-					resolvedImages,
-				);
-
-				const response = await fetch(`${ANTIGRAVITY_ENDPOINT}/v1internal:streamGenerateContent?alt=sse`, {
-					method: "POST",
-					headers: {
-						Authorization: `Bearer ${apiKey.apiKey}`,
-						"Content-Type": "application/json",
-						Accept: "text/event-stream",
-						"User-Agent": getAntigravityUserAgent(),
-					},
-					body: JSON.stringify(requestBody),
-					signal: requestSignal,
+				const antigravityKey: ApiKey = ctx.modelRegistry.resolver("google-antigravity", {
+					sessionId,
+					modelId: DEFAULT_ANTIGRAVITY_MODEL,
 				});
 
-				if (!response.ok) {
-					const errorText = await response.text();
-					let message = errorText;
-					try {
-						const parsed = JSON.parse(errorText) as { error?: { message?: string } };
-						message = parsed.error?.message ?? message;
-					} catch {
-						// Keep raw text.
-					}
-					throw new Error(`Antigravity image request failed (${response.status}): ${message}`);
-				}
+				const response = await withAuth(
+					antigravityKey,
+					async key => {
+						// On a retry the resolver yields the raw stored credential JSON
+						// ({ token, projectId }); the initial seed is the already-parsed
+						// access token. Tolerate both, falling back to the seed projectId.
+						const rotated = parseAntigravityCredentials(key);
+						const bearer = rotated?.accessToken ?? key;
+						const projectId = rotated?.projectId ?? apiKey.projectId!;
+						const requestBody = buildAntigravityRequest(
+							prompt,
+							model,
+							projectId,
+							params.aspect_ratio,
+							params.image_size,
+							resolvedImages,
+						);
+
+						const resp = await fetchImpl(`${ANTIGRAVITY_ENDPOINT}/v1internal:streamGenerateContent?alt=sse`, {
+							method: "POST",
+							headers: {
+								Authorization: `Bearer ${bearer}`,
+								"Content-Type": "application/json",
+								Accept: "text/event-stream",
+								"User-Agent": getAntigravityUserAgent(),
+							},
+							body: JSON.stringify(requestBody),
+							signal: requestSignal,
+						});
+
+						if (!resp.ok) {
+							const errorText = await resp.text();
+							let message = errorText;
+							try {
+								const parsedErr = JSON.parse(errorText) as { error?: { message?: string } };
+								message = parsedErr.error?.message ?? message;
+							} catch {
+								// Keep raw text.
+							}
+							throw new ProviderHttpError(
+								`Antigravity image request failed (${resp.status}): ${message}`,
+								resp.status,
+								{ headers: resp.headers },
+							);
+						}
+						return resp;
+					},
+					{ signal: requestSignal },
+				);
 
 				const parsed = await parseAntigravitySseForImage(response, requestSignal);
 				const responseText = parsed.text.length > 0 ? parsed.text.join(" ") : undefined;
@@ -1192,28 +1266,41 @@ export const imageGenTool: CustomTool<typeof imageGenSchema, ImageGenToolDetails
 					: xaiBaseBody;
 				const xaiEndpoint = isEdit ? "/images/edits" : "/images/generations";
 
-				const xaiResponse = await fetch(`${xaiCreds.baseURL}${xaiEndpoint}`, {
-					method: "POST",
-					headers: {
-						Authorization: `Bearer ${xaiCreds.apiKey}`,
-						"Content-Type": "application/json",
-						"User-Agent": prometheusXAIUserAgent(),
-					},
-					body: JSON.stringify(xaiBody),
-					signal: requestSignal,
+				const xaiKey: ApiKey = ctx.modelRegistry.resolver(xaiCreds.provider, {
+					sessionId,
+					baseUrl: xaiCreds.baseURL,
 				});
 
-				const xaiRawText = await xaiResponse.text();
-				if (!xaiResponse.ok) {
-					let message = xaiRawText;
-					try {
-						const parsedErr = JSON.parse(xaiRawText) as { error?: { message?: string } };
-						message = parsedErr.error?.message ?? message;
-					} catch {
-						// Keep raw text.
-					}
-					throw new Error(`xAI image request failed (${xaiResponse.status}): ${message}`);
-				}
+				const xaiRawText = await withAuth(
+					xaiKey,
+					async key => {
+						const resp = await fetchImpl(`${xaiCreds.baseURL}${xaiEndpoint}`, {
+							method: "POST",
+							headers: {
+								Authorization: `Bearer ${key}`,
+								"Content-Type": "application/json",
+								"User-Agent": getPrometheusXAIUserAgent(),
+							},
+							body: JSON.stringify(xaiBody),
+							signal: requestSignal,
+						});
+						const rawText = await resp.text();
+						if (!resp.ok) {
+							let message = rawText;
+							try {
+								const parsedErr = JSON.parse(rawText) as { error?: { message?: string } };
+								message = parsedErr.error?.message ?? message;
+							} catch {
+								// Keep raw text.
+							}
+							throw new ProviderHttpError(`xAI image request failed (${resp.status}): ${message}`, resp.status, {
+								headers: resp.headers,
+							});
+						}
+						return rawText;
+					},
+					{ signal: requestSignal },
+				);
 
 				const xaiData = JSON.parse(xaiRawText) as {
 					data?: Array<{ b64_json?: string; url?: string }>;
@@ -1225,7 +1312,7 @@ export const imageGenTool: CustomTool<typeof imageGenSchema, ImageGenToolDetails
 						const mimeType = parseImageMetadata(bytes)?.mimeType ?? "image/png";
 						xaiInlineImages.push({ data: entry.b64_json, mimeType });
 					} else if (entry.url) {
-						xaiInlineImages.push(await loadImageFromUrl(entry.url, requestSignal));
+						xaiInlineImages.push(await loadImageFromUrl(entry.url, fetchImpl, requestSignal));
 					}
 				}
 
@@ -1270,30 +1357,40 @@ export const imageGenTool: CustomTool<typeof imageGenSchema, ImageGenToolDetails
 					messages: [{ role: "user" as const, content: contentParts }],
 				};
 
-				const response = await fetch("https://openrouter.ai/api/v1/chat/completions", {
-					method: "POST",
-					headers: {
-						"Content-Type": "application/json",
-						Authorization: `Bearer ${apiKey.apiKey}`,
-						"HTTP-Referer": "https://prometheus.trivlab.com/",
-						"X-OpenRouter-Title": APP_DISPLAY_NAME,
-						"X-OpenRouter-Categories": "cli-agent",
+				const rawText = await withAuth(
+					apiKey.apiKey,
+					async key => {
+						const resp = await fetchImpl("https://openrouter.ai/api/v1/chat/completions", {
+							method: "POST",
+							headers: {
+								"Content-Type": "application/json",
+								Authorization: `Bearer ${key}`,
+								"HTTP-Referer": "https://prometheus.sh/",
+								"X-OpenRouter-Title": "Oh-My-Pi",
+								"X-OpenRouter-Categories": "cli-agent",
+							},
+							body: JSON.stringify(requestBody),
+							signal: requestSignal,
+						});
+						const text = await resp.text();
+						if (!resp.ok) {
+							let message = text;
+							try {
+								const parsed = JSON.parse(text) as { error?: { message?: string } };
+								message = parsed.error?.message ?? message;
+							} catch {
+								// Keep raw text.
+							}
+							throw new ProviderHttpError(
+								`OpenRouter image request failed (${resp.status}): ${message}`,
+								resp.status,
+								{ headers: resp.headers },
+							);
+						}
+						return text;
 					},
-					body: JSON.stringify(requestBody),
-					signal: requestSignal,
-				});
-
-				const rawText = await response.text();
-				if (!response.ok) {
-					let message = rawText;
-					try {
-						const parsed = JSON.parse(rawText) as { error?: { message?: string } };
-						message = parsed.error?.message ?? message;
-					} catch {
-						// Keep raw text.
-					}
-					throw new Error(`OpenRouter image request failed (${response.status}): ${message}`);
-				}
+					{ signal: requestSignal },
+				);
 
 				const data = JSON.parse(rawText) as OpenRouterResponse;
 				const message = data.choices?.[0]?.message;
@@ -1301,7 +1398,7 @@ export const imageGenTool: CustomTool<typeof imageGenSchema, ImageGenToolDetails
 				const imageUrls = extractOpenRouterImageUrls(message);
 				const inlineImages: InlineImageData[] = [];
 				for (const imageUrl of imageUrls) {
-					inlineImages.push(await loadImageFromUrl(imageUrl, requestSignal));
+					inlineImages.push(await loadImageFromUrl(imageUrl, fetchImpl, requestSignal));
 				}
 
 				if (inlineImages.length === 0) {
@@ -1361,31 +1458,39 @@ export const imageGenTool: CustomTool<typeof imageGenSchema, ImageGenToolDetails
 				generationConfig,
 			};
 
-			const response = await fetch(
-				`https://generativelanguage.googleapis.com/v1beta/models/${encodeURIComponent(model)}:generateContent`,
-				{
-					method: "POST",
-					headers: {
-						"Content-Type": "application/json",
-						"x-goog-api-key": apiKey.apiKey,
-					},
-					body: JSON.stringify(requestBody),
-					signal: requestSignal,
+			const rawText = await withAuth(
+				apiKey.apiKey,
+				async key => {
+					const resp = await fetchImpl(
+						`https://generativelanguage.googleapis.com/v1beta/models/${encodeURIComponent(model)}:generateContent`,
+						{
+							method: "POST",
+							headers: {
+								"Content-Type": "application/json",
+								"x-goog-api-key": key,
+							},
+							body: JSON.stringify(requestBody),
+							signal: requestSignal,
+						},
+					);
+					const text = await resp.text();
+					if (!resp.ok) {
+						let message = text;
+						try {
+							const parsed = JSON.parse(text) as { error?: { message?: string } };
+							message = parsed.error?.message ?? message;
+						} catch {
+							// Keep raw text.
+						}
+						throw new ProviderHttpError(`Gemini image request failed (${resp.status}): ${message}`, resp.status, {
+							headers: resp.headers,
+						});
+					}
+					return text;
 				},
+				{ signal: requestSignal },
 			);
 
-			const rawText = await response.text();
-			if (!response.ok) {
-				let message = rawText;
-				try {
-					const parsed = JSON.parse(rawText) as { error?: { message?: string } };
-					message = parsed.error?.message ?? message;
-				} catch {
-					// Keep raw text.
-				}
-				throw new Error(`Gemini image request failed (${response.status}): ${message}`);
-			}
-
 			const data = JSON.parse(rawText) as GeminiGenerateContentResponse;
 			const responseParts = combineParts(data);
 			const responseText = collectResponseText(responseParts);