@astrofoundry/pi-astro 0.5.0 → 0.6.0

package/extensions/gemini-image/models.test.ts

@@ -0,0 +1,45 @@
+import { describe, expect, it } from "vitest";
+import {
+    ALL_MODELS,
+    DEFAULT_MODEL,
+    GEMINI_ASPECT_RATIOS,
+    GEMINI_NATIVE_MODELS,
+    IMAGEN_ASPECT_RATIOS,
+    IMAGEN_MODELS,
+    isGeminiNative,
+    isImagen,
+} from "./models.ts";
+
+describe("models", () => {
+    it("has 3 Gemini native + 3 Imagen = 6 models", () => {
+        expect(GEMINI_NATIVE_MODELS).toHaveLength(3);
+        expect(IMAGEN_MODELS).toHaveLength(3);
+        expect(ALL_MODELS).toHaveLength(6);
+    });
+
+    it("default is gemini-2.5-flash-image", () => {
+        expect(DEFAULT_MODEL).toBe("gemini-2.5-flash-image");
+    });
+
+    it("Gemini aspect ratios include 4:5 and 5:4 (missing from the junior's impl)", () => {
+        expect(GEMINI_ASPECT_RATIOS).toContain("4:5");
+        expect(GEMINI_ASPECT_RATIOS).toContain("5:4");
+    });
+
+    it("Imagen has 5 aspect ratios (no 2:3 / 3:2 / 4:5 / 5:4 / 21:9)", () => {
+        expect(IMAGEN_ASPECT_RATIOS).toHaveLength(5);
+        expect(IMAGEN_ASPECT_RATIOS).not.toContain("2:3");
+    });
+
+    it("isGeminiNative narrows correctly", () => {
+        expect(isGeminiNative("gemini-2.5-flash-image")).toBe(true);
+        expect(isGeminiNative("imagen-4.0-generate-001")).toBe(false);
+        expect(isGeminiNative("unknown")).toBe(false);
+    });
+
+    it("isImagen narrows correctly", () => {
+        expect(isImagen("imagen-4.0-ultra-generate-001")).toBe(true);
+        expect(isImagen("gemini-3-pro-image-preview")).toBe(false);
+        expect(isImagen("unknown")).toBe(false);
+    });
+});

package/extensions/gemini-image/models.ts

@@ -0,0 +1,50 @@
+export const GEMINI_NATIVE_MODELS = [
+    "gemini-2.5-flash-image",
+    "gemini-3.1-flash-image-preview",
+    "gemini-3-pro-image-preview",
+] as const;
+
+export const IMAGEN_MODELS = [
+    "imagen-4.0-fast-generate-001",
+    "imagen-4.0-generate-001",
+    "imagen-4.0-ultra-generate-001",
+] as const;
+
+export type GeminiNativeModel = (typeof GEMINI_NATIVE_MODELS)[number];
+export type ImagenModel = (typeof IMAGEN_MODELS)[number];
+export type ImageModel = GeminiNativeModel | ImagenModel;
+
+export const ALL_MODELS: readonly ImageModel[] = [
+    ...GEMINI_NATIVE_MODELS,
+    ...IMAGEN_MODELS,
+];
+
+export const DEFAULT_MODEL: ImageModel = "gemini-2.5-flash-image";
+
+export function isGeminiNative(model: string): model is GeminiNativeModel {
+    return (GEMINI_NATIVE_MODELS as readonly string[]).includes(model);
+}
+
+export function isImagen(model: string): model is ImagenModel {
+    return (IMAGEN_MODELS as readonly string[]).includes(model);
+}
+
+export const GEMINI_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 const IMAGEN_ASPECT_RATIOS = ["1:1", "3:4", "4:3", "9:16", "16:9"] as const;
+
+export const GEMINI_IMAGE_SIZES = ["1K", "2K", "4K"] as const;
+export const IMAGEN_IMAGE_SIZES = ["1K", "2K"] as const;
+
+export const PERSON_GENERATION_VALUES = ["dont_allow", "allow_adult", "allow_all"] as const;

package/extensions/gemini-image/pricing.test.ts

@@ -0,0 +1,95 @@
+import { describe, expect, it } from "vitest";
+import { estimateCost, formatUsd, hasPricing } from "./pricing.ts";
+
+describe("estimateCost", () => {
+    it("Imagen 4 Fast: $0.02 × N images", () => {
+        const e = estimateCost({
+            model: "imagen-4.0-fast-generate-001",
+            prompt: "cat",
+            numberOfImages: 3,
+        });
+        expect(e.estimatedUsd).toBeCloseTo(0.06, 5);
+        expect(e.breakdown).toContain("0.02");
+    });
+
+    it("Imagen 4 Standard single image", () => {
+        const e = estimateCost({
+            model: "imagen-4.0-generate-001",
+            prompt: "x",
+            numberOfImages: 1,
+        });
+        expect(e.estimatedUsd).toBeCloseTo(0.04, 5);
+    });
+
+    it("Imagen 4 Ultra", () => {
+        const e = estimateCost({
+            model: "imagen-4.0-ultra-generate-001",
+            prompt: "x",
+            numberOfImages: 4,
+        });
+        expect(e.estimatedUsd).toBeCloseTo(0.24, 5);
+    });
+
+    it("Gemini 3 Pro 4K is more expensive than 2K", () => {
+        const at4k = estimateCost({
+            model: "gemini-3-pro-image-preview",
+            prompt: "small",
+            numberOfImages: 1,
+            imageSize: "4K",
+        });
+        const at2k = estimateCost({
+            model: "gemini-3-pro-image-preview",
+            prompt: "small",
+            numberOfImages: 1,
+            imageSize: "2K",
+        });
+        expect(at4k.estimatedUsd).toBeGreaterThan(at2k.estimatedUsd);
+    });
+
+    it("Gemini 3 Pro prompt tokens contribute to cost (long prompt > short)", () => {
+        const longPrompt = "x".repeat(10_000);
+        const shortEst = estimateCost({
+            model: "gemini-3-pro-image-preview",
+            prompt: "x",
+            numberOfImages: 1,
+        });
+        const longEst = estimateCost({
+            model: "gemini-3-pro-image-preview",
+            prompt: longPrompt,
+            numberOfImages: 1,
+        });
+        expect(longEst.estimatedUsd).toBeGreaterThan(shortEst.estimatedUsd);
+    });
+
+    it("Flash estimate scales linearly with numberOfImages", () => {
+        const one = estimateCost({
+            model: "gemini-2.5-flash-image",
+            prompt: "x",
+            numberOfImages: 1,
+        });
+        const four = estimateCost({
+            model: "gemini-2.5-flash-image",
+            prompt: "x",
+            numberOfImages: 4,
+        });
+        // Output dominates; roughly 4x
+        expect(four.estimatedUsd).toBeGreaterThan(3.5 * one.estimatedUsd * 0.99);
+    });
+});
+
+describe("formatUsd", () => {
+    it("renders sub-cent as <$0.01", () => {
+        expect(formatUsd(0.002)).toBe("<$0.01");
+    });
+
+    it("renders 4 decimals for normal amounts", () => {
+        expect(formatUsd(0.1234)).toBe("$0.1234");
+    });
+});
+
+describe("hasPricing", () => {
+    it("returns true for all known models", () => {
+        expect(hasPricing("gemini-2.5-flash-image")).toBe(true);
+        expect(hasPricing("imagen-4.0-fast-generate-001")).toBe(true);
+    });
+});

package/extensions/gemini-image/pricing.ts

@@ -0,0 +1,102 @@
+/**
+ * Gemini image-generation price table.
+ *
+ * Last verified: 2026-04-24 via ai.google.dev/gemini-api/docs/pricing.
+ * Prices drift — edit when re-checked.
+ */
+
+import type { ImageModel } from "./models.ts";
+
+interface ImagenPrice {
+    kind: "flat";
+    perImage: number; // USD
+}
+
+interface GeminiNativePrice {
+    kind: "token";
+    // $ per 1M input tokens
+    inputPer1M: number;
+    // $ per output image at 1K/2K resolution
+    outputPerImage1Kor2K: number;
+    // $ per output image at 4K resolution
+    outputPerImage4K: number;
+    // Approx output tokens per image (for reference; not used in estimate)
+    tokensPerImage: number;
+}
+
+type ModelPrice = ImagenPrice | GeminiNativePrice;
+
+const PRICES: Record<ImageModel, ModelPrice> = {
+    "imagen-4.0-fast-generate-001": { kind: "flat", perImage: 0.02 },
+    "imagen-4.0-generate-001": { kind: "flat", perImage: 0.04 },
+    "imagen-4.0-ultra-generate-001": { kind: "flat", perImage: 0.06 },
+    "gemini-2.5-flash-image": {
+        kind: "token",
+        inputPer1M: 0.3,
+        outputPerImage1Kor2K: 0.039,
+        outputPerImage4K: 0.039,
+        tokensPerImage: 1290,
+    },
+    "gemini-3.1-flash-image-preview": {
+        kind: "token",
+        inputPer1M: 0.3,
+        outputPerImage1Kor2K: 0.039,
+        outputPerImage4K: 0.078,
+        tokensPerImage: 1290,
+    },
+    "gemini-3-pro-image-preview": {
+        kind: "token",
+        inputPer1M: 2.0,
+        outputPerImage1Kor2K: 0.134,
+        outputPerImage4K: 0.24,
+        tokensPerImage: 1290,
+    },
+};
+
+export interface EstimateInput {
+    model: ImageModel;
+    prompt: string;
+    numberOfImages: number;
+    imageSize?: "1K" | "2K" | "4K";
+    inputImagesBytes?: number;
+}
+
+export interface EstimateResult {
+    estimatedUsd: number;
+    breakdown: string;
+}
+
+const CHARS_PER_TOKEN_APPROX = 4;
+
+export function estimateCost(input: EstimateInput): EstimateResult {
+    const price = PRICES[input.model];
+    if (price.kind === "flat") {
+        const usd = price.perImage * input.numberOfImages;
+        return {
+            estimatedUsd: usd,
+            breakdown: `${input.numberOfImages} × $${price.perImage.toFixed(2)}/image`,
+        };
+    }
+
+    const inputTokens = Math.ceil(input.prompt.length / CHARS_PER_TOKEN_APPROX);
+    const inputCost = (inputTokens / 1_000_000) * price.inputPer1M;
+    const perImage =
+        input.imageSize === "4K" ? price.outputPerImage4K : price.outputPerImage1Kor2K;
+    const outputCost = perImage * input.numberOfImages;
+    const total = inputCost + outputCost;
+    return {
+        estimatedUsd: total,
+        breakdown:
+            `${input.numberOfImages} × $${perImage.toFixed(3)} (${input.imageSize ?? "1K/2K"}) + ` +
+            `~${inputTokens} input tokens`,
+    };
+}
+
+export function formatUsd(usd: number): string {
+    if (usd < 0.01) return `<$0.01`;
+    return `$${usd.toFixed(4)}`;
+}
+
+export function hasPricing(model: ImageModel): boolean {
+    return model in PRICES;
+}

package/extensions/grimoire/index.test.ts

@@ -0,0 +1,244 @@
+import { EventEmitter } from "node:events";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+// Mock node:child_process BEFORE importing the module under test.
+vi.mock("node:child_process", () => {
+    const mod = {
+        execFileSync: vi.fn(),
+        spawn: vi.fn(),
+    };
+    return mod;
+});
+
+type SpawnFactory = () => { proc: EventEmitter & { stdout: EventEmitter; stderr: EventEmitter; kill: (sig: string) => void }; emitClose: (code: number) => void; emitError: (err: Error) => void; emitData: (where: "stdout" | "stderr", chunk: string) => void };
+
+function makeFakeProc(): ReturnType<SpawnFactory> {
+    const proc = Object.assign(new EventEmitter(), {
+        stdout: new EventEmitter(),
+        stderr: new EventEmitter(),
+        kill: vi.fn(),
+    });
+    return {
+        proc,
+        emitClose: (code: number) => proc.emit("close", code),
+        emitError: (err: Error) => proc.emit("error", err),
+        emitData: (where, chunk) => proc[where].emit("data", Buffer.from(chunk)),
+    };
+}
+
+interface CapturedTool {
+    name: string;
+    description: string;
+    promptGuidelines?: string[];
+    execute: (id: string, input: { query: string; source?: string; top?: number }, signal?: AbortSignal, onUpdate?: unknown, ctx?: unknown) => Promise<unknown>;
+}
+
+interface CapturedPi {
+    registerTool: (tool: CapturedTool) => void;
+    tools: CapturedTool[];
+}
+
+function makePi(): CapturedPi {
+    const tools: CapturedTool[] = [];
+    return {
+        tools,
+        registerTool: (tool) => {
+            tools.push(tool);
+        },
+    };
+}
+
+function makeCtx(): { ui: { notify: ReturnType<typeof vi.fn> }; notifications: Array<[string, string | undefined]> } {
+    const notifications: Array<[string, string | undefined]> = [];
+    return {
+        notifications,
+        ui: {
+            notify: vi.fn((msg: string, type?: string) => {
+                notifications.push([msg, type]);
+            }),
+        },
+    };
+}
+
+describe("grimoire extension", () => {
+    beforeEach(() => {
+        vi.resetModules();
+        vi.clearAllMocks();
+    });
+
+    afterEach(() => {
+        vi.clearAllMocks();
+    });
+
+    it("loads --help into promptGuidelines when grimoire is available", async () => {
+        const { execFileSync } = await import("node:child_process");
+        vi.mocked(execFileSync).mockReturnValue("FAKE HELP OUTPUT" as unknown as ReturnType<typeof execFileSync>);
+
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+
+        const tool = pi.tools[0];
+        expect(tool.name).toBe("grimoire");
+        expect(tool.promptGuidelines?.some((g) => g.includes("FAKE HELP OUTPUT"))).toBe(true);
+    });
+
+    it("falls back to 'not available' promptGuidelines when grimoire CLI is missing", async () => {
+        const { execFileSync } = await import("node:child_process");
+        vi.mocked(execFileSync).mockImplementation(() => {
+            throw new Error("ENOENT");
+        });
+
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+
+        expect(pi.tools[0].promptGuidelines?.[0]).toMatch(/NOT available/i);
+    });
+
+    it("execute throws when CLI was not available at load time", async () => {
+        const { execFileSync } = await import("node:child_process");
+        vi.mocked(execFileSync).mockImplementation(() => {
+            throw new Error("ENOENT");
+        });
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+        const ctx = makeCtx();
+        await expect(
+            pi.tools[0].execute("t1", { query: "x" }, undefined, undefined, ctx),
+        ).rejects.toThrow(/not installed/i);
+        expect(ctx.notifications.some((n) => n[1] === "error")).toBe(true);
+    });
+
+    it("execute returns stdout when grimoire exits 0 with results", async () => {
+        const { execFileSync, spawn } = await import("node:child_process");
+        vi.mocked(execFileSync).mockReturnValue("help" as unknown as ReturnType<typeof execFileSync>);
+        const fake = makeFakeProc();
+        vi.mocked(spawn).mockReturnValue(fake.proc as unknown as ReturnType<typeof spawn>);
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+        const ctx = makeCtx();
+        const promise = pi.tools[0].execute("t1", { query: "react hooks" }, undefined, undefined, ctx);
+        fake.emitData("stdout", "result line 1\nresult line 2\n");
+        fake.emitClose(0);
+        const result = (await promise) as { content: Array<{ text: string }> };
+        expect(result.content[0].text).toContain("result line 1");
+    });
+
+    it("execute with source + top passes them through as CLI args", async () => {
+        const { execFileSync, spawn } = await import("node:child_process");
+        vi.mocked(execFileSync).mockReturnValue("help" as unknown as ReturnType<typeof execFileSync>);
+        const fake = makeFakeProc();
+        vi.mocked(spawn).mockReturnValue(fake.proc as unknown as ReturnType<typeof spawn>);
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+        const ctx = makeCtx();
+        const promise = pi.tools[0].execute("t", { query: "q", source: "react", top: 3 }, undefined, undefined, ctx);
+        fake.emitData("stdout", "out");
+        fake.emitClose(0);
+        await promise;
+        const callArgs = vi.mocked(spawn).mock.calls[0][1] as string[];
+        expect(callArgs).toContain("--source");
+        expect(callArgs).toContain("react");
+        expect(callArgs).toContain("--top");
+        expect(callArgs).toContain("3");
+    });
+
+    it("execute notifies warning and returns soft result on no-results", async () => {
+        const { execFileSync, spawn } = await import("node:child_process");
+        vi.mocked(execFileSync).mockReturnValue("help" as unknown as ReturnType<typeof execFileSync>);
+        const fake = makeFakeProc();
+        vi.mocked(spawn).mockReturnValue(fake.proc as unknown as ReturnType<typeof spawn>);
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+        const ctx = makeCtx();
+        const promise = pi.tools[0].execute("t", { query: "zzz" }, undefined, undefined, ctx);
+        fake.emitData("stdout", "No results found.\n");
+        fake.emitClose(0);
+        const result = (await promise) as { content: Array<{ text: string }> };
+        expect(result.content[0].text).toMatch(/No results found/);
+        expect(ctx.notifications.some((n) => n[1] === "warning")).toBe(true);
+    });
+
+    it("execute also treats empty stdout as no-results and suggests widening scope when source given", async () => {
+        const { execFileSync, spawn } = await import("node:child_process");
+        vi.mocked(execFileSync).mockReturnValue("help" as unknown as ReturnType<typeof execFileSync>);
+        const fake = makeFakeProc();
+        vi.mocked(spawn).mockReturnValue(fake.proc as unknown as ReturnType<typeof spawn>);
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+        const ctx = makeCtx();
+        const promise = pi.tools[0].execute("t", { query: "zzz", source: "react" }, undefined, undefined, ctx);
+        fake.emitData("stdout", "");
+        fake.emitClose(0);
+        const result = (await promise) as { content: Array<{ text: string }> };
+        expect(result.content[0].text).toMatch(/widening scope/);
+    });
+
+    it("execute rejects on non-zero exit code", async () => {
+        const { execFileSync, spawn } = await import("node:child_process");
+        vi.mocked(execFileSync).mockReturnValue("help" as unknown as ReturnType<typeof execFileSync>);
+        const fake = makeFakeProc();
+        vi.mocked(spawn).mockReturnValue(fake.proc as unknown as ReturnType<typeof spawn>);
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+        const ctx = makeCtx();
+        const promise = pi.tools[0].execute("t", { query: "q" }, undefined, undefined, ctx);
+        fake.emitData("stderr", "bad source");
+        fake.emitClose(1);
+        await expect(promise).rejects.toThrow(/exited with code 1/);
+    });
+
+    it("execute rejects on ENOENT at spawn time with install guidance + notify", async () => {
+        const { execFileSync, spawn } = await import("node:child_process");
+        vi.mocked(execFileSync).mockReturnValue("help" as unknown as ReturnType<typeof execFileSync>);
+        const fake = makeFakeProc();
+        vi.mocked(spawn).mockReturnValue(fake.proc as unknown as ReturnType<typeof spawn>);
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+        const ctx = makeCtx();
+        const promise = pi.tools[0].execute("t", { query: "q" }, undefined, undefined, ctx);
+        const err = Object.assign(new Error("ENOENT"), { code: "ENOENT" });
+        fake.emitError(err);
+        await expect(promise).rejects.toThrow(/not installed/i);
+        expect(ctx.notifications.some((n) => n[1] === "error")).toBe(true);
+    });
+
+    it("execute rejects on non-ENOENT spawn error", async () => {
+        const { execFileSync, spawn } = await import("node:child_process");
+        vi.mocked(execFileSync).mockReturnValue("help" as unknown as ReturnType<typeof execFileSync>);
+        const fake = makeFakeProc();
+        vi.mocked(spawn).mockReturnValue(fake.proc as unknown as ReturnType<typeof spawn>);
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+        const ctx = makeCtx();
+        const promise = pi.tools[0].execute("t", { query: "q" }, undefined, undefined, ctx);
+        fake.emitError(new Error("mystery"));
+        await expect(promise).rejects.toThrow(/failed to run/i);
+    });
+
+    it("execute kills process on abort signal", async () => {
+        const { execFileSync, spawn } = await import("node:child_process");
+        vi.mocked(execFileSync).mockReturnValue("help" as unknown as ReturnType<typeof execFileSync>);
+        const fake = makeFakeProc();
+        vi.mocked(spawn).mockReturnValue(fake.proc as unknown as ReturnType<typeof spawn>);
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+        const ctx = makeCtx();
+        const ac = new AbortController();
+        const promise = pi.tools[0].execute("t", { query: "q" }, ac.signal, undefined, ctx);
+        ac.abort();
+        expect(fake.proc.kill).toHaveBeenCalledWith("SIGTERM");
+        fake.emitClose(0);
+        await expect(promise).resolves.toBeDefined();
+    });
+});