fix(coding-agent): make model refresh cancellation caller-owned

Add generation-safe catalog publication, cancellable auth and storage operations, locally consistent credential mutations, and cache-first interactive model flows.
This commit is contained in:
Mario Zechner 2026-08-04 00:00:12 +02:00
parent b06dc76fd7
commit fed6009cc9
86 changed files with 3426 additions and 870 deletions

View file

@ -2,13 +2,83 @@
## [Unreleased]
### Breaking Changes
- Required dynamic model providers to accept a concrete `RefreshModelsContext.signal`; `Models.refresh()` remains unbounded when callers omit its optional signal.
- Required provider login, API-key check/resolution, and OAuth refresh implementations to accept a concrete abort signal; public auth and credential operations remain unbounded when callers omit their optional signal.
- Replaced raw `RefreshModelsContext.store` access with the read-only `context.stored` snapshot and generation-checked `context.publish()` transaction.
**`createProvider({ fetchModels })`:** no catalog-publication migration is required. Before and after, return the fetched list; `createProvider()` restores stored models and publishes and persists refreshed models itself. `signal` is now guaranteed to be present.
```ts
// Before
const beforeProvider = createProvider({
// ...
fetchModels: async ({ signal }) => {
const response = await fetch(catalogUrl, { signal });
return parseModels(await response.json());
},
});
// After: unchanged
const afterProvider = createProvider({
// ...
fetchModels: async ({ signal }) => {
const response = await fetch(catalogUrl, { signal });
return parseModels(await response.json());
},
});
```
**Handwritten `Provider.refreshModels()`:** replace direct store access and pre-publication mutation with generation-guarded publications.
```ts
// Before
refreshModels: async (context) => {
const stored = await context.store.read();
if (stored) currentModels = stored.models;
if (!context.allowNetwork) return;
const refreshed = await fetchModels(context.signal);
currentModels = refreshed;
await context.store.write({ models: refreshed, checkedAt: Date.now() });
},
// After
refreshModels: async (context) => {
if (context.stored) {
const restored = context.stored.models;
if (!(await context.publish({
update: () => { currentModels = restored; },
}))) return;
}
if (!context.allowNetwork) return;
const refreshed = await fetchModels(context.signal);
if (context.signal.aborted) return;
await context.publish({
persist: { models: refreshed, checkedAt: Date.now() },
update: () => { currentModels = refreshed; },
});
},
```
In `publish()`, omit `persist` to leave storage unchanged, pass a `ModelsStoreEntry` to write it, or pass `persist: null` to delete it. Omit `update` for metadata-only persistence; omit `persist` for an ephemeral in-memory publication.
### Added
- Added Baseten as a built-in OpenAI-compatible provider with models.dev catalog generation and native `chat_template_args` reasoning controls.
### Changed
- Added optional cancellation to `ModelsStore` reads, writes, and deletions; catalog orchestration binds these waits to the provider refresh signal.
### Fixed
- Fixed tool argument validation to preserve values that already match an `anyOf`/`oneOf` union arm before attempting coercion, avoiding nullable unions converting `null` to another primitive value ([#7328](https://github.com/earendil-works/pi/issues/7328)).
- Fixed cancellation of model catalog refreshes so callers stop waiting even when a custom provider ignores its abort signal ([#7027](https://github.com/earendil-works/pi/issues/7027)).
- Fixed auth resolution, availability checks, OAuth refreshes, provider login, and in-memory credential queue waits to honor caller cancellation.
- Fixed newer provider refreshes being blocked by or overwritten by an older stalled generation, including persisted catalog publication.
- Updated GPT-5.6 Terra and Luna pricing across OpenAI and passthrough model catalogs.
- Fixed Fireworks Kimi K3 models to use the OpenAI-compatible API with native reasoning-effort levels and deferred tools ([#7199](https://github.com/earendil-works/pi/issues/7199), [#7230](https://github.com/earendil-works/pi/pull/7230) by [@XBeg9](https://github.com/XBeg9)).

View file

@ -313,8 +313,8 @@ Providers may have dynamic model lists (a llama.cpp server, a live OpenRouter li
```typescript
// getModels() returns the last-known list (empty before the first refresh)
await models.refresh('llamacpp'); // fetch one provider's list; rejects on failure
await models.refresh(); // refresh all providers concurrently, best-effort
await models.refresh({ providers: ['llamacpp'] }); // refresh one provider
await models.refresh(); // refresh all providers concurrently, best-effort
const fresh = models.getModel('llamacpp', 'qwen3-30b');
```
@ -352,6 +352,8 @@ if (modelAuth) {
Both overloads resolve credentials, refresh expired OAuth when necessary, and may return an auth-derived `apiKey`, `headers`, or `baseUrl`. `getAuth()` resolves `undefined` for unconfigured providers and rejects with `ModelsError` when something is actually broken (`"oauth"`: token refresh failed, credential preserved for re-login; `"auth"`: key resolution or credential store failure). Request paths surface the same failures as stream errors.
`getAuth()`, `checkAuth()`, `getAvailable()`, login, and logout accept optional caller cancellation through their existing options or interaction objects and remain unbounded when no signal is supplied. Provider `login`, `ApiKeyAuth.check`, `ApiKeyAuth.resolve`, and `OAuthAuth.refresh` implementations always receive a concrete signal and must honor it for blocking work.
### Transforming Request Headers
`Models.stream()`, `complete()`, `streamSimple()`, and `completeSimple()` accept a Models-only `transformHeaders` option. It runs once after provider auth, `model.headers`, and explicit `options.headers` have been merged, but before provider dispatch:
@ -388,7 +390,7 @@ const models = createModels({ credentials: myFileBackedStore });
// const models = builtinModels({ credentials: myFileBackedStore });
```
The contract is small: `read(providerId)`, `list()` for non-secret `{ providerId, type }` metadata, `modify(providerId, fn)` (the only write path — a serialized read-modify-write), and `delete(providerId)`. Enumeration must not resolve secrets or execute configured key commands. OAuth token refresh runs inside `modify`, so concurrent requests and processes cannot double-refresh a rotated token. A stored credential *owns* its provider: environment variables are only consulted when nothing is stored, and a failed refresh never silently falls back to an env key.
The contract is small: `read(providerId)`, `list()` for non-secret `{ providerId, type }` metadata, `modify(providerId, fn)` (the only write path — a serialized read-modify-write), and `delete(providerId)`. Each operation accepts optional cancellation options. Enumeration must not resolve secrets or execute configured key commands. OAuth token refresh runs inside `modify`, so concurrent requests and processes cannot double-refresh a rotated token. A stored credential *owns* its provider: environment variables are only consulted when nothing is stored, and a failed refresh never silently falls back to an env key.
API-key credentials use the same discriminator as pi's `auth.json` and can carry provider-scoped env/config values:
@ -1068,7 +1070,7 @@ const tenantGateway = createProvider({
});
```
Dynamic model lists use `fetchModels`. `Models.refresh()` refreshes every configured dynamic provider, passing its effective API-key or refreshed OAuth credential. A `ModelsStore` persists dynamic catalogs; both stores default to in-memory implementations.
Dynamic model lists use `fetchModels`. `Models.refresh()` refreshes every configured dynamic provider, passing its effective API-key or refreshed OAuth credential. A `ModelsStore` persists dynamic catalogs; both stores default to in-memory implementations. Its `read`, `write`, and `delete` operations accept optional cancellation, and `Models` binds those waits to the provider refresh signal.
```typescript
const models = createModels({ credentials, modelsStore });
@ -1086,7 +1088,11 @@ if (result.aborted) console.log('refresh cancelled');
for (const [provider, error] of result.errors) console.error(provider, error);
```
Use `models.refresh({ allowNetwork: false })` to restore persisted catalogs without network access, or `models.refresh({ force: true })` to bypass provider freshness checks. Model reads stay synchronous and return the last restored or refreshed list.
`Models.refresh()` is unbounded when its optional signal is omitted. Providers always receive a concrete `RefreshModelsContext.signal` and must honor it for network requests and other blocking work. When a caller supplies a signal, `Models.refresh()` returns promptly with `aborted: true` after cancellation even if a custom provider fails to cooperate; the provider must still honor the signal to stop its underlying work.
Use `models.refresh({ providers: ['openrouter'] })` to restrict work to selected providers, `models.refresh({ allowNetwork: false })` to restore persisted catalogs without network access, or `models.refresh({ force: true })` to bypass provider freshness checks. Model reads stay synchronous and return the last restored or refreshed list.
`createProvider()` handles dynamic publication and persistence automatically. Handwritten `Provider.refreshModels()` implementations receive the read-only `context.stored` snapshot and publish through `context.publish({ persist?, update? })`. Omit `persist` to leave storage unchanged, pass a `ModelsStoreEntry` to write it, or pass `persist: null` to delete it. Publication is generation-checked; put synchronous in-memory catalog changes in `update` rather than mutating state before publication.
Custom models can carry `headers` (e.g. proxies behind bot detection) and `compat` flags. `Models.getAuth(model)` includes those model headers, and stream methods merge them before explicit request headers and `transformHeaders`. See [OpenAI Compatibility Settings](#openai-compatibility-settings).
@ -1465,7 +1471,7 @@ Several providers support OAuth authentication instead of static API keys:
- **GitHub Copilot** (Copilot subscription)
- **OpenRouter** (OAuth PKCE that mints a user-controlled API key)
Each of these providers carries an `OAuthAuth` on `provider.auth.oauth` with three operations: `login(interaction)` uses the provider-neutral `AuthInteraction.prompt()`/`notify()` protocol and returns a credential, `refresh(credential)` refreshes expiring credentials when applicable, and `toAuth(credential)` derives request auth (GitHub Copilot's per-account base URL comes from here). Refresh is automatic: `models.getAuth(providerId)` and request paths refresh expired tokens under a credential-store lock, so concurrent requests and processes cannot double-refresh. OpenRouter's OAuth flow instead returns a permanent API key, so its refresh operation is a no-op.
Each of these providers carries an `OAuthAuth` on `provider.auth.oauth` with three operations: `login(interaction)` uses the provider-neutral `AuthInteraction.prompt()`/`notify()` protocol and returns a credential, `refresh(credential, signal)` refreshes expiring credentials when applicable, and `toAuth(credential)` derives request auth (GitHub Copilot's per-account base URL comes from here). Provider login interactions and refresh calls always carry a concrete abort signal. Refresh is automatic: `models.getAuth(providerId)` and request paths refresh expired tokens under a credential-store lock, so concurrent requests and processes cannot double-refresh. OpenRouter's OAuth flow instead returns a permanent API key, so its refresh operation is a no-op.
```typescript
import { createModels } from '@earendil-works/pi-ai';

View file

@ -1,4 +1,5 @@
import type { Credential, CredentialInfo, CredentialStore } from "./types.ts";
import { operationSignal, raceWithAbortSignal } from "../utils/abort.ts";
import type { AuthOperationOptions, Credential, CredentialInfo, CredentialStore } from "./types.ts";
/**
* Default in-memory credential store. Apps inject persistent stores.
@ -9,43 +10,58 @@ export class InMemoryCredentialStore implements CredentialStore {
private credentials = new Map<string, Credential>();
private chains = new Map<string, Promise<unknown>>();
/** Serialize tasks per provider id. */
private enqueue<T>(providerId: string, task: () => Promise<T>): Promise<T> {
/** Serialize tasks per provider id without releasing the chain before active work settles. */
private enqueue<T>(providerId: string, task: () => Promise<T>, options?: AuthOperationOptions): Promise<T> {
const signal = operationSignal(options?.signal);
const previous = this.chains.get(providerId) ?? Promise.resolve();
const next = (async () => {
const queued = (async () => {
await previous.catch(() => {});
signal.throwIfAborted();
return task();
})();
this.chains.set(
providerId,
next.catch(() => {}),
);
return next;
const tail = queued.catch(() => {});
this.chains.set(providerId, tail);
void tail.then(() => {
if (this.chains.get(providerId) === tail) this.chains.delete(providerId);
});
return raceWithAbortSignal(queued, signal);
}
async read(providerId: string): Promise<Credential | undefined> {
async read(providerId: string, options?: AuthOperationOptions): Promise<Credential | undefined> {
options?.signal?.throwIfAborted();
return this.credentials.get(providerId);
}
async list(): Promise<readonly CredentialInfo[]> {
async list(options?: AuthOperationOptions): Promise<readonly CredentialInfo[]> {
options?.signal?.throwIfAborted();
return [...this.credentials].map(([providerId, credential]) => ({ providerId, type: credential.type }));
}
modify(
providerId: string,
fn: (current: Credential | undefined) => Promise<Credential | undefined>,
options?: AuthOperationOptions,
): Promise<Credential | undefined> {
return this.enqueue(providerId, async () => {
const current = this.credentials.get(providerId);
const next = await fn(current);
if (next !== undefined) this.credentials.set(providerId, next);
return next ?? current;
});
return this.enqueue(
providerId,
async () => {
const current = this.credentials.get(providerId);
const next = await fn(current);
options?.signal?.throwIfAborted();
if (next !== undefined) this.credentials.set(providerId, next);
return next ?? current;
},
options,
);
}
delete(providerId: string): Promise<void> {
return this.enqueue(providerId, async () => {
this.credentials.delete(providerId);
});
delete(providerId: string, options?: AuthOperationOptions): Promise<void> {
return this.enqueue(
providerId,
async () => {
this.credentials.delete(providerId);
},
options,
);
}
}

View file

@ -10,15 +10,19 @@ export function envApiKeyAuth(name: string, envVars: readonly string[]): ApiKeyA
return {
name,
login: async (interaction) => {
interaction.signal.throwIfAborted();
const key = await interaction.prompt({ type: "secret", message: `Enter ${name}` });
interaction.signal.throwIfAborted();
return { type: "api_key", key };
},
resolve: async ({ ctx, credential }) => {
resolve: async ({ ctx, credential, signal }) => {
signal.throwIfAborted();
if (credential?.key) {
return { auth: { apiKey: credential.key }, env: credential.env, source: "stored credential" };
}
for (const envVar of envVars) {
const value = await ctx.env(envVar);
signal.throwIfAborted();
if (value) return { auth: { apiKey: value }, source: envVar };
}
return undefined;
@ -43,7 +47,7 @@ export function lazyOAuth(input: { name: string; loginLabel?: string; load: () =
name: input.name,
loginLabel: input.loginLabel,
login: async (interaction) => (await loaded()).login(interaction),
refresh: async (credential) => (await loaded()).refresh(credential),
refresh: async (credential, signal) => (await loaded()).refresh(credential, signal),
toAuth: async (credential) => (await loaded()).toAuth(credential),
};
}

View file

@ -7,7 +7,7 @@
import type { Server } from "node:http";
import { getProviderEnvValue } from "../../utils/provider-env.ts";
import type { AuthInteraction, OAuthAuth, OAuthCredential } from "../types.ts";
import type { OAuthAuth, OAuthCredential, ProviderAuthInteraction } from "../types.ts";
import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.ts";
import { generatePKCE } from "./pkce.ts";
@ -167,7 +167,7 @@ async function startCallbackServer(expectedState: string): Promise<CallbackServe
});
}
async function postJson(url: string, body: Record<string, string | number>): Promise<string> {
async function postJson(url: string, body: Record<string, string | number>, signal: AbortSignal): Promise<string> {
const response = await fetch(url, {
method: "POST",
headers: {
@ -175,7 +175,7 @@ async function postJson(url: string, body: Record<string, string | number>): Pro
Accept: "application/json",
},
body: JSON.stringify(body),
signal: AbortSignal.timeout(30_000),
signal: AbortSignal.any([signal, AbortSignal.timeout(30_000)]),
});
const responseBody = await response.text();
@ -192,17 +192,22 @@ async function exchangeAuthorizationCode(
state: string,
verifier: string,
redirectUri: string,
signal: AbortSignal,
): Promise<OAuthCredential> {
let responseBody: string;
try {
responseBody = await postJson(TOKEN_URL, {
grant_type: "authorization_code",
client_id: CLIENT_ID,
code,
state,
redirect_uri: redirectUri,
code_verifier: verifier,
});
responseBody = await postJson(
TOKEN_URL,
{
grant_type: "authorization_code",
client_id: CLIENT_ID,
code,
state,
redirect_uri: redirectUri,
code_verifier: verifier,
},
signal,
);
} catch (error) {
throw new Error(
`Token exchange request failed. url=${TOKEN_URL}; redirect_uri=${redirectUri}; response_type=authorization_code; details=${formatErrorDetails(error)}`,
@ -226,10 +231,13 @@ async function exchangeAuthorizationCode(
};
}
async function loginAnthropic(interaction: AuthInteraction): Promise<OAuthCredential> {
async function loginAnthropic(interaction: ProviderAuthInteraction): Promise<OAuthCredential> {
const { verifier, challenge } = await generatePKCE();
const server = await startCallbackServer(verifier);
const manualAbort = new AbortController();
const onAbort = () => server.cancelWait();
interaction.signal.addEventListener("abort", onAbort, { once: true });
if (interaction.signal.aborted) onAbort();
let code: string | undefined;
let state: string | undefined;
let manualInput: string | undefined;
@ -295,8 +303,9 @@ async function loginAnthropic(interaction: AuthInteraction): Promise<OAuthCreden
if (!code) throw new Error("Missing authorization code");
if (!state) throw new Error("Missing OAuth state");
interaction.notify({ type: "progress", message: "Exchanging authorization code for tokens..." });
return exchangeAuthorizationCode(code, state, verifier, REDIRECT_URI);
return exchangeAuthorizationCode(code, state, verifier, REDIRECT_URI, interaction.signal);
} finally {
interaction.signal.removeEventListener("abort", onAbort);
manualAbort.abort();
server.server.close();
}
@ -305,14 +314,18 @@ async function loginAnthropic(interaction: AuthInteraction): Promise<OAuthCreden
/**
* Refresh Anthropic OAuth token
*/
async function refreshAnthropicToken(refreshToken: string): Promise<OAuthCredential> {
async function refreshAnthropicToken(refreshToken: string, signal: AbortSignal): Promise<OAuthCredential> {
let responseBody: string;
try {
responseBody = await postJson(TOKEN_URL, {
grant_type: "refresh_token",
client_id: CLIENT_ID,
refresh_token: refreshToken,
});
responseBody = await postJson(
TOKEN_URL,
{
grant_type: "refresh_token",
client_id: CLIENT_ID,
refresh_token: refreshToken,
},
signal,
);
} catch (error) {
throw new Error(`Anthropic token refresh request failed. url=${TOKEN_URL}; details=${formatErrorDetails(error)}`);
}
@ -342,7 +355,7 @@ async function refreshAnthropicToken(refreshToken: string): Promise<OAuthCredent
export const anthropicOAuth: OAuthAuth = {
name: "Anthropic (Claude Pro/Max)",
login: loginAnthropic,
refresh: (credential) => refreshAnthropicToken(credential.refresh),
refresh: (credential, signal) => refreshAnthropicToken(credential.refresh, signal),
async toAuth(credential) {
return { apiKey: credential.access };

View file

@ -20,12 +20,12 @@ export type OAuthDeviceCodePollOptions<T> = {
expiresInSeconds?: number;
waitBeforeFirstPoll?: boolean;
poll: () => Promise<OAuthDeviceCodePollResult<T>>;
signal?: AbortSignal;
signal: AbortSignal;
};
function abortableSleep(ms: number, signal: AbortSignal | undefined, cancelMessage: string): Promise<void> {
function abortableSleep(ms: number, signal: AbortSignal, cancelMessage: string): Promise<void> {
return new Promise((resolve, reject) => {
if (signal?.aborted) {
if (signal.aborted) {
reject(new Error(cancelMessage));
return;
}
@ -35,11 +35,11 @@ function abortableSleep(ms: number, signal: AbortSignal | undefined, cancelMessa
reject(new Error(cancelMessage));
};
const timeout = setTimeout(() => {
signal?.removeEventListener("abort", onAbort);
signal.removeEventListener("abort", onAbort);
resolve();
}, ms);
signal?.addEventListener("abort", onAbort, { once: true });
signal.addEventListener("abort", onAbort, { once: true });
});
}
@ -62,7 +62,7 @@ export async function pollOAuthDeviceCodeFlow<T>(options: OAuthDeviceCodePollOpt
}
while (Date.now() < deadline) {
if (options.signal?.aborted) {
if (options.signal.aborted) {
throw new Error(CANCEL_MESSAGE);
}

View file

@ -3,7 +3,7 @@
*/
import { GITHUB_COPILOT_MODELS } from "../../providers/github-copilot.models.ts";
import type { AuthInteraction, OAuthAuth, OAuthCredential } from "../types.ts";
import type { OAuthAuth, OAuthCredential, ProviderAuthInteraction } from "../types.ts";
import { pollOAuthDeviceCodeFlow } from "./device-code.ts";
const decode = (s: string) => atob(s);
@ -113,7 +113,11 @@ function parseAvailableCopilotModelIds(raw: unknown): string[] {
return ids;
}
async function fetchAvailableGitHubCopilotModelIds(copilotToken: string, enterpriseDomain?: string): Promise<string[]> {
async function fetchAvailableGitHubCopilotModelIds(
copilotToken: string,
enterpriseDomain: string | undefined,
signal: AbortSignal,
): Promise<string[]> {
const baseUrl = getGitHubCopilotBaseUrl(copilotToken, enterpriseDomain);
const raw = await fetchJson(`${baseUrl}/models`, {
headers: {
@ -122,7 +126,7 @@ async function fetchAvailableGitHubCopilotModelIds(copilotToken: string, enterpr
...COPILOT_HEADERS,
"X-GitHub-Api-Version": COPILOT_API_VERSION,
},
signal: AbortSignal.timeout(5000),
signal: AbortSignal.any([signal, AbortSignal.timeout(5000)]),
});
return parseAvailableCopilotModelIds(raw);
}
@ -136,7 +140,7 @@ async function fetchJson(url: string, init: RequestInit): Promise<unknown> {
return response.json();
}
async function startDeviceFlow(domain: string): Promise<DeviceCodeResponse> {
async function startDeviceFlow(domain: string, signal: AbortSignal): Promise<DeviceCodeResponse> {
const urls = getUrls(domain);
const data = await fetchJson(urls.deviceCodeUrl, {
method: "POST",
@ -149,6 +153,7 @@ async function startDeviceFlow(domain: string): Promise<DeviceCodeResponse> {
client_id: CLIENT_ID,
scope: "read:user",
}),
signal,
});
if (!data || typeof data !== "object") {
@ -195,7 +200,7 @@ async function startDeviceFlow(domain: string): Promise<DeviceCodeResponse> {
async function pollForGitHubAccessToken(
domain: string,
device: DeviceCodeResponse,
signal?: AbortSignal,
signal: AbortSignal,
): Promise<string> {
const urls = getUrls(domain);
return pollOAuthDeviceCodeFlow<string>({
@ -216,6 +221,7 @@ async function pollForGitHubAccessToken(
device_code: device.device_code,
grant_type: "urn:ietf:params:oauth:grant-type:device_code",
}),
signal,
});
if (raw && typeof raw === "object" && typeof (raw as DeviceTokenSuccessResponse).access_token === "string") {
@ -243,7 +249,8 @@ async function pollForGitHubAccessToken(
async function refreshGitHubCopilotAccessToken(
refreshToken: string,
enterpriseDomain?: string,
enterpriseDomain: string | undefined,
signal: AbortSignal,
): Promise<OAuthCredential> {
const domain = enterpriseDomain || "github.com";
const urls = getUrls(domain);
@ -254,6 +261,7 @@ async function refreshGitHubCopilotAccessToken(
Authorization: `Bearer ${refreshToken}`,
...COPILOT_HEADERS,
},
signal,
});
if (!raw || typeof raw !== "object") {
@ -279,11 +287,15 @@ async function refreshGitHubCopilotAccessToken(
/**
* Refresh GitHub Copilot token
*/
async function refreshGitHubCopilotToken(refreshToken: string, enterpriseDomain?: string): Promise<OAuthCredential> {
const credentials = await refreshGitHubCopilotAccessToken(refreshToken, enterpriseDomain);
async function refreshGitHubCopilotToken(
refreshToken: string,
enterpriseDomain: string | undefined,
signal: AbortSignal,
): Promise<OAuthCredential> {
const credentials = await refreshGitHubCopilotAccessToken(refreshToken, enterpriseDomain, signal);
return {
...credentials,
availableModelIds: await fetchAvailableGitHubCopilotModelIds(credentials.access, enterpriseDomain),
availableModelIds: await fetchAvailableGitHubCopilotModelIds(credentials.access, enterpriseDomain, signal),
};
}
@ -291,7 +303,12 @@ async function refreshGitHubCopilotToken(refreshToken: string, enterpriseDomain?
* Enable a model for the user's GitHub Copilot account.
* This is required for some models (like Claude, Grok) before they can be used.
*/
async function enableGitHubCopilotModel(token: string, modelId: string, enterpriseDomain?: string): Promise<boolean> {
async function enableGitHubCopilotModel(
token: string,
modelId: string,
enterpriseDomain: string | undefined,
signal: AbortSignal,
): Promise<boolean> {
const baseUrl = getGitHubCopilotBaseUrl(token, enterpriseDomain);
const url = `${baseUrl}/models/${modelId}/policy`;
@ -306,9 +323,11 @@ async function enableGitHubCopilotModel(token: string, modelId: string, enterpri
"x-interaction-type": "chat-policy",
},
body: JSON.stringify({ state: "enabled" }),
signal,
});
return response.ok;
} catch {
} catch (error) {
if (signal.aborted) throw error;
return false;
}
}
@ -317,29 +336,33 @@ async function enableGitHubCopilotModel(token: string, modelId: string, enterpri
* Enable all known GitHub Copilot models that may require policy acceptance.
* Called after successful login to ensure all models are available.
*/
async function enableAllGitHubCopilotModels(token: string, enterpriseDomain?: string): Promise<void> {
async function enableAllGitHubCopilotModels(
token: string,
enterpriseDomain: string | undefined,
signal: AbortSignal,
): Promise<void> {
const models = Object.values(GITHUB_COPILOT_MODELS);
await Promise.all(
models.map(async (model) => {
await enableGitHubCopilotModel(token, model.id, enterpriseDomain);
await enableGitHubCopilotModel(token, model.id, enterpriseDomain, signal);
}),
);
}
async function loginGitHubCopilot(interaction: AuthInteraction): Promise<OAuthCredential> {
async function loginGitHubCopilot(interaction: ProviderAuthInteraction): Promise<OAuthCredential> {
const input = await interaction.prompt({
type: "text",
message: "GitHub Enterprise URL/domain (blank for github.com)",
placeholder: "company.ghe.com",
});
if (interaction.signal?.aborted) throw new Error("Login cancelled");
if (interaction.signal.aborted) throw new Error("Login cancelled");
const trimmed = input.trim();
const enterpriseDomain = normalizeDomain(input);
if (trimmed && !enterpriseDomain) throw new Error("Invalid GitHub Enterprise URL/domain");
const domain = enterpriseDomain || "github.com";
const device = await startDeviceFlow(domain);
const device = await startDeviceFlow(domain, interaction.signal);
interaction.notify({
type: "device_code",
userCode: device.user_code,
@ -349,12 +372,20 @@ async function loginGitHubCopilot(interaction: AuthInteraction): Promise<OAuthCr
});
const githubAccessToken = await pollForGitHubAccessToken(domain, device, interaction.signal);
const credentials = await refreshGitHubCopilotAccessToken(githubAccessToken, enterpriseDomain ?? undefined);
const credentials = await refreshGitHubCopilotAccessToken(
githubAccessToken,
enterpriseDomain ?? undefined,
interaction.signal,
);
interaction.notify({ type: "progress", message: "Enabling models..." });
await enableAllGitHubCopilotModels(credentials.access, enterpriseDomain ?? undefined);
await enableAllGitHubCopilotModels(credentials.access, enterpriseDomain ?? undefined, interaction.signal);
return {
...credentials,
availableModelIds: await fetchAvailableGitHubCopilotModelIds(credentials.access, enterpriseDomain ?? undefined),
availableModelIds: await fetchAvailableGitHubCopilotModelIds(
credentials.access,
enterpriseDomain ?? undefined,
interaction.signal,
),
};
}
@ -367,7 +398,8 @@ function copilotEnterpriseDomain(credential: OAuthCredential): string | undefine
export const githubCopilotOAuth: OAuthAuth = {
name: "GitHub Copilot",
login: loginGitHubCopilot,
refresh: (credential) => refreshGitHubCopilotToken(credential.refresh, copilotEnterpriseDomain(credential)),
refresh: (credential, signal) =>
refreshGitHubCopilotToken(credential.refresh, copilotEnterpriseDomain(credential), signal),
/** Derive the credential-specific proxy endpoint for each request. */
async toAuth(credential) {

View file

@ -7,7 +7,7 @@
*/
import { getProviderEnvValue } from "../../utils/provider-env.ts";
import type { AuthInteraction, OAuthAuth, OAuthCredential } from "../types.ts";
import type { OAuthAuth, OAuthCredential, ProviderAuthInteraction } from "../types.ts";
import { pollOAuthDeviceCodeFlow } from "./device-code.ts";
const CLIENT_ID = "17e5f671-d194-4dfb-9706-5516cb48c098";
@ -37,8 +37,8 @@ function getOauthHost(): string {
return (override || DEFAULT_OAUTH_HOST).replace(/\/+$/, "");
}
function requestSignal(signal?: AbortSignal): AbortSignal {
return AbortSignal.any([AbortSignal.timeout(REQUEST_TIMEOUT_MS), ...(signal ? [signal] : [])]);
function requestSignal(signal: AbortSignal): AbortSignal {
return AbortSignal.any([AbortSignal.timeout(REQUEST_TIMEOUT_MS), signal]);
}
function formUrlEncode(fields: Record<string, string>): string {
@ -66,7 +66,7 @@ function trustedHttpUrl(value: unknown): string | null {
}
}
async function startDeviceAuthorization(oauthHost: string, signal?: AbortSignal): Promise<DeviceAuthorization> {
async function startDeviceAuthorization(oauthHost: string, signal: AbortSignal): Promise<DeviceAuthorization> {
const response = await fetch(`${oauthHost}/api/oauth/device_authorization`, {
method: "POST",
headers: {
@ -141,7 +141,7 @@ function parseTokenResponse(json: Record<string, unknown> | null, operation: str
async function pollForToken(
oauthHost: string,
device: DeviceAuthorization,
signal?: AbortSignal,
signal: AbortSignal,
): Promise<TokenResponse> {
return pollOAuthDeviceCodeFlow<TokenResponse>({
intervalSeconds: device.intervalSeconds,
@ -206,25 +206,32 @@ async function pollForToken(
});
}
function sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
function sleep(ms: number, signal: AbortSignal): Promise<void> {
return new Promise((resolve, reject) => {
signal.throwIfAborted();
const onAbort = () => {
clearTimeout(timeout);
reject(signal.reason);
};
const timeout = setTimeout(() => {
signal.removeEventListener("abort", onAbort);
resolve();
}, ms);
signal.addEventListener("abort", onAbort, { once: true });
});
}
function isRetryableRefreshFailure(response: Response): boolean {
return response.status === 429 || response.status >= 500;
}
async function refreshToken(
oauthHost: string,
refreshTokenValue: string,
signal?: AbortSignal,
): Promise<TokenResponse> {
async function refreshToken(oauthHost: string, refreshTokenValue: string, signal: AbortSignal): Promise<TokenResponse> {
let lastError: Error | undefined;
for (let attempt = 0; attempt <= REFRESH_MAX_RETRIES; attempt++) {
if (attempt > 0) {
await sleep(1000 * 2 ** (attempt - 1));
await sleep(1000 * 2 ** (attempt - 1), signal);
}
if (signal?.aborted) {
if (signal.aborted) {
throw new Error("Kimi Code token refresh aborted");
}
@ -271,7 +278,7 @@ async function refreshToken(
throw lastError ?? new Error("Kimi Code token refresh failed");
}
async function loginKimiCoding(interaction: AuthInteraction): Promise<OAuthCredential> {
async function loginKimiCoding(interaction: ProviderAuthInteraction): Promise<OAuthCredential> {
const oauthHost = getOauthHost();
const device = await startDeviceAuthorization(oauthHost, interaction.signal);
interaction.notify({

View file

@ -18,7 +18,7 @@ if (typeof process !== "undefined" && (process.versions?.node || process.version
}
import { getProviderEnvValue } from "../../utils/provider-env.ts";
import type { AuthInteraction, OAuthAuth, OAuthCredential } from "../types.ts";
import type { OAuthAuth, OAuthCredential, ProviderAuthInteraction } from "../types.ts";
import { pollOAuthDeviceCodeFlow } from "./device-code.ts";
import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.ts";
import { generatePKCE } from "./pkce.ts";
@ -149,8 +149,8 @@ async function readTokenResponse(response: Response, operation: TokenOperation):
async function exchangeAuthorizationCode(
code: string,
verifier: string,
redirectUri: string = REDIRECT_URI,
signal?: AbortSignal,
redirectUri: string,
signal: AbortSignal,
): Promise<OAuthToken> {
const response = await fetchWithLoginCancellation(TOKEN_URL, {
method: "POST",
@ -168,7 +168,7 @@ async function exchangeAuthorizationCode(
return readTokenResponse(response, "exchange");
}
async function refreshAccessToken(refreshToken: string): Promise<OAuthToken> {
async function refreshAccessToken(refreshToken: string, signal: AbortSignal): Promise<OAuthToken> {
let response: Response;
try {
response = await fetch(TOKEN_URL, {
@ -179,6 +179,7 @@ async function refreshAccessToken(refreshToken: string): Promise<OAuthToken> {
refresh_token: refreshToken,
client_id: CLIENT_ID,
}),
signal,
});
} catch (error) {
throw new Error(`OpenAI Codex token refresh error: ${error instanceof Error ? error.message : String(error)}`);
@ -187,7 +188,7 @@ async function refreshAccessToken(refreshToken: string): Promise<OAuthToken> {
return readTokenResponse(response, "refresh");
}
async function startOpenAICodexDeviceAuth(signal?: AbortSignal): Promise<DeviceAuthInfo> {
async function startOpenAICodexDeviceAuth(signal: AbortSignal): Promise<DeviceAuthInfo> {
const response = await fetchWithLoginCancellation(DEVICE_USER_CODE_URL, {
method: "POST",
headers: { "Content-Type": "application/json" },
@ -231,7 +232,7 @@ async function startOpenAICodexDeviceAuth(signal?: AbortSignal): Promise<DeviceA
};
}
async function pollOpenAICodexDeviceAuth(device: DeviceAuthInfo, signal?: AbortSignal): Promise<DeviceTokenSuccess> {
async function pollOpenAICodexDeviceAuth(device: DeviceAuthInfo, signal: AbortSignal): Promise<DeviceTokenSuccess> {
return pollOAuthDeviceCodeFlow<DeviceTokenSuccess>({
intervalSeconds: device.intervalSeconds,
expiresInSeconds: DEVICE_CODE_TIMEOUT_SECONDS,
@ -418,12 +419,12 @@ async function exchangeAuthorizationCodeForCredentials(
code: string,
verifier: string,
redirectUri: string,
signal?: AbortSignal,
signal: AbortSignal,
): Promise<OAuthCredential> {
return credentialsFromToken(await exchangeAuthorizationCode(code, verifier, redirectUri, signal));
}
async function loginOpenAICodexDeviceCode(interaction: AuthInteraction): Promise<OAuthCredential> {
async function loginOpenAICodexDeviceCode(interaction: ProviderAuthInteraction): Promise<OAuthCredential> {
const device = await startOpenAICodexDeviceAuth(interaction.signal);
interaction.notify({
type: "device_code",
@ -441,10 +442,13 @@ async function loginOpenAICodexDeviceCode(interaction: AuthInteraction): Promise
);
}
async function loginOpenAICodex(interaction: AuthInteraction): Promise<OAuthCredential> {
async function loginOpenAICodex(interaction: ProviderAuthInteraction): Promise<OAuthCredential> {
const { verifier, state, url } = await createAuthorizationFlow();
const server = await startLocalOAuthServer(state);
const manualAbort = new AbortController();
const onAbort = () => server.cancelWait();
interaction.signal.addEventListener("abort", onAbort, { once: true });
if (interaction.signal.aborted) onAbort();
let code: string | undefined;
let manualCode: string | undefined;
let manualError: Error | undefined;
@ -495,6 +499,7 @@ async function loginOpenAICodex(interaction: AuthInteraction): Promise<OAuthCred
if (!code) throw new Error("Missing authorization code");
return exchangeAuthorizationCodeForCredentials(code, verifier, REDIRECT_URI, interaction.signal);
} finally {
interaction.signal.removeEventListener("abort", onAbort);
manualAbort.abort();
server.close();
}
@ -503,8 +508,8 @@ async function loginOpenAICodex(interaction: AuthInteraction): Promise<OAuthCred
/**
* Refresh OpenAI Codex OAuth token
*/
async function refreshOpenAICodexToken(refreshToken: string): Promise<OAuthCredential> {
return credentialsFromToken(await refreshAccessToken(refreshToken));
async function refreshOpenAICodexToken(refreshToken: string, signal: AbortSignal): Promise<OAuthCredential> {
return credentialsFromToken(await refreshAccessToken(refreshToken, signal));
}
export const openaiCodexOAuth: OAuthAuth = {
@ -530,7 +535,7 @@ export const openaiCodexOAuth: OAuthAuth = {
return loginOpenAICodex(interaction);
},
refresh: (credential) => refreshOpenAICodexToken(credential.refresh),
refresh: (credential, signal) => refreshOpenAICodexToken(credential.refresh, signal),
async toAuth(credential) {
return { apiKey: credential.access };

View file

@ -13,7 +13,7 @@
import { createServer, type Server, type ServerResponse } from "node:http";
import { getProviderEnvValue } from "../../utils/provider-env.ts";
import type { AuthInteraction, OAuthAuth, OAuthCredential } from "../types.ts";
import type { OAuthAuth, OAuthCredential, ProviderAuthInteraction } from "../types.ts";
import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.ts";
import { generatePKCE } from "./pkce.ts";
@ -80,12 +80,12 @@ function errorDetail(body: JsonObject): string | undefined {
async function exchangeAuthorizationCode(
code: string,
verifier: string,
signal?: AbortSignal,
signal: AbortSignal,
): Promise<OAuthCredential> {
if (signal?.aborted) throw new Error("Login cancelled");
if (signal.aborted) throw new Error("Login cancelled");
const controller = new AbortController();
const onAbort = () => controller.abort(signal?.reason);
signal?.addEventListener("abort", onAbort, { once: true });
const onAbort = () => controller.abort(signal.reason);
signal.addEventListener("abort", onAbort, { once: true });
const timeout = setTimeout(
() => controller.abort(new Error("OpenRouter OAuth token exchange timed out")),
TOKEN_EXCHANGE_TIMEOUT_MS,
@ -107,12 +107,12 @@ async function exchangeAuthorizationCode(
if (response.ok) throw new Error("OpenRouter OAuth returned invalid JSON");
}
} catch (error) {
if (signal?.aborted) throw new Error("Login cancelled");
if (signal.aborted) throw new Error("Login cancelled");
if (controller.signal.aborted) throw new Error("OpenRouter OAuth token exchange timed out");
throw error;
} finally {
clearTimeout(timeout);
signal?.removeEventListener("abort", onAbort);
signal.removeEventListener("abort", onAbort);
}
if (!response.ok) {
@ -135,9 +135,9 @@ async function exchangeAuthorizationCode(
async function startCallbackServer(
callbackPath: string,
verifier: string,
signal?: AbortSignal,
signal: AbortSignal,
): Promise<OpenRouterCallbackServer> {
if (signal?.aborted) throw new Error("Login cancelled");
if (signal.aborted) throw new Error("Login cancelled");
const callbackHost = getCallbackHost();
let resolveCredential: (credential: OAuthCredential | null) => void = () => {};
let rejectCredential: (error: Error) => void = () => {};
@ -154,7 +154,7 @@ async function startCallbackServer(
const close = (): void => {
if (timeout) clearTimeout(timeout);
if (onAbort) signal?.removeEventListener("abort", onAbort);
if (onAbort) signal.removeEventListener("abort", onAbort);
server.close();
};
@ -215,8 +215,8 @@ async function startCallbackServer(
server.on("error", (error) => finish({ error }));
onAbort = () => finish({ error: new Error("Login cancelled") });
signal?.addEventListener("abort", onAbort, { once: true });
if (signal?.aborted) {
signal.addEventListener("abort", onAbort, { once: true });
if (signal.aborted) {
close();
throw new Error("Login cancelled");
}
@ -239,7 +239,7 @@ async function startCallbackServer(
};
}
async function loginOpenRouter(interaction: AuthInteraction): Promise<OAuthCredential> {
async function loginOpenRouter(interaction: ProviderAuthInteraction): Promise<OAuthCredential> {
const { verifier, challenge } = await generatePKCE();
const callbackPath = `/oauth/callback/${crypto.randomUUID()}`;
const callback = await startCallbackServer(callbackPath, verifier, interaction.signal);
@ -302,7 +302,7 @@ export const openRouterOAuth: OAuthAuth = {
name: "OpenRouter OAuth",
loginLabel: "Sign in with OpenRouter",
login: loginOpenRouter,
async refresh(credential) {
async refresh(credential, _signal) {
return credential;
},
async toAuth(credential) {

View file

@ -18,7 +18,7 @@ if (typeof process !== "undefined" && (process.versions?.node || process.version
}
import { normalizeRadiusGatewayUrl } from "../../providers/radius-config.ts";
import type { AuthInteraction, OAuthAuth, OAuthCredential } from "../types.ts";
import type { OAuthAuth, OAuthCredential, ProviderAuthInteraction } from "../types.ts";
import { pollOAuthDeviceCodeFlow } from "./device-code.ts";
import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.ts";
import { generatePKCE } from "./pkce.ts";
@ -46,9 +46,10 @@ type DeviceAuthorizationResponse = {
interval?: number;
};
async function loadRadiusOAuthDiscovery(gateway: string): Promise<RadiusOAuthDiscovery> {
async function loadRadiusOAuthDiscovery(gateway: string, signal: AbortSignal): Promise<RadiusOAuthDiscovery> {
const response = await fetch(new URL("/v1/oauth", gateway), {
headers: { accept: "application/json" },
signal,
});
if (!response.ok) {
@ -101,7 +102,7 @@ async function readOAuthResponseError(response: Response, message: string): Prom
async function requestOAuthToken(
gateway: string,
body: URLSearchParams,
signal?: AbortSignal,
signal: AbortSignal,
): Promise<OAuthCredential> {
let response: Response;
try {
@ -112,7 +113,7 @@ async function requestOAuthToken(
signal,
});
} catch (error) {
if (signal?.aborted) {
if (signal.aborted) {
throw new Error("Login cancelled");
}
throw error;
@ -143,10 +144,7 @@ type OAuthCallbackServer = {
close(): void;
};
function startOAuthCallbackServer(
expectedState: string,
signal: AbortSignal | undefined,
): Promise<OAuthCallbackServer> {
function startOAuthCallbackServer(expectedState: string, signal: AbortSignal): Promise<OAuthCallbackServer> {
if (!_http) {
throw new Error("Radius OAuth is only available in Node.js environments");
}
@ -161,11 +159,11 @@ function startOAuthCallbackServer(
return;
}
settled = true;
signal?.removeEventListener("abort", onAbort);
signal.removeEventListener("abort", onAbort);
settle(code);
};
const onAbort = () => finish(null);
signal?.addEventListener("abort", onAbort, { once: true });
signal.addEventListener("abort", onAbort, { once: true });
const sendPage = (response: import("node:http").ServerResponse, status: number, html: string) => {
response.statusCode = status;
@ -222,7 +220,7 @@ function startOAuthCallbackServer(
async function loginWithBrowser(
gateway: string,
authorizationEndpoint: string,
interaction: AuthInteraction,
interaction: ProviderAuthInteraction,
): Promise<OAuthCredential> {
const { verifier, challenge } = await generatePKCE();
const state = crypto.randomUUID();
@ -249,7 +247,7 @@ async function loginWithBrowser(
try {
const code = await callbackServer.waitForCode();
if (!code) {
if (interaction.signal?.aborted) {
if (interaction.signal.aborted) {
throw new Error("Login cancelled");
}
throw new Error("OAuth callback did not complete.");
@ -270,10 +268,7 @@ async function loginWithBrowser(
}
}
async function requestDeviceAuthorization(
gateway: string,
signal: AbortSignal | undefined,
): Promise<DeviceAuthorizationResponse> {
async function requestDeviceAuthorization(gateway: string, signal: AbortSignal): Promise<DeviceAuthorizationResponse> {
let response: Response;
try {
response = await fetch(new URL("/v1/oauth/device", gateway), {
@ -283,7 +278,7 @@ async function requestDeviceAuthorization(
signal,
});
} catch (error) {
if (signal?.aborted) {
if (signal.aborted) {
throw new Error("Login cancelled");
}
throw error;
@ -307,7 +302,7 @@ async function requestDeviceAuthorization(
};
}
async function loginWithDeviceCode(gateway: string, interaction: AuthInteraction): Promise<OAuthCredential> {
async function loginWithDeviceCode(gateway: string, interaction: ProviderAuthInteraction): Promise<OAuthCredential> {
const device = await requestDeviceAuthorization(gateway, interaction.signal);
interaction.notify({
type: "device_code",
@ -382,7 +377,7 @@ export function createRadiusOAuth(options: RadiusOAuthOptions): OAuthAuth {
return loginWithDeviceCode(gateway, interaction);
}
if (loginMethod === LOGIN_METHOD_BROWSER) {
const discovery = await loadRadiusOAuthDiscovery(gateway);
const discovery = await loadRadiusOAuthDiscovery(gateway, interaction.signal);
return loginWithBrowser(gateway, discovery.authorizationEndpoint, interaction);
}
throw new Error(`Unknown ${options.name} sign-in method: ${loginMethod}`);

View file

@ -2,7 +2,7 @@
* xAI OAuth device-code flow.
*/
import type { AuthInteraction, OAuthAuth, OAuthCredential } from "../types.ts";
import type { OAuthAuth, OAuthCredential, ProviderAuthInteraction } from "../types.ts";
import { pollOAuthDeviceCodeFlow } from "./device-code.ts";
const XAI_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828";
@ -61,7 +61,7 @@ function validateVerificationUri(raw: string): string {
return url.href;
}
async function postForm(url: string, fields: Record<string, string>, signal?: AbortSignal): Promise<OAuthHttpResponse> {
async function postForm(url: string, fields: Record<string, string>, signal: AbortSignal): Promise<OAuthHttpResponse> {
let response: Response;
try {
response = await fetch(url, {
@ -74,7 +74,7 @@ async function postForm(url: string, fields: Record<string, string>, signal?: Ab
signal,
});
} catch (error) {
if (signal?.aborted) {
if (signal.aborted) {
throw new Error("Login cancelled");
}
throw error;
@ -85,7 +85,7 @@ async function postForm(url: string, fields: Record<string, string>, signal?: Ab
const parsed = (await response.json()) as unknown;
body = parsed && typeof parsed === "object" && !Array.isArray(parsed) ? (parsed as JsonObject) : {};
} catch {
if (signal?.aborted) {
if (signal.aborted) {
throw new Error("Login cancelled");
}
throw new Error(`xAI OAuth returned invalid JSON (HTTP ${response.status})`);
@ -142,7 +142,7 @@ function credentialsFromTokenResponse(body: JsonObject, previousRefreshToken?: s
};
}
async function requestDeviceCode(signal?: AbortSignal): Promise<XaiDeviceCode> {
async function requestDeviceCode(signal: AbortSignal): Promise<XaiDeviceCode> {
const response = await postForm(
XAI_DEVICE_CODE_URL,
{
@ -158,7 +158,7 @@ async function requestDeviceCode(signal?: AbortSignal): Promise<XaiDeviceCode> {
return parseDeviceCode(response.body);
}
async function pollForTokens(device: XaiDeviceCode, signal?: AbortSignal): Promise<OAuthCredential> {
async function pollForTokens(device: XaiDeviceCode, signal: AbortSignal): Promise<OAuthCredential> {
return pollOAuthDeviceCodeFlow<OAuthCredential>({
intervalSeconds: device.intervalSeconds,
expiresInSeconds: device.expiresInSeconds,
@ -198,7 +198,7 @@ async function pollForTokens(device: XaiDeviceCode, signal?: AbortSignal): Promi
});
}
async function loginXai(interaction: AuthInteraction): Promise<OAuthCredential> {
async function loginXai(interaction: ProviderAuthInteraction): Promise<OAuthCredential> {
const device = await requestDeviceCode(interaction.signal);
interaction.notify({
type: "device_code",
@ -210,7 +210,7 @@ async function loginXai(interaction: AuthInteraction): Promise<OAuthCredential>
return pollForTokens(device, interaction.signal);
}
async function refreshXaiToken(refreshToken: string, signal?: AbortSignal): Promise<OAuthCredential> {
async function refreshXaiToken(refreshToken: string, signal: AbortSignal): Promise<OAuthCredential> {
const response = await postForm(
XAI_TOKEN_URL,
{

View file

@ -1,4 +1,5 @@
import type { ProviderEnv } from "../types.ts";
import { operationSignal, raceWithAbortSignal } from "../utils/abort.ts";
import { formatThrownValue } from "../utils/diagnostics.ts";
import type {
ApiKeyAuth,
@ -19,6 +20,7 @@ export interface AuthResolutionOverrides {
env?: ProviderEnv;
/** Require this much remaining OAuth-token validity; defaults to five minutes. */
minOAuthValidityMs?: number;
signal?: AbortSignal;
}
export class ModelsError extends Error {
@ -45,23 +47,44 @@ function withCauseDetail(message: string, cause: unknown): string {
* nothing is stored. No silent env fallback after a failed refresh or for a
* credential type without a matching handler.
*/
export async function resolveProviderAuth(
export function resolveProviderAuth(
provider: { id: string; auth: ProviderAuth },
credentials: CredentialStore,
authContext: AuthContext,
overrides?: AuthResolutionOverrides,
): Promise<AuthResult | undefined> {
const signal = operationSignal(overrides?.signal);
return raceWithAbortSignal(
resolveProviderAuthWithSignal(provider, credentials, authContext, overrides, signal),
signal,
);
}
async function resolveProviderAuthWithSignal(
provider: { id: string; auth: ProviderAuth },
credentials: CredentialStore,
authContext: AuthContext,
overrides: AuthResolutionOverrides | undefined,
signal: AbortSignal,
): Promise<AuthResult | undefined> {
signal.throwIfAborted();
const requestAuthContext = overrides?.env ? overlayEnvAuthContext(authContext, overrides.env) : authContext;
if (overrides?.apiKey !== undefined && provider.auth.apiKey) {
return resolveApiKey(requestAuthContext, provider.auth.apiKey, provider.id, {
type: "api_key",
key: overrides.apiKey,
env: overrides.env,
});
return resolveApiKey(
requestAuthContext,
provider.auth.apiKey,
provider.id,
{
type: "api_key",
key: overrides.apiKey,
env: overrides.env,
},
signal,
);
}
const stored = await readCredential(credentials, provider.id);
const stored = await readCredential(credentials, provider.id, signal);
if (stored) {
if (stored.type === "oauth" && provider.auth.oauth) {
return resolveStoredOAuth(
@ -69,19 +92,20 @@ export async function resolveProviderAuth(
provider.id,
provider.auth.oauth,
stored,
signal,
overrides?.minOAuthValidityMs,
);
}
if (stored.type === "api_key" && provider.auth.apiKey) {
const credential = overrides?.env ? { ...stored, env: { ...stored.env, ...overrides.env } } : stored;
return resolveApiKey(requestAuthContext, provider.auth.apiKey, provider.id, credential);
return resolveApiKey(requestAuthContext, provider.auth.apiKey, provider.id, credential, signal);
}
return undefined;
}
// Ambient (env vars, AWS profiles, ADC files).
return provider.auth.apiKey
? resolveApiKey(requestAuthContext, provider.auth.apiKey, provider.id, undefined)
? resolveApiKey(requestAuthContext, provider.auth.apiKey, provider.id, undefined, signal)
: undefined;
}
@ -104,6 +128,7 @@ async function resolveStoredOAuth(
providerId: string,
oauth: OAuthAuth,
stored: OAuthCredential,
signal: AbortSignal,
minOAuthValidityMs?: number,
): Promise<AuthResult | undefined> {
const minimumValidityMs = Math.max(DEFAULT_OAUTH_MINIMUM_VALIDITY_MS, minOAuthValidityMs ?? 0);
@ -114,15 +139,19 @@ async function resolveStoredOAuth(
// Optimistic check said expired; the authoritative check runs under the lock.
let post: Credential | undefined;
try {
post = await credentials.modify(providerId, async (current) => {
if (current?.type !== "oauth") return undefined; // logged out meanwhile
if (!expiresSoon(current)) return undefined; // another process/request refreshed
try {
return await oauth.refresh(current);
} catch (error) {
throw new ModelsError("oauth", `OAuth refresh failed for ${providerId}`, { cause: error });
}
});
post = await credentials.modify(
providerId,
async (current) => {
if (current?.type !== "oauth") return undefined; // logged out meanwhile
if (!expiresSoon(current)) return undefined; // another process/request refreshed
try {
return await oauth.refresh(current, signal);
} catch (error) {
throw new ModelsError("oauth", `OAuth refresh failed for ${providerId}`, { cause: error });
}
},
{ signal },
);
} catch (error) {
if (error instanceof ModelsError) throw error;
throw new ModelsError("auth", `Credential store modify failed for ${providerId}`, { cause: error });
@ -149,17 +178,22 @@ async function resolveApiKey(
apiKey: ApiKeyAuth,
providerId: string,
credential: ApiKeyCredential | undefined,
signal: AbortSignal,
): Promise<AuthResult | undefined> {
try {
return await apiKey.resolve({ ctx: authContext, credential });
return await apiKey.resolve({ ctx: authContext, credential, signal });
} catch (error) {
throw new ModelsError("auth", `API key auth failed for provider ${providerId}`, { cause: error });
}
}
async function readCredential(credentials: CredentialStore, providerId: string): Promise<Credential | undefined> {
async function readCredential(
credentials: CredentialStore,
providerId: string,
signal: AbortSignal,
): Promise<Credential | undefined> {
try {
return await credentials.read(providerId);
return await credentials.read(providerId, { signal });
} catch (error) {
throw new ModelsError("auth", `Credential store read failed for ${providerId}`, { cause: error });
}

View file

@ -42,6 +42,11 @@ export interface CredentialInfo {
type: Credential["type"];
}
/** Optional cancellation for public auth and credential operations. */
export interface AuthOperationOptions {
signal?: AbortSignal;
}
/**
* App-owned credential storage, keyed by `Provider.id`, one credential per
* provider. `modify` is the only write path, so every mutation is a
@ -62,13 +67,13 @@ export interface CredentialStore {
* Read the stored credential, possibly expired. Display/status use;
* resolved request auth comes from `Models.getAuth()`.
*/
read(providerId: string): Promise<Credential | undefined>;
read(providerId: string, options?: AuthOperationOptions): Promise<Credential | undefined>;
/**
* List stored credential metadata without resolving or exposing secrets.
* Implementations must not execute configured API-key commands while listing.
*/
list(): Promise<readonly CredentialInfo[]>;
list(options?: AuthOperationOptions): Promise<readonly CredentialInfo[]>;
/**
* Serialized write the only write path. `fn` sees the current credential
@ -81,10 +86,11 @@ export interface CredentialStore {
modify(
providerId: string,
fn: (current: Credential | undefined) => Promise<Credential | undefined>,
options?: AuthOperationOptions,
): Promise<Credential | undefined>;
/** Remove a credential (logout). Implementations serialize this against `modify`. */
delete(providerId: string): Promise<void>;
delete(providerId: string, options?: AuthOperationOptions): Promise<void>;
}
/** Environment access for auth resolution. Injectable for tests and browsers. */
@ -154,6 +160,9 @@ export interface AuthInteraction {
notify(event: AuthEvent): void;
}
/** Normalized interaction passed to provider login implementations. */
export type ProviderAuthInteraction = AuthInteraction & { signal: AbortSignal };
/**
* Api-key auth: stored key/provider env plus ambient sources (env vars, AWS
* profiles, ADC files). Ambient-only providers omit `login`.
@ -163,14 +172,18 @@ export interface ApiKeyAuth {
name: string;
/** Interactive setup (prompt for key/provider env). Absent = ambient-only. */
login?(interaction: AuthInteraction): Promise<ApiKeyCredential>;
login?(interaction: ProviderAuthInteraction): Promise<ApiKeyCredential>;
/**
* Optional side-effect-free availability check. Use this when `resolve()` may
* execute commands or perform other request-time work. Missing means Models
* checks availability by resolving auth.
*/
check?(input: { ctx: AuthContext; credential?: ApiKeyCredential }): Promise<AuthCheck | undefined>;
check?(input: {
ctx: AuthContext;
credential?: ApiKeyCredential;
signal: AbortSignal;
}): Promise<AuthCheck | undefined>;
/**
* Resolve auth from the stored credential and/or ambient sources, merging
@ -178,7 +191,11 @@ export interface ApiKeyAuth {
* undefined = not configured. Resolution is provider-scoped; model-specific
* endpoint preparation happens after auth has been resolved.
*/
resolve(input: { ctx: AuthContext; credential?: ApiKeyCredential }): Promise<AuthResult | undefined>;
resolve(input: {
ctx: AuthContext;
credential?: ApiKeyCredential;
signal: AbortSignal;
}): Promise<AuthResult | undefined>;
}
/**
@ -193,13 +210,13 @@ export interface OAuthAuth {
/** Selector label for the subscription login option, e.g. "Sign in with SuperGrok or X Premium". */
loginLabel?: string;
login(interaction: AuthInteraction): Promise<OAuthCredential>;
login(interaction: ProviderAuthInteraction): Promise<OAuthCredential>;
/**
* Exchange the refresh token. Network call; throws on failure
* (invalid_grant etc.). `Models` runs this under the store lock.
*/
refresh(credential: OAuthCredential, signal?: AbortSignal): Promise<OAuthCredential>;
refresh(credential: OAuthCredential, signal: AbortSignal): Promise<OAuthCredential>;
/**
* Side-effect-free derivation of request auth from a valid credential.

View file

@ -48,6 +48,7 @@ async function login(providerId: string): Promise<void> {
const rl = createInterface({ input: process.stdin, output: process.stdout });
try {
const credential = await provider.auth.oauth.login({
signal: new AbortController().signal,
prompt: (authPrompt) => answerPrompt(rl, authPrompt),
notify: (event) => {
switch (event.type) {

View file

@ -194,6 +194,7 @@ class ImagesModelsImpl implements MutableImagesModels {
const resolution = await this.getAuth(model, {
apiKey: options?.apiKey,
env: options?.env,
signal: options?.signal,
});
const auth = resolution?.auth;
if (!auth) {

View file

@ -13,33 +13,33 @@ export interface ModelsStoreEntry {
etag?: string;
}
/** Persistent model catalogs keyed by provider ID. */
export interface ModelsStore {
read(providerId: string): Promise<ModelsStoreEntry | undefined>;
write(providerId: string, entry: ModelsStoreEntry): Promise<void>;
delete(providerId: string): Promise<void>;
export interface ModelsStoreOperationOptions {
signal?: AbortSignal;
}
/** ModelsStore scoped to one provider. Providers cannot access other providers' catalogs. */
export interface ProviderModelsStore {
read(): Promise<ModelsStoreEntry | undefined>;
write(entry: ModelsStoreEntry): Promise<void>;
delete(): Promise<void>;
/** Persistent model catalogs keyed by provider ID. */
export interface ModelsStore {
read(providerId: string, options?: ModelsStoreOperationOptions): Promise<ModelsStoreEntry | undefined>;
write(providerId: string, entry: ModelsStoreEntry, options?: ModelsStoreOperationOptions): Promise<void>;
delete(providerId: string, options?: ModelsStoreOperationOptions): Promise<void>;
}
export class InMemoryModelsStore implements ModelsStore {
private readonly entries = new Map<string, ModelsStoreEntry>();
async read(providerId: string): Promise<ModelsStoreEntry | undefined> {
async read(providerId: string, options?: ModelsStoreOperationOptions): Promise<ModelsStoreEntry | undefined> {
options?.signal?.throwIfAborted();
const entry = this.entries.get(providerId);
return entry ? structuredClone(entry) : undefined;
}
async write(providerId: string, entry: ModelsStoreEntry): Promise<void> {
async write(providerId: string, entry: ModelsStoreEntry, options?: ModelsStoreOperationOptions): Promise<void> {
options?.signal?.throwIfAborted();
this.entries.set(providerId, structuredClone(entry));
}
async delete(providerId: string): Promise<void> {
async delete(providerId: string, options?: ModelsStoreOperationOptions): Promise<void> {
options?.signal?.throwIfAborted();
this.entries.delete(providerId);
}
}

View file

@ -6,13 +6,14 @@ import type {
AuthCheck,
AuthContext,
AuthInteraction,
AuthOperationOptions,
AuthResult,
AuthType,
Credential,
CredentialStore,
ProviderAuth,
} from "./auth/types.ts";
import { InMemoryModelsStore, type ModelsStore, type ProviderModelsStore } from "./models-store.ts";
import { InMemoryModelsStore, type ModelsStore, type ModelsStoreEntry } from "./models-store.ts";
import type {
Api,
ApiStreamOptions,
@ -28,23 +29,39 @@ import type {
StreamOptions,
Usage,
} from "./types.ts";
import { operationSignal, raceWithAbortSignal } from "./utils/abort.ts";
export { ModelsError, type ModelsErrorCode } from "./auth/resolve.ts";
export interface ModelsPublication {
/** Provider-selected persisted catalog. Omit to leave storage unchanged; null deletes it. */
persist?: ModelsStoreEntry | null;
/** Optional synchronous update of provider-private in-memory catalog state. */
update?: () => void;
}
export interface RefreshModelsContext {
/** Effective configured credential. OAuth credentials are refreshed before network access. */
credential?: Credential;
/** Persistent model storage scoped to this provider ID. */
store: ProviderModelsStore;
/** Immutable provider-scoped catalog snapshot captured before this refresh phase. */
stored?: Readonly<ModelsStoreEntry>;
/**
* Generation-checked publication. Persistence policy remains provider-owned;
* the update runs synchronously only after the selected persistence mutation.
*/
publish(publication: ModelsPublication): Promise<boolean>;
/** False during offline/cache-only initialization. */
allowNetwork: boolean;
/** Bypass provider freshness checks and fetch immediately when network access is allowed. */
force?: boolean;
signal?: AbortSignal;
/** Always present, including when the public refresh caller omits its optional signal. */
signal: AbortSignal;
}
export interface ModelsRefreshOptions {
allowNetwork?: boolean;
/** Restrict refresh to these provider IDs. Unknown and static providers are ignored. */
providers?: readonly string[];
/** Bypass provider freshness checks and fetch immediately when network access is allowed. */
force?: boolean;
signal?: AbortSignal;
@ -97,9 +114,10 @@ export interface Provider<TApi extends Api = Api> {
getModels(): readonly Model<TApi>[];
/**
* Dynamic providers only: restore the provider-scoped stored catalog and optionally fetch
* a newer list using the effective credential. Implementations must retain their previous
* list on failure and honor the shared abort signal for network requests.
* Dynamic providers only: restore `context.stored` and optionally fetch a newer list using
* the effective credential. Implementations retain their previous list on failure, publish
* persistence and synchronous state changes through `context.publish()`, and honor the
* shared abort signal for blocking work.
*/
refreshModels?(context: RefreshModelsContext): Promise<void>;
@ -141,16 +159,17 @@ export interface Models {
getModel(provider: string, id: string): Model<Api> | undefined;
/**
* Refresh every configured dynamic provider concurrently. Provider errors and cancellation
* are returned without rejecting; static and unconfigured providers are skipped.
* Refresh selected configured dynamic providers concurrently (all when `providers` is omitted).
* Provider errors and cancellation are returned without rejecting; static, unknown, and
* unconfigured providers are skipped.
*/
refresh(options?: ModelsRefreshOptions): Promise<ModelsRefreshResult>;
/** Check whether a provider has complete auth configuration without refreshing OAuth. */
checkAuth(providerId: string): Promise<AuthCheck | undefined>;
checkAuth(providerId: string, options?: AuthOperationOptions): Promise<AuthCheck | undefined>;
/** Return models whose providers have complete auth configuration. */
getAvailable(providerId?: string): Promise<readonly Model<Api>[]>;
getAvailable(providerId?: string, options?: AuthOperationOptions): Promise<readonly Model<Api>[]>;
/**
* Resolve provider-scoped auth by provider id, or provider auth plus static
@ -168,7 +187,7 @@ export interface Models {
login(providerId: string, type: AuthType, interaction: AuthInteraction): Promise<Credential>;
/** Remove the stored credential for a provider. */
logout(providerId: string): Promise<void>;
logout(providerId: string, options?: AuthOperationOptions): Promise<void>;
stream<TApi extends Api>(
model: Model<TApi>,
@ -220,6 +239,9 @@ class ModelsImpl implements MutableModels {
private credentials: CredentialStore;
private modelsStore: ModelsStore;
private authContext: AuthContext;
private refreshGenerations = new Map<string, number>();
private refreshControllers = new Map<string, AbortController>();
private publicationChains = new Map<string, Promise<unknown>>();
constructor(options?: CreateModelsOptions) {
this.credentials = options?.credentials ?? new InMemoryCredentialStore();
@ -228,14 +250,19 @@ class ModelsImpl implements MutableModels {
}
setProvider(provider: Provider): void {
this.supersedeProviderRefresh(provider.id);
this.providers.set(provider.id, provider);
}
deleteProvider(id: string): void {
this.supersedeProviderRefresh(id);
this.providers.delete(id);
}
clearProviders(): void {
for (const id of new Set([...this.providers.keys(), ...this.refreshControllers.keys()])) {
this.supersedeProviderRefresh(id);
}
this.providers.clear();
}
@ -273,36 +300,110 @@ class ModelsImpl implements MutableModels {
return this.getModels(provider).find((model) => model.id === id);
}
private supersedeProviderRefresh(providerId: string): number {
const generation = (this.refreshGenerations.get(providerId) ?? 0) + 1;
this.refreshGenerations.set(providerId, generation);
const previous = this.refreshControllers.get(providerId);
if (previous) {
this.refreshControllers.delete(providerId);
previous.abort();
}
return generation;
}
private beginProviderRefresh(providerId: string): { generation: number; controller: AbortController } {
const generation = this.supersedeProviderRefresh(providerId);
const controller = new AbortController();
this.refreshControllers.set(providerId, controller);
return { generation, controller };
}
private publishProviderModels(
providerId: string,
generation: number,
signal: AbortSignal,
publication: ModelsPublication,
): Promise<boolean> {
const previous = this.publicationChains.get(providerId) ?? Promise.resolve();
const queued = (async () => {
await previous.catch(() => {});
if (signal.aborted || this.refreshGenerations.get(providerId) !== generation) return false;
if (publication.persist === null) {
await this.modelsStore.delete(providerId, { signal });
} else if (publication.persist !== undefined) {
await this.modelsStore.write(providerId, structuredClone(publication.persist), { signal });
}
if (signal.aborted || this.refreshGenerations.get(providerId) !== generation) return false;
publication.update?.();
return true;
})();
const tail = queued.catch(() => {});
this.publicationChains.set(providerId, tail);
void tail.then(() => {
if (this.publicationChains.get(providerId) === tail) this.publicationChains.delete(providerId);
});
return raceWithAbortSignal(queued, signal);
}
private async runProviderRefreshPhase(
provider: Provider & Required<Pick<Provider, "refreshModels">>,
credential: Credential | undefined,
allowNetwork: boolean,
force: boolean | undefined,
generation: number,
signal: AbortSignal,
): Promise<void> {
const stored = await this.modelsStore.read(provider.id, { signal });
await provider.refreshModels({
credential,
stored: stored ? structuredClone(stored) : undefined,
publish: (publication) => this.publishProviderModels(provider.id, generation, signal, publication),
allowNetwork,
force: allowNetwork ? force : undefined,
signal,
});
}
async refresh(options: ModelsRefreshOptions = {}): Promise<ModelsRefreshResult> {
const allowNetwork = options.allowNetwork ?? true;
const callerSignal = operationSignal(options.signal);
const errors = new Map<string, Error>();
if (callerSignal.aborted) return { aborted: true, errors };
const selected = options.providers ? new Set(options.providers) : undefined;
const refreshable = Array.from(this.providers.values()).filter(
(provider): provider is Provider & Required<Pick<Provider, "refreshModels">> =>
provider.refreshModels !== undefined,
provider.refreshModels !== undefined && (!selected || selected.has(provider.id)),
);
await Promise.all(
const refresh = Promise.all(
refreshable.map(async (provider) => {
if (options.signal?.aborted) return;
const store: ProviderModelsStore = {
read: () => this.modelsStore.read(provider.id),
write: (entry) => this.modelsStore.write(provider.id, entry),
delete: () => this.modelsStore.delete(provider.id),
};
let stored: Credential | undefined;
try {
stored = await this.readCredential(provider.id);
const credential = await this.resolveRefreshCredential(provider, stored, allowNetwork, options.signal);
const { generation, controller } = this.beginProviderRefresh(provider.id);
const signal = AbortSignal.any([callerSignal, controller.signal]);
const operation = (async () => {
let storedCredential: Credential | undefined;
let credentialError: unknown;
try {
storedCredential = await this.readCredential(provider.id, signal);
} catch (error) {
credentialError = error;
}
// Restore cached provider state before auth resolution or network access.
await this.runProviderRefreshPhase(provider, storedCredential, false, undefined, generation, signal);
if (credentialError !== undefined) throw credentialError;
if (!allowNetwork || signal.aborted) return;
const credential = await this.resolveRefreshCredential(provider, storedCredential, signal);
if (!credential) return;
await provider.refreshModels({
credential,
store,
allowNetwork,
force: options.force,
signal: options.signal,
});
await this.runProviderRefreshPhase(provider, credential, true, options.force, generation, signal);
})();
try {
await raceWithAbortSignal(operation, signal);
} catch (error) {
if (!options.signal?.aborted) {
if (!signal.aborted) {
errors.set(
provider.id,
error instanceof Error
@ -310,52 +411,55 @@ class ModelsImpl implements MutableModels {
: new ModelsError("model_source", `Model refresh failed for ${provider.id}`, { cause: error }),
);
}
try {
await provider.refreshModels({
credential: stored,
store,
allowNetwork: false,
signal: options.signal,
});
} catch {
// Preserve the original auth/network error; cache restoration is best-effort here.
} finally {
if (this.refreshControllers.get(provider.id) === controller) {
this.refreshControllers.delete(provider.id);
}
}
}),
);
return { aborted: options.signal?.aborted ?? false, errors };
try {
await raceWithAbortSignal(refresh, callerSignal);
} catch (error) {
if (!callerSignal.aborted) throw error;
}
return { aborted: callerSignal.aborted, errors: new Map(errors) };
}
private async resolveRefreshCredential(
provider: Provider,
stored: Credential | undefined,
allowNetwork: boolean,
signal?: AbortSignal,
signal: AbortSignal,
): Promise<Credential | undefined> {
if (stored?.type === "oauth") {
const oauth = provider.auth.oauth;
if (!oauth) return undefined;
if (!allowNetwork || Date.now() < stored.expires) return stored;
if (signal?.aborted) return undefined;
const post = await this.credentials.modify(provider.id, async (current) => {
if (current?.type !== "oauth" || Date.now() < current.expires) return undefined;
return oauth.refresh(current, signal);
});
if (Date.now() < stored.expires) return stored;
if (signal.aborted) return undefined;
const post = await this.credentials.modify(
provider.id,
async (current) => {
if (current?.type !== "oauth" || Date.now() < current.expires) return undefined;
return oauth.refresh(current, signal);
},
{ signal },
);
return post?.type === "oauth" ? post : undefined;
}
const apiKey = provider.auth.apiKey;
if (!apiKey) return undefined;
const credential = stored?.type === "api_key" ? stored : undefined;
const result = await apiKey.resolve({ ctx: this.authContext, credential });
const result = await apiKey.resolve({ ctx: this.authContext, credential, signal });
if (!result) return undefined;
return { type: "api_key", key: result.auth.apiKey, env: result.env };
}
private async readCredential(providerId: string): Promise<Credential | undefined> {
private async readCredential(providerId: string, signal: AbortSignal): Promise<Credential | undefined> {
try {
return await this.credentials.read(providerId);
return await this.credentials.read(providerId, { signal });
} catch (error) {
throw new ModelsError("auth", `Credential store read failed for ${providerId}`, { cause: error });
}
@ -364,6 +468,7 @@ class ModelsImpl implements MutableModels {
private async checkProviderAuth(
provider: Provider,
credential: Credential | undefined,
signal: AbortSignal,
): Promise<AuthCheck | undefined> {
if (credential?.type === "oauth") {
return provider.auth.oauth ? { source: "OAuth", type: "oauth" } : undefined;
@ -375,37 +480,48 @@ class ModelsImpl implements MutableModels {
return await apiKey.check({
ctx: this.authContext,
credential: credential?.type === "api_key" ? credential : undefined,
signal,
});
} catch (error) {
throw new ModelsError("auth", `API key auth check failed for provider ${provider.id}`, { cause: error });
}
}
const resolution = await resolveProviderAuth(provider, this.credentials, this.authContext);
const resolution = await resolveProviderAuth(provider, this.credentials, this.authContext, { signal });
return resolution ? { source: resolution.source, type: "api_key" } : undefined;
}
async checkAuth(providerId: string): Promise<AuthCheck | undefined> {
const provider = this.providers.get(providerId);
if (!provider) return undefined;
return this.checkProviderAuth(provider, await this.readCredential(providerId));
checkAuth(providerId: string, options?: AuthOperationOptions): Promise<AuthCheck | undefined> {
const signal = operationSignal(options?.signal);
const check = (async () => {
signal.throwIfAborted();
const provider = this.providers.get(providerId);
if (!provider) return undefined;
return this.checkProviderAuth(provider, await this.readCredential(providerId, signal), signal);
})();
return raceWithAbortSignal(check, signal);
}
async getAvailable(providerId?: string): Promise<readonly Model<Api>[]> {
const providers = providerId
? [this.providers.get(providerId)].filter((entry) => entry !== undefined)
: this.getProviders();
const checks = await Promise.all(
providers.map(async (provider) => {
const credential = await this.readCredential(provider.id);
return { provider, credential, auth: await this.checkProviderAuth(provider, credential) };
}),
);
return checks.flatMap(({ provider, credential, auth }) => {
if (!auth) return [];
const models = provider.getModels();
return provider.filterModels?.(models, credential) ?? models;
});
getAvailable(providerId?: string, options?: AuthOperationOptions): Promise<readonly Model<Api>[]> {
const signal = operationSignal(options?.signal);
const available = (async () => {
signal.throwIfAborted();
const providers = providerId
? [this.providers.get(providerId)].filter((entry) => entry !== undefined)
: this.getProviders();
const checks = await Promise.all(
providers.map(async (provider) => {
const credential = await this.readCredential(provider.id, signal);
return { provider, credential, auth: await this.checkProviderAuth(provider, credential, signal) };
}),
);
return checks.flatMap(({ provider, credential, auth }) => {
if (!auth) return [];
const models = provider.getModels();
return provider.filterModels?.(models, credential) ?? models;
});
})();
return raceWithAbortSignal(available, signal);
}
getAuth(providerId: string, overrides?: AuthResolutionOverrides): Promise<AuthResult | undefined>;
@ -414,10 +530,11 @@ class ModelsImpl implements MutableModels {
providerOrModel: string | Model<Api>,
overrides?: AuthResolutionOverrides,
): Promise<AuthResult | undefined> {
const signal = operationSignal(overrides?.signal);
const providerId = typeof providerOrModel === "string" ? providerOrModel : providerOrModel.provider;
const provider = this.providers.get(providerId);
if (!provider) return undefined;
const result = await resolveProviderAuth(provider, this.credentials, this.authContext, overrides);
const result = await resolveProviderAuth(provider, this.credentials, this.authContext, { ...overrides, signal });
if (!result || typeof providerOrModel === "string" || !providerOrModel.headers) return result;
return {
...result,
@ -429,25 +546,64 @@ class ModelsImpl implements MutableModels {
}
async login(providerId: string, type: AuthType, interaction: AuthInteraction): Promise<Credential> {
const signal = operationSignal(interaction.signal);
signal.throwIfAborted();
const provider = this.providers.get(providerId);
if (!provider) throw new ModelsError("provider", `Unknown provider: ${providerId}`);
const method = type === "oauth" ? provider.auth.oauth : provider.auth.apiKey;
if (!method?.login) {
throw new ModelsError("auth", `${provider.name} does not support ${type} login`);
}
const credential = await method.login(interaction);
const loginOperation: Promise<Credential> = method.login({ ...interaction, signal });
const credential = await raceWithAbortSignal(loginOperation, signal);
let mutationStarted = false;
let markMutationStarted: (() => void) | undefined;
const started = new Promise<void>((resolve) => {
markMutationStarted = resolve;
});
const mutation = this.credentials.modify(
providerId,
async () => {
mutationStarted = true;
markMutationStarted?.();
return credential;
},
{ signal },
);
void mutation.catch(() => {});
try {
await this.credentials.modify(providerId, async () => credential);
await new Promise<void>((resolve, reject) => {
const onAbort = () => {
if (!mutationStarted) reject(signal.reason);
};
signal.addEventListener("abort", onAbort, { once: true });
void Promise.race([started, mutation]).then(
() => {
signal.removeEventListener("abort", onAbort);
resolve();
},
(error: unknown) => {
signal.removeEventListener("abort", onAbort);
reject(error);
},
);
if (signal.aborted) onAbort();
});
await mutation;
} catch (error) {
signal.throwIfAborted();
throw new ModelsError("auth", `Credential store modify failed for ${providerId}`, { cause: error });
}
return credential;
}
async logout(providerId: string): Promise<void> {
async logout(providerId: string, options?: AuthOperationOptions): Promise<void> {
const signal = operationSignal(options?.signal);
signal.throwIfAborted();
try {
await this.credentials.delete(providerId);
await this.credentials.delete(providerId, { signal });
} catch (error) {
signal.throwIfAborted();
throw new ModelsError("auth", `Credential store delete failed for ${providerId}`, { cause: error });
}
}
@ -468,6 +624,7 @@ class ModelsImpl implements MutableModels {
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}`);
@ -540,7 +697,7 @@ export interface CreateProviderOptions<TApi extends Api = Api> {
auth: ProviderAuth;
/** Static baseline model list (empty for purely dynamic providers). */
models: readonly Model<TApi>[];
/** Fetch a dynamic model overlay. createProvider restores/persists it through ModelsStore. */
/** Fetch a dynamic model overlay. createProvider restores and publishes it transactionally. */
fetchModels?: (context: RefreshModelsContext) => Promise<readonly Model<TApi>[]>;
filterModels?: (models: readonly Model<TApi>[], credential: Credential | undefined) => readonly Model<TApi>[];
/** Single implementation, or map keyed by `model.api` for mixed-API providers. */
@ -556,7 +713,6 @@ export interface CreateProviderOptions<TApi extends Api = Api> {
export function createProvider<TApi extends Api = Api>(input: CreateProviderOptions<TApi>): Provider<TApi> {
const baselineModels = input.models;
let dynamicModels: readonly Model<TApi>[] = [];
let inflightRefresh: Promise<void> | undefined;
const fetchModels = input.fetchModels;
const currentModels = (): readonly Model<TApi>[] => {
const merged = [...baselineModels];
@ -594,25 +750,30 @@ export function createProvider<TApi extends Api = Api>(input: CreateProviderOpti
auth: input.auth,
getModels: currentModels,
refreshModels: fetchModels
? (context) => {
inflightRefresh ??= (async () => {
try {
const stored = await context.store.read();
if (stored) {
dynamicModels = stored.models
.filter((model) => model.provider === input.id)
.map((model) => model as Model<TApi>);
}
if (!context.allowNetwork || context.signal?.aborted) return;
const refreshed = await fetchModels(context);
if (context.signal?.aborted) return;
dynamicModels = refreshed;
await context.store.write({ models: refreshed, checkedAt: Date.now() });
} finally {
inflightRefresh = undefined;
? async (context) => {
if (context.stored) {
const restored = context.stored.models
.filter((model) => model.provider === input.id)
.map((model) => model as Model<TApi>);
if (
!(await context.publish({
update: () => {
dynamicModels = restored;
},
}))
) {
return;
}
})();
return inflightRefresh;
}
if (!context.allowNetwork || context.signal.aborted) return;
const refreshed = await fetchModels(context);
if (context.signal.aborted) return;
await context.publish({
persist: { models: refreshed, checkedAt: Date.now() },
update: () => {
dynamicModels = refreshed;
},
});
}
: undefined,
filterModels: input.filterModels,

View file

@ -11,6 +11,7 @@ import { AMAZON_BEDROCK_MODELS } from "./amazon-bedrock.models.ts";
const bedrockAuth: ApiKeyAuth = {
name: "AWS credentials or bearer token",
login: async (interaction) => {
interaction.signal.throwIfAborted();
const method = await interaction.prompt({
type: "select",
message: "Select Amazon Bedrock authentication method:",
@ -20,6 +21,7 @@ const bedrockAuth: ApiKeyAuth = {
{ id: "credential-chain", label: "Existing AWS credential chain" },
],
});
interaction.signal.throwIfAborted();
if (method === "bearer-token") {
return {
type: "api_key",
@ -49,24 +51,30 @@ const bedrockAuth: ApiKeyAuth = {
});
return { type: "api_key" };
},
resolve: async ({ ctx, credential }) => {
resolve: async ({ ctx, credential, signal }) => {
const env = async (name: string) => {
signal.throwIfAborted();
const value = await ctx.env(name);
signal.throwIfAborted();
return value;
};
if (credential?.key) {
return { auth: { apiKey: credential.key }, env: credential.env, source: "stored credential" };
}
if (await ctx.env("AWS_BEARER_TOKEN_BEDROCK")) return { auth: {}, source: "AWS_BEARER_TOKEN_BEDROCK" };
if (credential?.env?.AWS_PROFILE ?? (await ctx.env("AWS_PROFILE"))) {
if (await env("AWS_BEARER_TOKEN_BEDROCK")) return { auth: {}, source: "AWS_BEARER_TOKEN_BEDROCK" };
if (credential?.env?.AWS_PROFILE ?? (await env("AWS_PROFILE"))) {
return {
auth: {},
env: credential?.env,
source: credential?.env?.AWS_PROFILE ? "stored credential" : "AWS_PROFILE",
};
}
if ((await ctx.env("AWS_ACCESS_KEY_ID")) && (await ctx.env("AWS_SECRET_ACCESS_KEY"))) {
if ((await env("AWS_ACCESS_KEY_ID")) && (await env("AWS_SECRET_ACCESS_KEY"))) {
return { auth: {}, source: "AWS access keys" };
}
if (await ctx.env("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI")) return { auth: {}, source: "ECS task role" };
if (await ctx.env("AWS_CONTAINER_CREDENTIALS_FULL_URI")) return { auth: {}, source: "ECS task role" };
if (await ctx.env("AWS_WEB_IDENTITY_TOKEN_FILE")) return { auth: {}, source: "web identity token" };
if (await env("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI")) return { auth: {}, source: "ECS task role" };
if (await env("AWS_CONTAINER_CREDENTIALS_FULL_URI")) return { auth: {}, source: "ECS task role" };
if (await env("AWS_WEB_IDENTITY_TOKEN_FILE")) return { auth: {}, source: "web identity token" };
return undefined;
},
};

View file

@ -9,16 +9,20 @@ import { ANTHROPIC_MODELS } from "./anthropic.models.ts";
function anthropicApiKeyAuth(): ApiKeyAuth {
return {
name: "Anthropic API key",
login: async (interaction) => ({
type: "api_key",
key: await interaction.prompt({ type: "secret", message: "Enter Anthropic API key" }),
}),
resolve: async ({ ctx, credential }) => {
login: async (interaction) => {
interaction.signal.throwIfAborted();
const key = await interaction.prompt({ type: "secret", message: "Enter Anthropic API key" });
interaction.signal.throwIfAborted();
return { type: "api_key", key };
},
resolve: async ({ ctx, credential, signal }) => {
signal.throwIfAborted();
if (credential?.key) {
return { auth: { apiKey: credential.key }, env: credential.env, source: "stored credential" };
}
const authToken = await ctx.env(ANTHROPIC_AUTH_TOKEN_ENV);
signal.throwIfAborted();
if (authToken) {
return {
auth: { headers: { Authorization: `Bearer ${authToken}` } },
@ -28,6 +32,7 @@ function anthropicApiKeyAuth(): ApiKeyAuth {
for (const envVar of [ANTHROPIC_OAUTH_TOKEN_ENV, ANTHROPIC_API_KEY_ENV]) {
const apiKey = await ctx.env(envVar);
signal.throwIfAborted();
if (apiKey) return { auth: { apiKey }, source: envVar };
}
return undefined;

View file

@ -11,6 +11,7 @@ async function resolveValue(
name: string,
ctx: AuthContext,
credential: ApiKeyCredential | undefined,
signal: AbortSignal,
): Promise<string | undefined> {
// Per-field merge: prefer the credential value, fall back to ambient env.
// A credential carrying only the API key must still pick up the account /
@ -20,17 +21,23 @@ async function resolveValue(
? credential.key
: credential.env?.[name]
: undefined;
return fromCredential ?? (await ctx.env(name));
if (fromCredential !== undefined) return fromCredential;
signal.throwIfAborted();
const value = await ctx.env(name);
signal.throwIfAborted();
return value;
}
async function resolveCloudflareEnv(
kind: CloudflareAuthKind,
ctx: AuthContext,
credential: ApiKeyCredential | undefined,
signal: AbortSignal,
): Promise<{ apiKey: string; env: ProviderEnv; source: string } | undefined> {
const apiKey = await resolveValue(CLOUDFLARE_API_KEY, ctx, credential);
const accountId = await resolveValue(CLOUDFLARE_ACCOUNT_ID, ctx, credential);
const gatewayId = kind === "ai-gateway" ? await resolveValue(CLOUDFLARE_GATEWAY_ID, ctx, credential) : undefined;
const apiKey = await resolveValue(CLOUDFLARE_API_KEY, ctx, credential, signal);
const accountId = await resolveValue(CLOUDFLARE_ACCOUNT_ID, ctx, credential, signal);
const gatewayId =
kind === "ai-gateway" ? await resolveValue(CLOUDFLARE_GATEWAY_ID, ctx, credential, signal) : undefined;
if (!apiKey || !accountId || (kind === "ai-gateway" && !gatewayId)) return undefined;
@ -52,8 +59,8 @@ export function cloudflareWorkersAIAuth(): ApiKeyAuth {
const accountId = await interaction.prompt({ type: "text", message: "Enter Cloudflare account ID" });
return { type: "api_key", key, env: { CLOUDFLARE_ACCOUNT_ID: accountId } };
},
resolve: async ({ ctx, credential }) => {
const resolved = await resolveCloudflareEnv("workers-ai", ctx, credential);
resolve: async ({ ctx, credential, signal }) => {
const resolved = await resolveCloudflareEnv("workers-ai", ctx, credential, signal);
if (!resolved) return undefined;
return {
auth: { apiKey: resolved.apiKey },
@ -77,8 +84,8 @@ export function cloudflareAIGatewayAuth(): ApiKeyAuth {
env: { CLOUDFLARE_ACCOUNT_ID: accountId, CLOUDFLARE_GATEWAY_ID: gatewayId },
};
},
resolve: async ({ ctx, credential }) => {
const resolved = await resolveCloudflareEnv("ai-gateway", ctx, credential);
resolve: async ({ ctx, credential, signal }) => {
const resolved = await resolveCloudflareEnv("ai-gateway", ctx, credential, signal);
if (!resolved) return undefined;
return {
auth: {

View file

@ -13,6 +13,7 @@ const VERTEX_ADC_PATH = "~/.config/gcloud/application_default_credentials.json";
const vertexAuth: ApiKeyAuth = {
name: "Google Cloud credentials",
login: async (interaction) => {
interaction.signal.throwIfAborted();
const method = await interaction.prompt({
type: "select",
message: "Select Google Vertex AI authentication method:",
@ -22,6 +23,7 @@ const vertexAuth: ApiKeyAuth = {
{ id: "service-account", label: "Service account credentials file" },
],
});
interaction.signal.throwIfAborted();
if (method === "api-key") {
return {
type: "api_key",
@ -59,18 +61,23 @@ const vertexAuth: ApiKeyAuth = {
},
};
},
resolve: async ({ ctx, credential }) => {
const key = credential?.key ?? (await ctx.env("GOOGLE_CLOUD_API_KEY"));
resolve: async ({ ctx, credential, signal }) => {
const env = async (name: string) => {
signal.throwIfAborted();
const value = await ctx.env(name);
signal.throwIfAborted();
return value;
};
const key = credential?.key ?? (await env("GOOGLE_CLOUD_API_KEY"));
if (key) return { auth: { apiKey: key }, source: credential?.key ? "stored credential" : "GOOGLE_CLOUD_API_KEY" };
const adcPath =
credential?.env?.GOOGLE_APPLICATION_CREDENTIALS ?? (await ctx.env("GOOGLE_APPLICATION_CREDENTIALS"));
const adcPath = credential?.env?.GOOGLE_APPLICATION_CREDENTIALS ?? (await env("GOOGLE_APPLICATION_CREDENTIALS"));
signal.throwIfAborted();
const hasCredentials = await ctx.fileExists(adcPath ?? VERTEX_ADC_PATH);
signal.throwIfAborted();
const project =
credential?.env?.GOOGLE_CLOUD_PROJECT ??
(await ctx.env("GOOGLE_CLOUD_PROJECT")) ??
(await ctx.env("GCLOUD_PROJECT"));
const location = credential?.env?.GOOGLE_CLOUD_LOCATION ?? (await ctx.env("GOOGLE_CLOUD_LOCATION"));
credential?.env?.GOOGLE_CLOUD_PROJECT ?? (await env("GOOGLE_CLOUD_PROJECT")) ?? (await env("GCLOUD_PROJECT"));
const location = credential?.env?.GOOGLE_CLOUD_LOCATION ?? (await env("GOOGLE_CLOUD_LOCATION"));
if (hasCredentials && project && location) {
return {
auth: {},

View file

@ -22,7 +22,6 @@ export function radiusProvider(options: RadiusProviderOptions = {}): Provider<"p
const name = options.name ?? "Radius";
const gateway = normalizeRadiusGatewayUrl(options.gateway ?? DEFAULT_RADIUS_GATEWAY);
let models = getRadiusModels(id, undefined);
let inflightRefresh: Promise<void> | undefined;
const streams = piMessagesApi();
return {
@ -33,33 +32,49 @@ export function radiusProvider(options: RadiusProviderOptions = {}): Provider<"p
oauth: lazyOAuth({ name, load: () => loadRadiusOAuth({ name, gateway }) }),
},
getModels: () => models,
refreshModels: (context) => {
inflightRefresh ??= (async () => {
try {
const stored = await context.store.read();
if (stored) models = stored.models.filter((model) => model.provider === id) as typeof models;
// Import catalogs cached by the pre-ModelsStore Radius implementation.
if (!stored && context.credential?.type === "oauth") {
const legacy = getRadiusModels(id, context.credential);
if (legacy.length > 0) {
models = legacy;
await context.store.write({ models: legacy, checkedAt: Date.now() });
}
}
if (!context.allowNetwork || context.signal?.aborted) return;
const apiKey =
context.credential?.type === "oauth" ? context.credential.access : context.credential?.key;
const config = await loadRadiusGatewayConfig(gateway, apiKey, context.signal);
if (context.signal?.aborted) return;
models = getRadiusModelsFromConfig(id, config);
await context.store.write({ models, checkedAt: Date.now() });
} finally {
inflightRefresh = undefined;
refreshModels: async (context) => {
const stored = context.stored;
if (stored) {
const restored = stored.models.filter((model) => model.provider === id) as typeof models;
if (
!(await context.publish({
update: () => {
models = restored;
},
}))
) {
return;
}
})();
return inflightRefresh;
}
// Import catalogs cached by the pre-ModelsStore Radius implementation.
if (!stored && context.credential?.type === "oauth") {
const legacy = getRadiusModels(id, context.credential);
if (legacy.length > 0) {
if (
!(await context.publish({
persist: { models: legacy, checkedAt: Date.now() },
update: () => {
models = legacy;
},
}))
) {
return;
}
}
}
if (!context.allowNetwork || context.signal.aborted) return;
const apiKey = context.credential?.type === "oauth" ? context.credential.access : context.credential?.key;
const config = await loadRadiusGatewayConfig(gateway, apiKey, context.signal);
if (context.signal.aborted) return;
const refreshed = getRadiusModelsFromConfig(id, config);
await context.publish({
persist: { models: refreshed, checkedAt: Date.now() },
update: () => {
models = refreshed;
},
});
},
stream: (model, context, streamOptions) => streams.stream(model, context, streamOptions),
streamSimple: (model, context, streamOptions) => streams.streamSimple(model, context, streamOptions),

View file

@ -0,0 +1,50 @@
function abortReason(signal: AbortSignal): unknown {
if (signal.reason !== undefined) return signal.reason;
const error = new Error("The operation was aborted");
error.name = "AbortError";
return error;
}
/** Create an operation-local signal for public APIs whose signal is optional. */
export function operationSignal(signal?: AbortSignal): AbortSignal {
return signal ?? new AbortController().signal;
}
/**
* Stop waiting for an operation when its signal aborts while continuing to
* observe the abandoned promise so a later rejection is always handled.
*/
export function raceWithAbortSignal<T>(operation: Promise<T>, signal: AbortSignal): Promise<T> {
if (signal.aborted) {
void operation.catch(() => {});
return Promise.reject(abortReason(signal));
}
return new Promise<T>((resolve, reject) => {
let settled = false;
const cleanup = () => signal.removeEventListener("abort", onAbort);
const onAbort = () => {
if (settled) return;
settled = true;
cleanup();
reject(abortReason(signal));
};
signal.addEventListener("abort", onAbort, { once: true });
void operation.then(
(value) => {
if (settled) return;
settled = true;
cleanup();
resolve(value);
},
(error: unknown) => {
if (settled) return;
settled = true;
cleanup();
reject(error);
},
);
if (signal.aborted) onAbort();
});
}

View file

@ -51,6 +51,8 @@ vi.mock("@anthropic-ai/sdk", () => {
return { default: FakeAnthropic };
});
const neverAbortedSignal = new AbortController().signal;
const context: Context = {
systemPrompt: "System prompt.",
messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
@ -87,6 +89,7 @@ describe("Anthropic auth token env", () => {
})[name],
fileExists: async () => false,
},
signal: neverAbortedSignal,
});
expect(auth).toEqual({
@ -106,6 +109,7 @@ describe("Anthropic auth token env", () => {
})[name],
fileExists: async () => false,
},
signal: neverAbortedSignal,
});
expect(auth).toEqual({

View file

@ -2,6 +2,8 @@ import { afterEach, describe, expect, it, vi } from "vitest";
import { anthropicOAuth } from "../src/auth/oauth/anthropic.ts";
import type { AuthEvent, AuthPrompt } from "../src/auth/types.ts";
const neverAbortedSignal = new AbortController().signal;
function jsonResponse(body: unknown, status: number = 200): Response {
return new Response(JSON.stringify(body), {
status,
@ -54,6 +56,7 @@ describe.sequential("Anthropic OAuth", () => {
vi.stubGlobal("fetch", fetchMock);
const credentials = await anthropicOAuth.login({
signal: neverAbortedSignal,
notify: (event) => {
if (event.type === "auth_url") authUrl = event.url;
},
@ -89,12 +92,15 @@ describe.sequential("Anthropic OAuth", () => {
});
vi.stubGlobal("fetch", fetchMock);
const credentials = await anthropicOAuth.refresh({
type: "oauth",
access: "old-access-token",
refresh: "refresh-token",
expires: 0,
});
const credentials = await anthropicOAuth.refresh(
{
type: "oauth",
access: "old-access-token",
refresh: "refresh-token",
expires: 0,
},
neverAbortedSignal,
);
expect(credentials.access).toBe("new-access-token");
expect(credentials.refresh).toBe("new-refresh-token");
@ -116,6 +122,7 @@ describe.sequential("Anthropic OAuth", () => {
let manualSignal: AbortSignal | undefined;
const credential = await anthropicOAuth.login({
signal: neverAbortedSignal,
notify: (event) => events.push(event),
prompt: async (prompt) => {
prompts.push(prompt);

View file

@ -4,6 +4,8 @@ import { githubCopilotOAuth } from "../src/auth/oauth/github-copilot.ts";
import { createModels } from "../src/models.ts";
import { githubCopilotProvider } from "../src/providers/github-copilot.ts";
const neverAbortedSignal = new AbortController().signal;
function jsonResponse(body: unknown, status: number = 200): Response {
return new Response(JSON.stringify(body), {
status,
@ -38,7 +40,7 @@ function loginGitHubCopilotForTest(options: {
signal?: AbortSignal;
}) {
return githubCopilotOAuth.login({
signal: options.signal,
signal: options.signal ?? neverAbortedSignal,
prompt: (prompt) => {
if (prompt.type !== "text") throw new Error(`Unexpected prompt: ${prompt.type}`);
return options.onPrompt({ message: prompt.message, placeholder: prompt.placeholder, allowEmpty: true });
@ -101,12 +103,15 @@ describe("GitHub Copilot OAuth device flow", () => {
vi.stubGlobal("fetch", fetchMock);
const credentials = await githubCopilotOAuth.refresh({
type: "oauth",
access: "old-access-token",
refresh: "ghu_refresh_token",
expires: 0,
});
const credentials = await githubCopilotOAuth.refresh(
{
type: "oauth",
access: "old-access-token",
refresh: "ghu_refresh_token",
expires: 0,
},
neverAbortedSignal,
);
expect(credentials.availableModelIds).toEqual(["gpt-4.1"]);
const store = new InMemoryCredentialStore();

View file

@ -1,6 +1,6 @@
import { afterEach, describe, expect, it, vi } from "vitest";
import { kimiCodingOAuth } from "../src/auth/oauth/kimi-coding.ts";
import type { AuthInteraction } from "../src/auth/types.ts";
import type { ProviderAuthInteraction } from "../src/auth/types.ts";
const CLIENT_ID = "17e5f671-d194-4dfb-9706-5516cb48c098";
const OAUTH_HOST = "https://auth.kimi.com";
@ -31,8 +31,9 @@ function deviceAuthorizationResponse(overrides?: Record<string, unknown>): Respo
});
}
function createInteraction(events: Array<Record<string, unknown>>): AuthInteraction {
function createInteraction(events: Array<Record<string, unknown>>): ProviderAuthInteraction {
return {
signal: new AbortController().signal,
prompt: async () => {
throw new Error("Kimi Code login should not prompt");
},
@ -205,12 +206,15 @@ describe("Kimi Code OAuth", () => {
);
const before = Date.now();
const credential = await kimiCodingOAuth.refresh({
type: "oauth",
access: "old-access",
refresh: "old-refresh",
expires: before,
});
const credential = await kimiCodingOAuth.refresh(
{
type: "oauth",
access: "old-access",
refresh: "old-refresh",
expires: before,
},
new AbortController().signal,
);
expect(credential).toEqual({
type: "oauth",
access: "new-access",
@ -238,12 +242,15 @@ describe("Kimi Code OAuth", () => {
}),
);
const refreshPromise = kimiCodingOAuth.refresh({
type: "oauth",
access: "old",
refresh: "old",
expires: 0,
});
const refreshPromise = kimiCodingOAuth.refresh(
{
type: "oauth",
access: "old",
refresh: "old",
expires: 0,
},
new AbortController().signal,
);
await vi.advanceTimersByTimeAsync(1000);
await expect(refreshPromise).resolves.toMatchObject({ access: "a" });
expect(calls).toBe(2);
@ -254,7 +261,10 @@ describe("Kimi Code OAuth", () => {
vi.fn(async (): Promise<Response> => jsonResponse({ error: "invalid_grant" }, 400)),
);
await expect(
kimiCodingOAuth.refresh({ type: "oauth", access: "old", refresh: "old", expires: 0 }),
kimiCodingOAuth.refresh(
{ type: "oauth", access: "old", refresh: "old", expires: 0 },
new AbortController().signal,
),
).rejects.toThrow("unauthorized");
});
});

View file

@ -1,8 +1,8 @@
import { describe, expect, it, vi } from "vitest";
import { InMemoryCredentialStore } from "../src/auth/credential-store.ts";
import type { ApiKeyAuth, CredentialStore, OAuthAuth, ProviderAuth } from "../src/auth/types.ts";
import type { ApiKeyAuth, CredentialStore, OAuthAuth, OAuthCredential, ProviderAuth } from "../src/auth/types.ts";
import { calculateCost, createModels, createProvider, hasApi, type Provider } from "../src/models.ts";
import { InMemoryModelsStore } from "../src/models-store.ts";
import { InMemoryModelsStore, type ModelsStore, type ModelsStoreEntry } from "../src/models-store.ts";
import type { Api, AssistantMessage, Context, Model, SimpleStreamOptions, StreamOptions, Usage } from "../src/types.ts";
import { AssistantMessageEventStream } from "../src/utils/event-stream.ts";
@ -224,9 +224,14 @@ describe("Models runtime", () => {
testProvider({
id: "dyn",
getModels: () => list,
refreshModels: async () => {
refreshModels: async (refresh) => {
if (!refresh.allowNetwork) return;
refreshes++;
list = [testModel("dyn", "after")];
await refresh.publish({
update: () => {
list = [testModel("dyn", "after")];
},
});
},
}),
);
@ -242,8 +247,8 @@ describe("Models runtime", () => {
models.setProvider(
testProvider({
id: "flaky",
refreshModels: async () => {
throw new Error("fetch failed");
refreshModels: async ({ allowNetwork }) => {
if (allowNetwork) throw new Error("fetch failed");
},
}),
);
@ -252,6 +257,111 @@ describe("Models runtime", () => {
expect(second.errors.get("flaky")?.message).toBe("fetch failed");
});
it("restricts refresh work to selected providers", async () => {
const calls: string[] = [];
const models = createModels();
for (const id of ["one", "two"]) {
models.setProvider(
testProvider({
id,
refreshModels: async ({ allowNetwork }) => {
calls.push(`${id}:${allowNetwork ? "network" : "cache"}`);
},
}),
);
}
const result = await models.refresh({ providers: ["two", "unknown"] });
expect(result.errors.size).toBe(0);
expect(calls).toEqual(["two:cache", "two:network"]);
});
it("restores cached models before waiting for network auth", async () => {
const store = new InMemoryModelsStore();
await store.write("dynamic", { models: [testModel("dynamic", "cached")] });
let markAuthStarted: (() => void) | undefined;
let finishAuth: (() => void) | undefined;
const authStarted = new Promise<void>((resolve) => {
markAuthStarted = resolve;
});
const blockedAuth = new Promise<void>((resolve) => {
finishAuth = resolve;
});
const provider = createProvider({
id: "dynamic",
auth: {
apiKey: {
name: "Blocked auth",
resolve: async () => {
markAuthStarted?.();
await blockedAuth;
return { auth: { apiKey: "key" } };
},
},
},
models: [],
fetchModels: async () => {
throw new Error("must not fetch");
},
api: {
stream: () => new AssistantMessageEventStream(),
streamSimple: () => new AssistantMessageEventStream(),
},
});
const models = createModels({ modelsStore: store });
models.setProvider(provider);
const controller = new AbortController();
const pending = models.refresh({ providers: ["dynamic"], signal: controller.signal });
await authStarted;
expect(models.getModel("dynamic", "cached")).toBeDefined();
controller.abort();
expect(await pending).toMatchObject({ aborted: true });
finishAuth?.();
});
it("lets providers choose persistent deletion and ephemeral publication atomically", async () => {
let entry: ModelsStoreEntry | undefined = { models: [testModel("dynamic", "stored")] };
const store: ModelsStore = {
read: async () => entry,
write: async (_providerId, next) => {
entry = next;
},
delete: async () => {
entry = undefined;
},
};
let state = "initial";
const models = createModels({ modelsStore: store });
models.setProvider(
testProvider({
id: "dynamic",
refreshModels: async (context) => {
expect(context.stored?.models[0]?.id).toBe("stored");
await context.publish({
persist: null,
update: () => {
expect(entry).toBeUndefined();
state = "deleted";
},
});
await context.publish({
update: () => {
state = "ephemeral";
},
});
},
}),
);
const result = await models.refresh({ allowNetwork: false });
expect(result.errors.size).toBe(0);
expect(entry).toBeUndefined();
expect(state).toBe("ephemeral");
});
it("persists dynamic catalogs and restores them without network access", async () => {
const credentials = new InMemoryCredentialStore();
const modelsStore = new InMemoryModelsStore();
@ -293,6 +403,7 @@ describe("Models runtime", () => {
id: "configured",
auth: { apiKey: envKeyAuth("ambient-key") },
refreshModels: async (context) => {
if (!context.allowNetwork) return;
effectiveCredential = context.credential;
forceRefresh = context.force;
},
@ -302,8 +413,8 @@ describe("Models runtime", () => {
testProvider({
id: "unconfigured",
auth: { apiKey: envKeyAuth(undefined) },
refreshModels: async () => {
unconfiguredRefreshes++;
refreshModels: async ({ allowNetwork }) => {
if (allowNetwork) unconfiguredRefreshes++;
},
}),
);
@ -338,7 +449,7 @@ describe("Models runtime", () => {
}),
},
refreshModels: async (context) => {
modelRefreshCredential = context.credential;
if (context.allowNetwork) modelRefreshCredential = context.credential;
},
}),
);
@ -348,6 +459,59 @@ describe("Models runtime", () => {
expect(await credentials.read("oauth-dynamic")).toMatchObject({ access: "fresh", refresh: "rotated" });
});
it("always gives providers a concrete signal", async () => {
let receivedSignal: AbortSignal | undefined;
const models = createModels();
models.setProvider(
testProvider({
id: "dynamic",
refreshModels: async ({ signal }) => {
receivedSignal = signal;
},
}),
);
const result = await models.refresh();
expect(result.aborted).toBe(false);
expect(receivedSignal).toBeInstanceOf(AbortSignal);
expect(receivedSignal?.aborted).toBe(false);
});
it("binds model-store waits to the provider refresh signal", async () => {
const storageSignals: (AbortSignal | undefined)[] = [];
const store: ModelsStore = {
read: async (_providerId, options) => {
storageSignals.push(options?.signal);
return undefined;
},
write: async (_providerId, _entry, options) => {
storageSignals.push(options?.signal);
},
delete: async (_providerId, options) => {
storageSignals.push(options?.signal);
},
};
let providerSignal: AbortSignal | undefined;
const models = createModels({ modelsStore: store });
models.setProvider(
testProvider({
id: "dynamic",
auth: { apiKey: envKeyAuth("key") },
refreshModels: async (context) => {
providerSignal = context.signal;
if (!context.allowNetwork) return;
await context.publish({ persist: { models: [testModel("dynamic", "fresh")] } });
},
}),
);
const result = await models.refresh({ providers: ["dynamic"] });
expect(result.errors.size).toBe(0);
expect(storageSignals).toHaveLength(3);
expect(storageSignals.every((signal) => signal === providerSignal)).toBe(true);
});
it("returns aborted state without reporting cancellation as a provider error", async () => {
const controller = new AbortController();
const models = createModels();
@ -356,7 +520,7 @@ describe("Models runtime", () => {
id: "dynamic",
refreshModels: async ({ signal }) => {
controller.abort();
if (signal?.aborted) return;
if (signal.aborted) return;
},
}),
);
@ -366,6 +530,248 @@ describe("Models runtime", () => {
expect(result.errors.size).toBe(0);
});
it("stops waiting on abort when a provider ignores its signal", async () => {
const controller = new AbortController();
let markStarted: (() => void) | undefined;
let rejectRefresh: ((error: Error) => void) | undefined;
const started = new Promise<void>((resolve) => {
markStarted = resolve;
});
const stalled = new Promise<void>((_resolve, reject) => {
rejectRefresh = reject;
});
let calls = 0;
const models = createModels();
models.setProvider(
testProvider({
id: "dynamic",
refreshModels: async () => {
calls++;
if (calls !== 1) return;
markStarted?.();
await stalled;
},
}),
);
const pending = models.refresh({ signal: controller.signal });
await started;
controller.abort();
const result = await pending;
expect(result.aborted).toBe(true);
expect(result.errors.size).toBe(0);
rejectRefresh?.(new Error("late provider failure"));
await new Promise((resolve) => setTimeout(resolve, 0));
expect(result.errors.size).toBe(0);
});
it("rejects late publication from a superseded non-cooperative provider", async () => {
const store = new InMemoryModelsStore();
let state = "initial";
let calls = 0;
let markFirstStarted: (() => void) | undefined;
let finishFirst: (() => void) | undefined;
const firstStarted = new Promise<void>((resolve) => {
markFirstStarted = resolve;
});
const firstBlocked = new Promise<void>((resolve) => {
finishFirst = resolve;
});
const models = createModels({ modelsStore: store });
models.setProvider(
testProvider({
id: "dynamic",
refreshModels: async (context) => {
if (!context.allowNetwork) return;
calls++;
const current = calls;
if (current === 1) {
markFirstStarted?.();
await firstBlocked;
}
const value = `generation-${current}`;
await context.publish({
persist: { models: [testModel("dynamic", value)] },
update: () => {
state = value;
},
});
},
}),
);
const first = models.refresh({ providers: ["dynamic"] });
await firstStarted;
const second = models.refresh({ providers: ["dynamic"] });
await second;
await first;
finishFirst?.();
await new Promise((resolve) => setTimeout(resolve, 0));
expect(state).toBe("generation-2");
expect((await store.read("dynamic"))?.models[0]?.id).toBe("generation-2");
});
it("passes caller signals to provider auth callbacks", async () => {
const controller = new AbortController();
const received: AbortSignal[] = [];
const apiKey: ApiKeyAuth = {
name: "Signal auth",
login: async (interaction) => {
received.push(interaction.signal);
return { type: "api_key", key: "saved" };
},
check: async ({ signal }) => {
received.push(signal);
return { type: "api_key" };
},
resolve: async ({ signal }) => {
received.push(signal);
return { auth: { apiKey: "resolved" } };
},
};
const models = createModels();
models.setProvider(testProvider({ id: "p1", auth: { apiKey } }));
await models.checkAuth("p1", { signal: controller.signal });
await models.getAuth("p1", { signal: controller.signal });
await models.login("p1", "api_key", {
signal: controller.signal,
prompt: async () => "unused",
notify: () => {},
});
expect(received).toEqual([controller.signal, controller.signal, controller.signal]);
});
it("stops waiting for non-cooperative auth callbacks", async () => {
let startCheck: (() => void) | undefined;
let finishCheck: (() => void) | undefined;
const checkStarted = new Promise<void>((resolve) => {
startCheck = resolve;
});
const blockedCheck = new Promise<void>((resolve) => {
finishCheck = resolve;
});
let startResolve: (() => void) | undefined;
let finishResolve: (() => void) | undefined;
const resolveStarted = new Promise<void>((resolve) => {
startResolve = resolve;
});
const blockedResolve = new Promise<void>((resolve) => {
finishResolve = resolve;
});
const models = createModels();
models.setProvider(
testProvider({
id: "p1",
auth: {
apiKey: {
name: "Blocked auth",
check: async () => {
startCheck?.();
await blockedCheck;
return { type: "api_key" };
},
resolve: async () => {
startResolve?.();
await blockedResolve;
return { auth: { apiKey: "key" } };
},
},
},
}),
);
const availableController = new AbortController();
const available = models.getAvailable(undefined, { signal: availableController.signal });
await checkStarted;
availableController.abort();
await expect(available).rejects.toMatchObject({ name: "AbortError" });
const authController = new AbortController();
const auth = models.getAuth("p1", { signal: authController.signal });
await resolveStarted;
authController.abort();
await expect(auth).rejects.toMatchObject({ name: "AbortError" });
finishCheck?.();
finishResolve?.();
});
it("cancels queued credential mutations without running them later", async () => {
const credentials = new InMemoryCredentialStore();
let finishFirst: (() => void) | undefined;
const firstBlocked = new Promise<void>((resolve) => {
finishFirst = resolve;
});
let secondRan = false;
const first = credentials.modify("p1", async () => {
await firstBlocked;
return { type: "api_key", key: "first" };
});
const controller = new AbortController();
const second = credentials.modify(
"p1",
async () => {
secondRan = true;
return { type: "api_key", key: "second" };
},
{ signal: controller.signal },
);
controller.abort();
await expect(second).rejects.toMatchObject({ name: "AbortError" });
finishFirst?.();
await first;
await new Promise((resolve) => setTimeout(resolve, 0));
expect(secondRan).toBe(false);
expect(await credentials.read("p1")).toEqual({ type: "api_key", key: "first" });
});
it("passes cancellation to OAuth refresh and preserves the previous credential", async () => {
const credentials = new InMemoryCredentialStore();
const previous: OAuthCredential = { type: "oauth", access: "old", refresh: "old-refresh", expires: 0 };
await credentials.modify("p1", async () => previous);
let startRefresh: (() => void) | undefined;
let finishRefresh: ((credential: typeof previous) => void) | undefined;
const refreshStarted = new Promise<void>((resolve) => {
startRefresh = resolve;
});
const blockedRefresh = new Promise<typeof previous>((resolve) => {
finishRefresh = resolve;
});
let receivedSignal: AbortSignal | undefined;
const models = createModels({ credentials });
models.setProvider(
testProvider({
id: "p1",
auth: {
oauth: testOAuth({
refresh: async (_credential, signal) => {
receivedSignal = signal;
startRefresh?.();
return blockedRefresh;
},
}),
},
}),
);
const controller = new AbortController();
const auth = models.getAuth("p1", { signal: controller.signal });
await refreshStarted;
controller.abort();
await expect(auth).rejects.toMatchObject({ name: "AbortError" });
expect(receivedSignal).toBe(controller.signal);
finishRefresh?.({ ...previous, access: "new", expires: Date.now() + 60_000 });
await new Promise((resolve) => setTimeout(resolve, 0));
expect(await credentials.read("p1")).toEqual(previous);
});
it("resolves auth: stored credential owns the provider, ambient only when nothing stored", async () => {
const credentials = new InMemoryCredentialStore();
const models = createModels({ credentials });

View file

@ -10,6 +10,8 @@ import * as extensionOAuthCompatibility from "../src/oauth.ts";
import { anthropicProvider } from "../src/providers/anthropic.ts";
import { githubCopilotProvider } from "../src/providers/github-copilot.ts";
const neverAbortedSignal = new AbortController().signal;
function jsonResponse(body: unknown, status = 200): Response {
return new Response(JSON.stringify(body), { status, headers: { "Content-Type": "application/json" } });
}
@ -37,7 +39,7 @@ describe.sequential("OAuthAuth adapters", () => {
it("openrouter derives the api key and keeps the permanent credential on refresh", async () => {
const credential = { type: "oauth" as const, access: "token", refresh: "", expires: Number.MAX_SAFE_INTEGER };
expect(await openRouterOAuth.toAuth(credential)).toEqual({ apiKey: "token" });
expect(await openRouterOAuth.refresh(credential)).toBe(credential);
expect(await openRouterOAuth.refresh(credential, neverAbortedSignal)).toBe(credential);
});
it("xAI toAuth derives the api key from the access token", async () => {
@ -78,7 +80,10 @@ describe.sequential("OAuthAuth adapters", () => {
),
);
const refreshed = await anthropicOAuth.refresh({ type: "oauth", access: "old", refresh: "old-r", expires: 0 });
const refreshed = await anthropicOAuth.refresh(
{ type: "oauth", access: "old", refresh: "old-r", expires: 0 },
neverAbortedSignal,
);
expect(refreshed.type).toBe("oauth");
expect(refreshed.access).toBe("new-access");
expect(refreshed.refresh).toBe("new-refresh");
@ -97,13 +102,16 @@ describe.sequential("OAuthAuth adapters", () => {
});
vi.stubGlobal("fetch", fetchMock);
const refreshed = await githubCopilotOAuth.refresh({
type: "oauth",
access: "old",
refresh: "gh-token",
expires: 0,
enterpriseUrl: "company.ghe.com",
});
const refreshed = await githubCopilotOAuth.refresh(
{
type: "oauth",
access: "old",
refresh: "gh-token",
expires: 0,
enterpriseUrl: "company.ghe.com",
},
neverAbortedSignal,
);
expect(refreshed.access).toBe("new-token");
expect(refreshed.enterpriseUrl).toBe("company.ghe.com");
expect(fetchedUrls[0]).toContain("api.company.ghe.com");

View file

@ -1,6 +1,8 @@
import { afterEach, describe, expect, it, vi } from "vitest";
import { pollOAuthDeviceCodeFlow } from "../src/auth/oauth/device-code.ts";
const neverAbortedSignal = new AbortController().signal;
describe("OAuth device-code polling", () => {
afterEach(() => {
vi.useRealTimers();
@ -22,6 +24,7 @@ describe("OAuth device-code polling", () => {
intervalSeconds: 2,
expiresInSeconds: 30,
poll,
signal: neverAbortedSignal,
});
await vi.advanceTimersByTimeAsync(0);
@ -51,6 +54,7 @@ describe("OAuth device-code polling", () => {
pollTimes.push(Date.now());
return { status: "complete" as const, value: "token" };
},
signal: neverAbortedSignal,
});
await vi.advanceTimersByTimeAsync(1999);
@ -77,6 +81,7 @@ describe("OAuth device-code polling", () => {
if (!result) throw new Error("Unexpected extra poll");
return result;
},
signal: neverAbortedSignal,
});
await vi.advanceTimersByTimeAsync(0);
@ -109,6 +114,7 @@ describe("OAuth device-code polling", () => {
if (!result) throw new Error("Unexpected extra poll");
return result;
},
signal: neverAbortedSignal,
});
await vi.advanceTimersByTimeAsync(0);

View file

@ -69,7 +69,9 @@ export async function resolveApiKey(provider: string): Promise<string | undefine
if (!oauth) return undefined;
let credential = entry;
try {
if (Date.now() >= credential.expires) credential = await oauth.refresh(credential);
if (Date.now() >= credential.expires) {
credential = await oauth.refresh(credential, new AbortController().signal);
}
} catch (error) {
console.log(JSON.stringify(error));
return undefined;

View file

@ -1,6 +1,8 @@
import { afterEach, describe, expect, it, vi } from "vitest";
import { openaiCodexOAuth } from "../src/auth/oauth/openai-codex.ts";
const neverAbortedSignal = new AbortController().signal;
function jsonResponse(body: unknown, status: number = 200): Response {
return new Response(JSON.stringify(body), {
status,
@ -51,7 +53,7 @@ function loginOpenAICodexDeviceCodeForTest(options: {
signal?: AbortSignal;
}) {
return openaiCodexOAuth.login({
signal: options.signal,
signal: options.signal ?? neverAbortedSignal,
prompt: async (prompt) => {
if (prompt.type !== "select") throw new Error(`Unexpected prompt: ${prompt.type}`);
return "device_code";
@ -220,6 +222,7 @@ describe("OpenAI Codex OAuth", () => {
await expect(
openaiCodexOAuth.login({
signal: neverAbortedSignal,
prompt: async (prompt) => {
if (prompt.type !== "select") throw new Error("Text prompt should not be used");
selectPrompts.push(prompt);
@ -263,6 +266,7 @@ describe("OpenAI Codex OAuth", () => {
it("cancels when OpenAI Codex login method selection is cancelled", async () => {
await expect(
openaiCodexOAuth.login({
signal: neverAbortedSignal,
prompt: async () => {
throw new Error("Login cancelled");
},
@ -467,12 +471,15 @@ describe("OpenAI Codex OAuth", () => {
);
await expect(
openaiCodexOAuth.refresh({
type: "oauth",
access: "invalid-access-token",
refresh: "invalid-refresh-token",
expires: 0,
}),
openaiCodexOAuth.refresh(
{
type: "oauth",
access: "invalid-access-token",
refresh: "invalid-refresh-token",
expires: 0,
},
neverAbortedSignal,
),
).rejects.toThrow(/OpenAI Codex token refresh failed \(401\).*Could not validate your token/);
expect(consoleError).not.toHaveBeenCalled();
});

View file

@ -8,6 +8,7 @@ import { openrouterImagesProvider } from "../src/providers/openrouter-images.ts"
const TOKEN_URL = "https://openrouter.ai/api/v1/auth/keys";
const nativeFetch = globalThis.fetch;
const neverAbortedSignal = new AbortController().signal;
function jsonResponse(body: unknown, status = 200): Response {
return new Response(JSON.stringify(body), { status, headers: { "content-type": "application/json" } });
@ -65,6 +66,7 @@ describe.sequential("OpenRouter OAuth", () => {
let callbackResponse: Promise<Response> | undefined;
let manualSignal: AbortSignal | undefined;
const credential = await openRouterOAuth.login({
signal: neverAbortedSignal,
prompt: (prompt) => {
manualSignal = prompt.signal;
return new Promise<string>(() => {});
@ -113,6 +115,7 @@ describe.sequential("OpenRouter OAuth", () => {
let callbackResponse: Promise<Response> | undefined;
const login = openRouterOAuth.login({
signal: neverAbortedSignal,
prompt: () => new Promise<string>(() => {}),
notify: (event) => {
if (event.type !== "auth_url") return;
@ -141,6 +144,7 @@ describe.sequential("OpenRouter OAuth", () => {
let callbackUrl: URL | undefined;
let firstCallback: Promise<Response> | undefined;
const login = openRouterOAuth.login({
signal: neverAbortedSignal,
prompt: () => new Promise<string>(() => {}),
notify: (event) => {
if (event.type !== "auth_url") return;
@ -168,6 +172,7 @@ describe.sequential("OpenRouter OAuth", () => {
let callbackResponse: Promise<Response> | undefined;
const login = openRouterOAuth.login({
signal: neverAbortedSignal,
prompt: () => new Promise<string>(() => {}),
notify: (event) => {
if (event.type !== "auth_url") return;
@ -193,6 +198,7 @@ describe.sequential("OpenRouter OAuth", () => {
let callbackUrl: string | undefined;
const credential = await openRouterOAuth.login({
signal: neverAbortedSignal,
prompt: async (prompt) => {
if (prompt.type !== "manual_code") throw new Error(`Unexpected prompt: ${prompt.type}`);
return `${callbackUrl}?code=manual-code`;
@ -224,6 +230,7 @@ describe.sequential("OpenRouter OAuth", () => {
vi.stubGlobal("fetch", fetchMock);
const credential = await openRouterOAuth.login({
signal: neverAbortedSignal,
prompt: async () => " manual-code ",
notify: () => {},
});
@ -238,6 +245,7 @@ describe.sequential("OpenRouter OAuth", () => {
await expect(
openRouterOAuth.login({
signal: neverAbortedSignal,
prompt: async () => {
throw new Error("Login cancelled");
},
@ -253,6 +261,7 @@ describe.sequential("OpenRouter OAuth", () => {
await expect(
openRouterOAuth.login({
signal: neverAbortedSignal,
prompt: async () => " ",
notify: () => {},
}),

View file

@ -2,7 +2,7 @@ import { describe, expect, it } from "vitest";
import { envApiKeyAuth } from "../src/auth/helpers.ts";
import type { AuthContext, AuthEvent } from "../src/auth/types.ts";
import { createModels, createProvider } from "../src/models.ts";
import { InMemoryModelsStore, type ModelsStoreEntry } from "../src/models-store.ts";
import { InMemoryModelsStore } from "../src/models-store.ts";
import { builtinModels, builtinProviders, getBuiltinModel } from "../src/providers/all.ts";
import { amazonBedrockProvider } from "../src/providers/amazon-bedrock.ts";
import { anthropicProvider } from "../src/providers/anthropic.ts";
@ -20,6 +20,8 @@ function fakeAuthContext(env: Record<string, string>, files: string[] = []): Aut
};
}
const neverAbortedSignal = new AbortController().signal;
const context: Context = { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] };
describe("builtin providers", () => {
@ -111,6 +113,7 @@ describe("builtin providers", () => {
const bearerAnswers = ["bearer-token", "bedrock-token"];
expect(
await auth.login?.({
signal: neverAbortedSignal,
prompt: async () => bearerAnswers.shift()!,
notify: () => {},
}),
@ -120,6 +123,7 @@ describe("builtin providers", () => {
const events: AuthEvent[] = [];
expect(
await auth.login?.({
signal: neverAbortedSignal,
prompt: async () => profileAnswers.shift()!,
notify: (event) => events.push(event),
}),
@ -134,6 +138,7 @@ describe("builtin providers", () => {
await auth.resolve({
ctx: fakeAuthContext({}),
credential: { type: "api_key", env: { AWS_PROFILE: "work" } },
signal: neverAbortedSignal,
}),
).toMatchObject({ auth: {}, env: { AWS_PROFILE: "work" } });
});
@ -202,6 +207,7 @@ describe("builtin providers", () => {
const keyAnswers = ["api-key", "vertex-key"];
expect(
await auth.login?.({
signal: neverAbortedSignal,
prompt: async () => keyAnswers.shift()!,
notify: () => {},
}),
@ -211,6 +217,7 @@ describe("builtin providers", () => {
const events: AuthEvent[] = [];
expect(
await auth.login?.({
signal: neverAbortedSignal,
prompt: async () => adcAnswers.shift()!,
notify: (event) => events.push(event),
}),
@ -231,6 +238,7 @@ describe("builtin providers", () => {
type: "api_key",
env: { GOOGLE_CLOUD_PROJECT: "project-id", GOOGLE_CLOUD_LOCATION: "us-central1" },
},
signal: neverAbortedSignal,
}),
).toMatchObject({
auth: {},
@ -269,20 +277,22 @@ describe("envApiKeyAuth", () => {
const stored = await auth.resolve({
ctx: fakeAuthContext({ FIRST_KEY: "env" }),
credential: { type: "api_key", key: "stored" },
signal: neverAbortedSignal,
});
expect(stored?.auth.apiKey).toBe("stored");
expect(stored?.source).toBe("stored credential");
const second = await auth.resolve({ ctx: fakeAuthContext({ SECOND_KEY: "second" }) });
const second = await auth.resolve({ ctx: fakeAuthContext({ SECOND_KEY: "second" }), signal: neverAbortedSignal });
expect(second?.auth.apiKey).toBe("second");
expect(second?.source).toBe("SECOND_KEY");
expect(await auth.resolve({ ctx: fakeAuthContext({}) })).toBeUndefined();
expect(await auth.resolve({ ctx: fakeAuthContext({}), signal: neverAbortedSignal })).toBeUndefined();
});
it("login prompts for a secret and returns an api-key credential", async () => {
const auth = envApiKeyAuth("Test key", ["TEST_KEY"]);
const credential = await auth.login?.({
signal: neverAbortedSignal,
prompt: async (prompt) => {
expect(prompt.type).toBe("secret");
return "entered-key";
@ -391,38 +401,50 @@ describe("createProvider", () => {
expect(result.errorMessage).toContain("no API implementation");
});
it("supports dynamic providers: empty until refreshed, in-flight refreshes deduped", async () => {
it("lets a newer dynamic refresh bypass and supersede older network work", async () => {
let fetches = 0;
let markFirstStarted: (() => void) | undefined;
let finishFirst: (() => void) | undefined;
const firstStarted = new Promise<void>((resolve) => {
markFirstStarted = resolve;
});
const firstBlocked = new Promise<void>((resolve) => {
finishFirst = resolve;
});
const provider = createProvider({
id: "dynamic",
auth: { apiKey: { name: "Test", resolve: async () => ({ auth: {} }) } },
models: [],
fetchModels: async () => {
fetches++;
await new Promise((resolve) => setTimeout(resolve, 5));
return [testModel("api-a", "listed")];
const current = fetches;
if (current === 1) {
markFirstStarted?.();
await firstBlocked;
}
return [testModel("api-a", `listed-${current}`)];
},
api: recordingStreams("a", []),
});
const store = new InMemoryModelsStore();
const refreshContext = {
credential: { type: "api_key" as const },
store: {
read: () => store.read("dynamic"),
write: (entry: ModelsStoreEntry) => store.write("dynamic", entry),
delete: () => store.delete("dynamic"),
},
allowNetwork: true,
};
const models = createModels({ modelsStore: store });
models.setProvider(provider);
expect(provider.getModels()).toEqual([]);
await Promise.all([provider.refreshModels?.(refreshContext), provider.refreshModels?.(refreshContext)]);
expect(fetches).toBe(1);
expect(provider.getModels().map((m) => m.id)).toEqual(["listed"]);
// a later refresh fetches again
await provider.refreshModels?.(refreshContext);
const first = models.refresh({ providers: ["dynamic"] });
await firstStarted;
const second = models.refresh({ providers: ["dynamic"] });
await expect(second).resolves.toMatchObject({ aborted: false });
await expect(first).resolves.toMatchObject({ aborted: false });
expect(fetches).toBe(2);
expect(provider.getModels().map((model) => model.id)).toEqual(["listed-2"]);
expect((await store.read("dynamic"))?.models.map((model) => model.id)).toEqual(["listed-2"]);
finishFirst?.();
await new Promise((resolve) => setTimeout(resolve, 0));
expect(provider.getModels().map((model) => model.id)).toEqual(["listed-2"]);
expect((await store.read("dynamic"))?.models.map((model) => model.id)).toEqual(["listed-2"]);
});
});

View file

@ -1,6 +1,6 @@
import { afterEach, describe, expect, it, vi } from "vitest";
import { createRadiusOAuth } from "../src/auth/oauth/radius.ts";
import type { AuthEvent, AuthInteraction } from "../src/auth/types.ts";
import type { AuthEvent, ProviderAuthInteraction } from "../src/auth/types.ts";
const GATEWAY = "https://radius.example";
@ -18,8 +18,9 @@ function requestUrl(input: unknown): string {
throw new Error(`Unsupported request input: ${String(input)}`);
}
function interaction(loginMethod: "browser" | "device-code", events: AuthEvent[] = []): AuthInteraction {
function interaction(loginMethod: "browser" | "device-code", events: AuthEvent[] = []): ProviderAuthInteraction {
return {
signal: new AbortController().signal,
prompt: async () => loginMethod,
notify: (event) => events.push(event),
};
@ -106,7 +107,10 @@ describe("Radius OAuth", () => {
const oauth = createRadiusOAuth({ name: "Radius", gateway: GATEWAY });
await expect(
oauth.refresh({ type: "oauth", access: "old-access", refresh: "old-refresh", expires: 0 }),
oauth.refresh(
{ type: "oauth", access: "old-access", refresh: "old-refresh", expires: 0 },
new AbortController().signal,
),
).resolves.toMatchObject({ access: "new-access", refresh: "new-refresh" });
expect(fetchMock).toHaveBeenCalledOnce();
});

View file

@ -2,6 +2,8 @@ import { afterEach, describe, expect, it, vi } from "vitest";
import { xaiOAuth } from "../src/auth/oauth/xai.ts";
import type { OAuthCredential } from "../src/auth/types.ts";
const neverAbortedSignal = new AbortController().signal;
function jsonResponse(body: unknown, status = 200): Response {
return new Response(JSON.stringify(body), {
status,
@ -53,7 +55,7 @@ function loginXaiForTest(options: {
signal?: AbortSignal;
}): Promise<OAuthCredential> {
return xaiOAuth.login({
signal: options.signal,
signal: options.signal ?? neverAbortedSignal,
prompt: () => {
throw new Error("Unexpected prompt");
},
@ -67,7 +69,10 @@ function loginXaiForTest(options: {
}
function refreshXaiForTest(refreshToken: string): Promise<OAuthCredential> {
return xaiOAuth.refresh({ type: "oauth", access: "old-access", refresh: refreshToken, expires: 0 });
return xaiOAuth.refresh(
{ type: "oauth", access: "old-access", refresh: refreshToken, expires: 0 },
neverAbortedSignal,
);
}
describe("xAI OAuth device flow", () => {

View file

@ -6,10 +6,74 @@
- Changed JSON and RPC `message_update` events to emit only `assistantMessageEvent` deltas, removing the cumulative `message` and `assistantMessageEvent.partial` fields that caused quadratic output growth. Clients that need partial messages must assemble deltas between `message_start` and `message_end`; the latter remains authoritative ([#7290](https://github.com/earendil-works/pi/issues/7290)).
- `ModelRegistry.getApiKeyAndHeaders()` now returns `ProviderHeaders` with `string | null` values and preserves `null` header-deletion markers. Extensions that inspect returned headers must handle `null`; extensions forwarding them to pi-ai streams should pass them through unchanged. This prevents placeholder OpenAI credentials from being sent through Cloudflare AI Gateway ([#7030](https://github.com/earendil-works/pi/issues/7030)).
- Changed `ModelRegistry.refresh()` to accept `ModelsRefreshOptions` and return `ModelsRefreshResult` instead of discarding cancellation and provider errors.
- Changed `ModelRuntime.setRuntimeApiKey()` to accept auth cancellation options rather than catalog refresh options. Call `refresh({ providers: [providerId], signal })` separately when remote freshness is required.
- Required config-form extension OAuth `refreshToken(credentials, signal)` callbacks to accept and honor a concrete abort signal.
- Replaced dynamic provider refresh context store access with the read-only `context.stored` snapshot and generation-checked `context.publish()` transaction.
**Providers built with `createProvider({ fetchModels })`:** no catalog-publication migration is required. Before and after, return the fetched models and register the resulting provider; `createProvider()` owns restoration, persistence, and in-memory publication.
```ts
// Before
const beforeProvider = createProvider({
// ...
fetchModels: async ({ signal }) => {
const response = await fetch(catalogUrl, { signal });
return parseModels(await response.json());
},
});
pi.registerProvider(beforeProvider);
// After: unchanged
const afterProvider = createProvider({
// ...
fetchModels: async ({ signal }) => {
const response = await fetch(catalogUrl, { signal });
return parseModels(await response.json());
},
});
pi.registerProvider(afterProvider);
```
**Handwritten native `Provider.refreshModels()`:** replace direct store access and pre-publication mutation with generation-guarded publications.
```ts
// Before
refreshModels: async (context) => {
const stored = await context.store.read();
if (stored) currentModels = stored.models;
if (!context.allowNetwork) return;
const refreshed = await fetchModels(context.signal);
currentModels = refreshed;
await context.store.write({ models: refreshed, checkedAt: Date.now() });
},
// After
refreshModels: async (context) => {
if (context.stored) {
const restored = context.stored.models;
if (!(await context.publish({
update: () => { currentModels = restored; },
}))) return;
}
if (!context.allowNetwork) return;
const refreshed = await fetchModels(context.signal);
if (context.signal.aborted) return;
await context.publish({
persist: { models: refreshed, checkedAt: Date.now() },
update: () => { currentModels = refreshed; },
});
},
```
For the config-form `pi.registerProvider(name, { refreshModels })`, callbacks that only return models remain unchanged; pi publishes the returned list. If such a callback previously used `context.store` for custom persistence, read `context.stored` and call `context.publish({ persist: entry })`. In `publish()`, omit `persist` to leave storage unchanged, pass a `ModelsStoreEntry` to write it, or pass `persist: null` to delete it.
### Added
- Added built-in Baseten provider support with `BASETEN_API_KEY` authentication and `zai-org/GLM-5.2` as the default model.
- Added `CredentialSynchronizationError` for credential changes that commit successfully but fail to synchronize local model state.
- Added chainable `pi.registerMarkdownTransformer()` hooks for display-only transformation of user and assistant Markdown.
- Added an experimental fullscreen UI mode, selectable through `--ui-mode fullscreen` or `/settings` ([#7304](https://github.com/earendil-works/pi/issues/7304)).
- Added runtime switching between regular and fullscreen UI modes through `/settings`.
@ -36,7 +100,13 @@
- Updated the packaged `brace-expansion` dependency to 5.0.8 to address GHSA-mh99-v99m-4gvg ([#7316](https://github.com/earendil-works/pi/issues/7316)).
- Fixed forced model availability refreshes remaining blocked behind a stalled earlier refresh ([#7301](https://github.com/earendil-works/pi/issues/7301), [#7421](https://github.com/earendil-works/pi/pull/7421) by [@a-yeyang](https://github.com/a-yeyang)).
- Fixed `/model` catalog refresh failures to identify every catalog that failed.
- Fixed provider login remaining stuck after saving credentials when a model catalog refresh stalls ([#7027](https://github.com/earendil-works/pi/issues/7027)).
- Fixed provider login remaining stuck after saving credentials when a model catalog refresh stalls by separating local credential consistency from bounded background freshness ([#7027](https://github.com/earendil-works/pi/issues/7027), [#7113](https://github.com/earendil-works/pi/issues/7113), [#7418](https://github.com/earendil-works/pi/issues/7418)).
- Fixed `/scoped-models` waiting for remote catalogs before rendering instead of showing cached models and cancelling refresh on close ([#7153](https://github.com/earendil-works/pi/issues/7153)).
- Fixed `/model <name>` waiting for catalog refresh before checking cached model matches ([#7443](https://github.com/earendil-works/pi/issues/7443)).
- Fixed stale availability snapshots and errors publishing after a newer availability pass.
- Fixed stale pi.dev, Radius, llama.cpp, and extension catalog refreshes publishing after a newer provider refresh.
- Fixed cancellation while waiting for file-backed credential or model-catalog locks, preventing cancelled mutations from running or committing later.
- Fixed concurrent in-memory credential mutations losing unrelated provider updates by serializing their read-modify-write sections.
## [0.83.0] - 2026-07-29

View file

@ -331,8 +331,8 @@ pi.registerProvider("corporate-ai", {
};
},
async refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
const tokens = await refreshAccessToken(credentials.refresh);
async refreshToken(credentials: OAuthCredentials, signal: AbortSignal): Promise<OAuthCredentials> {
const tokens = await refreshAccessToken(credentials.refresh, signal);
return {
refresh: tokens.refreshToken ?? credentials.refresh,
access: tokens.accessToken,
@ -688,7 +688,7 @@ interface ProviderConfig {
oauth?: {
name: string;
login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials>;
refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials>;
refreshToken(credentials: OAuthCredentials, signal: AbortSignal): Promise<OAuthCredentials>;
getApiKey(credentials: OAuthCredentials): string;
};
}

View file

@ -1707,7 +1707,9 @@ Register or override a model provider dynamically. Useful for proxies, custom en
Calls made during the extension factory function are queued and applied once the runner initialises. Calls made after that — for example from a command handler following a user setup flow — take effect immediately without requiring a `/reload`.
Dynamic providers can implement `refreshModels`. Pi calls it during model refresh, publishes the returned list synchronously through the provider, and passes the canonical credential/store/network/signal context. The extension decides whether to persist the catalog through `context.store`; live servers such as llama.cpp can ignore it.
Dynamic providers can implement `refreshModels`. Pi calls it during model refresh, publishes the returned list synchronously through the provider, and passes the canonical credential/stored-catalog/network/signal context. The extension decides whether to persist catalog metadata through generation-checked `context.publish({ persist: entry })`; live servers such as llama.cpp can return models without persisting them.
`context.signal` is always a concrete signal and provider callbacks must pass it to blocking I/O. Public `ModelRuntime.refresh()` and `ModelRegistry.refresh()` calls accept an optional signal and are unbounded when it is omitted; extensions and applications choose their own deadlines. Cancellation stops the caller waiting even if a provider ignores the signal, but cooperation is still required to stop the underlying work.
Extensions that need native provider auth, filtering, refresh, or stream behavior can register a complete `Provider` from `@earendil-works/pi-ai`. The provider becomes the composition base and `models.json` overrides still apply above it.
@ -1797,7 +1799,8 @@ pi.registerProvider("corporate-ai", {
const code = await callbacks.onPrompt({ message: "Enter code:" });
return { refresh: code, access: code, expires: Date.now() + 3600000 };
},
async refreshToken(credentials) {
async refreshToken(credentials, signal) {
signal.throwIfAborted();
// Refresh logic
return credentials;
},
@ -1818,7 +1821,7 @@ The object form accepts a complete pi-ai `Provider`, including native `auth`, `g
- `headers` - Custom headers to include in requests.
- `authHeader` - If true, adds `Authorization: Bearer` header automatically.
- `models` - Array of model definitions. If provided, replaces all existing models for this provider. Model definitions can set `baseUrl` to override the provider endpoint for that model.
- `refreshModels` - Async dynamic discovery callback. Its returned models replace extension-provided models. Use the scoped `context.store` only when results should persist.
- `refreshModels` - Async dynamic discovery callback. Its returned models replace extension-provided models. `context.stored` contains the persisted provider snapshot; use generation-checked `context.publish({ persist: entry })` only when updated catalog data should persist. Use `persist: null` to delete that snapshot.
- `oauth` - OAuth provider config for `/login` support. When provided, the provider appears in the login menu.
- `streamSimple` - Custom streaming implementation for non-standard APIs.

View file

@ -452,7 +452,7 @@ for (const provider of modelRuntime.getProviders()) {
}
// Runtime API key override (not persisted to disk)
modelRuntime.setRuntimeApiKey("anthropic", "sk-my-temp-key");
await modelRuntime.setRuntimeApiKey("anthropic", "sk-my-temp-key");
// Custom credential and model locations
const customRuntime = await ModelRuntime.create({
@ -469,6 +469,24 @@ const { session } = await createAgentSession({
});
```
`login()`, `logout()`, `setRuntimeApiKey()`, and `removeRuntimeApiKey()` resolve after the affected provider's cached/built-in catalog, composition, and availability snapshot are locally consistent. They do not wait for remote catalog freshness. If credentials were committed but local synchronization fails, they reject with the exported `CredentialSynchronizationError`; inspect its `providerId`, `operation`, `credential`, and `cause` fields instead of retrying the credential mutation blindly.
Public model/auth operations and `ModelRuntime.create({ signal })` accept optional abort signals and are unbounded when omitted. SDK applications own deadline policy for remote catalog freshness:
```typescript
const signal = AbortSignal.timeout(15_000);
const result = await modelRuntime.refresh({
providers: ["anthropic"],
signal,
});
if (result.aborted) console.warn("Catalog refresh timed out; using cached models");
for (const [providerId, error] of result.errors) {
console.warn(`Could not refresh ${providerId}:`, error);
}
```
A failed or timed-out network refresh does not undo a successful credential operation. `refresh()` starts a new provider generation, so it does not wait behind an older stalled refresh and stale generations cannot publish afterward.
> See [examples/sdk/09-api-keys-and-oauth.ts](../examples/sdk/09-api-keys-and-oauth.ts)
### System Prompt
@ -938,7 +956,7 @@ const modelRuntime = await ModelRuntime.create({
modelsPath: "/custom/agent/models.json",
});
if (process.env.MY_KEY) {
modelRuntime.setRuntimeApiKey("anthropic", process.env.MY_KEY);
await modelRuntime.setRuntimeApiKey("anthropic", process.env.MY_KEY);
}
// Inline tool
@ -1144,6 +1162,7 @@ AgentSessionRuntime
// Auth and Models
ModelRuntime // implements pi-ai Models and owns credential storage
ModelRegistry // synchronous extension compatibility facade
CredentialSynchronizationError
resolveCliModel
resolveModelScopeWithDiagnostics

View file

@ -124,7 +124,7 @@ async function loginAnthropic(callbacks: OAuthLoginCallbacks): Promise<OAuthCred
};
}
async function refreshAnthropicToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
async function refreshAnthropicToken(credentials: OAuthCredentials, signal: AbortSignal): Promise<OAuthCredentials> {
const response = await fetch(TOKEN_URL, {
method: "POST",
headers: { "Content-Type": "application/json" },
@ -133,6 +133,7 @@ async function refreshAnthropicToken(credentials: OAuthCredentials): Promise<OAu
client_id: CLIENT_ID,
refresh_token: credentials.refresh,
}),
signal,
});
if (!response.ok) {

View file

@ -274,7 +274,7 @@ async function loginGitLab(callbacks: OAuthLoginCallbacks): Promise<OAuthCredent
};
}
async function refreshGitLabToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
async function refreshGitLabToken(credentials: OAuthCredentials, signal: AbortSignal): Promise<OAuthCredentials> {
const response = await fetch(`${GITLAB_COM_URL}/oauth/token`, {
method: "POST",
headers: { "Content-Type": "application/x-www-form-urlencoded" },
@ -283,6 +283,7 @@ async function refreshGitLabToken(credentials: OAuthCredentials): Promise<OAuthC
grant_type: "refresh_token",
refresh_token: credentials.refresh,
}).toString(),
signal,
});
if (!response.ok) throw new Error(`Token refresh failed: ${await response.text()}`);
const data = (await response.json()) as {

View file

@ -25,7 +25,7 @@ const { session: customAuthSession } = await createAgentSession({
console.log("Session with custom auth and models locations");
customAuthSession.dispose();
modelRuntime.setRuntimeApiKey("anthropic", "sk-my-temp-key");
await modelRuntime.setRuntimeApiKey("anthropic", "sk-my-temp-key");
const { session: runtimeKeySession } = await createAgentSession({
sessionManager: SessionManager.inMemory(),
modelRuntime,

View file

@ -19,7 +19,7 @@ const modelRuntime = await ModelRuntime.create({
modelsPath: "/tmp/my-agent/models.json",
});
if (process.env.MY_ANTHROPIC_KEY) {
modelRuntime.setRuntimeApiKey("anthropic", process.env.MY_ANTHROPIC_KEY);
await modelRuntime.setRuntimeApiKey("anthropic", process.env.MY_ANTHROPIC_KEY);
}
const model = getModel("anthropic", "claude-sonnet-4-5");

View file

@ -71,7 +71,7 @@ const customRuntime = await ModelRuntime.create({
authPath: "/my/app/auth.json",
modelsPath: "/my/app/models.json",
});
customRuntime.setRuntimeApiKey("anthropic", process.env.MY_KEY!);
await customRuntime.setRuntimeApiKey("anthropic", process.env.MY_KEY!);
const resourceLoader = new DefaultResourceLoader({
systemPromptOverride: () => "You are helpful.",

View file

@ -86,11 +86,12 @@ export async function resolveCredentialForPrint(
modelRuntime: ModelRuntime,
kind: CredentialPrintKind,
minExpiryMs?: number,
signal?: AbortSignal,
): Promise<string> {
validateCredentialPrintArgs(args);
const credentialTypes = new Map<string, CredentialInfo["type"]>(
(await modelRuntime.listCredentials()).map((credential) => [credential.providerId, credential.type]),
(await modelRuntime.listCredentials({ signal })).map((credential) => [credential.providerId, credential.type]),
);
const models: Model<Api>[] = [];
if (args.provider) {
@ -118,12 +119,10 @@ export async function resolveCredentialForPrint(
if (kind === "api_key" && type === "oauth") continue;
if (kind === "bearer_token" && type !== "oauth") continue;
const auth = await modelRuntime.getAuth(
model,
kind === "bearer_token"
? { minOAuthValidityMs: minExpiryMs ?? DEFAULT_BEARER_TOKEN_MIN_EXPIRY_MS }
: undefined,
);
const auth = await modelRuntime.getAuth(model, {
...(kind === "bearer_token" ? { minOAuthValidityMs: minExpiryMs ?? DEFAULT_BEARER_TOKEN_MIN_EXPIRY_MS } : {}),
signal,
});
const authorization = Object.entries(auth?.auth.headers ?? {}).find(
([name]) => name.toLowerCase() === "authorization",
)?.[1];

View file

@ -26,13 +26,17 @@ function formatTokenCount(count: number): string {
/**
* List available models, optionally filtered by search pattern
*/
export async function listModels(modelRuntime: ModelRuntime, searchPattern?: string): Promise<void> {
export async function listModels(
modelRuntime: ModelRuntime,
searchPattern?: string,
signal?: AbortSignal,
): Promise<void> {
const loadError = modelRuntime.getError();
if (loadError) {
console.error(chalk.yellow(`Warning: errors loading models.json:\n${loadError}`));
}
const models = [...(await modelRuntime.getAvailable())];
const models = [...(await modelRuntime.getAvailable(undefined, { signal }))];
if (models.length === 0) {
console.log(formatNoModelsAvailableMessage());

View file

@ -39,6 +39,7 @@ export interface CreateAgentSessionServicesOptions {
agentDir?: string;
settingsManager?: SettingsManager;
modelRuntime?: ModelRuntime;
modelRuntimeSignal?: AbortSignal;
extensionFlagValues?: Map<string, boolean | string>;
resourceLoaderOptions?: Omit<DefaultResourceLoaderOptions, "cwd" | "agentDir" | "settingsManager">;
resourceLoaderReloadOptions?: ResourceLoaderReloadOptions;
@ -141,6 +142,7 @@ export async function createAgentSessionServices(
(await ModelRuntime.create({
authPath: join(agentDir, "auth.json"),
modelsPath: join(agentDir, "models.json"),
signal: options.modelRuntimeSignal,
}));
const settingsManager = options.settingsManager ?? SettingsManager.create(cwd, agentDir);
const resourceLoader = new DefaultResourceLoader({

View file

@ -1605,13 +1605,12 @@ export class AgentSession {
}
private async _cycleScopedModel(direction: "forward" | "backward"): Promise<ModelCycleResult | undefined> {
const checks = await Promise.all(
this._scopedModels.map(async (scoped) => ({
scoped,
auth: await this._modelRuntime.checkAuth(scoped.model.provider),
})),
const availableIds = new Set(
this._modelRuntime.getAvailableSnapshot().map((model) => `${model.provider}\0${model.id}`),
);
const scopedModels = this._scopedModels.filter((scoped) =>
availableIds.has(`${scoped.model.provider}\0${scoped.model.id}`),
);
const scopedModels = checks.filter(({ auth }) => auth !== undefined).map(({ scoped }) => scoped);
if (scopedModels.length <= 1) return undefined;
const currentModel = this.model;
@ -1640,7 +1639,7 @@ export class AgentSession {
}
private async _cycleAvailableModel(direction: "forward" | "backward"): Promise<ModelCycleResult | undefined> {
const availableModels = await this._modelRuntime.getAvailable();
const availableModels = this._modelRuntime.getAvailableSnapshot();
if (availableModels.length <= 1) return undefined;
const currentModel = this.model;

View file

@ -3,11 +3,13 @@
* Provider auth orchestration belongs to ModelRuntime and pi-ai Models.
*/
import type { Credential, CredentialInfo, CredentialStore } from "@earendil-works/pi-ai";
import type { AuthOperationOptions, Credential, CredentialInfo, CredentialStore } from "@earendil-works/pi-ai";
import { chmodSync, existsSync, mkdirSync, readFileSync, statSync, writeFileSync } from "fs";
import { dirname, join } from "path";
import lockfile from "proper-lockfile";
import { setTimeout as sleep } from "timers/promises";
import { getAgentDir } from "../config.ts";
import { raceWithAbortSignal } from "../utils/abort.ts";
import { normalizePath } from "../utils/paths.ts";
import { resolveConfigValue } from "./resolve-config-value.ts";
@ -39,7 +41,10 @@ function getFileRevision(path: string): string | undefined {
export interface AuthStorageBackend {
withLock<T>(fn: (current: string | undefined) => LockResult<T>): T;
withLockAsync<T>(fn: (current: string | undefined) => Promise<LockResult<T>>): Promise<T>;
withLockAsync<T>(
fn: (current: string | undefined) => Promise<LockResult<T>>,
options?: AuthOperationOptions,
): Promise<T>;
}
export class FileAuthStorageBackend implements AuthStorageBackend {
@ -111,7 +116,47 @@ export class FileAuthStorageBackend implements AuthStorageBackend {
}
}
async withLockAsync<T>(fn: (current: string | undefined) => Promise<LockResult<T>>): Promise<T> {
private async acquireLockAsync(
signal: AbortSignal | undefined,
onCompromised: (error: Error) => void,
): Promise<() => Promise<void>> {
const maxRetries = 10;
for (let attempt = 0; attempt <= maxRetries; attempt++) {
signal?.throwIfAborted();
let release: (() => Promise<void>) | undefined;
try {
release = await lockfile.lock(this.authPath, {
realpath: false,
retries: 0,
stale: 30000,
onCompromised,
});
} catch (error) {
signal?.throwIfAborted();
const code =
typeof error === "object" && error !== null && "code" in error
? String((error as { code?: unknown }).code)
: undefined;
if (code !== "ELOCKED" || attempt === maxRetries) throw error;
const delayMs = Math.min(Math.round((Math.random() + 1) * 100 * 2 ** attempt), 10000);
if (signal) await sleep(delayMs, undefined, { signal });
else await sleep(delayMs);
continue;
}
if (signal?.aborted) {
await release();
signal.throwIfAborted();
}
return release;
}
throw new Error("Failed to acquire auth storage lock");
}
async withLockAsync<T>(
fn: (current: string | undefined) => Promise<LockResult<T>>,
options?: AuthOperationOptions,
): Promise<T> {
options?.signal?.throwIfAborted();
this.ensureParentDir();
this.ensureFileExists();
@ -125,25 +170,17 @@ export class FileAuthStorageBackend implements AuthStorageBackend {
};
try {
release = await lockfile.lock(this.authPath, {
retries: {
retries: 10,
factor: 2,
minTimeout: 100,
maxTimeout: 10000,
randomize: true,
},
stale: 30000,
onCompromised: (err) => {
lockCompromised = true;
lockCompromisedError = err;
},
release = await this.acquireLockAsync(options?.signal, (error) => {
lockCompromised = true;
lockCompromisedError = error;
});
throwIfCompromised();
options?.signal?.throwIfAborted();
const current = existsSync(this.authPath) ? readFileSync(this.authPath, "utf-8") : undefined;
const { result, next } = await fn(current);
throwIfCompromised();
options?.signal?.throwIfAborted();
if (next !== undefined) {
writeFileSync(this.authPath, next, AUTH_FILE_WRITE_OPTIONS);
chmodSync(this.authPath, 0o600);
@ -164,6 +201,7 @@ export class FileAuthStorageBackend implements AuthStorageBackend {
export class InMemoryAuthStorageBackend implements AuthStorageBackend {
private value: string | undefined;
private asyncChain: Promise<unknown> = Promise.resolve();
withLock<T>(fn: (current: string | undefined) => LockResult<T>): T {
const { result, next } = fn(this.value);
@ -173,12 +211,23 @@ export class InMemoryAuthStorageBackend implements AuthStorageBackend {
return result;
}
async withLockAsync<T>(fn: (current: string | undefined) => Promise<LockResult<T>>): Promise<T> {
const { result, next } = await fn(this.value);
if (next !== undefined) {
this.value = next;
}
return result;
withLockAsync<T>(
fn: (current: string | undefined) => Promise<LockResult<T>>,
options?: AuthOperationOptions,
): Promise<T> {
const previous = this.asyncChain;
const operation = (async () => {
await previous.catch(() => {});
options?.signal?.throwIfAborted();
const { result, next } = await fn(this.value);
options?.signal?.throwIfAborted();
if (next !== undefined) {
this.value = next;
}
return result;
})();
this.asyncChain = operation.catch(() => {});
return raceWithAbortSignal(operation, options?.signal);
}
}
@ -248,32 +297,37 @@ export class AuthStorage implements CredentialStore {
}
}
private async reloadFromStorageAsync(): Promise<AuthStorageData> {
private async reloadFromStorageAsync(options?: AuthOperationOptions): Promise<AuthStorageData> {
return this.storage.withLockAsync(async (content) => {
const currentData = this.parseStorageData(content);
const revision = this.authPath ? getFileRevision(this.authPath) : undefined;
this.updateReadState(currentData, revision);
return { result: currentData };
});
}, options);
}
private async readLatestData(): Promise<AuthStorageData> {
private readLatestData(options?: AuthOperationOptions): Promise<AuthStorageData> {
options?.signal?.throwIfAborted();
if (this.authPath) {
const revision = getFileRevision(this.authPath);
if (revision !== undefined && revision === this.readState.revision) return this.readState.data;
if (revision !== undefined && revision === this.readState.revision) {
return Promise.resolve(this.readState.data);
}
}
if (options?.signal) return this.reloadFromStorageAsync(options);
if (!this.readState.reload) {
this.readState.reload = this.reloadFromStorageAsync().catch(() => this.readState.data);
}
try {
return await this.readState.reload;
} finally {
this.readState.reload = undefined;
const reload = this.reloadFromStorageAsync().catch(() => this.readState.data);
this.readState.reload = reload;
void reload.then(() => {
if (this.readState.reload === reload) this.readState.reload = undefined;
});
}
return this.readState.reload;
}
async read(provider: string): Promise<Credential | undefined> {
const credential = (await this.readLatestData())[provider];
async read(provider: string, options?: AuthOperationOptions): Promise<Credential | undefined> {
const credential = (await this.readLatestData(options))[provider];
options?.signal?.throwIfAborted();
if (credential?.type !== "api_key") return credential;
if (credential.key === undefined) return credential;
return { ...credential, key: resolveConfigValue(credential.key, credential.env) };
@ -282,6 +336,7 @@ export class AuthStorage implements CredentialStore {
async modify(
provider: string,
fn: (current: Credential | undefined) => Promise<Credential | undefined>,
options?: AuthOperationOptions,
): Promise<Credential | undefined> {
let latestData = this.readState.data;
let revision: string | undefined;
@ -297,28 +352,27 @@ export class AuthStorage implements CredentialStore {
const merged: AuthStorageData = { ...currentData, [provider]: next };
latestData = merged;
return { result: next, next: JSON.stringify(merged, null, 2) };
});
}, options);
this.updateReadState(latestData, revision);
return result;
}
async delete(provider: string): Promise<void> {
async delete(provider: string, options?: AuthOperationOptions): Promise<void> {
let latestData = this.readState.data;
await this.storage.withLockAsync(async (content) => {
const currentData = this.parseStorageData(content);
delete currentData[provider];
latestData = currentData;
return { result: undefined, next: JSON.stringify(currentData, null, 2) };
});
}, options);
this.updateReadState(latestData);
}
/** List credential metadata without resolving configured key values. */
async list(): Promise<readonly CredentialInfo[]> {
return Object.entries(await this.readLatestData()).map(([providerId, credential]) => ({
providerId,
type: credential.type,
}));
async list(options?: AuthOperationOptions): Promise<readonly CredentialInfo[]> {
const entries = Object.entries(await this.readLatestData(options));
options?.signal?.throwIfAborted();
return entries.map(([providerId, credential]) => ({ providerId, type: credential.type }));
}
}

View file

@ -1454,7 +1454,7 @@ export interface ProviderConfig {
models?: ProviderModelConfig[];
/**
* Refresh this provider's model list. The returned list replaces extension-provided models.
* Use context.store explicitly when the catalog should persist across sessions.
* Use context.publish({ persist: entry }) when the catalog should persist across sessions.
*/
refreshModels?(context: RefreshModelsContext): Promise<ProviderModelConfig[]>;
/** OAuth provider for /login support. The `id` is set automatically from the provider name. */
@ -1466,7 +1466,7 @@ export interface ProviderConfig {
/** Run the login flow, return credentials to persist. */
login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials>;
/** Refresh expired credentials, return updated credentials to persist. */
refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials>;
refreshToken(credentials: OAuthCredentials, signal: AbortSignal): Promise<OAuthCredentials>;
/** Convert credentials to API key string for the provider. */
getApiKey(credentials: OAuthCredentials): string;
/** Legacy synchronous credential-dependent model projection. */

View file

@ -5,6 +5,8 @@ import type {
Context,
Model,
ModelsApiStreamOptions,
ModelsRefreshOptions,
ModelsRefreshResult,
Provider,
ProviderHeaders,
} from "@earendil-works/pi-ai";
@ -34,8 +36,8 @@ export class ModelRegistry {
}
/** Reload models.json asynchronously. Await before making synchronous registry reads. */
async refresh(): Promise<void> {
await this.runtime.refresh();
refresh(options?: ModelsRefreshOptions): Promise<ModelsRefreshResult> {
return this.runtime.refresh(options);
}
getError(): string | undefined {

View file

@ -3,7 +3,13 @@
*/
import type { ThinkingLevel } from "@earendil-works/pi-agent-core";
import { type Api, type KnownProvider, type Model, modelsAreEqual } from "@earendil-works/pi-ai";
import {
type Api,
type AuthOperationOptions,
type KnownProvider,
type Model,
modelsAreEqual,
} from "@earendil-works/pi-ai";
import chalk from "chalk";
import { minimatch } from "minimatch";
import { isValidThinkingLevel } from "../cli/args.ts";
@ -271,11 +277,11 @@ export interface ResolveModelScopeResult {
diagnostics: ModelScopeDiagnostic[];
}
export async function resolveModelScopeWithDiagnostics(
export function resolveModelScopeFromModels(
patterns: string[],
modelRuntime: ModelRuntime,
): Promise<ResolveModelScopeResult> {
const availableModels = [...(await modelRuntime.getAvailable())];
models: readonly Model<Api>[],
): ResolveModelScopeResult {
const availableModels = [...models];
const scopedModels: ScopedModel[] = [];
const diagnostics: ModelScopeDiagnostic[] = [];
@ -353,8 +359,20 @@ export async function resolveModelScopeWithDiagnostics(
return { scopedModels, diagnostics };
}
export async function resolveModelScope(patterns: string[], modelRuntime: ModelRuntime): Promise<ScopedModel[]> {
const { scopedModels, diagnostics } = await resolveModelScopeWithDiagnostics(patterns, modelRuntime);
export async function resolveModelScopeWithDiagnostics(
patterns: string[],
modelRuntime: ModelRuntime,
options?: AuthOperationOptions,
): Promise<ResolveModelScopeResult> {
return resolveModelScopeFromModels(patterns, await modelRuntime.getAvailable(undefined, options));
}
export async function resolveModelScope(
patterns: string[],
modelRuntime: ModelRuntime,
options?: AuthOperationOptions,
): Promise<ScopedModel[]> {
const { scopedModels, diagnostics } = await resolveModelScopeWithDiagnostics(patterns, modelRuntime, options);
for (const diagnostic of diagnostics) {
console.warn(chalk.yellow(`Warning: ${diagnostic.message}`));
}
@ -661,7 +679,7 @@ export async function findInitialModel(options: {
}
// 4. Try first available model with valid API key
const availableModels = [...(await modelRuntime.getAvailable())];
const availableModels = [...modelRuntime.getAvailableSnapshot()];
if (availableModels.length > 0) {
// Try to find a default model from known providers
@ -722,7 +740,7 @@ export async function restoreModelFromSession(
}
// Try to find any available model
const availableModels = [...(await modelRuntime.getAvailable())];
const availableModels = [...modelRuntime.getAvailableSnapshot()];
if (availableModels.length > 0) {
// Try to find a default model from known providers

View file

@ -6,6 +6,7 @@ import {
type AssistantMessageEventStream,
type AuthCheck,
type AuthInteraction,
type AuthOperationOptions,
type AuthResult,
type AuthType,
type Context,
@ -31,6 +32,7 @@ import {
} from "@earendil-works/pi-ai";
import * as builtinProviderCatalog from "@earendil-works/pi-ai/providers/all";
import { getAgentDir } from "../config.ts";
import { operationSignal, raceWithAbortSignal } from "../utils/abort.ts";
import { AuthStorage as DefaultAuthStorage } from "./auth-storage.ts";
import { ModelConfig } from "./model-config.ts";
import { FileModelsStore, InMemoryCodingAgentModelsStore } from "./models-store.ts";
@ -67,16 +69,38 @@ export interface CreateModelRuntimeOptions {
/** Timeout for the create-time network model refresh. */
modelRefreshTimeoutMs?: number;
catalogBaseUrl?: string;
/** Optional caller cancellation for initial cache restoration and availability checks. */
signal?: AbortSignal;
}
export interface ModelRuntimeAuthOverrides {
export interface ModelRuntimeAuthOverrides extends AuthOperationOptions {
apiKey?: string;
env?: Record<string, string>;
/** Require this much remaining OAuth-token validity; defaults to five minutes. */
minOAuthValidityMs?: number;
}
const DEFAULT_MODEL_REFRESH_TIMEOUT_MS = 15_000;
export type CredentialSynchronizationOperation = "login" | "logout" | "setRuntimeApiKey" | "removeRuntimeApiKey";
/** Credentials changed successfully, but the local model/auth snapshot could not be synchronized. */
export class CredentialSynchronizationError extends Error {
readonly providerId: string;
readonly operation: CredentialSynchronizationOperation;
readonly credential: Credential | undefined;
constructor(
providerId: string,
operation: CredentialSynchronizationOperation,
credential: Credential | undefined,
options: ErrorOptions,
) {
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: ProviderHeaders | undefined,
@ -113,9 +137,11 @@ export class ModelRuntime implements Models {
storedProviders: new Set(),
auth: new Map(),
};
private availabilityRefresh: Promise<void> | undefined;
private availabilityRefreshSeq = 0;
private availabilityErrorSeq = 0;
private readonly providerAvailabilitySeq = new Map<string, number>();
private availabilityError: string | undefined;
private readonly credentialOperations = new Map<string, Promise<unknown>>();
private constructor(
credentials: RuntimeCredentials,
@ -164,12 +190,16 @@ export class ModelRuntime implements Models {
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 ?? DEFAULT_MODEL_REFRESH_TIMEOUT_MS)
: 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) clearTimeout(timeout);
}
@ -242,23 +272,20 @@ export class ModelRuntime implements Models {
};
}
private async runAvailabilityRefresh(seq: number): Promise<void> {
private async runAvailabilityRefresh(seq: number, errorSeq: number, signal: AbortSignal): Promise<void> {
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): Promise<[string, AuthCheck | undefined]> => [
provider.id,
await this.models.checkAuth(provider.id),
await this.models.checkAuth(provider.id, { signal }),
],
),
),
this.credentials.list(),
this.credentials.list({ signal }),
]);
// A newer rebuild was requested while this one was in flight; drop this
// result so a slow, superseded refresh cannot clobber the snapshot with
// stale data.
if (seq !== this.availabilityRefreshSeq) return;
const auth = new Map(checks);
const configuredProviders = new Set(
@ -273,39 +300,75 @@ export class ModelRuntime implements Models {
storedProviders: new Set(credentials.map((entry) => entry.providerId)),
auth,
};
this.availabilityError = undefined;
if (errorSeq === this.availabilityErrorSeq) this.availabilityError = undefined;
}
private queueAvailabilityRefresh(): Promise<void> {
private queueAvailabilityRefresh(signal?: AbortSignal): Promise<void> {
const seq = ++this.availabilityRefreshSeq;
const refresh = this.runAvailabilityRefresh(seq);
const recorded = refresh.catch((error) => {
// Only the latest requested rebuild owns the error state.
if (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. */
private refreshAvailability(): Promise<void> {
return this.availabilityRefresh ?? this.queueAvailabilityRefresh();
}
/**
* Mutations must observe a rebuild that starts after their state change, and a
* stuck in-flight refresh must not block them. Start a fresh, independent
* rebuild instead of chaining onto the pending one. The sequence guard in
* runAvailabilityRefresh ensures a superseded rebuild cannot clobber its result.
*/
private forceRefreshAvailability(): Promise<void> {
return this.queueAvailabilityRefresh();
private async refreshProviderAvailability(providerId: string, signal: AbortSignal): Promise<void> {
// 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(): readonly Provider[] {
@ -324,24 +387,25 @@ export class ModelRuntime implements Models {
return this.models.getModel(providerId, modelId);
}
async checkAuth(providerId: string): Promise<AuthCheck | undefined> {
return this.models.checkAuth(providerId);
async checkAuth(providerId: string, options?: AuthOperationOptions): Promise<AuthCheck | undefined> {
return this.models.checkAuth(providerId, options);
}
async getAvailable(providerId?: string): Promise<readonly Model<Api>[]> {
async getAvailable(providerId?: string, options?: AuthOperationOptions): Promise<readonly Model<Api>[]> {
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;
}
@ -413,32 +477,71 @@ export class ModelRuntime implements Models {
};
}
async setRuntimeApiKey(
private enqueueCredentialOperation<T>(providerId: string, signal: AbortSignal, task: () => Promise<T>): Promise<T> {
const previous = this.credentialOperations.get(providerId) ?? Promise.resolve();
let markStarted: (() => void) | undefined;
const started = new Promise<void>((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);
}
private async synchronizeCredentialState(
providerId: string,
apiKey: string,
refreshOptions: ModelsRefreshOptions = {},
operation: CredentialSynchronizationOperation,
credential: Credential | undefined,
signal: AbortSignal,
): Promise<void> {
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);
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 });
}
}
async removeRuntimeApiKey(providerId: string): Promise<void> {
this.credentials.removeRuntimeApiKey(providerId);
await this.refresh({ allowNetwork: this.modelNetworkEnabled });
setRuntimeApiKey(providerId: string, apiKey: string, options: AuthOperationOptions = {}): Promise<void> {
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,
);
});
}
listCredentials(): Promise<readonly CredentialInfo[]> {
return this.credentials.list();
removeRuntimeApiKey(providerId: string, options: AuthOperationOptions = {}): Promise<void> {
const signal = operationSignal(options.signal);
return this.enqueueCredentialOperation(providerId, signal, async () => {
this.credentials.removeRuntimeApiKey(providerId);
await this.synchronizeCredentialState(providerId, "removeRuntimeApiKey", undefined, signal);
});
}
listCredentials(options?: AuthOperationOptions): Promise<readonly CredentialInfo[]> {
return this.credentials.list(options);
}
getProviderAuthStatus(providerId: string): AuthStatus {
@ -459,7 +562,11 @@ export class ModelRuntime implements Models {
): Promise<{ provider: Provider; model: Model<Api>; options: StreamOptions }> {
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 ?? {};
@ -518,27 +625,32 @@ export class ModelRuntime implements Models {
return this.streamSimple(model, context, options).result();
}
async login(providerId: string, type: AuthType, interaction: AuthInteraction): Promise<Credential> {
const credential = await this.models.login(providerId, type, interaction);
const timeoutSignal = AbortSignal.timeout(DEFAULT_MODEL_REFRESH_TIMEOUT_MS);
await this.refresh({
allowNetwork: this.modelNetworkEnabled,
signal: interaction.signal ? AbortSignal.any([interaction.signal, timeoutSignal]) : timeoutSignal,
login(providerId: string, type: AuthType, interaction: AuthInteraction): Promise<Credential> {
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;
});
return credential;
}
async logout(providerId: string): Promise<void> {
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 });
logout(providerId: string, options: AuthOperationOptions = {}): Promise<void> {
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: ModelsRefreshOptions = {}): Promise<ModelsRefreshResult> {
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,
@ -549,13 +661,28 @@ export class ModelRuntime implements Models {
aborted: refreshOptions.signal?.aborted ?? false,
errors: new Map(),
};
const errors = new Map(result.errors);
this.updateModelSnapshot();
try {
await this.forceRefreshAvailability();
} catch {
// Availability errors are recorded by forceRefreshAvailability; refreshed models remain usable.
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)));
}
}
}),
);
} 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: Provider): void {

View file

@ -1,5 +1,5 @@
import { join } from "node:path";
import type { ModelsStore, ModelsStoreEntry } from "@earendil-works/pi-ai";
import type { ModelsStore, ModelsStoreEntry, ModelsStoreOperationOptions } from "@earendil-works/pi-ai";
import { getAgentDir } from "../config.ts";
import { type AuthStorageBackend, FileAuthStorageBackend } from "./auth-storage.ts";
@ -8,15 +8,19 @@ type StoredModels = Record<string, ModelsStoreEntry>;
export class InMemoryCodingAgentModelsStore implements ModelsStore {
private readonly entries = new Map<string, ModelsStoreEntry>();
async read(providerId: string): Promise<ModelsStoreEntry | undefined> {
return this.entries.get(providerId);
async read(providerId: string, options?: ModelsStoreOperationOptions): Promise<ModelsStoreEntry | undefined> {
options?.signal?.throwIfAborted();
const entry = this.entries.get(providerId);
return entry ? structuredClone(entry) : undefined;
}
async write(providerId: string, entry: ModelsStoreEntry): Promise<void> {
this.entries.set(providerId, entry);
async write(providerId: string, entry: ModelsStoreEntry, options?: ModelsStoreOperationOptions): Promise<void> {
options?.signal?.throwIfAborted();
this.entries.set(providerId, structuredClone(entry));
}
async delete(providerId: string): Promise<void> {
async delete(providerId: string, options?: ModelsStoreOperationOptions): Promise<void> {
options?.signal?.throwIfAborted();
this.entries.delete(providerId);
}
}
@ -33,25 +37,26 @@ export class FileModelsStore implements ModelsStore {
return content ? (JSON.parse(content) as StoredModels) : {};
}
async read(providerId: string): Promise<ModelsStoreEntry | undefined> {
return this.storage.withLock((content) => ({
result: structuredClone(this.parse(content)[providerId]),
}));
async read(providerId: string, options?: ModelsStoreOperationOptions): Promise<ModelsStoreEntry | undefined> {
return this.storage.withLockAsync(
async (content) => ({ result: structuredClone(this.parse(content)[providerId]) }),
options,
);
}
async write(providerId: string, entry: ModelsStoreEntry): Promise<void> {
async write(providerId: string, entry: ModelsStoreEntry, options?: ModelsStoreOperationOptions): Promise<void> {
await this.storage.withLockAsync(async (content) => {
const current = this.parse(content);
current[providerId] = structuredClone(entry);
return { result: undefined, next: JSON.stringify(current, null, 2) };
});
}, options);
}
async delete(providerId: string): Promise<void> {
async delete(providerId: string, options?: ModelsStoreOperationOptions): Promise<void> {
await this.storage.withLockAsync(async (content) => {
const current = this.parse(content);
delete current[providerId];
return { result: undefined, next: JSON.stringify(current, null, 2) };
});
}, options);
}
}

View file

@ -35,7 +35,7 @@ export interface ExtensionOAuthConfig {
/** @deprecated Retained for extension source compatibility; ignored by canonical auth flows. */
usesCallbackServer?: boolean;
login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials>;
refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials>;
refreshToken(credentials: OAuthCredentials, signal: AbortSignal): Promise<OAuthCredentials>;
getApiKey(credentials: OAuthCredentials): string;
modifyModels?(models: Model<Api>[], credentials: OAuthCredentials): Model<Api>[];
}
@ -242,7 +242,7 @@ function adaptOAuth(config: ExtensionOAuthConfig): OAuthAuth {
});
return { ...credential, type: "oauth" };
},
refresh: async (credential) => ({ ...(await config.refreshToken(credential)), type: "oauth" }),
refresh: async (credential, signal) => ({ ...(await config.refreshToken(credential, signal)), type: "oauth" }),
toAuth: async (credential) => ({ apiKey: config.getApiKey(credential) }),
};
}
@ -476,18 +476,23 @@ export function composeModelProvider(
base?.refreshModels || extension?.refreshModels || extension?.oauth?.modifyModels
? async (context) => {
await base?.refreshModels?.(context);
if (extension?.refreshModels) {
const refreshed = await extension.refreshModels(context);
if (!context.signal?.aborted) {
// Validate before publishing the new synchronous list.
applyExtension(providerId, applyModelsJson(providerId, base?.getModels() ?? [], config), {
...extension,
models: refreshed,
});
refreshedExtensionModels = refreshed;
}
}
extensionOAuthCredential = context.credential?.type === "oauth" ? context.credential : undefined;
let refreshed: NonNullable<ProviderConfigInput["models"]> | undefined;
if (extension?.refreshModels) refreshed = await extension.refreshModels(context);
if (context.signal.aborted) return;
const oauthCredential = context.credential?.type === "oauth" ? context.credential : undefined;
await context.publish({
update: () => {
if (refreshed) {
// Validate before publishing the new synchronous list.
applyExtension(providerId, applyModelsJson(providerId, base?.getModels() ?? [], config), {
...extension,
models: refreshed,
});
refreshedExtensionModels = refreshed;
}
extensionOAuthCredential = oauthCredential;
},
});
}
: undefined,
filterModels: base?.filterModels

View file

@ -47,77 +47,85 @@ export function withRemoteCatalog(
localGeneratedAt?: number,
): Provider {
let dynamicModels: readonly Model<Api>[] = [];
let inflightRefresh: Promise<void> | undefined;
return {
...provider,
getModels: () => mergeModels(provider.getModels(), dynamicModels),
refreshModels: (context) => {
inflightRefresh ??= (async () => {
try {
const stored = await context.store.read();
dynamicModels = remoteModels(stored, localGeneratedAt).filter((model) => model.provider === provider.id);
if (!context.allowNetwork || context.signal?.aborted) return;
if (
!context.force &&
stored?.checkedAt !== undefined &&
stored.lastModified !== undefined &&
Date.now() - stored.checkedAt < REMOTE_CATALOG_REFRESH_INTERVAL_MS
) {
return;
}
refreshModels: async (context) => {
const stored = context.stored;
const restored = remoteModels(stored, localGeneratedAt).filter((model) => model.provider === provider.id);
if (
!(await context.publish({
update: () => {
dynamicModels = restored;
},
}))
) {
return;
}
if (!context.allowNetwork || context.signal.aborted) return;
if (
!context.force &&
stored?.checkedAt !== undefined &&
stored.lastModified !== undefined &&
Date.now() - stored.checkedAt < REMOTE_CATALOG_REFRESH_INTERVAL_MS
) {
return;
}
// Only revalidate when a cached body backs the validator, so a 304 can never
// leave the overlay empty.
const validator = stored?.models.length ? stored.etag : undefined;
const url = new URL(`/api/models/providers/${encodeURIComponent(provider.id)}`, catalogBaseUrl);
const response = await fetch(url, {
headers: {
accept: "application/json",
"User-Agent": getPiUserAgent(VERSION),
...(validator ? { "if-none-match": validator } : {}),
},
signal: context.signal,
});
if (context.signal?.aborted) return;
const checkedAt = Date.now();
// Unchanged: dynamicModels already holds the stored overlay, so only the
// freshness window moves.
if (response.status === 304 && stored) {
await context.store.write({ ...stored, checkedAt });
return;
}
if (response.status === 404 || response.status === 501) {
await context.store.write({
...(stored ?? { models: [] }),
checkedAt,
lastModified: 0,
etag: undefined,
});
return;
}
if (!response.ok) {
// Transient failure: the cached body and its validator stay valid, so keep the
// etag and let the next refresh revalidate instead of downloading the catalog.
await context.store.write({ ...(stored ?? { models: [] }), checkedAt });
throw new Error(`Model catalog request failed for ${provider.id}: ${response.status}`);
}
const refreshed = parseCatalog(provider.id, await response.json());
const lastModified = Date.parse(response.headers.get("last-modified") ?? "");
if (context.signal?.aborted) return;
const entry = {
models: refreshed,
// Only revalidate when a cached body backs the validator, so a 304 can never
// leave the overlay empty.
const validator = stored?.models.length ? stored.etag : undefined;
const url = new URL(`/api/models/providers/${encodeURIComponent(provider.id)}`, catalogBaseUrl);
const response = await fetch(url, {
headers: {
accept: "application/json",
"User-Agent": getPiUserAgent(VERSION),
...(validator ? { "if-none-match": validator } : {}),
},
signal: context.signal,
});
if (context.signal.aborted) return;
const checkedAt = Date.now();
// Unchanged: dynamicModels already holds the stored overlay, so only the
// freshness window moves.
if (response.status === 304 && stored) {
await context.publish({ persist: { ...stored, checkedAt } });
return;
}
if (response.status === 404 || response.status === 501) {
await context.publish({
persist: {
...(stored ?? { models: [] }),
checkedAt,
lastModified: Number.isNaN(lastModified) ? 0 : lastModified,
etag: response.headers.get("etag") ?? undefined,
};
dynamicModels = remoteModels(entry, localGeneratedAt);
await context.store.write(entry);
} finally {
inflightRefresh = undefined;
}
})();
return inflightRefresh;
lastModified: 0,
etag: undefined,
},
});
return;
}
if (!response.ok) {
// Transient failure: the cached body and its validator stay valid, so keep the
// etag and let the next refresh revalidate instead of downloading the catalog.
await context.publish({ persist: { ...(stored ?? { models: [] }), checkedAt } });
throw new Error(`Model catalog request failed for ${provider.id}: ${response.status}`);
}
const refreshed = parseCatalog(provider.id, await response.json());
const lastModified = Date.parse(response.headers.get("last-modified") ?? "");
if (context.signal.aborted) return;
const entry = {
models: refreshed,
checkedAt,
lastModified: Number.isNaN(lastModified) ? 0 : lastModified,
etag: response.headers.get("etag") ?? undefined,
};
const published = remoteModels(entry, localGeneratedAt);
await context.publish({
persist: entry,
update: () => {
dynamicModels = published;
},
});
},
};
}

View file

@ -1,4 +1,4 @@
import type { Credential, CredentialInfo, CredentialStore } from "@earendil-works/pi-ai";
import type { AuthOperationOptions, Credential, CredentialInfo, CredentialStore } from "@earendil-works/pi-ai";
/** Async credential store overlay for non-persistent runtime API keys. */
export class RuntimeCredentials implements CredentialStore {
@ -21,13 +21,15 @@ export class RuntimeCredentials implements CredentialStore {
return this.overrides.has(providerId);
}
async read(providerId: string): Promise<Credential | undefined> {
async read(providerId: string, options?: AuthOperationOptions): Promise<Credential | undefined> {
options?.signal?.throwIfAborted();
const override = this.overrides.get(providerId);
return override ? { type: "api_key", key: override } : this.store.read(providerId);
return override ? { type: "api_key", key: override } : this.store.read(providerId, options);
}
async list(): Promise<readonly CredentialInfo[]> {
const entries = new Map((await this.store.list()).map((entry) => [entry.providerId, entry]));
async list(options?: AuthOperationOptions): Promise<readonly CredentialInfo[]> {
const entries = new Map((await this.store.list(options)).map((entry) => [entry.providerId, entry]));
options?.signal?.throwIfAborted();
for (const providerId of this.overrides.keys()) {
entries.set(providerId, { providerId, type: "api_key" });
}
@ -37,12 +39,14 @@ export class RuntimeCredentials implements CredentialStore {
modify(
providerId: string,
fn: (current: Credential | undefined) => Promise<Credential | undefined>,
options?: AuthOperationOptions,
): Promise<Credential | undefined> {
return this.store.modify(providerId, fn);
return this.store.modify(providerId, fn, options);
}
async delete(providerId: string): Promise<void> {
async delete(providerId: string, options?: AuthOperationOptions): Promise<void> {
options?.signal?.throwIfAborted();
await this.store.delete(providerId, options);
this.overrides.delete(providerId);
await this.store.delete(providerId);
}
}

View file

@ -48,9 +48,16 @@ export default function llamaExtension(pi: ExtensionAPI): void {
client: LlamaClient,
catalog?: LlamaModelInfo[],
): Promise<LlamaModelInfo[]> => {
const current = catalog ?? (await client.list());
const signal = AbortSignal.timeout(15_000);
const current = catalog ?? (await client.list({ signal }));
provider.setCatalog(current, client.serverUrl);
await ctx.modelRegistry.refresh();
const result = await ctx.modelRegistry.refresh({
providers: [LLAMA_PROVIDER_ID],
signal,
});
if (result.aborted) throw new Error("Model catalog refresh timed out.");
const refreshError = result.errors.get(LLAMA_PROVIDER_ID);
if (refreshError) throw refreshError;
return current;
};

View file

@ -111,20 +111,36 @@ export function createLlamaProvider(): LlamaProviderController {
},
getModels: () => models,
refreshModels: async (context: RefreshModelsContext): Promise<void> => {
const stored = await context.store.read();
if (stored) {
models = stored.models.filter(
if (context.stored) {
const restored = context.stored.models.filter(
(model): model is Model<"openai-completions"> =>
model.provider === LLAMA_PROVIDER_ID && model.api === "openai-completions",
);
if (
!(await context.publish({
update: () => {
models = restored;
},
}))
) {
return;
}
}
if (!context.allowNetwork || context.signal?.aborted || context.credential?.type !== "api_key") return;
if (!context.allowNetwork || context.signal.aborted || context.credential?.type !== "api_key") return;
const serverUrl = credentialServerUrl(context.credential);
if (!serverUrl) return;
const catalog = await new LlamaClient(serverUrl, context.credential.key).list({ signal: context.signal });
setCatalog(catalog, serverUrl);
if (!context.signal?.aborted) await context.store.write({ models, checkedAt: Date.now() });
if (context.signal.aborted) return;
const refreshed = catalog
.filter((model) => model.status.value === "loaded")
.map((model) => toPiModel(model, serverUrl));
await context.publish({
persist: { models: refreshed, checkedAt: Date.now() },
update: () => {
models = refreshed;
},
});
},
stream: (model, context, options) => stream(model, context, options as ProviderStreamOptions | undefined),
streamSimple: (model, context, options) => streamSimple(model, context, options),

View file

@ -179,6 +179,8 @@ export {
} from "./core/model-resolver.ts";
export {
type CreateModelRuntimeOptions,
CredentialSynchronizationError,
type CredentialSynchronizationOperation,
ModelRuntime,
type ModelRuntimeAuthOverrides,
} from "./core/model-runtime.ts";

View file

@ -155,8 +155,15 @@ async function runCredentialPrintCommand(args: string[]): Promise<boolean> {
try {
validateCredentialPrintArgs(parsed);
const modelRuntime = await ModelRuntime.create({ allowModelNetwork: false });
const credential = await resolveCredentialForPrint(parsed, modelRuntime, command.kind, command.minExpiryMs);
const signal = AbortSignal.timeout(15_000);
const modelRuntime = await ModelRuntime.create({ allowModelNetwork: false, signal });
const credential = await resolveCredentialForPrint(
parsed,
modelRuntime,
command.kind,
command.minExpiryMs,
signal,
);
process.stdout.write(`${credential}\n`);
} catch (error) {
const message = error instanceof CredentialPrintError ? error.message : "Failed to resolve credential";
@ -687,6 +694,7 @@ export async function main(args: string[], options?: MainOptions) {
cwd,
agentDir,
settingsManager: runtimeSettingsManager,
modelRuntimeSignal: AbortSignal.timeout(15_000),
extensionFlagValues: parsed.unknownFlags,
resourceLoaderReloadOptions: shouldResolveProjectTrust
? {
@ -740,7 +748,9 @@ export async function main(args: string[], options?: MainOptions) {
const modelPatterns = parsed.models ?? settingsManager.getEnabledModels();
const scopedModels =
modelPatterns && modelPatterns.length > 0 ? await resolveModelScope(modelPatterns, modelRuntime) : [];
modelPatterns && modelPatterns.length > 0
? await resolveModelScope(modelPatterns, modelRuntime, { signal: AbortSignal.timeout(15_000) })
: [];
const {
options: sessionOptions,
cliThinkingFromModel,
@ -761,8 +771,7 @@ export async function main(args: string[], options?: MainOptions) {
message: "--api-key requires a model to be specified via --model, --provider/--model, or --models",
});
} else {
await modelRuntime.setRuntimeApiKey(sessionOptions.model.provider, parsed.apiKey, { allowNetwork: false });
await services.modelRuntime.getAvailable();
await modelRuntime.setRuntimeApiKey(sessionOptions.model.provider, parsed.apiKey);
}
}
@ -811,7 +820,7 @@ export async function main(args: string[], options?: MainOptions) {
if (parsed.listModels !== undefined) {
const searchPattern = typeof parsed.listModels === "string" ? parsed.listModels : undefined;
await listModels(modelRuntime, searchPattern);
await listModels(modelRuntime, searchPattern, AbortSignal.timeout(15_000));
process.exit(0);
}
@ -862,7 +871,12 @@ export async function main(args: string[], options?: MainOptions) {
// RPC refreshes catalogs here in the background; interactive mode starts its refresh after TUI initialization.
if (!offlineMode && appMode === "rpc") {
void modelRuntime.refresh().catch(() => {});
const controller = new AbortController();
const timeout = setTimeout(() => controller.abort(), 15_000);
void modelRuntime
.refresh({ signal: controller.signal })
.catch(() => {})
.finally(() => clearTimeout(timeout));
}
if (appMode === "rpc") {

View file

@ -186,12 +186,21 @@ export class ModelSelectorComponent extends Container implements Focusable {
this.loadModelsFromSnapshot();
this.filterModels(this.searchInput.getValue());
this.tui.requestRender();
} catch (error) {
if (this.closed) return;
this.refreshStatusMessage = "";
this.errorMessage = timedOut
? "Model refresh timed out; showing cached models."
: `Could not refresh model catalogs: ${error instanceof Error ? error.message : String(error)}`;
this.updateList();
this.tui.requestRender();
} finally {
if (this.refreshTimeout) clearTimeout(this.refreshTimeout);
}
}
private close(): void {
dispose(): void {
if (this.closed) return;
this.closed = true;
if (this.refreshTimeout) clearTimeout(this.refreshTimeout);
this.refreshAbortController.abort();
@ -341,7 +350,7 @@ export class ModelSelectorComponent extends Container implements Focusable {
}
// Escape or Ctrl+C
else if (kb.matches(keyData, "tui.select.cancel")) {
this.close();
this.dispose();
this.onCancelCallback();
}
// Pass everything else to search input
@ -352,7 +361,7 @@ export class ModelSelectorComponent extends Container implements Focusable {
}
private handleSelect(model: Model<any>): void {
this.close();
this.dispose();
// Save as new default
this.settingsManager.setDefaultModelAndProvider(model.provider, model.id);
this.onSelectCallback(model);

View file

@ -74,6 +74,7 @@ interface ModelItem {
export interface ModelsConfig {
allModels: Model<any>[];
enabledModelIds: string[] | null;
refreshStatus?: string;
}
export interface ModelsCallbacks {
@ -110,6 +111,7 @@ export class ScopedModelsSelectorComponent extends Container implements Focusabl
private callbacks: ModelsCallbacks;
private maxVisible = 8;
private isDirty = false;
private refreshStatusText?: Text;
constructor(config: ModelsConfig, callbacks: ModelsCallbacks) {
super();
@ -144,6 +146,10 @@ export class ScopedModelsSelectorComponent extends Container implements Focusabl
// Footer hint
this.addChild(new Spacer(1));
if (config.refreshStatus) {
this.refreshStatusText = new Text(theme.fg("muted", ` ${config.refreshStatus}`), 0, 0);
this.addChild(this.refreshStatusText);
}
this.footerText = new Text(this.getFooterText(), 0, 0);
this.addChild(this.footerText);
@ -151,6 +157,28 @@ export class ScopedModelsSelectorComponent extends Container implements Focusabl
this.updateList();
}
updateModels(models: readonly Model<any>[], enabledModelIds?: string[] | null): void {
const selectedId = this.filteredItems[this.selectedIndex]?.fullId;
if (enabledModelIds !== undefined) this.enabledIds = enabledModelIds === null ? null : [...enabledModelIds];
this.modelsById.clear();
this.allIds = [];
for (const model of models) {
const fullId = `${model.provider}/${model.id}`;
this.modelsById.set(fullId, model);
this.allIds.push(fullId);
}
this.refresh();
const refreshedIndex = selectedId ? this.filteredItems.findIndex((item) => item.fullId === selectedId) : -1;
if (refreshedIndex >= 0) {
this.selectedIndex = refreshedIndex;
this.updateList();
}
}
setRefreshStatus(message: string, kind: "muted" | "success" | "warning"): void {
this.refreshStatusText?.setText(theme.fg(kind, ` ${message}`));
}
private buildItems(): ModelItem[] {
return getSortedIds(this.enabledIds, this.allIds).map((id) => ({
fullId: id,

View file

@ -85,9 +85,9 @@ import { createCompactionSummaryMessage } from "../../core/messages.ts";
import {
defaultModelPerProvider,
findExactModelReferenceMatch,
resolveModelScope,
resolveModelScopeWithDiagnostics,
resolveModelScopeFromModels,
} from "../../core/model-resolver.ts";
import { CredentialSynchronizationError } from "../../core/model-runtime.ts";
import { DefaultPackageManager } from "../../core/package-manager.ts";
import type { ResourceDiagnostic } from "../../core/resource-loader.ts";
import { formatMissingSessionCwdPrompt, MissingSessionCwdError } from "../../core/session-cwd.ts";
@ -388,6 +388,8 @@ export class InteractiveMode {
private autocompleteProviderWrappers: AutocompleteProviderFactory[] = [];
private fdPath: string | undefined;
private editorContainer: Container;
private activeSelectorToken?: object;
private activeSelectorDispose?: () => void;
private footer: FooterComponent;
private footerContainer: Container;
private footerDataProvider: FooterDataProvider;
@ -622,12 +624,11 @@ export class InteractiveMode {
const modelCommand = slashCommands.find((command) => command.name === "model");
if (modelCommand) {
modelCommand.getArgumentCompletions = async (prefix: string): Promise<AutocompleteItem[] | null> => {
// Get available models (scoped or from registry)
modelCommand.getArgumentCompletions = (prefix: string): AutocompleteItem[] | null => {
const models =
this.session.scopedModels.length > 0
? this.session.scopedModels.map((s) => s.model)
: await this.session.modelRuntime.getAvailable();
: this.session.modelRuntime.getAvailableSnapshot();
if (models.length === 0) return null;
@ -987,10 +988,13 @@ export class InteractiveMode {
await this.init();
if (!process.env.PI_OFFLINE) {
const controller = new AbortController();
const timeout = setTimeout(() => controller.abort(), 15_000);
void this.session.modelRuntime
.refresh()
.refresh({ signal: controller.signal })
.then(() => this.updateAvailableProviderCount())
.catch(() => {});
.catch(() => {})
.finally(() => clearTimeout(timeout));
}
// Start version check asynchronously
@ -2405,6 +2409,7 @@ export class InteractiveMode {
{ tui: this.ui, timeout: opts?.timeout, onToggleToolsExpanded: () => this.toggleToolOutputExpansion() },
);
this.disposeActiveSelector();
this.editorContainer.clear();
this.editorContainer.addChild(this.extensionSelector);
this.ui.setFocus(this.extensionSelector);
@ -2480,6 +2485,7 @@ export class InteractiveMode {
{ tui: this.ui, timeout: opts?.timeout },
);
this.disposeActiveSelector();
this.editorContainer.clear();
this.editorContainer.addChild(this.extensionInput);
this.ui.setFocus(this.extensionInput);
@ -2521,6 +2527,7 @@ export class InteractiveMode {
this.settingsManager.getExternalEditorCommand(),
);
this.disposeActiveSelector();
this.editorContainer.clear();
this.editorContainer.addChild(this.extensionEditor);
this.ui.setFocus(this.extensionEditor);
@ -2549,6 +2556,7 @@ export class InteractiveMode {
// Save text from current editor before switching
const currentText = this.editor.getText();
this.disposeActiveSelector();
this.editorContainer.clear();
if (factory) {
@ -2690,6 +2698,7 @@ export class InteractiveMode {
// Expose handle to caller for visibility control
options?.onHandle?.(handle);
} else {
this.disposeActiveSelector();
this.editorContainer.clear();
this.editorContainer.addChild(component);
this.ui.setFocus(component);
@ -4292,20 +4301,39 @@ export class InteractiveMode {
// Selectors
// =========================================================================
private disposeActiveSelector(): void {
const dispose = this.activeSelectorDispose;
this.activeSelectorToken = undefined;
this.activeSelectorDispose = undefined;
dispose?.();
}
/**
* Shows a selector component in place of the editor.
* @param create Factory that receives a `done` callback and returns the component and focus target
*/
private showSelector(create: (done: () => void) => { component: Component; focus: Component }): void {
private showSelector(
create: (done: () => void) => { component: Component; focus: Component; dispose?: () => void },
): void {
const token = {};
let dispose: (() => void) | undefined;
const done = () => {
dispose?.();
if (this.activeSelectorToken !== token) return;
this.activeSelectorToken = undefined;
this.activeSelectorDispose = undefined;
this.editorContainer.clear();
this.editorContainer.addChild(this.editor);
this.ui.setFocus(this.editor);
};
const { component, focus } = create(done);
const created = create(done);
dispose = created.dispose;
this.disposeActiveSelector();
this.activeSelectorToken = token;
this.activeSelectorDispose = dispose;
this.editorContainer.clear();
this.editorContainer.addChild(component);
this.ui.setFocus(focus);
this.editorContainer.addChild(created.component);
this.ui.setFocus(created.focus);
this.ui.requestRender();
}
@ -4537,21 +4565,37 @@ export class InteractiveMode {
}
private async findExactModelMatch(searchTerm: string): Promise<Model<any> | undefined> {
const models = await this.getModelCandidates();
return findExactModelReferenceMatch(searchTerm, models);
}
private async getModelCandidates(): Promise<Model<any>[]> {
if (this.session.scopedModels.length > 0) {
return this.session.scopedModels.map((scoped) => scoped.model);
}
const cachedModels =
this.session.scopedModels.length > 0
? this.session.scopedModels.map((scoped) => scoped.model)
: [...this.session.modelRuntime.getAvailableSnapshot()];
const cachedMatch = findExactModelReferenceMatch(searchTerm, cachedModels);
if (cachedMatch || this.session.scopedModels.length > 0) return cachedMatch;
this.showStatus("Refreshing model catalogs…");
const controller = new AbortController();
let timedOut = false;
const timeout = setTimeout(() => {
timedOut = true;
controller.abort();
}, 15_000);
try {
await this.session.modelRuntime.refresh();
return [...(await this.session.modelRuntime.getAvailable())];
} catch {
return [];
const result = await this.session.modelRuntime.refresh({ signal: controller.signal });
if (result.aborted && timedOut) {
this.showWarning("Model refresh timed out; searching cached models.");
} else if (result.errors.size > 0) {
this.showWarning(`Could not refresh ${[...result.errors.keys()].join(", ")}; searching cached models.`);
}
} catch (error) {
this.showWarning(
timedOut
? "Model refresh timed out; searching cached models."
: `Could not refresh model catalogs: ${error instanceof Error ? error.message : String(error)}`,
);
} finally {
clearTimeout(timeout);
}
return findExactModelReferenceMatch(searchTerm, [...this.session.modelRuntime.getAvailableSnapshot()]);
}
/** Update the footer's available provider count from the current snapshot without refreshing catalogs. */
@ -4673,86 +4717,75 @@ export class InteractiveMode {
},
initialSearchInput,
);
return { component: selector, focus: selector };
return { component: selector, focus: selector, dispose: () => selector.dispose() };
});
}
private async showModelsSelector(): Promise<void> {
// Get all available models
await this.session.modelRuntime.refresh();
const allModels = [...(await this.session.modelRuntime.getAvailable())];
const allModelIds = new Set(allModels.map((model) => `${model.provider}/${model.id}`));
private showModelsSelector(): void {
let availableModels = [...this.session.modelRuntime.getAvailableSnapshot()];
let availableModelIds = new Set(availableModels.map((model) => `${model.provider}/${model.id}`));
const configuredPatterns = this.settingsManager.getEnabledModels();
const sessionScopedModels = this.session.scopedModels;
const configuredEnabledIds = (models: readonly Model<any>[]): string[] | null => {
if (!configuredPatterns?.length) return null;
const resolved = resolveModelScopeFromModels(configuredPatterns, models);
const ids = resolved.scopedModels.map((scoped) => `${scoped.model.provider}/${scoped.model.id}`);
for (const diagnostic of resolved.diagnostics) {
if (diagnostic.code === "no-match" && !ids.includes(diagnostic.pattern)) ids.push(diagnostic.pattern);
}
return ids;
};
if (allModels.length === 0 && !configuredPatterns?.length && sessionScopedModels.length === 0) {
this.showStatus("No models available");
return;
}
let currentEnabledIds =
sessionScopedModels.length > 0
? sessionScopedModels.map((scoped) => `${scoped.model.provider}/${scoped.model.id}`)
: configuredEnabledIds(availableModels);
let selectionChanged = false;
const configuredScope = configuredPatterns?.length
? await resolveModelScopeWithDiagnostics(configuredPatterns, this.session.modelRuntime)
: undefined;
// Check if session has scoped models (from previous session-only changes or CLI --models)
const hasSessionScope = sessionScopedModels.length > 0;
// Build enabled model IDs from session state or settings
let currentEnabledIds: string[] | null = null;
if (hasSessionScope) {
// Use current session's scoped models
currentEnabledIds = sessionScopedModels.map((scoped) => `${scoped.model.provider}/${scoped.model.id}`);
} else if (configuredScope) {
currentEnabledIds = configuredScope.scopedModels.map(
(scoped) => `${scoped.model.provider}/${scoped.model.id}`,
);
}
for (const diagnostic of configuredScope?.diagnostics ?? []) {
if (diagnostic.code !== "no-match") continue;
currentEnabledIds ??= [];
if (!currentEnabledIds.includes(diagnostic.pattern)) currentEnabledIds.push(diagnostic.pattern);
}
// Helper to update session's scoped models (session-only, no persist)
const updateSessionModels = async (enabledIds: string[] | null) => {
const updateSessionModels = (enabledIds: string[] | null): void => {
currentEnabledIds = enabledIds === null ? null : [...enabledIds];
const hasEnabledAvailableModel = enabledIds?.some((id) => allModelIds.has(id)) ?? false;
const hasEnabledAvailableModel = enabledIds?.some((id) => availableModelIds.has(id)) ?? false;
const allAvailableModelsEnabled =
enabledIds !== null && [...allModelIds].every((id) => enabledIds.includes(id));
enabledIds !== null && [...availableModelIds].every((id) => enabledIds.includes(id));
if (enabledIds && hasEnabledAvailableModel && !allAvailableModelsEnabled) {
const newScopedModels = await resolveModelScope(enabledIds, this.session.modelRuntime);
const newScopedModels = resolveModelScopeFromModels(enabledIds, availableModels).scopedModels;
this.session.setScopedModels(
newScopedModels.map((sm) => ({
model: sm.model,
thinkingLevel: sm.thinkingLevel,
newScopedModels.map((scoped) => ({
model: scoped.model,
thinkingLevel: scoped.thinkingLevel,
})),
);
} else {
// All enabled or none enabled = no filter
this.session.setScopedModels([]);
}
await this.updateAvailableProviderCount();
this.updateAvailableProviderCount();
this.ui.requestRender();
};
this.showSelector((done) => {
let disposed = false;
let timedOut = false;
const controller = new AbortController();
const timeout = setTimeout(() => {
timedOut = true;
controller.abort();
}, 15_000);
const selector = new ScopedModelsSelectorComponent(
{
allModels,
allModels: availableModels,
enabledModelIds: currentEnabledIds,
refreshStatus: "Refreshing model catalogs…",
},
{
onChange: async (enabledIds) => {
await updateSessionModels(enabledIds);
onChange: (enabledIds) => {
selectionChanged = true;
updateSessionModels(enabledIds);
},
onPersist: (enabledIds) => {
// Persist to settings
const allEnabled =
enabledIds !== null &&
enabledIds.length === allModels.length &&
enabledIds.every((id) => allModelIds.has(id));
enabledIds.length === availableModels.length &&
enabledIds.every((id) => availableModelIds.has(id));
const newPatterns = enabledIds === null || allEnabled ? undefined : enabledIds;
this.settingsManager.setEnabledModels(newPatterns ? [...newPatterns] : undefined);
this.showStatus("Model selection saved to settings");
@ -4763,7 +4796,51 @@ export class InteractiveMode {
},
},
);
return { component: selector, focus: selector };
void this.session.modelRuntime
.refresh({ signal: controller.signal })
.then((result) => {
if (disposed) return;
availableModels = [...this.session.modelRuntime.getAvailableSnapshot()];
availableModelIds = new Set(availableModels.map((model) => `${model.provider}/${model.id}`));
if (!selectionChanged && sessionScopedModels.length === 0) {
currentEnabledIds = configuredEnabledIds(availableModels);
selector.updateModels(availableModels, currentEnabledIds);
} else {
selector.updateModels(availableModels);
}
if (currentEnabledIds !== null) updateSessionModels(currentEnabledIds);
if (result.aborted && timedOut) {
selector.setRefreshStatus("Model refresh timed out; showing cached models.", "warning");
} else if (result.errors.size > 0) {
selector.setRefreshStatus(
`Could not refresh ${[...result.errors.keys()].join(", ")}; showing cached models.`,
"warning",
);
} else {
selector.setRefreshStatus("Model catalogs refreshed.", "success");
}
this.ui.requestRender();
})
.catch((error: unknown) => {
if (disposed) return;
selector.setRefreshStatus(
timedOut
? "Model refresh timed out; showing cached models."
: `Could not refresh model catalogs: ${error instanceof Error ? error.message : String(error)}`,
"warning",
);
this.ui.requestRender();
})
.finally(() => clearTimeout(timeout));
return {
component: selector,
focus: selector,
dispose: () => {
disposed = true;
clearTimeout(timeout);
controller.abort();
},
};
});
}
@ -5075,7 +5152,7 @@ export class InteractiveMode {
}
private async getLogoutProviderOptions(): Promise<AuthSelectorProvider[]> {
return (await this.session.modelRuntime.listCredentials())
return (await this.session.modelRuntime.listCredentials({ signal: AbortSignal.timeout(15_000) }))
.map(({ providerId, type }) => ({
id: providerId,
name: this.session.modelRuntime.getProvider(providerId)?.name ?? providerId,
@ -5099,7 +5176,6 @@ export class InteractiveMode {
}
private async handleLoginCommand(providerRef?: string): Promise<void> {
await this.session.modelRuntime.getAvailable();
if (!providerRef) {
this.showLoginAuthTypeSelector();
return;
@ -5239,7 +5315,13 @@ export class InteractiveMode {
return;
}
const providerOptions = await this.getLogoutProviderOptions();
let providerOptions: AuthSelectorProvider[];
try {
providerOptions = await this.getLogoutProviderOptions();
} catch (error) {
this.showError(`Could not read stored credentials: ${error instanceof Error ? error.message : String(error)}`);
return;
}
if (providerOptions.length === 0) {
this.showStatus(
"No stored credentials to remove. /logout only removes credentials saved by /login; environment variables and models.json config are unchanged.",
@ -5260,7 +5342,9 @@ export class InteractiveMode {
}
try {
await this.session.modelRuntime.logout(providerOption.id);
await this.session.modelRuntime.logout(providerOption.id, {
signal: AbortSignal.timeout(15_000),
});
await this.updateAvailableProviderCount();
const message =
providerOption.authType === "oauth"
@ -5268,7 +5352,12 @@ export class InteractiveMode {
: `Removed stored API key for ${providerOption.name}. Environment variables and models.json config are unchanged.`;
this.showStatus(message);
} catch (error: unknown) {
this.showError(`Logout failed: ${error instanceof Error ? error.message : String(error)}`);
const message = error instanceof Error ? error.message : String(error);
this.showError(
error instanceof CredentialSynchronizationError
? `Credentials removed for ${providerOption.name}, but local model state could not be synchronized: ${message}`
: `Logout failed: ${message}`,
);
}
},
() => {
@ -5286,14 +5375,12 @@ export class InteractiveMode {
authType: "oauth" | "api_key",
previousModel: Model<any> | undefined,
): Promise<void> {
await this.session.modelRuntime.getAvailable();
const actionLabel = authType === "oauth" ? `Logged in to ${providerName}` : `Saved API key for ${providerName}`;
let selectedModel: Model<any> | undefined;
let selectionError: string | undefined;
if (isUnknownModel(previousModel)) {
const availableModels = await this.session.modelRuntime.getAvailable();
const availableModels = this.session.modelRuntime.getAvailableSnapshot();
const providerModels = availableModels.filter((model) => model.provider === providerId);
if (!hasDefaultModelProvider(providerId)) {
selectionError = `${actionLabel}, but no default model is configured for provider "${providerId}". Use /model to select a model.`;
@ -5331,6 +5418,27 @@ export class InteractiveMode {
void this.maybeWarnAboutAnthropicSubscriptionAuth();
}
}
const controller = new AbortController();
const timeout = setTimeout(() => controller.abort(), 15_000);
void this.session.modelRuntime
.refresh({ providers: [providerId], signal: controller.signal })
.then((result) => {
if (result.aborted) {
this.showWarning(`${actionLabel}, but its model catalog refresh timed out; using cached models.`);
} else if (result.errors.size > 0) {
this.showWarning(`${actionLabel}, but its model catalog could not be refreshed; using cached models.`);
}
this.updateAvailableProviderCount();
this.footer.invalidate();
this.ui.requestRender();
})
.catch((error: unknown) => {
this.showWarning(
`${actionLabel}, but its model catalog could not be refreshed: ${error instanceof Error ? error.message : String(error)}`,
);
})
.finally(() => clearTimeout(timeout));
}
private showAmbientAuthDialog(providerOption: AuthSelectorProvider): void {
@ -5395,7 +5503,11 @@ export class InteractiveMode {
} catch (error: unknown) {
restoreEditor();
const errorMsg = error instanceof Error ? error.message : String(error);
if (errorMsg !== "Login cancelled") {
if (error instanceof CredentialSynchronizationError) {
this.showError(
`Saved API key for ${providerName}, but local model state could not be synchronized: ${errorMsg}`,
);
} else if (errorMsg !== "Login cancelled") {
this.showError(`Failed to save API key for ${providerName}: ${errorMsg}`);
}
}
@ -5505,7 +5617,11 @@ export class InteractiveMode {
} catch (error: unknown) {
restoreEditor();
const errorMsg = error instanceof Error ? error.message : String(error);
if (errorMsg !== "Login cancelled") {
if (error instanceof CredentialSynchronizationError) {
this.showError(
`Logged in to ${providerName}, but local model state could not be synchronized: ${errorMsg}`,
);
} else if (errorMsg !== "Login cancelled") {
this.showError(`Failed to login to ${providerName}: ${errorMsg}`);
}
}
@ -6216,6 +6332,7 @@ export class InteractiveMode {
}
stop(): void {
this.disposeActiveSelector();
if (this.settingsManager.getShowTerminalProgress()) {
this.ui.terminal.setProgress(false);
}

View file

@ -466,7 +466,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
// =================================================================
case "set_model": {
const models = await session.modelRuntime.getAvailable();
const models = session.modelRuntime.getAvailableSnapshot();
const model = models.find((m) => m.provider === command.provider && m.id === command.modelId);
if (!model) {
return error(id, "set_model", `Model not found: ${command.provider}/${command.modelId}`);
@ -484,7 +484,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
}
case "get_available_models": {
const models = await session.modelRuntime.getAvailable();
const models = session.modelRuntime.getAvailableSnapshot();
return success(id, "get_available_models", { models });
}

View file

@ -395,14 +395,15 @@ function updateTargetIncludesExtensions(target: UpdateTarget): boolean {
}
async function refreshModelCatalogs(agentDir: string): Promise<void> {
const modelRuntime = await ModelRuntime.create({
authPath: join(agentDir, "auth.json"),
modelsPath: join(agentDir, "models.json"),
allowModelNetwork: false,
});
const controller = new AbortController();
const timeout = setTimeout(() => controller.abort(), 15_000);
try {
const modelRuntime = await ModelRuntime.create({
authPath: join(agentDir, "auth.json"),
modelsPath: join(agentDir, "models.json"),
allowModelNetwork: false,
signal: controller.signal,
});
const result = await modelRuntime.refresh({
allowNetwork: true,
force: true,

View file

@ -0,0 +1,48 @@
function abortReason(signal: AbortSignal): unknown {
if (signal.reason !== undefined) return signal.reason;
const error = new Error("The operation was aborted");
error.name = "AbortError";
return error;
}
/** Normalize an optional public signal without imposing a deadline. */
export function operationSignal(signal?: AbortSignal): AbortSignal {
return signal ?? new AbortController().signal;
}
/** Stop waiting on abort while observing the abandoned operation through settlement. */
export function raceWithAbortSignal<T>(operation: Promise<T>, signal: AbortSignal | undefined): Promise<T> {
if (!signal) return operation;
if (signal.aborted) {
void operation.catch(() => {});
return Promise.reject(abortReason(signal));
}
return new Promise<T>((resolve, reject) => {
let settled = false;
const cleanup = () => signal.removeEventListener("abort", onAbort);
const onAbort = () => {
if (settled) return;
settled = true;
cleanup();
reject(abortReason(signal));
};
signal.addEventListener("abort", onAbort, { once: true });
void operation.then(
(value) => {
if (settled) return;
settled = true;
cleanup();
resolve(value);
},
(error: unknown) => {
if (settled) return;
settled = true;
cleanup();
reject(error);
},
);
if (signal.aborted) onAbort();
});
}

View file

@ -206,6 +206,215 @@ describe("AuthStorage", () => {
});
});
test("pre-aborted file operations do not create the backing file or run the mutation", async () => {
const backend = new FileAuthStorageBackend(authJsonPath);
const controller = new AbortController();
controller.abort();
const update = vi.fn(async () => ({ result: undefined, next: JSON.stringify({}) }));
await expect(backend.withLockAsync(update, { signal: controller.signal })).rejects.toMatchObject({
name: "AbortError",
});
expect(update).not.toHaveBeenCalled();
expect(existsSync(authJsonPath)).toBe(false);
});
test("aborts while waiting for a held file lock without running the mutation later", async () => {
writeAuthJson({ anthropic: { type: "api_key", key: "stored" } });
const release = await lockfile.lock(authJsonPath, { realpath: false });
const backend = new FileAuthStorageBackend(authJsonPath);
const controller = new AbortController();
const update = vi.fn(async () => ({ result: undefined, next: JSON.stringify({}) }));
const pending = backend.withLockAsync(update, { signal: controller.signal });
await new Promise((resolve) => setTimeout(resolve, 10));
controller.abort();
await expect(pending).rejects.toMatchObject({ name: "AbortError" });
expect(update).not.toHaveBeenCalled();
await release();
await new Promise((resolve) => setTimeout(resolve, 150));
expect(update).not.toHaveBeenCalled();
expect(JSON.parse(readFileSync(authJsonPath, "utf8"))).toEqual({
anthropic: { type: "api_key", key: "stored" },
});
});
test("releases a file lock acquired concurrently with cancellation before mutation", async () => {
writeAuthJson({ anthropic: { type: "api_key", key: "stored" } });
const backend = new FileAuthStorageBackend(authJsonPath);
const controller = new AbortController();
const release = vi.fn(async () => {});
vi.spyOn(lockfile, "lock").mockImplementation(async () => {
controller.abort();
return release;
});
const update = vi.fn(async () => ({ result: undefined, next: JSON.stringify({}) }));
await expect(backend.withLockAsync(update, { signal: controller.signal })).rejects.toMatchObject({
name: "AbortError",
});
await new Promise((resolve) => setTimeout(resolve, 0));
expect(update).not.toHaveBeenCalled();
expect(release).toHaveBeenCalledTimes(1);
});
test("holds the file lock until a cancelled active callback settles without committing it", async () => {
writeAuthJson({ anthropic: { type: "api_key", key: "stored" } });
const backend = new FileAuthStorageBackend(authJsonPath);
const controller = new AbortController();
let markStarted: (() => void) | undefined;
let finish: (() => void) | undefined;
const started = new Promise<void>((resolve) => {
markStarted = resolve;
});
const blocked = new Promise<void>((resolve) => {
finish = resolve;
});
const pending = backend.withLockAsync(
async () => {
markStarted?.();
await blocked;
return { result: undefined, next: JSON.stringify({ openai: { type: "api_key", key: "cancelled" } }) };
},
{ signal: controller.signal },
);
await started;
controller.abort();
const competingMutation = vi.fn(async () => ({
result: undefined,
next: JSON.stringify({ google: { type: "api_key", key: "committed" } }),
}));
const competing = backend.withLockAsync(competingMutation);
await new Promise((resolve) => setTimeout(resolve, 20));
expect(competingMutation).not.toHaveBeenCalled();
finish?.();
await expect(pending).rejects.toMatchObject({ name: "AbortError" });
await competing;
expect(competingMutation).toHaveBeenCalledTimes(1);
expect(JSON.parse(readFileSync(authJsonPath, "utf8"))).toEqual({
google: { type: "api_key", key: "committed" },
});
});
test("cancels a signalled credential read waiting for a held file lock", async () => {
writeAuthJson({ anthropic: { type: "api_key", key: "old" } });
const storage = AuthStorage.create(authJsonPath);
writeAuthJson({ anthropic: { type: "api_key", key: "new-value" } });
const release = await lockfile.lock(authJsonPath, { realpath: false });
const lockSpy = vi.spyOn(lockfile, "lock");
const controller = new AbortController();
const pending = storage.read("anthropic", { signal: controller.signal });
await new Promise((resolve) => setTimeout(resolve, 10));
controller.abort();
await expect(pending).rejects.toMatchObject({ name: "AbortError" });
await release();
await new Promise((resolve) => setTimeout(resolve, 150));
expect(lockSpy).toHaveBeenCalledTimes(1);
await expect(storage.read("anthropic")).resolves.toEqual({ type: "api_key", key: "new-value" });
});
test("serializes in-memory mutations across providers", async () => {
const storage = AuthStorage.inMemory();
let markStarted: (() => void) | undefined;
let finish: (() => void) | undefined;
const started = new Promise<void>((resolve) => {
markStarted = resolve;
});
const blocked = new Promise<void>((resolve) => {
finish = resolve;
});
const first = storage.modify("anthropic", async () => {
markStarted?.();
await blocked;
return { type: "api_key", key: "anthropic-key" };
});
await started;
const secondMutation = vi.fn(async () => ({ type: "api_key" as const, key: "openai-key" }));
const second = storage.modify("openai", secondMutation);
await new Promise((resolve) => setTimeout(resolve, 0));
expect(secondMutation).not.toHaveBeenCalled();
finish?.();
await Promise.all([first, second]);
expect(await storage.read("anthropic")).toEqual({ type: "api_key", key: "anthropic-key" });
expect(await storage.read("openai")).toEqual({ type: "api_key", key: "openai-key" });
});
test("cancels a queued in-memory mutation without running it later", async () => {
const storage = AuthStorage.inMemory();
let markStarted: (() => void) | undefined;
let finish: (() => void) | undefined;
const started = new Promise<void>((resolve) => {
markStarted = resolve;
});
const blocked = new Promise<void>((resolve) => {
finish = resolve;
});
const first = storage.modify("anthropic", async () => {
markStarted?.();
await blocked;
return { type: "api_key", key: "anthropic-key" };
});
await started;
const controller = new AbortController();
const secondMutation = vi.fn(async () => ({ type: "api_key" as const, key: "openai-key" }));
const second = storage.modify("openai", secondMutation, { signal: controller.signal });
controller.abort();
await expect(second).rejects.toMatchObject({ name: "AbortError" });
expect(secondMutation).not.toHaveBeenCalled();
finish?.();
await first;
await new Promise((resolve) => setTimeout(resolve, 0));
expect(secondMutation).not.toHaveBeenCalled();
expect(await storage.read("openai")).toBeUndefined();
});
test("preserves the stored credential after cancelling an active refresh mutation", async () => {
const previous = {
type: "oauth" as const,
access: "expired",
refresh: "refresh-token",
expires: 0,
};
const storage = AuthStorage.inMemory({ oauth: previous });
const controller = new AbortController();
let markStarted: (() => void) | undefined;
let finish: (() => void) | undefined;
const started = new Promise<void>((resolve) => {
markStarted = resolve;
});
const blocked = new Promise<void>((resolve) => {
finish = resolve;
});
const pending = storage.modify(
"oauth",
async () => {
markStarted?.();
await blocked;
return { ...previous, access: "refreshed", expires: Date.now() + 60_000 };
},
{ signal: controller.signal },
);
await started;
controller.abort();
await expect(pending).rejects.toMatchObject({ name: "AbortError" });
const competingMutation = vi.fn(async () => ({ type: "api_key" as const, key: "other" }));
const competing = storage.modify("other", competingMutation);
await new Promise((resolve) => setTimeout(resolve, 0));
expect(competingMutation).not.toHaveBeenCalled();
finish?.();
await competing;
expect(competingMutation).toHaveBeenCalledTimes(1);
expect(await storage.read("oauth")).toEqual(previous);
});
test("translates a credential-store refresh failure and allows a later retry", async () => {
const providerId = "oauth-provider";
const base = AuthStorage.inMemory({

View file

@ -871,7 +871,7 @@ describe("ExtensionRunner", () => {
expect(errors).toEqual([
'/tmp/broken-extension.ts: Provider broken-provider: "api" is required when registering streamSimple.',
]);
await expect(modelRegistry.refresh()).resolves.toBeUndefined();
await expect(modelRegistry.refresh()).resolves.toMatchObject({ aborted: false });
});
it("pre-bind unregister removes all queued registrations for a provider", () => {

View file

@ -1,7 +1,7 @@
import { once } from "node:events";
import { createServer, type RequestListener, type Server, type ServerResponse } from "node:http";
import type { AddressInfo } from "node:net";
import type { AuthContext, AuthPrompt, ModelsStoreEntry } from "@earendil-works/pi-ai";
import type { AuthContext, AuthPrompt, ModelsPublication, ModelsStoreEntry } from "@earendil-works/pi-ai";
import { afterEach, describe, expect, it } from "vitest";
import { createEventBus } from "../src/core/event-bus.ts";
import { createExtensionRuntime, loadExtensionFromFactory } from "../src/core/extensions/loader.ts";
@ -88,15 +88,6 @@ describe("llama.cpp extension", () => {
it("persists and restores loaded models for cache-only startup refreshes", async () => {
let cachedEntry: ModelsStoreEntry | undefined;
const store = {
read: async () => cachedEntry,
write: async (entry: ModelsStoreEntry) => {
cachedEntry = structuredClone(entry);
},
delete: async () => {
cachedEntry = undefined;
},
};
const { url } = await listen((request, response) => {
if (request.url === "/models") {
json(response, {
@ -110,11 +101,19 @@ describe("llama.cpp extension", () => {
response.writeHead(404).end();
});
const publish = async (publication: ModelsPublication): Promise<boolean> => {
if (publication.persist === null) cachedEntry = undefined;
else if (publication.persist !== undefined) cachedEntry = structuredClone(publication.persist);
publication.update?.();
return true;
};
const first = createLlamaProvider();
await first.provider.refreshModels?.({
credential: { type: "api_key", key: "local", env: { LLAMA_BASE_URL: url } },
store,
stored: cachedEntry,
publish,
allowNetwork: true,
signal: new AbortController().signal,
});
expect(first.provider.getModels().map((model) => model.id)).toEqual(["loaded"]);
expect(cachedEntry?.models.map((model) => model.id)).toEqual(["loaded"]);
@ -122,8 +121,10 @@ describe("llama.cpp extension", () => {
const second = createLlamaProvider();
await second.provider.refreshModels?.({
credential: { type: "api_key", key: "local", env: { LLAMA_BASE_URL: url } },
store,
stored: cachedEntry,
publish,
allowNetwork: false,
signal: new AbortController().signal,
});
expect(second.provider.getModels()).toEqual([
expect.objectContaining({ id: "loaded", baseUrl: `${url}/v1`, contextWindow: 32768 }),
@ -137,8 +138,9 @@ describe("llama.cpp extension", () => {
env: async () => undefined,
fileExists: async () => false,
};
expect(await auth.check?.({ ctx: emptyContext })).toBeUndefined();
expect(await auth.resolve({ ctx: emptyContext })).toBeUndefined();
const signal = new AbortController().signal;
expect(await auth.check?.({ ctx: emptyContext, signal })).toBeUndefined();
expect(await auth.resolve({ ctx: emptyContext, signal })).toBeUndefined();
const { url } = await listen((request, response) => {
expect(request.headers.authorization).toBe("Bearer secret");
@ -146,6 +148,7 @@ describe("llama.cpp extension", () => {
});
const answers = [url, "secret"];
const credential = await auth.login!({
signal,
prompt: async (_prompt: AuthPrompt) => answers.shift()!,
notify: () => {},
});
@ -154,7 +157,7 @@ describe("llama.cpp extension", () => {
key: "secret",
env: { LLAMA_BASE_URL: url },
});
expect(await auth.resolve({ ctx: emptyContext, credential })).toEqual({
expect(await auth.resolve({ ctx: emptyContext, credential, signal })).toEqual({
auth: { apiKey: "secret", baseUrl: `${url}/v1` },
env: { LLAMA_BASE_URL: url },
source: "stored credential",

View file

@ -1144,7 +1144,7 @@ describe("ModelRegistry", () => {
}),
).toThrow('Provider broken-provider: "api" is required when registering streamSimple.');
await expect(registry.refresh()).resolves.toBeUndefined();
await expect(registry.refresh()).resolves.toMatchObject({ aborted: false });
});
test("failed registerProvider does not remove existing provider models", async () => {
@ -1188,7 +1188,7 @@ describe("ModelRegistry", () => {
).toThrow('Provider demo-provider, model broken-model: no "api" specified.');
expect(registry.find("demo-provider", "demo-model")).toBeDefined();
await expect(registry.refresh()).resolves.toBeUndefined();
await expect(registry.refresh()).resolves.toMatchObject({ aborted: false });
expect(registry.find("demo-provider", "demo-model")).toBeDefined();
});

View file

@ -743,7 +743,7 @@ describe("default model selection", () => {
};
const registry = {
getAvailable: async () => [aiGatewayModel],
getAvailableSnapshot: () => [aiGatewayModel],
} as unknown as Parameters<typeof findInitialModel>[0]["modelRuntime"];
const result = await findInitialModel({
@ -780,7 +780,7 @@ describe("default model selection", () => {
? savedDeepSeekModel
: undefined,
hasConfiguredAuth: (provider: string) => provider === "spark-two",
getAvailable: async () => [localDeepSeekModel],
getAvailableSnapshot: () => [localDeepSeekModel],
} as unknown as Parameters<typeof findInitialModel>[0]["modelRuntime"];
const result = await findInitialModel({

View file

@ -230,6 +230,38 @@ describe("ModelRuntime auth options", () => {
});
});
it("forwards cancellation to extension OAuth refresh", async () => {
const credentials = AuthStorage.inMemory({
"extension-oauth": {
type: "oauth",
access: "expired",
refresh: "refresh",
expires: 0,
},
});
const runtime = await ModelRuntime.create({ credentials, modelsPath: null });
let refreshSignal: AbortSignal | undefined;
runtime.registerProvider("extension-oauth", {
name: "Extension OAuth",
baseUrl: "https://example.test/v1",
api: "openai-completions",
oauth: {
name: "Extension subscription",
login: async () => ({ access: "access", refresh: "refresh", expires: Date.now() + 60_000 }),
refreshToken: async (credential, signal) => {
refreshSignal = signal;
return { ...credential, expires: Date.now() + 60_000 };
},
getApiKey: (credential) => credential.access,
},
models: [testModel("extension-model")],
});
const controller = new AbortController();
await runtime.getAuth("extension-oauth", { signal: controller.signal });
expect(refreshSignal).toBe(controller.signal);
});
it("does not fabricate an API key method for an extension OAuth-only provider", async () => {
const runtime = await ModelRuntime.create({ credentials: AuthStorage.inMemory(), modelsPath: null });
runtime.registerProvider("extension-oauth", {

View file

@ -0,0 +1,375 @@
import type { ApiKeyCredential, Credential, CredentialStore, Model, Provider } from "@earendil-works/pi-ai";
import { describe, expect, it, vi } from "vitest";
import { AuthStorage } from "../src/core/auth-storage.ts";
import { CredentialSynchronizationError, ModelRuntime } from "../src/core/model-runtime.ts";
function model(provider: string): Model<"openai-completions"> {
return {
id: "dynamic",
name: "Dynamic",
api: "openai-completions",
provider,
baseUrl: "https://example.test/v1",
reasoning: false,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 1000,
maxTokens: 100,
};
}
function provider(
id: string,
options: {
login?: () => Promise<ApiKeyCredential>;
refreshModels?: Provider["refreshModels"];
} = {},
): Provider<"openai-completions"> {
const providerModel = model(id);
return {
id,
name: id,
auth: {
apiKey: {
name: "API key",
login: async () => options.login?.() ?? { type: "api_key", key: `${id}-key` },
check: async ({ credential }) => (credential ? { type: "api_key", source: "stored" } : undefined),
resolve: async ({ credential }) =>
credential ? { auth: { apiKey: credential.key }, source: "stored" } : undefined,
},
},
getModels: () => [providerModel],
refreshModels: options.refreshModels,
stream: () => {
throw new Error("unused");
},
streamSimple: () => {
throw new Error("unused");
},
};
}
async function runtimeWithProvider(
registered: Provider,
credentials: AuthStorage = AuthStorage.inMemory(),
): Promise<ModelRuntime> {
const runtime = await ModelRuntime.create({ credentials, modelsPath: null, allowModelNetwork: false });
runtime.registerNativeProvider(registered);
await runtime.refresh({ allowNetwork: false, providers: [registered.id] });
return runtime;
}
describe("ModelRuntime credential synchronization", () => {
it("publishes locally consistent availability before login and logout resolve", async () => {
const credentials = AuthStorage.inMemory();
const runtime = await runtimeWithProvider(provider("dynamic"), credentials);
await runtime.login("dynamic", "api_key", { prompt: async () => "unused", notify: () => {} });
expect(runtime.hasConfiguredAuth("dynamic")).toBe(true);
expect(runtime.getAvailableSnapshot().map((entry) => entry.id)).toContain("dynamic");
expect(await credentials.read("dynamic")).toEqual({ type: "api_key", key: "dynamic-key" });
await runtime.logout("dynamic");
expect(runtime.hasConfiguredAuth("dynamic")).toBe(false);
expect(runtime.getAvailableSnapshot().some((entry) => entry.provider === "dynamic")).toBe(false);
expect(await credentials.read("dynamic")).toBeUndefined();
});
it("orders same-provider credential operations through local synchronization", async () => {
let markLoginStarted: (() => void) | undefined;
let finishLogin: (() => void) | undefined;
const loginStarted = new Promise<void>((resolve) => {
markLoginStarted = resolve;
});
const blockedLogin = new Promise<void>((resolve) => {
finishLogin = resolve;
});
const credentials = AuthStorage.inMemory();
const runtime = await runtimeWithProvider(
provider("ordered", {
login: async () => {
markLoginStarted?.();
await blockedLogin;
return { type: "api_key", key: "ordered-key" };
},
}),
credentials,
);
const login = runtime.login("ordered", "api_key", { prompt: async () => "unused", notify: () => {} });
await loginStarted;
const logout = runtime.logout("ordered");
await new Promise((resolve) => setTimeout(resolve, 0));
expect(await credentials.read("ordered")).toBeUndefined();
finishLogin?.();
await Promise.all([login, logout]);
expect(await credentials.read("ordered")).toBeUndefined();
expect(runtime.hasConfiguredAuth("ordered")).toBe(false);
});
it("allows different providers to run credential flows concurrently", async () => {
let firstStarted: (() => void) | undefined;
let secondStarted: (() => void) | undefined;
let finish: (() => void) | undefined;
const first = new Promise<void>((resolve) => {
firstStarted = resolve;
});
const second = new Promise<void>((resolve) => {
secondStarted = resolve;
});
const blocked = new Promise<void>((resolve) => {
finish = resolve;
});
const runtime = await ModelRuntime.create({ credentials: AuthStorage.inMemory(), modelsPath: null });
runtime.registerNativeProvider(
provider("one", {
login: async () => {
firstStarted?.();
await blocked;
return { type: "api_key", key: "one" };
},
}),
);
runtime.registerNativeProvider(
provider("two", {
login: async () => {
secondStarted?.();
await blocked;
return { type: "api_key", key: "two" };
},
}),
);
await runtime.refresh({ allowNetwork: false, providers: ["one", "two"] });
const one = runtime.login("one", "api_key", { prompt: async () => "unused", notify: () => {} });
const two = runtime.login("two", "api_key", { prompt: async () => "unused", notify: () => {} });
await Promise.all([first, second]);
finish?.();
await Promise.all([one, two]);
});
it("does not wait for unrelated provider availability during local synchronization", async () => {
let stallUnrelated = false;
const runtime = await ModelRuntime.create({ credentials: AuthStorage.inMemory(), modelsPath: null });
runtime.registerNativeProvider(provider("target"));
const unrelated = provider("unrelated");
if (unrelated.auth.apiKey) {
unrelated.auth.apiKey.check = async () => {
if (stallUnrelated) await new Promise<void>(() => {});
return undefined;
};
}
runtime.registerNativeProvider(unrelated);
await runtime.refresh({ allowNetwork: false, providers: ["target", "unrelated"] });
stallUnrelated = true;
await runtime.login("target", "api_key", { prompt: async () => "unused", notify: () => {} });
expect(runtime.hasConfiguredAuth("target")).toBe(true);
await expect(runtime.refresh({ allowNetwork: false, providers: ["target"] })).resolves.toMatchObject({
aborted: false,
});
});
it("reports cancellation that occurs during provider-scoped availability", async () => {
let blockAvailability = false;
let markStarted: (() => void) | undefined;
const started = new Promise<void>((resolve) => {
markStarted = resolve;
});
const registered = provider("cancelled-availability");
if (registered.auth.apiKey) {
registered.auth.apiKey.check = async ({ credential }) => {
if (blockAvailability) {
markStarted?.();
await new Promise<void>(() => {});
}
return credential ? { type: "api_key", source: "stored" } : undefined;
};
}
const runtime = await runtimeWithProvider(registered);
await runtime.setRuntimeApiKey(registered.id, "key");
blockAvailability = true;
const controller = new AbortController();
const refresh = runtime.refresh({
allowNetwork: false,
providers: [registered.id],
signal: controller.signal,
});
await started;
controller.abort();
await expect(refresh).resolves.toMatchObject({ aborted: true });
});
it("does not run network refresh inside the credential operation chain", async () => {
const networkRefresh = vi.fn(async () => new Promise<void>(() => {}));
const runtime = await runtimeWithProvider(
provider("local-only", {
refreshModels: async (context) => {
if (context.allowNetwork) await networkRefresh();
},
}),
);
await runtime.login("local-only", "api_key", { prompt: async () => "unused", notify: () => {} });
expect(networkRefresh).not.toHaveBeenCalled();
expect(runtime.hasConfiguredAuth("local-only")).toBe(true);
});
it("keeps provider-scoped refreshes from superseding unrelated providers", async () => {
let markStarted: (() => void) | undefined;
let finish: (() => void) | undefined;
let firstSignal: AbortSignal | undefined;
const started = new Promise<void>((resolve) => {
markStarted = resolve;
});
const blocked = new Promise<void>((resolve) => {
finish = resolve;
});
const runtime = await ModelRuntime.create({ credentials: AuthStorage.inMemory(), modelsPath: null });
runtime.registerNativeProvider(
provider("one", {
refreshModels: async (context) => {
if (!context.allowNetwork) return;
firstSignal = context.signal;
markStarted?.();
await blocked;
},
}),
);
runtime.registerNativeProvider(provider("two"));
await runtime.refresh({ allowNetwork: false, providers: ["one", "two"] });
await runtime.setRuntimeApiKey("one", "one-key");
await runtime.setRuntimeApiKey("two", "two-key");
const first = runtime.refresh({ allowNetwork: true, providers: ["one"] });
await started;
await runtime.refresh({ allowNetwork: true, providers: ["two"] });
expect(firstSignal?.aborted).toBe(false);
finish?.();
await first;
});
it("waits for a committed credential mutation to settle before reporting cancellation", async () => {
let stored: Credential | undefined;
let markCommitted: (() => void) | undefined;
let finishMutation: (() => void) | undefined;
const committed = new Promise<void>((resolve) => {
markCommitted = resolve;
});
const mutationFinished = new Promise<void>((resolve) => {
finishMutation = resolve;
});
const credentials: CredentialStore = {
read: async () => stored,
list: async () => (stored ? [{ providerId: "delayed-commit", type: stored.type }] : []),
modify: async (_providerId, update) => {
const next = await update(stored);
if (next) stored = next;
markCommitted?.();
await mutationFinished;
return stored;
},
delete: async () => {
stored = undefined;
},
};
const runtime = await ModelRuntime.create({ credentials, modelsPath: null });
runtime.registerNativeProvider(provider("delayed-commit"));
await runtime.refresh({ allowNetwork: false, providers: ["delayed-commit"] });
const controller = new AbortController();
let settled = false;
const login = runtime.login("delayed-commit", "api_key", {
signal: controller.signal,
prompt: async () => "unused",
notify: () => {},
});
const outcome = login.then(
() => {
settled = true;
return undefined;
},
(error: unknown) => {
settled = true;
return error;
},
);
await committed;
controller.abort();
await new Promise((resolve) => setTimeout(resolve, 0));
expect(settled).toBe(false);
finishMutation?.();
await expect(outcome).resolves.toMatchObject({
name: "CredentialSynchronizationError",
credential: { type: "api_key", key: "delayed-commit-key" },
});
expect(stored).toEqual({ type: "api_key", key: "delayed-commit-key" });
});
it("reports a typed error when cancellation interrupts post-commit synchronization", async () => {
let blockCacheRefresh = false;
let markCacheRefreshStarted: (() => void) | undefined;
const cacheRefreshStarted = new Promise<void>((resolve) => {
markCacheRefreshStarted = resolve;
});
const credentials = AuthStorage.inMemory();
const runtime = await runtimeWithProvider(
provider("cancelled-sync", {
refreshModels: async (context) => {
if (!context.allowNetwork && blockCacheRefresh) {
markCacheRefreshStarted?.();
await new Promise<void>(() => {});
}
},
}),
credentials,
);
blockCacheRefresh = true;
const controller = new AbortController();
const login = runtime.login("cancelled-sync", "api_key", {
signal: controller.signal,
prompt: async () => "unused",
notify: () => {},
});
await cacheRefreshStarted;
controller.abort();
await expect(login).rejects.toMatchObject({
name: "CredentialSynchronizationError",
providerId: "cancelled-sync",
operation: "login",
credential: { type: "api_key", key: "cancelled-sync-key" },
});
expect(await credentials.read("cancelled-sync")).toEqual({
type: "api_key",
key: "cancelled-sync-key",
});
});
it("reports committed credentials when local synchronization fails", async () => {
let failCacheRefresh = false;
const credentials = AuthStorage.inMemory();
const runtime = await runtimeWithProvider(
provider("broken-sync", {
refreshModels: async (context) => {
if (!context.allowNetwork && failCacheRefresh) throw new Error("cache restore failed");
},
}),
credentials,
);
failCacheRefresh = true;
const login = runtime.login("broken-sync", "api_key", { prompt: async () => "unused", notify: () => {} });
await expect(login).rejects.toMatchObject({
name: "CredentialSynchronizationError",
providerId: "broken-sync",
operation: "login",
credential: { type: "api_key", key: "broken-sync-key" },
});
await expect(login).rejects.toBeInstanceOf(CredentialSynchronizationError);
expect(await credentials.read("broken-sync")).toEqual({ type: "api_key", key: "broken-sync-key" });
});
});

View file

@ -1,7 +1,8 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { Model } from "@earendil-works/pi-ai";
import lockfile from "proper-lockfile";
import { afterEach, describe, expect, it } from "vitest";
import { FileModelsStore } from "../src/core/models-store.ts";
@ -48,4 +49,26 @@ describe("FileModelsStore", () => {
expect(await reloaded.read("one")).toBeUndefined();
expect((await reloaded.read("two"))?.models.map((entry) => entry.id)).toEqual(["m2"]);
});
it("cancels a catalog write waiting for a held file lock without writing later", async () => {
const dir = join(tmpdir(), `pi-models-store-${Date.now()}-${Math.random().toString(36).slice(2)}`);
tempDirs.push(dir);
mkdirSync(dir, { recursive: true });
const path = join(dir, "models-store.json");
writeFileSync(path, JSON.stringify({ one: { models: [model("one", "existing")] } }));
const store = new FileModelsStore(path);
const release = await lockfile.lock(path, { realpath: false });
const controller = new AbortController();
const pending = store.write("two", { models: [model("two", "cancelled")] }, { signal: controller.signal });
await new Promise((resolve) => setTimeout(resolve, 10));
controller.abort();
await expect(pending).rejects.toMatchObject({ name: "AbortError" });
await release();
await new Promise((resolve) => setTimeout(resolve, 150));
const stored = JSON.parse(readFileSync(path, "utf8")) as Record<string, unknown>;
expect(stored.one).toBeDefined();
expect(stored.two).toBeUndefined();
});
});

View file

@ -66,11 +66,12 @@ describe("Radius provider", () => {
});
it("fetches and stores the catalog for configured Radius auth", async () => {
vi.spyOn(globalThis, "fetch").mockResolvedValue(
new Response(JSON.stringify(radiusConfig("https://radius.example.com/v1")), {
status: 200,
headers: { "content-type": "application/json" },
}),
vi.spyOn(globalThis, "fetch").mockImplementation(
async () =>
new Response(JSON.stringify(radiusConfig("https://radius.example.com/v1")), {
status: 200,
headers: { "content-type": "application/json" },
}),
);
const modelsStore = new InMemoryModelsStore();
const credentials = AuthStorage.inMemory({
@ -90,7 +91,10 @@ describe("Radius provider", () => {
expect(runtime.getModel(RADIUS_PROVIDER_ID, "auto")).toBeDefined();
expect((await modelsStore.read(RADIUS_PROVIDER_ID))?.models).toHaveLength(1);
expect(vi.mocked(fetch).mock.calls[0]?.[1]?.headers).toMatchObject({ authorization: "Bearer access-token" });
const radiusRequest = vi
.mocked(fetch)
.mock.calls.find(([url]) => String(url) === "https://radius.pi.dev/v1/config");
expect(radiusRequest?.[1]?.headers).toMatchObject({ authorization: "Bearer access-token" });
});
it("does not refresh catalogs over the network by default", async () => {
@ -121,8 +125,8 @@ describe("Radius provider", () => {
});
it("supports custom Radius gateways from models.json", async () => {
vi.spyOn(globalThis, "fetch").mockResolvedValue(
new Response(JSON.stringify(radiusConfig("http://localhost:8788/v1")), { status: 200 }),
vi.spyOn(globalThis, "fetch").mockImplementation(
async () => new Response(JSON.stringify(radiusConfig("http://localhost:8788/v1")), { status: 200 }),
);
const modelsPath = join(tempDir, "models.json");
writeFileSync(

View file

@ -1,14 +1,18 @@
import {
createModels,
createProvider,
InMemoryModelsStore,
type Model,
type ModelsStoreEntry,
type ProviderModelsStore,
type ModelsPublication,
type Provider,
type RefreshModelsContext,
} from "@earendil-works/pi-ai";
import { afterEach, describe, expect, it, vi } from "vitest";
import { VERSION } from "../src/config.ts";
import { withRemoteCatalog } from "../src/core/remote-catalog-provider.ts";
const neverAbortedSignal = new AbortController().signal;
function model(id: string): Model<"openai-completions"> {
return {
id,
@ -44,12 +48,25 @@ function testProvider(localGeneratedAt?: number) {
);
}
function scopedStore(store: InMemoryModelsStore): ProviderModelsStore {
return {
read: () => store.read("test-provider"),
write: (entry: ModelsStoreEntry) => store.write("test-provider", entry),
delete: () => store.delete("test-provider"),
async function refreshProvider(
provider: Provider,
store: InMemoryModelsStore,
overrides: Partial<Pick<RefreshModelsContext, "allowNetwork" | "force" | "signal">> = {},
): Promise<void> {
const publish = async (publication: ModelsPublication): Promise<boolean> => {
if (publication.persist === null) await store.delete(provider.id);
else if (publication.persist !== undefined) await store.write(provider.id, publication.persist);
publication.update?.();
return true;
};
await provider.refreshModels?.({
credential: { type: "api_key" },
stored: await store.read(provider.id),
publish,
allowNetwork: overrides.allowNetwork ?? true,
force: overrides.force,
signal: overrides.signal ?? neverAbortedSignal,
});
}
afterEach(() => vi.restoreAllMocks());
@ -65,10 +82,9 @@ describe("remote catalog provider", () => {
);
const provider = testProvider();
const store = new InMemoryModelsStore();
const refresh = { credential: { type: "api_key" } as const, store: scopedStore(store), allowNetwork: true };
await provider.refreshModels?.(refresh);
await provider.refreshModels?.(refresh);
await provider.refreshModels?.({ ...refresh, force: true });
await refreshProvider(provider, store);
await refreshProvider(provider, store);
await refreshProvider(provider, store, { force: true });
expect(provider.getModels().map((entry) => entry.id)).toEqual(["static", "dynamic"]);
expect((await store.read(provider.id))?.models.map((entry) => entry.id)).toEqual(["dynamic"]);
@ -92,12 +108,11 @@ describe("remote catalog provider", () => {
vi.spyOn(globalThis, "fetch").mockImplementation(async () => responses.shift() as Response);
const provider = testProvider(localGeneratedAt);
const store = new InMemoryModelsStore();
const refresh = { credential: { type: "api_key" } as const, store: scopedStore(store), allowNetwork: true };
await provider.refreshModels?.(refresh);
await refreshProvider(provider, store);
expect(provider.getModels().map((entry) => entry.id)).toEqual(["static"]);
await provider.refreshModels?.({ ...refresh, force: true });
await refreshProvider(provider, store, { force: true });
expect(provider.getModels().map((entry) => entry.id)).toEqual(["static", "newer"]);
expect(await store.read(provider.id)).toMatchObject({ lastModified: Date.parse(newerHeader) });
});
@ -112,14 +127,13 @@ describe("remote catalog provider", () => {
const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async () => responses.shift() as Response);
const provider = testProvider();
const store = new InMemoryModelsStore();
const refresh = { credential: { type: "api_key" } as const, store: scopedStore(store), allowNetwork: true };
await provider.refreshModels?.(refresh);
await refreshProvider(provider, store);
expect(fetchSpy.mock.calls[0]?.[1]?.headers).not.toHaveProperty("if-none-match");
expect(await store.read(provider.id)).toMatchObject({ etag: '"catalog-1"' });
const checkedAt = (await store.read(provider.id))?.checkedAt;
await provider.refreshModels?.({ ...refresh, force: true });
await refreshProvider(provider, store, { force: true });
expect(fetchSpy.mock.calls[1]?.[1]?.headers).toMatchObject({ "if-none-match": '"catalog-1"' });
expect(provider.getModels().map((entry) => entry.id)).toEqual(["static", "dynamic"]);
@ -139,10 +153,9 @@ describe("remote catalog provider", () => {
vi.spyOn(globalThis, "fetch").mockImplementation(async () => responses.shift() as Response);
const provider = testProvider();
const store = new InMemoryModelsStore();
const refresh = { credential: { type: "api_key" } as const, store: scopedStore(store), allowNetwork: true };
await provider.refreshModels?.(refresh);
await provider.refreshModels?.({ ...refresh, force: true });
await refreshProvider(provider, store);
await refreshProvider(provider, store, { force: true });
expect((await store.read(provider.id))?.etag).toBeUndefined();
});
@ -158,32 +171,68 @@ describe("remote catalog provider", () => {
const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async () => responses.shift() as Response);
const provider = testProvider();
const store = new InMemoryModelsStore();
const refresh = { credential: { type: "api_key" } as const, store: scopedStore(store), allowNetwork: true };
await provider.refreshModels?.(refresh);
await expect(provider.refreshModels?.({ ...refresh, force: true })).rejects.toThrow(/429/);
await refreshProvider(provider, store);
await expect(refreshProvider(provider, store, { force: true })).rejects.toThrow(/429/);
const stored = await store.read(provider.id);
expect(stored?.etag).toBe('"catalog-1"');
expect(stored?.models.map((entry) => entry.id)).toEqual(["dynamic"]);
await provider.refreshModels?.({ ...refresh, force: true });
await refreshProvider(provider, store, { force: true });
expect(fetchSpy.mock.calls[2]?.[1]?.headers).toMatchObject({ "if-none-match": '"catalog-1"' });
expect(provider.getModels().map((entry) => entry.id)).toEqual(["static", "dynamic"]);
});
it("lets a newer catalog request bypass a stalled older request without stale publication", async () => {
let calls = 0;
let markFirstStarted: (() => void) | undefined;
let finishFirst: ((response: Response) => void) | undefined;
const firstStarted = new Promise<void>((resolve) => {
markFirstStarted = resolve;
});
const firstResponse = new Promise<Response>((resolve) => {
finishFirst = resolve;
});
vi.spyOn(globalThis, "fetch").mockImplementation(async () => {
calls++;
if (calls === 1) {
markFirstStarted?.();
return firstResponse;
}
return new Response(JSON.stringify({ newer: model("newer") }), {
headers: { "content-type": "application/json" },
});
});
const provider = testProvider();
const store = new InMemoryModelsStore();
const models = createModels({ modelsStore: store });
models.setProvider(provider);
const first = models.refresh({ providers: [provider.id], force: true });
await firstStarted;
const second = models.refresh({ providers: [provider.id], force: true });
await second;
await first;
expect(provider.getModels().map((entry) => entry.id)).toEqual(["static", "newer"]);
finishFirst?.(
new Response(JSON.stringify({ older: model("older") }), {
headers: { "content-type": "application/json" },
}),
);
await new Promise((resolve) => setTimeout(resolve, 0));
expect(provider.getModels().map((entry) => entry.id)).toEqual(["static", "newer"]);
expect((await store.read(provider.id))?.models.map((entry) => entry.id)).toEqual(["newer"]);
});
it("treats unimplemented pi.dev catalog routes as an unavailable overlay", async () => {
vi.spyOn(globalThis, "fetch").mockResolvedValue(new Response("not implemented", { status: 501 }));
const provider = testProvider();
const store = new InMemoryModelsStore();
await expect(
provider.refreshModels?.({
credential: { type: "api_key" },
store: scopedStore(store),
allowNetwork: true,
}),
).resolves.toBeUndefined();
await expect(refreshProvider(provider, store)).resolves.toBeUndefined();
expect(provider.getModels().map((entry) => entry.id)).toEqual(["static"]);
expect(await store.read(provider.id)).toMatchObject({ models: [], checkedAt: expect.any(Number) });
});

View file

@ -1,4 +1,5 @@
import { describe, expect, test } from "vitest";
import type { CredentialStore } from "@earendil-works/pi-ai";
import { describe, expect, test, vi } from "vitest";
import { AuthStorage } from "../src/core/auth-storage.ts";
import { RuntimeCredentials } from "../src/core/runtime-credentials.ts";
@ -29,6 +30,49 @@ describe("RuntimeCredentials", () => {
]);
});
test("forwards operation signals to the persistent store", async () => {
const controller = new AbortController();
const received: AbortSignal[] = [];
const storage: CredentialStore = {
read: async (_providerId, options) => {
received.push(options?.signal as AbortSignal);
return undefined;
},
list: async (options) => {
received.push(options?.signal as AbortSignal);
return [];
},
modify: async (_providerId, _fn, options) => {
received.push(options?.signal as AbortSignal);
return undefined;
},
delete: async (_providerId, options) => {
received.push(options?.signal as AbortSignal);
},
};
const credentials = new RuntimeCredentials(storage);
await credentials.read("anthropic", { signal: controller.signal });
await credentials.list({ signal: controller.signal });
await credentials.modify("anthropic", async () => undefined, { signal: controller.signal });
await credentials.delete("anthropic", { signal: controller.signal });
expect(received).toEqual([controller.signal, controller.signal, controller.signal, controller.signal]);
});
test("keeps a runtime override when persistent deletion is cancelled", async () => {
const aborted = new Error("cancelled");
aborted.name = "AbortError";
const storage = AuthStorage.inMemory({ anthropic: { type: "api_key", key: "stored-key" } });
const deleteSpy = vi.spyOn(storage, "delete").mockRejectedValueOnce(aborted);
const credentials = new RuntimeCredentials(storage);
credentials.setRuntimeApiKey("anthropic", "runtime-key");
await expect(credentials.delete("anthropic", { signal: new AbortController().signal })).rejects.toBe(aborted);
expect(deleteSpy).toHaveBeenCalledTimes(1);
expect(await credentials.read("anthropic")).toEqual({ type: "api_key", key: "runtime-key" });
});
test("delete clears both the override and persisted credential", async () => {
const storage = AuthStorage.inMemory({ anthropic: { type: "api_key", key: "stored-key" } });
const credentials = new RuntimeCredentials(storage);

View file

@ -15,12 +15,12 @@ function createInteractiveContext(options: {
}) {
let selector: ScopedModelsSelectorComponent | undefined;
const setScopedModels = vi.fn();
const getAvailable = vi.fn().mockResolvedValue(options.allModels);
const getAvailableSnapshot = vi.fn(() => options.allModels);
const context = {
session: {
modelRuntime: {
refresh: vi.fn(),
getAvailable,
refresh: vi.fn().mockResolvedValue({ aborted: false, errors: new Map() }),
getAvailableSnapshot,
},
scopedModels: options.scopedModels ?? [],
setScopedModels,
@ -36,7 +36,7 @@ function createInteractiveContext(options: {
updateAvailableProviderCount: vi.fn(),
ui: { requestRender: vi.fn() },
};
return { context, getAvailable, getSelector: () => selector, setScopedModels };
return { context, getAvailableSnapshot, getSelector: () => selector, setScopedModels };
}
async function showModelsSelector(context: object): Promise<void> {
@ -93,7 +93,7 @@ describe("issue #6949 unavailable scoped models", () => {
const harness = await createHarness({ models: [{ id: "available", name: "Available" }] });
harnesses.push(harness);
const unavailableIds = ["unavailable-one", "unavailable-two"].map((id) => `${harness.models[0].provider}/${id}`);
const { context, getAvailable, getSelector } = createInteractiveContext({
const { context, getAvailableSnapshot, getSelector } = createInteractiveContext({
allModels: [],
enabledModelIds: unavailableIds,
});
@ -106,7 +106,7 @@ describe("issue #6949 unavailable scoped models", () => {
for (const unavailableId of unavailableIds) {
expect(rendered).toContain(`${unavailableId} [unavailable] ✗`);
}
expect(getAvailable).toHaveBeenCalledTimes(2);
expect(getAvailableSnapshot).toHaveBeenCalled();
});
it("opens when only a session-scoped model is unavailable", async () => {

View file

@ -0,0 +1,125 @@
import type { Api, Model, Provider } from "@earendil-works/pi-ai";
import { afterEach, describe, expect, it, vi } from "vitest";
import { AuthStorage } from "../../../src/core/auth-storage.ts";
import { ModelRuntime } from "../../../src/core/model-runtime.ts";
import { InteractiveMode } from "../../../src/modes/interactive/interactive-mode.ts";
import { createHarness, type Harness } from "../harness.ts";
const dynamicModel: Model<"openai-completions"> = {
id: "dynamic",
name: "Dynamic",
api: "openai-completions",
provider: "stalled-login",
baseUrl: "https://example.test/v1",
reasoning: false,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 1000,
maxTokens: 100,
};
describe("issues #7027 and #7113 credential refresh hang", () => {
let harness: Harness | undefined;
afterEach(() => {
vi.useRealTimers();
harness?.cleanup();
harness = undefined;
vi.restoreAllMocks();
});
it("does not hold login behind an older stalled network catalog refresh", async () => {
let markNetworkStarted: (() => void) | undefined;
const networkStarted = new Promise<void>((resolve) => {
markNetworkStarted = resolve;
});
const provider: Provider<"openai-completions"> = {
id: "stalled-login",
name: "Stalled Login",
auth: {
apiKey: {
name: "API key",
login: async () => ({ type: "api_key", key: "secret" }),
check: async ({ credential }) =>
credential?.key ? { type: "api_key", source: "stored key" } : undefined,
resolve: async ({ credential }) => ({
auth: { apiKey: credential?.key ?? "ambient-key" },
source: credential?.key ? "stored key" : "ambient key",
}),
},
},
getModels: () => [dynamicModel],
refreshModels: async ({ allowNetwork }) => {
if (!allowNetwork) return;
markNetworkStarted?.();
await new Promise<void>(() => {});
},
stream: () => {
throw new Error("unused");
},
streamSimple: () => {
throw new Error("unused");
},
};
const credentials = AuthStorage.inMemory();
const runtime = await ModelRuntime.create({ credentials, modelsPath: null, allowModelNetwork: false });
runtime.registerNativeProvider(provider);
await runtime.refresh({ allowNetwork: false, providers: [provider.id] });
const stalledRefresh = runtime.refresh({ allowNetwork: true, providers: [provider.id] });
await networkStarted;
await expect(
runtime.login(provider.id, "api_key", { prompt: async () => "unused", notify: () => {} }),
).resolves.toEqual({ type: "api_key", key: "secret" });
expect(runtime.getAvailableSnapshot().map((model) => model.id)).toContain(dynamicModel.id);
expect(await credentials.read(provider.id)).toEqual({ type: "api_key", key: "secret" });
await expect(stalledRefresh).resolves.toMatchObject({ aborted: false });
});
it("completes interactive login before its bounded background refresh", async () => {
harness = await createHarness();
vi.useFakeTimers();
const runtime = harness.session.modelRuntime;
vi.spyOn(runtime, "refresh").mockImplementation(
(options) =>
new Promise((resolve) => {
options?.signal?.addEventListener("abort", () => resolve({ aborted: true, errors: new Map() }), {
once: true,
});
}),
);
const showWarning = vi.fn();
const context = {
session: harness.session,
updateAvailableProviderCount: vi.fn(),
footer: { invalidate: vi.fn() },
updateEditorBorderColor: vi.fn(),
showStatus: vi.fn(),
showError: vi.fn(),
showWarning,
maybeWarnAboutAnthropicSubscriptionAuth: vi.fn(),
checkDaxnutsEasterEgg: vi.fn(),
ui: { requestRender: vi.fn() },
};
const complete = Reflect.get(InteractiveMode.prototype, "completeProviderAuthentication") as (
this: object,
providerId: string,
providerName: string,
authType: "oauth" | "api_key",
previousModel: Model<Api>,
) => Promise<void>;
await complete.call(context, dynamicModel.provider, "Stalled Login", "api_key", harness.getModel());
expect(runtime.refresh).toHaveBeenCalledWith({
providers: [dynamicModel.provider],
signal: expect.any(AbortSignal),
});
expect(showWarning).not.toHaveBeenCalled();
await vi.advanceTimersByTimeAsync(15_000);
expect(showWarning).toHaveBeenCalledWith(
"Saved API key for Stalled Login, but its model catalog refresh timed out; using cached models.",
);
});
});

View file

@ -0,0 +1,107 @@
import type { Api, Model, ModelsRefreshResult } from "@earendil-works/pi-ai";
import { setKeybindings, type TUI } from "@earendil-works/pi-tui";
import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
import { KeybindingsManager } from "../../../src/core/keybindings.ts";
import type { ScopedModelsSelectorComponent } from "../../../src/modes/interactive/components/scoped-models-selector.ts";
import { InteractiveMode } from "../../../src/modes/interactive/interactive-mode.ts";
import { initTheme } from "../../../src/modes/interactive/theme/theme.ts";
import { stripAnsi } from "../../../src/utils/ansi.ts";
import { createHarness, type Harness } from "../harness.ts";
const showModelsSelector = Reflect.get(InteractiveMode.prototype, "showModelsSelector") as (this: object) => void;
function openSelector(harness: Harness, initialModels: readonly Model<Api>[]) {
let snapshot = initialModels;
let finishRefresh: ((result: ModelsRefreshResult) => void) | undefined;
let refreshSignal: AbortSignal | undefined;
let selector: ScopedModelsSelectorComponent | undefined;
let dispose: (() => void) | undefined;
const done = vi.fn();
vi.spyOn(harness.session.modelRuntime, "getAvailableSnapshot").mockImplementation(() => snapshot);
vi.spyOn(harness.session.modelRuntime, "refresh").mockImplementation(
(options) =>
new Promise((resolve) => {
refreshSignal = options?.signal;
finishRefresh = resolve;
}),
);
const context = {
session: harness.session,
settingsManager: harness.settingsManager,
showSelector: (
factory: (close: () => void) => {
component: ScopedModelsSelectorComponent;
dispose?: () => void;
},
) => {
const close = () => {
dispose?.();
done();
};
const created = factory(close);
selector = created.component;
dispose = created.dispose;
},
updateAvailableProviderCount: vi.fn(),
ui: { requestRender: vi.fn() } as unknown as TUI,
};
showModelsSelector.call(context);
if (!selector) throw new Error("Expected scoped-model selector to open");
return {
done,
get refreshSignal() {
return refreshSignal;
},
selector,
complete(models: readonly Model<Api>[], result: ModelsRefreshResult) {
snapshot = models;
if (!finishRefresh) throw new Error("Expected model refresh to start");
finishRefresh(result);
},
};
}
describe("issue #7153 scoped models refresh", () => {
let harness: Harness | undefined;
beforeAll(() => initTheme("dark"));
beforeEach(() => setKeybindings(new KeybindingsManager()));
afterEach(() => {
harness?.cleanup();
harness = undefined;
vi.restoreAllMocks();
});
it("renders cached models immediately and updates after background refresh", async () => {
harness = await createHarness({
models: [
{ id: "cached", name: "Cached" },
{ id: "refreshed", name: "Refreshed" },
],
});
const refresh = openSelector(harness, [harness.models[0]]);
const initial = stripAnsi(refresh.selector.render(100).join("\n"));
expect(initial).toContain("cached");
expect(initial).toContain("Refreshing model catalogs…");
expect(initial).not.toContain("refreshed");
refresh.complete(harness.models, { aborted: false, errors: new Map() });
await vi.waitFor(() => {
const rendered = stripAnsi(refresh.selector.render(100).join("\n"));
expect(rendered).toContain("refreshed");
expect(rendered).toContain("Model catalogs refreshed.");
});
});
it("cancels the background refresh when the selector closes", async () => {
harness = await createHarness({ models: [{ id: "cached", name: "Cached" }] });
const refresh = openSelector(harness, harness.models);
expect(refresh.refreshSignal).toBeDefined();
refresh.selector.handleInput("\x1b");
expect(refresh.refreshSignal?.aborted).toBe(true);
expect(refresh.done).toHaveBeenCalledOnce();
});
});

View file

@ -1,3 +1,4 @@
import type { Models } from "@earendil-works/pi-ai";
import { afterEach, describe, expect, it, vi } from "vitest";
import { createHarness, type Harness } from "../harness.ts";
@ -103,4 +104,30 @@ describe("issue #7301 stalled availability refresh", () => {
await expect(staleRefresh).rejects.toThrow("stale credential list failure");
expect(runtime.getError()).toBeUndefined();
});
it("does not let a stale provider-scoped failure overwrite a newer availability pass", async () => {
harness = await createHarness();
const runtime = harness.session.modelRuntime;
const models = Reflect.get(runtime, "models") as Models;
const originalGetAvailable = models.getAvailable.bind(models);
const started = createDeferred();
const gate = createDeferred();
let stall = true;
models.getAvailable = async (providerId, options) => {
if (!providerId || !stall) return originalGetAvailable(providerId, options);
stall = false;
started.resolve();
await gate.promise;
throw new Error("stale provider availability failure");
};
const staleRefresh = runtime.getAvailable(harness.getModel().provider);
await started.promise;
await runtime.getAvailable();
expect(runtime.getError()).toBeUndefined();
gate.resolve();
await expect(staleRefresh).rejects.toThrow("stale provider availability failure");
expect(runtime.getError()).toBeUndefined();
});
});

View file

@ -0,0 +1,46 @@
import type { Api, Model } from "@earendil-works/pi-ai";
import { afterEach, describe, expect, it, vi } from "vitest";
import { InteractiveMode } from "../../../src/modes/interactive/interactive-mode.ts";
import { createHarness, type Harness } from "../harness.ts";
const findExactModelMatch = Reflect.get(InteractiveMode.prototype, "findExactModelMatch") as (
this: object,
searchTerm: string,
) => Promise<Model<Api> | undefined>;
describe("issue #7443 /model cached match", () => {
let harness: Harness | undefined;
afterEach(() => {
harness?.cleanup();
harness = undefined;
vi.restoreAllMocks();
});
it("matches the availability snapshot without starting a catalog refresh", async () => {
harness = await createHarness({ models: [{ id: "cached", name: "Cached" }] });
const refresh = vi.spyOn(harness.session.modelRuntime, "refresh").mockImplementation(() => new Promise(() => {}));
const context = { session: harness.session, showStatus: vi.fn(), showWarning: vi.fn() };
const model = await findExactModelMatch.call(context, harness.models[0].id);
expect(model?.id).toBe("cached");
expect(refresh).not.toHaveBeenCalled();
expect(context.showStatus).not.toHaveBeenCalled();
});
it("uses a caller-owned deadline only after a cache miss", async () => {
harness = await createHarness({ models: [{ id: "cached", name: "Cached" }] });
const refresh = vi.spyOn(harness.session.modelRuntime, "refresh").mockResolvedValue({
aborted: true,
errors: new Map(),
});
const context = { session: harness.session, showStatus: vi.fn(), showWarning: vi.fn() };
await expect(findExactModelMatch.call(context, "not-cached")).resolves.toBeUndefined();
expect(refresh).toHaveBeenCalledOnce();
expect(refresh.mock.calls[0]?.[0]?.signal).toBeInstanceOf(AbortSignal);
expect(context.showStatus).toHaveBeenCalledWith("Refreshing model catalogs…");
});
});

View file

@ -93,7 +93,7 @@ export async function resolveApiKey(provider: string): Promise<string | undefine
if (!oauth) return undefined;
let credential = entry;
if (Date.now() >= credential.expires) {
credential = await oauth.refresh(credential);
credential = await oauth.refresh(credential, new AbortController().signal);
storage[provider] = credential;
saveAuthStorage(storage);
}