@astrofoundry/pi-astro 0.5.1 → 0.6.1

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,400 @@
+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) };
+}
+
+type InputResponse = string | undefined;
+
+function makeCtx(opts: { hasUI?: boolean; cwd?: string; inputs?: InputResponse[] } = {}): {
+    cwd: string;
+    hasUI: boolean;
+    ui: {
+        notify: ReturnType<typeof vi.fn>;
+        confirm: ReturnType<typeof vi.fn>;
+        input: ReturnType<typeof vi.fn>;
+    };
+} {
+    const responses = [...(opts.inputs ?? [])];
+    return {
+        cwd: opts.cwd ?? "/tmp/cwd",
+        hasUI: opts.hasUI ?? true,
+        ui: {
+            notify: vi.fn(),
+            confirm: vi.fn().mockResolvedValue(true),
+            input: vi.fn().mockImplementation(() => Promise.resolve(responses.shift())),
+        },
+    };
+}
+
+describe("gemini-image extension", () => {
+    let fakeHome: string;
+    let fakeCwd: string;
+    const originalHome = process.env.HOME;
+    const originalKey = process.env.GEMINI_API_KEY;
+
+    beforeEach(() => {
+        fakeHome = mkdtempSync(join(tmpdir(), "gem-ext-"));
+        fakeCwd = mkdtempSync(join(tmpdir(), "gem-cwd-"));
+        process.env.HOME = fakeHome;
+        process.env.GEMINI_API_KEY = "test-key";
+        vi.resetModules();
+        generateContentMock.mockReset();
+        generateImagesMock.mockReset();
+    });
+
+    afterEach(() => {
+        rmSync(fakeHome, { recursive: true, force: true });
+        rmSync(fakeCwd, { 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] };
+    }
+
+    function pngImageResponse(): void {
+        generateContentMock.mockResolvedValue({
+            candidates: [
+                { content: { parts: [{ inlineData: { mimeType: "image/png", data: "BASE64DATA" } }] } },
+            ],
+        });
+    }
+
+    it("registers tool named gemini_image", async () => {
+        const { tool } = await load();
+        expect(tool.name).toBe("gemini_image");
+    });
+
+    it("skip_confirm bypasses review and calls generateContent", async () => {
+        pngImageResponse();
+        const { tool } = await load();
+        const ctx = makeCtx({ cwd: fakeCwd });
+        const res = await tool.execute("t", { prompt: "a cat", skip_confirm: true }, undefined, undefined, ctx);
+        expect(res.details.model).toBe("gemini-2.5-flash-image");
+        // ui.input should NOT be called — skip_confirm bypasses the review loop
+        expect(ctx.ui.input).not.toHaveBeenCalled();
+    });
+
+    it("images + text returned in content; image AFTER text", async () => {
+        pngImageResponse();
+        const { tool } = await load();
+        const res = await tool.execute(
+            "t",
+            { prompt: "x", skip_confirm: true },
+            undefined,
+            undefined,
+            makeCtx({ cwd: fakeCwd }),
+        );
+        expect((res.content as Array<{ type: string }>)[0].type).toBe("text");
+        expect((res.content as Array<{ type: string }>)[1].type).toBe("image");
+    });
+
+    it("always saves to <cwd>/.gemini-images by default", async () => {
+        pngImageResponse();
+        const { tool } = await load();
+        const res = await tool.execute(
+            "t",
+            { prompt: "a-thing", skip_confirm: true },
+            undefined,
+            undefined,
+            makeCtx({ cwd: fakeCwd }),
+        );
+        const paths = res.details.savedPaths as string[];
+        expect(paths).toHaveLength(1);
+        expect(paths[0]).toContain(".gemini-images");
+        expect(readdirSync(join(fakeCwd, ".gemini-images"))).toHaveLength(1);
+    });
+
+    it("save_to override writes into that directory (resolved against cwd)", async () => {
+        pngImageResponse();
+        const { tool } = await load();
+        const res = await tool.execute(
+            "t",
+            { prompt: "x", skip_confirm: true, save_to: "out" },
+            undefined,
+            undefined,
+            makeCtx({ cwd: fakeCwd }),
+        );
+        const paths = res.details.savedPaths as string[];
+        expect(paths[0]).toContain(join(fakeCwd, "out"));
+    });
+
+    it("silently saves API key on first paste (no confirm prompt)", async () => {
+        delete process.env.GEMINI_API_KEY;
+        pngImageResponse();
+        const { tool } = await load();
+        const ctx = makeCtx({ cwd: fakeCwd, inputs: ["pasted-key"] });
+        await tool.execute("t", { prompt: "x", skip_confirm: true }, undefined, undefined, ctx);
+        // Only ONE input call (the key paste). confirm never used.
+        expect(ctx.ui.input).toHaveBeenCalledTimes(1);
+        expect(ctx.ui.confirm).not.toHaveBeenCalled();
+        expect(ctx.ui.notify).toHaveBeenCalledWith(expect.stringMatching(/saved/i), "info");
+    });
+
+    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, cwd: fakeCwd });
+        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({ cwd: fakeCwd, inputs: ["   "] });
+        await expect(tool.execute("t", { prompt: "x" }, undefined, undefined, ctx)).rejects.toThrow(/No API key/);
+    });
+
+    it("review loop: empty input confirms and proceeds", async () => {
+        pngImageResponse();
+        const { tool } = await load();
+        const ctx = makeCtx({ cwd: fakeCwd, inputs: [""] });
+        await tool.execute("t", { prompt: "hello" }, undefined, undefined, ctx);
+        expect(ctx.ui.input).toHaveBeenCalledTimes(1);
+    });
+
+    it("review loop: 'cancel' throws cancelled", async () => {
+        const { tool } = await load();
+        const ctx = makeCtx({ cwd: fakeCwd, inputs: ["cancel"] });
+        await expect(tool.execute("t", { prompt: "x" }, undefined, undefined, ctx)).rejects.toThrow(/cancelled/i);
+    });
+
+    it("review loop: undefined input (escape) throws cancelled", async () => {
+        const { tool } = await load();
+        const ctx = makeCtx({ cwd: fakeCwd, inputs: [undefined] });
+        await expect(tool.execute("t", { prompt: "x" }, undefined, undefined, ctx)).rejects.toThrow(/cancelled/i);
+    });
+
+    it("review loop: tweak '2K' updates size then Enter confirms", async () => {
+        pngImageResponse();
+        const { tool } = await load();
+        const ctx = makeCtx({ cwd: fakeCwd, inputs: ["2K", ""] });
+        await tool.execute("t", { prompt: "x" }, undefined, undefined, ctx);
+        // Two input calls: one for 2K tweak, one empty confirm
+        expect(ctx.ui.input).toHaveBeenCalledTimes(2);
+        // The second input's preview message should mention 2K
+        const secondPreview = ctx.ui.input.mock.calls[1][1] as string;
+        expect(secondPreview).toContain("2K");
+    });
+
+    it("review loop: tweak 'use imagen ultra, 16:9, 3 images' + empty confirms and calls Imagen", async () => {
+        generateImagesMock.mockResolvedValue({
+            generatedImages: [
+                { image: { imageBytes: "IMG1", mimeType: "image/png" } },
+                { image: { imageBytes: "IMG2", mimeType: "image/png" } },
+                { image: { imageBytes: "IMG3", mimeType: "image/png" } },
+            ],
+        });
+        const { tool } = await load();
+        const ctx = makeCtx({ cwd: fakeCwd, inputs: ["use imagen ultra 16:9 3 images", ""] });
+        const res = await tool.execute("t", { prompt: "landscape" }, undefined, undefined, ctx);
+        expect(res.details.model).toBe("imagen-4.0-ultra-generate-001");
+        expect(res.details.images).toBe(3);
+        const cfg = generateImagesMock.mock.calls[0][0].config;
+        expect(cfg.numberOfImages).toBe(3);
+        expect(cfg.aspectRatio).toBe("16:9");
+    });
+
+    it("review loop: tweak that fails validation is rejected with notify warning, user can re-tweak", async () => {
+        pngImageResponse();
+        generateImagesMock.mockResolvedValue({
+            generatedImages: [{ image: { imageBytes: "IMG", mimeType: "image/png" } }],
+        });
+        const { tool } = await load();
+        // user tries 4K with Imagen → rejected; then switches back and confirms
+        const ctx = makeCtx({ cwd: fakeCwd, inputs: ["imagen fast 4K", "imagen fast 2K", ""] });
+        await tool.execute("t", { prompt: "x" }, undefined, undefined, ctx);
+        expect(ctx.ui.notify).toHaveBeenCalledWith(expect.stringMatching(/Change rejected/), "warning");
+    });
+
+    it("review loop: unrecognized token surfaces as warning", async () => {
+        pngImageResponse();
+        const { tool } = await load();
+        const ctx = makeCtx({ cwd: fakeCwd, inputs: ["frobnicate plz", ""] });
+        await tool.execute("t", { prompt: "x" }, undefined, undefined, ctx);
+        expect(ctx.ui.notify).toHaveBeenCalledWith(expect.stringMatching(/Didn't understand/), "warning");
+    });
+
+    it("rejects unknown model at call time", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute(
+                "t",
+                { prompt: "x", model: "bogus-model", skip_confirm: true },
+                undefined,
+                undefined,
+                makeCtx({ cwd: fakeCwd }),
+            ),
+        ).rejects.toThrow(/Unknown model/);
+    });
+
+    it("rejects invalid aspect_ratio for Imagen at validate time", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute(
+                "t",
+                {
+                    prompt: "x",
+                    model: "imagen-4.0-generate-001",
+                    aspect_ratio: "21:9",
+                    skip_confirm: true,
+                },
+                undefined,
+                undefined,
+                makeCtx({ cwd: fakeCwd }),
+            ),
+        ).rejects.toThrow(/aspect_ratio .* Imagen/);
+    });
+
+    it("rejects 4K for Imagen", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute(
+                "t",
+                {
+                    prompt: "x",
+                    model: "imagen-4.0-generate-001",
+                    image_size: "4K",
+                    skip_confirm: true,
+                },
+                undefined,
+                undefined,
+                makeCtx({ cwd: fakeCwd }),
+            ),
+        ).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"],
+                    skip_confirm: true,
+                },
+                undefined,
+                undefined,
+                makeCtx({ cwd: fakeCwd }),
+            ),
+        ).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", skip_confirm: true },
+                undefined,
+                undefined,
+                makeCtx({ cwd: fakeCwd }),
+            ),
+        ).rejects.toThrow(/negative_prompt/);
+    });
+
+    it("rejects seed on Gemini native", async () => {
+        const { tool } = await load();
+        await expect(
+            tool.execute("t", { prompt: "x", seed: 42, skip_confirm: true }, undefined, undefined, makeCtx({ cwd: fakeCwd })),
+        ).rejects.toThrow(/seed/);
+    });
+
+    it("throws when model returns no images", async () => {
+        generateContentMock.mockResolvedValue({ candidates: [] });
+        const { tool } = await load();
+        await expect(
+            tool.execute("t", { prompt: "x", skip_confirm: true }, undefined, undefined, makeCtx({ cwd: fakeCwd })),
+        ).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");
+        pngImageResponse();
+        const { tool } = await load();
+        const ctx = makeCtx({ cwd: fakeCwd });
+        await tool.execute(
+            "t",
+            { prompt: "edit", input_images: [imgPath], skip_confirm: true },
+            undefined,
+            undefined,
+            ctx,
+        );
+        const callArgs = generateContentMock.mock.calls[0][0];
+        expect(callArgs.contents.some((c: Record<string, unknown>) => "inlineData" in c)).toBe(true);
+        rmSync(sourceDir, { recursive: true, force: true });
+    });
+
+    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",
+                skip_confirm: true,
+            },
+            undefined,
+            undefined,
+            makeCtx({ cwd: fakeCwd }),
+        );
+        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");
+    });
+});