@earendil-works/pi-coding-agent
Version:
Coding agent CLI with read, bash, edit, write tools and session management
439 lines • 19.6 kB
JavaScript
import { dirname, join } from "node:path";
import { createModels, lazyStream, ModelsError, } from "@earendil-works/pi-ai";
import * as builtinProviderCatalog from "@earendil-works/pi-ai/providers/all";
import { getAgentDir } from "../config.js";
import { AuthStorage as DefaultAuthStorage } from "./auth-storage.js";
import { ModelConfig } from "./model-config.js";
import { FileModelsStore, InMemoryCodingAgentModelsStore } from "./models-store.js";
import { composeModelProvider, configuredRequestAuthStatus, resolveCompatibilityRequestConfig, resolveConfiguredModelHeaders, validateExtensionProvider, } from "./provider-composer.js";
import { withRemoteCatalog } from "./remote-catalog-provider.js";
import { RuntimeCredentials } from "./runtime-credentials.js";
function mergeHeaders(base, override) {
if (!base && !override)
return undefined;
const merged = { ...base };
for (const [name, value] of Object.entries(override ?? {})) {
const lowerName = name.toLowerCase();
for (const existingName of Object.keys(merged)) {
if (existingName.toLowerCase() === lowerName)
delete merged[existingName];
}
merged[name] = value;
}
return merged;
}
/** Configured pi-ai Models collection used by coding-agent and SDK consumers. */
export class ModelRuntime {
models;
credentials;
defaultBuiltins;
builtins = new Map();
nativeExtensionProviders = new Map();
extensionProviders = new Map();
compositionErrors = new Map();
modelsPath;
modelNetworkEnabled;
config;
snapshot = {
all: [],
available: [],
configuredProviders: new Set(),
storedProviders: new Set(),
auth: new Map(),
};
availabilityRefresh;
availabilityError;
constructor(credentials, config, modelsPath, modelsStore, providers, modelNetworkEnabled) {
this.credentials = credentials;
this.config = config;
this.modelsPath = modelsPath;
this.modelNetworkEnabled = modelNetworkEnabled;
this.defaultBuiltins = new Map(providers.map((provider) => [provider.id, provider]));
for (const [providerId, provider] of this.defaultBuiltins)
this.builtins.set(providerId, provider);
this.models = createModels({ credentials, modelsStore });
this.rebuildProviders();
}
static async create(options = {}) {
const credentials = new RuntimeCredentials(options.credentials ?? DefaultAuthStorage.create(options.authPath));
const modelsPath = options.modelsPath === null ? undefined : (options.modelsPath ?? join(getAgentDir(), "models.json"));
const config = await ModelConfig.load(modelsPath);
const modelsStore = options.modelsStore ??
(modelsPath
? new FileModelsStore(options.modelsStorePath ?? join(dirname(modelsPath), "models-store.json"))
: new InMemoryCodingAgentModelsStore());
const builtinModelDataGeneratedAt = builtinProviderCatalog.getBuiltinModelDataGeneratedAt();
const providers = builtinProviderCatalog
.builtinProviders()
.map((provider) => provider.id === "radius"
? provider
: withRemoteCatalog(provider, options.catalogBaseUrl, builtinModelDataGeneratedAt));
const runtime = new ModelRuntime(credentials, config, modelsPath, modelsStore, providers, process.env.PI_OFFLINE === undefined);
runtime.configureRadiusProviders();
runtime.rebuildProviders();
const refreshFromNetwork = runtime.modelNetworkEnabled && options.allowModelNetwork === true;
const controller = refreshFromNetwork ? new AbortController() : undefined;
const timeout = controller
? setTimeout(() => controller.abort(), options.modelRefreshTimeoutMs ?? 15_000)
: undefined;
try {
await runtime.refresh({ allowNetwork: refreshFromNetwork, signal: controller?.signal });
}
finally {
if (timeout)
clearTimeout(timeout);
}
return runtime;
}
configureRadiusProviders() {
this.builtins.clear();
for (const [providerId, provider] of this.defaultBuiltins)
this.builtins.set(providerId, provider);
for (const providerId of this.config.getProviderIds()) {
const config = this.config.getProvider(providerId);
if (config?.oauth !== "radius" || !config.baseUrl)
continue;
this.builtins.set(providerId, builtinProviderCatalog.radiusProvider({
id: providerId,
name: config.name ?? providerId,
gateway: config.baseUrl.replace(/\/v1\/?$/u, ""),
}));
}
}
providerIds() {
return new Set([
...this.builtins.keys(),
...this.nativeExtensionProviders.keys(),
...this.config.getProviderIds(),
...this.extensionProviders.keys(),
]);
}
recomposeProvider(providerId) {
const base = this.nativeExtensionProviders.get(providerId) ?? this.builtins.get(providerId);
const extension = this.extensionProviders.get(providerId);
if (!base && !this.config.getProvider(providerId) && !extension) {
this.models.deleteProvider(providerId);
this.compositionErrors.delete(providerId);
return;
}
if (base && !this.config.getProvider(providerId) && !extension) {
// No overlays: use the builtin untouched so its auth/login/stream behavior is exact.
this.models.setProvider(base);
this.compositionErrors.delete(providerId);
return;
}
try {
this.models.setProvider(composeModelProvider(providerId, base, this.config, extension));
this.compositionErrors.delete(providerId);
}
catch (error) {
this.compositionErrors.set(providerId, error instanceof Error ? error.message : String(error));
if (base)
this.models.setProvider(base);
else
this.models.deleteProvider(providerId);
}
}
rebuildProviders() {
this.models.clearProviders();
this.compositionErrors.clear();
for (const providerId of this.providerIds())
this.recomposeProvider(providerId);
this.updateModelSnapshot();
}
updateModelSnapshot() {
const all = [...this.models.getModels()];
this.snapshot = {
...this.snapshot,
all,
available: all.filter((model) => this.snapshot.configuredProviders.has(model.provider)),
};
}
async runAvailabilityRefresh() {
const providers = this.models.getProviders();
const [available, checks, credentials] = await Promise.all([
this.models.getAvailable(),
Promise.all(providers.map(async (provider) => [
provider.id,
await this.models.checkAuth(provider.id),
])),
this.credentials.list(),
]);
const auth = new Map(checks);
const configuredProviders = new Set(checks
.filter((entry) => entry[1] !== undefined)
.map(([providerId]) => providerId));
this.snapshot = {
all: [...this.models.getModels()],
available: [...available],
configuredProviders,
storedProviders: new Set(credentials.map((entry) => entry.providerId)),
auth,
};
this.availabilityError = undefined;
}
queueAvailabilityRefresh(after) {
const refresh = (after ?? Promise.resolve()).catch(() => { }).then(() => this.runAvailabilityRefresh());
const recorded = refresh.catch((error) => {
this.availabilityError = error instanceof Error ? error.message : String(error);
throw error;
});
const tracked = recorded.finally(() => {
if (this.availabilityRefresh === tracked)
this.availabilityRefresh = undefined;
});
this.availabilityRefresh = tracked;
return tracked;
}
/** Coalesce concurrent readers onto the pending refresh. */
refreshAvailability() {
return this.availabilityRefresh ?? this.queueAvailabilityRefresh(undefined);
}
/** Mutations must not observe an in-flight refresh started before them. */
forceRefreshAvailability() {
return this.queueAvailabilityRefresh(this.availabilityRefresh);
}
getProviders() {
return this.models.getProviders();
}
getProvider(providerId) {
return this.models.getProvider(providerId);
}
getModels(providerId) {
return this.models.getModels(providerId);
}
getModel(providerId, modelId) {
return this.models.getModel(providerId, modelId);
}
async checkAuth(providerId) {
return this.models.checkAuth(providerId);
}
async getAvailable(providerId) {
if (providerId) {
if (this.availabilityRefresh) {
await this.availabilityRefresh;
return this.snapshot.available.filter((model) => model.provider === providerId);
}
try {
return await this.models.getAvailable(providerId);
}
catch (error) {
this.availabilityError = error instanceof Error ? error.message : String(error);
throw error;
}
}
await this.refreshAvailability();
return this.snapshot.available;
}
getAvailableSnapshot() {
return this.snapshot.available;
}
getError() {
const errors = [];
const configError = this.config.getError();
if (configError)
errors.push(configError);
for (const [providerId, error] of this.compositionErrors) {
errors.push(`Provider "${providerId}": ${error}`);
}
if (this.availabilityError)
errors.push(`Availability refresh: ${this.availabilityError}`);
return errors.length > 0 ? errors.join("\n\n") : undefined;
}
getRegisteredProviderConfig(providerId) {
return this.extensionProviders.get(providerId);
}
getRegisteredProviderIds() {
return [...new Set([...this.extensionProviders.keys(), ...this.nativeExtensionProviders.keys()])];
}
getRegisteredNativeProvider(providerId) {
return this.nativeExtensionProviders.get(providerId);
}
/** @internal Compatibility fallback for ModelRegistry when provider auth is unconfigured. */
getCompatibilityRequestConfig(model) {
return resolveCompatibilityRequestConfig(model, this.config.getProvider(model.provider), this.extensionProviders.get(model.provider));
}
isUsingOAuth(providerId) {
return this.snapshot.auth.get(providerId)?.type === "oauth";
}
hasConfiguredAuth(providerId) {
return this.snapshot.configuredProviders.has(providerId);
}
async getAuth(providerOrModel, overrides = {}) {
if (typeof providerOrModel === "string")
return this.models.getAuth(providerOrModel, overrides);
const resolution = await this.models.getAuth(providerOrModel, overrides);
if (!resolution)
return undefined;
const configuredHeaders = resolveConfiguredModelHeaders(providerOrModel, this.config.getProvider(providerOrModel.provider), this.extensionProviders.get(providerOrModel.provider), { ...(resolution.env ?? {}), ...(overrides.env ?? {}) });
return {
...resolution,
auth: {
...resolution.auth,
headers: mergeHeaders(resolution.auth.headers, configuredHeaders),
},
};
}
async setRuntimeApiKey(providerId, apiKey, refreshOptions = {}) {
this.credentials.setRuntimeApiKey(providerId, apiKey);
const auth = new Map(this.snapshot.auth).set(providerId, { type: "api_key", source: "runtime API key" });
const configuredProviders = new Set(this.snapshot.configuredProviders).add(providerId);
const storedProviders = new Set(this.snapshot.storedProviders).add(providerId);
this.snapshot = {
...this.snapshot,
auth,
configuredProviders,
storedProviders,
available: this.snapshot.all.filter((model) => configuredProviders.has(model.provider)),
};
await this.refresh(refreshOptions);
}
async removeRuntimeApiKey(providerId) {
this.credentials.removeRuntimeApiKey(providerId);
await this.refresh({ allowNetwork: this.modelNetworkEnabled });
}
listCredentials() {
return this.credentials.list();
}
getProviderAuthStatus(providerId) {
if (this.credentials.hasRuntimeApiKey(providerId))
return { configured: true, source: "runtime" };
if (this.snapshot.storedProviders.has(providerId))
return { configured: true, source: "stored" };
const configured = configuredRequestAuthStatus(this.config.getProvider(providerId), this.extensionProviders.get(providerId));
if (configured)
return configured;
const check = this.snapshot.auth.get(providerId);
return check ? { configured: true, source: "environment", label: check.source } : { configured: false };
}
async prepareRequest(model, options) {
const provider = this.models.getProvider(model.provider);
if (!provider)
throw new ModelsError("provider", `Unknown provider: ${model.provider}`);
const resolution = await this.getAuth(model, { apiKey: options?.apiKey, env: options?.env });
if (!resolution)
throw new ModelsError("auth", `Provider is not configured: ${model.provider}`);
const { transformHeaders, ...providerOptions } = options ?? {};
let headers = mergeHeaders(resolution.auth.headers, providerOptions.headers);
if (transformHeaders)
headers = await transformHeaders(headers ?? {});
const env = resolution.env || providerOptions.env
? { ...(resolution.env ?? {}), ...(providerOptions.env ?? {}) }
: undefined;
return {
provider,
model: resolution.auth.baseUrl ? { ...model, baseUrl: resolution.auth.baseUrl } : model,
options: {
...providerOptions,
apiKey: providerOptions.apiKey ?? resolution.auth.apiKey,
headers,
env,
},
};
}
stream(model, context, options) {
return lazyStream(model, async () => {
const prepared = await this.prepareRequest(model, options);
return prepared.provider.stream(prepared.model, context, prepared.options);
});
}
complete(model, context, options) {
return this.stream(model, context, options).result();
}
streamSimple(model, context, options) {
return lazyStream(model, async () => {
const prepared = await this.prepareRequest(model, options);
return prepared.provider.streamSimple(prepared.model, context, prepared.options);
});
}
completeSimple(model, context, options) {
return this.streamSimple(model, context, options).result();
}
async login(providerId, type, interaction) {
const credential = await this.models.login(providerId, type, interaction);
await this.refresh({ allowNetwork: this.modelNetworkEnabled });
return credential;
}
async logout(providerId) {
await this.models.logout(providerId);
// Reset credential-dependent compatibility projections before the unconfigured provider is skipped by refresh.
this.recomposeProvider(providerId);
await this.refresh({ allowNetwork: this.modelNetworkEnabled });
}
async refresh(options = {}) {
this.config = await ModelConfig.load(this.modelsPath);
this.configureRadiusProviders();
this.rebuildProviders();
const refreshOptions = {
...options,
allowNetwork: options.allowNetwork ?? this.modelNetworkEnabled,
};
// Published pi-ai builds before ModelsStore returned void and accepted a provider ID.
// The fallback keeps source-mode CLI tests working without rebuilding workspace dependencies.
const result = (await this.models.refresh(refreshOptions)) ?? {
aborted: refreshOptions.signal?.aborted ?? false,
errors: new Map(),
};
this.updateModelSnapshot();
try {
await this.forceRefreshAvailability();
}
catch {
// Availability errors are recorded by forceRefreshAvailability; refreshed models remain usable.
}
return result;
}
registerNativeProvider(provider) {
if (!provider.id.trim())
throw new Error("Provider id must not be empty.");
this.extensionProviders.delete(provider.id);
this.nativeExtensionProviders.set(provider.id, provider);
this.recomposeProvider(provider.id);
this.updateModelSnapshot();
void this.refresh({ allowNetwork: false });
}
registerProvider(providerId, config) {
// Validate the incoming registration on its own, like the legacy registry:
// a broken re-registration must throw without touching the stored config.
validateExtensionProvider(providerId, this.builtins.get(providerId), this.config.getProvider(providerId), config);
this.nativeExtensionProviders.delete(providerId);
// Re-registration merges defined values over the previous registration and
// preserves undefined ones, matching the legacy ModelRegistry contract.
const previous = this.extensionProviders.get(providerId);
const effective = { ...previous };
for (const [key, value] of Object.entries(config)) {
if (value !== undefined)
effective[key] = value;
}
this.extensionProviders.set(providerId, effective);
this.recomposeProvider(providerId);
this.updateModelSnapshot();
if (this.snapshot.storedProviders.has(providerId) ||
configuredRequestAuthStatus(this.config.getProvider(providerId), effective)?.configured) {
const configuredProviders = new Set(this.snapshot.configuredProviders).add(providerId);
const auth = new Map(this.snapshot.auth);
// Provisional entry until the async refresh lands; never clobber a real check result.
if (!auth.get(providerId)) {
auth.set(providerId, {
type: effective.oauth && !effective.apiKey ? "oauth" : "api_key",
source: "configured provider",
});
}
this.snapshot = {
...this.snapshot,
auth,
configuredProviders,
available: this.snapshot.all.filter((model) => configuredProviders.has(model.provider)),
};
}
void this.refresh({ allowNetwork: false });
}
unregisterProvider(providerId) {
this.extensionProviders.delete(providerId);
this.nativeExtensionProviders.delete(providerId);
this.recomposeProvider(providerId);
this.updateModelSnapshot();
void this.refresh({ allowNetwork: false });
}
}
//# sourceMappingURL=model-runtime.js.map