mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-08-20 14:14:00 +00:00
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:
parent
b06dc76fd7
commit
fed6009cc9
86 changed files with 3426 additions and 870 deletions
|
|
@ -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)).
|
||||
|
||||
|
|
|
|||
|
|
@ -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';
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 };
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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({
|
||||
|
|
|
|||
|
|
@ -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 };
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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}`);
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
},
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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: {
|
||||
|
|
|
|||
|
|
@ -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: {},
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
50
packages/ai/src/utils/abort.ts
Normal file
50
packages/ai/src/utils/abort.ts
Normal 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();
|
||||
});
|
||||
}
|
||||
|
|
@ -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({
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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: () => {},
|
||||
}),
|
||||
|
|
|
|||
|
|
@ -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"]);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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", () => {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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];
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
|
|
|
|||
|
|
@ -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({
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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 }));
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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. */
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
},
|
||||
});
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -179,6 +179,8 @@ export {
|
|||
} from "./core/model-resolver.ts";
|
||||
export {
|
||||
type CreateModelRuntimeOptions,
|
||||
CredentialSynchronizationError,
|
||||
type CredentialSynchronizationOperation,
|
||||
ModelRuntime,
|
||||
type ModelRuntimeAuthOverrides,
|
||||
} from "./core/model-runtime.ts";
|
||||
|
|
|
|||
|
|
@ -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") {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
48
packages/coding-agent/src/utils/abort.ts
Normal file
48
packages/coding-agent/src/utils/abort.ts
Normal 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();
|
||||
});
|
||||
}
|
||||
|
|
@ -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({
|
||||
|
|
|
|||
|
|
@ -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", () => {
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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({
|
||||
|
|
|
|||
|
|
@ -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", {
|
||||
|
|
|
|||
375
packages/coding-agent/test/model-runtime-credential-sync.test.ts
Normal file
375
packages/coding-agent/test/model-runtime-credential-sync.test.ts
Normal 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" });
|
||||
});
|
||||
});
|
||||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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) });
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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…");
|
||||
});
|
||||
});
|
||||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue