From def9e4e9a9675ea4c69761e5e0aeec063b5a0153 Mon Sep 17 00:00:00 2001 From: Carlos Villela Date: Sun, 18 Jan 2026 10:48:06 -0800 Subject: [PATCH] Support shell command execution for API key resolution in models.json (#762) * Support shell command execution for API key resolution in models.json Add ! prefix support to apiKey field in models.json to execute shell commands and use stdout as the API key. This allows users to store API keys in secure credential managers like macOS Keychain, 1Password, Bitwarden, or HashiCorp Vault. Example: "apiKey": "!security find-generic-password -ws 'anthropic'" The apiKey field now supports three formats: - !command - executes shell command, uses trimmed stdout - ENV_VAR_NAME - uses environment variable value - literal - uses value directly fixes #697 * feat(coding-agent): cache API key command results for process lifetime Shell commands (! prefix) are now executed once and cached. Environment variables and literal values are not cached, so changes are picked up. Addresses review feedback on #762. --------- Co-authored-by: Mario Zechner --- package-lock.json | 9 +- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/README.md | 5 +- .../coding-agent/src/core/model-registry.ts | 40 ++- .../coding-agent/test/model-registry.test.ts | 274 +++++++++++++++++- 5 files changed, 315 insertions(+), 14 deletions(-) diff --git a/package-lock.json b/package-lock.json index c957758d..289f9955 100644 --- a/package-lock.json +++ b/package-lock.json @@ -6409,7 +6409,6 @@ "resolved": "https://registry.npmjs.org/lit/-/lit-3.3.2.tgz", "integrity": "sha512-NF9zbsP79l4ao2SNrH3NkfmFgN/hBYSQo90saIVI1o5GpjAdCPVstVzO1MrLOakHoEhYkrtRjPK6Ob521aoYWQ==", "license": "BSD-3-Clause", - "peer": true, "dependencies": { "@lit/reactive-element": "^2.1.0", "lit-element": "^4.2.0", @@ -7764,7 +7763,6 @@ "resolved": "https://registry.npmjs.org/tailwind-merge/-/tailwind-merge-3.4.0.tgz", "integrity": "sha512-uSaO4gnW+b3Y2aWoWfFpX62vn2sR3skfhbjsEnaBI81WD1wBLlHZe5sWf0AqjksNdYTbGBEd0UasQMT3SNV15g==", "license": "MIT", - "peer": true, "funding": { "type": "github", "url": "https://github.com/sponsors/dcastil" @@ -7793,8 +7791,7 @@ "version": "4.1.18", "resolved": "https://registry.npmjs.org/tailwindcss/-/tailwindcss-4.1.18.tgz", "integrity": "sha512-4+Z+0yiYyEtUVCScyfHCxOYP06L5Ne+JiHhY2IjR2KWMIWhJOYZKLSGZaP5HkZ8+bY0cxfzwDE5uOmzFXyIwxw==", - "license": "MIT", - "peer": true + "license": "MIT" }, "node_modules/tapable": { "version": "2.3.0", @@ -7912,7 +7909,6 @@ "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, @@ -8009,7 +8005,6 @@ "integrity": "sha512-5C1sg4USs1lfG0GFb2RLXsdpXqBSEhAaA/0kPL01wxzpMqLILNxIxIOKiILz+cdg/pLnOUxFYOR5yhHU666wbw==", "devOptional": true, "license": "MIT", - "peer": true, "dependencies": { "esbuild": "~0.27.0", "get-tsconfig": "^4.7.5" @@ -8089,7 +8084,6 @@ "resolved": "https://registry.npmjs.org/vite/-/vite-7.3.1.tgz", "integrity": "sha512-w+N7Hifpc3gRjZ63vYBXA56dvvRlNWRczTdmCBBa+CotUzAPf5b7YMdMR/8CQoeYE5LX3W4wj6RYTgonm1b9DA==", "license": "MIT", - "peer": true, "dependencies": { "esbuild": "^0.27.0", "fdir": "^6.5.0", @@ -8204,7 +8198,6 @@ "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 1fbc1dda..3f42ad27 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -5,6 +5,7 @@ ### Changed - Share URLs now use hash fragments (`#`) instead of query strings (`?`) to prevent session IDs from being sent to buildwithpi.ai ([#828](https://github.com/badlogic/pi-mono/issues/828)) +- API keys in `models.json` can now be retrieved via shell command using `!` prefix (e.g., `"apiKey": "!security find-generic-password -ws 'anthropic'"` for macOS Keychain) ([#697](https://github.com/badlogic/pi-mono/issues/697)) ### Fixed diff --git a/packages/coding-agent/README.md b/packages/coding-agent/README.md index 7ad4e94c..653b6164 100644 --- a/packages/coding-agent/README.md +++ b/packages/coding-agent/README.md @@ -674,7 +674,10 @@ Add custom models (Ollama, vLLM, LM Studio, etc.) via `~/.pi/agent/models.json`: **Supported APIs:** `openai-completions`, `openai-responses`, `openai-codex-responses`, `anthropic-messages`, `google-generative-ai` -**API key resolution:** The `apiKey` field is checked as environment variable name first, then used as literal value. +**API key resolution:** The `apiKey` field supports three formats: +- `"!command"` - Executes the command and uses stdout (e.g., `"!security find-generic-password -ws 'anthropic'"` for macOS Keychain, `"!op read 'op://vault/item/credential'"` for 1Password) +- Environment variable name (e.g., `"MY_API_KEY"`) - Uses the value of the environment variable +- Literal value - Used directly as the API key **API override:** Set `api` at provider level (default for all models) or model level (override per model). diff --git a/packages/coding-agent/src/core/model-registry.ts b/packages/coding-agent/src/core/model-registry.ts index 9e2fc094..1d07c5d8 100644 --- a/packages/coding-agent/src/core/model-registry.ts +++ b/packages/coding-agent/src/core/model-registry.ts @@ -13,6 +13,7 @@ import { } from "@mariozechner/pi-ai"; import { type Static, Type } from "@sinclair/typebox"; import AjvModule from "ajv"; +import { execSync } from "child_process"; import { existsSync, readFileSync } from "fs"; import type { AuthStorage } from "./auth-storage.js"; @@ -99,14 +100,47 @@ function emptyCustomModelsResult(error?: string): CustomModelsResult { return { models: [], replacedProviders: new Set(), overrides: new Map(), error }; } +// Cache for shell command results (persists for process lifetime) +const commandResultCache = new Map(); + /** * Resolve an API key config value to an actual key. - * Checks environment variable first, then treats as literal. + * - If starts with "!", executes the rest as a shell command and uses stdout (cached) + * - Otherwise checks environment variable first, then treats as literal (not cached) */ function resolveApiKeyConfig(keyConfig: string): string | undefined { + if (keyConfig.startsWith("!")) { + return executeApiKeyCommand(keyConfig); + } const envValue = process.env[keyConfig]; - if (envValue) return envValue; - return keyConfig; + return envValue || keyConfig; +} + +function executeApiKeyCommand(commandConfig: string): string | undefined { + if (commandResultCache.has(commandConfig)) { + return commandResultCache.get(commandConfig); + } + + const command = commandConfig.slice(1); + let result: string | undefined; + try { + const output = execSync(command, { + encoding: "utf-8", + timeout: 10000, + stdio: ["ignore", "pipe", "ignore"], + }); + result = output.trim() || undefined; + } catch { + result = undefined; + } + + commandResultCache.set(commandConfig, result); + return result; +} + +/** Clear the API key command cache. Exported for testing. */ +export function clearApiKeyCache(): void { + commandResultCache.clear(); } /** diff --git a/packages/coding-agent/test/model-registry.test.ts b/packages/coding-agent/test/model-registry.test.ts index 26bad48e..cc5f4dee 100644 --- a/packages/coding-agent/test/model-registry.test.ts +++ b/packages/coding-agent/test/model-registry.test.ts @@ -1,9 +1,9 @@ -import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs"; +import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, beforeEach, describe, expect, test } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.js"; -import { ModelRegistry } from "../src/core/model-registry.js"; +import { clearApiKeyCache, ModelRegistry } from "../src/core/model-registry.js"; describe("ModelRegistry", () => { let tempDir: string; @@ -21,6 +21,7 @@ describe("ModelRegistry", () => { if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); } + clearApiKeyCache(); }); /** Create minimal provider config */ @@ -246,4 +247,273 @@ describe("ModelRegistry", () => { expect(anthropicModels.some((m) => m.id.includes("claude"))).toBe(true); }); }); + + describe("API key resolution", () => { + /** Create provider config with custom apiKey */ + function providerWithApiKey(apiKey: string) { + return { + baseUrl: "https://example.com/v1", + apiKey, + api: "anthropic-messages", + models: [ + { + id: "test-model", + name: "Test Model", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 100000, + maxTokens: 8000, + }, + ], + }; + } + + test("apiKey with ! prefix executes command and uses stdout", async () => { + writeRawModelsJson({ + "custom-provider": providerWithApiKey("!echo test-api-key-from-command"), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const apiKey = await registry.getApiKeyForProvider("custom-provider"); + + expect(apiKey).toBe("test-api-key-from-command"); + }); + + test("apiKey with ! prefix trims whitespace from command output", async () => { + writeRawModelsJson({ + "custom-provider": providerWithApiKey("!echo ' spaced-key '"), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const apiKey = await registry.getApiKeyForProvider("custom-provider"); + + expect(apiKey).toBe("spaced-key"); + }); + + test("apiKey with ! prefix handles multiline output (uses trimmed result)", async () => { + writeRawModelsJson({ + "custom-provider": providerWithApiKey("!printf 'line1\\nline2'"), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const apiKey = await registry.getApiKeyForProvider("custom-provider"); + + expect(apiKey).toBe("line1\nline2"); + }); + + test("apiKey with ! prefix returns undefined on command failure", async () => { + writeRawModelsJson({ + "custom-provider": providerWithApiKey("!exit 1"), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const apiKey = await registry.getApiKeyForProvider("custom-provider"); + + expect(apiKey).toBeUndefined(); + }); + + test("apiKey with ! prefix returns undefined on nonexistent command", async () => { + writeRawModelsJson({ + "custom-provider": providerWithApiKey("!nonexistent-command-12345"), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const apiKey = await registry.getApiKeyForProvider("custom-provider"); + + expect(apiKey).toBeUndefined(); + }); + + test("apiKey with ! prefix returns undefined on empty output", async () => { + writeRawModelsJson({ + "custom-provider": providerWithApiKey("!printf ''"), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const apiKey = await registry.getApiKeyForProvider("custom-provider"); + + expect(apiKey).toBeUndefined(); + }); + + test("apiKey as environment variable name resolves to env value", async () => { + const originalEnv = process.env.TEST_API_KEY_12345; + process.env.TEST_API_KEY_12345 = "env-api-key-value"; + + try { + writeRawModelsJson({ + "custom-provider": providerWithApiKey("TEST_API_KEY_12345"), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const apiKey = await registry.getApiKeyForProvider("custom-provider"); + + expect(apiKey).toBe("env-api-key-value"); + } finally { + if (originalEnv === undefined) { + delete process.env.TEST_API_KEY_12345; + } else { + process.env.TEST_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; + + writeRawModelsJson({ + "custom-provider": providerWithApiKey("literal_api_key_value"), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const apiKey = await registry.getApiKeyForProvider("custom-provider"); + + expect(apiKey).toBe("literal_api_key_value"); + }); + + test("apiKey command can use shell features like pipes", async () => { + writeRawModelsJson({ + "custom-provider": providerWithApiKey("!echo 'hello world' | tr ' ' '-'"), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + const apiKey = await registry.getApiKeyForProvider("custom-provider"); + + 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"'`; + writeRawModelsJson({ + "custom-provider": providerWithApiKey(command), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + + // Call multiple times + await registry.getApiKeyForProvider("custom-provider"); + await registry.getApiKeyForProvider("custom-provider"); + await registry.getApiKeyForProvider("custom-provider"); + + // Command should have only run once + const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); + expect(count).toBe(1); + }); + + test("cache persists across registry instances", async () => { + const counterFile = join(tempDir, "counter"); + writeFileSync(counterFile, "0"); + + const command = `!sh -c 'count=$(cat ${counterFile}); echo $((count + 1)) > ${counterFile}; echo "key-value"'`; + writeRawModelsJson({ + "custom-provider": providerWithApiKey(command), + }); + + // Create multiple registry instances + const registry1 = new ModelRegistry(authStorage, modelsJsonPath); + await registry1.getApiKeyForProvider("custom-provider"); + + const registry2 = new ModelRegistry(authStorage, modelsJsonPath); + await registry2.getApiKeyForProvider("custom-provider"); + + // Command should still have only run once + const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); + expect(count).toBe(1); + }); + + test("clearApiKeyCache 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"'`; + writeRawModelsJson({ + "custom-provider": providerWithApiKey(command), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + await registry.getApiKeyForProvider("custom-provider"); + + // Clear cache and call again + clearApiKeyCache(); + await registry.getApiKeyForProvider("custom-provider"); + + // 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 () => { + writeRawModelsJson({ + "provider-a": providerWithApiKey("!echo key-a"), + "provider-b": providerWithApiKey("!echo key-b"), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + + const keyA = await registry.getApiKeyForProvider("provider-a"); + const keyB = await registry.getApiKeyForProvider("provider-b"); + + expect(keyA).toBe("key-a"); + expect(keyB).toBe("key-b"); + }); + + 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'`; + writeRawModelsJson({ + "custom-provider": providerWithApiKey(command), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + + // Call multiple times - all should return undefined + const key1 = await registry.getApiKeyForProvider("custom-provider"); + const key2 = await registry.getApiKeyForProvider("custom-provider"); + + 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_API_KEY_CACHE_TEST_98765"; + const originalEnv = process.env[envVarName]; + + try { + process.env[envVarName] = "first-value"; + + writeRawModelsJson({ + "custom-provider": providerWithApiKey(envVarName), + }); + + const registry = new ModelRegistry(authStorage, modelsJsonPath); + + const key1 = await registry.getApiKeyForProvider("custom-provider"); + expect(key1).toBe("first-value"); + + // Change env var + process.env[envVarName] = "second-value"; + + const key2 = await registry.getApiKeyForProvider("custom-provider"); + expect(key2).toBe("second-value"); + } finally { + if (originalEnv === undefined) { + delete process.env[envVarName]; + } else { + process.env[envVarName] = originalEnv; + } + } + }); + }); + }); });