@astralyn/pi 0.83.0 → 0.84.0

package/dist/core/model-runtime.js

@@ -2,12 +2,23 @@ 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 { operationSignal, raceWithAbortSignal } from "../utils/abort.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";
+/** Credentials changed successfully, but the local model/auth snapshot could not be synchronized. */
+export class CredentialSynchronizationError extends Error {
+    constructor(providerId, operation, credential, options) {
+        super(`Credential ${operation} committed for ${providerId}, but local synchronization failed`, options);
+        this.name = "CredentialSynchronizationError";
+        this.providerId = providerId;
+        this.operation = operation;
+        this.credential = credential;
+    }
+}
 function mergeHeaders(base, override) {
     if (!base && !override)
         return undefined;
@@ -24,26 +35,22 @@ function mergeHeaders(base, override) {
 }
 /** 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.builtins = new Map();
+        this.nativeExtensionProviders = new Map();
+        this.extensionProviders = new Map();
+        this.compositionErrors = new Map();
+        this.snapshot = {
+            all: [],
+            available: [],
+            configuredProviders: new Set(),
+            storedProviders: new Set(),
+            auth: new Map(),
+        };
+        this.availabilityRefreshSeq = 0;
+        this.availabilityErrorSeq = 0;
+        this.providerAvailabilitySeq = new Map();
+        this.credentialOperations = new Map();
         this.credentials = credentials;
         this.config = config;
         this.modelsPath = modelsPath;
@@ -72,12 +79,15 @@ export class ModelRuntime {
         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;
+        const controller = refreshFromNetwork && options.modelRefreshTimeoutMs !== undefined ? new AbortController() : undefined;
+        const timeout = controller ? setTimeout(() => controller.abort(), options.modelRefreshTimeoutMs) : undefined;
+        const signal = controller
+            ? options.signal
+                ? AbortSignal.any([options.signal, controller.signal])
+                : controller.signal
+            : options.signal;
         try {
-            await runtime.refresh({ allowNetwork: refreshFromNetwork, signal: controller?.signal });
+            await runtime.refresh({ allowNetwork: refreshFromNetwork, signal });
         }
         finally {
             if (timeout)
@@ -149,16 +159,18 @@ export class ModelRuntime {
             available: all.filter((model) => this.snapshot.configuredProviders.has(model.provider)),
         };
     }
-    async runAvailabilityRefresh() {
+    async runAvailabilityRefresh(seq, errorSeq, signal) {
         const providers = this.models.getProviders();
         const [available, checks, credentials] = await Promise.all([
-            this.models.getAvailable(),
+            this.models.getAvailable(undefined, { signal }),
             Promise.all(providers.map(async (provider) => [
                 provider.id,
-                await this.models.checkAuth(provider.id),
+                await this.models.checkAuth(provider.id, { signal }),
             ])),
-            this.credentials.list(),
+            this.credentials.list({ signal }),
         ]);
+        if (seq !== this.availabilityRefreshSeq)
+            return;
         const auth = new Map(checks);
         const configuredProviders = new Set(checks
             .filter((entry) => entry[1] !== undefined)
@@ -170,28 +182,76 @@ export class ModelRuntime {
             storedProviders: new Set(credentials.map((entry) => entry.providerId)),
             auth,
         };
-        this.availabilityError = undefined;
+        if (errorSeq === this.availabilityErrorSeq)
+            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);
+    queueAvailabilityRefresh(signal) {
+        const seq = ++this.availabilityRefreshSeq;
+        for (const [providerId, providerSeq] of this.providerAvailabilitySeq) {
+            this.providerAvailabilitySeq.set(providerId, providerSeq + 1);
+        }
+        const errorSeq = ++this.availabilityErrorSeq;
+        const effectiveSignal = operationSignal(signal);
+        return this.runAvailabilityRefresh(seq, errorSeq, effectiveSignal).catch((error) => {
+            if (errorSeq === this.availabilityErrorSeq && !effectiveSignal.aborted) {
+                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);
+    async refreshProviderAvailability(providerId, signal) {
+        // Invalidate any full availability pass that started before this credential change.
+        ++this.availabilityRefreshSeq;
+        const providerSeq = (this.providerAvailabilitySeq.get(providerId) ?? 0) + 1;
+        this.providerAvailabilitySeq.set(providerId, providerSeq);
+        const errorSeq = ++this.availabilityErrorSeq;
+        try {
+            const [available, auth, credential] = await Promise.all([
+                this.models.getAvailable(providerId, { signal }),
+                this.models.checkAuth(providerId, { signal }),
+                this.credentials.read(providerId, { signal }),
+            ]);
+            signal.throwIfAborted();
+            if (this.providerAvailabilitySeq.get(providerId) !== providerSeq)
+                return;
+            const configuredProviders = new Set(this.snapshot.configuredProviders);
+            const storedProviders = new Set(this.snapshot.storedProviders);
+            const authByProvider = new Map(this.snapshot.auth);
+            if (auth) {
+                configuredProviders.add(providerId);
+                authByProvider.set(providerId, auth);
+            }
+            else {
+                configuredProviders.delete(providerId);
+                authByProvider.delete(providerId);
+            }
+            if (credential)
+                storedProviders.add(providerId);
+            else
+                storedProviders.delete(providerId);
+            const all = [...this.models.getModels()];
+            const availableById = new Map([...this.snapshot.available.filter((model) => model.provider !== providerId), ...available].map((model) => [
+                `${model.provider}\0${model.id}`,
+                model,
+            ]));
+            this.snapshot = {
+                all,
+                available: all.flatMap((model) => availableById.get(`${model.provider}\0${model.id}`) ?? []),
+                configuredProviders,
+                storedProviders,
+                auth: authByProvider,
+            };
+            if (errorSeq === this.availabilityErrorSeq)
+                this.availabilityError = undefined;
+        }
+        catch (error) {
+            if (this.providerAvailabilitySeq.get(providerId) === providerSeq &&
+                errorSeq === this.availabilityErrorSeq &&
+                !signal.aborted) {
+                this.availabilityError = error instanceof Error ? error.message : String(error);
+            }
+            throw error;
+        }
     }
     getProviders() {
         return this.models.getProviders();
@@ -205,24 +265,26 @@ export class ModelRuntime {
     getModel(providerId, modelId) {
         return this.models.getModel(providerId, modelId);
     }
-    async checkAuth(providerId) {
-        return this.models.checkAuth(providerId);
+    async checkAuth(providerId, options) {
+        return this.models.checkAuth(providerId, options);
     }
-    async getAvailable(providerId) {
+    async getAvailable(providerId, options) {
         if (providerId) {
-            if (this.availabilityRefresh) {
-                await this.availabilityRefresh;
-                return this.snapshot.available.filter((model) => model.provider === providerId);
-            }
+            const errorSeq = ++this.availabilityErrorSeq;
             try {
-                return await this.models.getAvailable(providerId);
+                const available = await this.models.getAvailable(providerId, options);
+                if (errorSeq === this.availabilityErrorSeq)
+                    this.availabilityError = undefined;
+                return available;
             }
             catch (error) {
-                this.availabilityError = error instanceof Error ? error.message : String(error);
+                if (errorSeq === this.availabilityErrorSeq && !options?.signal?.aborted) {
+                    this.availabilityError = error instanceof Error ? error.message : String(error);
+                }
                 throw error;
             }
         }
-        await this.refreshAvailability();
+        await this.queueAvailabilityRefresh(options?.signal);
         return this.snapshot.available;
     }
     getAvailableSnapshot() {
@@ -256,6 +318,9 @@ export class ModelRuntime {
     isUsingOAuth(providerId) {
         return this.snapshot.auth.get(providerId)?.type === "oauth";
     }
+    isUsingSubscription(providerId) {
+        return this.isUsingOAuth(providerId) && this.models.getProvider(providerId)?.auth.oauth?.isSubscription === true;
+    }
     hasConfiguredAuth(providerId) {
         return this.snapshot.configuredProviders.has(providerId);
     }
@@ -274,26 +339,62 @@ export class ModelRuntime {
             },
         };
     }
-    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);
+    enqueueCredentialOperation(providerId, signal, task) {
+        const previous = this.credentialOperations.get(providerId) ?? Promise.resolve();
+        let markStarted;
+        const started = new Promise((resolve) => {
+            markStarted = resolve;
+        });
+        const operation = (async () => {
+            await previous.catch(() => { });
+            signal.throwIfAborted();
+            markStarted?.();
+            return task();
+        })();
+        const tail = operation.catch(() => { });
+        this.credentialOperations.set(providerId, tail);
+        void tail.then(() => {
+            if (this.credentialOperations.get(providerId) === tail)
+                this.credentialOperations.delete(providerId);
+        });
+        return raceWithAbortSignal(started, signal).then(() => operation);
     }
-    async removeRuntimeApiKey(providerId) {
-        this.credentials.removeRuntimeApiKey(providerId);
-        await this.refresh({ allowNetwork: this.modelNetworkEnabled });
+    async synchronizeCredentialState(providerId, operation, credential, signal) {
+        try {
+            signal.throwIfAborted();
+            this.recomposeProvider(providerId);
+            const compositionError = this.compositionErrors.get(providerId);
+            if (compositionError)
+                throw new Error(compositionError);
+            const result = await this.models.refresh({ allowNetwork: false, providers: [providerId], signal });
+            if (result.aborted)
+                signal.throwIfAborted();
+            const refreshError = result.errors.get(providerId);
+            if (refreshError)
+                throw refreshError;
+            this.updateModelSnapshot();
+            await this.refreshProviderAvailability(providerId, signal);
+        }
+        catch (cause) {
+            throw new CredentialSynchronizationError(providerId, operation, credential, { cause });
+        }
     }
-    listCredentials() {
-        return this.credentials.list();
+    setRuntimeApiKey(providerId, apiKey, options = {}) {
+        const signal = operationSignal(options.signal);
+        return this.enqueueCredentialOperation(providerId, signal, async () => {
+            this.credentials.setRuntimeApiKey(providerId, apiKey);
+            await this.synchronizeCredentialState(providerId, "setRuntimeApiKey", { type: "api_key", key: apiKey }, signal);
+        });
+    }
+    removeRuntimeApiKey(providerId, options = {}) {
+        const signal = operationSignal(options.signal);
+        return this.enqueueCredentialOperation(providerId, signal, async () => {
+            this.credentials.removeRuntimeApiKey(providerId);
+            await this.synchronizeCredentialState(providerId, "removeRuntimeApiKey", undefined, signal);
+        });
+    }
+    listCredentials(options) {
+        return this.credentials.list(options);
     }
     getProviderAuthStatus(providerId) {
         if (this.credentials.hasRuntimeApiKey(providerId))
@@ -310,10 +411,15 @@ export class ModelRuntime {
         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 });
+        const resolution = await this.getAuth(model, {
+            apiKey: options?.apiKey,
+            env: options?.env,
+            signal: options?.signal,
+        });
         if (!resolution)
             throw new ModelsError("auth", `Provider is not configured: ${model.provider}`);
-        const { transformHeaders, ...providerOptions } = options ?? {};
+        const { transformHeaders, ...rawProviderOptions } = options ?? {};
+        const providerOptions = rawProviderOptions;
         let headers = mergeHeaders(resolution.auth.headers, providerOptions.headers);
         if (transformHeaders)
             headers = await transformHeaders(headers ?? {});
@@ -349,21 +455,48 @@ export class ModelRuntime {
     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 fetchDeferred(model, handle, options) {
+        return lazyStream(model, async () => {
+            const prepared = await this.prepareRequest(model, options);
+            if (!prepared.provider.fetchDeferred) {
+                throw new ModelsError("provider", `Provider ${model.provider} does not support deferred responses`);
+            }
+            return prepared.provider.fetchDeferred(prepared.model, handle, prepared.options);
+        }).result();
     }
-    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 cancelDeferred(model, handle, options) {
+        const prepared = await this.prepareRequest(model, options);
+        if (!prepared.provider.cancelDeferred) {
+            throw new ModelsError("provider", `Provider ${model.provider} does not support deferred responses`);
+        }
+        await prepared.provider.cancelDeferred(prepared.model, handle, prepared.options);
+    }
+    login(providerId, type, interaction) {
+        const signal = operationSignal(interaction.signal);
+        return this.enqueueCredentialOperation(providerId, signal, async () => {
+            const credential = await this.models.login(providerId, type, { ...interaction, signal });
+            await this.synchronizeCredentialState(providerId, "login", credential, signal);
+            return credential;
+        });
+    }
+    logout(providerId, options = {}) {
+        const signal = operationSignal(options.signal);
+        return this.enqueueCredentialOperation(providerId, signal, async () => {
+            await this.models.logout(providerId, { signal });
+            await this.synchronizeCredentialState(providerId, "logout", undefined, signal);
+        });
     }
     async refresh(options = {}) {
         this.config = await ModelConfig.load(this.modelsPath);
         this.configureRadiusProviders();
-        this.rebuildProviders();
+        if (options.providers) {
+            for (const providerId of new Set(options.providers))
+                this.recomposeProvider(providerId);
+            this.updateModelSnapshot();
+        }
+        else {
+            this.rebuildProviders();
+        }
         const refreshOptions = {
             ...options,
             allowNetwork: options.allowNetwork ?? this.modelNetworkEnabled,
@@ -374,14 +507,29 @@ export class ModelRuntime {
             aborted: refreshOptions.signal?.aborted ?? false,
             errors: new Map(),
         };
+        const errors = new Map(result.errors);
         this.updateModelSnapshot();
-        try {
-            await this.forceRefreshAvailability();
+        if (options.providers) {
+            await Promise.all([...new Set(options.providers)].map(async (providerId) => {
+                try {
+                    await this.refreshProviderAvailability(providerId, operationSignal(options.signal));
+                }
+                catch (error) {
+                    if (!options.signal?.aborted) {
+                        errors.set(providerId, error instanceof Error ? error : new Error(String(error)));
+                    }
+                }
+            }));
         }
-        catch {
-            // Availability errors are recorded by forceRefreshAvailability; refreshed models remain usable.
+        else {
+            try {
+                await this.queueAvailabilityRefresh(options.signal);
+            }
+            catch {
+                // Availability errors are recorded by the latest pass; refreshed models remain usable.
+            }
         }
-        return result;
+        return { aborted: result.aborted || (options.signal?.aborted ?? false), errors };
     }
     registerNativeProvider(provider) {
         if (!provider.id.trim())