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.
68 kB · 1877 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878/** * Model registry - manages built-in and custom models, provides API key resolution. */
import { Buffer } from "node:buffer";import { createHash } from "node:crypto";import { type AnthropicMessagesCompat, type Api, type AssistantMessageEventStream, type Context, fetchModelsDevModels, getLogger, getModels, getProviders, type KnownProvider, type Model, type OAuthProviderInterface, type OpenAICompletionsCompat, type OpenAIResponsesCompat, parsePrimeInferenceModelCatalog, registerApiProvider, registerLiveModels, resetApiProviders, type SimpleStreamOptions,} from "@earendil-works/pi-ai";import { registerBuiltinMcpOAuthProviders } from "@earendil-works/pi-ai/mcp";import { 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 { 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>; } | { 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 `find_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 { return createHash("sha256").update(apiKey).update("\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";}
/** * Model registry - loads and manages models, resolves API keys via AuthStorage. */export class ModelRegistry { private models: Model<Api>[] = []; 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 loadError: string | undefined = undefined; private modelsDevLiveModels: Model<Api>[] = []; private modelsDevRefreshTimer?: NodeJS.Timeout;
/** 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; }
/** 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; }
/** * Get only models that have auth configured. * This is a fast check that doesn't refresh OAuth tokens. */ getAvailable(): Model<Api>[] { return this.models.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) { void refreshPrimeInferenceModels(cachePath, this.bundledPrimeInferenceModels(), { offline: isOfflineModeEnabled(), }).then((models) => { if (!models) return; this.livePrimeInferenceModels = models; this.reloadModelsAfterCatalogChange(); }); } await this.refreshPrivatePrimeInferenceAuthorization( previousPrivateModelIds, previousTeamId, previousPrivateModels, ); return this.getAvailable(); }); }
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.models.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 = createHash("sha256").update(auth.apiKey).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 { return this.models.find((m) => m.provider === provider && m.id === modelId); }
/** * 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): boolean { if (this.authStorage.markAuthStale(provider)) { return true; }
const token = this.getCurrentProviderAuthSourceToken(provider); 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): boolean { return this.authStorage.blockLastSelectedCredential( provider, blockedUntilMs ?? Date.now() + DEFAULT_RATE_LIMIT_BLOCK_MS, "rate_limit", ); }
/** 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): AuthSourceToken | undefined { 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>, 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); } } this.setLastProviderAuthSourceToken(model.provider, apiKey === undefined ? undefined : authSourceToken);
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}"`, );
let headers = model.headers || authStorageHeaders || providerHeaders || modelHeaders ? { ...model.headers, ...authStorageHeaders, ...providerHeaders, ...modelHeaders } : undefined;
if (providerConfig?.authHeader) { if (!apiKey) { return { ok: false, error: `No API key found for "${model.provider}"` }; } headers = { ...headers, Authorization: `Bearer ${apiKey}` }; }
return { ok: true, apiKey, headers: headers && Object.keys(headers).length > 0 ? headers : undefined, }; } 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 { 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); }
/** * 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); 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({ 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); } } } 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. */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">; 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"]; }>;}