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.
54 kB · 1639 lines
TypeScript
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640/** * 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 { 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 { clearPrimeCliCredentials, 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;
export type AuthCredential = ApiKeyCredential | OAuthCredential;
/** * 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;};
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);}
/** * 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. */export function credentialIdentity(credential: AuthCredential): string { if (credential.type === "api_key") { return credential.key.startsWith("!") ? `command:${credential.key}` : `key:${shortDigest(credential.key)}`; } const accountId = (credential as { accountId?: unknown }).accountId; if (typeof accountId === "string" && accountId.length > 0) { 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; } const email = (credential as { email?: unknown }).email; if (typeof email === "string" && email.length > 0) return email; const accountId = (credential as { accountId?: unknown }).accountId; if (typeof accountId === "string" && accountId.length > 0) { return `account ${accountId.slice(0, 8)}`; } return 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 }; 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): boolean { const identity = 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) ?? ""}`; } 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 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 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): boolean { const token = this.getCurrentAuthSourceToken(provider); 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): AuthSourceToken | undefined { // Prefer the credential actually handed out last, so callers can mark // the failing account stale instead of whichever sorts first. const lastIdentity = 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; } } } } const { candidate } = this.getAvailableAuthCandidate(provider); if (!candidate) { return undefined; } return this.getAuthSourceTokenForCandidate(provider, candidate); }
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 {}; } return JSON.parse(content) 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. */ removeVerified(provider: string): void { this.storage.withLock((current) => { const currentData = this.parseStorageData(current); if (!(provider in currentData)) return { result: undefined }; const merged: AuthStorageData = { ...currentData }; delete merged[provider]; return { result: undefined, next: JSON.stringify(merged, null, 2) }; }); delete this.data[provider]; // Post-success only: a failed removal must not make a stale-marked credential selectable again. this.clearStaleAuthSource(provider, "stored"); }
/** * 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; }
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.isPrimeCliConfigEnabled()) { try { clearPrimeCliCredentials(this.getEnabledPrimeCliConfigPath()); this.clearStaleAuthSource(provider, "prime_cli"); } catch (error) { this.recordError(error); throw error; } } 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, Prime CLI config, 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), }; }
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), }; } }
const stored = await this.selectStoredApiKey(providerId, options?.sessionId); if (stored) { return stored; } // 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 }; } 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, 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), 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), 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(); }
setPrimeInferenceTeamSelection(team: PrimeTeam | null): void { if (this.isPrimeCliConfigEnabled()) { try { savePrimeCliTeamSelection(team, this.getEnabledPrimeCliConfigPath()); } catch (error) { this.recordError(error); throw error; } return; }
const credential = this.get(PRIME_INFERENCE_PROVIDER_ID); if (credential?.type !== "api_key") { return; } this.set(PRIME_INFERENCE_PROVIDER_ID, { ...credential, primeTeam: team ? this.toPrimeTeamCredential(team) : null, }); }
setPrimeInferenceApiKey(apiKey: string): void { if (this.isPrimeCliConfigEnabled()) { try { const configPath = this.getEnabledPrimeCliConfigPath(); const config = loadPrimeCliConfig(configPath); const existingCredential = this.get(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; } if (this.has(PRIME_INFERENCE_PROVIDER_ID)) { this.remove(PRIME_INFERENCE_PROVIDER_ID); } return; }
const existingCredential = this.get(PRIME_INFERENCE_PROVIDER_ID); const existingPrimeTeam = existingCredential?.type === "api_key" ? existingCredential.primeTeam : undefined; this.set(PRIME_INFERENCE_PROVIDER_ID, { type: "api_key", key: apiKey, ...(existingPrimeTeam !== undefined ? { primeTeam: existingPrimeTeam } : {}), }); }
getPrimeInferenceTeamSelection(): PrimeTeamCredential | null | undefined { let config: PrimeCliConfig | undefined; if (this.isPrimeCliConfigEnabled()) { config = this.getPrimeCliConfig(PRIME_INFERENCE_PROVIDER_ID); if (config?.teamIdFromEnv) { return undefined; } }
const credential = this.get(PRIME_INFERENCE_PROVIDER_ID); const authSource = this.getAuthStatus(PRIME_INFERENCE_PROVIDER_ID).source; if (authSource === "runtime" || authSource === "environment") { return undefined; } // A stale CLI key must not erase the selected team used to validate cached model access. if (authSource === "prime_cli" || (authSource === "stale" && config?.apiKey)) { if (credential?.type === "api_key" && credential.primeTeam === null) { return null; } if (config?.teamId) { return this.toPrimeTeamCredential({ teamId: config.teamId, name: config.teamName ?? "Prime CLI team", ...(config.teamRole ? { role: config.teamRole } : {}), }); } if (credential?.type === "api_key" && credential.primeTeam) { return credential.primeTeam; } return null; } if (credential?.type === "api_key" && credential.primeTeam !== undefined) { return credential.primeTeam; } if (!config?.apiKey && config?.teamId) { return this.toPrimeTeamCredential({ teamId: config.teamId, name: config.teamName ?? "Prime CLI team", ...(config.teamRole ? { role: config.teamRole } : {}), }); } return undefined; }
getProviderHeaders(providerId: string): Record<string, string> | undefined { if (providerId !== PRIME_INFERENCE_PROVIDER_ID) { return undefined; }
const primeCliConfig = this.getPrimeCliConfig(providerId); if (primeCliConfig?.teamIdFromEnv) { return primeCliConfig.teamId ? { "X-Prime-Team-ID": primeCliConfig.teamId } : undefined; }
const teamId = 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; }
private getPrimeCliConfig(providerId: string): PrimeCliConfig | undefined { if (providerId !== PRIME_INFERENCE_PROVIDER_ID) { return undefined; } if (!this.isPrimeCliConfigEnabled()) { return undefined; } return loadPrimeCliConfig(this.options.primeCliConfigPath); }
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 isPrimeCliConfigEnabled(): boolean { return Boolean(this.options.usePrimeCliConfig || this.options.primeCliConfigPath); }}