@centerforagenticai/pi-multi-account 0.1.1

package/packages/pi-antigravity/src/diagnostics/diagnostics.ts

@@ -0,0 +1,96 @@
+import { AsyncLocalStorage } from "node:async_hooks";
+import { redactSecrets } from "../utils/security.js";
+
+export type DiagnosticsSnapshot = {
+  status?: number;
+  endpoint?: string;
+  error?: string;
+  projectId?: string;
+  resolvedRuntimeModel?: string;
+  availableModels?: string;
+  matchedModelDebug?: string;
+  latencyMs?: number;
+  maskedEmail?: string;
+  tokenExpiry?: string;
+  toolSchemaWarnings?: string;
+};
+
+const storage = new AsyncLocalStorage<DiagnosticsSnapshot>();
+
+/** Last completed request snapshot for `/antigravity.doctor`. */
+let lastSnapshot: DiagnosticsSnapshot = {};
+
+function currentBag(): DiagnosticsSnapshot {
+  return storage.getStore() ?? lastSnapshot;
+}
+
+/** Run work with an isolated diagnostics bag; commits it to `lastSnapshot` when done. */
+export async function runWithDiagnostics<T>(fn: () => Promise<T>): Promise<T> {
+  const bag: DiagnosticsSnapshot = {};
+  return storage.run(bag, async () => {
+    try {
+      return await fn();
+    } finally {
+      lastSnapshot = { ...bag };
+    }
+  });
+}
+
+export function getLastDiagnostics(): Readonly<DiagnosticsSnapshot> {
+  return lastSnapshot;
+}
+
+/** Read endpoint from the active request bag (or last snapshot outside a request). */
+export function getCurrentEndpoint(): string | undefined {
+  return currentBag().endpoint;
+}
+
+export function getCurrentMatchedModelDebug(): string | undefined {
+  return currentBag().matchedModelDebug;
+}
+
+export function getCurrentAvailableModels(): string | undefined {
+  return currentBag().availableModels;
+}
+
+export function setLastStatus(status: number | undefined): void {
+  currentBag().status = status;
+}
+export function setLastEndpoint(endpoint: string | undefined): void {
+  currentBag().endpoint = endpoint;
+}
+export function setLastError(error: string | undefined): void {
+  currentBag().error = error === undefined ? undefined : redactSecrets(error).slice(0, 800);
+}
+export function setLastProjectId(projectId: string | undefined): void {
+  currentBag().projectId = projectId;
+}
+export function setLastResolvedRuntimeModel(model: string | undefined): void {
+  currentBag().resolvedRuntimeModel = model;
+}
+export function setLastAvailableModels(models: string | undefined): void {
+  currentBag().availableModels = models;
+}
+export function setLastMatchedModelDebug(debug: string | undefined): void {
+  currentBag().matchedModelDebug =
+    debug === undefined ? undefined : redactSecrets(debug).slice(0, 1200);
+}
+export function setLastLatencyMs(ms: number | undefined): void {
+  currentBag().latencyMs = ms;
+}
+export function setLastMaskedEmail(email: string | undefined): void {
+  currentBag().maskedEmail = email;
+}
+export function setLastTokenExpiry(expiry: string | undefined): void {
+  currentBag().tokenExpiry = expiry;
+}
+/** Store sanitized tool-schema omissions for the next doctor report. */
+export function setLastToolSchemaWarnings(warnings: string[] | undefined): void {
+  currentBag().toolSchemaWarnings =
+    warnings === undefined ? undefined : redactSecrets(warnings.join(" | ")).slice(0, 1200);
+}
+
+/** Test helper: reset last snapshot between cases. */
+export function resetDiagnosticsForTests(): void {
+  lastSnapshot = {};
+}

package/packages/pi-antigravity/src/diagnostics/index.ts

@@ -0,0 +1 @@
+export * from "./diagnostics.js";

package/packages/pi-antigravity/src/image/image.ts

@@ -0,0 +1,336 @@
+import { mkdir, writeFile } from "node:fs/promises";
+import { dirname, extname, isAbsolute, join, relative, resolve, sep } from "node:path";
+import {
+  antigravityHeaders,
+  endpointCandidates,
+  jsonOrTextError,
+  parseApiKey,
+} from "../client/client.js";
+import { AntigravityRequestType, AntigravityUserAgent, GeminiRole } from "../types/enums.js";
+import { antigravityFetch } from "../utils/http.js";
+import { safeError } from "../utils/security.js";
+import { antigravityRequestEnvelope, sanitizeText } from "../utils/util.js";
+
+export const DEFAULT_IMAGE_MODEL = "gemini-3-pro-image";
+export const IMAGE_ASPECT_RATIOS = [
+  "1:1",
+  "2:3",
+  "3:2",
+  "3:4",
+  "4:3",
+  "4:5",
+  "5:4",
+  "9:16",
+  "16:9",
+  "21:9",
+] as const;
+export type ImageAspectRatio = (typeof IMAGE_ASPECT_RATIOS)[number];
+
+const IMAGE_MODEL_FALLBACKS = [
+  DEFAULT_IMAGE_MODEL,
+  "gemini-3.1-flash-image",
+  "gemini-3-pro-image-preview",
+];
+const IMAGE_SYSTEM_INSTRUCTION =
+  "You are an AI image generator. Generate images based on user descriptions. Focus on creating high-quality, visually appealing images that match the user's request.";
+const DEFAULT_IMAGE_DIR = join(".pi", "generated-images");
+const MAX_PROMPT_CHARS = 8000;
+
+export type GeneratedImage = { data: string; mimeType: string };
+
+export type ImageGenerateRequest = {
+  project: string;
+  model: string;
+  request: {
+    contents: Array<{ role: GeminiRole.User; parts: Array<{ text: string }> }>;
+    systemInstruction: { role: GeminiRole.User; parts: Array<{ text: string }> };
+    generationConfig: {
+      imageConfig: { aspectRatio: string };
+      candidateCount: number;
+    };
+  };
+  requestType: AntigravityRequestType.Agent;
+  userAgent: AntigravityUserAgent.Antigravity;
+  requestId: string;
+};
+
+export type ImageCommandArgs = {
+  prompt: string;
+  aspectRatio?: string;
+  model?: string;
+  path?: string;
+};
+
+export type GenerateImageOptions = ImageCommandArgs & {
+  apiKey: string;
+  cwd: string;
+  signal?: AbortSignal;
+};
+
+export type GenerateImageResult = {
+  images: GeneratedImage[];
+  savedPaths: string[];
+  text: string[];
+  model: string;
+};
+
+type ImageStreamChunk = {
+  error?: { message?: string };
+  response?: {
+    candidates?: Array<{
+      content?: {
+        parts?: Array<{ text?: string; inlineData?: { mimeType?: string; data?: string } }>;
+      };
+    }>;
+  };
+  candidates?: Array<{
+    content?: {
+      parts?: Array<{ text?: string; inlineData?: { mimeType?: string; data?: string } }>;
+    };
+  }>;
+};
+
+function imageExtension(mimeType: string): string {
+  const lower = mimeType.toLowerCase();
+  if (lower.includes("jpeg") || lower.includes("jpg")) return "jpg";
+  if (lower.includes("webp")) return "webp";
+  if (lower.includes("gif")) return "gif";
+  return "png";
+}
+
+export function assertSafeImageModel(modelId: string): string {
+  const id = modelId.trim();
+  if (id.length === 0 || id.length > 80) {
+    throw new Error("Unsupported image model id.");
+  }
+  if (!/^(gemini-[a-z0-9.+-]*image[a-z0-9.+-]*|imagen-[a-z0-9.+-]+)$/i.test(id)) {
+    throw new Error(`Unsupported image model: ${id}`);
+  }
+  return id;
+}
+
+export function assertSafeAspectRatio(ratio: string): ImageAspectRatio {
+  const value = ratio.trim();
+  for (const allowed of IMAGE_ASPECT_RATIOS) {
+    if (allowed === value) return allowed;
+  }
+  throw new Error(
+    `Unsupported aspect ratio: ${value}. Use one of ${IMAGE_ASPECT_RATIOS.join(", ")}.`,
+  );
+}
+
+export function parseImageCommandArgs(args: string): ImageCommandArgs {
+  const tokens = args.trim().split(/\s+/).filter(Boolean);
+  const out: ImageCommandArgs = { prompt: "" };
+  const rest: string[] = [];
+  for (let i = 0; i < tokens.length; i++) {
+    const token = tokens[i];
+    if (token === undefined) continue;
+    const next = tokens[i + 1];
+    if ((token === "--ratio" || token === "--aspect-ratio") && next) {
+      out.aspectRatio = next;
+      i += 1;
+      continue;
+    }
+    if (token === "--model" && next) {
+      out.model = next;
+      i += 1;
+      continue;
+    }
+    if (token === "--path" && next) {
+      out.path = next;
+      i += 1;
+      continue;
+    }
+    rest.push(token);
+  }
+  out.prompt = rest.join(" ");
+  return out;
+}
+
+export function resolveImageSavePath(
+  cwd: string,
+  requested?: string,
+  mimeType = "image/png",
+  index?: number,
+): string {
+  const ext = imageExtension(mimeType);
+  const stamp = new Date().toISOString().replace(/[:.]/g, "-");
+  const suffix = index === undefined ? "" : `-${index + 1}`;
+  const defaultName = `image-${stamp}${suffix}.${ext}`;
+  const root = resolve(cwd);
+  const target = requested?.trim()
+    ? resolve(root, requested.trim())
+    : resolve(root, DEFAULT_IMAGE_DIR, defaultName);
+  const rel = relative(root, target);
+  if (!rel || rel === ".." || rel.startsWith(`..${sep}`) || isAbsolute(rel)) {
+    throw new Error("Image save path must be inside the working directory.");
+  }
+  if (!extname(target)) return join(target, defaultName);
+  if (index === undefined) return target;
+  const currentExt = extname(target);
+  return `${target.slice(0, -currentExt.length)}${suffix}${currentExt}`;
+}
+
+export function buildImageGenerateRequest(
+  prompt: string,
+  model: string,
+  projectId: string,
+  aspectRatio: string,
+): ImageGenerateRequest {
+  const envelope = antigravityRequestEnvelope(model, false);
+  return {
+    project: projectId,
+    model,
+    request: {
+      contents: [{ role: GeminiRole.User, parts: [{ text: sanitizeText(prompt) }] }],
+      systemInstruction: {
+        role: GeminiRole.User,
+        parts: [{ text: IMAGE_SYSTEM_INSTRUCTION }],
+      },
+      generationConfig: {
+        imageConfig: { aspectRatio },
+        candidateCount: 1,
+      },
+    },
+    requestType: AntigravityRequestType.Agent,
+    userAgent: AntigravityUserAgent.Antigravity,
+    requestId: envelope.requestId,
+  };
+}
+
+function collectImagesFromParts(
+  parts: Array<{ text?: string; inlineData?: { mimeType?: string; data?: string } }> | undefined,
+  images: GeneratedImage[],
+  text: string[],
+): void {
+  for (const part of parts || []) {
+    if (part.text) text.push(part.text);
+    if (part.inlineData?.data) {
+      images.push({
+        data: part.inlineData.data,
+        mimeType: part.inlineData.mimeType || "image/png",
+      });
+    }
+  }
+}
+
+export async function collectImagesFromSse(
+  response: Response,
+  signal?: AbortSignal,
+): Promise<{ images: GeneratedImage[]; text: string[] }> {
+  if (!response.body) throw new Error("No response body");
+  const reader = response.body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+  const images: GeneratedImage[] = [];
+  const text: string[] = [];
+  try {
+    while (true) {
+      if (signal?.aborted) throw new Error("Request was aborted");
+      const result = await reader.read();
+      if (result.done) break;
+      if (!(result.value instanceof Uint8Array)) continue;
+      buffer += decoder.decode(result.value, { stream: true });
+      const lines = buffer.split("\n");
+      buffer = lines.pop() || "";
+      for (const line of lines) {
+        if (!line.startsWith("data:")) continue;
+        const json = line.slice(5).trim();
+        if (!json || json === "[DONE]") continue;
+        let chunk: ImageStreamChunk;
+        try {
+          chunk = JSON.parse(json) as ImageStreamChunk;
+        } catch {
+          continue;
+        }
+        if (chunk.error?.message) throw new Error(chunk.error.message);
+        const responseData = chunk.response || chunk;
+        for (const candidate of responseData.candidates || []) {
+          collectImagesFromParts(candidate.content?.parts, images, text);
+        }
+      }
+    }
+  } finally {
+    reader.releaseLock();
+  }
+  return { images, text };
+}
+
+async function writeImage(filePath: string, image: GeneratedImage): Promise<string> {
+  await mkdir(dirname(filePath), { recursive: true });
+  await writeFile(filePath, Buffer.from(image.data, "base64"));
+  return filePath;
+}
+
+export async function generateAntigravityImage(
+  options: GenerateImageOptions,
+): Promise<GenerateImageResult> {
+  const prompt = options.prompt.trim();
+  if (!prompt) throw new Error("Image prompt is required.");
+  if (prompt.length > MAX_PROMPT_CHARS) {
+    throw new Error(`Image prompt is too long (max ${MAX_PROMPT_CHARS} characters).`);
+  }
+  const aspectRatio = assertSafeAspectRatio(options.aspectRatio || "1:1");
+  const preferred = assertSafeImageModel(options.model || DEFAULT_IMAGE_MODEL);
+  const models = [preferred, ...IMAGE_MODEL_FALLBACKS.filter((id) => id !== preferred)];
+  const creds = parseApiKey(options.apiKey);
+  const headers = antigravityHeaders(creds.token);
+
+  let lastError = "no endpoint available";
+  for (const model of models) {
+    const body = JSON.stringify(
+      buildImageGenerateRequest(prompt, model, creds.projectId, aspectRatio),
+    );
+    for (const endpoint of endpointCandidates()) {
+      if (options.signal?.aborted) throw new Error("Request was aborted");
+      try {
+        const response = await antigravityFetch(
+          `${endpoint}/v1internal:streamGenerateContent?alt=sse`,
+          {
+            method: "POST",
+            headers,
+            body,
+            signal: options.signal,
+          },
+        );
+        if (!response.ok) {
+          lastError = jsonOrTextError(await response.text()).slice(0, 400);
+          if (response.status === 404 || [403, 429, 500, 502, 503, 504].includes(response.status)) {
+            continue;
+          }
+          throw new Error(
+            `Antigravity image request failed (${response.status}): ${safeError(lastError)}`,
+          );
+        }
+        const parsed = await collectImagesFromSse(response, options.signal);
+        if (!parsed.images.length) {
+          lastError = parsed.text.join(" ").trim() || "No image data returned.";
+          continue;
+        }
+        const savedPaths: string[] = [];
+        const many = parsed.images.length > 1;
+        for (const [index, image] of parsed.images.entries()) {
+          savedPaths.push(
+            await writeImage(
+              resolveImageSavePath(
+                options.cwd,
+                options.path,
+                image.mimeType,
+                many ? index : undefined,
+              ),
+              image,
+            ),
+          );
+        }
+        return { images: parsed.images, savedPaths, text: parsed.text, model };
+      } catch (error) {
+        lastError = safeError(error);
+        if (options.signal?.aborted) {
+          throw new Error("Request was aborted", { cause: error });
+        }
+      }
+    }
+  }
+  throw new Error(`Antigravity image generation failed: ${safeError(lastError)}`);
+}

package/packages/pi-antigravity/src/image/index.ts

@@ -0,0 +1 @@
+export * from "./image.js";