pi-better-openai 0.1.18 → 0.1.21

package/src/image.ts

@@ -1,16 +1,29 @@
 import { randomUUID } from "node:crypto";
-import { mkdir, readFile, writeFile } from "node:fs/promises";
+import { mkdir, readFile, realpath, stat, writeFile } from "node:fs/promises";
 import { homedir } from "node:os";
 import { extname, isAbsolute, join, resolve, sep } from "node:path";
-import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { Box, Container, Image, Text } from "@earendil-works/pi-tui";
+import sharp from "sharp";
 import type { ResolvedConfig } from "./config.ts";
 import { isRecord } from "./config.ts";
-import { readCodexAuth } from "./usage.ts";
+import {
+  extractAccountIdFromJwt,
+  getCodexCredentials,
+  type CodexCredentialsWithSource,
+} from "./codex-auth.ts";
+import { maskIdentifier, sanitizeDiagnosticError } from "./format.ts";
+import { piAgentDir, resolveUserPath } from "./paths.ts";
 
 const OPENAI_IMAGE_TOOL = "openai_image";
 const OPENAI_IMAGE_COMMAND = "openai-image";
 const CODEX_RESPONSES_URL = "https://chatgpt.com/backend-api/codex/responses";
 const DEFAULT_TIMEOUT_MS = 180_000;
+const MAX_IMAGE_INPUT_BYTES = 20 * 1024 * 1024;
+const MAX_IMAGE_INPUTS = 5;
+const MAX_TOTAL_IMAGE_INPUT_BYTES = 50 * 1024 * 1024;
+const SUPPORTED_INPUT_IMAGE_FORMATS = new Set(["png", "jpeg", "jpg", "webp", "gif"]);
+const SSE_EVENT_BOUNDARY = /\r?\n\r?\n/;
 
 export const IMAGE_SAVE_MODES = ["none", "project", "global", "custom"] as const;
 export const IMAGE_ACTIONS = ["auto", "generate", "edit"] as const;
@@ -36,6 +49,7 @@ const TOOL_PARAMS = {
     },
     images: {
       type: "array",
+      maxItems: MAX_IMAGE_INPUTS,
       items: { type: "string" },
       description: "Local image paths to use as edit targets or references.",
     },
@@ -70,11 +84,7 @@ type ToolParams = {
   saveDir?: string;
 };
 
-type CodexImageCredentials = {
-  accessToken: string;
-  accountId: string;
-  source: "modelRegistry" | "authFile";
-};
+type CodexImageCredentials = CodexCredentialsWithSource;
 
 type ImageInput = {
   path: string;
@@ -112,65 +122,12 @@ export type ImageGenerationDebug = {
   lastError?: string;
 };
 
-function decodeBase64Url(value: string): string {
-  const normalized = value.replace(/-/g, "+").replace(/_/g, "/");
-  const padded = normalized + "=".repeat((4 - (normalized.length % 4)) % 4);
-  return Buffer.from(padded, "base64").toString("utf8");
-}
-
-export function extractAccountIdFromJwt(token: string): string | undefined {
-  try {
-    const [, payload] = token.split(".");
-    if (!payload) return undefined;
-    const parsed = JSON.parse(decodeBase64Url(payload)) as unknown;
-    if (!isRecord(parsed)) return undefined;
-    const auth = parsed["https://api.openai.com/auth"];
-    if (!isRecord(auth)) return undefined;
-    const accountId = auth.chatgpt_account_id;
-    return typeof accountId === "string" && accountId.trim() ? accountId.trim() : undefined;
-  } catch {
-    return undefined;
-  }
-}
-
-function parseRegistryCredentials(
-  raw: string | undefined,
-): Omit<CodexImageCredentials, "source"> | undefined {
-  const value = raw?.trim();
-  if (!value) return undefined;
-  try {
-    const parsed = JSON.parse(value) as unknown;
-    if (isRecord(parsed)) {
-      const accessToken =
-        typeof parsed.access === "string"
-          ? parsed.access
-          : typeof parsed.token === "string"
-            ? parsed.token
-            : undefined;
-      const accountId =
-        typeof parsed.accountId === "string"
-          ? parsed.accountId
-          : typeof parsed.account_id === "string"
-            ? parsed.account_id
-            : undefined;
-      if (accessToken?.trim() && accountId?.trim())
-        return { accessToken: accessToken.trim(), accountId: accountId.trim() };
-    }
-  } catch {
-    // Plain bearer token is expected for openai-codex in pi.
-  }
-  const accountId = extractAccountIdFromJwt(value);
-  return accountId ? { accessToken: value, accountId } : undefined;
-}
-
-async function getCredentials(ctx: ExtensionContext): Promise<CodexImageCredentials> {
-  const registryToken = await ctx.modelRegistry
-    .getApiKeyForProvider("openai-codex")
-    .catch(() => undefined);
-  const registryCredentials = parseRegistryCredentials(registryToken);
-  if (registryCredentials) return { ...registryCredentials, source: "modelRegistry" };
-  const auth = readCodexAuth();
-  if (auth) return { ...auth, source: "authFile" };
+async function getCredentials(
+  ctx: ExtensionContext,
+  signal?: AbortSignal,
+): Promise<CodexImageCredentials> {
+  const credentials = await getCodexCredentials(ctx, signal);
+  if (credentials) return credentials;
   throw new Error("Missing openai-codex OAuth credentials. Run /login openai-codex.");
 }
 
@@ -193,12 +150,14 @@ function resolveImageConfig(cfg: ResolvedConfig, params: ToolParams) {
 }
 
 function imageMimeType(path: string, outputFormat?: string): string {
+  if (outputFormat === "jpeg" || outputFormat === "jpg") return "image/jpeg";
+  if (outputFormat === "webp") return "image/webp";
+  if (outputFormat === "gif") return "image/gif";
+  if (outputFormat === "png") return "image/png";
   const ext = extname(path).toLowerCase();
   if (ext === ".jpg" || ext === ".jpeg") return "image/jpeg";
   if (ext === ".webp") return "image/webp";
   if (ext === ".gif") return "image/gif";
-  if (outputFormat === "jpeg") return "image/jpeg";
-  if (outputFormat === "webp") return "image/webp";
   return "image/png";
 }
 
@@ -206,16 +165,75 @@ function extensionForFormat(format: ImageOutputFormat): string {
   return format === "jpeg" ? "jpg" : format;
 }
 
+function isInsideDirectory(root: string, child: string): boolean {
+  const normalizedRoot = resolve(root);
+  const normalizedChild = resolve(child);
+  return (
+    normalizedChild !== normalizedRoot && normalizedChild.startsWith(`${normalizedRoot}${sep}`)
+  );
+}
+
+async function validateImageInput(
+  path: string,
+  realWorkspaceRoot: string,
+): Promise<{ mimeType: string; path: string; size: number }> {
+  const realInputPath = await realpath(path).catch(() => undefined);
+  if (!realInputPath || !isInsideDirectory(realWorkspaceRoot, realInputPath))
+    throw new Error(
+      `Image input must be a file inside the current workspace: ${displayPath(path)}`,
+    );
+
+  const pathStats = await stat(realInputPath).catch(() => undefined);
+  if (!pathStats?.isFile())
+    throw new Error(
+      `Image input must be a file inside the current workspace: ${displayPath(path)}`,
+    );
+  if (pathStats.size > MAX_IMAGE_INPUT_BYTES)
+    throw new Error(`Image input is too large (max 20 MB): ${displayPath(path)}`);
+
+  const metadata = await sharp(realInputPath, { animated: false })
+    .metadata()
+    .catch(() => undefined);
+  if (!metadata?.format || !SUPPORTED_INPUT_IMAGE_FORMATS.has(metadata.format))
+    throw new Error(`Image input is not a readable image: ${displayPath(path)}`);
+  return {
+    mimeType: imageMimeType(path, metadata.format),
+    path: realInputPath,
+    size: pathStats.size,
+  };
+}
+
 async function readImageInputs(paths: string[] | undefined, cwd: string): Promise<ImageInput[]> {
-  const inputs: ImageInput[] = [];
+  const validatedInputs: Array<{ path: string; mimeType: string }> = [];
+  const seenPaths = new Set<string>();
+  let totalBytes = 0;
+  const workspaceRoot = resolve(cwd);
+  let realWorkspaceRoot: string | undefined;
   for (const rawPath of paths ?? []) {
     const trimmed = rawPath.trim();
     if (!trimmed) continue;
-    const path = isAbsolute(trimmed) ? trimmed : resolve(cwd, trimmed);
-    const data = (await readFile(path)).toString("base64");
-    inputs.push({ path, data, mimeType: imageMimeType(path) });
+    const path = isAbsolute(trimmed) ? resolve(trimmed) : resolve(workspaceRoot, trimmed);
+    if (!isInsideDirectory(workspaceRoot, path))
+      throw new Error(
+        `Image input must be a file inside the current workspace: ${displayPath(path)}`,
+      );
+    realWorkspaceRoot ??= await realpath(workspaceRoot).catch(() => workspaceRoot);
+    const input = await validateImageInput(path, realWorkspaceRoot);
+    if (seenPaths.has(input.path)) continue;
+    if (validatedInputs.length >= MAX_IMAGE_INPUTS)
+      throw new Error(`Too many image inputs (max ${MAX_IMAGE_INPUTS}).`);
+    totalBytes += input.size;
+    if (totalBytes > MAX_TOTAL_IMAGE_INPUT_BYTES)
+      throw new Error("Image inputs are too large in total (max 50 MB).");
+    seenPaths.add(input.path);
+    validatedInputs.push({ path: input.path, mimeType: input.mimeType });
   }
-  return inputs;
+  return Promise.all(
+    validatedInputs.map(async (input) => ({
+      ...input,
+      data: (await readFile(input.path)).toString("base64"),
+    })),
+  );
 }
 
 function resolveSaveDir(
@@ -225,14 +243,10 @@ function resolveSaveDir(
 ): string | undefined {
   if (mode === "none") return undefined;
   if (mode === "project") return join(cwd, ".pi", "generated-images");
-  if (mode === "global")
-    return join(
-      process.env.PI_CODING_AGENT_DIR?.trim() || join(homedir(), ".pi", "agent"),
-      "generated-images",
-    );
+  if (mode === "global") return join(piAgentDir(), "generated-images");
   const dir = params.saveDir?.trim() || process.env.PI_IMAGE_SAVE_DIR?.trim();
   if (!dir) throw new Error("save=custom requires saveDir or PI_IMAGE_SAVE_DIR.");
-  return dir;
+  return resolveUserPath(dir, cwd);
 }
 
 async function saveImage(
@@ -335,10 +349,15 @@ function extractImageFromEvent(
       mimeType,
     };
   }
-  const partial = event.partial_image_b64 ?? event.b64_json;
+  const partial =
+    typeof event.partial_image_b64 === "string"
+      ? event.partial_image_b64
+      : typeof event.b64_json === "string"
+        ? event.b64_json
+        : undefined;
   if (typeof partial === "string" && partial.trim()) {
     const { data, mimeType } = dataUrlParts(partial, fallbackMimeType);
-    return { id: `ig_${randomUUID().slice(0, 8)}`, status: "completed", data, mimeType };
+    return { id: `ig_${randomUUID().slice(0, 8)}`, status: "partial", data, mimeType };
   }
   return undefined;
 }
@@ -352,19 +371,19 @@ async function parseSseForImage(
   const reader = response.body.getReader();
   const decoder = new TextDecoder();
   let buffer = "";
-  let lastImage: ExtractedImageResult | undefined;
   try {
     while (true) {
       if (signal?.aborted) throw new Error("Image request was aborted.");
       const { done, value } = await reader.read();
       if (done) break;
       buffer += decoder.decode(value, { stream: true });
-      let idx = buffer.indexOf("\n\n");
-      while (idx !== -1) {
+      let boundary = SSE_EVENT_BOUNDARY.exec(buffer);
+      while (boundary) {
+        const idx = boundary.index;
         const chunk = buffer.slice(0, idx);
-        buffer = buffer.slice(idx + 2);
+        buffer = buffer.slice(idx + boundary[0].length);
         const data = chunk
-          .split("\n")
+          .split(/\r?\n/)
           .filter((line) => line.startsWith("data:"))
           .map((line) => line.slice(5).trim())
           .join("\n")
@@ -377,36 +396,34 @@ async function parseSseForImage(
             event = undefined;
           }
           const image = extractImageFromEvent(event, fallbackMimeType);
-          if (image?.data) {
-            lastImage = image;
-            if (image.status === "completed") {
-              await reader.cancel().catch(() => undefined);
-              return image;
-            }
+          if (image?.data && image.status === "completed") {
+            await reader.cancel().catch(() => undefined);
+            return image;
           }
           if (isRecord(event) && event.type === "response.failed") {
             const error =
               isRecord(event.response) && isRecord(event.response.error)
                 ? event.response.error
                 : undefined;
-            const message =
-              typeof error?.message === "string" ? error.message : "Codex image request failed.";
+            const message = sanitizeDiagnosticError(
+              typeof error?.message === "string" ? error.message : "Codex image request failed.",
+            );
             throw new Error(message);
           }
           if (isRecord(event) && event.type === "error") {
-            const message =
-              typeof event.message === "string" ? event.message : JSON.stringify(event);
+            const message = sanitizeDiagnosticError(
+              typeof event.message === "string" ? event.message : JSON.stringify(event),
+            );
             throw new Error(`Codex image error: ${message}`);
           }
         }
-        idx = buffer.indexOf("\n\n");
+        boundary = SSE_EVENT_BOUNDARY.exec(buffer);
       }
     }
   } finally {
     reader.releaseLock();
   }
-  if (lastImage) return lastImage;
-  throw new Error("No image_generation_call result returned by Codex.");
+  throw new Error("No completed image_generation_call result returned by Codex.");
 }
 
 async function requestCodexImage(
@@ -416,14 +433,16 @@ async function requestCodexImage(
   requestSignal?: AbortSignal,
 ): Promise<CodexImageResult> {
   if (!cfg.image.enabled) throw new Error("OpenAI image generation is disabled in config.");
-  const credentials = await getCredentials(ctx);
+  const cwd = ctx.cwd || process.cwd();
   const model = resolveModel(params, ctx, cfg);
   const { action, outputFormat, save } = resolveImageConfig(cfg, params);
-  const images = await readImageInputs(params.images, ctx.cwd || process.cwd());
-  const request = buildRequest(params, model, cfg, images);
+  const saveDir = resolveSaveDir(save, params, cwd);
   const timeoutSignal = AbortSignal.timeout(cfg.image.timeoutMs);
   const baseSignal = requestSignal ?? ctx.signal;
   const signal = baseSignal ? AbortSignal.any([baseSignal, timeoutSignal]) : timeoutSignal;
+  const credentials = await getCredentials(ctx, signal);
+  const images = await readImageInputs(params.images, cwd);
+  const request = buildRequest(params, model, cfg, images);
   const response = await fetch(CODEX_RESPONSES_URL, {
     method: "POST",
     headers: {
@@ -439,17 +458,16 @@ async function requestCodexImage(
     signal,
   });
   if (!response.ok) {
-    const text = await response.text().catch(() => "");
-    throw new Error(
-      `Codex image request failed (${response.status}): ${text || response.statusText}`,
-    );
+    const statusText = response.statusText
+      ? ` ${sanitizeDiagnosticError(response.statusText, 120)}`
+      : "";
+    throw new Error(`Codex image request failed (${response.status}${statusText}).`);
   }
   const parsed = await parseSseForImage(
     response,
     imageMimeType(`image.${outputFormat}`, outputFormat),
     signal,
   );
-  const saveDir = resolveSaveDir(save, params, ctx.cwd || process.cwd());
   const savedPath = saveDir
     ? await saveImage(parsed.data, outputFormat, saveDir, parsed.id)
     : undefined;
@@ -464,37 +482,6 @@ function displayPath(path: string): string {
   return path.startsWith(homePrefix) ? `~/${path.slice(homePrefix.length)}` : path;
 }
 
-function textFromMessageContent(content: unknown): string | undefined {
-  if (typeof content === "string") return content.trim() || undefined;
-  if (!Array.isArray(content)) return undefined;
-  const text = content
-    .filter((part) => isRecord(part) && part.type === "text" && typeof part.text === "string")
-    .map((part) => (part as { text: string }).text)
-    .join("\n")
-    .trim();
-  return text || undefined;
-}
-
-function latestUserPromptFromEntries(entries: unknown[]): string | undefined {
-  for (let i = entries.length - 1; i >= 0; i--) {
-    const entry = entries[i];
-    if (
-      !isRecord(entry) ||
-      entry.type !== "message" ||
-      !isRecord(entry.message) ||
-      entry.message.role !== "user"
-    )
-      continue;
-    const text = textFromMessageContent(entry.message.content);
-    if (text) return text;
-  }
-  return undefined;
-}
-
-function resolveToolPrompt(params: ToolParams, ctx: ExtensionContext): string {
-  return latestUserPromptFromEntries(ctx.sessionManager.getEntries()) ?? params.prompt;
-}
-
 function resultText(result: CodexImageResult): string {
   const parts = [
     `Generated image using OpenAI image_generation tool via openai-codex/${result.model}.`,
@@ -526,7 +513,7 @@ export function registerOpenAIImage(
       return result;
     } catch (error) {
       lastStatus = "error";
-      lastError = error instanceof Error ? error.message : String(error);
+      lastError = sanitizeDiagnosticError(error instanceof Error ? error.message : String(error));
       throw error;
     }
   }
@@ -542,7 +529,7 @@ export function registerOpenAIImage(
     return {
       authFound: credentials !== undefined,
       authSource: credentials?.source,
-      accountId: credentials?.accountId,
+      accountId: maskIdentifier(credentials?.accountId),
       endpoint: CODEX_RESPONSES_URL,
       defaultModel: ctx.model?.provider === "openai-codex" ? ctx.model.id : cfg.image.defaultModel,
       defaultSave: cfg.image.defaultSave,
@@ -552,63 +539,57 @@ export function registerOpenAIImage(
     };
   }
 
-  void import("@mariozechner/pi-tui")
-    .then(({ Box, Container, Image, Text }) => {
-      pi.registerMessageRenderer<CodexImageResult>("openai-image", (message, _options, theme) => {
-        const result = message.details;
-        const text =
-          result && isRecord(result)
-            ? resultText(result as CodexImageResult)
-            : typeof message.content === "string"
-              ? message.content
-              : message.content
-                  .filter((part) => part.type === "text")
-                  .map((part) => part.text)
-                  .join("\n");
-        let image: { data: string; mimeType: string; savedPath?: string } | undefined;
-        if (
-          result &&
-          isRecord(result) &&
-          typeof result.data === "string" &&
-          typeof result.mimeType === "string"
-        ) {
-          image = {
-            data: result.data,
-            mimeType: result.mimeType,
-            savedPath: typeof result.savedPath === "string" ? result.savedPath : undefined,
-          };
-        } else if (Array.isArray(message.content)) {
-          const imagePart = message.content.find(isImageContent);
-          if (imagePart) image = { data: imagePart.data, mimeType: imagePart.mimeType };
-        }
+  pi.registerMessageRenderer<CodexImageResult>("openai-image", (message, _options, theme) => {
+    const result = message.details;
+    const text =
+      result && isRecord(result)
+        ? resultText(result as CodexImageResult)
+        : typeof message.content === "string"
+          ? message.content
+          : message.content
+              .filter((part) => part.type === "text")
+              .map((part) => part.text)
+              .join("\n");
+    let image: { data: string; mimeType: string; savedPath?: string } | undefined;
+    if (
+      result &&
+      isRecord(result) &&
+      typeof result.data === "string" &&
+      typeof result.mimeType === "string"
+    ) {
+      image = {
+        data: result.data,
+        mimeType: result.mimeType,
+        savedPath: typeof result.savedPath === "string" ? result.savedPath : undefined,
+      };
+    } else if (Array.isArray(message.content)) {
+      const imagePart = message.content.find(isImageContent);
+      if (imagePart) image = { data: imagePart.data, mimeType: imagePart.mimeType };
+    }
 
-        const container = new Container();
-        const box = new Box(1, 1, (line) => theme.bg("customMessageBg", line));
-        box.addChild(
-          new Text(`${theme.fg("accent", theme.bold("[openai-image]"))}\n\n${text}`, 0, 0),
-        );
-        if (image) {
-          box.addChild(
-            new Image(
-              image.data,
-              image.mimeType,
-              { fallbackColor: (line) => theme.fg("dim", line) },
-              {
-                maxWidthCells: 80,
-                maxHeightCells: 24,
-                filename:
-                  "savedPath" in image && typeof image.savedPath === "string"
-                    ? image.savedPath
-                    : undefined,
-              },
-            ),
-          );
-        }
-        container.addChild(box);
-        return container;
-      });
-    })
-    .catch(() => undefined);
+    const container = new Container();
+    const box = new Box(1, 1, (line) => theme.bg("customMessageBg", line));
+    box.addChild(new Text(`${theme.fg("accent", theme.bold("[openai-image]"))}\n\n${text}`, 0, 0));
+    if (image) {
+      box.addChild(
+        new Image(
+          image.data,
+          image.mimeType,
+          { fallbackColor: (line) => theme.fg("dim", line) },
+          {
+            maxWidthCells: 80,
+            maxHeightCells: 24,
+            filename:
+              "savedPath" in image && typeof image.savedPath === "string"
+                ? image.savedPath
+                : undefined,
+          },
+        ),
+      );
+    }
+    container.addChild(box);
+    return container;
+  });
 
   pi.registerCommand(OPENAI_IMAGE_COMMAND, {
     description: "Generate an image with OpenAI Codex image generation",
@@ -647,7 +628,6 @@ export function registerOpenAIImage(
     async execute(_toolCallId, params: ToolParams, signal, onUpdate, ctx) {
       const cfg = getConfig(ctx);
       const model = resolveModel(params, ctx, cfg);
-      const requestParams = { ...params, prompt: resolveToolPrompt(params, ctx) };
       onUpdate?.({
         content: [
           {
@@ -657,7 +637,7 @@ export function registerOpenAIImage(
         ],
         details: undefined,
       });
-      const result = await generate(requestParams, ctx, signal);
+      const result = await generate(params, ctx, signal);
       return {
         content: [
           { type: "text", text: resultText(result) },
@@ -676,11 +656,13 @@ export const _imageTest = {
   DEFAULT_TIMEOUT_MS,
   OPENAI_IMAGE_TOOL,
   OPENAI_IMAGE_COMMAND,
+  MAX_IMAGE_INPUT_BYTES,
+  MAX_IMAGE_INPUTS,
+  MAX_TOTAL_IMAGE_INPUT_BYTES,
   extractAccountIdFromJwt,
   imageMimeType,
   dataUrlParts,
   extractImageFromEvent,
   displayPath,
-  latestUserPromptFromEntries,
   buildRequest,
 };

package/src/paths.ts

@@ -0,0 +1,17 @@
+import { homedir } from "node:os";
+import { join, resolve } from "node:path";
+
+export function expandTildePath(path: string, home = homedir()): string {
+  if (path === "~") return home;
+  if (path.startsWith("~/") || path.startsWith("~\\")) return join(home, path.slice(2));
+  return path;
+}
+
+export function piAgentDir(env = process.env, home = homedir()): string {
+  const configuredDir = env.PI_CODING_AGENT_DIR?.trim();
+  return configuredDir ? expandTildePath(configuredDir, home) : join(home, ".pi", "agent");
+}
+
+export function resolveUserPath(path: string, cwd: string, home = homedir()): string {
+  return resolve(cwd, expandTildePath(path, home));
+}