@astrofoundry/pi-astro 0.5.1 → 0.6.0

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

@@ -0,0 +1,130 @@
+import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+describe("credentials", () => {
+    let fakeHome: string;
+    const originalHome = process.env.HOME;
+    const originalGemini = process.env.GEMINI_API_KEY;
+
+    beforeEach(() => {
+        fakeHome = mkdtempSync(join(tmpdir(), "gem-cred-"));
+        process.env.HOME = fakeHome;
+        delete process.env.GEMINI_API_KEY;
+        vi.resetModules();
+    });
+
+    afterEach(() => {
+        rmSync(fakeHome, { recursive: true, force: true });
+        process.env.HOME = originalHome;
+        if (originalGemini === undefined) delete process.env.GEMINI_API_KEY;
+        else process.env.GEMINI_API_KEY = originalGemini;
+    });
+
+    it("resolves env var first when set", async () => {
+        process.env.GEMINI_API_KEY = "from-env";
+        const { resolveExistingApiKey } = await import("./credentials.ts");
+        expect(resolveExistingApiKey()).toBe("from-env");
+    });
+
+    it("returns undefined when neither env nor file set", async () => {
+        const { resolveExistingApiKey } = await import("./credentials.ts");
+        expect(resolveExistingApiKey()).toBeUndefined();
+    });
+
+    it("reads auth.json google.key when env is absent", async () => {
+        const dir = join(fakeHome, ".pi", "agent");
+        mkdirSync(dir, { recursive: true });
+        writeFileSync(
+            join(dir, "auth.json"),
+            JSON.stringify({ google: { type: "api_key", key: "from-file" } }),
+            "utf-8",
+        );
+        const { resolveExistingApiKey } = await import("./credentials.ts");
+        expect(resolveExistingApiKey()).toBe("from-file");
+    });
+
+    it("returns undefined when auth.json google entry is malformed", async () => {
+        const dir = join(fakeHome, ".pi", "agent");
+        mkdirSync(dir, { recursive: true });
+        writeFileSync(join(dir, "auth.json"), JSON.stringify({ google: "not-an-object" }), "utf-8");
+        const { resolveExistingApiKey } = await import("./credentials.ts");
+        expect(resolveExistingApiKey()).toBeUndefined();
+    });
+
+    it("returns undefined when auth.json is corrupt JSON", async () => {
+        const dir = join(fakeHome, ".pi", "agent");
+        mkdirSync(dir, { recursive: true });
+        writeFileSync(join(dir, "auth.json"), "{not-json", "utf-8");
+        const { resolveExistingApiKey } = await import("./credentials.ts");
+        expect(resolveExistingApiKey()).toBeUndefined();
+    });
+
+    it("persistApiKey creates auth.json with google entry", async () => {
+        const { persistApiKey } = await import("./credentials.ts");
+        persistApiKey("new-key");
+        const data = JSON.parse(readFileSync(join(fakeHome, ".pi", "agent", "auth.json"), "utf-8"));
+        expect(data.google).toEqual({ type: "api_key", key: "new-key" });
+    });
+
+    it("persistApiKey merges with existing entries", async () => {
+        const dir = join(fakeHome, ".pi", "agent");
+        mkdirSync(dir, { recursive: true });
+        writeFileSync(
+            join(dir, "auth.json"),
+            JSON.stringify({ anthropic: { type: "api_key", key: "ant-k" } }),
+            "utf-8",
+        );
+        const { persistApiKey } = await import("./credentials.ts");
+        persistApiKey("new-google-key");
+        const data = JSON.parse(readFileSync(join(dir, "auth.json"), "utf-8"));
+        expect(data.anthropic).toEqual({ type: "api_key", key: "ant-k" });
+        expect(data.google).toEqual({ type: "api_key", key: "new-google-key" });
+    });
+
+    it("persistApiKey overwrites an existing google entry", async () => {
+        const dir = join(fakeHome, ".pi", "agent");
+        mkdirSync(dir, { recursive: true });
+        writeFileSync(
+            join(dir, "auth.json"),
+            JSON.stringify({ google: { type: "api_key", key: "old" } }),
+            "utf-8",
+        );
+        const { persistApiKey } = await import("./credentials.ts");
+        persistApiKey("replacement");
+        const data = JSON.parse(readFileSync(join(dir, "auth.json"), "utf-8"));
+        expect(data.google.key).toBe("replacement");
+    });
+
+    it("env var with whitespace is trimmed", async () => {
+        process.env.GEMINI_API_KEY = "   spaced   ";
+        const { resolveExistingApiKey } = await import("./credentials.ts");
+        expect(resolveExistingApiKey()).toBe("spaced");
+    });
+
+    it("empty string env var is treated as unset", async () => {
+        process.env.GEMINI_API_KEY = "   ";
+        const { resolveExistingApiKey } = await import("./credentials.ts");
+        expect(resolveExistingApiKey()).toBeUndefined();
+    });
+
+    it("auth file exists but lacks google entry -> undefined", async () => {
+        const dir = join(fakeHome, ".pi", "agent");
+        mkdirSync(dir, { recursive: true });
+        writeFileSync(
+            join(dir, "auth.json"),
+            JSON.stringify({ anthropic: { type: "api_key", key: "x" } }),
+            "utf-8",
+        );
+        const { resolveExistingApiKey } = await import("./credentials.ts");
+        expect(resolveExistingApiKey()).toBeUndefined();
+    });
+
+    it("persistApiKey creates parent dirs implicitly if needed", async () => {
+        expect(existsSync(join(fakeHome, ".pi", "agent"))).toBe(false);
+        const { persistApiKey } = await import("./credentials.ts");
+        persistApiKey("x");
+        expect(existsSync(join(fakeHome, ".pi", "agent", "auth.json"))).toBe(true);
+    });
+});

package/extensions/gemini-image/credentials.ts

@@ -0,0 +1,53 @@
+import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
+import { homedir } from "node:os";
+import { dirname, join } from "node:path";
+
+const AUTH_FILE = join(homedir(), ".pi", "agent", "auth.json");
+const PROVIDER_KEY = "google";
+
+interface AuthEntry {
+    type: "api_key";
+    key: string;
+}
+
+type AuthFile = Record<string, AuthEntry | unknown>;
+
+function readAuthFile(): AuthFile {
+    if (!existsSync(AUTH_FILE)) return {};
+    try {
+        const raw = readFileSync(AUTH_FILE, "utf-8");
+        const parsed = JSON.parse(raw) as unknown;
+        return typeof parsed === "object" && parsed !== null ? (parsed as AuthFile) : {};
+    } catch {
+        return {};
+    }
+}
+
+function writeAuthFile(data: AuthFile): void {
+    mkdirSync(dirname(AUTH_FILE), { recursive: true });
+    writeFileSync(AUTH_FILE, JSON.stringify(data, null, 2) + "\n", "utf-8");
+    try {
+        chmodSync(AUTH_FILE, 0o600);
+    } catch {
+        // non-fatal: filesystem may not support chmod (e.g. some Windows envs)
+    }
+}
+
+export function resolveExistingApiKey(): string | undefined {
+    const fromEnv = process.env.GEMINI_API_KEY?.trim();
+    if (fromEnv) return fromEnv;
+
+    const authFile = readAuthFile();
+    const entry = authFile[PROVIDER_KEY];
+    if (entry && typeof entry === "object" && "type" in entry && entry.type === "api_key") {
+        const { key } = entry as AuthEntry;
+        if (typeof key === "string" && key.trim().length > 0) return key.trim();
+    }
+    return undefined;
+}
+
+export function persistApiKey(key: string): void {
+    const authFile = readAuthFile();
+    authFile[PROVIDER_KEY] = { type: "api_key", key };
+    writeAuthFile(authFile);
+}

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

@@ -0,0 +1,369 @@
+import { mkdtempSync, readdirSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+const generateContentMock = vi.fn();
+const generateImagesMock = vi.fn();
+
+vi.mock("@google/genai", () => {
+    class FakeGoogleGenAI {
+        models = {
+            generateContent: generateContentMock,
+            generateImages: generateImagesMock,
+        };
+    }
+    return { GoogleGenAI: FakeGoogleGenAI };
+});
+
+type ToolDef = {
+    name: string;
+    execute: (
+        id: string,
+        input: Record<string, unknown>,
+        signal: AbortSignal | undefined,
+        onUpdate: unknown,
+        ctx: {
+            cwd: string;
+            hasUI: boolean;
+            ui: {
+                notify: ReturnType<typeof vi.fn>;
+                confirm: ReturnType<typeof vi.fn>;
+                input: ReturnType<typeof vi.fn>;
+            };
+        },
+    ) => Promise<{ content: unknown[]; details: Record<string, unknown> }>;
+};
+
+function makePi(): { tools: ToolDef[]; registerTool: (t: ToolDef) => void } {
+    const tools: ToolDef[] = [];
+    return { tools, registerTool: (t) => tools.push(t) };
+}
+
+function makeCtx(override: Partial<{ hasUI: boolean; confirm: boolean; inputValue: string | undefined }> = {}): {
+    cwd: string;
+    hasUI: boolean;
+    ui: {
+        notify: ReturnType<typeof vi.fn>;
+        confirm: ReturnType<typeof vi.fn>;
+        input: ReturnType<typeof vi.fn>;
+    };
+} {
+    const hasUI = override.hasUI ?? true;
+    return {
+        cwd: "/tmp/cwd",
+        hasUI,
+        ui: {
+            notify: vi.fn(),
+            confirm: vi.fn().mockResolvedValue(override.confirm ?? true),
+            input: vi.fn().mockResolvedValue(override.inputValue ?? "abc-key-from-paste"),
+        },
+    };
+}
+
+describe("gemini-image extension", () => {
+    let fakeHome: string;
+    const originalHome = process.env.HOME;
+    const originalKey = process.env.GEMINI_API_KEY;
+
+    beforeEach(() => {
+        fakeHome = mkdtempSync(join(tmpdir(), "gem-ext-"));
+        process.env.HOME = fakeHome;
+        process.env.GEMINI_API_KEY = "test-key";
+        vi.resetModules();
+        generateContentMock.mockReset();
+        generateImagesMock.mockReset();
+    });
+
+    afterEach(() => {
+        rmSync(fakeHome, { recursive: true, force: true });
+        process.env.HOME = originalHome;
+        if (originalKey === undefined) delete process.env.GEMINI_API_KEY;
+        else process.env.GEMINI_API_KEY = originalKey;
+    });
+
+    async function load(): Promise<{ tool: ToolDef }> {
+        const mod = await import("./index.ts");
+        const pi = makePi();
+        mod.default(pi as unknown as Parameters<typeof mod.default>[0]);
+        return { tool: pi.tools[0] };
+    }
+
+    it("registers a tool named gemini_image", async () => {
+        const { tool } = await load();
+        expect(tool.name).toBe("gemini_image");
+    });
+
+    it("uses default model gemini-2.5-flash-image and calls generateContent", async () => {
+        generateContentMock.mockResolvedValue({
+            candidates: [
+                { content: { parts: [{ inlineData: { mimeType: "image/png", data: "BASE64DATA" } }] } },
+            ],
+        });
+        const { tool } = await load();
+        const ctx = makeCtx();
+        const res = await tool.execute("t", { prompt: "a cat" }, undefined, undefined, ctx);
+        expect(res.details.model).toBe("gemini-2.5-flash-image");
+        expect(res.content[0]).toMatchObject({ type: "image", data: "BASE64DATA" });
+    });
+
+    it("calls Imagen for imagen-4.0-generate-001 with numberOfImages", async () => {
+        generateImagesMock.mockResolvedValue({
+            generatedImages: [
+                { image: { imageBytes: "IMG1", mimeType: "image/png" } },
+                { image: { imageBytes: "IMG2", mimeType: "image/png" } },
+            ],
+        });
+        const { tool } = await load();
+        const ctx = makeCtx();
+        const res = await tool.execute(
+            "t",
+            {
+                prompt: "landscape",
+                model: "imagen-4.0-generate-001",
+                number_of_images: 2,
+            },
+            undefined,
+            undefined,
+            ctx,
+        );
+        expect(res.details.images).toBe(2);
+        const args = generateImagesMock.mock.calls[0][0];
+        expect(args.config.numberOfImages).toBe(2);
+    });
+
+    it("rejects unknown model", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute("t", { prompt: "x", model: "bogus-model" }, undefined, undefined, makeCtx()),
+        ).rejects.toThrow(/Unknown model/);
+    });
+
+    it("rejects invalid aspect_ratio for Imagen (e.g. 21:9)", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute(
+                "t",
+                { prompt: "x", model: "imagen-4.0-generate-001", aspect_ratio: "21:9" },
+                undefined,
+                undefined,
+                makeCtx(),
+            ),
+        ).rejects.toThrow(/aspect_ratio .* Imagen/);
+    });
+
+    it("rejects 4K image_size for Imagen", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute(
+                "t",
+                { prompt: "x", model: "imagen-4.0-generate-001", image_size: "4K" },
+                undefined,
+                undefined,
+                makeCtx(),
+            ),
+        ).rejects.toThrow(/image_size .* Imagen/);
+    });
+
+    it("rejects input_images on Imagen", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute(
+                "t",
+                {
+                    prompt: "x",
+                    model: "imagen-4.0-generate-001",
+                    input_images: ["/tmp/x.png"],
+                },
+                undefined,
+                undefined,
+                makeCtx(),
+            ),
+        ).rejects.toThrow(/input_images/);
+    });
+
+    it("rejects negative_prompt on Gemini native", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute(
+                "t",
+                { prompt: "x", negative_prompt: "ugly" },
+                undefined,
+                undefined,
+                makeCtx(),
+            ),
+        ).rejects.toThrow(/negative_prompt/);
+    });
+
+    it("rejects seed on Gemini native", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute("t", { prompt: "x", seed: 42 }, undefined, undefined, makeCtx()),
+        ).rejects.toThrow(/seed/);
+    });
+
+    it("rejects number_of_images > 1 on Gemini native", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute("t", { prompt: "x", number_of_images: 3 }, undefined, undefined, makeCtx()),
+        ).rejects.toThrow(/number_of_images/);
+    });
+
+    it("prompts for API key when none configured and saves if confirmed", async () => {
+        delete process.env.GEMINI_API_KEY;
+        generateContentMock.mockResolvedValue({
+            candidates: [{ content: { parts: [{ inlineData: { mimeType: "image/png", data: "X" } }] } }],
+        });
+        const { tool } = await load();
+        const ctx = makeCtx({ inputValue: "pasted-key" });
+        await tool.execute("t", { prompt: "x" }, undefined, undefined, ctx);
+        expect(ctx.ui.input).toHaveBeenCalled();
+        // Confirm was called for both the save prompt and the cost confirmation
+        expect(ctx.ui.confirm.mock.calls.length).toBeGreaterThanOrEqual(1);
+    });
+
+    it("errors when API key missing and no UI available", async () => {
+        delete process.env.GEMINI_API_KEY;
+        const { tool } = await load();
+        const ctx = makeCtx({ hasUI: false });
+        await expect(tool.execute("t", { prompt: "x" }, undefined, undefined, ctx)).rejects.toThrow(/no interactive UI/);
+    });
+
+    it("errors when user provides empty key at prompt", async () => {
+        delete process.env.GEMINI_API_KEY;
+        const { tool } = await load();
+        const ctx = makeCtx({ inputValue: "   " });
+        await expect(tool.execute("t", { prompt: "x" }, undefined, undefined, ctx)).rejects.toThrow(/No API key/);
+    });
+
+    it("cost confirmation: throws when user rejects", async () => {
+        generateContentMock.mockResolvedValue({
+            candidates: [{ content: { parts: [{ inlineData: { mimeType: "image/png", data: "X" } }] } }],
+        });
+        const { tool } = await load();
+        const ctx = makeCtx({ confirm: false });
+        await expect(tool.execute("t", { prompt: "x" }, undefined, undefined, ctx)).rejects.toThrow(/User cancelled/);
+    });
+
+    it("skip_confirm bypasses cost prompt", async () => {
+        generateContentMock.mockResolvedValue({
+            candidates: [{ content: { parts: [{ inlineData: { mimeType: "image/png", data: "X" } }] } }],
+        });
+        const { tool } = await load();
+        const ctx = makeCtx();
+        await tool.execute("t", { prompt: "x", skip_confirm: true }, undefined, undefined, ctx);
+        // confirm should not be called for cost
+        expect(ctx.ui.confirm).not.toHaveBeenCalled();
+    });
+
+    it("save_to writes PNG files and reports paths", async () => {
+        generateImagesMock.mockResolvedValue({
+            generatedImages: [
+                { image: { imageBytes: Buffer.from("fakeimage").toString("base64"), mimeType: "image/png" } },
+            ],
+        });
+        const saveDir = mkdtempSync(join(tmpdir(), "gem-save-"));
+        const { tool } = await load();
+        const ctx = makeCtx();
+        const res = await tool.execute(
+            "t",
+            {
+                prompt: "a-landscape",
+                model: "imagen-4.0-fast-generate-001",
+                save_to: saveDir,
+            },
+            undefined,
+            undefined,
+            ctx,
+        );
+        const paths = res.details.savedPaths as string[];
+        expect(paths).toHaveLength(1);
+        expect(readdirSync(saveDir)).toHaveLength(1);
+        rmSync(saveDir, { recursive: true, force: true });
+    });
+
+    it("throws when model returns no images", async () => {
+        generateContentMock.mockResolvedValue({ candidates: [] });
+        const { tool } = await load();
+        await expect(tool.execute("t", { prompt: "x" }, undefined, undefined, makeCtx())).rejects.toThrow(/no images/);
+    });
+
+    it("input_images are read from disk and passed as inlineData", async () => {
+        const sourceDir = mkdtempSync(join(tmpdir(), "gem-input-"));
+        const imgPath = join(sourceDir, "input.png");
+        writeFileSync(imgPath, "binary-image");
+        generateContentMock.mockResolvedValue({
+            candidates: [{ content: { parts: [{ inlineData: { mimeType: "image/png", data: "OUT" } }] } }],
+        });
+        const { tool } = await load();
+        const ctx = makeCtx();
+        await tool.execute(
+            "t",
+            { prompt: "edit", input_images: [imgPath] },
+            undefined,
+            undefined,
+            ctx,
+        );
+        const callArgs = generateContentMock.mock.calls[0][0];
+        // First content is the text, subsequent are inlineData parts
+        expect(callArgs.contents.some((c: Record<string, unknown>) => "inlineData" in c)).toBe(true);
+        rmSync(sourceDir, { recursive: true, force: true });
+    });
+
+    it("passes aspect_ratio + image_size into imageConfig for Gemini native", async () => {
+        generateContentMock.mockResolvedValue({
+            candidates: [{ content: { parts: [{ inlineData: { mimeType: "image/png", data: "X" } }] } }],
+        });
+        const { tool } = await load();
+        await tool.execute(
+            "t",
+            { prompt: "x", aspect_ratio: "16:9", image_size: "2K" },
+            undefined,
+            undefined,
+            makeCtx(),
+        );
+        const cfg = generateContentMock.mock.calls[0][0].config;
+        expect(cfg.imageConfig.aspectRatio).toBe("16:9");
+        expect(cfg.imageConfig.imageSize).toBe("2K");
+    });
+
+    it("Imagen config: numberOfImages + negative_prompt + seed + person_generation all pass through", async () => {
+        generateImagesMock.mockResolvedValue({
+            generatedImages: [{ image: { imageBytes: "X", mimeType: "image/png" } }],
+        });
+        const { tool } = await load();
+        await tool.execute(
+            "t",
+            {
+                prompt: "x",
+                model: "imagen-4.0-ultra-generate-001",
+                number_of_images: 2,
+                negative_prompt: "blurry",
+                seed: 1337,
+                person_generation: "allow_adult",
+            },
+            undefined,
+            undefined,
+            makeCtx(),
+        );
+        const cfg = generateImagesMock.mock.calls[0][0].config;
+        expect(cfg.numberOfImages).toBe(2);
+        expect(cfg.negativePrompt).toBe("blurry");
+        expect(cfg.seed).toBe(1337);
+        expect(cfg.personGeneration).toBe("allow_adult");
+    });
+
+    it("rejects invalid Gemini aspect_ratio", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute("t", { prompt: "x", aspect_ratio: "17:3" }, undefined, undefined, makeCtx()),
+        ).rejects.toThrow(/aspect_ratio .* Gemini native/);
+    });
+
+    it("rejects invalid Gemini image_size", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute("t", { prompt: "x", image_size: "8K" as "4K" }, undefined, undefined, makeCtx()),
+        ).rejects.toThrow(/image_size/);
+    });
+});