From fed6009cc9ccfb15088bc55244120918b5fe23f2 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 4 Aug 2026 00:00:12 +0200 Subject: [PATCH] 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. --- packages/ai/CHANGELOG.md | 70 +++ packages/ai/README.md | 18 +- packages/ai/src/auth/credential-store.ts | 58 ++- packages/ai/src/auth/helpers.ts | 8 +- packages/ai/src/auth/oauth/anthropic.ts | 53 ++- packages/ai/src/auth/oauth/device-code.ts | 12 +- packages/ai/src/auth/oauth/github-copilot.ts | 72 ++- packages/ai/src/auth/oauth/kimi-coding.ts | 37 +- packages/ai/src/auth/oauth/openai-codex.ts | 29 +- packages/ai/src/auth/oauth/openrouter.ts | 28 +- packages/ai/src/auth/oauth/radius.ts | 33 +- packages/ai/src/auth/oauth/xai.ts | 16 +- packages/ai/src/auth/resolve.ts | 76 +++- packages/ai/src/auth/types.ts | 33 +- packages/ai/src/cli.ts | 1 + packages/ai/src/images-models.ts | 1 + packages/ai/src/models-store.ts | 26 +- packages/ai/src/models.ts | 357 +++++++++++---- packages/ai/src/providers/amazon-bedrock.ts | 22 +- packages/ai/src/providers/anthropic.ts | 15 +- packages/ai/src/providers/cloudflare-auth.ts | 23 +- packages/ai/src/providers/google-vertex.ts | 23 +- packages/ai/src/providers/radius.ts | 69 +-- packages/ai/src/utils/abort.ts | 50 ++ packages/ai/test/anthropic-auth-token.test.ts | 4 + packages/ai/test/anthropic-oauth.test.ts | 19 +- packages/ai/test/github-copilot-oauth.test.ts | 19 +- packages/ai/test/kimi-coding-oauth.test.ts | 40 +- packages/ai/test/models-runtime.test.ts | 426 +++++++++++++++++- packages/ai/test/oauth-auth.test.ts | 26 +- packages/ai/test/oauth-device-code.test.ts | 6 + packages/ai/test/oauth.ts | 4 +- packages/ai/test/openai-codex-oauth.test.ts | 21 +- packages/ai/test/openrouter-oauth.test.ts | 9 + packages/ai/test/providers.test.ts | 62 ++- packages/ai/test/radius-oauth.test.ts | 10 +- packages/ai/test/xai-oauth.test.ts | 9 +- packages/coding-agent/CHANGELOG.md | 72 ++- packages/coding-agent/docs/custom-provider.md | 6 +- packages/coding-agent/docs/extensions.md | 9 +- packages/coding-agent/docs/sdk.md | 23 +- .../custom-provider-anthropic/index.ts | 3 +- .../custom-provider-gitlab-duo/index.ts | 3 +- .../examples/sdk/09-api-keys-and-oauth.ts | 2 +- .../examples/sdk/12-full-control.ts | 2 +- packages/coding-agent/examples/sdk/README.md | 2 +- .../coding-agent/src/cli/credential-print.ts | 13 +- packages/coding-agent/src/cli/list-models.ts | 8 +- .../src/core/agent-session-services.ts | 2 + .../coding-agent/src/core/agent-session.ts | 13 +- .../coding-agent/src/core/auth-storage.ts | 138 ++++-- .../coding-agent/src/core/extensions/types.ts | 4 +- .../coding-agent/src/core/model-registry.ts | 6 +- .../coding-agent/src/core/model-resolver.ts | 36 +- .../coding-agent/src/core/model-runtime.ts | 303 +++++++++---- .../coding-agent/src/core/models-store.ts | 33 +- .../src/core/provider-composer.ts | 33 +- .../src/core/remote-catalog-provider.ts | 138 +++--- .../src/core/runtime-credentials.ts | 20 +- .../src/extensions/llama/index.ts | 11 +- .../src/extensions/llama/provider.ts | 28 +- packages/coding-agent/src/index.ts | 2 + packages/coding-agent/src/main.ts | 28 +- .../interactive/components/model-selector.ts | 15 +- .../components/scoped-models-selector.ts | 28 ++ .../src/modes/interactive/interactive-mode.ts | 287 ++++++++---- .../coding-agent/src/modes/rpc/rpc-mode.ts | 4 +- .../coding-agent/src/package-manager-cli.ts | 11 +- packages/coding-agent/src/utils/abort.ts | 48 ++ .../coding-agent/test/auth-storage.test.ts | 209 +++++++++ .../test/extensions-runner.test.ts | 2 +- .../coding-agent/test/llama-extension.test.ts | 33 +- .../coding-agent/test/model-registry.test.ts | 4 +- .../coding-agent/test/model-resolver.test.ts | 4 +- .../test/model-runtime-auth-options.test.ts | 32 ++ .../model-runtime-credential-sync.test.ts | 375 +++++++++++++++ .../coding-agent/test/models-store.test.ts | 25 +- packages/coding-agent/test/radius.test.ts | 20 +- .../test/remote-catalog-provider.test.ts | 111 +++-- .../test/runtime-credentials.test.ts | 46 +- .../6949-unavailable-scoped-model.test.ts | 12 +- .../7027-credential-refresh-hang.test.ts | 125 +++++ .../7153-scoped-models-refresh.test.ts | 107 +++++ .../7301-stalled-availability-refresh.test.ts | 27 ++ .../7443-model-command-cached-match.test.ts | 46 ++ packages/coding-agent/test/utilities.ts | 2 +- 86 files changed, 3426 insertions(+), 870 deletions(-) create mode 100644 packages/ai/src/utils/abort.ts create mode 100644 packages/coding-agent/src/utils/abort.ts create mode 100644 packages/coding-agent/test/model-runtime-credential-sync.test.ts create mode 100644 packages/coding-agent/test/suite/regressions/7027-credential-refresh-hang.test.ts create mode 100644 packages/coding-agent/test/suite/regressions/7153-scoped-models-refresh.test.ts create mode 100644 packages/coding-agent/test/suite/regressions/7443-model-command-cached-match.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 153431255..2863586e5 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -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)). diff --git a/packages/ai/README.md b/packages/ai/README.md index da7f20124..48bee3586 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -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'; diff --git a/packages/ai/src/auth/credential-store.ts b/packages/ai/src/auth/credential-store.ts index f2a09e1a3..fc17f6cc6 100644 --- a/packages/ai/src/auth/credential-store.ts +++ b/packages/ai/src/auth/credential-store.ts @@ -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(); private chains = new Map>(); - /** Serialize tasks per provider id. */ - private enqueue(providerId: string, task: () => Promise): Promise { + /** Serialize tasks per provider id without releasing the chain before active work settles. */ + private enqueue(providerId: string, task: () => Promise, options?: AuthOperationOptions): Promise { + 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 { + async read(providerId: string, options?: AuthOperationOptions): Promise { + options?.signal?.throwIfAborted(); return this.credentials.get(providerId); } - async list(): Promise { + async list(options?: AuthOperationOptions): Promise { + options?.signal?.throwIfAborted(); return [...this.credentials].map(([providerId, credential]) => ({ providerId, type: credential.type })); } modify( providerId: string, fn: (current: Credential | undefined) => Promise, + options?: AuthOperationOptions, ): Promise { - 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 { - return this.enqueue(providerId, async () => { - this.credentials.delete(providerId); - }); + delete(providerId: string, options?: AuthOperationOptions): Promise { + return this.enqueue( + providerId, + async () => { + this.credentials.delete(providerId); + }, + options, + ); } } diff --git a/packages/ai/src/auth/helpers.ts b/packages/ai/src/auth/helpers.ts index 251c60da1..cd3640dc0 100644 --- a/packages/ai/src/auth/helpers.ts +++ b/packages/ai/src/auth/helpers.ts @@ -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), }; } diff --git a/packages/ai/src/auth/oauth/anthropic.ts b/packages/ai/src/auth/oauth/anthropic.ts index dbaaa3001..ad20d68c3 100644 --- a/packages/ai/src/auth/oauth/anthropic.ts +++ b/packages/ai/src/auth/oauth/anthropic.ts @@ -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): Promise { +async function postJson(url: string, body: Record, signal: AbortSignal): Promise { const response = await fetch(url, { method: "POST", headers: { @@ -175,7 +175,7 @@ async function postJson(url: string, body: Record): 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 { 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 { +async function loginAnthropic(interaction: ProviderAuthInteraction): Promise { 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 { +async function refreshAnthropicToken(refreshToken: string, signal: AbortSignal): Promise { 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 refreshAnthropicToken(credential.refresh), + refresh: (credential, signal) => refreshAnthropicToken(credential.refresh, signal), async toAuth(credential) { return { apiKey: credential.access }; diff --git a/packages/ai/src/auth/oauth/device-code.ts b/packages/ai/src/auth/oauth/device-code.ts index eac240837..f87c5e3b5 100644 --- a/packages/ai/src/auth/oauth/device-code.ts +++ b/packages/ai/src/auth/oauth/device-code.ts @@ -20,12 +20,12 @@ export type OAuthDeviceCodePollOptions = { expiresInSeconds?: number; waitBeforeFirstPoll?: boolean; poll: () => Promise>; - signal?: AbortSignal; + signal: AbortSignal; }; -function abortableSleep(ms: number, signal: AbortSignal | undefined, cancelMessage: string): Promise { +function abortableSleep(ms: number, signal: AbortSignal, cancelMessage: string): Promise { 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(options: OAuthDeviceCodePollOpt } while (Date.now() < deadline) { - if (options.signal?.aborted) { + if (options.signal.aborted) { throw new Error(CANCEL_MESSAGE); } diff --git a/packages/ai/src/auth/oauth/github-copilot.ts b/packages/ai/src/auth/oauth/github-copilot.ts index 0fe1a04e7..6261972d2 100644 --- a/packages/ai/src/auth/oauth/github-copilot.ts +++ b/packages/ai/src/auth/oauth/github-copilot.ts @@ -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 { +async function fetchAvailableGitHubCopilotModelIds( + copilotToken: string, + enterpriseDomain: string | undefined, + signal: AbortSignal, +): Promise { 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 { return response.json(); } -async function startDeviceFlow(domain: string): Promise { +async function startDeviceFlow(domain: string, signal: AbortSignal): Promise { const urls = getUrls(domain); const data = await fetchJson(urls.deviceCodeUrl, { method: "POST", @@ -149,6 +153,7 @@ async function startDeviceFlow(domain: string): Promise { client_id: CLIENT_ID, scope: "read:user", }), + signal, }); if (!data || typeof data !== "object") { @@ -195,7 +200,7 @@ async function startDeviceFlow(domain: string): Promise { async function pollForGitHubAccessToken( domain: string, device: DeviceCodeResponse, - signal?: AbortSignal, + signal: AbortSignal, ): Promise { const urls = getUrls(domain); return pollOAuthDeviceCodeFlow({ @@ -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 { 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 { - const credentials = await refreshGitHubCopilotAccessToken(refreshToken, enterpriseDomain); +async function refreshGitHubCopilotToken( + refreshToken: string, + enterpriseDomain: string | undefined, + signal: AbortSignal, +): Promise { + 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 { +async function enableGitHubCopilotModel( + token: string, + modelId: string, + enterpriseDomain: string | undefined, + signal: AbortSignal, +): Promise { 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 { +async function enableAllGitHubCopilotModels( + token: string, + enterpriseDomain: string | undefined, + signal: AbortSignal, +): Promise { 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 { +async function loginGitHubCopilot(interaction: ProviderAuthInteraction): Promise { 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 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) { diff --git a/packages/ai/src/auth/oauth/kimi-coding.ts b/packages/ai/src/auth/oauth/kimi-coding.ts index 93e7062de..e98493e94 100644 --- a/packages/ai/src/auth/oauth/kimi-coding.ts +++ b/packages/ai/src/auth/oauth/kimi-coding.ts @@ -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 { @@ -66,7 +66,7 @@ function trustedHttpUrl(value: unknown): string | null { } } -async function startDeviceAuthorization(oauthHost: string, signal?: AbortSignal): Promise { +async function startDeviceAuthorization(oauthHost: string, signal: AbortSignal): Promise { const response = await fetch(`${oauthHost}/api/oauth/device_authorization`, { method: "POST", headers: { @@ -141,7 +141,7 @@ function parseTokenResponse(json: Record | null, operation: str async function pollForToken( oauthHost: string, device: DeviceAuthorization, - signal?: AbortSignal, + signal: AbortSignal, ): Promise { return pollOAuthDeviceCodeFlow({ intervalSeconds: device.intervalSeconds, @@ -206,25 +206,32 @@ async function pollForToken( }); } -function sleep(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)); +function sleep(ms: number, signal: AbortSignal): Promise { + 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 { +async function refreshToken(oauthHost: string, refreshTokenValue: string, signal: AbortSignal): Promise { 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 { +async function loginKimiCoding(interaction: ProviderAuthInteraction): Promise { const oauthHost = getOauthHost(); const device = await startDeviceAuthorization(oauthHost, interaction.signal); interaction.notify({ diff --git a/packages/ai/src/auth/oauth/openai-codex.ts b/packages/ai/src/auth/oauth/openai-codex.ts index f1b27c3ad..fc558c7cc 100644 --- a/packages/ai/src/auth/oauth/openai-codex.ts +++ b/packages/ai/src/auth/oauth/openai-codex.ts @@ -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 { 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 { +async function refreshAccessToken(refreshToken: string, signal: AbortSignal): Promise { let response: Response; try { response = await fetch(TOKEN_URL, { @@ -179,6 +179,7 @@ async function refreshAccessToken(refreshToken: string): Promise { 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 { return readTokenResponse(response, "refresh"); } -async function startOpenAICodexDeviceAuth(signal?: AbortSignal): Promise { +async function startOpenAICodexDeviceAuth(signal: AbortSignal): Promise { 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 { +async function pollOpenAICodexDeviceAuth(device: DeviceAuthInfo, signal: AbortSignal): Promise { return pollOAuthDeviceCodeFlow({ 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 { return credentialsFromToken(await exchangeAuthorizationCode(code, verifier, redirectUri, signal)); } -async function loginOpenAICodexDeviceCode(interaction: AuthInteraction): Promise { +async function loginOpenAICodexDeviceCode(interaction: ProviderAuthInteraction): Promise { 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 { +async function loginOpenAICodex(interaction: ProviderAuthInteraction): Promise { 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 { - return credentialsFromToken(await refreshAccessToken(refreshToken)); +async function refreshOpenAICodexToken(refreshToken: string, signal: AbortSignal): Promise { + 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 }; diff --git a/packages/ai/src/auth/oauth/openrouter.ts b/packages/ai/src/auth/oauth/openrouter.ts index ee066250a..b16be2a5b 100644 --- a/packages/ai/src/auth/oauth/openrouter.ts +++ b/packages/ai/src/auth/oauth/openrouter.ts @@ -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 { - 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 { - 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 { +async function loginOpenRouter(interaction: ProviderAuthInteraction): Promise { 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) { diff --git a/packages/ai/src/auth/oauth/radius.ts b/packages/ai/src/auth/oauth/radius.ts index 655316910..776ec6d65 100644 --- a/packages/ai/src/auth/oauth/radius.ts +++ b/packages/ai/src/auth/oauth/radius.ts @@ -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 { +async function loadRadiusOAuthDiscovery(gateway: string, signal: AbortSignal): Promise { 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 { 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 { +function startOAuthCallbackServer(expectedState: string, signal: AbortSignal): Promise { 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 { 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 { +async function requestDeviceAuthorization(gateway: string, signal: AbortSignal): Promise { 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 { +async function loginWithDeviceCode(gateway: string, interaction: ProviderAuthInteraction): Promise { 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}`); diff --git a/packages/ai/src/auth/oauth/xai.ts b/packages/ai/src/auth/oauth/xai.ts index 5432043d1..5b169785d 100644 --- a/packages/ai/src/auth/oauth/xai.ts +++ b/packages/ai/src/auth/oauth/xai.ts @@ -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, signal?: AbortSignal): Promise { +async function postForm(url: string, fields: Record, signal: AbortSignal): Promise { let response: Response; try { response = await fetch(url, { @@ -74,7 +74,7 @@ async function postForm(url: string, fields: Record, 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, 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 { +async function requestDeviceCode(signal: AbortSignal): Promise { const response = await postForm( XAI_DEVICE_CODE_URL, { @@ -158,7 +158,7 @@ async function requestDeviceCode(signal?: AbortSignal): Promise { return parseDeviceCode(response.body); } -async function pollForTokens(device: XaiDeviceCode, signal?: AbortSignal): Promise { +async function pollForTokens(device: XaiDeviceCode, signal: AbortSignal): Promise { return pollOAuthDeviceCodeFlow({ intervalSeconds: device.intervalSeconds, expiresInSeconds: device.expiresInSeconds, @@ -198,7 +198,7 @@ async function pollForTokens(device: XaiDeviceCode, signal?: AbortSignal): Promi }); } -async function loginXai(interaction: AuthInteraction): Promise { +async function loginXai(interaction: ProviderAuthInteraction): Promise { const device = await requestDeviceCode(interaction.signal); interaction.notify({ type: "device_code", @@ -210,7 +210,7 @@ async function loginXai(interaction: AuthInteraction): Promise return pollForTokens(device, interaction.signal); } -async function refreshXaiToken(refreshToken: string, signal?: AbortSignal): Promise { +async function refreshXaiToken(refreshToken: string, signal: AbortSignal): Promise { const response = await postForm( XAI_TOKEN_URL, { diff --git a/packages/ai/src/auth/resolve.ts b/packages/ai/src/auth/resolve.ts index 678332a8a..55a654918 100644 --- a/packages/ai/src/auth/resolve.ts +++ b/packages/ai/src/auth/resolve.ts @@ -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 { + 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 { + 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 { 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 { 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 { +async function readCredential( + credentials: CredentialStore, + providerId: string, + signal: AbortSignal, +): Promise { 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 }); } diff --git a/packages/ai/src/auth/types.ts b/packages/ai/src/auth/types.ts index 076388f3c..6fbdd0bb9 100644 --- a/packages/ai/src/auth/types.ts +++ b/packages/ai/src/auth/types.ts @@ -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; + read(providerId: string, options?: AuthOperationOptions): Promise; /** * List stored credential metadata without resolving or exposing secrets. * Implementations must not execute configured API-key commands while listing. */ - list(): Promise; + list(options?: AuthOperationOptions): Promise; /** * 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, + options?: AuthOperationOptions, ): Promise; /** Remove a credential (logout). Implementations serialize this against `modify`. */ - delete(providerId: string): Promise; + delete(providerId: string, options?: AuthOperationOptions): Promise; } /** 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; + login?(interaction: ProviderAuthInteraction): Promise; /** * 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; + check?(input: { + ctx: AuthContext; + credential?: ApiKeyCredential; + signal: AbortSignal; + }): Promise; /** * 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; + resolve(input: { + ctx: AuthContext; + credential?: ApiKeyCredential; + signal: AbortSignal; + }): Promise; } /** @@ -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; + login(interaction: ProviderAuthInteraction): Promise; /** * 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; + refresh(credential: OAuthCredential, signal: AbortSignal): Promise; /** * Side-effect-free derivation of request auth from a valid credential. diff --git a/packages/ai/src/cli.ts b/packages/ai/src/cli.ts index e5fff7a5e..14403946a 100644 --- a/packages/ai/src/cli.ts +++ b/packages/ai/src/cli.ts @@ -48,6 +48,7 @@ async function login(providerId: string): Promise { 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) { diff --git a/packages/ai/src/images-models.ts b/packages/ai/src/images-models.ts index 4afe0f704..8121eb9da 100644 --- a/packages/ai/src/images-models.ts +++ b/packages/ai/src/images-models.ts @@ -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) { diff --git a/packages/ai/src/models-store.ts b/packages/ai/src/models-store.ts index 8fe8fb9f5..96dd8b3b0 100644 --- a/packages/ai/src/models-store.ts +++ b/packages/ai/src/models-store.ts @@ -13,33 +13,33 @@ export interface ModelsStoreEntry { etag?: string; } -/** Persistent model catalogs keyed by provider ID. */ -export interface ModelsStore { - read(providerId: string): Promise; - write(providerId: string, entry: ModelsStoreEntry): Promise; - delete(providerId: string): Promise; +export interface ModelsStoreOperationOptions { + signal?: AbortSignal; } -/** ModelsStore scoped to one provider. Providers cannot access other providers' catalogs. */ -export interface ProviderModelsStore { - read(): Promise; - write(entry: ModelsStoreEntry): Promise; - delete(): Promise; +/** Persistent model catalogs keyed by provider ID. */ +export interface ModelsStore { + read(providerId: string, options?: ModelsStoreOperationOptions): Promise; + write(providerId: string, entry: ModelsStoreEntry, options?: ModelsStoreOperationOptions): Promise; + delete(providerId: string, options?: ModelsStoreOperationOptions): Promise; } export class InMemoryModelsStore implements ModelsStore { private readonly entries = new Map(); - async read(providerId: string): Promise { + async read(providerId: string, options?: ModelsStoreOperationOptions): Promise { + options?.signal?.throwIfAborted(); const entry = this.entries.get(providerId); return entry ? structuredClone(entry) : undefined; } - async write(providerId: string, entry: ModelsStoreEntry): Promise { + async write(providerId: string, entry: ModelsStoreEntry, options?: ModelsStoreOperationOptions): Promise { + options?.signal?.throwIfAborted(); this.entries.set(providerId, structuredClone(entry)); } - async delete(providerId: string): Promise { + async delete(providerId: string, options?: ModelsStoreOperationOptions): Promise { + options?.signal?.throwIfAborted(); this.entries.delete(providerId); } } diff --git a/packages/ai/src/models.ts b/packages/ai/src/models.ts index d5eed1a42..988b12f38 100644 --- a/packages/ai/src/models.ts +++ b/packages/ai/src/models.ts @@ -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; + /** + * Generation-checked publication. Persistence policy remains provider-owned; + * the update runs synchronously only after the selected persistence mutation. + */ + publish(publication: ModelsPublication): Promise; /** 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 { getModels(): readonly Model[]; /** - * 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; @@ -141,16 +159,17 @@ export interface Models { getModel(provider: string, id: string): Model | 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; /** Check whether a provider has complete auth configuration without refreshing OAuth. */ - checkAuth(providerId: string): Promise; + checkAuth(providerId: string, options?: AuthOperationOptions): Promise; /** Return models whose providers have complete auth configuration. */ - getAvailable(providerId?: string): Promise[]>; + getAvailable(providerId?: string, options?: AuthOperationOptions): Promise[]>; /** * 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; /** Remove the stored credential for a provider. */ - logout(providerId: string): Promise; + logout(providerId: string, options?: AuthOperationOptions): Promise; stream( model: Model, @@ -220,6 +239,9 @@ class ModelsImpl implements MutableModels { private credentials: CredentialStore; private modelsStore: ModelsStore; private authContext: AuthContext; + private refreshGenerations = new Map(); + private refreshControllers = new Map(); + private publicationChains = new Map>(); 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 { + 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>, + credential: Credential | undefined, + allowNetwork: boolean, + force: boolean | undefined, + generation: number, + signal: AbortSignal, + ): Promise { + 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 { const allowNetwork = options.allowNetwork ?? true; + const callerSignal = operationSignal(options.signal); const errors = new Map(); + 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> => - 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 { 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 { + private async readCredential(providerId: string, signal: AbortSignal): Promise { 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 { 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 { - const provider = this.providers.get(providerId); - if (!provider) return undefined; - return this.checkProviderAuth(provider, await this.readCredential(providerId)); + checkAuth(providerId: string, options?: AuthOperationOptions): Promise { + 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[]> { - 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[]> { + 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; @@ -414,10 +530,11 @@ class ModelsImpl implements MutableModels { providerOrModel: string | Model, overrides?: AuthResolutionOverrides, ): Promise { + 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 { + 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 = method.login({ ...interaction, signal }); + const credential = await raceWithAbortSignal(loginOperation, signal); + let mutationStarted = false; + let markMutationStarted: (() => void) | undefined; + const started = new Promise((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((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 { + async logout(providerId: string, options?: AuthOperationOptions): Promise { + 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 { auth: ProviderAuth; /** Static baseline model list (empty for purely dynamic providers). */ models: readonly Model[]; - /** 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[]>; filterModels?: (models: readonly Model[], credential: Credential | undefined) => readonly Model[]; /** Single implementation, or map keyed by `model.api` for mixed-API providers. */ @@ -556,7 +713,6 @@ export interface CreateProviderOptions { export function createProvider(input: CreateProviderOptions): Provider { const baselineModels = input.models; let dynamicModels: readonly Model[] = []; - let inflightRefresh: Promise | undefined; const fetchModels = input.fetchModels; const currentModels = (): readonly Model[] => { const merged = [...baselineModels]; @@ -594,25 +750,30 @@ export function createProvider(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); - } - 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); + 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, diff --git a/packages/ai/src/providers/amazon-bedrock.ts b/packages/ai/src/providers/amazon-bedrock.ts index 61deef5cb..d629f1e1c 100644 --- a/packages/ai/src/providers/amazon-bedrock.ts +++ b/packages/ai/src/providers/amazon-bedrock.ts @@ -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; }, }; diff --git a/packages/ai/src/providers/anthropic.ts b/packages/ai/src/providers/anthropic.ts index a3d73c8a5..d1ddce0f8 100644 --- a/packages/ai/src/providers/anthropic.ts +++ b/packages/ai/src/providers/anthropic.ts @@ -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; diff --git a/packages/ai/src/providers/cloudflare-auth.ts b/packages/ai/src/providers/cloudflare-auth.ts index 9fc50df47..72fd23d07 100644 --- a/packages/ai/src/providers/cloudflare-auth.ts +++ b/packages/ai/src/providers/cloudflare-auth.ts @@ -11,6 +11,7 @@ async function resolveValue( name: string, ctx: AuthContext, credential: ApiKeyCredential | undefined, + signal: AbortSignal, ): Promise { // 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: { diff --git a/packages/ai/src/providers/google-vertex.ts b/packages/ai/src/providers/google-vertex.ts index e66aa6eff..83c7d0fe8 100644 --- a/packages/ai/src/providers/google-vertex.ts +++ b/packages/ai/src/providers/google-vertex.ts @@ -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: {}, diff --git a/packages/ai/src/providers/radius.ts b/packages/ai/src/providers/radius.ts index 1321cd187..21db78527 100644 --- a/packages/ai/src/providers/radius.ts +++ b/packages/ai/src/providers/radius.ts @@ -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 | 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), diff --git a/packages/ai/src/utils/abort.ts b/packages/ai/src/utils/abort.ts new file mode 100644 index 000000000..5d76d19f4 --- /dev/null +++ b/packages/ai/src/utils/abort.ts @@ -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(operation: Promise, signal: AbortSignal): Promise { + if (signal.aborted) { + void operation.catch(() => {}); + return Promise.reject(abortReason(signal)); + } + + return new Promise((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(); + }); +} diff --git a/packages/ai/test/anthropic-auth-token.test.ts b/packages/ai/test/anthropic-auth-token.test.ts index 47c41e53a..e2780a175 100644 --- a/packages/ai/test/anthropic-auth-token.test.ts +++ b/packages/ai/test/anthropic-auth-token.test.ts @@ -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({ diff --git a/packages/ai/test/anthropic-oauth.test.ts b/packages/ai/test/anthropic-oauth.test.ts index 679dd7d5a..9dc227b93 100644 --- a/packages/ai/test/anthropic-oauth.test.ts +++ b/packages/ai/test/anthropic-oauth.test.ts @@ -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); diff --git a/packages/ai/test/github-copilot-oauth.test.ts b/packages/ai/test/github-copilot-oauth.test.ts index 26738f29e..cac3964a3 100644 --- a/packages/ai/test/github-copilot-oauth.test.ts +++ b/packages/ai/test/github-copilot-oauth.test.ts @@ -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(); diff --git a/packages/ai/test/kimi-coding-oauth.test.ts b/packages/ai/test/kimi-coding-oauth.test.ts index 07c56360f..d64bbf1bb 100644 --- a/packages/ai/test/kimi-coding-oauth.test.ts +++ b/packages/ai/test/kimi-coding-oauth.test.ts @@ -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): Respo }); } -function createInteraction(events: Array>): AuthInteraction { +function createInteraction(events: Array>): 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 => 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"); }); }); diff --git a/packages/ai/test/models-runtime.test.ts b/packages/ai/test/models-runtime.test.ts index f582b80e2..2e5626797 100644 --- a/packages/ai/test/models-runtime.test.ts +++ b/packages/ai/test/models-runtime.test.ts @@ -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((resolve) => { + markAuthStarted = resolve; + }); + const blockedAuth = new Promise((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((resolve) => { + markStarted = resolve; + }); + const stalled = new Promise((_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((resolve) => { + markFirstStarted = resolve; + }); + const firstBlocked = new Promise((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((resolve) => { + startCheck = resolve; + }); + const blockedCheck = new Promise((resolve) => { + finishCheck = resolve; + }); + let startResolve: (() => void) | undefined; + let finishResolve: (() => void) | undefined; + const resolveStarted = new Promise((resolve) => { + startResolve = resolve; + }); + const blockedResolve = new Promise((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((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((resolve) => { + startRefresh = resolve; + }); + const blockedRefresh = new Promise((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 }); diff --git a/packages/ai/test/oauth-auth.test.ts b/packages/ai/test/oauth-auth.test.ts index 0ba080886..65d5a386e 100644 --- a/packages/ai/test/oauth-auth.test.ts +++ b/packages/ai/test/oauth-auth.test.ts @@ -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"); diff --git a/packages/ai/test/oauth-device-code.test.ts b/packages/ai/test/oauth-device-code.test.ts index fe146f586..2bda9f5e6 100644 --- a/packages/ai/test/oauth-device-code.test.ts +++ b/packages/ai/test/oauth-device-code.test.ts @@ -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); diff --git a/packages/ai/test/oauth.ts b/packages/ai/test/oauth.ts index 57060b894..81c35d042 100644 --- a/packages/ai/test/oauth.ts +++ b/packages/ai/test/oauth.ts @@ -69,7 +69,9 @@ export async function resolveApiKey(provider: string): Promise= 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; diff --git a/packages/ai/test/openai-codex-oauth.test.ts b/packages/ai/test/openai-codex-oauth.test.ts index 64c0e8c79..f45d4aeea 100644 --- a/packages/ai/test/openai-codex-oauth.test.ts +++ b/packages/ai/test/openai-codex-oauth.test.ts @@ -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(); }); diff --git a/packages/ai/test/openrouter-oauth.test.ts b/packages/ai/test/openrouter-oauth.test.ts index 08e115d93..47ff090f0 100644 --- a/packages/ai/test/openrouter-oauth.test.ts +++ b/packages/ai/test/openrouter-oauth.test.ts @@ -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 | undefined; let manualSignal: AbortSignal | undefined; const credential = await openRouterOAuth.login({ + signal: neverAbortedSignal, prompt: (prompt) => { manualSignal = prompt.signal; return new Promise(() => {}); @@ -113,6 +115,7 @@ describe.sequential("OpenRouter OAuth", () => { let callbackResponse: Promise | undefined; const login = openRouterOAuth.login({ + signal: neverAbortedSignal, prompt: () => new Promise(() => {}), notify: (event) => { if (event.type !== "auth_url") return; @@ -141,6 +144,7 @@ describe.sequential("OpenRouter OAuth", () => { let callbackUrl: URL | undefined; let firstCallback: Promise | undefined; const login = openRouterOAuth.login({ + signal: neverAbortedSignal, prompt: () => new Promise(() => {}), notify: (event) => { if (event.type !== "auth_url") return; @@ -168,6 +172,7 @@ describe.sequential("OpenRouter OAuth", () => { let callbackResponse: Promise | undefined; const login = openRouterOAuth.login({ + signal: neverAbortedSignal, prompt: () => new Promise(() => {}), 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: () => {}, }), diff --git a/packages/ai/test/providers.test.ts b/packages/ai/test/providers.test.ts index c16b89cca..fda4422fa 100644 --- a/packages/ai/test/providers.test.ts +++ b/packages/ai/test/providers.test.ts @@ -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, 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((resolve) => { + markFirstStarted = resolve; + }); + const firstBlocked = new Promise((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"]); }); }); diff --git a/packages/ai/test/radius-oauth.test.ts b/packages/ai/test/radius-oauth.test.ts index 762c8c25d..ac613520d 100644 --- a/packages/ai/test/radius-oauth.test.ts +++ b/packages/ai/test/radius-oauth.test.ts @@ -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(); }); diff --git a/packages/ai/test/xai-oauth.test.ts b/packages/ai/test/xai-oauth.test.ts index 6e707ef84..08843f87f 100644 --- a/packages/ai/test/xai-oauth.test.ts +++ b/packages/ai/test/xai-oauth.test.ts @@ -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 { 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 { - 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", () => { diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index ec55235be..3390e6de9 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -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 ` 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 diff --git a/packages/coding-agent/docs/custom-provider.md b/packages/coding-agent/docs/custom-provider.md index 1b9b00504..027aa6115 100644 --- a/packages/coding-agent/docs/custom-provider.md +++ b/packages/coding-agent/docs/custom-provider.md @@ -331,8 +331,8 @@ pi.registerProvider("corporate-ai", { }; }, - async refreshToken(credentials: OAuthCredentials): Promise { - const tokens = await refreshAccessToken(credentials.refresh); + async refreshToken(credentials: OAuthCredentials, signal: AbortSignal): Promise { + 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; - refreshToken(credentials: OAuthCredentials): Promise; + refreshToken(credentials: OAuthCredentials, signal: AbortSignal): Promise; getApiKey(credentials: OAuthCredentials): string; }; } diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 9ec30321a..48adbba9c 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -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. diff --git a/packages/coding-agent/docs/sdk.md b/packages/coding-agent/docs/sdk.md index 1dc32a14f..edced9186 100644 --- a/packages/coding-agent/docs/sdk.md +++ b/packages/coding-agent/docs/sdk.md @@ -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 diff --git a/packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts b/packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts index 45304be42..ca17958cd 100644 --- a/packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts +++ b/packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts @@ -124,7 +124,7 @@ async function loginAnthropic(callbacks: OAuthLoginCallbacks): Promise { +async function refreshAnthropicToken(credentials: OAuthCredentials, signal: AbortSignal): Promise { const response = await fetch(TOKEN_URL, { method: "POST", headers: { "Content-Type": "application/json" }, @@ -133,6 +133,7 @@ async function refreshAnthropicToken(credentials: OAuthCredentials): Promise { +async function refreshGitLabToken(credentials: OAuthCredentials, signal: AbortSignal): Promise { 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 "You are helpful.", diff --git a/packages/coding-agent/src/cli/credential-print.ts b/packages/coding-agent/src/cli/credential-print.ts index e26287100..8f069781d 100644 --- a/packages/coding-agent/src/cli/credential-print.ts +++ b/packages/coding-agent/src/cli/credential-print.ts @@ -86,11 +86,12 @@ export async function resolveCredentialForPrint( modelRuntime: ModelRuntime, kind: CredentialPrintKind, minExpiryMs?: number, + signal?: AbortSignal, ): Promise { validateCredentialPrintArgs(args); const credentialTypes = new Map( - (await modelRuntime.listCredentials()).map((credential) => [credential.providerId, credential.type]), + (await modelRuntime.listCredentials({ signal })).map((credential) => [credential.providerId, credential.type]), ); const models: Model[] = []; 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]; diff --git a/packages/coding-agent/src/cli/list-models.ts b/packages/coding-agent/src/cli/list-models.ts index 10c1cf369..d24d1e845 100644 --- a/packages/coding-agent/src/cli/list-models.ts +++ b/packages/coding-agent/src/cli/list-models.ts @@ -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 { +export async function listModels( + modelRuntime: ModelRuntime, + searchPattern?: string, + signal?: AbortSignal, +): Promise { 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()); diff --git a/packages/coding-agent/src/core/agent-session-services.ts b/packages/coding-agent/src/core/agent-session-services.ts index 495154b29..531f51d61 100644 --- a/packages/coding-agent/src/core/agent-session-services.ts +++ b/packages/coding-agent/src/core/agent-session-services.ts @@ -39,6 +39,7 @@ export interface CreateAgentSessionServicesOptions { agentDir?: string; settingsManager?: SettingsManager; modelRuntime?: ModelRuntime; + modelRuntimeSignal?: AbortSignal; extensionFlagValues?: Map; resourceLoaderOptions?: Omit; 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({ diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 0b5ceff1e..f26df44d5 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1605,13 +1605,12 @@ export class AgentSession { } private async _cycleScopedModel(direction: "forward" | "backward"): Promise { - 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 { - const availableModels = await this._modelRuntime.getAvailable(); + const availableModels = this._modelRuntime.getAvailableSnapshot(); if (availableModels.length <= 1) return undefined; const currentModel = this.model; diff --git a/packages/coding-agent/src/core/auth-storage.ts b/packages/coding-agent/src/core/auth-storage.ts index 58c0eaa1a..0b7e96392 100644 --- a/packages/coding-agent/src/core/auth-storage.ts +++ b/packages/coding-agent/src/core/auth-storage.ts @@ -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(fn: (current: string | undefined) => LockResult): T; - withLockAsync(fn: (current: string | undefined) => Promise>): Promise; + withLockAsync( + fn: (current: string | undefined) => Promise>, + options?: AuthOperationOptions, + ): Promise; } export class FileAuthStorageBackend implements AuthStorageBackend { @@ -111,7 +116,47 @@ export class FileAuthStorageBackend implements AuthStorageBackend { } } - async withLockAsync(fn: (current: string | undefined) => Promise>): Promise { + private async acquireLockAsync( + signal: AbortSignal | undefined, + onCompromised: (error: Error) => void, + ): Promise<() => Promise> { + const maxRetries = 10; + for (let attempt = 0; attempt <= maxRetries; attempt++) { + signal?.throwIfAborted(); + let release: (() => Promise) | 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( + fn: (current: string | undefined) => Promise>, + options?: AuthOperationOptions, + ): Promise { + 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 = Promise.resolve(); withLock(fn: (current: string | undefined) => LockResult): T { const { result, next } = fn(this.value); @@ -173,12 +211,23 @@ export class InMemoryAuthStorageBackend implements AuthStorageBackend { return result; } - async withLockAsync(fn: (current: string | undefined) => Promise>): Promise { - const { result, next } = await fn(this.value); - if (next !== undefined) { - this.value = next; - } - return result; + withLockAsync( + fn: (current: string | undefined) => Promise>, + options?: AuthOperationOptions, + ): Promise { + 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 { + private async reloadFromStorageAsync(options?: AuthOperationOptions): Promise { 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 { + private readLatestData(options?: AuthOperationOptions): Promise { + 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 { - const credential = (await this.readLatestData())[provider]; + async read(provider: string, options?: AuthOperationOptions): Promise { + 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, + options?: AuthOperationOptions, ): Promise { 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 { + async delete(provider: string, options?: AuthOperationOptions): Promise { 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 { - return Object.entries(await this.readLatestData()).map(([providerId, credential]) => ({ - providerId, - type: credential.type, - })); + async list(options?: AuthOperationOptions): Promise { + const entries = Object.entries(await this.readLatestData(options)); + options?.signal?.throwIfAborted(); + return entries.map(([providerId, credential]) => ({ providerId, type: credential.type })); } } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index f8d13b199..0abb501f5 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -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; /** 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; /** Refresh expired credentials, return updated credentials to persist. */ - refreshToken(credentials: OAuthCredentials): Promise; + refreshToken(credentials: OAuthCredentials, signal: AbortSignal): Promise; /** Convert credentials to API key string for the provider. */ getApiKey(credentials: OAuthCredentials): string; /** Legacy synchronous credential-dependent model projection. */ diff --git a/packages/coding-agent/src/core/model-registry.ts b/packages/coding-agent/src/core/model-registry.ts index 8d6651af4..d01ad3a03 100644 --- a/packages/coding-agent/src/core/model-registry.ts +++ b/packages/coding-agent/src/core/model-registry.ts @@ -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 { - await this.runtime.refresh(); + refresh(options?: ModelsRefreshOptions): Promise { + return this.runtime.refresh(options); } getError(): string | undefined { diff --git a/packages/coding-agent/src/core/model-resolver.ts b/packages/coding-agent/src/core/model-resolver.ts index f27a89af3..7a35558e9 100644 --- a/packages/coding-agent/src/core/model-resolver.ts +++ b/packages/coding-agent/src/core/model-resolver.ts @@ -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 { - const availableModels = [...(await modelRuntime.getAvailable())]; + models: readonly Model[], +): 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 { - const { scopedModels, diagnostics } = await resolveModelScopeWithDiagnostics(patterns, modelRuntime); +export async function resolveModelScopeWithDiagnostics( + patterns: string[], + modelRuntime: ModelRuntime, + options?: AuthOperationOptions, +): Promise { + return resolveModelScopeFromModels(patterns, await modelRuntime.getAvailable(undefined, options)); +} + +export async function resolveModelScope( + patterns: string[], + modelRuntime: ModelRuntime, + options?: AuthOperationOptions, +): Promise { + 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 diff --git a/packages/coding-agent/src/core/model-runtime.ts b/packages/coding-agent/src/core/model-runtime.ts index c31102f5c..c530d29a6 100644 --- a/packages/coding-agent/src/core/model-runtime.ts +++ b/packages/coding-agent/src/core/model-runtime.ts @@ -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; /** 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 | undefined; private availabilityRefreshSeq = 0; + private availabilityErrorSeq = 0; + private readonly providerAvailabilitySeq = new Map(); private availabilityError: string | undefined; + private readonly credentialOperations = new Map>(); 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 { + private async runAvailabilityRefresh(seq: number, errorSeq: number, signal: AbortSignal): Promise { 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 { + private queueAvailabilityRefresh(signal?: AbortSignal): Promise { 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 { - 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 { - return this.queueAvailabilityRefresh(); + private async refreshProviderAvailability(providerId: string, signal: AbortSignal): Promise { + // 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 { - return this.models.checkAuth(providerId); + async checkAuth(providerId: string, options?: AuthOperationOptions): Promise { + return this.models.checkAuth(providerId, options); } - async getAvailable(providerId?: string): Promise[]> { + async getAvailable(providerId?: string, options?: AuthOperationOptions): Promise[]> { 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(providerId: string, signal: AbortSignal, task: () => Promise): Promise { + const previous = this.credentialOperations.get(providerId) ?? Promise.resolve(); + let markStarted: (() => void) | undefined; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + const operation = (async () => { + await previous.catch(() => {}); + signal.throwIfAborted(); + markStarted?.(); + return task(); + })(); + const tail = operation.catch(() => {}); + this.credentialOperations.set(providerId, tail); + void tail.then(() => { + if (this.credentialOperations.get(providerId) === tail) this.credentialOperations.delete(providerId); + }); + return raceWithAbortSignal(started, signal).then(() => operation); + } + + private async synchronizeCredentialState( providerId: string, - apiKey: string, - refreshOptions: ModelsRefreshOptions = {}, + operation: CredentialSynchronizationOperation, + credential: Credential | undefined, + signal: AbortSignal, ): Promise { - 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 { - this.credentials.removeRuntimeApiKey(providerId); - await this.refresh({ allowNetwork: this.modelNetworkEnabled }); + setRuntimeApiKey(providerId: string, apiKey: string, options: AuthOperationOptions = {}): Promise { + 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 { - return this.credentials.list(); + removeRuntimeApiKey(providerId: string, options: AuthOperationOptions = {}): Promise { + 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 { + return this.credentials.list(options); } getProviderAuthStatus(providerId: string): AuthStatus { @@ -459,7 +562,11 @@ export class ModelRuntime implements Models { ): Promise<{ provider: Provider; model: Model; 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 { - 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 { + 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 { - 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 { + 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 { 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 { diff --git a/packages/coding-agent/src/core/models-store.ts b/packages/coding-agent/src/core/models-store.ts index dca579952..649e54054 100644 --- a/packages/coding-agent/src/core/models-store.ts +++ b/packages/coding-agent/src/core/models-store.ts @@ -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; export class InMemoryCodingAgentModelsStore implements ModelsStore { private readonly entries = new Map(); - async read(providerId: string): Promise { - return this.entries.get(providerId); + async read(providerId: string, options?: ModelsStoreOperationOptions): Promise { + options?.signal?.throwIfAborted(); + const entry = this.entries.get(providerId); + return entry ? structuredClone(entry) : undefined; } - async write(providerId: string, entry: ModelsStoreEntry): Promise { - this.entries.set(providerId, entry); + async write(providerId: string, entry: ModelsStoreEntry, options?: ModelsStoreOperationOptions): Promise { + options?.signal?.throwIfAborted(); + this.entries.set(providerId, structuredClone(entry)); } - async delete(providerId: string): Promise { + async delete(providerId: string, options?: ModelsStoreOperationOptions): Promise { + 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 { - return this.storage.withLock((content) => ({ - result: structuredClone(this.parse(content)[providerId]), - })); + async read(providerId: string, options?: ModelsStoreOperationOptions): Promise { + return this.storage.withLockAsync( + async (content) => ({ result: structuredClone(this.parse(content)[providerId]) }), + options, + ); } - async write(providerId: string, entry: ModelsStoreEntry): Promise { + async write(providerId: string, entry: ModelsStoreEntry, options?: ModelsStoreOperationOptions): Promise { 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 { + async delete(providerId: string, options?: ModelsStoreOperationOptions): Promise { await this.storage.withLockAsync(async (content) => { const current = this.parse(content); delete current[providerId]; return { result: undefined, next: JSON.stringify(current, null, 2) }; - }); + }, options); } } diff --git a/packages/coding-agent/src/core/provider-composer.ts b/packages/coding-agent/src/core/provider-composer.ts index 5cbdea5fb..e05878329 100644 --- a/packages/coding-agent/src/core/provider-composer.ts +++ b/packages/coding-agent/src/core/provider-composer.ts @@ -35,7 +35,7 @@ export interface ExtensionOAuthConfig { /** @deprecated Retained for extension source compatibility; ignored by canonical auth flows. */ usesCallbackServer?: boolean; login(callbacks: OAuthLoginCallbacks): Promise; - refreshToken(credentials: OAuthCredentials): Promise; + refreshToken(credentials: OAuthCredentials, signal: AbortSignal): Promise; getApiKey(credentials: OAuthCredentials): string; modifyModels?(models: Model[], credentials: OAuthCredentials): Model[]; } @@ -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 | 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 diff --git a/packages/coding-agent/src/core/remote-catalog-provider.ts b/packages/coding-agent/src/core/remote-catalog-provider.ts index 56147035d..682c4ab5d 100644 --- a/packages/coding-agent/src/core/remote-catalog-provider.ts +++ b/packages/coding-agent/src/core/remote-catalog-provider.ts @@ -47,77 +47,85 @@ export function withRemoteCatalog( localGeneratedAt?: number, ): Provider { let dynamicModels: readonly Model[] = []; - let inflightRefresh: Promise | 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; + }, + }); }, }; } diff --git a/packages/coding-agent/src/core/runtime-credentials.ts b/packages/coding-agent/src/core/runtime-credentials.ts index 022b145e0..65f238c83 100644 --- a/packages/coding-agent/src/core/runtime-credentials.ts +++ b/packages/coding-agent/src/core/runtime-credentials.ts @@ -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 { + async read(providerId: string, options?: AuthOperationOptions): Promise { + 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 { - const entries = new Map((await this.store.list()).map((entry) => [entry.providerId, entry])); + async list(options?: AuthOperationOptions): Promise { + 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, + options?: AuthOperationOptions, ): Promise { - return this.store.modify(providerId, fn); + return this.store.modify(providerId, fn, options); } - async delete(providerId: string): Promise { + async delete(providerId: string, options?: AuthOperationOptions): Promise { + options?.signal?.throwIfAborted(); + await this.store.delete(providerId, options); this.overrides.delete(providerId); - await this.store.delete(providerId); } } diff --git a/packages/coding-agent/src/extensions/llama/index.ts b/packages/coding-agent/src/extensions/llama/index.ts index 5fc17b997..0cf72141f 100644 --- a/packages/coding-agent/src/extensions/llama/index.ts +++ b/packages/coding-agent/src/extensions/llama/index.ts @@ -48,9 +48,16 @@ export default function llamaExtension(pi: ExtensionAPI): void { client: LlamaClient, catalog?: LlamaModelInfo[], ): Promise => { - 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; }; diff --git a/packages/coding-agent/src/extensions/llama/provider.ts b/packages/coding-agent/src/extensions/llama/provider.ts index 64d45d7bc..518007cb4 100644 --- a/packages/coding-agent/src/extensions/llama/provider.ts +++ b/packages/coding-agent/src/extensions/llama/provider.ts @@ -111,20 +111,36 @@ export function createLlamaProvider(): LlamaProviderController { }, getModels: () => models, refreshModels: async (context: RefreshModelsContext): Promise => { - 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), diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 394851162..3d4ca9831 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -179,6 +179,8 @@ export { } from "./core/model-resolver.ts"; export { type CreateModelRuntimeOptions, + CredentialSynchronizationError, + type CredentialSynchronizationOperation, ModelRuntime, type ModelRuntimeAuthOverrides, } from "./core/model-runtime.ts"; diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index ca515d7a0..e737f3b47 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -155,8 +155,15 @@ async function runCredentialPrintCommand(args: string[]): Promise { 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") { diff --git a/packages/coding-agent/src/modes/interactive/components/model-selector.ts b/packages/coding-agent/src/modes/interactive/components/model-selector.ts index 7be56e3ab..6f2668a49 100644 --- a/packages/coding-agent/src/modes/interactive/components/model-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/model-selector.ts @@ -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): void { - this.close(); + this.dispose(); // Save as new default this.settingsManager.setDefaultModelAndProvider(model.provider, model.id); this.onSelectCallback(model); diff --git a/packages/coding-agent/src/modes/interactive/components/scoped-models-selector.ts b/packages/coding-agent/src/modes/interactive/components/scoped-models-selector.ts index fc5c61133..e7dfc2714 100644 --- a/packages/coding-agent/src/modes/interactive/components/scoped-models-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/scoped-models-selector.ts @@ -74,6 +74,7 @@ interface ModelItem { export interface ModelsConfig { allModels: Model[]; 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[], 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, diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index e0d67dd9b..141d04df4 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -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 => { - // 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 | undefined> { - const models = await this.getModelCandidates(); - return findExactModelReferenceMatch(searchTerm, models); - } - - private async getModelCandidates(): Promise[]> { - 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 { - // 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[]): 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 { - 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 { - 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 | undefined, ): Promise { - await this.session.modelRuntime.getAvailable(); - const actionLabel = authType === "oauth" ? `Logged in to ${providerName}` : `Saved API key for ${providerName}`; let selectedModel: Model | 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); } diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index a2e2e65d2..1b5d8a656 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -466,7 +466,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise 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 { - 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, diff --git a/packages/coding-agent/src/utils/abort.ts b/packages/coding-agent/src/utils/abort.ts new file mode 100644 index 000000000..cca32a11e --- /dev/null +++ b/packages/coding-agent/src/utils/abort.ts @@ -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(operation: Promise, signal: AbortSignal | undefined): Promise { + if (!signal) return operation; + if (signal.aborted) { + void operation.catch(() => {}); + return Promise.reject(abortReason(signal)); + } + + return new Promise((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(); + }); +} diff --git a/packages/coding-agent/test/auth-storage.test.ts b/packages/coding-agent/test/auth-storage.test.ts index 51464bb1b..67cb5edf3 100644 --- a/packages/coding-agent/test/auth-storage.test.ts +++ b/packages/coding-agent/test/auth-storage.test.ts @@ -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((resolve) => { + markStarted = resolve; + }); + const blocked = new Promise((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((resolve) => { + markStarted = resolve; + }); + const blocked = new Promise((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((resolve) => { + markStarted = resolve; + }); + const blocked = new Promise((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((resolve) => { + markStarted = resolve; + }); + const blocked = new Promise((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({ diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 115b2df52..fd4cccc77 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -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", () => { diff --git a/packages/coding-agent/test/llama-extension.test.ts b/packages/coding-agent/test/llama-extension.test.ts index 3af25e563..a89a0940e 100644 --- a/packages/coding-agent/test/llama-extension.test.ts +++ b/packages/coding-agent/test/llama-extension.test.ts @@ -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 => { + 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", diff --git a/packages/coding-agent/test/model-registry.test.ts b/packages/coding-agent/test/model-registry.test.ts index bd040101f..e492bc622 100644 --- a/packages/coding-agent/test/model-registry.test.ts +++ b/packages/coding-agent/test/model-registry.test.ts @@ -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(); }); diff --git a/packages/coding-agent/test/model-resolver.test.ts b/packages/coding-agent/test/model-resolver.test.ts index 6840a282a..508c949ba 100644 --- a/packages/coding-agent/test/model-resolver.test.ts +++ b/packages/coding-agent/test/model-resolver.test.ts @@ -743,7 +743,7 @@ describe("default model selection", () => { }; const registry = { - getAvailable: async () => [aiGatewayModel], + getAvailableSnapshot: () => [aiGatewayModel], } as unknown as Parameters[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[0]["modelRuntime"]; const result = await findInitialModel({ diff --git a/packages/coding-agent/test/model-runtime-auth-options.test.ts b/packages/coding-agent/test/model-runtime-auth-options.test.ts index 74f65be45..387e22cef 100644 --- a/packages/coding-agent/test/model-runtime-auth-options.test.ts +++ b/packages/coding-agent/test/model-runtime-auth-options.test.ts @@ -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", { diff --git a/packages/coding-agent/test/model-runtime-credential-sync.test.ts b/packages/coding-agent/test/model-runtime-credential-sync.test.ts new file mode 100644 index 000000000..8f38a9c18 --- /dev/null +++ b/packages/coding-agent/test/model-runtime-credential-sync.test.ts @@ -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; + 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 { + 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((resolve) => { + markLoginStarted = resolve; + }); + const blockedLogin = new Promise((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((resolve) => { + firstStarted = resolve; + }); + const second = new Promise((resolve) => { + secondStarted = resolve; + }); + const blocked = new Promise((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(() => {}); + 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((resolve) => { + markStarted = resolve; + }); + const registered = provider("cancelled-availability"); + if (registered.auth.apiKey) { + registered.auth.apiKey.check = async ({ credential }) => { + if (blockAvailability) { + markStarted?.(); + await new Promise(() => {}); + } + 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(() => {})); + 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((resolve) => { + markStarted = resolve; + }); + const blocked = new Promise((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((resolve) => { + markCommitted = resolve; + }); + const mutationFinished = new Promise((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((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(() => {}); + } + }, + }), + 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" }); + }); +}); diff --git a/packages/coding-agent/test/models-store.test.ts b/packages/coding-agent/test/models-store.test.ts index 483153974..2d42211e1 100644 --- a/packages/coding-agent/test/models-store.test.ts +++ b/packages/coding-agent/test/models-store.test.ts @@ -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; + expect(stored.one).toBeDefined(); + expect(stored.two).toBeUndefined(); + }); }); diff --git a/packages/coding-agent/test/radius.test.ts b/packages/coding-agent/test/radius.test.ts index 630bcecf2..cec1a15e8 100644 --- a/packages/coding-agent/test/radius.test.ts +++ b/packages/coding-agent/test/radius.test.ts @@ -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( diff --git a/packages/coding-agent/test/remote-catalog-provider.test.ts b/packages/coding-agent/test/remote-catalog-provider.test.ts index f4fecd853..8b47d0d01 100644 --- a/packages/coding-agent/test/remote-catalog-provider.test.ts +++ b/packages/coding-agent/test/remote-catalog-provider.test.ts @@ -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> = {}, +): Promise { + const publish = async (publication: ModelsPublication): Promise => { + 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((resolve) => { + markFirstStarted = resolve; + }); + const firstResponse = new Promise((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) }); }); diff --git a/packages/coding-agent/test/runtime-credentials.test.ts b/packages/coding-agent/test/runtime-credentials.test.ts index a044f715e..fe50503cb 100644 --- a/packages/coding-agent/test/runtime-credentials.test.ts +++ b/packages/coding-agent/test/runtime-credentials.test.ts @@ -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); diff --git a/packages/coding-agent/test/suite/regressions/6949-unavailable-scoped-model.test.ts b/packages/coding-agent/test/suite/regressions/6949-unavailable-scoped-model.test.ts index 212235406..a2630b300 100644 --- a/packages/coding-agent/test/suite/regressions/6949-unavailable-scoped-model.test.ts +++ b/packages/coding-agent/test/suite/regressions/6949-unavailable-scoped-model.test.ts @@ -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 { @@ -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 () => { diff --git a/packages/coding-agent/test/suite/regressions/7027-credential-refresh-hang.test.ts b/packages/coding-agent/test/suite/regressions/7027-credential-refresh-hang.test.ts new file mode 100644 index 000000000..d1e7f8404 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/7027-credential-refresh-hang.test.ts @@ -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((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(() => {}); + }, + 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, + ) => Promise; + + 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.", + ); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/7153-scoped-models-refresh.test.ts b/packages/coding-agent/test/suite/regressions/7153-scoped-models-refresh.test.ts new file mode 100644 index 000000000..0fb55a306 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/7153-scoped-models-refresh.test.ts @@ -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[]) { + 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[], 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(); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/7301-stalled-availability-refresh.test.ts b/packages/coding-agent/test/suite/regressions/7301-stalled-availability-refresh.test.ts index 356f5b939..fc424d9ff 100644 --- a/packages/coding-agent/test/suite/regressions/7301-stalled-availability-refresh.test.ts +++ b/packages/coding-agent/test/suite/regressions/7301-stalled-availability-refresh.test.ts @@ -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(); + }); }); diff --git a/packages/coding-agent/test/suite/regressions/7443-model-command-cached-match.test.ts b/packages/coding-agent/test/suite/regressions/7443-model-command-cached-match.test.ts new file mode 100644 index 000000000..288bb8264 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/7443-model-command-cached-match.test.ts @@ -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 | 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…"); + }); +}); diff --git a/packages/coding-agent/test/utilities.ts b/packages/coding-agent/test/utilities.ts index 27053ef56..41641db7c 100644 --- a/packages/coding-agent/test/utilities.ts +++ b/packages/coding-agent/test/utilities.ts @@ -93,7 +93,7 @@ export async function resolveApiKey(provider: string): Promise= credential.expires) { - credential = await oauth.refresh(credential); + credential = await oauth.refresh(credential, new AbortController().signal); storage[provider] = credential; saveAuthStorage(storage); }