Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
67 kB · 1925 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926/** * Credential storage for API keys and OAuth tokens. * Handles loading, saving, and refreshing credentials from auth.json. * * Uses file locking to prevent race conditions when multiple pi instances * try to refresh tokens simultaneously. */
import { Buffer } from "node:buffer";import { createHash } from "node:crypto";import { findEnvKeys, getEnvApiKey, type OAuthCredentials, type OAuthLoginCallbacks, type OAuthProviderId,} from "@earendil-works/pi-ai";import { getOAuthApiKey, getOAuthProvider, getOAuthProviders } from "@earendil-works/pi-ai/oauth";import { closeSync, existsSync, fchmodSync, mkdirSync, openSync, readFileSync, writeSync } from "fs";import { dirname, join } from "path";import lockfile from "proper-lockfile";import { getAgentDir } from "../config.js";import { realpathIfPresentSync, writeFileAtomicSync } from "../utils/atomic-file.js";import { getPrimeCliConfigPath, loadPrimeCliConfig, PRIME_INFERENCE_PROVIDER_ID, type PrimeCliConfig, type PrimeTeam, savePrimeCliApiKey, savePrimeCliTeamSelection,} from "./prime-inference-auth.js";import { resolveConfigValue, resolveConfigValueUncached } from "./resolve-config-value.js";
export type PrimeTeamCredential = { teamId: string; name: string; slug?: string; role?: string; createdAt?: string;};
export type ApiKeyCredential = { type: "api_key"; key: string; primeTeam?: PrimeTeamCredential | null;};
export type OAuthCredential = { type: "oauth";} & OAuthCredentials;
/** * A static token pasted for one MCP connection through the inline paste flow. * Deliberately NOT the OAuth shape: there is no refresh token, no expiry, and * no client identity to fake — the handshake sends `bearer` as * `Authorization: Bearer`. Exactly ONE credential per connection (the paste * flow prompts once; multiple catalog fields may only be alternative names for * that one credential). Bound to the exact endpoint it was pasted for, stored * only in the credential store under the owning connection's * `mcp:<connectionId>` key — never in settings.json. */export type McpStaticTokenCredential = { type: "mcp_static_token"; /** The endpoint the pasted token is bound to; a retargeted entry fails closed. */ endpoint: string; /** The value the MCP handshake sends as the bearer. */ bearer: string; /** The catalog setup field id the token was collected for (the first alternative name). */ bearerFieldId: string; createdAt: number;};
export type AuthCredential = ApiKeyCredential | OAuthCredential | McpStaticTokenCredential;
/** * One provider may store several accounts. Older auth.json files store a bare * credential object; newer files store an array once more than one account is * logged in. Both forms load interchangeably. */export type AuthCredentialEntry = AuthCredential | AuthCredential[];
export type AuthStorageData = Record<string, AuthCredentialEntry>;
/** Default cooldown applied to a credential after a usage-limit response. */export const DEFAULT_RATE_LIMIT_BLOCK_MS = 60_000;
export type StoredCredentialSummary = { provider: string; type: AuthCredential["type"]; /** Stable identity: account id/email when known, otherwise a token/key digest. */ identity: string; /** Human-readable label for UI, e.g. the OAuth account email. */ label?: string;};
export type AuthStatus = { configured: boolean; source?: | "stored" | "runtime" | "environment" | "prime_cli" | "fallback" | "models_json_key" | "models_json_command" | "stale"; label?: string;};
export type AuthStorageOptions = { primeCliConfigPath?: string; usePrimeCliConfig?: boolean;};
type LockResult<T> = { result: T; next?: string;};
type ActiveAuthStatusSource = Exclude<NonNullable<AuthStatus["source"]>, "stale">;
export type AuthSourceToken = { provider: string; source: ActiveAuthStatusSource; identityFingerprint: string; valueFingerprint: string;};
type AuthSourceCandidate = { source: ActiveAuthStatusSource; configured: boolean; label?: string; identityFingerprint: string; valueFingerprint?: string; resolveValueFingerprint?: () => string | undefined;};
type AuthApiKeyResult = { apiKey?: string; sourceToken?: AuthSourceToken; credentialType?: AuthCredential["type"];};
export interface AuthStorageBackend { withLock<T>(fn: (current: string | undefined) => LockResult<T>): T; withLockAsync<T>(fn: (current: string | undefined) => Promise<LockResult<T>>): Promise<T>;}
export class FileAuthStorageBackend implements AuthStorageBackend { constructor(private authPath: string = join(getAgentDir(), "auth.json")) {}
/** Path of the backing auth.json (rate-limit blocks live beside it). */ get path(): string { return this.authPath; }
private ensureParentDir(): void { const dir = dirname(this.authPath); if (!existsSync(dir)) { mkdirSync(dir, { recursive: true, mode: 0o700 }); } }
private ensureFileExists(): void { let descriptor: number; try { // Exclusive create: a racing initializer must never replace saved credentials. descriptor = openSync(this.authPath, "wx", 0o600); } catch (error) { if ((error as NodeJS.ErrnoException).code !== "EEXIST") { throw error; } return; } try { const bytes = Buffer.from("{}"); let offset = 0; while (offset < bytes.length) { const written = writeSync(descriptor, bytes, offset, bytes.length - offset); if (written <= 0) throw new Error(`Short write initializing ${this.authPath}`); offset += written; } fchmodSync(descriptor, 0o600); // Exact bits despite the umask. } finally { closeSync(descriptor); } }
private acquireLockSyncWithRetry(path: string): () => void { const maxAttempts = 10; const delayMs = 20; let lastError: unknown; let compromisedError: Error | undefined;
for (let attempt = 1; attempt <= maxAttempts; attempt++) { try { const release = lockfile.lockSync(path, { realpath: false, onCompromised: (error) => { compromisedError ??= error; }, }); if (compromisedError) { release(); throw compromisedError; } return release; } catch (error) { if (compromisedError) throw compromisedError; const code = typeof error === "object" && error !== null && "code" in error ? String((error as { code?: unknown }).code) : undefined; if (code !== "ELOCKED" || attempt === maxAttempts) { throw error; } lastError = error; const start = Date.now(); while (Date.now() - start < delayMs) { // Sleep synchronously to avoid changing callers to async. } } }
throw (lastError as Error) ?? new Error("Failed to acquire auth storage lock"); }
withLock<T>(fn: (current: string | undefined) => LockResult<T>): T { this.ensureParentDir(); this.ensureFileExists();
let release: (() => void) | undefined; try { release = this.acquireLockSyncWithRetry(this.authPath); const current = existsSync(this.authPath) ? readFileSync(this.authPath, "utf-8") : undefined; const { result, next } = fn(current); if (next !== undefined) { writeFileAtomicSync(realpathIfPresentSync(this.authPath), next, { mode: 0o600 }); } return result; } finally { if (release) { release(); } } }
async withLockAsync<T>(fn: (current: string | undefined) => Promise<LockResult<T>>): Promise<T> { this.ensureParentDir(); this.ensureFileExists();
let release: (() => Promise<void>) | undefined; let lockCompromised = false; let lockCompromisedError: Error | undefined; const throwIfCompromised = () => { if (lockCompromised) { throw lockCompromisedError ?? new Error("Auth storage lock was compromised"); } };
try { release = await lockfile.lock(this.authPath, { retries: { retries: 10, factor: 2, minTimeout: 100, maxTimeout: 10000, randomize: true, }, stale: 30000, onCompromised: (err) => { lockCompromised = true; lockCompromisedError = err; }, });
throwIfCompromised(); const current = existsSync(this.authPath) ? readFileSync(this.authPath, "utf-8") : undefined; const { result, next } = await fn(current); throwIfCompromised(); if (next !== undefined) { writeFileAtomicSync(realpathIfPresentSync(this.authPath), next, { mode: 0o600 }); } throwIfCompromised(); return result; } finally { if (release) { if (lockCompromised) await release().catch(() => undefined); else await release(); } } }}
export class InMemoryAuthStorageBackend implements AuthStorageBackend { private value: string | undefined;
withLock<T>(fn: (current: string | undefined) => LockResult<T>): T { const { result, next } = fn(this.value); if (next !== undefined) { this.value = next; } return result; }
async withLockAsync<T>(fn: (current: string | undefined) => Promise<LockResult<T>>): Promise<T> { const { result, next } = await fn(this.value); if (next !== undefined) { this.value = next; } return result; }}
/** * Credential storage backed by a JSON file. *//** FNV-1a hash for deterministic session-to-credential mapping. */function hashSessionId(sessionId: string): number { let hash = 0x811c9dc5; for (let i = 0; i < sessionId.length; i++) { hash ^= sessionId.charCodeAt(i); hash = Math.imul(hash, 0x01000193) >>> 0; } return hash >>> 0;}
function shortDigest(material: string): string { return createHash("sha256").update(material).digest("hex").slice(0, 16);}
function oauthAccountId(credential: OAuthCredential): string | undefined { const explicit = (credential as { accountId?: unknown }).accountId; if (typeof explicit === "string" && explicit.length > 0) return explicit;
try { const parts = credential.access.split("."); if (parts.length !== 3) return undefined; const payload = JSON.parse(Buffer.from(parts[1], "base64url").toString("utf8")) as Record<string, unknown>; const auth = payload["https://api.openai.com/auth"]; if (!auth || typeof auth !== "object") return undefined; const accountId = (auth as Record<string, unknown>).chatgpt_account_id; return typeof accountId === "string" && accountId.length > 0 ? accountId : undefined; } catch { return undefined; }}
/** * Stable identity for one stored credential, used to dedupe logins of the * same account and to mark a single credential stale when only it fails. *//** * One provider may store several accounts: older auth.json files hold a bare * credential, newer ones an array. Callers that work on a single credential read * the first account and fall back to the bare form. */export function firstCredential(entry: AuthCredentialEntry | undefined): AuthCredential | undefined { return Array.isArray(entry) ? entry[0] : entry;}
export function credentialIdentity(credential: AuthCredential): string { if (credential.type === "api_key") { return credential.key.startsWith("!") ? `command:${credential.key}` : `key:${shortDigest(credential.key)}`; } if (credential.type === "mcp_static_token") return `mcp:${shortDigest(credential.bearer)}`; const accountId = oauthAccountId(credential); if (accountId) return `account:${accountId}`; const email = (credential as { email?: unknown }).email; if (typeof email === "string" && email.length > 0) { return `email:${email.toLowerCase()}`; } return `refresh:${shortDigest(credential.refresh)}`;}
function credentialLabel(credential: AuthCredential): string | undefined { if (credential.type === "api_key") return undefined; if (credential.type === "mcp_static_token") return undefined; const email = (credential as { email?: unknown }).email; if (typeof email === "string" && email.length > 0) return email; const accountId = oauthAccountId(credential); return accountId ? `account ${accountId.slice(0, 8)}` : undefined;}
type PersistedCredentialBlock = { provider: string; identity: string; blockedUntilMs: number; cause?: string;};
export class AuthStorage { private data: AuthStorageData = {}; private runtimeOverrides: Map<string, string> = new Map(); private staleAuthSources: Map<string, AuthSourceToken[]> = new Map(); private fallbackResolver?: (provider: string) => string | undefined; private loadError: Error | null = null; private errors: Error[] = []; /** Round-robin cursor per provider, for load distribution across accounts. */ private roundRobinIndex = new Map<string, number>(); /** Last credential used per (provider, session), so a session keeps its account. */ private sessionCredentials = new Map<string, string>(); /** Rate-limit blocks: provider -> credential identity -> unblock timestamp. */ private credentialBlocks = new Map<string, Map<string, PersistedCredentialBlock>>(); /** Last stored credential selected per provider, for targeted block/stale marking. */ private lastSelectedIdentities = new Map<string, string>(); private blockStorePath?: string;
private constructor( private storage: AuthStorageBackend, private options: AuthStorageOptions = {}, ) { this.blockStorePath = storage instanceof FileAuthStorageBackend ? join(dirname(storage.path), "auth-blocks.json") : undefined; this.reload(); this.loadPersistedBlocks(); }
static create(authPath?: string, options?: AuthStorageOptions): AuthStorage { const authOptions = options ?? { usePrimeCliConfig: authPath === undefined }; return new AuthStorage(new FileAuthStorageBackend(authPath ?? join(getAgentDir(), "auth.json")), authOptions); }
static fromStorage(storage: AuthStorageBackend, options?: AuthStorageOptions): AuthStorage { return new AuthStorage(storage, options); }
static inMemory(data: AuthStorageData = {}, options?: AuthStorageOptions): AuthStorage { const storage = new InMemoryAuthStorageBackend(); storage.withLock(() => ({ result: undefined, next: JSON.stringify(data, null, 2) })); return AuthStorage.fromStorage(storage, options); }
/** * Set a runtime API key override (not persisted to disk). * Used for CLI --api-key flag. */ setRuntimeApiKey(provider: string, apiKey: string): void { this.clearStaleAuthSource(provider, "runtime"); this.runtimeOverrides.set(provider, apiKey); }
/** * Remove a runtime API key override. */ removeRuntimeApiKey(provider: string): void { this.clearStaleAuthSource(provider, "runtime"); this.runtimeOverrides.delete(provider); }
/** * Set a fallback resolver for API keys not found in auth.json or env vars. * Used for custom provider keys from models.json. */ setFallbackResolver(resolver: (provider: string) => string | undefined): void { this.fallbackResolver = resolver; }
private recordError(error: unknown): void { const normalizedError = error instanceof Error ? error : new Error(String(error)); this.errors.push(normalizedError); }
// ── Multi-account credential pool ─────────────────────────────────────────
/** All stored credentials for a provider, normalizing the legacy single form. */ getCredentialList(provider: string): AuthCredential[] { const entry = this.data[provider]; if (!entry) return []; return Array.isArray(entry) ? entry : [entry]; }
private setCredentialList(provider: string, credentials: AuthCredential[]): void { if (credentials.length === 0) { delete this.data[provider]; } else if (credentials.length === 1) { this.data[provider] = credentials[0]; } else { this.data[provider] = credentials; } }
/** * Resolve one stored credential by identity, refreshing OAuth when necessary. * This is intentionally per-account so multi-account usage probes never * accidentally refresh or report a sibling account. */ async getApiKeyForCredential( providerId: string, identity: string, ): Promise<{ apiKey: string; credential: AuthCredential } | undefined> { const credentials = this.getCredentialList(providerId); const credentialIndex = credentials.findIndex((credential) => credentialIdentity(credential) === identity); if (credentialIndex < 0) return undefined; const credential = credentials[credentialIndex]; if (credential.type === "api_key") return { apiKey: credential.key, credential }; if (credential.type === "mcp_static_token") return { apiKey: credential.bearer, credential }; const provider = getOAuthProvider(providerId as OAuthProviderId); if (!provider) return undefined; if (Date.now() < credential.expires) return { apiKey: provider.getApiKey(credential), credential }; const refreshed = await this.refreshOAuthTokenWithLock(providerId as OAuthProviderId, credentialIndex); if (!refreshed) return undefined; return { apiKey: refreshed.apiKey, credential: { type: "oauth", ...refreshed.newCredentials } }; }
/** Insert or replace a credential, matching by stable account identity. */ upsertCredential(provider: string, credential: AuthCredential): void { const identity = credentialIdentity(credential); const credentials = this.getCredentialList(provider); const existingIndex = credentials.findIndex((existing) => credentialIdentity(existing) === identity); if (existingIndex >= 0) { credentials[existingIndex] = credential; } else { credentials.push(credential); } this.setCredentialList(provider, credentials); this.clearStaleStoredCredential(provider, credential); this.persistProviderEntry(provider); this.clearCredentialBlock(provider, identity); }
/** Remove one credential by identity. Returns true when a credential was removed. */ removeCredential(provider: string, identity: string): boolean { const credentials = this.getCredentialList(provider); const next = credentials.filter((credential) => credentialIdentity(credential) !== identity); if (next.length === credentials.length) return false; this.setCredentialList(provider, next); this.persistProviderEntry(provider); this.clearCredentialBlock(provider, identity); this.resetRotation(provider); return true; }
/** Descriptive list of stored accounts for UI (never returns key/token bytes). */ listStoredCredentials(provider: string): StoredCredentialSummary[] { return this.getCredentialList(provider).map((credential) => ({ provider, type: credential.type, identity: credentialIdentity(credential), ...(credentialLabel(credential) ? { label: credentialLabel(credential) } : {}), })); }
private persistProviderEntry(provider: string): void { this.persistProviderChange(provider, this.data[provider]); }
private clearStaleStoredCredential(provider: string, credential: AuthCredential): void { const candidate = this.getStoredAuthCandidate(provider, credential); if (!candidate) return; const stale = this.staleAuthSources.get(provider); if (!stale) return; const next = stale.filter((token) => token.identityFingerprint !== candidate.identityFingerprint); if (next.length === 0) this.staleAuthSources.delete(provider); else this.staleAuthSources.set(provider, next); }
private resetRotation(provider: string): void { this.roundRobinIndex.delete(provider); for (const key of this.sessionCredentials.keys()) { if (key === provider || key.startsWith(`${provider}\0`)) { this.sessionCredentials.delete(key); } } }
/** * Credential order for selection: sessions start at a stable hashed index * (and stick to their last-used account), anonymous callers advance a * round-robin cursor so load spreads across accounts. */ private getCredentialOrder(provider: string, sessionId: string | undefined, total: number): number[] { if (total <= 1) return [0]; const pinned = sessionId ? this.sessionCredentials.get(`${provider}\0${sessionId}`) : undefined; let start: number; if (pinned) { start = this.identityIndex(provider, pinned, total) ?? 0; } else if (sessionId) { start = hashSessionId(sessionId) % total; } else { const current = this.roundRobinIndex.get(provider) ?? -1; const next = (current + 1) % total; this.roundRobinIndex.set(provider, next); start = next; } const order: number[] = []; for (let i = 0; i < total; i++) { order.push((start + i) % total); } return order; }
private identityIndex(provider: string, identity: string, total: number): number | undefined { const credentials = this.getCredentialList(provider); for (let i = 0; i < Math.min(total, credentials.length); i++) { if (credentialIdentity(credentials[i]) === identity) return i; } return undefined; }
private recordSessionCredential(provider: string, sessionId: string | undefined, identity: string): void { if (!sessionId) return; this.sessionCredentials.set(`${provider}\0${sessionId}`, identity); }
// ── Rate-limit blocks ─────────────────────────────────────────────────────
private loadPersistedBlocks(): void { if (!this.blockStorePath || !existsSync(this.blockStorePath)) return; try { const parsed = JSON.parse(readFileSync(this.blockStorePath, "utf-8")) as { blocks?: PersistedCredentialBlock[]; }; const nowMs = Date.now(); for (const block of parsed.blocks ?? []) { if ( typeof block.provider !== "string" || typeof block.identity !== "string" || typeof block.blockedUntilMs !== "number" || block.blockedUntilMs <= nowMs ) { continue; } const byIdentity = this.credentialBlocks.get(block.provider) ?? new Map(); byIdentity.set(block.identity, block); this.credentialBlocks.set(block.provider, byIdentity); } } catch { // Unreadable or malformed block store: start with no blocks. } }
private persistBlocks(): void { if (!this.blockStorePath) return; const nowMs = Date.now(); const blocks: PersistedCredentialBlock[] = []; for (const byIdentity of this.credentialBlocks.values()) { for (const block of byIdentity.values()) { if (block.blockedUntilMs > nowMs) blocks.push(block); } } try { this.storage.withLock(() => { mkdirSync(dirname(this.blockStorePath as string), { recursive: true, mode: 0o700 }); writeFileAtomicSync(this.blockStorePath as string, JSON.stringify({ blocks }, null, 2)); return { result: undefined }; }); } catch { // Persistence is best-effort; in-memory blocks still apply. } }
/** Block a credential (typically after a usage-limit response) until `untilMs`. */ markCredentialBlocked(provider: string, identity: string, blockedUntilMs: number, cause?: string): void { const nowMs = Date.now(); if (blockedUntilMs <= nowMs) return; const byIdentity = this.credentialBlocks.get(provider) ?? new Map(); const existing = byIdentity.get(identity); if (existing && existing.blockedUntilMs >= blockedUntilMs) return; byIdentity.set(identity, { provider, identity, blockedUntilMs, ...(cause ? { cause } : {}), }); this.credentialBlocks.set(provider, byIdentity); this.persistBlocks(); }
/** Clear a credential's rate-limit block (e.g. after a successful request). */ clearCredentialBlock(provider: string, identity: string): void { const byIdentity = this.credentialBlocks.get(provider); if (!byIdentity?.delete(identity)) return; if (byIdentity.size === 0) this.credentialBlocks.delete(provider); this.persistBlocks(); }
getCredentialBlockedUntil(provider: string, identity: string): number | undefined { const block = this.credentialBlocks.get(provider)?.get(identity); if (!block) return undefined; if (block.blockedUntilMs <= Date.now()) { this.clearCredentialBlock(provider, identity); return undefined; } return block.blockedUntilMs; }
/** * Block the credential a provider most recently used, as captured by the * last {@link AuthStorage.getApiKey} resolve. Used when a response arrives * outside the resolve path (stream failure diagnostics). */ blockLastSelectedCredential(provider: string, blockedUntilMs: number, cause?: string, sessionId?: string): boolean { const identity = sessionId ? this.sessionCredentials.get(`${provider}\0${sessionId}`) : this.lastSelectedIdentities.get(provider); if (!identity) return false; this.markCredentialBlocked(provider, identity, blockedUntilMs, cause); return true; }
/** Identity of the stored credential last handed out for a provider. */ getLastSelectedCredentialIdentity(provider: string): string | undefined { return this.lastSelectedIdentities.get(provider); }
private fingerprintAuthSource(source: ActiveAuthStatusSource, material: string): string { const digest = createHash("sha256").update(source).update("\0").update(material).digest("hex"); return `${source}:${digest}`; }
private createAuthSourceCandidate(options: { source: ActiveAuthStatusSource; configured: boolean; identityMaterial: string; valueMaterial?: string; label?: string; resolveValueMaterial?: () => string | undefined; }): AuthSourceCandidate { return { configured: options.configured, source: options.source, ...(options.label ? { label: options.label } : {}), identityFingerprint: this.fingerprintAuthSource(options.source, `identity:${options.identityMaterial}`), ...(options.valueMaterial !== undefined ? { valueFingerprint: this.fingerprintAuthSource( options.source, `value:${options.identityMaterial}\0${options.valueMaterial}`, ), } : {}), ...(options.resolveValueMaterial ? { resolveValueFingerprint: () => { const valueMaterial = options.resolveValueMaterial?.(); return valueMaterial === undefined ? undefined : this.fingerprintAuthSource( options.source, `value:${options.identityMaterial}\0${valueMaterial}`, ); }, } : {}), }; }
private getStoredCredentialValueMaterial(providerId: string, credential: AuthCredential): string | undefined { if (credential.type === "api_key") { if (credential.key.startsWith("!")) { const resolvedKey = resolveConfigValueUncached(credential.key); return resolvedKey === undefined ? undefined : `api_key:command:${credential.key}\0${resolvedKey}`; } return `api_key:${credential.key}\0${resolveConfigValue(credential.key) ?? ""}`; } // Static MCP tokens are not model-provider key material: they never // resolve to a provider API key value fingerprint. if (credential.type !== "oauth") return undefined; const provider = getOAuthProvider(providerId); const apiKey = provider?.getApiKey(credential) ?? credential.access; return `oauth:${apiKey}\0${credential.refresh}\0${credential.expires}`; }
private getRuntimeAuthCandidate(provider: string): AuthSourceCandidate | undefined { const apiKey = this.runtimeOverrides.get(provider); if (!apiKey) { return undefined; } return { label: "--api-key", ...this.createAuthSourceCandidate({ configured: false, source: "runtime", identityMaterial: provider, valueMaterial: apiKey, }), }; }
private getStoredAuthCandidate( provider: string, credential?: AuthCredential, options?: { resolveCommandValue?: boolean; resolvedCommandValue?: string }, ): AuthSourceCandidate | undefined { const target = credential ?? this.getCredentialList(provider)[0]; if (!target) { return undefined; } const isCommandApiKey = target.type === "api_key" && target.key.startsWith("!"); const identityMaterial = `${provider}:${target.type}:${credentialIdentity(target)}`; const commandValueMaterial = isCommandApiKey && options?.resolvedCommandValue !== undefined ? `api_key:command:${target.key}\0${options.resolvedCommandValue}` : undefined; return this.createAuthSourceCandidate({ configured: true, source: "stored", identityMaterial, valueMaterial: commandValueMaterial ?? (isCommandApiKey && !options?.resolveCommandValue ? undefined : this.getStoredCredentialValueMaterial(provider, target)), resolveValueMaterial: isCommandApiKey ? () => this.getStoredCredentialValueMaterial(provider, target) : undefined, }); }
private getEnvironmentAuthCandidate(provider: string): AuthSourceCandidate | undefined { const envKeys = findEnvKeys(provider); const envKey = envKeys?.[0]; const apiKey = getEnvApiKey(provider); if (!apiKey) { return undefined; } const label = envKey ?? "ambient credentials"; const identityMaterial = envKey ?? this.getAmbientEnvironmentIdentityMaterial(provider); return this.createAuthSourceCandidate({ configured: false, source: "environment", label, identityMaterial, valueMaterial: `${identityMaterial}\0${apiKey}`, }); }
private getAmbientEnvironmentIdentityMaterial(provider: string): string { if (provider === "amazon-bedrock") { if (process.env.AWS_PROFILE) return `amazon-bedrock:profile:${process.env.AWS_PROFILE}`; if (process.env.AWS_ACCESS_KEY_ID) { return `amazon-bedrock:access-key:${process.env.AWS_ACCESS_KEY_ID}:${process.env.AWS_SECRET_ACCESS_KEY ?? ""}:${process.env.AWS_SESSION_TOKEN ?? ""}`; } if (process.env.AWS_BEARER_TOKEN_BEDROCK) { return `amazon-bedrock:bearer:${process.env.AWS_BEARER_TOKEN_BEDROCK}`; } if (process.env.AWS_CONTAINER_CREDENTIALS_RELATIVE_URI) { return `amazon-bedrock:ecs-relative:${process.env.AWS_CONTAINER_CREDENTIALS_RELATIVE_URI}`; } if (process.env.AWS_CONTAINER_CREDENTIALS_FULL_URI) { return `amazon-bedrock:ecs-full:${process.env.AWS_CONTAINER_CREDENTIALS_FULL_URI}`; } if (process.env.AWS_WEB_IDENTITY_TOKEN_FILE) { return `amazon-bedrock:web-identity:${process.env.AWS_WEB_IDENTITY_TOKEN_FILE}`; } } if (provider === "google-vertex") { const project = process.env.GOOGLE_CLOUD_PROJECT ?? process.env.GCLOUD_PROJECT ?? ""; const location = process.env.GOOGLE_CLOUD_LOCATION ?? ""; const credentialsPath = process.env.GOOGLE_APPLICATION_CREDENTIALS ?? "application-default"; return `google-vertex:${project}:${location}:${credentialsPath}`; } return provider; }
private getFallbackAuthCandidate(provider: string): AuthSourceCandidate | undefined { const apiKey = this.fallbackResolver?.(provider); if (!apiKey) { return undefined; } return this.createAuthSourceCandidate({ configured: false, source: "fallback", label: "custom provider config", identityMaterial: provider, valueMaterial: apiKey, }); }
private getAuthSourceCandidates(provider: string, options?: { includeFallback?: boolean }): AuthSourceCandidate[] { const fallbackCandidate = options?.includeFallback === false ? undefined : this.getFallbackAuthCandidate(provider); // One stored candidate per account so each can be marked stale independently. const storedCandidates = this.getCredentialList(provider).map((credential) => this.getStoredAuthCandidate(provider, credential), ); const candidates = provider === PRIME_INFERENCE_PROVIDER_ID ? [ this.getRuntimeAuthCandidate(provider), this.getEnvironmentAuthCandidate(provider), this.getPrimeCliAuthCandidate(provider), ...storedCandidates, fallbackCandidate, ] : [ this.getRuntimeAuthCandidate(provider), ...storedCandidates, this.getEnvironmentAuthCandidate(provider), fallbackCandidate, ]; return candidates.filter((candidate): candidate is AuthSourceCandidate => candidate !== undefined); }
private isAuthSourceStale(provider: string, candidate: AuthSourceCandidate): boolean { const matchingStale = this.getMatchingStaleAuthSources(provider, candidate); if (matchingStale.length === 0) { return false; } const valueFingerprint = candidate.valueFingerprint ?? candidate.resolveValueFingerprint?.(); return Boolean(valueFingerprint && matchingStale.some((token) => token.valueFingerprint === valueFingerprint)); }
private getMatchingStaleAuthSources(provider: string, candidate: AuthSourceCandidate): AuthSourceToken[] { const stale = this.staleAuthSources.get(provider); if (!stale) { return []; } return stale.filter( (token) => token.source === candidate.source && token.identityFingerprint === candidate.identityFingerprint, ); }
private getAvailableAuthCandidate( provider: string, options?: { includeFallback?: boolean }, ): { candidate?: AuthSourceCandidate; hasStaleCandidate: boolean } { let hasStaleCandidate = false; for (const candidate of this.getAuthSourceCandidates(provider, options)) { if (this.isAuthSourceStale(provider, candidate)) { hasStaleCandidate = true; continue; } return { candidate, hasStaleCandidate }; } return { hasStaleCandidate }; }
private toAuthStatus(candidate: AuthSourceCandidate): AuthStatus { return { configured: candidate.configured, source: candidate.source, ...(candidate.label ? { label: candidate.label } : {}), }; }
private getAuthStatusFromCandidates(provider: string): AuthStatus { const { candidate, hasStaleCandidate } = this.getAvailableAuthCandidate(provider); if (candidate) { return this.toAuthStatus(candidate); } if (hasStaleCandidate) { return { configured: false, source: "stale", label: "expired" }; } return { configured: false }; }
markAuthStale(provider: string, sessionId?: string): boolean { const token = this.getCurrentAuthSourceToken(provider, sessionId); return token ? this.markAuthSourceStale(token) : false; }
private getAuthSourceTokenForCandidate( provider: string, candidate: AuthSourceCandidate, ): AuthSourceToken | undefined { const valueFingerprint = candidate.valueFingerprint ?? candidate.resolveValueFingerprint?.(); if (!valueFingerprint) { return undefined; } return { provider, source: candidate.source, identityFingerprint: candidate.identityFingerprint, valueFingerprint, }; }
getCurrentAuthSourceToken(provider: string, sessionId?: string): AuthSourceToken | undefined { // Prefer the credential actually handed to this session, so concurrent // failures cannot mark a sibling session's account stale. const lastIdentity = sessionId ? this.sessionCredentials.get(`${provider}\0${sessionId}`) : this.lastSelectedIdentities.get(provider); if (lastIdentity) { const credential = this.getCredentialList(provider).find( (candidate) => credentialIdentity(candidate) === lastIdentity, ); if (credential) { const candidate = this.getStoredAuthCandidate(provider, credential); if (candidate && !this.isAuthSourceStale(provider, candidate)) { const token = this.getAuthSourceTokenForCandidate(provider, candidate); if (token) return token; } } } if (sessionId) return undefined; const { candidate } = this.getAvailableAuthCandidate(provider); return candidate ? this.getAuthSourceTokenForCandidate(provider, candidate) : undefined; }
markAuthSourceStale(token: AuthSourceToken): boolean { if (token.provider.length === 0) { return false; } const stale = this.staleAuthSources.get(token.provider) ?? []; if ( !stale.some( (existing) => existing.source === token.source && existing.identityFingerprint === token.identityFingerprint && existing.valueFingerprint === token.valueFingerprint, ) ) { stale.push(token); } this.staleAuthSources.set(token.provider, stale); return true; }
/** Forget every stale marking for a provider (explicit user re-selection). */ clearAuthStale(provider: string): void { this.staleAuthSources.delete(provider); }
private clearStaleAuthSource(provider: string, source: ActiveAuthStatusSource): void { const stale = this.staleAuthSources.get(provider); if (!stale) { return; } const next = stale.filter((token) => token.source !== source); if (next.length === 0) { this.staleAuthSources.delete(provider); } else { this.staleAuthSources.set(provider, next); } }
private parseStorageData(content: string | undefined): AuthStorageData { if (!content) { return {}; } const data: unknown = JSON.parse(content); if (typeof data !== "object" || data === null || Array.isArray(data)) { throw new Error("Invalid auth storage: expected a JSON object"); } return data as AuthStorageData; }
/** * Reload credentials from storage. */ reload(): void { let content: string | undefined; try { this.storage.withLock((current) => { content = current; return { result: undefined }; }); this.data = this.parseStorageData(content); this.loadError = null; } catch (error) { this.loadError = error as Error; this.recordError(error); } }
private persistProviderChange(provider: string, credential: AuthCredentialEntry | undefined): void { if (this.loadError) { return; }
try { this.storage.withLock((current) => { const currentData = this.parseStorageData(current); const merged: AuthStorageData = { ...currentData }; if (credential) { merged[provider] = credential; } else { delete merged[provider]; } return { result: undefined, next: JSON.stringify(merged, null, 2) }; }); } catch (error) { this.recordError(error); } }
/** * Get the primary credential for a provider (first stored account). * Use {@link AuthStorage.getCredentialList} for all accounts. */ get(provider: string): AuthCredential | undefined { const entry = this.data[provider]; if (!entry) { return undefined; } return Array.isArray(entry) ? entry[0] : entry; }
/** * Replace the stored credentials for a provider with a single credential. * Use {@link AuthStorage.upsertCredential} to add an account without * removing existing ones. */ set(provider: string, credential: AuthCredential): void { this.clearStaleAuthSource(provider, "stored"); this.data[provider] = credential; this.resetRotation(provider); this.persistProviderChange(provider, credential); }
/** * Remove credential for a provider. */ remove(provider: string): void { this.clearStaleAuthSource(provider, "stored"); delete this.data[provider]; this.credentialBlocks.delete(provider); this.persistBlocks(); this.resetRotation(provider); this.persistProviderChange(provider, undefined); }
/** * Remove a provider's credential with the disk write verified: throws on any * load or write failure instead of recording it, so callers can refuse to * proceed while the credential may still exist on disk. Disk-authoritative * and idempotent — in-memory state is only updated after the write succeeds. * Returns whether a credential was actually removed from disk. */ removeVerified(provider: string): boolean { const removed = this.storage.withLock((current) => { const currentData = this.parseStorageData(current); if (!(provider in currentData)) return { result: false }; const merged: AuthStorageData = { ...currentData }; delete merged[provider]; return { result: true, next: JSON.stringify(merged, null, 2) }; }); if (removed) { delete this.data[provider]; // Post-success only: a failed removal must not make a stale-marked credential selectable again. this.clearStaleAuthSource(provider, "stored"); } return removed; }
/** * Disk-authoritative conditional move for staged MCP logins: move * `stagedProvider`'s credential to `provider` ONLY when no credential * exists at `provider` ON DISK, reading and writing under the backend's * own file lock. An ordinary login in another process — invisible to this * instance's cache — can never be clobbered by a race between the get and * the set. Returns "occupied" when the destination already holds a * credential, "nothing" when the staged slot is empty, or the exact * credential that moved (for exact-own rollback). */ moveStagedCredential( stagedProvider: string, provider: string, ): { status: "occupied" } | { status: "nothing" } | { status: "moved"; credential: AuthCredential } { type MoveOutcome = | { status: "occupied" } | { status: "nothing" } | { status: "moved"; credential: AuthCredential }; const outcome = this.storage.withLock<MoveOutcome>((current) => { const currentData = this.parseStorageData(current); if (provider in currentData) { return { result: { status: "occupied" } }; } const staged = firstCredential(currentData[stagedProvider]); if (!staged) { return { result: { status: "nothing" } }; } const merged: AuthStorageData = { ...currentData, [provider]: staged }; delete merged[stagedProvider]; return { result: { status: "moved", credential: staged }, next: JSON.stringify(merged, null, 2), }; }); // Post-success only: refresh the cache from disk under the same lock // discipline so no stale entry survives the move. if (outcome.status === "moved") { this.reload(); } return outcome; }
/** * Disk-authoritative conditional restore: write `credential` to * `provider` ONLY when the slot is empty ON DISK. A credential written by * anyone else is never overwritten. */ restoreCredentialIfAbsent(provider: string, credential: AuthCredential): boolean { const restored = this.storage.withLock((current) => { const currentData = this.parseStorageData(current); if (provider in currentData) { return { result: false }; } const merged: AuthStorageData = { ...currentData, [provider]: credential }; return { result: true, next: JSON.stringify(merged, null, 2) }; }); if (restored) { this.data[provider] = credential; this.clearStaleAuthSource(provider, "stored"); } return restored; }
/** * Disk-authoritative conditional removal: remove `provider`'s credential * ONLY when the ON-DISK value is exactly `expected` (full-object * comparison, not token equality) — a credential written by anyone else is * never deleted. The in-memory cache drops the key only after the write * succeeds. */ removeIfCredentialMatches(provider: string, expected: AuthCredential): boolean { const removed = this.storage.withLock((current) => { const currentData = this.parseStorageData(current); const currentCredential = firstCredential(currentData[provider]); if (currentCredential === undefined) { return { result: false }; } if (JSON.stringify(currentCredential) !== JSON.stringify(expected)) { return { result: false }; } const merged: AuthStorageData = { ...currentData }; delete merged[provider]; return { result: true, next: JSON.stringify(merged, null, 2) }; }); if (removed) { delete this.data[provider]; this.clearStaleAuthSource(provider, "stored"); } return removed; }
/** * Disk-authoritative credential read under the backend's own file lock — * a cross-instance writer is always visible, unlike the cached `get()`. * Used to capture the full identity a guarded login's compare-and-swap * expects to replace (legacy/credential-only accounts included). */ getVerified(provider: string): AuthCredential | undefined { return this.storage.withLock((current) => { const currentData = this.parseStorageData(current); return { result: firstCredential(currentData[provider]) }; }); }
/** * Atomic full-identity compare-and-swap move for guarded MCP logins: * move `stagedProvider`'s credential to `provider` ONLY when the on-disk * value at `provider` is exactly `expectedOld` — the comparison INCLUDES * absence (both present, or both absent) — read and written under the * backend's own file lock. A changed OR deleted grant refuses ("occupied"): * a logged-out account is never reactivated and a newer writer is never * clobbered. Returns the exact credential that moved (for full-identity * rollback). */ replaceStagedCredential( stagedProvider: string, provider: string, expectedOld: AuthCredential | undefined, ): { status: "occupied" } | { status: "nothing" } | { status: "replaced"; credential: AuthCredential } { type ReplaceOutcome = | { status: "occupied" } | { status: "nothing" } | { status: "replaced"; credential: AuthCredential }; const outcome = this.storage.withLock<ReplaceOutcome>((current) => { const currentData = this.parseStorageData(current); const existing = firstCredential(currentData[provider]); // FULL-IDENTITY comparison INCLUDING absence: the on-disk value must // be exactly `expectedOld` (both present, or both absent). A // changed OR deleted grant refuses — never reactivate a logged-out // account, never clobber a newer writer. if (JSON.stringify(existing) !== JSON.stringify(expectedOld)) { return { result: { status: "occupied" as const } }; } const staged = firstCredential(currentData[stagedProvider]); if (!staged) { return { result: { status: "nothing" as const } }; } const merged: AuthStorageData = { ...currentData, [provider]: staged }; delete merged[stagedProvider]; return { result: { status: "replaced" as const, credential: staged }, next: JSON.stringify(merged, null, 2), }; }); // Post-success only: refresh the cache under the same lock discipline. if (outcome.status === "replaced") { this.reload(); } return outcome; }
/** * Atomic full-identity compare-and-swap write: set `provider` to `next` * ONLY when the on-disk value is exactly `expected`. Used to roll back a * guarded replacement (restoring the PREVIOUS credential) and to undo * only this attempt's own write — a newer writer is never clobbered. */ replaceCredentialIfMatches(provider: string, expected: AuthCredential, next: AuthCredential): boolean { const replaced = this.storage.withLock((current) => { const currentData = this.parseStorageData(current); const existing = firstCredential(currentData[provider]); if (existing === undefined || JSON.stringify(existing) !== JSON.stringify(expected)) { return { result: false }; } const merged: AuthStorageData = { ...currentData, [provider]: next }; return { result: true, next: JSON.stringify(merged, null, 2) }; }); if (replaced) { this.data[provider] = next; this.clearStaleAuthSource(provider, "stored"); } return replaced; }
/** * List all providers with credentials. */ list(): string[] { return Object.keys(this.data); }
/** * Check if credentials exist for a provider in auth.json. */ has(provider: string): boolean { return provider in this.data; }
/** * Check if any form of auth is configured for a provider. * Unlike getApiKey(), this doesn't refresh OAuth tokens. */ hasAuth(provider: string): boolean { return this.getAvailableAuthCandidate(provider).candidate !== undefined; }
/** * Return auth status without exposing credential values or refreshing tokens. */ getAuthStatus(provider: string): AuthStatus { return this.getAuthStatusFromCandidates(provider); }
/** * Get all credentials (for passing to getOAuthApiKey). */ getAll(): AuthStorageData { return { ...this.data }; }
/** Number of stored accounts for a provider, optionally limited by credential type. */ getCredentialCount(provider: string, type?: AuthCredential["type"]): number { const credentials = this.getCredentialList(provider); return type ? credentials.filter((credential) => credential.type === type).length : credentials.length; }
/** Whether a stored credential is selectable without waiting for a rate-limit block. */ hasUnblockedCredential(provider: string): boolean { return this.getCredentialList(provider).some((credential) => { const candidate = this.getStoredAuthCandidate(provider, credential); return ( candidate !== undefined && !this.isAuthSourceStale(provider, candidate) && this.getCredentialBlockedUntil(provider, credentialIdentity(credential)) === undefined ); }); }
/** * Whether stored credentials exist for the provider and every one of them is * rate-limit blocked right now. * * False when nothing is stored: environment-provided and runtime-key auth carry no * stored identity to block, so they are never refused on this basis. */ hasOnlyBlockedCredentials(provider: string): boolean { return this.getCredentialList(provider).length > 0 && !this.hasUnblockedCredential(provider); }
/** * Providers whose stored credentials are all rate-limit blocked right now. * * The block store is the durable half of the quota state: the in-memory breaker is * deliberately not persisted, so without this a restart advertises a depleted * provider as available again. */ blockedProviders(): string[] { return [...this.credentialBlocks.keys()].filter((provider) => this.hasOnlyBlockedCredentials(provider)); }
/** Soonest moment any blocked credential for the provider becomes selectable again. */ nextCredentialUnblockMs(provider: string): number | undefined { const nowMs = Date.now(); let soonest: number | undefined; for (const block of this.credentialBlocks.get(provider)?.values() ?? []) { if (block.blockedUntilMs <= nowMs) continue; if (soonest === undefined || block.blockedUntilMs < soonest) soonest = block.blockedUntilMs; } return soonest; }
drainErrors(): Error[] { const drained = [...this.errors]; this.errors = []; return drained; }
/** * Login to an OAuth provider. Logging in again with a different account * adds it to the pool; re-authenticating the same account replaces it. */ async login(providerId: OAuthProviderId, callbacks: OAuthLoginCallbacks): Promise<void> { const provider = getOAuthProvider(providerId); if (!provider) { throw new Error(`Unknown OAuth provider: ${providerId}`); }
const credentials = await provider.login(callbacks); this.upsertCredential(providerId, { type: "oauth", ...credentials }); }
/** * Logout from a provider. */ logout(provider: string): void { if (provider === PRIME_INFERENCE_PROVIDER_ID) { this.removeVerified(provider); return; } this.remove(provider); }
/** * Logout a single stored account, identified by {@link credentialIdentity}. * Returns true when the account existed and was removed. */ logoutCredential(provider: string, identity: string): boolean { if (provider === PRIME_INFERENCE_PROVIDER_ID && this.isPrimeCliConfigEnabled()) { // Prime CLI credentials are single-key; fall back to a full logout. this.logout(provider); return true; } return this.removeCredential(provider, identity); }
/** * Refresh OAuth token with backend locking to prevent race conditions. * Multiple pi instances may try to refresh simultaneously when tokens expire. */ private async refreshOAuthTokenWithLock( providerId: OAuthProviderId, credentialIndex = 0, ): Promise<{ apiKey: string; newCredentials: OAuthCredentials } | null> { const provider = getOAuthProvider(providerId); if (!provider) { return null; }
const result = await this.storage.withLockAsync(async (current) => { const currentData = this.parseStorageData(current); this.data = currentData; this.loadError = null;
const entry = currentData[providerId]; const credentials = entry ? (Array.isArray(entry) ? entry : [entry]) : []; const cred = credentials[credentialIndex]; if (cred?.type !== "oauth") { return { result: null }; }
if (Date.now() < cred.expires) { return { result: { apiKey: provider.getApiKey(cred), newCredentials: cred } }; }
const refreshed = await getOAuthApiKey(providerId, { [providerId]: cred }); if (!refreshed) { return { result: null }; }
credentials[credentialIndex] = { type: "oauth", ...refreshed.newCredentials }; const merged: AuthStorageData = { ...currentData, [providerId]: credentials.length === 1 ? credentials[0] : credentials, }; this.data = merged; this.loadError = null; return { result: refreshed, next: JSON.stringify(merged, null, 2) }; });
return result; }
/** * Get API key for a provider. * Priority: * 1. Runtime override (CLI --api-key) * 2. Prime Inference: environment variable, auth.json * 3. Other providers: auth.json, environment variable * 4. Fallback resolver (models.json custom providers) */ async getApiKeyWithSourceToken( providerId: string, options?: { includeFallback?: boolean; sessionId?: string }, ): Promise<AuthApiKeyResult> { // Runtime overrides take precedence over stored credentials and environment keys. const runtimeCandidate = this.getRuntimeAuthCandidate(providerId); const runtimeKey = this.runtimeOverrides.get(providerId); if (runtimeKey && runtimeCandidate && !this.isAuthSourceStale(providerId, runtimeCandidate)) { return { apiKey: runtimeKey, sourceToken: this.getAuthSourceTokenForCandidate(providerId, runtimeCandidate), }; }
const envCandidate = this.getEnvironmentAuthCandidate(providerId); const envKey = getEnvApiKey(providerId); if ( providerId === PRIME_INFERENCE_PROVIDER_ID && envKey && envCandidate && !this.isAuthSourceStale(providerId, envCandidate) ) { return { apiKey: envKey, sourceToken: this.getAuthSourceTokenForCandidate(providerId, envCandidate), }; }
const stored = await this.selectStoredApiKey(providerId, options?.sessionId); if (stored) { return stored; }
if (providerId === PRIME_INFERENCE_PROVIDER_ID) { const primeCliCandidate = this.getPrimeCliAuthCandidate(providerId); const primeCliKey = this.getPrimeCliApiKey(providerId); if (primeCliKey && primeCliCandidate && !this.isAuthSourceStale(providerId, primeCliCandidate)) { return { apiKey: primeCliKey, sourceToken: this.getAuthSourceTokenForCandidate(providerId, primeCliCandidate), }; } } // A provider whose live stored credentials could not supply a key fails closed: // silently falling through to an ambient environment key would charge a different // credential after a failed refresh. Accounts the user or a failure already marked // stale are the deliberate exception, and an explicit login reopens this path. const hasLiveStoredCredential = this.getCredentialList(providerId).some((credential) => { const candidate = this.getStoredAuthCandidate(providerId, credential); return candidate !== undefined && !this.isAuthSourceStale(providerId, candidate); }); if (hasLiveStoredCredential) { return {}; } // Stored auth wins over environment variables for non-Prime-Inference providers. if ( providerId !== PRIME_INFERENCE_PROVIDER_ID && envKey && envCandidate && !this.isAuthSourceStale(providerId, envCandidate) ) { return { apiKey: envKey, sourceToken: this.getAuthSourceTokenForCandidate(providerId, envCandidate), }; } if (options?.includeFallback !== false) { const fallbackCandidate = this.getFallbackAuthCandidate(providerId); if (fallbackCandidate && !this.isAuthSourceStale(providerId, fallbackCandidate)) { return { apiKey: this.fallbackResolver?.(providerId) ?? undefined, sourceToken: this.getAuthSourceTokenForCandidate(providerId, fallbackCandidate), }; } }
return {}; }
async getApiKey( providerId: string, options?: { includeFallback?: boolean; sessionId?: string }, ): Promise<string | undefined> { const result = await this.getApiKeyWithSourceToken(providerId, options); return result.apiKey; }
/** * Pick a usable stored credential, rotating across accounts: * - credentials are visited in session-sticky or round-robin order, * - rate-limit-blocked credentials are deferred behind unblocked siblings, * - stale (failed-auth) credentials are skipped entirely, * - OAuth credentials are selected ahead of stored API keys. * When every credential of a kind is blocked, the first in rotation order is * still returned so the request is attempted rather than dropped. */ private async selectStoredApiKey(provider: string, sessionId?: string): Promise<AuthApiKeyResult | undefined> { const credentials = this.getCredentialList(provider); if (credentials.length === 0) { return undefined; } const order = this.getCredentialOrder(provider, sessionId, credentials.length); const oauthOrder = this.deferBlockedCredentials( provider, credentials, order.filter((index) => credentials[index]?.type === "oauth"), ); for (const index of oauthOrder) { const resolved = await this.resolveStoredOAuthCredential(provider, index, sessionId); if (resolved) { return resolved; } }
const apiKeyOrder = this.deferBlockedCredentials( provider, credentials, order.filter((index) => credentials[index]?.type === "api_key"), ); for (const index of apiKeyOrder) { const credential = credentials[index]; if (credential?.type !== "api_key") { continue; } const storedCandidate = this.getStoredAuthCandidate(provider, credential); if (!storedCandidate || this.isAuthSourceStale(provider, storedCandidate)) { continue; } const hasStaleRecord = this.getMatchingStaleAuthSources(provider, storedCandidate).length > 0; const apiKey = credential.key.startsWith("!") && hasStaleRecord ? resolveConfigValueUncached(credential.key) : resolveConfigValue(credential.key); if (apiKey === undefined) { continue; } const sourceToken = this.getAuthSourceTokenForCandidate( provider, credential.key.startsWith("!") ? (this.getStoredAuthCandidate(provider, credential, { resolvedCommandValue: apiKey }) ?? storedCandidate) : storedCandidate, ); this.recordSelection(provider, credential, sessionId); return { apiKey, sourceToken, credentialType: "api_key" }; } return undefined; }
private async resolveStoredOAuthCredential( provider: string, index: number, sessionId?: string, ): Promise<AuthApiKeyResult | undefined> { const credential = this.getCredentialList(provider)[index]; if (credential?.type !== "oauth") { return undefined; } const storedCandidate = this.getStoredAuthCandidate(provider, credential); if (!storedCandidate || this.isAuthSourceStale(provider, storedCandidate)) { return undefined; } const oauthProvider = getOAuthProvider(provider); if (!oauthProvider) { return undefined; } const identity = credentialIdentity(credential);
if (Date.now() >= credential.expires) { // Lock refreshes so concurrent instances cannot race on the credential file. try { const result = await this.refreshOAuthTokenWithLock(provider, index); if (!result) { return undefined; } const refreshedCandidate = this.getStoredAuthCandidate(provider, this.getCredentialList(provider)[index]) ?? storedCandidate; this.recordSelection(provider, this.getCredentialList(provider)[index] ?? credential, sessionId); return { apiKey: result.apiKey, credentialType: "oauth", sourceToken: this.getAuthSourceTokenForCandidate(provider, refreshedCandidate), }; } catch (error) { this.recordError(error); // A peer may have refreshed successfully; reload before treating this refresh as failed. this.reload(); const updated = this.getCredentialList(provider).find( (candidate) => candidate.type === "oauth" && credentialIdentity(candidate) === identity, ); if (updated?.type === "oauth" && Date.now() < updated.expires) { this.recordSelection(provider, updated, sessionId); return { apiKey: oauthProvider.getApiKey(updated), credentialType: "oauth", sourceToken: this.getAuthSourceTokenForCandidate( provider, this.getStoredAuthCandidate(provider, updated) ?? storedCandidate, ), }; } // Move on to a sibling account; credentials stay for a later /login retry. return undefined; } }
this.recordSelection(provider, credential, sessionId); return { apiKey: oauthProvider.getApiKey(credential), credentialType: "oauth", sourceToken: this.getAuthSourceTokenForCandidate(provider, storedCandidate), }; }
private recordSelection(provider: string, credential: AuthCredential, sessionId: string | undefined): void { const identity = credentialIdentity(credential); this.lastSelectedIdentities.set(provider, identity); this.recordSessionCredential(provider, sessionId, identity); }
/** Move blocked credentials behind unblocked siblings, preserving relative order. */ private deferBlockedCredentials(provider: string, credentials: AuthCredential[], indexes: number[]): number[] { if (indexes.length <= 1) { return indexes; } const unblocked = indexes.filter((index) => { const credential = credentials[index]; return !credential || this.getCredentialBlockedUntil(provider, credentialIdentity(credential)) === undefined; }); if (unblocked.length === 0) { return indexes; } if (unblocked.length === indexes.length) { return indexes; } const blocked = indexes.filter((index) => !unblocked.includes(index)); return [...unblocked, ...blocked]; }
/** * Get all registered OAuth providers */ getOAuthProviders() { return getOAuthProviders(); }
private updatePrimeInferenceCredential( update: (credential: AuthCredential | undefined) => ApiKeyCredential | undefined, ): void { try { const data = this.storage.withLock((current) => { const data = this.parseStorageData(current); const credential = update(firstCredential(data[PRIME_INFERENCE_PROVIDER_ID])); if (!credential) return { result: data }; data[PRIME_INFERENCE_PROVIDER_ID] = credential; return { result: data, next: JSON.stringify(data, null, 2) }; }); this.data = data; this.loadError = null; } catch (error) { this.recordError(error); throw error; } }
setPrimeInferenceTeamSelection(team: PrimeTeam | null, expectedApiKey?: string): void { if (this.isPrimeCliConfigEnabled()) { try { savePrimeCliTeamSelection(team, this.getEnabledPrimeCliConfigPath()); } catch (error) { this.recordError(error); throw error; } return; }
this.updatePrimeInferenceCredential((credential) => credential?.type === "api_key" && (expectedApiKey === undefined || credential.key === expectedApiKey) ? { ...credential, primeTeam: team ? this.toPrimeTeamCredential(team) : null } : undefined, ); }
setPrimeInferenceApiKey(apiKey: string, team?: PrimeTeam | null): void { if (this.isPrimeCliConfigEnabled()) { try { const configPath = this.getEnabledPrimeCliConfigPath(); const config = loadPrimeCliConfig(configPath); const existingCredential = firstCredential(this.data[PRIME_INFERENCE_PROVIDER_ID]); const legacyPrimeTeam = existingCredential?.type === "api_key" ? existingCredential.primeTeam : undefined; if (config.apiKey !== apiKey) { savePrimeCliApiKey(apiKey, configPath); } else if (!config.teamIdFromEnv && (legacyPrimeTeam === null || (!config.teamId && legacyPrimeTeam))) { savePrimeCliTeamSelection(legacyPrimeTeam, configPath); } this.clearStaleAuthSource(PRIME_INFERENCE_PROVIDER_ID, "prime_cli"); } catch (error) { this.recordError(error); throw error; } try { const data = this.storage.withLock((current) => { const data = this.parseStorageData(current); delete data[PRIME_INFERENCE_PROVIDER_ID]; return { result: data, next: JSON.stringify(data, null, 2) }; }); this.data = data; this.loadError = null; } catch (error) { this.recordError(error); throw error; } return; }
this.updatePrimeInferenceCredential((existing) => ({ type: "api_key", key: apiKey, primeTeam: team !== undefined ? team ? this.toPrimeTeamCredential(team) : null : existing?.type === "api_key" && existing.key === apiKey ? (existing.primeTeam ?? null) : null, })); this.clearStaleAuthSource(PRIME_INFERENCE_PROVIDER_ID, "stored"); }
getPrimeInferenceTeamSelection(): PrimeTeamCredential | null | undefined { if (process.env.PRIME_TEAM_ID?.trim()) return undefined; let config: PrimeCliConfig | undefined; if (this.isPrimeCliConfigEnabled()) { config = this.getPrimeCliConfig(PRIME_INFERENCE_PROVIDER_ID); if (config?.teamIdFromEnv) { return undefined; } } const authSource = this.getAuthStatus(PRIME_INFERENCE_PROVIDER_ID).source; if (authSource === "runtime" || authSource === "environment") return undefined; const credential = firstCredential(this.data[PRIME_INFERENCE_PROVIDER_ID]); return credential?.type === "api_key" ? credential.primeTeam : undefined; }
getProviderHeaders(providerId: string): Record<string, string> | undefined { if (providerId !== PRIME_INFERENCE_PROVIDER_ID) return undefined; const teamId = process.env.PRIME_TEAM_ID?.trim() || this.getPrimeInferenceTeamSelection()?.teamId; return teamId ? { "X-Prime-Team-ID": teamId } : undefined; }
getPrimeCliConfigPath(): string | undefined { if (!this.isPrimeCliConfigEnabled()) { return undefined; } return getPrimeCliConfigPath(this.options.primeCliConfigPath); }
private toPrimeTeamCredential(team: PrimeTeam): PrimeTeamCredential { const credential: PrimeTeamCredential = { teamId: team.teamId, name: team.name, }; if (team.slug) { credential.slug = team.slug; } if (team.role) { credential.role = team.role; } if (team.createdAt) { credential.createdAt = team.createdAt; } return credential; }
/** * The stored Prime CLI config for a provider, or undefined when the provider is * not Prime Inference or the config file is not enabled for this storage. */ private getPrimeCliConfig(providerId: string): PrimeCliConfig | undefined { if (providerId !== PRIME_INFERENCE_PROVIDER_ID) { return undefined; } if (!this.isPrimeCliConfigEnabled()) { return undefined; } try { return loadPrimeCliConfig(this.options.primeCliConfigPath); } catch (error) { this.recordError(error); return undefined; } }
private getPrimeCliApiKey(providerId: string): string | undefined { return this.getPrimeCliConfig(providerId)?.apiKey; }
private getEnabledPrimeCliConfigPath(): string { const configPath = this.getPrimeCliConfigPath(); if (!configPath) { throw new Error("Prime CLI config is not enabled"); } return configPath; }
private getPrimeCliAuthCandidate(provider: string): AuthSourceCandidate | undefined { const apiKey = this.getPrimeCliApiKey(provider); if (!apiKey) { return undefined; } return { label: "Prime CLI", ...this.createAuthSourceCandidate({ configured: false, source: "prime_cli", identityMaterial: provider, valueMaterial: apiKey, }), }; }
private isPrimeCliConfigEnabled(): boolean { return Boolean(this.options.usePrimeCliConfig || this.options.primeCliConfigPath); }}