import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { registerOAuthProvider } from "@mariozechner/pi-ai"; import lockfile from "proper-lockfile"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.js"; import { clearConfigValueCache } from "../src/core/resolve-config-value.js"; describe("AuthStorage", () => { let tempDir: string; let authJsonPath: string; let authStorage: AuthStorage; beforeEach(() => { tempDir = join(tmpdir(), `pi-test-auth-storage-${Date.now()}-${Math.random().toString(36).slice(2)}`); mkdirSync(tempDir, { recursive: true }); authJsonPath = join(tempDir, "auth.json"); }); afterEach(() => { if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); } clearConfigValueCache(); vi.restoreAllMocks(); }); function writeAuthJson(data: Record) { writeFileSync(authJsonPath, JSON.stringify(data)); } describe("API key resolution", () => { test("literal API key is returned directly", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "sk-ant-literal-key" }, }); authStorage = new AuthStorage(authJsonPath); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBe("sk-ant-literal-key"); }); test("apiKey with ! prefix executes command and uses stdout", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "!echo test-api-key-from-command" }, }); authStorage = new AuthStorage(authJsonPath); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBe("test-api-key-from-command"); }); test("apiKey with ! prefix trims whitespace from command output", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "!echo ' spaced-key '" }, }); authStorage = new AuthStorage(authJsonPath); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBe("spaced-key"); }); test("apiKey with ! prefix handles multiline output (uses trimmed result)", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "!printf 'line1\\nline2'" }, }); authStorage = new AuthStorage(authJsonPath); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBe("line1\nline2"); }); test("apiKey with ! prefix returns undefined on command failure", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "!exit 1" }, }); authStorage = new AuthStorage(authJsonPath); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBeUndefined(); }); test("apiKey with ! prefix returns undefined on nonexistent command", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "!nonexistent-command-12345" }, }); authStorage = new AuthStorage(authJsonPath); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBeUndefined(); }); test("apiKey with ! prefix returns undefined on empty output", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "!printf ''" }, }); authStorage = new AuthStorage(authJsonPath); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBeUndefined(); }); test("apiKey as environment variable name resolves to env value", async () => { const originalEnv = process.env.TEST_AUTH_API_KEY_12345; process.env.TEST_AUTH_API_KEY_12345 = "env-api-key-value"; try { writeAuthJson({ anthropic: { type: "api_key", key: "TEST_AUTH_API_KEY_12345" }, }); authStorage = new AuthStorage(authJsonPath); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBe("env-api-key-value"); } finally { if (originalEnv === undefined) { delete process.env.TEST_AUTH_API_KEY_12345; } else { process.env.TEST_AUTH_API_KEY_12345 = originalEnv; } } }); test("apiKey as literal value is used directly when not an env var", async () => { // Make sure this isn't an env var delete process.env.literal_api_key_value; writeAuthJson({ anthropic: { type: "api_key", key: "literal_api_key_value" }, }); authStorage = new AuthStorage(authJsonPath); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBe("literal_api_key_value"); }); test("apiKey command can use shell features like pipes", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "!echo 'hello world' | tr ' ' '-'" }, }); authStorage = new AuthStorage(authJsonPath); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBe("hello-world"); }); describe("caching", () => { test("command is only executed once per process", async () => { // Use a command that writes to a file to count invocations const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0"); const command = `!sh -c 'count=$(cat ${counterFile}); echo $((count + 1)) > ${counterFile}; echo "key-value"'`; writeAuthJson({ anthropic: { type: "api_key", key: command }, }); authStorage = new AuthStorage(authJsonPath); // Call multiple times await authStorage.getApiKey("anthropic"); await authStorage.getApiKey("anthropic"); await authStorage.getApiKey("anthropic"); // Command should have only run once const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); expect(count).toBe(1); }); test("cache persists across AuthStorage instances", async () => { const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0"); const command = `!sh -c 'count=$(cat ${counterFile}); echo $((count + 1)) > ${counterFile}; echo "key-value"'`; writeAuthJson({ anthropic: { type: "api_key", key: command }, }); // Create multiple AuthStorage instances const storage1 = new AuthStorage(authJsonPath); await storage1.getApiKey("anthropic"); const storage2 = new AuthStorage(authJsonPath); await storage2.getApiKey("anthropic"); // Command should still have only run once const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); expect(count).toBe(1); }); test("clearConfigValueCache allows command to run again", async () => { const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0"); const command = `!sh -c 'count=$(cat ${counterFile}); echo $((count + 1)) > ${counterFile}; echo "key-value"'`; writeAuthJson({ anthropic: { type: "api_key", key: command }, }); authStorage = new AuthStorage(authJsonPath); await authStorage.getApiKey("anthropic"); // Clear cache and call again clearConfigValueCache(); await authStorage.getApiKey("anthropic"); // Command should have run twice const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); expect(count).toBe(2); }); test("different commands are cached separately", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "!echo key-anthropic" }, openai: { type: "api_key", key: "!echo key-openai" }, }); authStorage = new AuthStorage(authJsonPath); const keyA = await authStorage.getApiKey("anthropic"); const keyB = await authStorage.getApiKey("openai"); expect(keyA).toBe("key-anthropic"); expect(keyB).toBe("key-openai"); }); test("failed commands are cached (not retried)", async () => { const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0"); const command = `!sh -c 'count=$(cat ${counterFile}); echo $((count + 1)) > ${counterFile}; exit 1'`; writeAuthJson({ anthropic: { type: "api_key", key: command }, }); authStorage = new AuthStorage(authJsonPath); // Call multiple times - all should return undefined const key1 = await authStorage.getApiKey("anthropic"); const key2 = await authStorage.getApiKey("anthropic"); expect(key1).toBeUndefined(); expect(key2).toBeUndefined(); // Command should have only run once despite failures const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); expect(count).toBe(1); }); test("environment variables are not cached (changes are picked up)", async () => { const envVarName = "TEST_AUTH_KEY_CACHE_TEST_98765"; const originalEnv = process.env[envVarName]; try { process.env[envVarName] = "first-value"; writeAuthJson({ anthropic: { type: "api_key", key: envVarName }, }); authStorage = new AuthStorage(authJsonPath); const key1 = await authStorage.getApiKey("anthropic"); expect(key1).toBe("first-value"); // Change env var process.env[envVarName] = "second-value"; const key2 = await authStorage.getApiKey("anthropic"); expect(key2).toBe("second-value"); } finally { if (originalEnv === undefined) { delete process.env[envVarName]; } else { process.env[envVarName] = originalEnv; } } }); }); }); describe("oauth lock compromise handling", () => { test("returns undefined on compromised lock and allows a later retry", async () => { const providerId = `test-oauth-provider-${Date.now()}-${Math.random().toString(36).slice(2)}`; registerOAuthProvider({ id: providerId, name: "Test OAuth Provider", async login() { throw new Error("Not used in this test"); }, async refreshToken(credentials) { return { ...credentials, access: "refreshed-access-token", expires: Date.now() + 60_000, }; }, getApiKey(credentials) { return `Bearer ${credentials.access}`; }, }); writeAuthJson({ [providerId]: { type: "oauth", refresh: "refresh-token", access: "expired-access-token", expires: Date.now() - 10_000, }, }); authStorage = new AuthStorage(authJsonPath); const realLock = lockfile.lock.bind(lockfile); const lockSpy = vi.spyOn(lockfile, "lock"); lockSpy.mockImplementationOnce(async (file, options) => { options?.onCompromised?.(new Error("Unable to update lock within the stale threshold")); return realLock(file, options); }); const firstTry = await authStorage.getApiKey(providerId); expect(firstTry).toBeUndefined(); lockSpy.mockRestore(); const secondTry = await authStorage.getApiKey(providerId); expect(secondTry).toBe("Bearer refreshed-access-token"); }); }); describe("runtime overrides", () => { test("runtime override takes priority over auth.json", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "!echo stored-key" }, }); authStorage = new AuthStorage(authJsonPath); authStorage.setRuntimeApiKey("anthropic", "runtime-key"); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBe("runtime-key"); }); test("removing runtime override falls back to auth.json", async () => { writeAuthJson({ anthropic: { type: "api_key", key: "!echo stored-key" }, }); authStorage = new AuthStorage(authJsonPath); authStorage.setRuntimeApiKey("anthropic", "runtime-key"); authStorage.removeRuntimeApiKey("anthropic"); const apiKey = await authStorage.getApiKey("anthropic"); expect(apiKey).toBe("stored-key"); }); }); });