pi-openai-toolkit 0.7.2 → 0.9.0

package/src/image-generation/render.ts

@@ -0,0 +1,252 @@
+import { readFileSync, realpathSync, statSync } from "node:fs";
+import { isAbsolute, relative, resolve, sep } from "node:path";
+import type { AgentToolResult } from "@earendil-works/pi-agent-core";
+import {
+	getAgentDir,
+	type Theme,
+	type ToolRenderResultOptions,
+} from "@earendil-works/pi-coding-agent";
+import {
+	Container,
+	getCapabilities,
+	Image,
+	Spacer,
+	Text,
+	type Component,
+} from "@earendil-works/pi-tui";
+import { getGeneratedImagesRoot } from "./artifacts";
+import { isValidPng, readPngDimensions } from "./protocol";
+import {
+	IMAGE_GENERATION_MIME_TYPE,
+	MAX_GENERATED_IMAGE_BYTES,
+	isImageGenerationDetails,
+	type GenerateImageParams,
+	type ImageGenerationDetails,
+} from "./types";
+
+export type ImageGenerationRenderState = Record<string, never>;
+
+export type ImageGenerationRenderContext = {
+	lastComponent: Component | undefined;
+	showImages: boolean;
+	isError: boolean;
+	isPartial: boolean;
+};
+
+export type ImageRendererDependencies = {
+	getAgentDir: () => string;
+	getCapabilities: typeof getCapabilities;
+	readFileSync: typeof readFileSync;
+	realpathSync: typeof realpathSync;
+	statSync: typeof statSync;
+};
+
+const DEFAULT_RENDER_DEPS: ImageRendererDependencies = {
+	getAgentDir,
+	getCapabilities,
+	readFileSync,
+	realpathSync,
+	statSync,
+};
+
+function isPathInsideOrEqual(root: string, candidate: string): boolean {
+	const rel = relative(root, candidate);
+	return rel === "" || (!rel.startsWith(`..${sep}`) && rel !== ".." && !isAbsolute(rel));
+}
+
+function resultText(result: AgentToolResult<ImageGenerationDetails>): string {
+	return result.content
+		.filter((block): block is { type: "text"; text: string } => block.type === "text")
+		.map((block) => block.text)
+		.join("\n")
+		.trim();
+}
+
+class DiskBackedImageResult extends Container {
+	readonly artifactPath?: string;
+	readonly fileSize?: number;
+	readonly mtimeMs?: number;
+	readonly showImages: boolean;
+	readonly textValue: string;
+
+	constructor(args: {
+		text: string;
+		artifactPath?: string;
+		fileSize?: number;
+		mtimeMs?: number;
+		showImages: boolean;
+		base64?: string;
+		theme: Theme;
+		previewWarning?: string;
+	}) {
+		super();
+		this.artifactPath = args.artifactPath;
+		this.fileSize = args.fileSize;
+		this.mtimeMs = args.mtimeMs;
+		this.showImages = args.showImages;
+		this.textValue = `${args.text}\n${args.previewWarning ?? ""}`;
+		this.addChild(new Text(args.theme.fg("toolOutput", args.text), 0, 0));
+		if (args.previewWarning) {
+			this.addChild(new Text(args.theme.fg("muted", args.previewWarning), 0, 0));
+		}
+		if (args.base64) {
+			this.addChild(new Spacer(1));
+			this.addChild(
+				new Image(
+					args.base64,
+					IMAGE_GENERATION_MIME_TYPE,
+					{ fallbackColor: (text) => args.theme.fg("toolOutput", text) },
+					{ maxWidthCells: 60, filename: args.artifactPath },
+				),
+			);
+		}
+	}
+
+	matches(args: {
+		text: string;
+		artifactPath?: string;
+		fileSize?: number;
+		mtimeMs?: number;
+		showImages: boolean;
+		previewWarning?: string;
+	}): boolean {
+		return (
+			this.textValue === `${args.text}\n${args.previewWarning ?? ""}` &&
+			this.artifactPath === args.artifactPath &&
+			this.fileSize === args.fileSize &&
+			this.mtimeMs === args.mtimeMs &&
+			this.showImages === args.showImages
+		);
+	}
+}
+
+function textOnly(args: {
+	text: string;
+	theme: Theme;
+	context: ImageGenerationRenderContext;
+	previewWarning?: string;
+}): Component {
+	const previous = args.context.lastComponent;
+	if (
+		previous instanceof DiskBackedImageResult &&
+		previous.matches({
+			text: args.text,
+			showImages: false,
+			previewWarning: args.previewWarning,
+		})
+	) {
+		return previous;
+	}
+	return new DiskBackedImageResult({
+		text: args.text,
+		showImages: false,
+		theme: args.theme,
+		previewWarning: args.previewWarning,
+	});
+}
+
+export function createImageGenerationResultRenderer(
+	deps: ImageRendererDependencies = DEFAULT_RENDER_DEPS,
+): (
+	result: AgentToolResult<ImageGenerationDetails>,
+	options: ToolRenderResultOptions,
+	theme: Theme,
+	context: ImageGenerationRenderContext,
+) => Component {
+	return (result, _options, theme, context) => {
+		const text = resultText(result) || "Image generation finished.";
+		if (
+			context.isError ||
+			context.isPartial ||
+			!context.showImages ||
+			!isImageGenerationDetails(result.details) ||
+			!deps.getCapabilities().images
+		) {
+			return textOnly({ text, theme, context });
+		}
+
+		const artifactPath = resolve(result.details.artifactPath);
+		let root: string;
+		let realArtifact: string;
+		let info: ReturnType<typeof statSync>;
+		try {
+			root = deps.realpathSync(getGeneratedImagesRoot(deps.getAgentDir()));
+			realArtifact = deps.realpathSync(artifactPath);
+			info = deps.statSync(realArtifact);
+		} catch {
+			return textOnly({
+				text,
+				theme,
+				context,
+				previewWarning: "Image preview unavailable: the saved artifact is missing.",
+			});
+		}
+
+		if (
+			!isPathInsideOrEqual(root, realArtifact) ||
+			!info.isFile() ||
+			info.size <= 0 ||
+			info.size > MAX_GENERATED_IMAGE_BYTES
+		) {
+			return textOnly({
+				text,
+				theme,
+				context,
+				previewWarning: "Image preview unavailable: the artifact path or file is not safe.",
+			});
+		}
+
+		const previous = context.lastComponent;
+		if (
+			previous instanceof DiskBackedImageResult &&
+			previous.matches({
+				text,
+				artifactPath: realArtifact,
+				fileSize: info.size,
+				mtimeMs: info.mtimeMs,
+				showImages: true,
+			})
+		) {
+			return previous;
+		}
+
+		let bytes: Buffer;
+		try {
+			bytes = deps.readFileSync(realArtifact);
+		} catch {
+			return textOnly({
+				text,
+				theme,
+				context,
+				previewWarning: "Image preview unavailable: the artifact could not be read.",
+			});
+		}
+		if (bytes.length !== info.size || !isValidPng(bytes) || !readPngDimensions(bytes)) {
+			bytes.fill(0);
+			return textOnly({
+				text,
+				theme,
+				context,
+				previewWarning: "Image preview unavailable: the artifact is not a valid PNG.",
+			});
+		}
+		const base64 = bytes.toString("base64");
+		bytes.fill(0);
+		return new DiskBackedImageResult({
+			text,
+			artifactPath: realArtifact,
+			fileSize: info.size,
+			mtimeMs: info.mtimeMs,
+			showImages: true,
+			base64,
+			theme,
+		});
+	};
+}
+
+export const renderImageGenerationResult = createImageGenerationResultRenderer();
+
+export const _renderTest = {
+	DiskBackedImageResult,
+	isPathInsideOrEqual,
+};

package/src/image-generation/service.ts

@@ -0,0 +1,243 @@
+import { getAgentDir, type ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { loadToolkitConfig } from "../config";
+import { resolveResponsesEnvironment } from "../runtime";
+import {
+	copyImageToExplicitPath,
+	prepareExplicitOutputPath,
+	saveCanonicalImage,
+} from "./artifacts";
+import { requestGeneratedImage } from "./client";
+import { isImageGenerationEnabledForModel } from "./eligibility";
+import {
+	buildImageGenerationRequest,
+	normalizeGenerateImageParams,
+} from "./protocol";
+import {
+	clearPreparedReferences,
+	prepareReferenceImages,
+} from "./references";
+import {
+	IMAGE_GENERATION_CAPABLE_APIS,
+	IMAGE_GENERATION_MIME_TYPE,
+	IMAGE_GENERATION_MODEL,
+	ImageGenerationError,
+	sanitizeImageDiagnostic,
+	type GenerateImageParams,
+	type ImageGenerationExecutionResult,
+	type ImageGenerationFailureCode,
+	type ParsedGeneratedImage,
+} from "./types";
+
+export type ImageGenerationServiceDependencies = {
+	loadConfig: typeof loadToolkitConfig;
+	resolveRuntime: typeof resolveResponsesEnvironment;
+	getAgentDir: typeof getAgentDir;
+	prepareOutput: typeof prepareExplicitOutputPath;
+	prepareReferences: typeof prepareReferenceImages;
+	requestImage: typeof requestGeneratedImage;
+	saveCanonical: typeof saveCanonicalImage;
+	copyExplicit: typeof copyImageToExplicitPath;
+	clearReferences: typeof clearPreparedReferences;
+};
+
+const DEFAULT_SERVICE_DEPS: ImageGenerationServiceDependencies = {
+	loadConfig: loadToolkitConfig,
+	resolveRuntime: resolveResponsesEnvironment,
+	getAgentDir,
+	prepareOutput: prepareExplicitOutputPath,
+	prepareReferences: prepareReferenceImages,
+	requestImage: requestGeneratedImage,
+	saveCanonical: saveCanonicalImage,
+	copyExplicit: copyImageToExplicitPath,
+	clearReferences: clearPreparedReferences,
+};
+
+function runtimeFailure(
+	reason: string,
+	errorMessage?: string,
+): ImageGenerationError {
+	switch (reason) {
+		case "disabled":
+			return new ImageGenerationError("disabled", "Image generation is disabled.");
+		case "unsupported-api":
+		case "missing-model":
+			return new ImageGenerationError(
+				"unsupported-model",
+				"Image generation is unavailable for the current model.",
+			);
+		case "missing-api-key":
+		case "auth-resolution-failed":
+			return new ImageGenerationError(
+				"authentication",
+				sanitizeImageDiagnostic(
+					errorMessage,
+					"Unable to resolve authentication for image generation with the current model.",
+				),
+			);
+		default:
+			return new ImageGenerationError(
+				"missing-runtime",
+				"Unable to resolve the current model's Responses endpoint for image generation.",
+			);
+	}
+}
+
+function throwIfAborted(signal: AbortSignal | undefined): void {
+	if (signal?.aborted) {
+		throw new ImageGenerationError("aborted", "Image generation was cancelled.");
+	}
+}
+
+function clientFailureCode(reason: string): ImageGenerationFailureCode {
+	switch (reason) {
+		case "aborted":
+		case "timeout":
+		case "authentication":
+		case "rate-limit":
+		case "request-rejected":
+		case "backend-unavailable":
+		case "network":
+		case "oversized-response":
+		case "malformed-response":
+		case "no-image":
+			return reason;
+		default:
+			return "network";
+	}
+}
+
+function formatResultText(args: {
+	artifactPath: string;
+	outputPath?: string;
+	warning?: string;
+	edited: boolean;
+	width: number;
+	height: number;
+}): string {
+	const lines = [
+		`${args.edited ? "Edited" : "Generated"} PNG image (${args.width}x${args.height}).`,
+		`Artifact: ${args.artifactPath}`,
+	];
+	if (args.outputPath) lines.push(`Copied to: ${args.outputPath}`);
+	if (args.warning) lines.push(`Warning: ${args.warning}`);
+	return lines.join("\n");
+}
+
+export function createImageGenerationExecutor(
+	deps: ImageGenerationServiceDependencies = DEFAULT_SERVICE_DEPS,
+): (args: {
+	params: GenerateImageParams;
+	toolCallId: string;
+	signal?: AbortSignal;
+	ctx: ExtensionContext;
+}) => Promise<ImageGenerationExecutionResult> {
+	return async (args) => {
+		const params = normalizeGenerateImageParams(args.params);
+		throwIfAborted(args.signal);
+		const { config } = deps.loadConfig();
+		if (!isImageGenerationEnabledForModel(args.ctx.model, config.imageGeneration)) {
+			throw new ImageGenerationError(
+				"unsupported-model",
+				"Image generation is not enabled for the current provider/model-id.",
+			);
+		}
+
+		const runtimeResolution = await deps.resolveRuntime(args.ctx, {
+			enabled: config.imageGeneration.enabled,
+			responsesApis: IMAGE_GENERATION_CAPABLE_APIS,
+		});
+		if (!runtimeResolution.ok) {
+			throw runtimeFailure(runtimeResolution.reason, runtimeResolution.errorMessage);
+		}
+
+		const agentDir = deps.getAgentDir();
+		const explicitOutput = await deps.prepareOutput({
+			rawPath: params.outputPath,
+			agentDir,
+			ctx: args.ctx,
+			signal: args.signal,
+		});
+
+		const references = await deps.prepareReferences({
+			paths: params.referenceImagePaths,
+			ctx: args.ctx,
+			signal: args.signal,
+		});
+		let generated: ParsedGeneratedImage | undefined;
+		try {
+			throwIfAborted(args.signal);
+			const body = buildImageGenerationRequest({
+				routingModel: runtimeResolution.runtime.model,
+				params,
+				references,
+			});
+			const response = await deps.requestImage({
+				runtime: runtimeResolution.runtime,
+				body,
+				signal: args.signal,
+			});
+			if (!response.ok) {
+				throw new ImageGenerationError(
+					clientFailureCode(response.reason),
+					response.errorMessage,
+				);
+			}
+			generated = response.image;
+
+			const artifactPath = await deps.saveCanonical({
+				bytes: generated.bytes,
+				agentDir,
+				sessionId: args.ctx.sessionManager.getSessionId(),
+				imageCallId: generated.imageCallId || args.toolCallId,
+			});
+
+			let outputPath: string | undefined;
+			let warning: string | undefined;
+			if (explicitOutput) {
+				try {
+					await deps.copyExplicit({ bytes: generated.bytes, plan: explicitOutput });
+					outputPath = explicitOutput.path;
+				} catch (error) {
+					warning = sanitizeImageDiagnostic(
+						error instanceof Error ? error.message : error,
+						`The generated image was saved to the canonical artifact, but copying to ${explicitOutput.path} failed.`,
+					);
+				}
+			}
+
+			const edited = references.length > 0;
+			const details = {
+				artifactPath,
+				...(outputPath ? { outputPath } : {}),
+				routingModel: `${runtimeResolution.runtime.provider}/${runtimeResolution.runtime.model}`,
+				imageModel: IMAGE_GENERATION_MODEL,
+				imageCallId: generated.imageCallId,
+				...(generated.responseId ? { responseId: generated.responseId } : {}),
+				mimeType: IMAGE_GENERATION_MIME_TYPE,
+				byteCount: generated.bytes.length,
+				width: generated.width,
+				height: generated.height,
+				edited,
+				referenceCount: references.length,
+				...(generated.revisedPrompt ? { revisedPrompt: generated.revisedPrompt } : {}),
+				...(warning ? { warning } : {}),
+			};
+			return {
+				details,
+				text: formatResultText({
+					artifactPath,
+					outputPath,
+					warning,
+					edited,
+					width: generated.width,
+					height: generated.height,
+				}),
+			};
+		} finally {
+			deps.clearReferences(references);
+			generated?.bytes.fill(0);
+		}
+	};
+}
+
+export const executeImageGeneration = createImageGenerationExecutor();

package/src/image-generation/types.ts

@@ -0,0 +1,184 @@
+export const IMAGE_GENERATION_TOOL_NAME = "openai_generate_image";
+export const IMAGE_GENERATION_MODEL = "gpt-image-2" as const;
+export const IMAGE_GENERATION_MIME_TYPE = "image/png" as const;
+export const IMAGE_GENERATION_CAPABLE_APIS = [
+	"openai-responses",
+	"openai-codex-responses",
+] as const;
+export const IMAGE_GENERATION_SIZES = [
+	"auto",
+	"1024x1024",
+	"1536x1024",
+	"1024x1536",
+] as const;
+export const IMAGE_GENERATION_QUALITIES = ["auto", "low", "medium", "high"] as const;
+
+export const MAX_IMAGE_PROMPT_CHARS = 20_000;
+export const MAX_REFERENCE_IMAGE_COUNT = 5;
+export const MAX_REFERENCE_IMAGE_BYTES = 20 * 1024 * 1024;
+export const MAX_TOTAL_REFERENCE_IMAGE_BYTES = 50 * 1024 * 1024;
+export const MAX_GENERATED_IMAGE_BYTES = 32 * 1024 * 1024;
+export const MAX_IMAGE_RESPONSE_BYTES = 48 * 1024 * 1024;
+export const MAX_IMAGE_ERROR_BYTES = 64 * 1024;
+export const IMAGE_REQUEST_TIMEOUT_MS = 5 * 60 * 1000;
+export const MAX_IMAGE_PATH_CHARS = 4096;
+export const MAX_IMAGE_DIAGNOSTIC_CHARS = 4096;
+export const MAX_IMAGE_IDENTIFIER_CHARS = 256;
+export const MAX_IMAGE_DIMENSION = 100_000;
+
+export type ImageGenerationCapableApi = (typeof IMAGE_GENERATION_CAPABLE_APIS)[number];
+export type ImageGenerationSize = (typeof IMAGE_GENERATION_SIZES)[number];
+export type ImageGenerationQuality = (typeof IMAGE_GENERATION_QUALITIES)[number];
+export type ImageGenerationAction = "generate" | "edit";
+export type ReferenceImageMimeType = "image/png" | "image/jpeg" | "image/webp";
+
+export type ImageGenerationModel = {
+	api?: string;
+	provider?: string;
+	id?: string;
+};
+
+export type GenerateImageParams = {
+	prompt: string;
+	referenceImagePaths?: string[] | null;
+	outputPath?: string | null;
+	size?: ImageGenerationSize;
+	quality?: ImageGenerationQuality;
+};
+
+export type NormalizedGenerateImageParams = {
+	prompt: string;
+	referenceImagePaths: string[];
+	outputPath?: string;
+	size: ImageGenerationSize;
+	quality: ImageGenerationQuality;
+	action: ImageGenerationAction;
+};
+
+export type PreparedReferenceImage = {
+	path: string;
+	mimeType: ReferenceImageMimeType;
+	bytes: Buffer;
+};
+
+export type ParsedGeneratedImage = {
+	bytes: Buffer;
+	imageCallId: string;
+	responseId?: string;
+	revisedPrompt?: string;
+	width: number;
+	height: number;
+};
+
+export type ImageGenerationDetails = {
+	artifactPath: string;
+	outputPath?: string;
+	routingModel: string;
+	imageModel: typeof IMAGE_GENERATION_MODEL;
+	imageCallId: string;
+	responseId?: string;
+	mimeType: typeof IMAGE_GENERATION_MIME_TYPE;
+	byteCount: number;
+	width: number;
+	height: number;
+	edited: boolean;
+	referenceCount: number;
+	revisedPrompt?: string;
+	warning?: string;
+};
+
+export type ImageGenerationExecutionResult = {
+	details: ImageGenerationDetails;
+	text: string;
+};
+
+export type ImageGenerationFailureCode =
+	| "aborted"
+	| "timeout"
+	| "disabled"
+	| "unsupported-model"
+	| "missing-runtime"
+	| "invalid-parameters"
+	| "reference-input-invalid"
+	| "reference-upload-approval-required"
+	| "reference-upload-declined"
+	| "output-path-invalid"
+	| "output-path-approval-required"
+	| "output-path-declined"
+	| "authentication"
+	| "rate-limit"
+	| "request-rejected"
+	| "backend-unavailable"
+	| "network"
+	| "oversized-response"
+	| "malformed-response"
+	| "no-image"
+	| "artifact-write-failed";
+
+export class ImageGenerationError extends Error {
+	readonly code: ImageGenerationFailureCode;
+
+	constructor(code: ImageGenerationFailureCode, message: string) {
+		super(message);
+		this.name = "ImageGenerationError";
+		this.code = code;
+	}
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return !!value && typeof value === "object" && !Array.isArray(value);
+}
+
+function isBoundedString(value: unknown, maxChars: number): value is string {
+	return typeof value === "string" && value.length > 0 && value.length <= maxChars;
+}
+
+export function isImageGenerationDetails(value: unknown): value is ImageGenerationDetails {
+	if (!isRecord(value)) return false;
+	return (
+		isBoundedString(value.artifactPath, MAX_IMAGE_PATH_CHARS) &&
+		(value.outputPath === undefined || isBoundedString(value.outputPath, MAX_IMAGE_PATH_CHARS)) &&
+		isBoundedString(value.routingModel, MAX_IMAGE_DIAGNOSTIC_CHARS) &&
+		value.imageModel === IMAGE_GENERATION_MODEL &&
+		isBoundedString(value.imageCallId, MAX_IMAGE_IDENTIFIER_CHARS) &&
+		(value.responseId === undefined || isBoundedString(value.responseId, MAX_IMAGE_IDENTIFIER_CHARS)) &&
+		value.mimeType === IMAGE_GENERATION_MIME_TYPE &&
+		typeof value.byteCount === "number" &&
+		Number.isInteger(value.byteCount) &&
+		value.byteCount > 0 &&
+		value.byteCount <= MAX_GENERATED_IMAGE_BYTES &&
+		typeof value.width === "number" &&
+		Number.isInteger(value.width) &&
+		value.width > 0 &&
+		value.width <= MAX_IMAGE_DIMENSION &&
+		typeof value.height === "number" &&
+		Number.isInteger(value.height) &&
+		value.height > 0 &&
+		value.height <= MAX_IMAGE_DIMENSION &&
+		typeof value.edited === "boolean" &&
+		typeof value.referenceCount === "number" &&
+		Number.isInteger(value.referenceCount) &&
+		value.referenceCount >= 0 &&
+		value.referenceCount <= MAX_REFERENCE_IMAGE_COUNT &&
+		(value.revisedPrompt === undefined || isBoundedString(value.revisedPrompt, MAX_IMAGE_DIAGNOSTIC_CHARS)) &&
+		(value.warning === undefined || isBoundedString(value.warning, MAX_IMAGE_DIAGNOSTIC_CHARS))
+	);
+}
+
+export function sanitizeImageDiagnostic(value: unknown, fallback: string): string {
+	if (typeof value !== "string") return fallback;
+	const normalized = value
+		.replace(/data:image\/[a-z0-9.+-]+;base64,[A-Za-z0-9+/=]+/gi, "[REDACTED_IMAGE_DATA]")
+		.replace(/[A-Za-z0-9+/]{128,}={0,2}/g, "[REDACTED_BASE64]")
+		.replace(/\bBearer\s+[A-Za-z0-9._~+/=-]+/gi, "Bearer [REDACTED]")
+		.replace(
+			/\b(authorization|api[-_ ]?key|access[-_ ]?token|refresh[-_ ]?token|account[-_ ]?id|cookie)\s*[:=]\s*[^\s,;]+/gi,
+			"$1: [REDACTED]",
+		)
+		.replace(/\bsk-[A-Za-z0-9_-]{8,}\b/g, "[REDACTED]")
+		.replace(/[\u0000-\u001f\u007f]+/g, " ")
+		.replace(/\s+/g, " ")
+		.trim();
+	if (!normalized) return fallback;
+	return normalized.slice(0, MAX_IMAGE_DIAGNOSTIC_CHARS);
+}

package/src/model-scope.ts

@@ -0,0 +1,17 @@
+export type ExactModelScopeModel = {
+	provider?: string;
+	id?: string;
+};
+
+export function getExactModelKey(model: ExactModelScopeModel | undefined): string | undefined {
+	if (!model?.provider || !model.id) return undefined;
+	return `${model.provider}/${model.id}`;
+}
+
+export function isExactModelAllowed(
+	model: ExactModelScopeModel | undefined,
+	allowlist: readonly string[],
+): boolean {
+	const modelKey = getExactModelKey(model);
+	return modelKey !== undefined && allowlist.includes(modelKey);
+}