pi-provider-litellm 1.2.3 → 1.2.5

package/README.md

@@ -61,6 +61,7 @@ Stored pi credentials for `litellm` take precedence over `LITELLM_API_KEY`; the
 
 | Variable | Default | Effect |
 |---|---|---|
+| `LITELLM_API_KEY_HELPER` | unset | Command that prints a fresh LiteLLM bearer token. Takes precedence over `LITELLM_API_KEY`. Registered as a `!command` provider key; Pi re-runs it on every request (the per-request auth path is uncached), so rotating/short-lived tokens stay fresh. |
 | `LITELLM_OFFLINE` | unset | If `1`, skip discovery on this start; use cache only |
 | `LITELLM_DISCOVERY_TIMEOUT_MS` | `5000` | Discovery fetch timeout in ms; `0` to skip discovery |
 
@@ -117,6 +118,7 @@ If the cache is older than 24 hours, the extension refreshes it in the backgroun
 | "discovered no models" | Proxy returned an empty list — check pi's startup log for the actual response |
 | `/model/info` returning 401/403/404 | Expected behavior with virtual keys — extension falls back to `/v1/models` |
 | Discovery times out | Increase `LITELLM_DISCOVERY_TIMEOUT_MS` or set `LITELLM_OFFLINE=1` to fall back on cached models |
+| `401 Token expired` | Set `LITELLM_API_KEY_HELPER`. |
 
 ## License
 

package/dist/index.js

@@ -1,3 +1,4 @@
+import { execSync } from "node:child_process";
 import { readFile } from "node:fs/promises";
 import { join } from "node:path";
 import { AuthStorage, getAgentDir } from "@earendil-works/pi-coding-agent";
@@ -8,12 +9,16 @@ import { getSessionIdFromFile } from "./litellm.js";
 const PROVIDER_NAME = "litellm";
 const ENV_BASE_URL = "LITELLM_BASE_URL";
 const ENV_API_KEY = "LITELLM_API_KEY";
+const ENV_API_KEY_HELPER = "LITELLM_API_KEY_HELPER";
 const ENV_TIMEOUT = "LITELLM_DISCOVERY_TIMEOUT_MS";
 const ENV_OFFLINE = "LITELLM_OFFLINE";
 const DEFAULT_TIMEOUT_MS = 5000;
 const LOGIN_TIMEOUT_MS = 10_000;
 const CACHE_FILENAME = "litellm-models.json";
 const CACHE_STALE_MS = 24 * 60 * 60 * 1000;
+const TOKEN_REFRESH_LEAD_MS = 5 * 60 * 1000;
+const PERMANENT_TOKEN_EXPIRES_AT = Number.MAX_SAFE_INTEGER;
+const EXPIRE_TOKEN_IMMEDIATELY = 0;
 function getAuthPath() {
     return join(getAgentDir(), "auth.json");
 }
@@ -30,20 +35,74 @@ async function readAuthEntry() {
         return undefined;
     }
 }
-async function resolveCredentials() {
+function cleanConfig(raw) {
+    const trimmed = raw?.trim();
+    return trimmed && trimmed !== "undefined" ? trimmed : undefined;
+}
+function normalizeCommand(raw) {
+    const trimmed = cleanConfig(raw);
+    if (!trimmed)
+        return undefined;
+    return trimmed.startsWith("!") ? trimmed : `!${trimmed}`;
+}
+function getApiKeyHelperCommand() {
+    return normalizeCommand(process.env[ENV_API_KEY_HELPER]);
+}
+function executeApiKeyCommand(commandConfig) {
+    const command = commandConfig.startsWith("!") ? commandConfig.slice(1) : commandConfig;
+    const output = execSync(command, { encoding: "utf8", stdio: ["ignore", "pipe", "pipe"], timeout: 10_000 }).trim();
+    if (!output)
+        throw new Error(`LiteLLM API key helper produced no output: ${command}`);
+    return output;
+}
+function tokenExpiresAt(apiKey, opaqueFallback = PERMANENT_TOKEN_EXPIRES_AT) {
+    const [, payload] = apiKey.split(".");
+    if (!payload)
+        return opaqueFallback;
+    try {
+        const claims = JSON.parse(Buffer.from(payload, "base64url").toString("utf8"));
+        return typeof claims.exp === "number"
+            ? Math.max(Date.now(), claims.exp * 1000 - TOKEN_REFRESH_LEAD_MS)
+            : opaqueFallback;
+    }
+    catch {
+        return opaqueFallback;
+    }
+}
+function getLiteLLMApiKey(credentials) {
+    return credentials.access;
+}
+function resolveOAuthApiKey(credentials) {
+    return credentials.refresh.startsWith("!")
+        ? executeApiKeyCommand(credentials.refresh)
+        : getLiteLLMApiKey(credentials);
+}
+async function resolveCredentials({ executeHelpers = true } = {}) {
     const entry = await readAuthEntry();
-    const envBase = process.env[ENV_BASE_URL]?.trim();
-    const envKey = process.env[ENV_API_KEY]?.trim();
+    const envBase = cleanConfig(process.env[ENV_BASE_URL]);
+    const envKey = cleanConfig(process.env[ENV_API_KEY]);
+    const envHelperCommand = getApiKeyHelperCommand();
     const authBase = entry?.type === "oauth" ? entry.baseUrl?.trim() : undefined;
     const authKey = entry?.type === "oauth"
-        ? entry.access?.trim()
+        ? (executeHelpers ? resolveOAuthApiKey(entry) : entry.access).trim()
         : entry?.type === "api_key"
             ? (await AuthStorage.create(getAuthPath()).getApiKey(PROVIDER_NAME, { includeFallback: false }))?.trim()
             : undefined;
+    const apiKey = authKey || (executeHelpers && envHelperCommand ? executeApiKeyCommand(envHelperCommand) : undefined) || envKey;
+    const apiKeyFingerprint = entry?.type === "oauth" && entry.refresh.startsWith("!")
+        ? fingerprint(entry.refresh)
+        : authKey
+            ? fingerprint(authKey)
+            : envHelperCommand
+                ? fingerprint(envHelperCommand)
+                : envKey
+                    ? fingerprint(envKey)
+                    : undefined;
     const rawBase = authBase || envBase;
     return {
         baseUrl: rawBase ? normalizeBaseUrl(rawBase) : undefined,
-        apiKey: authKey || envKey || undefined,
+        apiKey: apiKey || undefined,
+        apiKeyFingerprint,
     };
 }
 function getDiscoveryTimeoutMs() {
@@ -78,17 +137,19 @@ async function loginLiteLLM(callbacks, onCacheWrite) {
         message: "Enter LiteLLM proxy URL (no trailing /v1):",
         placeholder: "https://litellm.example.com",
     })).trim();
-    const apiKey = (await callbacks.onPrompt({ message: "Enter API key:" })).trim();
-    if (!rawBaseUrl || !apiKey)
+    const apiKeyInput = (await callbacks.onPrompt({ message: "Enter API key:" })).trim();
+    if (!rawBaseUrl || !apiKeyInput)
         throw new Error("Both base URL and API key are required");
     const baseUrl = normalizeBaseUrl(rawBaseUrl);
+    const refresh = apiKeyInput.startsWith("!") ? apiKeyInput : "";
+    const apiKey = refresh ? executeApiKeyCommand(refresh) : apiKeyInput;
     const { models, source } = await discoverModels(baseUrl, apiKey, {
         timeoutMs: LOGIN_TIMEOUT_MS,
         signal: callbacks.signal,
     });
     const cache = {
         baseUrl,
-        apiKeyFingerprint: fingerprint(apiKey),
+        apiKeyFingerprint: fingerprint(refresh || apiKey),
         fetchedAt: Date.now(),
         source,
         models,
@@ -98,13 +159,16 @@ async function loginLiteLLM(callbacks, onCacheWrite) {
     callbacks.onProgress?.(`LiteLLM: ${models.length} models discovered (source: ${source})`);
     return {
         access: apiKey,
-        refresh: "",
-        expires: Number.MAX_SAFE_INTEGER,
+        refresh,
+        expires: tokenExpiresAt(apiKey, refresh ? EXPIRE_TOKEN_IMMEDIATELY : PERMANENT_TOKEN_EXPIRES_AT),
         baseUrl,
     };
 }
 async function refreshLiteLLM(credentials) {
-    return credentials;
+    if (!credentials.refresh.startsWith("!"))
+        return credentials;
+    const access = executeApiKeyCommand(credentials.refresh);
+    return { ...credentials, access, expires: tokenExpiresAt(access, EXPIRE_TOKEN_IMMEDIATELY) };
 }
 function modifyLiteLLMModels(models, cred) {
     const baseUrl = cred.baseUrl;
@@ -194,9 +258,9 @@ function normalizeThinkTags(message) {
     return { ...message, content };
 }
 export default async function (pi) {
-    const creds = await resolveCredentials();
+    let creds = await resolveCredentials({ executeHelpers: false });
     const cache = await readCache(getCachePath());
-    const fp = creds.apiKey ? fingerprint(creds.apiKey) : undefined;
+    let fp = creds.apiKeyFingerprint;
     let cacheFetchedAt = cache?.fetchedAt ?? 0;
     const cacheValid = cache !== null &&
         creds.baseUrl !== undefined &&
@@ -204,61 +268,86 @@ export default async function (pi) {
         cache.baseUrl === creds.baseUrl &&
         cache.apiKeyFingerprint === fp;
     let models = cacheValid && cache ? cache.models : [];
-    const haveCreds = creds.baseUrl !== undefined && creds.apiKey !== undefined && fp !== undefined;
-    const shouldFetch = haveCreds && !isOffline() && (!cacheValid || isListModelsMode());
-    if (shouldFetch && creds.baseUrl && creds.apiKey && fp) {
+    const shouldFetch = creds.baseUrl !== undefined &&
+        fp !== undefined &&
+        !isOffline() &&
+        getDiscoveryTimeoutMs() > 0 &&
+        (!cacheValid || isListModelsMode());
+    let credentialWarning;
+    if (shouldFetch) {
+        try {
+            creds = await resolveCredentials();
+            fp = creds.apiKeyFingerprint;
+        }
+        catch (error) {
+            if (!cacheValid || !cache)
+                throw error;
+            credentialWarning = error instanceof Error ? error.message : String(error);
+            process.stderr.write(`LiteLLM: discovery failed (${credentialWarning}); using cached models.\n`);
+            models = cache.models;
+        }
+    }
+    if (shouldFetch && !credentialWarning && creds.baseUrl && creds.apiKey && fp) {
         const timeoutMs = getDiscoveryTimeoutMs();
-        if (timeoutMs > 0) {
-            const { result, warning } = await discoverWithFallback(creds.baseUrl, creds.apiKey, {
-                timeoutMs,
-            });
-            if (warning) {
-                if (cacheValid && cache) {
-                    process.stderr.write(`LiteLLM: discovery failed (${warning}); using cached models.\n`);
-                    models = cache.models;
-                }
-                else {
-                    process.stderr.write(`LiteLLM: discovery failed (${warning}); registering provider with no models.\n`);
-                    models = [];
-                }
+        const { result, warning } = await discoverWithFallback(creds.baseUrl, creds.apiKey, {
+            timeoutMs,
+        });
+        if (warning) {
+            if (cacheValid && cache) {
+                process.stderr.write(`LiteLLM: discovery failed (${warning}); using cached models.\n`);
+                models = cache.models;
             }
             else {
-                models = result.models;
-                const next = {
-                    baseUrl: creds.baseUrl,
-                    apiKeyFingerprint: fp,
-                    fetchedAt: Date.now(),
-                    source: result.source,
-                    models: result.models,
-                };
-                await writeCache(getCachePath(), next);
-                cacheFetchedAt = next.fetchedAt;
-                if (isListModelsMode()) {
-                    process.stderr.write(`LiteLLM: ${result.models.length} models discovered (source: ${result.source}).\n`);
-                }
+                process.stderr.write(`LiteLLM: discovery failed (${warning}); registering provider with no models.\n`);
+                models = [];
+            }
+        }
+        else {
+            models = result.models;
+            const next = {
+                baseUrl: creds.baseUrl,
+                apiKeyFingerprint: fp,
+                fetchedAt: Date.now(),
+                source: result.source,
+                models: result.models,
+            };
+            await writeCache(getCachePath(), next);
+            cacheFetchedAt = next.fetchedAt;
+            if (isListModelsMode()) {
+                process.stderr.write(`LiteLLM: ${result.models.length} models discovered (source: ${result.source}).\n`);
             }
         }
     }
+    let updateCosts = () => undefined;
     const oauth = {
         name: "LiteLLM",
         login: (callbacks) => loginLiteLLM(callbacks, (next) => {
             cacheFetchedAt = next.fetchedAt;
+            registerProvider(next.baseUrl, next.models);
+            updateCosts(next.models);
         }),
         refreshToken: refreshLiteLLM,
-        getApiKey: (cred) => cred.access,
+        getApiKey: getLiteLLMApiKey,
         modifyModels: modifyLiteLLMModels,
     };
     function registerProvider(baseUrl, models) {
         pi.registerProvider(PROVIDER_NAME, {
             baseUrl: baseUrl ? `${baseUrl}/v1` : "https://litellm.example.com/v1",
-            apiKey: ENV_API_KEY,
+            // When LITELLM_API_KEY_HELPER is set we register the helper as a `!command` provider key.
+            // Pi's per-request auth path (ModelRegistry.getApiKeyAndHeaders) resolves provider keys via
+            // resolveConfigValueOrThrow -> resolveConfigValueUncached, i.e. the command is re-executed on
+            // every request (it does NOT use the process-lifetime command cache, which only applies to
+            // resolveConfigValue). So a short-lived/rotating helper token stays fresh. The OAuth hooks
+            // remain registered for `/login litellm` users. See the regression test
+            // "re-runs the helper command on every request" in tests/index.test.ts.
+            apiKey: getApiKeyHelperCommand() ?? `$${ENV_API_KEY}`,
             api: "openai-completions",
             models,
             oauth,
         });
     }
     registerProvider(creds.baseUrl, models);
-    const updateCosts = setupLiteLLMCostTracking(pi, models);
+    updateCosts = setupLiteLLMCostTracking(pi, models);
     let refreshInProgress = null;
     function discoveryDisabledReason() {
         if (isOffline())
@@ -269,7 +358,7 @@ export default async function (pi) {
     }
     async function refreshModelsAndCosts() {
         const fresh = await resolveCredentials();
-        const freshFp = fresh.apiKey ? fingerprint(fresh.apiKey) : undefined;
+        const freshFp = fresh.apiKeyFingerprint;
         if (!fresh.baseUrl || !fresh.apiKey || !freshFp) {
             throw new Error("no credentials. Run /login litellm or set env vars.");
         }

package/dist/types.d.ts

@@ -52,4 +52,5 @@ export type AuthFileEntry = {
 export interface ResolvedCredentials {
     baseUrl?: string;
     apiKey?: string;
+    apiKeyFingerprint?: string;
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-provider-litellm",
-  "version": "1.2.3",
+  "version": "1.2.5",
   "description": "LiteLLM proxy provider extension for Pi",
   "type": "module",
   "license": "MIT",
@@ -55,8 +55,8 @@
     "node": ">=22.19.0"
   },
   "peerDependencies": {
-    "@earendil-works/pi-ai": "^0.75.4",
-    "@earendil-works/pi-coding-agent": "^0.75.4"
+    "@earendil-works/pi-ai": "*",
+    "@earendil-works/pi-coding-agent": "*"
   },
   "overrides": {
     "basic-ftp": "6.0.1",
@@ -64,11 +64,11 @@
     "protobufjs": "8.2.1"
   },
   "devDependencies": {
-    "@earendil-works/pi-ai": "^0.75.4",
-    "@earendil-works/pi-coding-agent": "^0.75.4",
+    "@earendil-works/pi-ai": "^0.78.1",
+    "@earendil-works/pi-coding-agent": "^0.78.1",
     "@biomejs/biome": "^2.4.15",
     "@types/node": "^25.9.1",
-    "@typescript/native-preview": "7.0.0-dev.20260519.1",
+    "@typescript/native-preview": "7.0.0-dev.20260607.1",
     "shx": "^0.4.0",
     "tsx": "^4.21.0",
     "typescript": "^6.0.3",