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.
84 kB · 2309 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310/** * Model registry - manages built-in and custom models, provides API key resolution. */
import { Buffer } from "node:buffer";import { createHash, createHmac } from "node:crypto";import { type AnthropicMessagesCompat, type Api, type AssistantMessageEventStream, type Context, type CredentialKind, fetchModelsDevModels, getLogger, getModels, getProviders, type KnownProvider, type Model, type OAuthProviderInterface, type OpenAICompletionsCompat, type OpenAIResponsesCompat, openCodeSessionHeaders, parsePrimeInferenceModelCatalog, registerApiProvider, registerLiveModels, registerUsageFeed, resetApiProviders, type SimpleStreamOptions, type UsageResponse, type UsageState, type UsageStep, unregisterUsageFeed,} from "@earendil-works/pi-ai";import { registerBuiltinMcpOAuthProviders } from "@earendil-works/pi-ai/mcp";import { getXaiSubscriptionModel, registerOAuthProvider, resetOAuthProviders } from "@earendil-works/pi-ai/oauth";import { existsSync, readFileSync } from "fs";import { dirname, join } from "path";import { type Static, type TProperties, Type } from "typebox";import type { Validator } from "typebox/compile";import type { TLocalizedValidationError } from "typebox/error";import { getAgentDir } from "../config.js";import { writeFileAtomicSync } from "../utils/atomic-file.js";import { type AuthSourceToken, type AuthStatus, type AuthStorage, DEFAULT_RATE_LIMIT_BLOCK_MS,} from "./auth-storage.js";import { PRIME_INFERENCE_PROVIDER_ID } from "./prime-inference-auth.js";import { buildPrimeInferenceModels, mergePrimeInferenceModels, readCachedPrimeInferenceModels, refreshPrimeInferenceModels,} from "./prime-inference-model-catalog.js";import { fetchAuthorizedPrivatePrimeInferenceModels, getPrivatePrimeInferenceModels, isPrivatePrimeInferenceModel,} from "./prime-inference-models.js";import { calculatePromptPayloadCeiling } from "./prompts/payload-ceiling.js";import { BUILT_IN_PROVIDER_DISPLAY_NAMES } from "./provider-display-names.js";import { resolveConfigValueOrThrow, resolveConfigValueUncached, resolveHeadersOrThrow,} from "./resolve-config-value.js";
const MODELS_DEV_TTL_MS = 60 * 60 * 1000;const log = getLogger("coding-agent.model-registry");
const PercentileCutoffsSchema = Type.Object({ p50: Type.Optional(Type.Number()), p75: Type.Optional(Type.Number()), p90: Type.Optional(Type.Number()), p99: Type.Optional(Type.Number()),});
const OpenRouterRoutingSchema = Type.Object({ allow_fallbacks: Type.Optional(Type.Boolean()), require_parameters: Type.Optional(Type.Boolean()), data_collection: Type.Optional(Type.Union([Type.Literal("deny"), Type.Literal("allow")])), zdr: Type.Optional(Type.Boolean()), enforce_distillable_text: Type.Optional(Type.Boolean()), order: Type.Optional(Type.Array(Type.String())), only: Type.Optional(Type.Array(Type.String())), ignore: Type.Optional(Type.Array(Type.String())), quantizations: Type.Optional(Type.Array(Type.String())), sort: Type.Optional( Type.Union([ Type.String(), Type.Object({ by: Type.Optional(Type.String()), partition: Type.Optional(Type.Union([Type.String(), Type.Null()])), }), ]), ), max_price: Type.Optional( Type.Object({ prompt: Type.Optional(Type.Union([Type.Number(), Type.String()])), completion: Type.Optional(Type.Union([Type.Number(), Type.String()])), image: Type.Optional(Type.Union([Type.Number(), Type.String()])), audio: Type.Optional(Type.Union([Type.Number(), Type.String()])), request: Type.Optional(Type.Union([Type.Number(), Type.String()])), }), ), preferred_min_throughput: Type.Optional(Type.Union([Type.Number(), PercentileCutoffsSchema])), preferred_max_latency: Type.Optional(Type.Union([Type.Number(), PercentileCutoffsSchema])),});
const VercelGatewayRoutingSchema = Type.Object({ only: Type.Optional(Type.Array(Type.String())), order: Type.Optional(Type.Array(Type.String())),});
const ThinkingLevelMapValueSchema = Type.Union([Type.String(), Type.Null()]);const ThinkingLevelMapSchema = Type.Object({ off: Type.Optional(ThinkingLevelMapValueSchema), minimal: Type.Optional(ThinkingLevelMapValueSchema), low: Type.Optional(ThinkingLevelMapValueSchema), medium: Type.Optional(ThinkingLevelMapValueSchema), high: Type.Optional(ThinkingLevelMapValueSchema), xhigh: Type.Optional(ThinkingLevelMapValueSchema), max: Type.Optional(ThinkingLevelMapValueSchema),});
const OpenAICompletionsCompatSchema = Type.Object({ supportsStore: Type.Optional(Type.Boolean()), supportsDeveloperRole: Type.Optional(Type.Boolean()), supportsReasoningEffort: Type.Optional(Type.Boolean()), supportsUsageInStreaming: Type.Optional(Type.Boolean()), maxTokensField: Type.Optional(Type.Union([Type.Literal("max_completion_tokens"), Type.Literal("max_tokens")])), requiresToolResultName: Type.Optional(Type.Boolean()), requiresAssistantAfterToolResult: Type.Optional(Type.Boolean()), requiresThinkingAsText: Type.Optional(Type.Boolean()), requiresReasoningContentOnAssistantMessages: Type.Optional(Type.Boolean()), thinkingFormat: Type.Optional( Type.Union([ Type.Literal("openai"), Type.Literal("openrouter"), Type.Literal("deepseek"), Type.Literal("zai"), Type.Literal("qwen"), Type.Literal("qwen-chat-template"), ]), ), cacheControlFormat: Type.Optional(Type.Literal("anthropic")), openRouterRouting: Type.Optional(OpenRouterRoutingSchema), vercelGatewayRouting: Type.Optional(VercelGatewayRoutingSchema), supportsStrictMode: Type.Optional(Type.Boolean()), supportsLongCacheRetention: Type.Optional(Type.Boolean()),});
const OpenAIResponsesCompatSchema = Type.Object({ sendSessionIdHeader: Type.Optional(Type.Boolean()), supportsLongCacheRetention: Type.Optional(Type.Boolean()),});
const AnthropicMessagesCompatSchema = Type.Object({ supportsEagerToolInputStreaming: Type.Optional(Type.Boolean()), supportsLongCacheRetention: Type.Optional(Type.Boolean()),});
const ProviderCompatSchema = Type.Union([ OpenAICompletionsCompatSchema, OpenAIResponsesCompatSchema, AnthropicMessagesCompatSchema,]);
// Most fields are optional with sensible defaults for local models (Ollama, LM Studio, etc.)const ModelDefinitionSchema = Type.Object({ id: Type.String({ minLength: 1 }), name: Type.Optional(Type.String({ minLength: 1 })), api: Type.Optional(Type.String({ minLength: 1 })), baseUrl: Type.Optional(Type.String({ minLength: 1 })), reasoning: Type.Optional(Type.Boolean()), thinkingLevelMap: Type.Optional(ThinkingLevelMapSchema), input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))), cost: Type.Optional( Type.Object({ input: Type.Number(), output: Type.Number(), cacheRead: Type.Number(), cacheWrite: Type.Number(), }), ), contextWindow: Type.Optional(Type.Number()), maxTokens: Type.Optional(Type.Number()), headers: Type.Optional(Type.Record(Type.String(), Type.String())), compat: Type.Optional(ProviderCompatSchema),});
const ModelOverrideSchema = Type.Object({ name: Type.Optional(Type.String({ minLength: 1 })), reasoning: Type.Optional(Type.Boolean()), thinkingLevelMap: Type.Optional(ThinkingLevelMapSchema), input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))), cost: Type.Optional( Type.Object({ input: Type.Optional(Type.Number()), output: Type.Optional(Type.Number()), cacheRead: Type.Optional(Type.Number()), cacheWrite: Type.Optional(Type.Number()), }), ), contextWindow: Type.Optional(Type.Number()), maxTokens: Type.Optional(Type.Number()), headers: Type.Optional(Type.Record(Type.String(), Type.String())), compat: Type.Optional(ProviderCompatSchema),});
type ModelOverride = Static<typeof ModelOverrideSchema>;
const ProviderConfigSchema = Type.Object({ name: Type.Optional(Type.String({ minLength: 1 })), baseUrl: Type.Optional(Type.String({ minLength: 1 })), apiKey: Type.Optional(Type.String({ minLength: 1 })), api: Type.Optional(Type.String({ minLength: 1 })), headers: Type.Optional(Type.Record(Type.String(), Type.String())), compat: Type.Optional(ProviderCompatSchema), authHeader: Type.Optional(Type.Boolean()), models: Type.Optional(Type.Array(ModelDefinitionSchema)), modelOverrides: Type.Optional(Type.Record(Type.String(), ModelOverrideSchema)),});
const ModelsConfigSchema = Type.Object({ providers: Type.Record(Type.String(), ProviderConfigSchema),});
// typebox/compile costs ~300ms to import, so the models.json validator loads// lazily. The first load of an existing models.json proceeds on JSON parsing +// validateConfig() and reports schema errors asynchronously once the validator// is ready; subsequent refreshes validate synchronously.let validateModelsConfig: Validator<TProperties, typeof ModelsConfigSchema> | undefined;let modelsValidatorPromise: Promise<void> | undefined;
function preloadModelsConfigValidator(): Promise<void> { modelsValidatorPromise ??= import("typebox/compile").then(({ Compile }) => { validateModelsConfig = Compile(ModelsConfigSchema); }); return modelsValidatorPromise;}
type ModelsConfig = Static<typeof ModelsConfigSchema>;
function formatValidationPath(error: TLocalizedValidationError): string { if (error.keyword === "required") { const requiredProperties = (error.params as { requiredProperties?: string[] }).requiredProperties; const requiredProperty = requiredProperties?.[0]; if (requiredProperty) { const basePath = error.instancePath.replace(/^\//, "").replace(/\//g, "."); return basePath ? `${basePath}.${requiredProperty}` : requiredProperty; } } const path = error.instancePath.replace(/^\//, "").replace(/\//g, "."); return path || "root";}
/** Strip `//` line comments and trailing commas from JSON, leaving string literals untouched. */function stripJsonComments(input: string): string { return input .replace(/"(?:\\.|[^"\\])*"|\/\/[^\n]*/g, (m) => (m[0] === '"' ? m : "")) .replace(/"(?:\\.|[^"\\])*"|,(\s*[}\]])/g, (m, tail) => tail ?? (m[0] === '"' ? m : ""));}
/** Provider override config (baseUrl, compat) without request auth/headers */interface ProviderOverride { baseUrl?: string; compat?: Model<Api>["compat"];}
interface ProviderRequestConfig { apiKey?: string; headers?: Record<string, string>; authHeader?: boolean;}
type ProviderRequestAuthSource = { source: "environment" | "models_json_key" | "models_json_command"; configured: true; label?: string; identityFingerprint: string; valueFingerprint?: string; resolveValueFingerprint?: () => string | undefined;};
export type ResolvedRequestAuth = | { ok: true; apiKey?: string; headers?: Record<string, string>; requestModel?: Model<Api>; } | { ok: false; error: string; };
export interface ModelCatalogSnapshot { models: Model<Api>[]; configuredProviders: string[];}
/** Result of loading custom models from models.json */interface CustomModelsResult { models: Model<Api>[]; /** Providers with baseUrl/headers/apiKey overrides for built-in models */ overrides: Map<string, ProviderOverride>; /** Per-model overrides: provider -> modelId -> override */ modelOverrides: Map<string, Map<string, ModelOverride>>; error: string | undefined;}
function emptyCustomModelsResult(error?: string): CustomModelsResult { return { models: [], overrides: new Map(), modelOverrides: new Map(), error };}
function mergeCompat( baseCompat: Model<Api>["compat"], overrideCompat: ModelOverride["compat"],): Model<Api>["compat"] | undefined { if (!overrideCompat) return baseCompat;
const base = baseCompat as OpenAICompletionsCompat | OpenAIResponsesCompat | AnthropicMessagesCompat | undefined; const override = overrideCompat as OpenAICompletionsCompat | OpenAIResponsesCompat | AnthropicMessagesCompat; const merged = { ...base, ...override } as OpenAICompletionsCompat | OpenAIResponsesCompat | AnthropicMessagesCompat;
const baseCompletions = base as OpenAICompletionsCompat | undefined; const overrideCompletions = override as OpenAICompletionsCompat; const mergedCompletions = merged as OpenAICompletionsCompat;
if (baseCompletions?.openRouterRouting || overrideCompletions.openRouterRouting) { mergedCompletions.openRouterRouting = { ...baseCompletions?.openRouterRouting, ...overrideCompletions.openRouterRouting, }; }
if (baseCompletions?.vercelGatewayRouting || overrideCompletions.vercelGatewayRouting) { mergedCompletions.vercelGatewayRouting = { ...baseCompletions?.vercelGatewayRouting, ...overrideCompletions.vercelGatewayRouting, }; }
return merged as Model<Api>["compat"];}
/** * Deep merge a model override into a model. * Handles nested objects (cost, compat) by merging rather than replacing. */function applyModelOverride(model: Model<Api>, override: ModelOverride): Model<Api> { const result = { ...model };
if (override.name !== undefined) result.name = override.name; if (override.reasoning !== undefined) result.reasoning = override.reasoning; if (override.thinkingLevelMap !== undefined) { result.thinkingLevelMap = { ...model.thinkingLevelMap, ...override.thinkingLevelMap }; } if (override.input !== undefined) result.input = override.input as ("text" | "image")[]; if (override.contextWindow !== undefined) result.contextWindow = override.contextWindow; if (override.maxTokens !== undefined) result.maxTokens = override.maxTokens;
if (override.cost) { result.cost = { input: override.cost.input ?? model.cost.input, output: override.cost.output ?? model.cost.output, cacheRead: override.cost.cacheRead ?? model.cost.cacheRead, cacheWrite: override.cost.cacheWrite ?? model.cost.cacheWrite, }; }
result.compat = mergeCompat(model.compat, override.compat);
return result;}
function readOpenAICodexAccountId(token: string): string | undefined { try { const payload = JSON.parse(Buffer.from(token.split(".")[1] ?? "", "base64url").toString("utf8")) as { "https://api.openai.com/auth"?: { chatgpt_account_id?: unknown }; }; const accountId = payload["https://api.openai.com/auth"]?.chatgpt_account_id; return typeof accountId === "string" && accountId.length > 0 ? accountId : undefined; } catch { return undefined; }}
/** * The Codex backend gates its model catalog on the reported client version: it answers HTTP 200 with a * catalog that grows as the version rises, so a low version yields a silently empty or partial list rather * than an error. Prime Agent's own package version is far below the Codex CLI's version line, so it must * report a supported Codex client version here instead. * * Shipping a new Codex model takes two edits, and both are required: * 1. Add the model to `codexModels` in `packages/ai/scripts/generate-models.ts` and regenerate. That list is * explicit, not fetched, so an unlisted model does not exist for Prime Agent at all. * 2. Raise this constant to a Codex CLI release whose catalog includes that model. `getExecutableModels()` * below intersects the registry with the discovered catalog, so a listed model the catalog omits is * dropped. * * Skipping step 2 fails silently and asymmetrically: `rlm` subagent delegation and `rlm.models()` resolve * through `getExecutableModels()` and lose the model, while `/model` reads the unfiltered `getAvailable()` * and keeps offering it. * * Catalog behaviour measured 2026-08-13; see #702. */const OPENAI_CODEX_CLIENT_VERSION = "0.153.4";
function openAICodexModelsUrl(baseUrl: string): string { const normalized = baseUrl.replace(/\/+$/, ""); let path: string; if (normalized.endsWith("/codex/responses")) { path = `${normalized.slice(0, -"/responses".length)}/models`; } else if (normalized.endsWith("/codex")) { path = `${normalized}/models`; } else { path = `${normalized}/codex/models`; } const url = new URL(path); url.searchParams.set("client_version", OPENAI_CODEX_CLIENT_VERSION); return url.toString();}
function readOpenAICodexModelIds(value: unknown): Set<string> { if (!value || typeof value !== "object" || !("models" in value) || !Array.isArray(value.models)) { throw new Error("Invalid OpenAI Codex model catalog"); } return new Set( value.models.flatMap((model) => { if (!model || typeof model !== "object" || !("slug" in model) || typeof model.slug !== "string") { return []; } return [model.slug]; }), );}
const PRIVATE_PRIME_AUTHORIZATION_CACHE_FILE = "prime-inference-private-models.json";const PRIVATE_PRIME_AUTHORIZATION_CACHE_TTL_MS = 5 * 60_000;const PRIVATE_PRIME_BACKGROUND_REFRESH_TIMEOUT_MS = 3_000;
interface PrivatePrimeAuthorizationCache { fingerprint: string; models: Model<"openai-completions">[]; refreshedAt: number;}
function privatePrimeAuthorizationFingerprint(apiKey: string, teamId: string): string { // Use the bearer token as a MAC key, not a password to hash. Keep the scope stable for disk cache reuse. return createHmac("sha256", apiKey) .update("prime-agent:private-prime-authorization:v1\0") .update(teamId) .digest("hex");}
function isOfflineModeEnabled(): boolean { const value = process.env.PI_OFFLINE; if (!value) return false; return value === "1" || value.toLowerCase() === "true" || value.toLowerCase() === "yes";}
/** Default breaker window when a provider sends no retry hint of its own. */export const DEFAULT_PROVIDER_BREAKER_BLOCK_MS = 60_000;
export type ProviderBreakerReason = "quota_exhausted" | "rate_limit";
export interface ProviderBreakerState { provider: string; reason: ProviderBreakerReason; /** Wall-clock ms until which later requests to this provider are refused. */ blockedUntilMs: number;}
export interface ProviderCircuitBreakerOptions { /** Window used when the failure carries neither a deadline nor a retry hint. */ defaultBlockMs?: number; /** Injectable clock, so tests can cross the window without sleeping. */ now?: () => number;}
/** * Rolling in-memory breaker for providers whose account is depleted. Depletion * is a wall, not a hiccup: once tripped, later requests to the same provider * are refused until the window expires, so a swarm cannot keep hammering a dead * account. Nothing is persisted: a restart re-probes the provider once. */export class ProviderCircuitBreaker { private readonly blockedUntil = new Map<string, ProviderBreakerState>(); private readonly defaultBlockMs: number; private readonly now: () => number;
constructor(options: ProviderCircuitBreakerOptions = {}) { this.defaultBlockMs = options.defaultBlockMs ?? DEFAULT_PROVIDER_BREAKER_BLOCK_MS; this.now = options.now ?? Date.now; }
/** Trip the breaker, preferring the provider's own deadline over the default. */ trip( provider: string, options: { blockedUntilMs?: number; retryAfterMs?: number; reason?: ProviderBreakerReason } = {}, ): ProviderBreakerState { const now = this.now(); const requested = options.blockedUntilMs ?? (options.retryAfterMs !== undefined ? now + Math.max(0, options.retryAfterMs) : undefined); const state: ProviderBreakerState = { provider, reason: options.reason ?? "quota_exhausted", blockedUntilMs: Math.max(requested ?? now + this.defaultBlockMs, now), }; this.blockedUntil.set(provider, state); return state; }
/** Open state for one provider; an expired window is forgotten as it is read. */ state(provider: string): ProviderBreakerState | undefined { const state = this.blockedUntil.get(provider); if (!state) { return undefined; } if (state.blockedUntilMs <= this.now()) { this.blockedUntil.delete(provider); return undefined; } return state; }
isOpen(provider: string): boolean { return this.state(provider) !== undefined; }
/** Local override: forget one provider's window before it expires. */ clear(provider: string): void { this.blockedUntil.delete(provider); }
/** Providers currently refused, with expired windows swept. */ open(): ProviderBreakerState[] { return [...this.blockedUntil.keys()].flatMap((provider) => this.state(provider) ?? []); }}
export type RlmSpawnAdmissionReason = ProviderBreakerReason | "auth_unconfigured";
/** * A spawn refused before any child was allocated. Callers switch on `reason` * instead of matching text; the message names what the caller should change. */export class RlmSpawnAdmissionError extends Error { readonly provider: string; readonly reason: RlmSpawnAdmissionReason; readonly blockedUntilMs?: number;
constructor( message: string, options: { provider: string; reason: RlmSpawnAdmissionReason; blockedUntilMs?: number }, ) { super(message); this.name = "RlmSpawnAdmissionError"; this.provider = options.provider; this.reason = options.reason; this.blockedUntilMs = options.blockedUntilMs; }}
/** * Pinned static model metadata. * Models retain these values for the life of a session once loaded or registered; * background catalog refreshes do not mutate active model metadata. */export interface StaticModelMetadata { id: string; provider: string; contextWindow: number; maxTokens: number; cost: { input: number; output: number; cacheRead: number; cacheWrite: number; }; reasoning: boolean; payloadCeiling: number;}
/** * Model registry - loads and manages models, resolves API keys via AuthStorage. */export class ModelRegistry { private models: Model<Api>[] = []; private xaiModelSources = new WeakMap<Model<Api>, Model<Api>>(); private pinnedModelMetadata: Map<string, StaticModelMetadata> = new Map();
private pinModelMetadata(model: Model<Api>): Model<Api> { const key = `${model.provider}/${model.id}`.toLowerCase(); const existing = this.pinnedModelMetadata.get(key); if (existing) { return { ...model, contextWindow: existing.contextWindow, maxTokens: existing.maxTokens, cost: { ...existing.cost }, reasoning: existing.reasoning, }; } const metadata: StaticModelMetadata = { id: model.id, provider: model.provider, contextWindow: model.contextWindow, maxTokens: model.maxTokens, cost: { ...model.cost }, reasoning: model.reasoning, payloadCeiling: calculatePromptPayloadCeiling(model.contextWindow), }; this.pinnedModelMetadata.set(key, metadata); return model; } private providerRequestConfigs: Map<string, ProviderRequestConfig> = new Map(); private staleProviderRequestAuthSources: Map<string, AuthSourceToken[]> = new Map(); private lastProviderAuthSourceTokens: Map<string, AuthSourceToken> = new Map(); private modelRequestHeaders: Map<string, Record<string, string>> = new Map(); private registeredProviders: Map<string, ProviderConfigInput> = new Map(); private authorizedPrivatePrimeInferenceModelIds = new Set<string>(); private authorizedPrivatePrimeInferenceModels: Model<"openai-completions">[] = []; private authorizedPrivatePrimeInferenceTeamId: string | undefined; private explicitPrivatePrimeInferenceModelIds = new Set<string>(); private openAICodexModelsCache: { authFingerprint: string; modelIds: Set<string>; refreshedAt: number } | undefined; private backgroundPrivatePrimeAuthorization: { fingerprint: string; promise: Promise<void> } | undefined; private livePrimeInferenceModels: Model<"openai-completions">[] | undefined; private pendingPrimeInferenceCatalogRefresh: Promise<void> | undefined; private loadError: string | undefined = undefined; private modelsDevLiveModels: Model<Api>[] = []; private modelsDevRefreshTimer?: NodeJS.Timeout; /** Local-only: providers whose quota is depleted, so later requests fail fast. */ private readonly providerCircuitBreaker = new ProviderCircuitBreaker();
/** Re-register dynamic OAuth providers (e.g. user MCP servers) after refresh() resets the registry. */ private onOAuthProvidersReset?: () => void;
private constructor( readonly authStorage: AuthStorage, private modelsJsonPath: string | undefined, private modelsDevCachePath: string | undefined = modelsJsonPath ? join(dirname(modelsJsonPath), "models-dev-cache.json") : undefined, ) { this.loadModelsDevCache(); this.loadModels(); this.initModelsDevBackgroundRefresh(); }
setOnOAuthProvidersReset(hook: () => void): void { this.onOAuthProvidersReset = hook; }
static create( authStorage: AuthStorage, modelsJsonPath: string = join(getAgentDir(), "models.json"), modelsDevCachePath?: string, ): ModelRegistry { return new ModelRegistry( authStorage, modelsJsonPath, modelsDevCachePath ?? (modelsJsonPath ? join(dirname(modelsJsonPath), "models-dev-cache.json") : undefined), ); }
static inMemory(authStorage: AuthStorage): ModelRegistry { return new ModelRegistry(authStorage, undefined, undefined); }
/** * Reload models from disk (built-in + custom from models.json). */ refresh(): void { this.providerRequestConfigs.clear(); this.modelRequestHeaders.clear(); this.lastProviderAuthSourceTokens.clear(); this.explicitPrivatePrimeInferenceModelIds.clear(); this.loadError = undefined;
// Credentials may have been written by another process (e.g. the UI // process saving a login while the session lives in the daemon). this.authStorage.reload(); const teamId = this.authStorage.getProviderHeaders(PRIME_INFERENCE_PROVIDER_ID)?.["X-Prime-Team-ID"]; // Direct refreshes must preserve same-team stale recovery, but invalidate changed auth immediately. if ( this.authStorage.getAuthStatus(PRIME_INFERENCE_PROVIDER_ID).source !== "stale" || !teamId || teamId !== this.authorizedPrivatePrimeInferenceTeamId ) { this.authorizedPrivatePrimeInferenceModelIds.clear(); this.authorizedPrivatePrimeInferenceModels = []; this.authorizedPrivatePrimeInferenceTeamId = undefined; } resetApiProviders(); resetOAuthProviders(); // reset drops everything but model-provider built-ins; re-add MCP integrations // (built-in catalog + this session's user-declared servers via the hook). // NOTE: the OAuth registry is process-global. Built-in MCP providers are // identical across sessions so they always survive; a user-declared server // unique to another daemon session is dropped here and re-registered on that // session's next refresh. Fully isolating it would require a session-scoped // registry in pi-ai (out of scope here). registerBuiltinMcpOAuthProviders(); this.onOAuthProvidersReset?.();
this.loadModelsDevCache(); if (this.modelsDevLiveModels.length > 0) { registerLiveModels(this.modelsDevLiveModels); } this.reloadModelsAfterCatalogChange(); }
private reloadModelsAfterCatalogChange(): void { this.loadModels(); this.reapplyRegisteredProviders(); }
private reapplyRegisteredProviders(): void { for (const [providerName, config] of this.registeredProviders.entries()) { this.applyProviderConfig(providerName, config); } }
private primeInferenceCatalogCachePath(): string | undefined { return this.modelsJsonPath ? join(dirname(this.modelsJsonPath), "prime-inference-models-cache.json") : undefined; }
private bundledPrimeInferenceModels(): Model<"openai-completions">[] { return getModels(PRIME_INFERENCE_PROVIDER_ID) as Model<"openai-completions">[]; }
/** * Get any error from loading models.json (undefined if no error). */ getError(): string | undefined { return this.loadError; }
private loadModelsDevCache(): void { if (!this.modelsDevCachePath || !existsSync(this.modelsDevCachePath)) { return; } try { const content = readFileSync(this.modelsDevCachePath, "utf-8"); const parsed = JSON.parse(content) as { updatedAt?: number; models?: Model<Api>[] }; if (parsed.models && Array.isArray(parsed.models)) { this.modelsDevLiveModels = parsed.models; registerLiveModels(this.modelsDevLiveModels); } } catch { // ignore corrupt cache } }
private initModelsDevBackgroundRefresh(): void { if (!this.modelsDevCachePath || isOfflineModeEnabled()) return;
let shouldRefresh = true; if (existsSync(this.modelsDevCachePath)) { try { const content = readFileSync(this.modelsDevCachePath, "utf-8"); const parsed = JSON.parse(content) as { updatedAt?: number }; if (parsed.updatedAt && Date.now() - parsed.updatedAt < MODELS_DEV_TTL_MS) { shouldRefresh = false; } } catch { shouldRefresh = true; } }
if (shouldRefresh) { void this.refreshModelsDev({ force: false }); }
this.modelsDevRefreshTimer = setInterval( () => { void this.refreshModelsDev({ force: false }); }, 15 * 60 * 1000, ); this.modelsDevRefreshTimer.unref?.(); }
dispose(): void { if (this.modelsDevRefreshTimer) { clearInterval(this.modelsDevRefreshTimer); this.modelsDevRefreshTimer = undefined; } }
async refreshModelsDev(options?: { force?: boolean; }): Promise<{ refreshed: boolean; count: number; error?: string }> { if (!this.modelsDevCachePath || isOfflineModeEnabled()) { return { refreshed: false, count: this.modelsDevLiveModels.length }; }
if (!options?.force && existsSync(this.modelsDevCachePath)) { try { const content = readFileSync(this.modelsDevCachePath, "utf-8"); const parsed = JSON.parse(content) as { updatedAt?: number }; if (parsed.updatedAt && Date.now() - parsed.updatedAt < MODELS_DEV_TTL_MS) { return { refreshed: false, count: this.modelsDevLiveModels.length }; } } catch { // cache corrupt, proceed to fetch } }
try { const models = await fetchModelsDevModels(); this.modelsDevLiveModels = models; registerLiveModels(models);
try { const cacheData = { updatedAt: Date.now(), models, }; writeFileAtomicSync(this.modelsDevCachePath, JSON.stringify(cacheData, null, 2)); } catch (writeError) { log.warn("failed to write models-dev cache", { cachePath: this.modelsDevCachePath, error: writeError instanceof Error ? writeError.message : String(writeError), }); }
this.loadModels(); return { refreshed: true, count: models.length }; } catch (error) { const msg = error instanceof Error ? error.message : String(error); log.warn("failed to refresh models from models.dev", { error: msg }); return { refreshed: false, count: this.modelsDevLiveModels.length, error: msg }; } }
private loadModels(): void { const { models: customModels, overrides, modelOverrides, error, } = this.modelsJsonPath ? this.loadCustomModels(this.modelsJsonPath) : emptyCustomModelsResult();
if (error) { this.loadError = error; }
this.explicitPrivatePrimeInferenceModelIds = new Set( customModels.filter(isPrivatePrimeInferenceModel).map((model) => model.id), ); const cachePath = this.primeInferenceCatalogCachePath(); this.livePrimeInferenceModels ??= cachePath ? readCachedPrimeInferenceModels(cachePath, this.bundledPrimeInferenceModels()) : undefined; const privateModels = new Map( [...getPrivatePrimeInferenceModels(), ...this.authorizedPrivatePrimeInferenceModels].map((model) => [ model.id, model, ]), ); const builtInModels = [ ...this.loadBuiltInModels(overrides, modelOverrides, this.livePrimeInferenceModels), ...privateModels.values(), ]; let combined = this.mergeCustomModels(builtInModels, customModels);
for (const oauthProvider of this.authStorage.getOAuthProviders()) { const cred = this.authStorage.get(oauthProvider.id); if (cred?.type === "oauth" && oauthProvider.modifyModels) { combined = oauthProvider.modifyModels(combined, cred); } }
this.models = combined.map((m) => this.pinModelMetadata(m)); }
/** Load built-in models and apply provider/model overrides */ private loadBuiltInModels( overrides: Map<string, ProviderOverride>, modelOverrides: Map<string, Map<string, ModelOverride>>, livePrimeInferenceModels?: Model<"openai-completions">[], ): Model<Api>[] { const bundledModels = getProviders().flatMap((provider) => getModels(provider as KnownProvider) as Model<Api>[]); return mergePrimeInferenceModels(bundledModels, livePrimeInferenceModels).map((model) => { const providerOverride = overrides.get(model.provider); const perModelOverrides = modelOverrides.get(model.provider); let configuredModel = model;
if (providerOverride) { configuredModel = { ...configuredModel, baseUrl: providerOverride.baseUrl ?? configuredModel.baseUrl, compat: mergeCompat(configuredModel.compat, providerOverride.compat), }; }
const modelOverride = perModelOverrides?.get(model.id); return modelOverride ? applyModelOverride(configuredModel, modelOverride) : configuredModel; }); }
/** Merge custom models into built-in list by provider+id (custom wins on conflicts). */ private mergeCustomModels(builtInModels: Model<Api>[], customModels: Model<Api>[]): Model<Api>[] { const merged = [...builtInModels]; for (const customModel of customModels) { const existingIndex = merged.findIndex((m) => m.provider === customModel.provider && m.id === customModel.id); if (existingIndex >= 0) { merged[existingIndex] = customModel; } else { merged.push(customModel); } } return merged; }
private loadCustomModels(modelsJsonPath: string): CustomModelsResult { if (!existsSync(modelsJsonPath)) { return emptyCustomModelsResult(); }
try { const content = readFileSync(modelsJsonPath, "utf-8"); const parsed = JSON.parse(stripJsonComments(content)) as unknown;
if (!validateModelsConfig) { // Validator not loaded yet (first refresh during startup): proceed on // validateConfig() below and report schema errors asynchronously once // the validator is ready. Later refreshes validate synchronously. void preloadModelsConfigValidator().then(() => { if (validateModelsConfig && !validateModelsConfig.Check(parsed)) { const errors = validateModelsConfig .Errors(parsed) .map((error) => ` - ${formatValidationPath(error)}: ${error.message}`) .join("\n") || "Unknown schema error"; log.error("invalid models.json schema", { errors, file: modelsJsonPath }); } }); } else if (!validateModelsConfig.Check(parsed)) { const errors = validateModelsConfig .Errors(parsed) .map((error) => ` - ${formatValidationPath(error)}: ${error.message}`) .join("\n") || "Unknown schema error"; return emptyCustomModelsResult(`Invalid models.json schema:\n${errors}\n\nFile: ${modelsJsonPath}`); }
const config = parsed as ModelsConfig;
this.validateConfig(config);
const overrides = new Map<string, ProviderOverride>(); const modelOverrides = new Map<string, Map<string, ModelOverride>>();
for (const [providerName, providerConfig] of Object.entries(config.providers)) { if (providerConfig.baseUrl || providerConfig.compat) { overrides.set(providerName, { baseUrl: providerConfig.baseUrl, compat: providerConfig.compat, }); }
this.storeProviderRequestConfig(providerName, providerConfig);
if (providerConfig.modelOverrides) { modelOverrides.set(providerName, new Map(Object.entries(providerConfig.modelOverrides))); for (const [modelId, modelOverride] of Object.entries(providerConfig.modelOverrides)) { this.storeModelHeaders(providerName, modelId, modelOverride.headers); } } }
return { models: this.parseModels(config), overrides, modelOverrides, error: undefined }; } catch (error) { if (error instanceof SyntaxError) { return emptyCustomModelsResult(`Failed to parse models.json: ${error.message}\n\nFile: ${modelsJsonPath}`); } return emptyCustomModelsResult( `Failed to load models.json: ${error instanceof Error ? error.message : error}\n\nFile: ${modelsJsonPath}`, ); } }
private validateConfig(config: ModelsConfig): void { const builtInProviders = new Set<string>(getProviders());
for (const [providerName, providerConfig] of Object.entries(config.providers)) { const isBuiltIn = builtInProviders.has(providerName); const hasProviderApi = !!providerConfig.api; const models = providerConfig.models ?? []; const hasModelOverrides = providerConfig.modelOverrides && Object.keys(providerConfig.modelOverrides).length > 0;
if (models.length === 0) { if (!providerConfig.baseUrl && !providerConfig.headers && !providerConfig.compat && !hasModelOverrides) { throw new Error( `Provider ${providerName}: must specify "baseUrl", "headers", "compat", "modelOverrides", or "models".`, ); } } else if (!isBuiltIn) { if (!providerConfig.baseUrl) { throw new Error(`Provider ${providerName}: "baseUrl" is required when defining custom models.`); } if (!providerConfig.apiKey) { throw new Error(`Provider ${providerName}: "apiKey" is required when defining custom models.`); } } // inherited from built-in models. Auth comes from env vars / auth storage.
for (const modelDef of models) { const hasModelApi = !!modelDef.api;
if (!hasProviderApi && !hasModelApi && !isBuiltIn) { throw new Error( `Provider ${providerName}, model ${modelDef.id}: no "api" specified. Set at provider or model level.`, ); }
if (!modelDef.id) throw new Error(`Provider ${providerName}: model missing "id"`); if (modelDef.contextWindow !== undefined && modelDef.contextWindow <= 0) throw new Error(`Provider ${providerName}, model ${modelDef.id}: invalid contextWindow`); if (modelDef.maxTokens !== undefined && modelDef.maxTokens <= 0) throw new Error(`Provider ${providerName}, model ${modelDef.id}: invalid maxTokens`); } } }
private parseModels(config: ModelsConfig): Model<Api>[] { const models: Model<Api>[] = []; const builtInProviders = new Set<string>(getProviders());
const builtInDefaultsCache = new Map<string, { api: string; baseUrl: string }>(); const getBuiltInDefaults = (providerName: string): { api: string; baseUrl: string } | undefined => { if (!builtInProviders.has(providerName)) return undefined; if (builtInDefaultsCache.has(providerName)) return builtInDefaultsCache.get(providerName); const builtIn = getModels(providerName as KnownProvider) as Model<Api>[]; if (builtIn.length === 0) return undefined; const defaults = { api: builtIn[0].api, baseUrl: builtIn[0].baseUrl }; builtInDefaultsCache.set(providerName, defaults); return defaults; };
for (const [providerName, providerConfig] of Object.entries(config.providers)) { const modelDefs = providerConfig.models ?? []; if (modelDefs.length === 0) continue; // Override-only, no custom models
const builtInDefaults = getBuiltInDefaults(providerName);
for (const modelDef of modelDefs) { const api = modelDef.api ?? providerConfig.api ?? builtInDefaults?.api; if (!api) continue;
const baseUrl = modelDef.baseUrl ?? providerConfig.baseUrl ?? builtInDefaults?.baseUrl; if (!baseUrl) continue;
const compat = mergeCompat(providerConfig.compat, modelDef.compat); this.storeModelHeaders(providerName, modelDef.id, modelDef.headers);
const defaultCost = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; models.push({ id: modelDef.id, name: modelDef.name ?? modelDef.id, api: api as Api, provider: providerName, baseUrl, reasoning: modelDef.reasoning ?? false, thinkingLevelMap: modelDef.thinkingLevelMap, input: (modelDef.input ?? ["text"]) as ("text" | "image")[], cost: modelDef.cost ?? defaultCost, contextWindow: modelDef.contextWindow ?? 128000, maxTokens: modelDef.maxTokens ?? 16384, headers: undefined, compat, } as Model<Api>); } }
return models; }
/** * Get all models (built-in + custom). * If models.json had errors, returns only built-in models. */ getAll(): Model<Api>[] { return this.models.map((model) => this.getModelForCurrentAuth(model)); }
getModelForCurrentAuth(model: Model<Api>): Model<Api> { if (model.provider !== "xai") return model; const source = this.xaiModelSources.get(model) ?? model; return this.isUsingXaiSubscription(source) ? (this.getXaiSubscriptionModel(source) ?? source) : source; }
private getXaiSubscriptionModel(model: Model<Api>): Model<Api> | undefined { const source = this.xaiModelSources.get(model) ?? model; const subscription = getXaiSubscriptionModel(source); if (subscription) this.xaiModelSources.set(subscription, source); return subscription; }
private isUsingXaiSubscription(model: Model<Api>): boolean { return ( model.provider === "xai" && this.authStorage.get("xai")?.type === "oauth" && this.getProviderAuthStatus("xai").source === "stored" ); }
/** * Get only models that have auth configured. * This is a fast check that doesn't refresh OAuth tokens. */ getAvailable(): Model<Api>[] { return this.getAll().filter((model) => { if (isPrivatePrimeInferenceModel(model) && !this.isAuthorizedPrivatePrimeInferenceModel(model)) { return false; } return this.hasConfiguredAuth(model); }); }
/** * Reload local state and private authorization. Public Prime Inference models * return from the disk/bundled fallback immediately and refresh in the background. */ async refreshAvailableModels(): Promise<Model<Api>[]> { return this.runSerializedEntitlementRefresh(async () => { const previousPrivateModelIds = new Set(this.authorizedPrivatePrimeInferenceModelIds); const previousTeamId = this.authorizedPrivatePrimeInferenceTeamId; const previousPrivateModels = this.authorizedPrivatePrimeInferenceModels; this.refresh(); const cachePath = this.primeInferenceCatalogCachePath(); if (cachePath) { // Track the in-flight catalog refresh so waitForPendingModelRefreshes() // can await it: session-model restore must observe the post-refresh // catalog, not the pre-refresh fallback (bundled/disk snapshot). this.pendingPrimeInferenceCatalogRefresh = refreshPrimeInferenceModels( cachePath, this.bundledPrimeInferenceModels(), { offline: isOfflineModeEnabled(), }, ) .then((models) => { if (!models) return; this.livePrimeInferenceModels = models; this.reloadModelsAfterCatalogChange(); }) // Settlement, not failure, is what the waiter observes. .catch(() => undefined); } await this.refreshPrivatePrimeInferenceAuthorization( previousPrivateModelIds, previousTeamId, previousPrivateModels, ); return this.getAvailable(); }); }
/** * Wait (bounded by timeoutMs) for in-flight Prime Inference catalog and * private-authorization refreshes to settle. * * refreshAvailableModels() returns while the catalog fetch and the stale-cache * authorization refresh keep running in the background, so a decision taken * right after it can still see pre-refresh state (e.g. restoring a session's * saved model on a freshly restarted daemon). Returns immediately when * nothing is pending; never waits longer than timeoutMs. */ async waitForPendingModelRefreshes(timeoutMs: number): Promise<void> { const pending: Promise<unknown>[] = []; if (this.pendingPrimeInferenceCatalogRefresh) pending.push(this.pendingPrimeInferenceCatalogRefresh); if (this.backgroundPrivatePrimeAuthorization?.promise) { pending.push(this.backgroundPrivatePrimeAuthorization.promise); } if (pending.length === 0) return; await Promise.race([ Promise.allSettled(pending), new Promise<void>((resolve) => { const timer = setTimeout(resolve, timeoutMs); // The bounded wait must never keep the process alive on its own. if (typeof timer.unref === "function") timer.unref(); }), ]); }
private entitlementRefreshChain: Promise<unknown> = Promise.resolve();
private runSerializedEntitlementRefresh<T>(task: () => Promise<T>): Promise<T> { const run = this.entitlementRefreshChain.then(task, task); this.entitlementRefreshChain = run.then( () => undefined, () => undefined, ); return run; }
private async refreshPrivatePrimeInferenceAuthorization( previousPrivateModelIds = new Set(this.authorizedPrivatePrimeInferenceModelIds), previousTeamId = this.authorizedPrivatePrimeInferenceTeamId, previousPrivateModels = this.authorizedPrivatePrimeInferenceModels, ): Promise<void> { const apiKey = await this.authStorage.getApiKey(PRIME_INFERENCE_PROVIDER_ID); const teamHeaders = this.authStorage.getProviderHeaders(PRIME_INFERENCE_PROVIDER_ID); const teamId = teamHeaders?.["X-Prime-Team-ID"]; if (!apiKey || !teamHeaders || !teamId) { // Stale is not logout: keep fetched entitlements for explicit re-selection // (the auth filter still hides the models while stale) — but only for // the team they were fetched for; a team switch invalidates them. if ( this.authStorage.getAuthStatus(PRIME_INFERENCE_PROVIDER_ID).source === "stale" && teamId && teamId === previousTeamId ) { this.authorizedPrivatePrimeInferenceModelIds = previousPrivateModelIds; this.authorizedPrivatePrimeInferenceModels = previousPrivateModels; this.authorizedPrivatePrimeInferenceTeamId = previousTeamId; this.reloadModelsAfterCatalogChange(); return; } this.authorizedPrivatePrimeInferenceModelIds.clear(); this.authorizedPrivatePrimeInferenceModels = []; this.authorizedPrivatePrimeInferenceTeamId = undefined; this.reloadModelsAfterCatalogChange(); return; }
const fingerprint = privatePrimeAuthorizationFingerprint(apiKey, teamId); const cached = this.readPrivatePrimeAuthorizationCache(); if (cached?.fingerprint === fingerprint) { // Serve the credential-scoped cache so startup and model lists don't // block on the network. Stale entries refresh in the background. this.authorizedPrivatePrimeInferenceModels = cached.models; this.authorizedPrivatePrimeInferenceModelIds = new Set(cached.models.map((model) => model.id)); this.authorizedPrivatePrimeInferenceTeamId = teamId; this.reloadModelsAfterCatalogChange(); const cacheIsFresh = Date.now() - cached.refreshedAt < PRIVATE_PRIME_AUTHORIZATION_CACHE_TTL_MS; if (isOfflineModeEnabled() || cacheIsFresh) return; this.startBackgroundPrivatePrimeAuthorizationRefresh(apiKey, teamHeaders, teamId, fingerprint); return; } if (isOfflineModeEnabled()) { this.authorizedPrivatePrimeInferenceModelIds.clear(); this.authorizedPrivatePrimeInferenceModels = []; this.authorizedPrivatePrimeInferenceTeamId = undefined; this.reloadModelsAfterCatalogChange(); return; }
let authorizedModels: Model<"openai-completions">[] | undefined; try { authorizedModels = await fetchAuthorizedPrivatePrimeInferenceModels( apiKey, teamHeaders, new Set((this.livePrimeInferenceModels ?? this.bundledPrimeInferenceModels()).map((model) => model.id)), ); } catch { // Fall back to the previous authorization below. } // Leave newer state untouched if the credentials changed while fetching. if ((await this.currentPrivatePrimeAuthorizationFingerprint()) !== fingerprint) { return; } if (authorizedModels) { this.authorizedPrivatePrimeInferenceModels = authorizedModels; this.authorizedPrivatePrimeInferenceModelIds = new Set(authorizedModels.map((model) => model.id)); this.authorizedPrivatePrimeInferenceTeamId = teamId; this.reloadModelsAfterCatalogChange(); this.writePrivatePrimeAuthorizationCache({ fingerprint, models: authorizedModels, refreshedAt: Date.now() }); } else if (teamId === previousTeamId) { this.authorizedPrivatePrimeInferenceModelIds = previousPrivateModelIds; this.authorizedPrivatePrimeInferenceModels = previousPrivateModels; this.authorizedPrivatePrimeInferenceTeamId = teamId; this.reloadModelsAfterCatalogChange(); } else { this.authorizedPrivatePrimeInferenceModelIds.clear(); this.authorizedPrivatePrimeInferenceModels = []; this.authorizedPrivatePrimeInferenceTeamId = undefined; this.reloadModelsAfterCatalogChange(); } }
/** * Stale cache hits refresh in the background; failures keep the cached ids. * Refreshes for the same credentials are deduped, a changed-credentials * refresh is queued after the in-flight one, and a result is only applied * if the credentials it was fetched with are still current. */ private startBackgroundPrivatePrimeAuthorizationRefresh( apiKey: string, teamHeaders: Record<string, string>, teamId: string, fingerprint: string, ): void { if (this.backgroundPrivatePrimeAuthorization?.fingerprint === fingerprint) { return; } const run = async () => { try { const authorizedModels = await fetchAuthorizedPrivatePrimeInferenceModels( apiKey, teamHeaders, new Set((this.livePrimeInferenceModels ?? this.bundledPrimeInferenceModels()).map((model) => model.id)), undefined, PRIVATE_PRIME_BACKGROUND_REFRESH_TIMEOUT_MS, ); if ((await this.currentPrivatePrimeAuthorizationFingerprint()) !== fingerprint) { return; } this.authorizedPrivatePrimeInferenceModels = authorizedModels; this.authorizedPrivatePrimeInferenceModelIds = new Set(authorizedModels.map((model) => model.id)); this.authorizedPrivatePrimeInferenceTeamId = teamId; this.reloadModelsAfterCatalogChange(); this.writePrivatePrimeAuthorizationCache({ fingerprint, models: authorizedModels, refreshedAt: Date.now(), }); } catch { // Keep the cached authorization. } }; const pending = this.backgroundPrivatePrimeAuthorization?.promise; const promise = (pending ?? Promise.resolve()).then(run); this.backgroundPrivatePrimeAuthorization = { fingerprint, promise }; void promise.finally(() => { if (this.backgroundPrivatePrimeAuthorization?.promise === promise) { this.backgroundPrivatePrimeAuthorization = undefined; } }); }
private async currentPrivatePrimeAuthorizationFingerprint(): Promise<string | undefined> { const apiKey = await this.authStorage.getApiKey(PRIME_INFERENCE_PROVIDER_ID); const teamId = this.authStorage.getProviderHeaders(PRIME_INFERENCE_PROVIDER_ID)?.["X-Prime-Team-ID"]; return apiKey && teamId ? privatePrimeAuthorizationFingerprint(apiKey, teamId) : undefined; }
private privatePrimeAuthorizationCachePath(): string | undefined { if (!this.modelsJsonPath) { return undefined; } return join(dirname(this.modelsJsonPath), PRIVATE_PRIME_AUTHORIZATION_CACHE_FILE); }
private readPrivatePrimeAuthorizationCache(): PrivatePrimeAuthorizationCache | undefined { const cachePath = this.privatePrimeAuthorizationCachePath(); if (!cachePath) return undefined; try { const parsed = JSON.parse(readFileSync(cachePath, "utf8")) as { fingerprint?: unknown; data?: unknown; refreshedAt?: unknown; }; if ( typeof parsed.fingerprint !== "string" || !Array.isArray(parsed.data) || typeof parsed.refreshedAt !== "number" ) { return undefined; } const entries = parsePrimeInferenceModelCatalog({ data: parsed.data }, { allowEmpty: true }).filter((entry) => isPrivatePrimeInferenceModel({ provider: PRIME_INFERENCE_PROVIDER_ID, id: entry.id }), ); const models = buildPrimeInferenceModels(getPrivatePrimeInferenceModels(), entries, { includePrivate: true, minimumModels: 0, }); return { fingerprint: parsed.fingerprint, models: models ?? [], refreshedAt: parsed.refreshedAt }; } catch { return undefined; } }
private writePrivatePrimeAuthorizationCache(cache: PrivatePrimeAuthorizationCache): void { const cachePath = this.privatePrimeAuthorizationCachePath(); if (!cachePath) return; const data = cache.models.map((model) => ({ id: model.id, display_name: model.name, pricing: { input_usd_per_mtok: model.cost.input, output_usd_per_mtok: model.cost.output, cache_read_usd_per_mtok: model.cost.cacheRead, cache_write_usd_per_mtok: model.cost.cacheWrite, }, specs: { context_window: model.contextWindow, max_output_tokens: model.maxTokens, modalities: { input: model.input, output: ["text"] }, supports_reasoning: model.reasoning, }, })); try { writeFileAtomicSync( cachePath, JSON.stringify({ fingerprint: cache.fingerprint, data, refreshedAt: cache.refreshedAt }), { mode: 0o600 }, ); } catch { // A failed cache write only requires a later refetch. } }
async refreshModelCatalog(): Promise<ModelCatalogSnapshot> { const availableModels = await this.refreshAvailableModels(); const availablePrivateModels = new Set( availableModels.filter(isPrivatePrimeInferenceModel).map((model) => `${model.provider}/${model.id}`), ); return { models: this.getAll().filter( (model) => !isPrivatePrimeInferenceModel(model) || availablePrivateModels.has(`${model.provider}/${model.id}`), ), configuredProviders: [...new Set(availableModels.map((model) => model.provider))], }; }
/** `assumeAuthConfigured` validates an explicit stale-provider selection BEFORE the clear commits. */ async canUseModel(model: Model<Api>, options?: { assumeAuthConfigured?: boolean }): Promise<boolean> { if (options?.assumeAuthConfigured) { // Must be side-effect-free: a keyless refresh would drop the cached entitlements it needs. return !isPrivatePrimeInferenceModel(model) || this.isAuthorizedPrivatePrimeInferenceModel(model); } if (!this.hasConfiguredAuth(model)) { return false; } if (!isPrivatePrimeInferenceModel(model)) { return true; }
const availableModels = await this.refreshAvailableModels(); return availableModels.some((candidate) => candidate.provider === model.provider && candidate.id === model.id); }
private isAuthorizedPrivatePrimeInferenceModel(model: Model<Api>): boolean { return ( this.explicitPrivatePrimeInferenceModelIds.has(model.id) || this.authorizedPrivatePrimeInferenceModelIds.has(model.id) ); }
async getExecutableModels(): Promise<Model<Api>[]> { await this.runSerializedEntitlementRefresh(() => this.refreshPrivatePrimeInferenceAuthorization()); const availableModels = this.getAvailable(); const codexModels = availableModels.filter((model) => model.provider === "openai-codex"); if (codexModels.length === 0) { return availableModels; }
const auth = await this.getApiKeyAndHeaders(codexModels[0]!); if (!auth.ok || !auth.apiKey) { return availableModels.filter((model) => model.provider !== "openai-codex"); } const authFingerprint = createHmac("sha256", auth.apiKey) .update("prime-agent:openai-codex-models:v1") .digest("hex"); const cached = this.openAICodexModelsCache; if (cached?.authFingerprint === authFingerprint && Date.now() - cached.refreshedAt < 300_000) { return availableModels.filter((model) => model.provider !== "openai-codex" || cached.modelIds.has(model.id)); }
const accountId = readOpenAICodexAccountId(auth.apiKey); if (!accountId) { return availableModels.filter((model) => model.provider !== "openai-codex"); } try { const response = await fetch(openAICodexModelsUrl(codexModels[0]!.baseUrl), { headers: { ...auth.headers, Authorization: `Bearer ${auth.apiKey}`, "chatgpt-account-id": accountId, originator: "pi", }, signal: AbortSignal.timeout(5_000), }); if (!response.ok) { throw new Error(`OpenAI Codex model discovery failed with HTTP ${response.status}`); } const modelIds = readOpenAICodexModelIds(await response.json()); this.openAICodexModelsCache = { authFingerprint, modelIds, refreshedAt: Date.now() }; return availableModels.filter((model) => model.provider !== "openai-codex" || modelIds.has(model.id)); } catch { if (cached?.authFingerprint === authFingerprint && Date.now() - cached.refreshedAt < 300_000) { return availableModels.filter( (model) => model.provider !== "openai-codex" || cached.modelIds.has(model.id), ); } return availableModels.filter((model) => model.provider !== "openai-codex"); } }
/** * Find a model by provider and ID. */ find(provider: string, modelId: string): Model<Api> | undefined { const model = this.models.find((m) => m.provider === provider && m.id === modelId); return model ? this.getModelForCurrentAuth(model) : undefined; }
/** * Get static pinned metadata for a model. * Static per model id for the life of the session; does not re-derive or mutate on remote refresh. */ getStaticModelMetadata(modelId: string, provider?: string): StaticModelMetadata | undefined { if (provider) { const key = `${provider}/${modelId}`.toLowerCase(); const pinned = this.pinnedModelMetadata.get(key); if (pinned) return pinned; } if (modelId.includes("/")) { const pinned = this.pinnedModelMetadata.get(modelId.toLowerCase()); if (pinned) return pinned; } for (const meta of this.pinnedModelMetadata.values()) { if (meta.id.toLowerCase() === modelId.toLowerCase()) { return meta; } } const found = this.models.find( (m) => (provider ? m.provider.toLowerCase() === provider.toLowerCase() : true) && (m.id.toLowerCase() === modelId.toLowerCase() || `${m.provider}/${m.id}`.toLowerCase() === modelId.toLowerCase()), ); if (found) { return { id: found.id, provider: found.provider, contextWindow: found.contextWindow, maxTokens: found.maxTokens, cost: { ...found.cost }, reasoning: found.reasoning, payloadCeiling: calculatePromptPayloadCeiling(found.contextWindow), }; } return undefined; }
/** Check if static metadata has been pinned for a model in this session. */ hasPinnedModelMetadata(provider: string, modelId: string): boolean { return this.pinnedModelMetadata.has(`${provider}/${modelId}`.toLowerCase()); }
/** Get static pinned metadata for all known models. */ getAllStaticModelMetadata(): StaticModelMetadata[] { return Array.from(this.pinnedModelMetadata.values()); }
/** * Get API key for a model. */ hasConfiguredAuth(model: Model<Api>): boolean { return this.authStorage.hasAuth(model.provider) || this.hasConfiguredProviderRequestAuth(model.provider); }
private fingerprintProviderRequestAuthSource(source: ProviderRequestAuthSource["source"], material: string): string { const digest = createHash("sha256").update(source).update("\0").update(material).digest("hex"); return `${source}:${digest}`; }
private createProviderRequestAuthSource(options: { source: ProviderRequestAuthSource["source"]; identityMaterial: string; valueMaterial?: string; label?: string; resolveValueMaterial?: () => string | undefined; }): ProviderRequestAuthSource { return { configured: true, source: options.source, ...(options.label ? { label: options.label } : {}), identityFingerprint: this.fingerprintProviderRequestAuthSource( options.source, `identity:${options.identityMaterial}`, ), ...(options.valueMaterial !== undefined ? { valueFingerprint: this.fingerprintProviderRequestAuthSource( options.source, `value:${options.identityMaterial}\0${options.valueMaterial}`, ), } : {}), ...(options.resolveValueMaterial ? { resolveValueFingerprint: () => { const valueMaterial = options.resolveValueMaterial?.(); return valueMaterial === undefined ? undefined : this.fingerprintProviderRequestAuthSource( options.source, `value:${options.identityMaterial}\0${valueMaterial}`, ); }, } : {}), }; }
private getProviderRequestAuthSource( provider: string, options?: { resolvedApiKey?: string }, ): ProviderRequestAuthSource | undefined { const providerApiKey = this.providerRequestConfigs.get(provider)?.apiKey; if (!providerApiKey) { return undefined; }
if (providerApiKey.startsWith("!")) { return this.createProviderRequestAuthSource({ source: "models_json_command", identityMaterial: providerApiKey, valueMaterial: options?.resolvedApiKey === undefined ? undefined : `${providerApiKey}\0${options.resolvedApiKey}`, resolveValueMaterial: () => { const resolved = resolveConfigValueUncached(providerApiKey); return resolved === undefined ? undefined : `${providerApiKey}\0${resolved}`; }, }); }
const envValue = process.env[providerApiKey]; if (envValue) { return this.createProviderRequestAuthSource({ source: "environment", label: providerApiKey, identityMaterial: providerApiKey, valueMaterial: `${providerApiKey}\0${envValue}`, }); }
return this.createProviderRequestAuthSource({ source: "models_json_key", identityMaterial: provider, valueMaterial: providerApiKey, }); }
private isProviderRequestAuthStale(provider: string, source: ProviderRequestAuthSource): boolean { const matchingStale = this.getMatchingStaleProviderRequestAuthSources(provider, source); if (matchingStale.length === 0) { return false; } const valueFingerprint = source.valueFingerprint ?? source.resolveValueFingerprint?.(); return Boolean(valueFingerprint && matchingStale.some((token) => token.valueFingerprint === valueFingerprint)); }
private isProviderRequestAuthStaleForStatus(provider: string, source: ProviderRequestAuthSource): boolean { const matchingStale = this.getMatchingStaleProviderRequestAuthSources(provider, source); if (matchingStale.length === 0) { return false; } if (!source.valueFingerprint) { return true; } const isStale = matchingStale.some((token) => token.valueFingerprint === source.valueFingerprint); if (!isStale) { this.clearStaleProviderRequestAuthSource(provider, source); } return isStale; }
private getMatchingStaleProviderRequestAuthSources( provider: string, source: ProviderRequestAuthSource, ): AuthSourceToken[] { const stale = this.staleProviderRequestAuthSources.get(provider); if (!stale) { return []; } return stale.filter( (token) => token.source === source.source && token.identityFingerprint === source.identityFingerprint, ); }
private clearStaleProviderRequestAuthSource(provider: string, source: ProviderRequestAuthSource): void { const stale = this.staleProviderRequestAuthSources.get(provider); if (!stale) { return; } const next = stale.filter( (token) => token.source !== source.source || token.identityFingerprint !== source.identityFingerprint, ); if (next.length === 0) { this.staleProviderRequestAuthSources.delete(provider); } else { this.staleProviderRequestAuthSources.set(provider, next); } }
private getProviderRequestAuthSourceToken( provider: string, source: ProviderRequestAuthSource, ): AuthSourceToken | undefined { const valueFingerprint = source.valueFingerprint ?? source.resolveValueFingerprint?.(); if (!valueFingerprint) { return undefined; } return { provider, source: source.source, identityFingerprint: source.identityFingerprint, valueFingerprint, }; }
private setLastProviderAuthSourceToken(provider: string, token: AuthSourceToken | undefined): void { if (token) { this.lastProviderAuthSourceTokens.set(provider, token); } else { this.lastProviderAuthSourceTokens.delete(provider); } }
private hasConfiguredProviderRequestAuth(provider: string): boolean { const source = this.getProviderRequestAuthSource(provider); return source !== undefined && !this.isProviderRequestAuthStaleForStatus(provider, source); }
markProviderAuthStale(provider: string, sessionId?: string): boolean { if (this.authStorage.markAuthStale(provider, sessionId)) { return true; }
const token = this.getCurrentProviderAuthSourceToken(provider, sessionId); return token ? this.markProviderAuthSourceStale(token) : false; }
/** * Rate-limit the credential a provider most recently used, so subsequent * requests rotate to a sibling account until the block expires. Falls back * to a short default cooldown when no duration is known. */ markProviderRateLimited(provider: string, blockedUntilMs?: number, sessionId?: string): boolean { return this.authStorage.blockLastSelectedCredential( provider, blockedUntilMs ?? Date.now() + DEFAULT_RATE_LIMIT_BLOCK_MS, "rate_limit", sessionId, ); }
/** * Trip the quota breaker for a provider. A provider-supplied retry hint wins * over the 60s default, so an explicit "come back in 5 min" is honoured * without a live usage probe. */ tripProviderQuotaBreaker( provider: string, options?: { blockedUntilMs?: number; retryAfterMs?: number; reason?: ProviderBreakerReason }, ): ProviderBreakerState { return this.providerCircuitBreaker.trip(provider, options); }
/** Open breaker state for one provider; undefined once the window expires. */ providerQuotaBreaker(provider: string): ProviderBreakerState | undefined { return this.providerCircuitBreaker.state(provider); }
/** * Providers to route around, for `rlm.models().exhausted`: the in-memory quota * breaker plus the persisted credential blocks. * * The breaker is deliberately not persisted, so on its own it advertises a * depleted lane as available after any restart, while the durable block store * keeps the record. Both belong in the same advisory list: a spawn on either is * expected to die at turn 0, which is exactly what the caller must not discover * only from a child that never answers. */ exhaustedProviders(): string[] { return [ ...new Set([ ...this.providerCircuitBreaker.open().map((state) => state.provider), ...this.authStorage.blockedProviders(), ]), ]; }
/** Local override for a provider that recovered: forget its breaker window. */ clearProviderQuotaBreaker(provider: string): void { this.providerCircuitBreaker.clear(provider); }
/** * Local admission gate for an RLM spawn: configured auth first, then the * quota breaker. Local state only, so a provider that would kill the child * at turn 0 is rejected before the child is allocated. */ assertRlmSpawnAdmission(model: Model<Api>, options?: { force?: boolean }): void { if (!this.hasConfiguredAuth(model)) { throw new RlmSpawnAdmissionError( `Provider "${model.provider}" has no configured credentials; finish /login or set its API key before spawning`, { provider: model.provider, reason: "auth_unconfigured" }, ); } if (options?.force) { return; } // Only the in-memory breaker refuses here. A durable credential block is // reported through `exhaustedProviders()` and leaves the spawn admissible: it is // a record of one past refusal, and a lane can serve again well before its // window expires, so refusing on it would brick a working provider. const state = this.providerCircuitBreaker.state(`${model.provider}/${model.id}`) ?? this.providerCircuitBreaker.state(model.provider); if (!state) { return; } const seconds = Math.max(1, Math.ceil((state.blockedUntilMs - Date.now()) / 1000)); const targetName = state.provider.includes("/") ? `Model "${state.provider}"` : `Provider "${state.provider}"`; const reasonDesc = state.reason === "rate_limit" ? "rate-limited" : "out of quota"; throw new RlmSpawnAdmissionError( `${targetName} is ${reasonDesc}; spawning is blocked for ~${seconds}s. Pick another provider or wait for the reset.`, { provider: model.provider, reason: state.reason, blockedUntilMs: state.blockedUntilMs }, ); }
hasUnblockedProviderCredential(provider: string): boolean { return this.authStorage.hasUnblockedCredential(provider); }
/** Forget stale-auth markings; a structured auth failure on the next request re-marks the provider. */ clearProviderAuthStale(provider: string): void { this.staleProviderRequestAuthSources.delete(provider); this.authStorage.clearAuthStale(provider); }
getCurrentProviderAuthSourceToken(provider: string, sessionId?: string): AuthSourceToken | undefined { // Stored credentials are session-pinned; consult them before the provider-global // request-config token so concurrent sessions cannot invalidate each other. if (sessionId) { const sessionToken = this.authStorage.getCurrentAuthSourceToken(provider, sessionId); if (sessionToken) return sessionToken; }
const lastRequestToken = this.lastProviderAuthSourceTokens.get(provider); if (lastRequestToken) return lastRequestToken;
const authStorageToken = this.authStorage.getCurrentAuthSourceToken(provider); if (authStorageToken) return authStorageToken;
const providerApiKey = this.providerRequestConfigs.get(provider)?.apiKey; if (!providerApiKey) { return undefined; } const resolvedApiKey = resolveConfigValueUncached(providerApiKey); const source = this.getProviderRequestAuthSource(provider, { resolvedApiKey }); const valueFingerprint = source?.valueFingerprint ?? source?.resolveValueFingerprint?.(); if (!source || !valueFingerprint || this.isProviderRequestAuthStale(provider, source)) { return undefined; }
return { provider, source: source.source, identityFingerprint: source.identityFingerprint, valueFingerprint, }; }
markProviderAuthSourceStale(token: AuthSourceToken): boolean { let marked = false; const providerRequestSource = this.getProviderRequestAuthSource(token.provider); if ( providerRequestSource?.source === token.source && providerRequestSource.identityFingerprint === token.identityFingerprint ) { const stale = this.staleProviderRequestAuthSources.get(token.provider) ?? []; if ( !stale.some( (existing) => existing.source === token.source && existing.identityFingerprint === token.identityFingerprint && existing.valueFingerprint === token.valueFingerprint, ) ) { stale.push(token); } this.staleProviderRequestAuthSources.set(token.provider, stale); marked = true; }
if (token.source !== "models_json_key" && token.source !== "models_json_command") { marked = this.authStorage.markAuthSourceStale(token) || marked; }
return marked; }
private getModelRequestKey(provider: string, modelId: string): string { return `${provider}:${modelId}`; }
private storeProviderRequestConfig( providerName: string, config: { apiKey?: string; headers?: Record<string, string>; authHeader?: boolean; }, ): void { if (!config.apiKey && !config.headers && !config.authHeader) { return; }
this.providerRequestConfigs.set(providerName, { apiKey: config.apiKey, headers: config.headers, authHeader: config.authHeader, }); }
private storeModelHeaders(providerName: string, modelId: string, headers?: Record<string, string>): void { const key = this.getModelRequestKey(providerName, modelId); if (!headers || Object.keys(headers).length === 0) { this.modelRequestHeaders.delete(key); return; } this.modelRequestHeaders.set(key, headers); }
/** * Get API key and request headers for a model. When `sessionId` is given, * stored multi-account credentials stay pinned to that session. */ async getApiKeyAndHeaders( model: Model<Api>, requestHeaders?: Record<string, string>, sessionId?: string, ): Promise<ResolvedRequestAuth> { try { const providerConfig = this.providerRequestConfigs.get(model.provider); const authStorageAuth = await this.authStorage.getApiKeyWithSourceToken(model.provider, { includeFallback: false, ...(sessionId ? { sessionId } : {}), }); let apiKey = authStorageAuth.apiKey; let authSourceToken = authStorageAuth.sourceToken; if (apiKey === undefined && providerConfig?.apiKey) { const resolvedApiKey = resolveConfigValueOrThrow( providerConfig.apiKey, `API key for provider "${model.provider}"`, ); const providerRequestAuthSource = this.getProviderRequestAuthSource(model.provider, { resolvedApiKey }); if ( providerRequestAuthSource && !this.isProviderRequestAuthStale(model.provider, providerRequestAuthSource) ) { this.clearStaleProviderRequestAuthSource(model.provider, providerRequestAuthSource); apiKey = resolvedApiKey; authSourceToken = this.getProviderRequestAuthSourceToken(model.provider, providerRequestAuthSource); } } if (!sessionId || authSourceToken?.source !== "stored") { this.setLastProviderAuthSourceToken(model.provider, apiKey === undefined ? undefined : authSourceToken); }
let requestModel: Model<Api> | undefined; if (model.provider === "xai") { if (!apiKey) { return { ok: false, error: "No usable xAI credential. Use /login and select an xAI entry to sign in again or configure an API key.", }; } const configuredModel = this.xaiModelSources.get(model) ?? model; const subscription = authStorageAuth.credentialType === "oauth" && authSourceToken?.source === "stored"; requestModel = subscription ? this.getXaiSubscriptionModel(configuredModel) : configuredModel; if (!requestModel) { return { ok: false, error: `Cannot configure "${model.id}" for Grok subscription requests. Select an xAI model or use /login and select the xAI API-key entry.`, }; } }
const providerHeaders = resolveHeadersOrThrow(providerConfig?.headers, `provider "${model.provider}"`); const authStorageHeaders = this.authStorage.getProviderHeaders(model.provider); const modelHeaders = resolveHeadersOrThrow( this.modelRequestHeaders.get(this.getModelRequestKey(model.provider, model.id)), `model "${model.provider}/${model.id}"`, );
// Providers that route by session identity need it on every request for this // session, including the auxiliary ones that never pass a StreamOptions.sessionId. const sessionHeaders = openCodeSessionHeaders(model.provider, sessionId); const configuredHeaders = model.headers || authStorageHeaders || providerHeaders || modelHeaders ? { ...model.headers, ...authStorageHeaders, ...providerHeaders, ...modelHeaders } : undefined; let headers = sessionHeaders || configuredHeaders ? { ...sessionHeaders, ...configuredHeaders } : undefined;
if (providerConfig?.authHeader) { if (!apiKey) { return { ok: false, error: `No API key found for "${model.provider}"` }; } headers = { ...headers, Authorization: `Bearer ${apiKey}` }; }
if (requestHeaders) headers = { ...headers, ...requestHeaders }; if ( model.provider === "xai" && authStorageAuth.credentialType === "oauth" && authSourceToken?.source === "stored" ) { for (const [name, value] of Object.entries(headers ?? {})) { if (name.toLowerCase() === "authorization" && value !== `Bearer ${apiKey}`) { return { ok: false, error: "Grok subscription cannot use a custom Authorization header. Remove the header or use /login and select the xAI API-key entry.", }; } } } return { ok: true, apiKey, headers: headers && Object.keys(headers).length > 0 ? headers : undefined, ...(requestModel ? { requestModel } : {}), }; } catch (error) { return { ok: false, error: error instanceof Error ? error.message : String(error), }; } }
/** * Return auth status for a provider, including request auth configured in models.json. * This intentionally does not execute command-backed config values. */ getProviderAuthStatus(provider: string): AuthStatus { const authStatus = this.authStorage.getAuthStatus(provider); if (authStatus.source && authStatus.source !== "stale") { return authStatus; }
const source = this.getProviderRequestAuthSource(provider); if (!source) { return authStatus; }
if (this.isProviderRequestAuthStaleForStatus(provider, source)) { return { configured: false, source: "stale", label: "expired" }; }
return { configured: true, source: source.source, ...(source.label ? { label: source.label } : {}), }; }
/** * Get display name for a provider. */ getProviderDisplayName(provider: string): string { const registeredProvider = this.registeredProviders.get(provider); const oauthProvider = this.authStorage.getOAuthProviders().find((p) => p.id === provider);
return ( registeredProvider?.name ?? registeredProvider?.oauth?.name ?? oauthProvider?.name ?? BUILT_IN_PROVIDER_DISPLAY_NAMES[provider] ?? provider ); }
/** * Get API key for a provider. */ async getApiKeyForProvider(provider: string): Promise<string | undefined> { const authStorageAuth = await this.authStorage.getApiKeyWithSourceToken(provider, { includeFallback: false }); if (authStorageAuth.apiKey !== undefined) { this.setLastProviderAuthSourceToken(provider, authStorageAuth.sourceToken); return authStorageAuth.apiKey; }
const providerApiKey = this.providerRequestConfigs.get(provider)?.apiKey; if (!providerApiKey) { this.setLastProviderAuthSourceToken(provider, undefined); return undefined; }
const resolvedApiKey = resolveConfigValueUncached(providerApiKey); if (resolvedApiKey === undefined) { this.setLastProviderAuthSourceToken(provider, undefined); return undefined; } const source = this.getProviderRequestAuthSource(provider, { resolvedApiKey }); if (!source || this.isProviderRequestAuthStale(provider, source)) { this.setLastProviderAuthSourceToken(provider, undefined); return undefined; } this.clearStaleProviderRequestAuthSource(provider, source); this.setLastProviderAuthSourceToken(provider, this.getProviderRequestAuthSourceToken(provider, source));
return resolvedApiKey; }
/** * Check if a model is using OAuth credentials (subscription). */ isUsingOAuth(model: Model<Api>): boolean { if (model.provider === "xai") return this.isUsingXaiSubscription(model); const cred = this.authStorage.get(model.provider); return cred?.type === "oauth"; }
/** * Register a provider dynamically (from extensions). * * If provider has models: replaces all existing models for this provider. * If provider has only baseUrl/headers: overrides existing models' URLs. * If provider has oauth: registers OAuth provider for /login support. */ registerProvider(providerName: string, config: ProviderConfigInput): void { this.validateProviderConfig(providerName, config); this.applyProviderConfig(providerName, config); this.upsertRegisteredProvider(providerName, config); // A provider that knows how to report its quota joins the same lookup the shipped // feeds live in, so /usage and `prime-agent usage --json` pick it up with no extra wiring. if (config.usage) { registerUsageFeed({ provider: providerName, providerName: config.name ?? providerName, credentialKinds: config.usage.credentialKinds ?? ["oauth", "apiKey", "session"], advance: config.usage.advance, }); } }
/** * Unregister a previously registered provider. * * Removes the provider from the registry and reloads models from disk so that * built-in models overridden by this provider are restored to their original state. * Also resets dynamic OAuth and API stream registrations before reapplying * remaining dynamic providers. * Has no effect if the provider was never registered. */ unregisterProvider(providerName: string): void { if (!this.registeredProviders.has(providerName)) return; this.registeredProviders.delete(providerName); unregisterUsageFeed(providerName); this.refresh(); }
/** * Upsert a provider config into registeredProviders. * If the provider is already registered, defined values in the incoming config * override existing ones; undefined values are preserved from the stored config. * If the provider is not registered, the incoming config is stored as-is. */ private upsertRegisteredProvider(providerName: string, config: ProviderConfigInput): void { const existing = this.registeredProviders.get(providerName); if (!existing) { this.registeredProviders.set(providerName, config); return; } for (const k of Object.keys(config) as (keyof ProviderConfigInput)[]) { if (config[k] !== undefined) { (existing as Record<string, unknown>)[k] = config[k]; } } }
private validateProviderConfig(providerName: string, config: ProviderConfigInput): void { if (config.streamSimple && !config.api) { throw new Error(`Provider ${providerName}: "api" is required when registering streamSimple.`); }
if (!config.models || config.models.length === 0) { return; }
if (!config.baseUrl) { throw new Error(`Provider ${providerName}: "baseUrl" is required when defining models.`); } if (!config.apiKey && !config.oauth) { throw new Error(`Provider ${providerName}: "apiKey" or "oauth" is required when defining models.`); }
for (const modelDef of config.models) { const api = modelDef.api || config.api; if (!api) { throw new Error(`Provider ${providerName}, model ${modelDef.id}: no "api" specified.`); } } }
private applyProviderConfig(providerName: string, config: ProviderConfigInput): void { if (config.oauth) { const oauthProvider: OAuthProviderInterface = { ...config.oauth, id: providerName, }; registerOAuthProvider(oauthProvider); }
if (config.streamSimple) { const streamSimple = config.streamSimple; registerApiProvider( { api: config.api!, stream: (model, context, options) => streamSimple(model, context, options as SimpleStreamOptions), streamSimple, }, `provider:${providerName}`, ); }
this.storeProviderRequestConfig(providerName, config);
if (config.models && config.models.length > 0) { this.models = this.models.filter((m) => m.provider !== providerName);
for (const modelDef of config.models) { const api = modelDef.api || config.api; this.storeModelHeaders(providerName, modelDef.id, modelDef.headers);
this.models.push( this.pinModelMetadata({ id: modelDef.id, name: modelDef.name, api: api as Api, provider: providerName, baseUrl: modelDef.baseUrl ?? config.baseUrl!, reasoning: modelDef.reasoning, thinkingLevelMap: modelDef.thinkingLevelMap, input: modelDef.input as ("text" | "image")[], cost: modelDef.cost, contextWindow: modelDef.contextWindow, maxTokens: modelDef.maxTokens, headers: undefined, compat: modelDef.compat, } as Model<Api>), ); } if (config.oauth?.modifyModels) { const cred = this.authStorage.get(providerName); if (cred?.type === "oauth") { this.models = config.oauth.modifyModels(this.models, cred).map((m) => this.pinModelMetadata(m)); } } } else if (config.baseUrl || config.headers) { this.models = this.models.map((m) => { if (m.provider !== providerName) return m; return { ...m, baseUrl: config.baseUrl ?? m.baseUrl, }; }); } }}
/** * Input type for registerProvider API. *//** * A quota feed contributed by a registered provider. The host drives it: `advance` returns either * the requests to send or the finished report, exactly like the shipped feeds * (see packages/ai/docs/usage-feeds.md). */export interface ProviderUsageConfig { /** Credential kinds the feed accepts. Defaults to all three. */ credentialKinds?: readonly CredentialKind[]; advance(state: UsageState, responses: readonly UsageResponse[]): UsageStep;}
export interface ProviderConfigInput { name?: string; baseUrl?: string; apiKey?: string; api?: Api; streamSimple?: (model: Model<Api>, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream; headers?: Record<string, string>; authHeader?: boolean; /** OAuth provider for /login support */ oauth?: Omit<OAuthProviderInterface, "id">; /** Quota feed for /usage. A provider without one is reported as skipped, with the reason. */ usage?: ProviderUsageConfig; models?: Array<{ id: string; name: string; api?: Api; baseUrl?: string; reasoning: boolean; thinkingLevelMap?: Model<Api>["thinkingLevelMap"]; input: ("text" | "image")[]; cost: { input: number; output: number; cacheRead: number; cacheWrite: number }; contextWindow: number; maxTokens: number; headers?: Record<string, string>; compat?: Model<Api>["compat"]; }>;}