From 9993c96907bb0c97260d2c353c31a3464f211122 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 14 Jul 2026 17:48:45 +0200 Subject: [PATCH] feat(coding-agent): replace model registry with model runtime Move provider auth and OAuth flows onto pi-ai Models, compose models.json and extension overlays through ModelRuntime, and retain ModelRegistry as an extension compatibility facade. --- packages/agent/docs/models.md | 40 +- packages/ai/CHANGELOG.md | 47 + packages/ai/README.md | 87 +- packages/ai/src/api/lazy.ts | 25 +- packages/ai/src/auth/credential-store.ts | 6 +- packages/ai/src/auth/helpers.ts | 6 +- .../ai/src/{utils => auth}/oauth/anthropic.ts | 188 +-- .../src/{utils => auth}/oauth/device-code.ts | 0 .../{utils => auth}/oauth/github-copilot.ts | 130 +-- packages/ai/src/{utils => auth}/oauth/load.ts | 2 +- .../src/{utils => auth}/oauth/oauth-page.ts | 0 .../src/{utils => auth}/oauth/openai-codex.ts | 252 +--- packages/ai/src/{utils => auth}/oauth/pkce.ts | 0 packages/ai/src/auth/resolve.ts | 20 +- packages/ai/src/auth/types.ts | 61 +- packages/ai/src/cli.ts | 155 +-- packages/ai/src/compat.ts | 28 +- .../ai/src/compat/extension-oauth-types.ts | 45 + packages/ai/src/images-models.ts | 21 +- packages/ai/src/index.ts | 21 +- packages/ai/src/models.ts | 231 +++- packages/ai/src/oauth.ts | 11 +- .../ai/src/providers/amazon-bedrock.models.ts | 54 + packages/ai/src/providers/amazon-bedrock.ts | 55 +- packages/ai/src/providers/anthropic.ts | 2 +- .../azure-openai-responses.models.ts | 17 + packages/ai/src/providers/cerebras.models.ts | 2 +- .../providers/cloudflare-ai-gateway.models.ts | 72 ++ .../ai/src/providers/cloudflare-ai-gateway.ts | 7 +- packages/ai/src/providers/cloudflare-auth.ts | 41 +- .../ai/src/providers/cloudflare-stream.ts | 28 + .../ai/src/providers/cloudflare-workers-ai.ts | 3 +- .../ai/src/providers/github-copilot.models.ts | 59 +- packages/ai/src/providers/github-copilot.ts | 11 +- packages/ai/src/providers/google-vertex.ts | 65 +- packages/ai/src/providers/openai-codex.ts | 2 +- packages/ai/src/providers/openai.models.ts | 17 + packages/ai/src/providers/opencode.models.ts | 56 +- .../ai/src/providers/openrouter.models.ts | 140 +-- .../src/providers/vercel-ai-gateway.models.ts | 174 +-- packages/ai/src/utils/oauth/index.ts | 160 --- packages/ai/src/utils/oauth/types.ts | 79 -- packages/ai/test/anthropic-oauth.test.ts | 23 +- packages/ai/test/cloudflare-stream.test.ts | 65 ++ .../ai/test/codex-websocket-cached-probe.ts | 7 +- packages/ai/test/github-copilot-oauth.test.ts | 63 +- packages/ai/test/images-models.test.ts | 8 +- packages/ai/test/models-runtime.test.ts | 131 ++- packages/ai/test/oauth-auth.test.ts | 16 +- packages/ai/test/oauth-device-code.test.ts | 2 +- packages/ai/test/oauth.ts | 32 +- packages/ai/test/openai-codex-oauth.test.ts | 84 +- .../openai-responses-copilot-provider.test.ts | 54 +- packages/ai/test/providers.test.ts | 107 +- packages/ai/test/scratch.ts | 2 +- packages/coding-agent/CHANGELOG.md | 65 ++ packages/coding-agent/README.md | 8 +- packages/coding-agent/docs/custom-provider.md | 17 +- packages/coding-agent/docs/sdk.md | 92 +- .../custom-provider-anthropic/index.ts | 2 +- .../examples/sdk/02-custom-model.ts | 13 +- .../examples/sdk/09-api-keys-and-oauth.ts | 44 +- .../examples/sdk/12-full-control.ts | 19 +- packages/coding-agent/examples/sdk/README.md | 32 +- packages/coding-agent/src/cli/list-models.ts | 8 +- .../src/core/agent-session-services.ts | 25 +- .../coding-agent/src/core/agent-session.ts | 103 +- .../coding-agent/src/core/auth-storage.ts | 368 +----- packages/coding-agent/src/core/cache-stats.ts | 6 +- .../src/core/extensions/loader.ts | 13 +- .../src/core/extensions/runner.ts | 4 + .../coding-agent/src/core/extensions/types.ts | 4 +- .../coding-agent/src/core/model-config.ts | 277 +++++ .../coding-agent/src/core/model-registry.ts | 1010 +---------------- .../coding-agent/src/core/model-resolver.ts | 40 +- .../coding-agent/src/core/model-runtime.ts | 489 ++++++++ .../src/core/provider-composer.ts | 513 +++++++++ .../src/core/provider-display-names.ts | 35 - .../src/core/runtime-credentials.ts | 48 + packages/coding-agent/src/core/sdk.ts | 56 +- packages/coding-agent/src/index.ts | 17 +- packages/coding-agent/src/main.ts | 26 +- .../modes/interactive/components/footer.ts | 2 +- .../interactive/components/login-dialog.ts | 25 +- .../interactive/components/model-selector.ts | 16 +- .../interactive/components/oauth-selector.ts | 49 +- .../src/modes/interactive/interactive-mode.ts | 327 +++--- .../coding-agent/src/modes/rpc/rpc-mode.ts | 4 +- ...gent-session-auto-compaction-queue.test.ts | 10 +- .../test/agent-session-branching.test.ts | 2 +- .../test/agent-session-compaction.test.ts | 20 +- .../test/agent-session-concurrent.test.ts | 42 +- .../agent-session-dynamic-provider.test.ts | 9 +- .../test/agent-session-retry.test.ts | 36 +- .../test/agent-session-runtime-events.test.ts | 27 +- .../test/agent-session-stats.test.ts | 24 +- .../agent-session-tree-navigation.test.ts | 8 +- .../coding-agent/test/auth-storage.test.ts | 826 +++----------- .../coding-agent/test/cache-stats.test.ts | 2 +- .../test/compaction-extensions.test.ts | 31 +- .../test/config-value-migration.test.ts | 9 +- .../test/extensions-discovery.test.ts | 36 + .../test/extensions-input-event.test.ts | 5 +- .../test/extensions-runner.test.ts | 11 +- .../coding-agent/test/footer-width.test.ts | 2 +- ...interactive-mode-anthropic-warning.test.ts | 57 +- .../test/interactive-mode-status.test.ts | 8 +- .../coding-agent/test/model-registry.test.ts | 388 ++++--- .../coding-agent/test/model-resolver.test.ts | 116 +- .../test/model-runtime-auth-options.test.ts | 256 +++++ .../model-runtime-cloudflare-compat.test.ts | 95 ++ .../test/model-runtime-test-utils.ts | 25 + .../coding-agent/test/oauth-selector.test.ts | 170 +-- .../test/resolve-config-value.test.ts | 121 ++ .../coding-agent/test/resource-loader.test.ts | 7 +- .../rpc-prompt-response-semantics.test.ts | 14 +- .../test/runtime-credentials.test.ts | 42 + .../test/sdk-codex-cache-probe-tool-loop.ts | 8 +- .../test/sdk-openrouter-attribution.test.ts | 29 +- .../test/sdk-stream-options.test.ts | 49 +- .../test/suite/agent-session-runtime.test.ts | 8 +- packages/coding-agent/test/suite/harness.ts | 8 +- ...753-reload-stale-resource-settings.test.ts | 9 +- .../2860-replaced-session-context.test.ts | 9 +- .../3217-scoped-model-order.test.ts | 12 +- .../5433-extension-oauth-prompt-input.test.ts | 14 + .../5596-missing-theme-export.test.ts | 8 +- .../5661-uppercase-header-values.test.ts | 4 +- .../coding-agent/test/test-harness.test.ts | 28 +- packages/coding-agent/test/test-harness.ts | 35 +- packages/coding-agent/test/utilities.ts | 37 +- packages/coding-agent/vitest.config.ts | 2 + packages/orchestrator/src/radius.ts | 13 +- 133 files changed, 5103 insertions(+), 4340 deletions(-) rename packages/ai/src/{utils => auth}/oauth/anthropic.ts (67%) rename packages/ai/src/{utils => auth}/oauth/device-code.ts (100%) rename packages/ai/src/{utils => auth}/oauth/github-copilot.ts (71%) rename packages/ai/src/{utils => auth}/oauth/load.ts (94%) rename packages/ai/src/{utils => auth}/oauth/oauth-page.ts (100%) rename packages/ai/src/{utils => auth}/oauth/openai-codex.ts (67%) rename packages/ai/src/{utils => auth}/oauth/pkce.ts (100%) create mode 100644 packages/ai/src/compat/extension-oauth-types.ts create mode 100644 packages/ai/src/providers/cloudflare-stream.ts delete mode 100644 packages/ai/src/utils/oauth/index.ts delete mode 100644 packages/ai/src/utils/oauth/types.ts create mode 100644 packages/ai/test/cloudflare-stream.test.ts create mode 100644 packages/coding-agent/src/core/model-config.ts create mode 100644 packages/coding-agent/src/core/model-runtime.ts create mode 100644 packages/coding-agent/src/core/provider-composer.ts delete mode 100644 packages/coding-agent/src/core/provider-display-names.ts create mode 100644 packages/coding-agent/src/core/runtime-credentials.ts create mode 100644 packages/coding-agent/test/model-runtime-auth-options.test.ts create mode 100644 packages/coding-agent/test/model-runtime-cloudflare-compat.test.ts create mode 100644 packages/coding-agent/test/model-runtime-test-utils.ts create mode 100644 packages/coding-agent/test/resolve-config-value.test.ts create mode 100644 packages/coding-agent/test/runtime-credentials.test.ts diff --git a/packages/agent/docs/models.md b/packages/agent/docs/models.md index 15914371..64a623c1 100644 --- a/packages/agent/docs/models.md +++ b/packages/agent/docs/models.md @@ -68,7 +68,7 @@ packages/ai/src/ openrouter-images.ts # image-generation provider factory faux.ts # test provider factory all.ts # explicit aggregate: builtinModels(), builtinImagesModels(), getBuiltin*() - utils/oauth/ # OAuth flow implementations (node), lazy-loaded + auth/oauth/ # Canonical OAuth implementations (node), lazy-loaded ``` `src/index.ts` must stay core-only. It must not import: @@ -407,7 +407,7 @@ export interface ApiKeyAuth { name: string; // "Anthropic API key" /** Interactive setup (prompt for key/provider env). Absent = ambient-only (env, ADC, IAM). */ - login?(callbacks: AuthLoginCallbacks): Promise; + login?(interaction: AuthInteraction): Promise; /** * Resolve auth from the stored credential and/or ambient sources, merging @@ -424,7 +424,7 @@ export interface ApiKeyAuth { export interface OAuthAuth { name: string; // "Anthropic (Claude Pro/Max)" - login(callbacks: AuthLoginCallbacks): Promise; + login(interaction: AuthInteraction): Promise; /** Exchange the refresh token. Network call; throws on failure (invalid_grant etc.). Runs under the store lock. */ refresh(credential: OAuthCredential): Promise; @@ -445,9 +445,9 @@ export interface AuthContext { } ``` -The OAuth split (`refresh` + `toAuth` instead of one `resolve`) matches the old `OAuthProviderInterface` (`refreshToken` + `getApiKey`) and lets `Models` own the locking pattern without closure gymnastics: refresh produces a credential, `toAuth` derives request auth from whatever credential ends up stored. +The `refresh`/`toAuth` split lets `Models` own the locked refresh pattern without closure gymnastics: refresh produces a credential, while `toAuth` derives request auth from whatever credential ends up stored. -There is no `usesCallbackServer` flag. With `prompt()/notify()` callbacks the flow self-describes at runtime: a flow that runs a callback server issues a `manual_code` prompt racing the server and aborts the prompt when the callback wins. The UI needs no static foreknowledge. +OAuth implementations use the provider-neutral `AuthInteraction` protocol directly. A callback-server flow issues a `manual_code` prompt racing the server and aborts the prompt when the callback wins, so the UI needs no provider-specific callback or static callback-server flag. ### Credentials @@ -564,7 +564,7 @@ FileCredentialStore ports AuthStorage's lock backend: read = memory snaps └─ withRuntimeOverrides --api-key └─ createModels({ credentials: store }) -login/logout UI provider.auth.{oauth,apiKey}.login(callbacks) + store.modify/delete +login/logout UI provider.auth.{oauth,apiKey}.login(interaction) + store.modify/delete status UI store.read(pid) + getAuth try/catch ("needs /login" on rejection) getOAuthProviders presence of provider.auth.oauth across registered providers ``` @@ -574,7 +574,7 @@ getOAuthProviders presence of provider.auth.oauth across registered pro One interface serves api-key and OAuth login: ```ts -export interface AuthLoginCallbacks { +export interface AuthInteraction { /** Aborts the whole login flow. Per-prompt cancellation uses AuthPrompt.signal. */ signal?: AbortSignal; @@ -612,7 +612,7 @@ export function anthropicProvider(): Provider { apiKey: envApiKeyAuth("Anthropic API key", ["ANTHROPIC_API_KEY"]), oauth: lazyOAuth({ name: "Anthropic (Claude Pro/Max)", - load: () => import("../utils/oauth/anthropic.ts").then((m) => m.anthropicOAuth), + load: () => import("../auth/oauth/anthropic.ts").then((m) => m.anthropicOAuth), }), }, models: ANTHROPIC_MODELS, @@ -632,7 +632,7 @@ export function lazyOAuth(input: { OAuth must not force Node-only code (`node:http`, `node:crypto`) into browser bundles: the dynamic import inside `lazyOAuth()` uses the same bundler-opaque variable-specifier trick as the bedrock lazy wrapper. Browser hosts never trigger the load (no stored node OAuth credentials, no login flow). If web OAuth lands later (sitegeist proved feasibility: Web Crypto PKCE, auth tab, fetch token exchange, device-code polling), it is just a different `OAuthAuth` implementation — no reserved option values. -The existing flows in `src/utils/oauth/` (anthropic, openai-codex, github-copilot) are adapted to `OAuthAuth` (`login`/`refresh`/`toAuth`, replacing `login`/`refreshToken`/`getApiKey`/`modifyModels`) with the new callbacks, staying Node-targeted and lazy-loaded. Copilot's `modifyModels` baseUrl rewriting becomes `toAuth` returning `ModelAuth.baseUrl`. +The built-in flows in `src/auth/oauth/` implement `OAuthAuth` and `AuthInteraction` directly while remaining Node-targeted and lazy-loaded. Copilot derives its credential-specific request endpoint through `toAuth().baseUrl`. ## Provider wrappers and models.json @@ -817,7 +817,7 @@ Check items off as they land. Keep this list current; it is the working state fo ### Phase 3 — provider factories + catalogs -- [x] Auth helpers in `src/auth/helpers.ts`: `envApiKeyAuth()` (with secret-prompt `login`), `lazyOAuth()`. OAuth flow loads go through `utils/oauth/load.ts` (bundler-opaque dynamic import); the `OAuthAuth` exports it references land in Phase 4. +- [x] Auth helpers in `src/auth/helpers.ts`: `envApiKeyAuth()` (with secret-prompt `login`), `lazyOAuth()`. OAuth flow loads go through `auth/oauth/load.ts` (bundler-opaque dynamic import); the `OAuthAuth` exports it references land in Phase 4. - [x] `createProvider()` in `models.ts` (single + mixed `api` map, dispatch on `model.api`, unknown api -> stream error). - [x] Per-provider factories under `src/providers/` for all built-in catalog providers; OAuth attached via `lazyOAuth()` (anthropic, openai-codex, github-copilot); ambient `ApiKeyAuth` for amazon-bedrock (AWS env/profile) and google-vertex (key or ADC+project+location). - [x] `providers/all.ts`: `builtinProviders()`, `builtinModels()`, `getBuiltinModel/getBuiltinModels/getBuiltinProviders` re-exports. @@ -826,8 +826,8 @@ Check items off as they land. Keep this list current; it is the working state fo ### Phase 4 — OAuth adaptation -- [x] Adapt `utils/oauth/anthropic.ts`, `openai-codex.ts`, `github-copilot.ts` to `OAuthAuth` (`login`/`refresh`/`toAuth`) + `prompt()/notify()`; `modifyModels` baseUrl rewriting becomes `toAuth().baseUrl`. New exports (`anthropicOAuth`, `openaiCodexOAuth`, `githubCopilotOAuth`) sit next to the old `OAuthProviderInterface` objects, which survive until Phase 7. -- [x] No `usesCallbackServer` on `OAuthAuth`: callback-server flows race a `manual_code` prompt (aborted via `AuthPrompt.signal` once the flow settles). The old interface keeps its flag until it dies with compat. +- [x] Built-in implementations live under `auth/oauth/` and implement `OAuthAuth` directly through `AuthInteraction.prompt()`/`notify()`. They are private provider implementations loaded lazily by provider factories. +- [x] Callback-server flows race a `manual_code` prompt, aborted through `AuthPrompt.signal` once the flow settles. The public `oauth` subpath retains only coding-agent extension compatibility types. ### Phase 5 — packaging @@ -876,7 +876,7 @@ Decisions: - Runtime `--api-key` overrides are an explicit store overlay (an override reads as an ephemeral stored api-key credential, masking stored OAuth — matches today's priority). Every registered provider is guaranteed an `apiKey` auth slot so overrides apply to OAuth-only providers too. - `ModelRegistry.getAll`/`find`/`getAvailable` stay sync for SDK and extension compatibility, delegating to the collection's last-known sync model lists and fast configured-looking status checks. Dynamic providers update through explicit async `refresh()`, and request auth remains async through `getApiKeyAndHeaders()`/`Models.getAuth()`. Extensions also get the collection itself as the forward API. - models.json keeps FULL feature parity, implemented as provider decoration: builtin factories wrapped so `getModels()` applies provider `baseUrl`/`compat` overlays, `modelOverrides`, and custom-model merges (async-safe); provider `apiKey`/`headers`/`authHeader` configs become that provider's `ApiKeyAuth` (config first, factory auth fallback); parse errors keep `getError()` semantics. -- Extension `ProviderConfig` parity: provider-keyed `streamSimple`, old-style `oauth` adapted to `OAuthAuth` (`modifyModels` -> `getModels` wrap + `toAuth`), full model replacement per provider. Legacy `registerApiProvider` writes stay compat-local for consumers that call global `complete()`; they die with compat. +- Extension `ProviderConfig` parity: provider-keyed `streamSimple`, legacy extension OAuth callbacks adapted to `OAuthAuth`, and full model replacement per provider. Legacy `registerApiProvider` writes stay compat-local for consumers that call global `complete()`; they die with compat. - Copilot: stored-credential baseUrl applied in the wrapped `getModels()` (extension-visible models stay correct) plus per-request `toAuth().baseUrl`. - Cloudflare: provider-auth substitution (key + `CLOUDFLARE_ACCOUNT_ID`/`CLOUDFLARE_GATEWAY_ID` from credential `env` or ambient `AuthContext.env()` -> `ModelAuth.baseUrl`). Built-in compat calls route through `Models`, so they use the same provider auth path. @@ -909,13 +909,11 @@ Ordering for new sessions: - Wrap Copilot's provider `getModels()` when an OAuth credential is present so extension/UI-visible model metadata also carries the authenticated account base URL. - Keep API-key/env-token Copilot behavior unchanged. - Add tests for model metadata before login, after OAuth credential, after refresh/baseUrl change, and logout. -7. [ ] Extension OAuth adapter. - - Adapt old extension `OAuthProviderInterface` configs to pi-ai `OAuthAuth`. - - `login` maps old callbacks/events to `prompt()/notify()`. - - `refreshToken` maps to `refresh`. - - `getApiKey` maps to `toAuth`. - - `modifyModels` becomes a provider `getModels()` wrapper plus `toAuth().baseUrl` where applicable. - - Preserve existing extension runtime compatibility through the `/compat` alias until Phase 10. +7. [x] Extension OAuth adapter. + - Keep only the legacy callback/credential declarations required by coding-agent `ProviderConfig.oauth`. + - `login` maps legacy callbacks/events to `AuthInteraction.prompt()`/`notify()`. + - `refreshToken` maps to `refresh`; `getApiKey` maps to `toAuth`. + - Preserve the type-only pi-ai `oauth` barrel and extension-loader aliases. 8. [ ] Rebuild coding-agent `ModelRegistry` over `MutableModels`. - It owns a `MutableModels` instance built from decorated built-ins + models.json custom providers + extension providers. - `getAll()`, `find()`, and `getAvailable()` remain sync compatibility methods over last-known model lists and fast configured-looking auth status. Do not break the extension-facing `modelRegistry` surface for these reads. @@ -939,7 +937,7 @@ Ordering for new sessions: - [ ] AgentSession -> AgentHarness; the registry facade dies in favor of harness `Models`. - [ ] Move ALL internal `/compat` imports to the new API: every package's src, all tests, and the example extensions (examples then demonstrate the new API). Nothing inside the repo may import `/compat` at that point. -- [ ] Delete `/compat`, `env-api-keys.ts`, the extension-loader root-to-compat alias, the old `pi-ai/oauth` registry and `OAuthProviderInterface` (incl. `usesCallbackServer`), and the compat-local legacy API registry. This is the extension-author breaking release; changelog carries the migration guide. +- [ ] Delete `/compat`, `env-api-keys.ts`, the extension-loader root-to-compat alias, and the compat-local legacy API registry. The old OAuth registry/provider interface is already gone; the type-only `oauth` barrel remains for extension compatibility. ### Deferred / follow-ups diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 124f0f53..dc8b0ccd 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,13 +2,60 @@ ## [Unreleased] +### Breaking Changes + +- Changed runtime authentication to provider-scoped `Models.checkAuth()`, `getAuth()`, `login()`, and `logout()` APIs. `checkAuth()` now returns `AuthCheck | undefined`, and API-key auth resolvers no longer receive a model. +- Removed the legacy built-in OAuth provider objects, global OAuth registry APIs, and public low-level built-in login/refresh functions. Use canonical `Provider.auth.oauth` methods instead; the `oauth` subpath now retains only extension compatibility types. +- Renamed the canonical login interaction interface from `AuthLoginCallbacks` to `AuthInteraction`; it exposes the provider-neutral `prompt()`/`notify()` protocol used by API-key and OAuth flows. +- Changed the `Models` request contract: `getAuth(model)` now includes model headers, while `getAuth(providerId)` remains provider-scoped, and Models stream options may include `transformHeaders`. Custom `Models` implementations must execute the transform after merging auth/model and explicit headers, then remove it before provider dispatch. + +#### Models migration + +Use the `getAuth()` argument to choose the desired scope: + +```typescript +// Provider auth only +const providerAuth = await models.getAuth(model.provider); + +// Provider auth plus model.headers +const modelAuth = await models.getAuth(model); +``` + +Use the Models-only `transformHeaders` option instead of resolving auth before streaming. It runs once on assembled headers and is not passed to `Provider.stream*()`: + +```typescript +models.streamSimple(model, context, { + transformHeaders: async (headers) => ({ + ...headers, + "X-Request-ID": requestId, + }), +}); +``` + +Custom `Models` implementations must apply the same order: + +```text +getAuth(model) -> explicit options.headers -> transformHeaders -> Provider.stream*() +``` + +`Provider.stream*()` continues to accept ordinary `ApiStreamOptions`/`SimpleStreamOptions`; providers do not handle `transformHeaders`. + ### Added +- Added provider-owned authentication and availability resolution to `Models`, including stored OAuth refresh and interactive login support through `CredentialStore`. +- Added async non-secret credential enumeration through `CredentialStore.list()` and credential-aware `Provider.filterModels()` availability policy. +- Added neutral auth-flow information/link events and provider-owned Amazon Bedrock and Google Vertex AI credential selection flows. - Added a separate opt-in `max` thinking level, including native `xhigh` and `max` support for GPT-5.6 and Anthropic adaptive-thinking effort metadata matching Anthropic's documentation: `max` on all adaptive Claude models, native `xhigh` on Opus 4.7/4.8, Sonnet 5, and Fable 5 only. - Added request-wide input-token pricing tiers to model cost metadata and usage cost calculation. +### Changed + +- Changed `Models.getAuth(model)` to include model headers and added a Models-only `transformHeaders` stream option that runs after auth and explicit header assembly but is not forwarded to providers. + ### Fixed +- Fixed Cloudflare Workers AI and AI Gateway streams to materialize account and gateway endpoint placeholders after auth resolution, including compat streaming with custom model objects. +- Fixed lazy provider streams to preserve their final assistant message when forwarding an inner stream. - Fixed post-compaction output-token budgeting to ignore stale assistant usage from before the compaction boundary ([#6464](https://github.com/earendil-works/pi/issues/6464)). - Fixed GPT-5.4 and GPT-5.5 long-context cost accounting while retaining the intentional 272K default context limit for models that require an explicit override. - Fixed GPT-5.6 metadata to keep direct OpenAI requests in the 272K short-context tier while exposing the Codex backend's 372K context window with long-context pricing. diff --git a/packages/ai/README.md b/packages/ai/README.md index c9db20ac..c58fb3a1 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -17,6 +17,7 @@ Unified LLM API with provider collections, automatic auth resolution, token and - [Dynamic Providers](#dynamic-providers) - [Auth](#auth) - [How Auth Resolves](#how-auth-resolves) + - [Transforming Request Headers](#transforming-request-headers) - [Credential Store](#credential-store) - [Environment Variables](#environment-variables) - [Tools](#tools) @@ -334,18 +335,45 @@ await models.complete(model, context); await models.complete(model, context, { apiKey: 'sk-explicit' }); ``` -You can inspect resolution without making a request — useful for status UIs: +You can inspect resolution without making a request. Pass a provider ID for provider-scoped auth, or a model to include its static `model.headers`: ```typescript -const auth = await models.getAuth(model); -if (auth) { - console.log(`configured via ${auth.source}`); // e.g. "ANTHROPIC_API_KEY", "OAuth", "stored credential" +const providerAuth = await models.getAuth(model.provider); +const modelAuth = await models.getAuth(model); + +if (modelAuth) { + console.log(`configured via ${modelAuth.source}`); // e.g. "ANTHROPIC_API_KEY", "OAuth", "stored credential" + console.log(modelAuth.auth.headers); // Provider auth headers + model.headers } else { console.log('not configured'); } ``` -`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. +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. + +### 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: + +```typescript +const response = await models.completeSimple(model, context, { + headers: { "X-Client": "my-app" }, + transformHeaders: async (headers) => ({ + ...headers, + "X-Request-ID": crypto.randomUUID(), + }), +}); +``` + +The ordering is: + +```text +provider auth headers -> model.headers -> explicit options.headers -> transformHeaders -> Provider.stream*() +``` + +Header names are merged case-insensitively. Explicit headers override auth/model headers, and the transform has final control; returning `null` for a header suppresses lower-level defaults that support deletion. + +`transformHeaders` belongs to `Models`, not `Provider`. A `Models` implementation must consume it and remove it before calling `Provider.stream*()`. Provider implementations continue receiving ordinary `ApiStreamOptions` or `SimpleStreamOptions` and never handle the transform themselves. Use this option instead of calling `getAuth(model)` before `stream*()`, which would resolve request auth twice. ### Credential Store @@ -359,7 +387,7 @@ const models = createModels({ credentials: myFileBackedStore }); // const models = builtinModels({ credentials: myFileBackedStore }); ``` -The contract is small: `read(providerId)`, `modify(providerId, fn)` (the only write path — a serialized read-modify-write), and `delete(providerId)`. 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)`. 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: @@ -412,7 +440,7 @@ Built-in providers resolve these env vars (Node.js; in browsers pass `apiKey` ex | Xiaomi MiMo Token Plan (Singapore) | `XIAOMI_TOKEN_PLAN_SGP_API_KEY` | | GitHub Copilot | `COPILOT_GITHUB_TOKEN` | -Amazon Bedrock resolves ambient AWS credentials (`AWS_PROFILE`, access key pairs, `AWS_BEARER_TOKEN_BEDROCK`, ECS task roles, web identity tokens). Vertex AI resolves either an explicit key or gcloud Application Default Credentials plus project/location. +Amazon Bedrock resolves ambient AWS credentials (`AWS_PROFILE`, access key pairs, `AWS_BEARER_TOKEN_BEDROCK`, ECS task roles, web identity tokens); its provider-owned login flow supports bearer tokens, AWS profiles, and the existing credential chain. Vertex AI resolves either an explicit key or gcloud Application Default Credentials plus project/location, with a provider-owned login flow for API keys, ADC, and service-account files. ## Tools @@ -660,7 +688,7 @@ for (const block of result.output) { } ``` -Like the chat side, you can build the collection from parts: `createImagesModels({ credentials?, authContext? })`, the `openrouterImagesProvider()` factory from `@earendil-works/pi-ai/providers/openrouter-images`, and `createImagesProvider({ id, auth, models, refreshModels?, api })` for custom image providers (with `imagesModels.refresh(provider?)` for dynamic lists). Failures never reject — they return an `AssistantImages` with `stopReason: "error"`. The collection's `getAuth(model)` works exactly like the chat-side one. +Like the chat side, you can build the collection from parts: `createImagesModels({ credentials?, authContext? })`, the `openrouterImagesProvider()` factory from `@earendil-works/pi-ai/providers/openrouter-images`, and `createImagesProvider({ id, auth, models, refreshModels?, api })` for custom image providers (with `imagesModels.refresh(provider?)` for dynamic lists). Failures never reject — they return an `AssistantImages` with `stopReason: "error"`. The collection's provider-scoped `getAuth(providerId)` works exactly like the chat-side one. The old global API (`getImageModel()` / `getImageModels()` / `getImageProviders()` / `generateImages()`) remains available on the [compat entrypoint](#migrating-from-the-old-global-api): @@ -982,6 +1010,25 @@ const gateway = createProvider({ }); ``` +Provider-wide endpoint or request transformations belong in the provider's API implementation: wrap the `ProviderStreams` you pass as `api` so every request goes through the transformation before dispatch. The Cloudflare providers do this to materialize account/gateway endpoint placeholders from the resolved provider env: + +```typescript +function tenantStreams(streams: ProviderStreams): ProviderStreams { + const withTenant = (model: Model) => ({ ...model, baseUrl: model.baseUrl.replace('{tenant}', tenantId) }); + return { + stream: (model, context, options) => streams.stream(withTenant(model), context, options), + streamSimple: (model, context, options) => streams.streamSimple(withTenant(model), context, options), + }; +} + +const tenantGateway = createProvider({ + id: 'tenant-gateway', + auth: { apiKey: envApiKeyAuth('Gateway key', ['GATEWAY_API_KEY']) }, + models: [/* ... */], + api: tenantStreams(openAICompletionsApi()), +}); +``` + Dynamic model lists use `refreshModels`; the provider lists empty until the first `models.refresh()`: ```typescript @@ -997,7 +1044,7 @@ models.setProvider(llamacpp); await models.refresh('llamacpp'); ``` -Custom models can carry `headers` (e.g. proxies behind bot detection) and `compat` flags — see [OpenAI Compatibility Settings](#openai-compatibility-settings). +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). Some OpenAI-compatible servers do not understand the `developer` role used for reasoning-capable models. For those providers, set `compat.supportsDeveloperRole` to `false` so the system prompt is sent as a `system` message instead. If the server also does not support `reasoning_effort`, set `compat.supportsReasoningEffort` to `false` too. This commonly applies to Ollama, vLLM, SGLang, and similar OpenAI-compatible servers. @@ -1368,7 +1415,7 @@ Several providers support OAuth authentication instead of static API keys: - **OpenAI Codex** (ChatGPT Plus/Pro subscription, access to GPT-5.x Codex models) - **GitHub Copilot** (Copilot subscription) -Each of these providers carries an `OAuthAuth` on `provider.auth.oauth` with three operations: `login(callbacks)` runs the interactive flow and returns a credential, `refresh(credential)` exchanges the refresh token, and `toAuth(credential)` derives request auth (GitHub Copilot's per-account base URL comes from here). Refresh is automatic: `models.getAuth()` and the request paths refresh expired tokens under a credential-store lock, so concurrent requests and processes cannot double-refresh. +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)` exchanges the refresh token, 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. ```typescript import { createModels } from '@earendil-works/pi-ai'; @@ -1377,34 +1424,36 @@ import { anthropicProvider } from '@earendil-works/pi-ai/providers/anthropic'; const models = createModels({ credentials: myStore }); // persistent CredentialStore models.setProvider(anthropicProvider()); -// Login: drive the flow with prompt()/notify() callbacks, persist the credential -const provider = models.getProvider('anthropic')!; -const credential = await provider.auth.oauth!.login({ +// Login: Models drives the flow and persists the credential +await models.login('anthropic', 'oauth', { prompt: async (p) => { // p.type: 'text' | 'secret' | 'select' | 'manual_code' // manual_code prompts race a local callback server; p.signal aborts them when the server wins return await askUser(p.message); }, notify: (event) => { - // event.type: 'auth_url' | 'device_code' | 'progress' + // event.type: 'info' | 'auth_url' | 'device_code' | 'progress' + if (event.type === 'info') { + console.log(event.message); + for (const link of event.links ?? []) console.log(`${link.label ?? 'More information'}: ${link.url}`); + } if (event.type === 'auth_url') console.log(`Open: ${event.url}`); if (event.type === 'device_code') console.log(`Code: ${event.userCode} at ${event.verificationUri}`); if (event.type === 'progress') console.log(event.message); }, }); -await myStore.modify('anthropic', async () => credential); // From here on, requests resolve and refresh the token automatically const model = models.getModel('anthropic', 'claude-sonnet-4-5')!; await models.complete(model, context); // Logout -await myStore.delete('anthropic'); +await models.logout('anthropic'); ``` ### Vertex AI -Vertex AI models support either a Google Cloud API key or Application Default Credentials (ADC): +Vertex AI models support either a Google Cloud API key or Application Default Credentials (ADC). Its provider-owned API-key login flow can configure either method: - **API key**: Set `GOOGLE_CLOUD_API_KEY` or pass `apiKey` in the call options. - **Local development (ADC)**: Run `gcloud auth application-default login` @@ -1438,7 +1487,7 @@ Credentials are saved to `auth.json` in the current directory. ### Programmatic OAuth -The legacy flow functions remain available via the `@earendil-works/pi-ai/oauth` entry point (`loginAnthropic`, `loginOpenAICodex`, `loginGitHubCopilot`, `refreshOAuthToken`, `getOAuthApiKey`); credential storage is the caller's responsibility there. New code should prefer the provider-owned `OAuthAuth` shown above — it composes with the credential store and gets locked auto-refresh for free. +Built-in login and refresh flows are private provider implementations. Use provider-owned `OAuthAuth`, which composes with `CredentialStore` and gets locked auto-refresh through `Models`. The `@earendil-works/pi-ai/oauth` entry point retains only type declarations required by coding-agent extension OAuth compatibility. Provider notes: @@ -1468,7 +1517,7 @@ Compat is a strict superset of the root entrypoint, so a file can switch its imp | `getModels('anthropic')` / `getProviders()` | `models.getModels('anthropic')` / `models.getProviders()` or `getBuiltin*` | | `stream(model, ctx, opts)` (env-key injection) | `models.stream(model, ctx, opts)` (provider auth resolution) | | `registerApiProvider({ api, stream, streamSimple })` | `createProvider({ id, auth, models, api })` + `models.setProvider()` | -| `getEnvApiKey('openai')` | `await models.getAuth(model)` | +| `getEnvApiKey('openai')` | `await models.getAuth(model.provider)` | | `streamAnthropic(model, ctx, opts)` | `stream` from `@earendil-works/pi-ai/api/anthropic-messages`, or a provider in a collection | | `registerFauxProvider()` | `fauxProvider()` + `models.setProvider()` | diff --git a/packages/ai/src/api/lazy.ts b/packages/ai/src/api/lazy.ts index fe1836ae..b30224e9 100644 --- a/packages/ai/src/api/lazy.ts +++ b/packages/ai/src/api/lazy.ts @@ -22,13 +22,20 @@ function createSetupErrorMessage(model: Model, error: unknown): AssistantMe }; } -function forwardStream(target: AssistantMessageEventStream, source: AsyncIterable): void { - (async () => { - for await (const event of source) { - target.push(event); - } - target.end(); - })(); +function hasResult( + source: AsyncIterable, +): source is AsyncIterable & { result(): Promise } { + return typeof (source as { result?: unknown }).result === "function"; +} + +async function forwardStream( + target: AssistantMessageEventStream, + source: AsyncIterable, +): Promise { + for await (const event of source) { + target.push(event); + } + target.end(hasResult(source) ? await source.result() : undefined); } /** @@ -43,9 +50,7 @@ export function lazyStream( const outer = new AssistantMessageEventStream(); setup() - .then((inner) => { - forwardStream(outer, inner); - }) + .then((inner) => forwardStream(outer, inner)) .catch((error) => { const message = createSetupErrorMessage(model, error); outer.push({ type: "error", reason: "error", error: message }); diff --git a/packages/ai/src/auth/credential-store.ts b/packages/ai/src/auth/credential-store.ts index beeb9d85..f2a09e1a 100644 --- a/packages/ai/src/auth/credential-store.ts +++ b/packages/ai/src/auth/credential-store.ts @@ -1,4 +1,4 @@ -import type { Credential, CredentialStore } from "./types.ts"; +import type { Credential, CredentialInfo, CredentialStore } from "./types.ts"; /** * Default in-memory credential store. Apps inject persistent stores. @@ -27,6 +27,10 @@ export class InMemoryCredentialStore implements CredentialStore { return this.credentials.get(providerId); } + async list(): Promise { + return [...this.credentials].map(([providerId, credential]) => ({ providerId, type: credential.type })); + } + modify( providerId: string, fn: (current: Credential | undefined) => Promise, diff --git a/packages/ai/src/auth/helpers.ts b/packages/ai/src/auth/helpers.ts index d9a34ad2..968b029b 100644 --- a/packages/ai/src/auth/helpers.ts +++ b/packages/ai/src/auth/helpers.ts @@ -9,8 +9,8 @@ import type { ApiKeyAuth, OAuthAuth } from "./types.ts"; export function envApiKeyAuth(name: string, envVars: readonly string[]): ApiKeyAuth { return { name, - login: async (callbacks) => { - const key = await callbacks.prompt({ type: "secret", message: `Enter ${name}` }); + login: async (interaction) => { + const key = await interaction.prompt({ type: "secret", message: `Enter ${name}` }); return { type: "api_key", key }; }, resolve: async ({ ctx, credential }) => { @@ -39,7 +39,7 @@ export function lazyOAuth(input: { name: string; load: () => Promise }; return { name: input.name, - login: async (callbacks) => (await loaded()).login(callbacks), + login: async (interaction) => (await loaded()).login(interaction), refresh: async (credential) => (await loaded()).refresh(credential), toAuth: async (credential) => (await loaded()).toAuth(credential), }; diff --git a/packages/ai/src/utils/oauth/anthropic.ts b/packages/ai/src/auth/oauth/anthropic.ts similarity index 67% rename from packages/ai/src/utils/oauth/anthropic.ts rename to packages/ai/src/auth/oauth/anthropic.ts index 591e9cde..dbaaa300 100644 --- a/packages/ai/src/utils/oauth/anthropic.ts +++ b/packages/ai/src/auth/oauth/anthropic.ts @@ -6,11 +6,10 @@ */ import type { Server } from "node:http"; -import type { OAuthAuth } from "../../auth/types.ts"; -import { getProviderEnvValue } from "../provider-env.ts"; +import { getProviderEnvValue } from "../../utils/provider-env.ts"; +import type { AuthInteraction, OAuthAuth, OAuthCredential } from "../types.ts"; import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.ts"; import { generatePKCE } from "./pkce.ts"; -import type { OAuthCredentials, OAuthLoginCallbacks, OAuthPrompt, OAuthProviderInterface } from "./types.ts"; type CallbackServerInfo = { server: Server; @@ -193,7 +192,7 @@ async function exchangeAuthorizationCode( state: string, verifier: string, redirectUri: string, -): Promise { +): Promise { let responseBody: string; try { responseBody = await postJson(TOKEN_URL, { @@ -220,27 +219,21 @@ async function exchangeAuthorizationCode( } return { + type: "oauth", refresh: tokenData.refresh_token, access: tokenData.access_token, expires: Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000, }; } -/** - * Login with Anthropic OAuth (authorization code + PKCE) - */ -export async function loginAnthropic(options: { - onAuth: (info: { url: string; instructions?: string }) => void; - onPrompt: (prompt: OAuthPrompt) => Promise; - onProgress?: (message: string) => void; - onManualCodeInput?: () => Promise; -}): Promise { +async function loginAnthropic(interaction: AuthInteraction): Promise { const { verifier, challenge } = await generatePKCE(); const server = await startCallbackServer(verifier); - + const manualAbort = new AbortController(); let code: string | undefined; let state: string | undefined; - let redirectUriForExchange = REDIRECT_URI; + let manualInput: string | undefined; + let manualError: Error | undefined; try { const authParams = new URLSearchParams({ @@ -253,93 +246,58 @@ export async function loginAnthropic(options: { code_challenge_method: "S256", state: verifier, }); - - options.onAuth({ + interaction.notify({ + type: "auth_url", url: `${AUTHORIZE_URL}?${authParams.toString()}`, instructions: "Complete login in your browser. If the browser is on another machine, paste the final redirect URL here.", }); - if (options.onManualCodeInput) { - let manualInput: string | undefined; - let manualError: Error | undefined; - const manualPromise = options - .onManualCodeInput() - .then((input) => { - manualInput = input; - server.cancelWait(); - }) - .catch((err) => { - manualError = err instanceof Error ? err : new Error(String(err)); - server.cancelWait(); - }); - - const result = await server.waitForCode(); - - if (manualError) { - throw manualError; - } - - if (result?.code) { - code = result.code; - state = result.state; - redirectUriForExchange = REDIRECT_URI; - } else if (manualInput) { - const parsed = parseAuthorizationInput(manualInput); - if (parsed.state && parsed.state !== verifier) { - throw new Error("OAuth state mismatch"); - } - code = parsed.code; - state = parsed.state ?? verifier; - } - - if (!code) { - await manualPromise; - if (manualError) { - throw manualError; - } - if (manualInput) { - const parsed = parseAuthorizationInput(manualInput); - if (parsed.state && parsed.state !== verifier) { - throw new Error("OAuth state mismatch"); - } - code = parsed.code; - state = parsed.state ?? verifier; - } - } - } else { - const result = await server.waitForCode(); - if (result?.code) { - code = result.code; - state = result.state; - redirectUriForExchange = REDIRECT_URI; - } - } - - if (!code) { - const input = await options.onPrompt({ - message: "Paste the authorization code or full redirect URL:", + const manualPromise = interaction + .prompt({ + type: "manual_code", + message: "Complete login in your browser, or paste the authorization code / redirect URL here:", placeholder: REDIRECT_URI, + signal: manualAbort.signal, + }) + .then((input) => { + manualInput = input; + server.cancelWait(); + }) + .catch((error) => { + manualError = error instanceof Error ? error : new Error(String(error)); + server.cancelWait(); }); - const parsed = parseAuthorizationInput(input); - if (parsed.state && parsed.state !== verifier) { - throw new Error("OAuth state mismatch"); - } + + const result = await server.waitForCode(); + if (manualError) throw manualError; + if (result?.code) { + code = result.code; + state = result.state; + } else if (manualInput) { + const parsed = parseAuthorizationInput(manualInput); + if (parsed.state && parsed.state !== verifier) throw new Error("OAuth state mismatch"); code = parsed.code; state = parsed.state ?? verifier; } if (!code) { - throw new Error("Missing authorization code"); + await manualPromise; + if (manualError) throw manualError; + if (manualInput) { + const parsed = parseAuthorizationInput(manualInput); + if (parsed.state && parsed.state !== verifier) throw new Error("OAuth state mismatch"); + code = parsed.code; + state = parsed.state ?? verifier; + } } - if (!state) { - throw new Error("Missing OAuth state"); - } - - options.onProgress?.("Exchanging authorization code for tokens..."); - return exchangeAuthorizationCode(code, state, verifier, redirectUriForExchange); + if (!code) throw new Error("Missing authorization code"); + if (!state) throw new Error("Missing OAuth state"); + interaction.notify({ type: "progress", message: "Exchanging authorization code for tokens..." }); + return exchangeAuthorizationCode(code, state, verifier, REDIRECT_URI); } finally { + manualAbort.abort(); server.server.close(); } } @@ -347,7 +305,7 @@ export async function loginAnthropic(options: { /** * Refresh Anthropic OAuth token */ -export async function refreshAnthropicToken(refreshToken: string): Promise { +async function refreshAnthropicToken(refreshToken: string): Promise { let responseBody: string; try { responseBody = await postJson(TOKEN_URL, { @@ -374,6 +332,7 @@ export async function refreshAnthropicToken(refreshToken: string): Promise callbacks.notify({ type: "auth_url", url: info.url, instructions: info.instructions }), - onProgress: (message) => callbacks.notify({ type: "progress", message }), - onPrompt: (prompt) => - callbacks.prompt({ type: "text", message: prompt.message, placeholder: prompt.placeholder }), - onManualCodeInput: () => - callbacks.prompt({ - type: "manual_code", - message: "Complete login in your browser, or paste the authorization code / redirect URL here:", - placeholder: REDIRECT_URI, - signal: manualAbort.signal, - }), - }); - return { ...credentials, type: "oauth" }; - } finally { - manualAbort.abort(); - } - }, - - async refresh(credential) { - return { ...(await refreshAnthropicToken(credential.refresh)), type: "oauth" }; - }, + login: loginAnthropic, + refresh: (credential) => refreshAnthropicToken(credential.refresh), async toAuth(credential) { return { apiKey: credential.access }; }, }; - -export const anthropicOAuthProvider: OAuthProviderInterface = { - id: "anthropic", - name: "Anthropic (Claude Pro/Max)", - usesCallbackServer: true, - - async login(callbacks: OAuthLoginCallbacks): Promise { - return loginAnthropic({ - onAuth: callbacks.onAuth, - onPrompt: callbacks.onPrompt, - onProgress: callbacks.onProgress, - onManualCodeInput: callbacks.onManualCodeInput, - }); - }, - - async refreshToken(credentials: OAuthCredentials): Promise { - return refreshAnthropicToken(credentials.refresh); - }, - - getApiKey(credentials: OAuthCredentials): string { - return credentials.access; - }, -}; diff --git a/packages/ai/src/utils/oauth/device-code.ts b/packages/ai/src/auth/oauth/device-code.ts similarity index 100% rename from packages/ai/src/utils/oauth/device-code.ts rename to packages/ai/src/auth/oauth/device-code.ts diff --git a/packages/ai/src/utils/oauth/github-copilot.ts b/packages/ai/src/auth/oauth/github-copilot.ts similarity index 71% rename from packages/ai/src/utils/oauth/github-copilot.ts rename to packages/ai/src/auth/oauth/github-copilot.ts index 01f86957..0fe1a04e 100644 --- a/packages/ai/src/utils/oauth/github-copilot.ts +++ b/packages/ai/src/auth/oauth/github-copilot.ts @@ -2,16 +2,9 @@ * GitHub Copilot OAuth flow */ -import type { OAuthAuth, OAuthCredential } from "../../auth/types.ts"; import { GITHUB_COPILOT_MODELS } from "../../providers/github-copilot.models.ts"; -import type { Api, Model } from "../../types.ts"; +import type { AuthInteraction, OAuthAuth, OAuthCredential } from "../types.ts"; import { pollOAuthDeviceCodeFlow } from "./device-code.ts"; -import type { OAuthCredentials, OAuthDeviceCodeInfo, OAuthLoginCallbacks, OAuthProviderInterface } from "./types.ts"; - -type CopilotCredentials = OAuthCredentials & { - enterpriseUrl?: string; - availableModelIds: string[]; -}; const decode = (s: string) => atob(s); const CLIENT_ID = decode("SXYxLmI1MDdhMDhjODdlY2ZlOTg="); @@ -44,7 +37,7 @@ type DeviceTokenErrorResponse = { interval?: number; }; -export function normalizeDomain(input: string): string | null { +function normalizeDomain(input: string): string | null { const trimmed = input.trim(); if (!trimmed) return null; try { @@ -81,7 +74,7 @@ function getBaseUrlFromToken(token: string): string | null { return `https://${apiHost}`; } -export function getGitHubCopilotBaseUrl(token?: string, enterpriseDomain?: string): string { +function getGitHubCopilotBaseUrl(token?: string, enterpriseDomain?: string): string { // If we have a token, extract the base URL from proxy-ep if (token) { const urlFromToken = getBaseUrlFromToken(token); @@ -251,7 +244,7 @@ async function pollForGitHubAccessToken( async function refreshGitHubCopilotAccessToken( refreshToken: string, enterpriseDomain?: string, -): Promise { +): Promise { const domain = enterpriseDomain || "github.com"; const urls = getUrls(domain); @@ -275,6 +268,7 @@ async function refreshGitHubCopilotAccessToken( } return { + type: "oauth", refresh: refreshToken, access: token, expires: expiresAt * 1000 - 5 * 60 * 1000, @@ -285,10 +279,7 @@ async function refreshGitHubCopilotAccessToken( /** * Refresh GitHub Copilot token */ -export async function refreshGitHubCopilotToken( - refreshToken: string, - enterpriseDomain?: string, -): Promise { +async function refreshGitHubCopilotToken(refreshToken: string, enterpriseDomain?: string): Promise { const credentials = await refreshGitHubCopilotAccessToken(refreshToken, enterpriseDomain); return { ...credentials, @@ -326,68 +317,41 @@ 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, - onProgress?: (model: string, success: boolean) => void, -): Promise { +async function enableAllGitHubCopilotModels(token: string, enterpriseDomain?: string): Promise { const models = Object.values(GITHUB_COPILOT_MODELS); await Promise.all( models.map(async (model) => { - const success = await enableGitHubCopilotModel(token, model.id, enterpriseDomain); - onProgress?.(model.id, success); + await enableGitHubCopilotModel(token, model.id, enterpriseDomain); }), ); } -/** - * Login with GitHub Copilot OAuth (device code flow) - * - * @param options.onDeviceCode - Callback with URL and user code - * @param options.onPrompt - Callback to prompt user for input - * @param options.onProgress - Optional progress callback - * @param options.signal - Optional AbortSignal for cancellation - */ -export async function loginGitHubCopilot(options: { - onDeviceCode: (info: OAuthDeviceCodeInfo) => void; - onPrompt: (prompt: { message: string; placeholder?: string; allowEmpty?: boolean }) => Promise; - onProgress?: (message: string) => void; - signal?: AbortSignal; -}): Promise { - const input = await options.onPrompt({ +async function loginGitHubCopilot(interaction: AuthInteraction): Promise { + const input = await interaction.prompt({ + type: "text", message: "GitHub Enterprise URL/domain (blank for github.com)", placeholder: "company.ghe.com", - allowEmpty: true, }); - - if (options.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"); - } + if (trimmed && !enterpriseDomain) throw new Error("Invalid GitHub Enterprise URL/domain"); const domain = enterpriseDomain || "github.com"; const device = await startDeviceFlow(domain); - options.onDeviceCode({ + interaction.notify({ + type: "device_code", userCode: device.user_code, verificationUri: device.verification_uri, intervalSeconds: device.interval, expiresInSeconds: device.expires_in, }); - const githubAccessToken = await pollForGitHubAccessToken(domain, device, options.signal); + const githubAccessToken = await pollForGitHubAccessToken(domain, device, interaction.signal); const credentials = await refreshGitHubCopilotAccessToken(githubAccessToken, enterpriseDomain ?? undefined); - - // Enable all models after successful login - options.onProgress?.("Enabling models..."); + interaction.notify({ type: "progress", message: "Enabling models..." }); await enableAllGitHubCopilotModels(credentials.access, enterpriseDomain ?? undefined); - - // Fetch availability after policy enable so newly enabled models are included, - // while unavailable models are still filtered out. return { ...credentials, availableModelIds: await fetchAvailableGitHubCopilotModelIds(credentials.access, enterpriseDomain ?? undefined), @@ -402,26 +366,10 @@ function copilotEnterpriseDomain(credential: OAuthCredential): string | undefine export const githubCopilotOAuth: OAuthAuth = { name: "GitHub Copilot", + login: loginGitHubCopilot, + refresh: (credential) => refreshGitHubCopilotToken(credential.refresh, copilotEnterpriseDomain(credential)), - async login(callbacks) { - const credentials = await loginGitHubCopilot({ - onDeviceCode: (info) => callbacks.notify({ type: "device_code", ...info }), - onPrompt: (prompt) => - callbacks.prompt({ type: "text", message: prompt.message, placeholder: prompt.placeholder }), - onProgress: (message) => callbacks.notify({ type: "progress", message }), - signal: callbacks.signal, - }); - return { ...credentials, type: "oauth" }; - }, - - async refresh(credential) { - return { - ...(await refreshGitHubCopilotToken(credential.refresh, copilotEnterpriseDomain(credential))), - type: "oauth", - }; - }, - - /** Per-credential baseUrl from the token's proxy endpoint replaces the old `modifyModels` rewriting. */ + /** Derive the credential-specific proxy endpoint for each request. */ async toAuth(credential) { return { apiKey: credential.access, @@ -429,41 +377,3 @@ export const githubCopilotOAuth: OAuthAuth = { }; }, }; - -export const githubCopilotOAuthProvider: OAuthProviderInterface = { - id: "github-copilot", - name: "GitHub Copilot", - - async login(callbacks: OAuthLoginCallbacks): Promise { - return loginGitHubCopilot({ - onDeviceCode: callbacks.onDeviceCode, - onPrompt: callbacks.onPrompt, - onProgress: callbacks.onProgress, - signal: callbacks.signal, - }); - }, - - async refreshToken(credentials: OAuthCredentials): Promise { - const creds = credentials as CopilotCredentials; - return refreshGitHubCopilotToken(creds.refresh, creds.enterpriseUrl); - }, - - getApiKey(credentials: OAuthCredentials): string { - return credentials.access; - }, - - modifyModels(models: Model[], credentials: OAuthCredentials): Model[] { - const creds = credentials as CopilotCredentials; - const domain = creds.enterpriseUrl ? (normalizeDomain(creds.enterpriseUrl) ?? undefined) : undefined; - const baseUrl = getGitHubCopilotBaseUrl(creds.access, domain); - // Older stored Pi auth entries do not have account-specific model IDs yet; - // keep their existing generated-catalog behavior until the next refresh/login. - const availableModelIds = "availableModelIds" in creds ? new Set(creds.availableModelIds) : undefined; - - return models.flatMap((m) => { - if (m.provider !== "github-copilot") return [m]; - if (availableModelIds && !availableModelIds.has(m.id)) return []; - return [{ ...m, baseUrl }]; - }); - }, -}; diff --git a/packages/ai/src/utils/oauth/load.ts b/packages/ai/src/auth/oauth/load.ts similarity index 94% rename from packages/ai/src/utils/oauth/load.ts rename to packages/ai/src/auth/oauth/load.ts index 11198853..8e3d023d 100644 --- a/packages/ai/src/utils/oauth/load.ts +++ b/packages/ai/src/auth/oauth/load.ts @@ -1,4 +1,4 @@ -import type { OAuthAuth } from "../../auth/types.ts"; +import type { OAuthAuth } from "../types.ts"; /** * Loads an OAuth flow module through a variable specifier so bundlers cannot diff --git a/packages/ai/src/utils/oauth/oauth-page.ts b/packages/ai/src/auth/oauth/oauth-page.ts similarity index 100% rename from packages/ai/src/utils/oauth/oauth-page.ts rename to packages/ai/src/auth/oauth/oauth-page.ts diff --git a/packages/ai/src/utils/oauth/openai-codex.ts b/packages/ai/src/auth/oauth/openai-codex.ts similarity index 67% rename from packages/ai/src/utils/oauth/openai-codex.ts rename to packages/ai/src/auth/oauth/openai-codex.ts index a2f7cd00..f1b27c3a 100644 --- a/packages/ai/src/utils/oauth/openai-codex.ts +++ b/packages/ai/src/auth/oauth/openai-codex.ts @@ -17,18 +17,11 @@ if (typeof process !== "undefined" && (process.versions?.node || process.version }); } -import type { OAuthAuth } from "../../auth/types.ts"; -import { getProviderEnvValue } from "../provider-env.ts"; +import { getProviderEnvValue } from "../../utils/provider-env.ts"; +import type { AuthInteraction, OAuthAuth, OAuthCredential } from "../types.ts"; import { pollOAuthDeviceCodeFlow } from "./device-code.ts"; import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.ts"; import { generatePKCE } from "./pkce.ts"; -import type { - OAuthCredentials, - OAuthDeviceCodeInfo, - OAuthLoginCallbacks, - OAuthPrompt, - OAuthProviderInterface, -} from "./types.ts"; const CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"; const AUTH_BASE_URL = "https://auth.openai.com"; @@ -40,8 +33,8 @@ const DEVICE_TOKEN_URL = `${AUTH_BASE_URL}/api/accounts/deviceauth/token`; const DEVICE_VERIFICATION_URI = `${AUTH_BASE_URL}/codex/device`; const DEVICE_REDIRECT_URI = `${AUTH_BASE_URL}/deviceauth/callback`; const DEVICE_CODE_TIMEOUT_SECONDS = 15 * 60; -export const OPENAI_CODEX_BROWSER_LOGIN_METHOD = "browser"; -export const OPENAI_CODEX_DEVICE_CODE_LOGIN_METHOD = "device_code"; +const OPENAI_CODEX_BROWSER_LOGIN_METHOD = "browser"; +const OPENAI_CODEX_DEVICE_CODE_LOGIN_METHOD = "device_code"; const SCOPE = "openid profile email offline_access"; const JWT_CLAIM_PATH = "https://api.openai.com/auth"; @@ -406,13 +399,14 @@ function getAccountId(accessToken: string): string | null { return typeof accountId === "string" && accountId.length > 0 ? accountId : null; } -function credentialsFromToken(token: OAuthToken): OAuthCredentials { +function credentialsFromToken(token: OAuthToken): OAuthCredential { const accountId = getAccountId(token.access); if (!accountId) { throw new Error("Failed to extract accountId from token"); } return { + type: "oauth", access: token.access, refresh: token.refresh, expires: token.expires, @@ -425,132 +419,83 @@ async function exchangeAuthorizationCodeForCredentials( verifier: string, redirectUri: string, signal?: AbortSignal, -): Promise { +): Promise { return credentialsFromToken(await exchangeAuthorizationCode(code, verifier, redirectUri, signal)); } -/** - * Login with OpenAI Codex OAuth using the Codex device-code flow. - */ -export async function loginOpenAICodexDeviceCode(options: { - onDeviceCode: (info: OAuthDeviceCodeInfo) => void; - signal?: AbortSignal; -}): Promise { - const device = await startOpenAICodexDeviceAuth(options.signal); - options.onDeviceCode({ +async function loginOpenAICodexDeviceCode(interaction: AuthInteraction): Promise { + const device = await startOpenAICodexDeviceAuth(interaction.signal); + interaction.notify({ + type: "device_code", userCode: device.userCode, verificationUri: DEVICE_VERIFICATION_URI, intervalSeconds: device.intervalSeconds, expiresInSeconds: DEVICE_CODE_TIMEOUT_SECONDS, }); - const code = await pollOpenAICodexDeviceAuth(device, options.signal); + const code = await pollOpenAICodexDeviceAuth(device, interaction.signal); return exchangeAuthorizationCodeForCredentials( code.authorizationCode, code.codeVerifier, DEVICE_REDIRECT_URI, - options.signal, + interaction.signal, ); } -/** - * Login with OpenAI Codex OAuth - * - * @param options.onAuth - Called with URL and instructions when auth starts - * @param options.onPrompt - Called to prompt user for manual code paste (fallback if no onManualCodeInput) - * @param options.onProgress - Optional progress messages - * @param options.onManualCodeInput - Optional promise that resolves with user-pasted code. - * Races with browser callback - whichever completes first wins. - * Useful for showing paste input immediately alongside browser flow. - * @param options.originator - OAuth originator parameter (defaults to "pi") - */ -export async function loginOpenAICodex(options: { - onAuth: (info: { url: string; instructions?: string }) => void; - onPrompt: (prompt: OAuthPrompt) => Promise; - onProgress?: (message: string) => void; - onManualCodeInput?: () => Promise; - originator?: string; -}): Promise { - const { verifier, state, url } = await createAuthorizationFlow(options.originator); +async function loginOpenAICodex(interaction: AuthInteraction): Promise { + const { verifier, state, url } = await createAuthorizationFlow(); const server = await startLocalOAuthServer(state); - - options.onAuth({ url, instructions: "A browser window should open. Complete login to finish." }); - + const manualAbort = new AbortController(); let code: string | undefined; + let manualCode: string | undefined; + let manualError: Error | undefined; + + interaction.notify({ + type: "auth_url", + url, + instructions: "A browser window should open. Complete login to finish.", + }); + try { - if (options.onManualCodeInput) { - // Race between browser callback and manual input - let manualCode: string | undefined; - let manualError: Error | undefined; - const manualPromise = options - .onManualCodeInput() - .then((input) => { - manualCode = input; - server.cancelWait(); - }) - .catch((err) => { - manualError = err instanceof Error ? err : new Error(String(err)); - server.cancelWait(); - }); - - const result = await server.waitForCode(); - - // If manual input was cancelled, throw that error - if (manualError) { - throw manualError; - } - - if (result?.code) { - // Browser callback won - code = result.code; - } else if (manualCode) { - // Manual input won (or callback timed out and user had entered code) - const parsed = parseAuthorizationInput(manualCode); - if (parsed.state && parsed.state !== state) { - throw new Error("State mismatch"); - } - code = parsed.code; - } - - // If still no code, wait for manual promise to complete and try that - if (!code) { - await manualPromise; - if (manualError) { - throw manualError; - } - if (manualCode) { - const parsed = parseAuthorizationInput(manualCode); - if (parsed.state && parsed.state !== state) { - throw new Error("State mismatch"); - } - code = parsed.code; - } - } - } else { - // Original flow: wait for callback, then prompt if needed - const result = await server.waitForCode(); - if (result?.code) { - code = result.code; - } - } - - // Fallback to onPrompt if still no code - if (!code) { - const input = await options.onPrompt({ - message: "Paste the authorization code (or full redirect URL):", + const manualPromise = interaction + .prompt({ + type: "manual_code", + message: "Complete login in your browser, or paste the authorization code / redirect URL here:", + placeholder: REDIRECT_URI, + signal: manualAbort.signal, + }) + .then((input) => { + manualCode = input; + server.cancelWait(); + }) + .catch((error) => { + manualError = error instanceof Error ? error : new Error(String(error)); + server.cancelWait(); }); - const parsed = parseAuthorizationInput(input); - if (parsed.state && parsed.state !== state) { - throw new Error("State mismatch"); - } + + const result = await server.waitForCode(); + if (manualError) throw manualError; + if (result?.code) { + code = result.code; + } else if (manualCode) { + const parsed = parseAuthorizationInput(manualCode); + if (parsed.state && parsed.state !== state) throw new Error("State mismatch"); code = parsed.code; } if (!code) { - throw new Error("Missing authorization code"); + await manualPromise; + if (manualError) throw manualError; + if (manualCode) { + const parsed = parseAuthorizationInput(manualCode); + if (parsed.state && parsed.state !== state) throw new Error("State mismatch"); + code = parsed.code; + } } - return exchangeAuthorizationCodeForCredentials(code, verifier, REDIRECT_URI); + if (!code) throw new Error("Missing authorization code"); + return exchangeAuthorizationCodeForCredentials(code, verifier, REDIRECT_URI, interaction.signal); } finally { + manualAbort.abort(); server.close(); } } @@ -558,15 +503,15 @@ export async function loginOpenAICodex(options: { /** * Refresh OpenAI Codex OAuth token */ -export async function refreshOpenAICodexToken(refreshToken: string): Promise { +async function refreshOpenAICodexToken(refreshToken: string): Promise { return credentialsFromToken(await refreshAccessToken(refreshToken)); } export const openaiCodexOAuth: OAuthAuth = { name: "OpenAI (ChatGPT Plus/Pro)", - async login(callbacks) { - const method = await callbacks.prompt({ + async login(interaction) { + const method = await interaction.prompt({ type: "select", message: "Select OpenAI Codex login method:", options: [ @@ -576,89 +521,18 @@ export const openaiCodexOAuth: OAuthAuth = { }); if (method === OPENAI_CODEX_DEVICE_CODE_LOGIN_METHOD) { - const credentials = await loginOpenAICodexDeviceCode({ - onDeviceCode: (info) => callbacks.notify({ type: "device_code", ...info }), - signal: callbacks.signal, - }); - return { ...credentials, type: "oauth" }; + return loginOpenAICodexDeviceCode(interaction); } if (method !== OPENAI_CODEX_BROWSER_LOGIN_METHOD) { throw new Error(`Unknown OpenAI Codex login method: ${method}`); } - // The manual_code prompt races the local callback server; abort it once - // the flow settles so the UI can dismiss the pending input. - const manualAbort = new AbortController(); - try { - const credentials = await loginOpenAICodex({ - onAuth: (info) => callbacks.notify({ type: "auth_url", url: info.url, instructions: info.instructions }), - onProgress: (message) => callbacks.notify({ type: "progress", message }), - onPrompt: (prompt) => - callbacks.prompt({ type: "text", message: prompt.message, placeholder: prompt.placeholder }), - onManualCodeInput: () => - callbacks.prompt({ - type: "manual_code", - message: "Complete login in your browser, or paste the authorization code / redirect URL here:", - placeholder: REDIRECT_URI, - signal: manualAbort.signal, - }), - }); - return { ...credentials, type: "oauth" }; - } finally { - manualAbort.abort(); - } + return loginOpenAICodex(interaction); }, - async refresh(credential) { - return { ...(await refreshOpenAICodexToken(credential.refresh)), type: "oauth" }; - }, + refresh: (credential) => refreshOpenAICodexToken(credential.refresh), async toAuth(credential) { return { apiKey: credential.access }; }, }; - -export const openaiCodexOAuthProvider: OAuthProviderInterface = { - id: "openai-codex", - name: "ChatGPT Plus/Pro (Codex Subscription)", - usesCallbackServer: true, - - async login(callbacks: OAuthLoginCallbacks): Promise { - const loginMethod = await callbacks.onSelect({ - message: "Select OpenAI Codex login method:", - options: [ - { id: OPENAI_CODEX_BROWSER_LOGIN_METHOD, label: "Browser login (default)" }, - { id: OPENAI_CODEX_DEVICE_CODE_LOGIN_METHOD, label: "Device code login (headless)" }, - ], - }); - if (!loginMethod) { - throw new Error("Login cancelled"); - } - - if (loginMethod === OPENAI_CODEX_DEVICE_CODE_LOGIN_METHOD) { - return loginOpenAICodexDeviceCode({ - onDeviceCode: callbacks.onDeviceCode, - signal: callbacks.signal, - }); - } - - if (loginMethod !== OPENAI_CODEX_BROWSER_LOGIN_METHOD) { - throw new Error(`Unknown OpenAI Codex login method: ${loginMethod}`); - } - - return loginOpenAICodex({ - onAuth: callbacks.onAuth, - onPrompt: callbacks.onPrompt, - onProgress: callbacks.onProgress, - onManualCodeInput: callbacks.onManualCodeInput, - }); - }, - - async refreshToken(credentials: OAuthCredentials): Promise { - return refreshOpenAICodexToken(credentials.refresh); - }, - - getApiKey(credentials: OAuthCredentials): string { - return credentials.access; - }, -}; diff --git a/packages/ai/src/utils/oauth/pkce.ts b/packages/ai/src/auth/oauth/pkce.ts similarity index 100% rename from packages/ai/src/utils/oauth/pkce.ts rename to packages/ai/src/auth/oauth/pkce.ts diff --git a/packages/ai/src/auth/resolve.ts b/packages/ai/src/auth/resolve.ts index 81d7a270..b5ae130f 100644 --- a/packages/ai/src/auth/resolve.ts +++ b/packages/ai/src/auth/resolve.ts @@ -1,4 +1,4 @@ -import type { Api, ImagesApi, ImagesModel, Model, ProviderEnv } from "../types.ts"; +import type { ProviderEnv } from "../types.ts"; import type { ApiKeyAuth, ApiKeyCredential, @@ -28,9 +28,6 @@ export class ModelsError extends Error { } } -/** Model shape auth resolution receives: chat or image-generation models. */ -export type AuthModel = Model | ImagesModel; - /** * Auth resolution shared by the `Models` and `ImagesModels` collections. * A stored credential owns the provider: ambient/env is consulted only when @@ -39,7 +36,6 @@ export type AuthModel = Model | ImagesModel; */ export async function resolveProviderAuth( provider: { id: string; auth: ProviderAuth }, - model: AuthModel, credentials: CredentialStore, authContext: AuthContext, overrides?: AuthResolutionOverrides, @@ -47,7 +43,7 @@ export async function resolveProviderAuth( const requestAuthContext = overrides?.env ? overlayEnvAuthContext(authContext, overrides.env) : authContext; if (overrides?.apiKey !== undefined && provider.auth.apiKey) { - return resolveApiKey(requestAuthContext, provider.auth.apiKey, model, { + return resolveApiKey(requestAuthContext, provider.auth.apiKey, provider.id, { type: "api_key", key: overrides.apiKey, env: overrides.env, @@ -61,13 +57,15 @@ export async function resolveProviderAuth( } 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, model, credential); + return resolveApiKey(requestAuthContext, provider.auth.apiKey, provider.id, credential); } return undefined; } // Ambient (env vars, AWS profiles, ADC files). - return provider.auth.apiKey ? resolveApiKey(requestAuthContext, provider.auth.apiKey, model, undefined) : undefined; + return provider.auth.apiKey + ? resolveApiKey(requestAuthContext, provider.auth.apiKey, provider.id, undefined) + : undefined; } function overlayEnvAuthContext(base: AuthContext, env: ProviderEnv): AuthContext { @@ -122,13 +120,13 @@ async function resolveStoredOAuth( async function resolveApiKey( authContext: AuthContext, apiKey: ApiKeyAuth, - model: AuthModel, + providerId: string, credential: ApiKeyCredential | undefined, ): Promise { try { - return await apiKey.resolve({ model, ctx: authContext, credential }); + return await apiKey.resolve({ ctx: authContext, credential }); } catch (error) { - throw new ModelsError("auth", `API key auth failed for provider ${model.provider}`, { cause: error }); + throw new ModelsError("auth", `API key auth failed for provider ${providerId}`, { cause: error }); } } diff --git a/packages/ai/src/auth/types.ts b/packages/ai/src/auth/types.ts index 9710dbcb..12c1e7a8 100644 --- a/packages/ai/src/auth/types.ts +++ b/packages/ai/src/auth/types.ts @@ -1,5 +1,4 @@ -import type { Api, ImagesApi, ImagesModel, Model, ProviderEnv, ProviderHeaders } from "../types.ts"; -import type { OAuthCredentials } from "../utils/oauth/types.ts"; +import type { ProviderEnv, ProviderHeaders } from "../types.ts"; /** * Request auth for a single model request. If a value cannot be expressed as @@ -21,7 +20,15 @@ export interface ApiKeyCredential { env?: ProviderEnv; } -/** Stored OAuth credential (`access`, `refresh`, `expires` from OAuthCredentials). */ +/** OAuth token data returned by extension compatibility flows. */ +export interface OAuthCredentials { + refresh: string; + access: string; + expires: number; + [key: string]: unknown; +} + +/** Stored canonical OAuth credential. */ export interface OAuthCredential extends OAuthCredentials { type: "oauth"; } @@ -29,6 +36,12 @@ export interface OAuthCredential extends OAuthCredentials { /** One type-tagged credential per provider — the shape of today's auth.json. */ export type Credential = ApiKeyCredential | OAuthCredential; +/** Non-secret credential metadata for account/status enumeration. */ +export interface CredentialInfo { + providerId: string; + type: Credential["type"]; +} + /** * App-owned credential storage, keyed by `Provider.id`, one credential per * provider. `modify` is the only write path, so every mutation is a @@ -51,6 +64,12 @@ export interface CredentialStore { */ read(providerId: string): Promise; + /** + * List stored credential metadata without resolving or exposing secrets. + * Implementations must not execute configured API-key commands while listing. + */ + list(): Promise; + /** * Serialized write — the only write path. `fn` sees the current credential * because correct writes (refresh, login-during-refresh) depend on it; @@ -84,6 +103,13 @@ export interface AuthResult { source?: string; } +export interface AuthCheck { + source?: string; + type: "api_key" | "oauth"; +} + +export type AuthType = "api_key" | "oauth"; + /** * Prompt shown to the user during login. `signal` lets the flow cancel a * pending prompt when an out-of-band event resolves the step, e.g. a @@ -97,7 +123,13 @@ export type AuthPrompt = { signal?: AbortSignal } & ( | { type: "manual_code"; message: string; placeholder?: string } ); +export interface AuthInfoLink { + url: string; + label?: string; +} + export type AuthEvent = + | { type: "info"; message: string; links?: readonly AuthInfoLink[] } | { type: "auth_url"; url: string; instructions?: string } | { type: "device_code"; @@ -115,7 +147,7 @@ export type AuthEvent = * id). Rejects on cancel/abort. `signal` aborts the whole login flow; * per-prompt cancellation uses `AuthPrompt.signal`. */ -export interface AuthLoginCallbacks { +export interface AuthInteraction { signal?: AbortSignal; prompt(prompt: AuthPrompt): Promise; @@ -131,19 +163,22 @@ export interface ApiKeyAuth { name: string; /** Interactive setup (prompt for key/provider env). Absent = ambient-only. */ - login?(callbacks: AuthLoginCallbacks): Promise; + login?(interaction: AuthInteraction): 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; /** * Resolve auth from the stored credential and/or ambient sources, merging * per field (`credential.key ?? env("...")`, `credential.env?.NAME ?? env("...")`). - * undefined = not configured. Receives the chat or image-generation model - * the request is for (both carry `provider` and `baseUrl`). + * undefined = not configured. Resolution is provider-scoped; model-specific + * endpoint preparation happens after auth has been resolved. */ - resolve(input: { - model: Model | ImagesModel; - ctx: AuthContext; - credential?: ApiKeyCredential; - }): Promise; + resolve(input: { ctx: AuthContext; credential?: ApiKeyCredential }): Promise; } /** @@ -155,7 +190,7 @@ export interface OAuthAuth { /** Display name, e.g. "Anthropic (Claude Pro/Max)". */ name: string; - login(callbacks: AuthLoginCallbacks): Promise; + login(interaction: AuthInteraction): Promise; /** * Exchange the refresh token. Network call; throws on failure diff --git a/packages/ai/src/cli.ts b/packages/ai/src/cli.ts index 21699dbd..e5fff7a5 100644 --- a/packages/ai/src/cli.ts +++ b/packages/ai/src/cli.ts @@ -1,71 +1,74 @@ #!/usr/bin/env node +import { existsSync, readFileSync, writeFileSync } from "node:fs"; import { createInterface } from "node:readline"; -import { existsSync, readFileSync, writeFileSync } from "fs"; -import { getOAuthProvider, getOAuthProviders } from "./utils/oauth/index.ts"; -import type { OAuthCredentials, OAuthProviderId } from "./utils/oauth/types.ts"; +import type { AuthPrompt, OAuthCredential, Provider } from "./index.ts"; +import { builtinProviders } from "./providers/all.ts"; const AUTH_FILE = "auth.json"; -const PROVIDERS = getOAuthProviders(); +const PROVIDERS = builtinProviders().filter( + (provider): provider is Provider & { auth: { oauth: NonNullable } } => + provider.auth.oauth !== undefined, +); function prompt(rl: ReturnType, question: string): Promise { return new Promise((resolve) => rl.question(question, resolve)); } -function loadAuth(): Record { +function loadAuth(): Record { if (!existsSync(AUTH_FILE)) return {}; try { - return JSON.parse(readFileSync(AUTH_FILE, "utf-8")); + return JSON.parse(readFileSync(AUTH_FILE, "utf-8")) as Record; } catch { return {}; } } -function saveAuth(auth: Record): void { +function saveAuth(auth: Record): void { writeFileSync(AUTH_FILE, JSON.stringify(auth, null, 2), "utf-8"); } -async function login(providerId: OAuthProviderId): Promise { - const provider = getOAuthProvider(providerId); - if (!provider) { - console.error(`Unknown provider: ${providerId}`); - process.exit(1); +async function answerPrompt(rl: ReturnType, authPrompt: AuthPrompt): Promise { + if (authPrompt.type === "select") { + console.log(`\n${authPrompt.message}`); + for (let index = 0; index < authPrompt.options.length; index++) { + console.log(` ${index + 1}. ${authPrompt.options[index].label}`); + } + const choice = Number.parseInt(await prompt(rl, `Enter number (1-${authPrompt.options.length}): `), 10) - 1; + const selected = authPrompt.options[choice]; + if (!selected) throw new Error("Invalid selection"); + return selected.id; } + return prompt(rl, `${authPrompt.message}${authPrompt.placeholder ? ` (${authPrompt.placeholder})` : ""}: `); +} +async function login(providerId: string): Promise { + const provider = PROVIDERS.find((entry) => entry.id === providerId); + if (!provider) throw new Error(`Unknown provider: ${providerId}`); const rl = createInterface({ input: process.stdin, output: process.stdout }); - const promptFn = (msg: string) => prompt(rl, `${msg} `); - try { - const credentials = await provider.login({ - onAuth: (info) => { - console.log(`\nOpen this URL in your browser:\n${info.url}`); - if (info.instructions) console.log(info.instructions); - console.log(); - }, - onDeviceCode: (info) => { - console.log(`\nOpen this URL in your browser:\n${info.verificationUri}`); - console.log(`Enter code: ${info.userCode}`); - console.log(); - }, - onPrompt: async (p) => { - return await promptFn(`${p.message}${p.placeholder ? ` (${p.placeholder})` : ""}:`); - }, - onSelect: async (p) => { - console.log(`\n${p.message}`); - for (let i = 0; i < p.options.length; i++) { - console.log(` ${i + 1}. ${p.options[i].label}`); + const credential = await provider.auth.oauth.login({ + prompt: (authPrompt) => answerPrompt(rl, authPrompt), + notify: (event) => { + switch (event.type) { + case "auth_url": + console.log(`\nOpen this URL in your browser:\n${event.url}`); + if (event.instructions) console.log(event.instructions); + break; + case "device_code": + console.log(`\nOpen this URL in your browser:\n${event.verificationUri}`); + console.log(`Enter code: ${event.userCode}`); + break; + case "info": + case "progress": + console.log(event.message); + break; } - const choice = await promptFn(`Enter number (1-${p.options.length}):`); - const index = parseInt(choice, 10) - 1; - return p.options[index]?.id; }, - onProgress: (msg) => console.log(msg), }); - const auth = loadAuth(); - auth[providerId] = { type: "oauth", ...credentials }; + auth[providerId] = credential; saveAuth(auth); - console.log(`\nCredentials saved to ${AUTH_FILE}`); } finally { rl.close(); @@ -75,73 +78,41 @@ async function login(providerId: OAuthProviderId): Promise { async function main(): Promise { const args = process.argv.slice(2); const command = args[0]; - if (!command || command === "help" || command === "--help" || command === "-h") { - const providerList = PROVIDERS.map((p) => ` ${p.id.padEnd(20)} ${p.name}`).join("\n"); - console.log(`Usage: npx @earendil-works/pi-ai [provider] - -Commands: - login [provider] Login to an OAuth provider - list List available providers - -Providers: -${providerList} - -Examples: - npx @earendil-works/pi-ai login # interactive provider selection - npx @earendil-works/pi-ai login anthropic # login to specific provider - npx @earendil-works/pi-ai list # list providers -`); + const providerList = PROVIDERS.map((provider) => ` ${provider.id.padEnd(20)} ${provider.name}`).join("\n"); + console.log( + `Usage: npx @earendil-works/pi-ai [provider]\n\nCommands:\n login [provider] Login to an OAuth provider\n list List available providers\n\nProviders:\n${providerList}`, + ); return; } - if (command === "list") { - console.log("Available OAuth providers:\n"); - for (const p of PROVIDERS) { - console.log(` ${p.id.padEnd(20)} ${p.name}`); - } + for (const provider of PROVIDERS) console.log(`${provider.id.padEnd(20)} ${provider.name}`); return; } - if (command === "login") { - let provider = args[1] as OAuthProviderId | undefined; - - if (!provider) { + let providerId = args[1]; + if (!providerId) { const rl = createInterface({ input: process.stdin, output: process.stdout }); - console.log("Select a provider:\n"); - for (let i = 0; i < PROVIDERS.length; i++) { - console.log(` ${i + 1}. ${PROVIDERS[i].name}`); + try { + for (let index = 0; index < PROVIDERS.length; index++) { + console.log(` ${index + 1}. ${PROVIDERS[index].name}`); + } + const index = Number.parseInt(await prompt(rl, `Enter number (1-${PROVIDERS.length}): `), 10) - 1; + providerId = PROVIDERS[index]?.id; + } finally { + rl.close(); } - console.log(); - - const choice = await prompt(rl, `Enter number (1-${PROVIDERS.length}): `); - rl.close(); - - const index = parseInt(choice, 10) - 1; - if (index < 0 || index >= PROVIDERS.length) { - console.error("Invalid selection"); - process.exit(1); - } - provider = PROVIDERS[index].id; } - - if (!PROVIDERS.some((p) => p.id === provider)) { - console.error(`Unknown provider: ${provider}`); - console.error(`Use 'npx @earendil-works/pi-ai list' to see available providers`); - process.exit(1); + if (!providerId || !PROVIDERS.some((provider) => provider.id === providerId)) { + throw new Error(`Unknown provider: ${providerId ?? ""}`); } - - console.log(`Logging in to ${provider}...`); - await login(provider); + await login(providerId); return; } - - console.error(`Unknown command: ${command}`); - console.error(`Use 'npx @earendil-works/pi-ai --help' for usage`); - process.exit(1); + throw new Error(`Unknown command: ${command}`); } -main().catch((err) => { - console.error("Error:", err.message); +main().catch((error: unknown) => { + console.error("Error:", error instanceof Error ? error.message : String(error)); process.exit(1); }); diff --git a/packages/ai/src/compat.ts b/packages/ai/src/compat.ts index 91ecda4f..d1465ddb 100644 --- a/packages/ai/src/compat.ts +++ b/packages/ai/src/compat.ts @@ -37,6 +37,7 @@ import { openAICodexResponsesApi } from "./api/openai-codex-responses.lazy.ts"; import { openAICompletionsApi } from "./api/openai-completions.lazy.ts"; import { openAIResponsesApi } from "./api/openai-responses.lazy.ts"; import { getEnvApiKey } from "./env-api-keys.ts"; +import type { ModelsApiStreamOptions } from "./models.ts"; import { builtinModels, getBuiltinModel, getBuiltinModels, getBuiltinProviders } from "./providers/all.ts"; import { createFauxCore, type FauxProviderRegistration, type RegisterFauxProviderOptions } from "./providers/faux.ts"; import type { @@ -221,9 +222,14 @@ function withEnvApiKey( return { ...options, apiKey } as TOptions; } -function shouldUseBuiltinModels(model: Model): boolean { - const builtin = compatModels.getModel(model.provider, model.id); - return builtin?.api === model.api && getApiProvider(model.api) === builtinApiProviderInstances.get(model.api); +function hasResolvedCloudflareAuth(options: StreamOptions | undefined): boolean { + return hasExplicitApiKey(options?.apiKey) || typeof options?.headers?.["cf-aig-authorization"] === "string"; +} + +function getBuiltinProviderForModel(model: Model) { + if (getApiProvider(model.api) !== builtinApiProviderInstances.get(model.api)) return undefined; + const provider = compatModels.getProvider(model.provider); + return provider?.getModels().some((candidate) => candidate.api === model.api) ? provider : undefined; } function resolveApiProvider(api: Api) { @@ -239,8 +245,12 @@ export function stream( context: Context, options?: ProviderStreamOptions, ): AssistantMessageEventStream { - if (shouldUseBuiltinModels(model)) { - return compatModels.stream(model, context, options as ApiStreamOptions | undefined); + const builtinProvider = getBuiltinProviderForModel(model); + if (builtinProvider) { + if (model.provider.startsWith("cloudflare-") && !hasResolvedCloudflareAuth(options)) { + return compatModels.stream(model, context, options as ModelsApiStreamOptions | undefined); + } + return builtinProvider.stream(model, context, withEnvApiKey(model, options) as ApiStreamOptions); } const provider = resolveApiProvider(model.api); return provider.stream(model, context, withEnvApiKey(model, options) as StreamOptions); @@ -260,8 +270,12 @@ export function streamSimple( context: Context, options?: SimpleStreamOptions, ): AssistantMessageEventStream { - if (shouldUseBuiltinModels(model)) { - return compatModels.streamSimple(model, context, options); + const builtinProvider = getBuiltinProviderForModel(model); + if (builtinProvider) { + if (model.provider.startsWith("cloudflare-") && !hasResolvedCloudflareAuth(options)) { + return compatModels.streamSimple(model, context, options); + } + return builtinProvider.streamSimple(model, context, withEnvApiKey(model, options)); } const provider = resolveApiProvider(model.api); return provider.streamSimple(model, context, withEnvApiKey(model, options)); diff --git a/packages/ai/src/compat/extension-oauth-types.ts b/packages/ai/src/compat/extension-oauth-types.ts new file mode 100644 index 00000000..f1bf4e5d --- /dev/null +++ b/packages/ai/src/compat/extension-oauth-types.ts @@ -0,0 +1,45 @@ +import type { OAuthCredentials } from "../auth/types.ts"; + +/** Legacy extension OAuth prompt. */ +export interface OAuthPrompt { + message: string; + placeholder?: string; + allowEmpty?: boolean; +} + +/** Legacy extension OAuth authorization link. */ +export interface OAuthAuthInfo { + url: string; + instructions?: string; +} + +/** Legacy extension OAuth device-code notification. */ +export interface OAuthDeviceCodeInfo { + userCode: string; + verificationUri: string; + intervalSeconds?: number; + expiresInSeconds?: number; +} + +export interface OAuthSelectOption { + id: string; + label: string; +} + +export interface OAuthSelectPrompt { + message: string; + options: OAuthSelectOption[]; +} + +/** Callback surface retained only for coding-agent extension compatibility. */ +export interface OAuthLoginCallbacks { + onAuth(info: OAuthAuthInfo): void; + onDeviceCode(info: OAuthDeviceCodeInfo): void; + onPrompt(prompt: OAuthPrompt): Promise; + onProgress?(message: string): void; + onManualCodeInput?(): Promise; + onSelect(prompt: OAuthSelectPrompt): Promise; + signal?: AbortSignal; +} + +export type { OAuthCredentials }; diff --git a/packages/ai/src/images-models.ts b/packages/ai/src/images-models.ts index 0ca5f2da..4afe0f70 100644 --- a/packages/ai/src/images-models.ts +++ b/packages/ai/src/images-models.ts @@ -1,6 +1,6 @@ import { defaultProviderAuthContext as defaultAuthContext } from "./auth/context.ts"; import { InMemoryCredentialStore } from "./auth/credential-store.ts"; -import { ModelsError, resolveProviderAuth } from "./auth/resolve.ts"; +import { type AuthResolutionOverrides, ModelsError, resolveProviderAuth } from "./auth/resolve.ts"; import type { AuthContext, AuthResult, CredentialStore, ProviderAuth } from "./auth/types.ts"; import type { CreateModelsOptions } from "./models.ts"; import type { AssistantImages, ImagesApi, ImagesContext, ImagesModel, ImagesOptions, ProviderImages } from "./types.ts"; @@ -68,11 +68,12 @@ export interface ImagesModels { refresh(provider?: string): Promise; /** - * Resolve request auth for an image model. Same contract as + * Resolve request auth by provider id or image model. Same contract as * `Models.getAuth()`: undefined when unknown/unconfigured, rejects with * `ModelsError` ("oauth"/"auth") on real failures. */ - getAuth(model: ImagesModel): Promise; + getAuth(providerId: string, overrides?: AuthResolutionOverrides): Promise; + getAuth(model: ImagesModel, overrides?: AuthResolutionOverrides): Promise; /** * Generate images through the owning provider with auth resolved and @@ -167,10 +168,16 @@ class ImagesModelsImpl implements MutableImagesModels { await Promise.allSettled(Array.from(this.providers.values(), async (entry) => entry.refreshModels?.())); } - async getAuth(model: ImagesModel): Promise { - const provider = this.providers.get(model.provider); + getAuth(providerId: string, overrides?: AuthResolutionOverrides): Promise; + getAuth(model: ImagesModel, overrides?: AuthResolutionOverrides): Promise; + async getAuth( + providerOrModel: string | ImagesModel, + overrides?: AuthResolutionOverrides, + ): Promise { + const providerId = typeof providerOrModel === "string" ? providerOrModel : providerOrModel.provider; + const provider = this.providers.get(providerId); if (!provider) return undefined; - return resolveProviderAuth(provider, model, this.credentials, this.authContext); + return resolveProviderAuth(provider, this.credentials, this.authContext, overrides); } async generateImages( @@ -184,7 +191,7 @@ class ImagesModelsImpl implements MutableImagesModels { throw new ModelsError("provider", `Unknown provider: ${model.provider}`); } - const resolution = await resolveProviderAuth(provider, model, this.credentials, this.authContext, { + const resolution = await this.getAuth(model, { apiKey: options?.apiKey, env: options?.env, }); diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 57c448cc..38533bbd 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -21,6 +21,14 @@ export * from "./auth/context.ts"; export * from "./auth/credential-store.ts"; export * from "./auth/helpers.ts"; export * from "./auth/types.ts"; +export type { + OAuthAuthInfo, + OAuthDeviceCodeInfo, + OAuthLoginCallbacks, + OAuthPrompt, + OAuthSelectOption, + OAuthSelectPrompt, +} from "./compat/extension-oauth-types.ts"; export * from "./images-models.ts"; export * from "./models.ts"; export * from "./providers/faux.ts"; @@ -29,19 +37,6 @@ export * from "./types.ts"; export * from "./utils/diagnostics.ts"; export * from "./utils/event-stream.ts"; export * from "./utils/json-parse.ts"; -export type { - OAuthAuthInfo, - OAuthCredentials, - OAuthDeviceCodeInfo, - OAuthLoginCallbacks, - OAuthPrompt, - OAuthProvider, - OAuthProviderId, - OAuthProviderInfo, - OAuthProviderInterface, - OAuthSelectOption, - OAuthSelectPrompt, -} from "./utils/oauth/types.ts"; export * from "./utils/overflow.ts"; export * from "./utils/retry.ts"; export * from "./utils/typebox-helpers.ts"; diff --git a/packages/ai/src/models.ts b/packages/ai/src/models.ts index 7288524e..3c6de74d 100644 --- a/packages/ai/src/models.ts +++ b/packages/ai/src/models.ts @@ -1,8 +1,17 @@ import { lazyStream } from "./api/lazy.ts"; import { defaultProviderAuthContext as defaultAuthContext } from "./auth/context.ts"; import { InMemoryCredentialStore } from "./auth/credential-store.ts"; -import { ModelsError, resolveProviderAuth } from "./auth/resolve.ts"; -import type { AuthContext, AuthResult, CredentialStore, ProviderAuth } from "./auth/types.ts"; +import { type AuthResolutionOverrides, ModelsError, resolveProviderAuth } from "./auth/resolve.ts"; +import type { + AuthCheck, + AuthContext, + AuthInteraction, + AuthResult, + AuthType, + Credential, + CredentialStore, + ProviderAuth, +} from "./auth/types.ts"; import type { Api, ApiStreamOptions, @@ -19,7 +28,15 @@ import type { Usage, } from "./types.ts"; -export { type AuthModel, ModelsError, type ModelsErrorCode } from "./auth/resolve.ts"; +export { ModelsError, type ModelsErrorCode } from "./auth/resolve.ts"; + +export interface ModelsStreamTransforms { + /** Transform fully assembled model/auth/request headers before provider dispatch. */ + transformHeaders?: (headers: ProviderHeaders) => ProviderHeaders | Promise; +} + +export type ModelsApiStreamOptions = ApiStreamOptions & ModelsStreamTransforms; +export type ModelsSimpleStreamOptions = SimpleStreamOptions & ModelsStreamTransforms; /** * A provider is the concrete runtime unit. It owns id/name/base metadata, @@ -63,6 +80,13 @@ export interface Provider { */ refreshModels?(): Promise; + /** + * Optional provider policy for credential-specific model availability. + * `getModels()` remains the complete synchronous catalog; `Models.getAvailable()` + * applies this filter after confirming that provider auth is configured. + */ + filterModels?(models: readonly Model[], credential: Credential | undefined): readonly Model[]; + stream( model: Model, context: Context, @@ -101,31 +125,44 @@ export interface Models { */ refresh(provider?: string): Promise; + /** Check whether a provider has complete auth configuration without refreshing OAuth. */ + checkAuth(providerId: string): Promise; + + /** Return models whose providers have complete auth configuration. */ + getAvailable(providerId?: string): Promise[]>; + /** - * Resolve request auth for a model. Includes a source label for status UI. + * Resolve provider-scoped auth by provider id, or provider auth plus static + * model headers when passed a model. Includes a source label for status UI. * Resolves `undefined` when the provider is unknown or unconfigured. * Rejects with `ModelsError`: code "oauth" when a token refresh fails (the * stored credential is preserved for retry; re-login fixes it), code "auth" * when api-key resolution or the credential store fails. Request paths - * surface rejections as stream errors; status/availability UIs catch them - * and render "needs re-login" instead of treating them as unconfigured. + * surface rejections as stream errors. */ - getAuth(model: Model): Promise; + getAuth(providerId: string, overrides?: AuthResolutionOverrides): Promise; + getAuth(model: Model, overrides?: AuthResolutionOverrides): Promise; + + /** Run a provider-owned login flow and persist its returned credential. */ + login(providerId: string, type: AuthType, interaction: AuthInteraction): Promise; + + /** Remove the stored credential for a provider. */ + logout(providerId: string): Promise; stream( model: Model, context: Context, - options?: ApiStreamOptions, + options?: ModelsApiStreamOptions, ): AssistantMessageEventStream; complete( model: Model, context: Context, - options?: ApiStreamOptions, + options?: ModelsApiStreamOptions, ): Promise; - streamSimple(model: Model, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream; - completeSimple(model: Model, context: Context, options?: SimpleStreamOptions): Promise; + streamSimple(model: Model, context: Context, options?: ModelsSimpleStreamOptions): AssistantMessageEventStream; + completeSimple(model: Model, context: Context, options?: ModelsSimpleStreamOptions): Promise; } export interface MutableModels extends Models { @@ -140,6 +177,22 @@ export interface CreateModelsOptions { authContext?: AuthContext; } +function mergeHeaders( + base: ProviderHeaders | undefined, + override: ProviderHeaders | undefined, +): ProviderHeaders | undefined { + if (!base && !override) return undefined; + const merged = { ...base }; + for (const [name, value] of Object.entries(override ?? {})) { + const lowerName = name.toLowerCase(); + for (const existingName of Object.keys(merged)) { + if (existingName.toLowerCase() === lowerName) delete merged[existingName]; + } + merged[name] = value; + } + return merged; +} + class ModelsImpl implements MutableModels { private providers = new Map(); private credentials: CredentialStore; @@ -214,10 +267,103 @@ class ModelsImpl implements MutableModels { await Promise.allSettled(Array.from(this.providers.values(), async (entry) => entry.refreshModels?.())); } - async getAuth(model: Model): Promise { - const provider = this.providers.get(model.provider); + private async readCredential(providerId: string): Promise { + try { + return await this.credentials.read(providerId); + } catch (error) { + throw new ModelsError("auth", `Credential store read failed for ${providerId}`, { cause: error }); + } + } + + private async checkProviderAuth( + provider: Provider, + credential: Credential | undefined, + ): Promise { + if (credential?.type === "oauth") { + return provider.auth.oauth ? { source: "OAuth", type: "oauth" } : undefined; + } + const apiKey = provider.auth.apiKey; + if (!apiKey) return undefined; + if (apiKey.check) { + try { + return await apiKey.check({ + ctx: this.authContext, + credential: credential?.type === "api_key" ? credential : undefined, + }); + } 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); + 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 resolveProviderAuth(provider, model, this.credentials, this.authContext); + return this.checkProviderAuth(provider, await this.readCredential(providerId)); + } + + 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; + }); + } + + getAuth(providerId: string, overrides?: AuthResolutionOverrides): Promise; + getAuth(model: Model, overrides?: AuthResolutionOverrides): Promise; + async getAuth( + providerOrModel: string | Model, + overrides?: AuthResolutionOverrides, + ): Promise { + 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); + if (!result || typeof providerOrModel === "string" || !providerOrModel.headers) return result; + return { + ...result, + auth: { + ...result.auth, + headers: mergeHeaders(result.auth.headers, providerOrModel.headers), + }, + }; + } + + async login(providerId: string, type: AuthType, interaction: AuthInteraction): Promise { + 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); + try { + await this.credentials.modify(providerId, async () => credential); + } catch (error) { + throw new ModelsError("auth", `Credential store modify failed for ${providerId}`, { cause: error }); + } + return credential; + } + + async logout(providerId: string): Promise { + try { + await this.credentials.delete(providerId); + } catch (error) { + throw new ModelsError("auth", `Credential store delete failed for ${providerId}`, { cause: error }); + } } private requireProvider(model: Model): Provider { @@ -228,30 +374,28 @@ class ModelsImpl implements MutableModels { return provider; } - private async applyAuth( + private async applyAuth( model: Model, options: TOptions | undefined, - ): Promise<{ requestModel: Model; requestOptions: TOptions | undefined }> { - const resolution = await resolveProviderAuth( - this.requireProvider(model), - model, - this.credentials, - this.authContext, - { - apiKey: options?.apiKey, - env: options?.env, - }, - ); - const auth = resolution?.auth; - if (!auth) return { requestModel: model, requestOptions: options }; + ): Promise<{ requestModel: Model; requestOptions: StreamOptions | undefined }> { + this.requireProvider(model); + const resolution = await this.getAuth(model, { + apiKey: options?.apiKey, + env: options?.env, + }); + if (!resolution) { + throw new ModelsError("auth", `Provider is not configured: ${model.provider}`); + } + const auth = resolution.auth; - const requestModel = auth.baseUrl ? { ...model, baseUrl: auth.baseUrl } : model; - - // Explicit request options win per-field; headers/env merge per key. + // Explicit request options win per-field; the Models-only transform runs last. const apiKey = options?.apiKey ?? auth.apiKey; - const headers = auth.headers || options?.headers ? { ...auth.headers, ...options?.headers } : undefined; + let headers = mergeHeaders(auth.headers, options?.headers); + if (options?.transformHeaders) headers = await options.transformHeaders(headers ?? {}); const env = resolution.env || options?.env ? { ...(resolution.env ?? {}), ...(options?.env ?? {}) } : undefined; - const requestOptions = { ...options, apiKey, headers, env } as TOptions; + const requestModel = auth.baseUrl ? { ...model, baseUrl: auth.baseUrl } : model; + const { transformHeaders: _transformHeaders, ...providerOptions } = options ?? {}; + const requestOptions = { ...providerOptions, apiKey, headers, env } as StreamOptions; return { requestModel, requestOptions }; } @@ -259,11 +403,14 @@ class ModelsImpl implements MutableModels { stream( model: Model, context: Context, - options?: ApiStreamOptions, + options?: ModelsApiStreamOptions, ): AssistantMessageEventStream { return lazyStream(model, async () => { const provider = this.requireProvider(model); - const { requestModel, requestOptions } = await this.applyAuth(model, options as StreamOptions | undefined); + const { requestModel, requestOptions } = await this.applyAuth( + model, + options as ModelsApiStreamOptions | undefined, + ); return provider.stream(requestModel as Model, context, requestOptions as ApiStreamOptions); }); } @@ -271,20 +418,24 @@ class ModelsImpl implements MutableModels { async complete( model: Model, context: Context, - options?: ApiStreamOptions, + options?: ModelsApiStreamOptions, ): Promise { return this.stream(model, context, options).result(); } - streamSimple(model: Model, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream { + streamSimple(model: Model, context: Context, options?: ModelsSimpleStreamOptions): AssistantMessageEventStream { return lazyStream(model, async () => { const provider = this.requireProvider(model); const { requestModel, requestOptions } = await this.applyAuth(model, options); - return provider.streamSimple(requestModel, context, requestOptions); + return provider.streamSimple(requestModel, context, requestOptions as SimpleStreamOptions); }); } - async completeSimple(model: Model, context: Context, options?: SimpleStreamOptions): Promise { + async completeSimple( + model: Model, + context: Context, + options?: ModelsSimpleStreamOptions, + ): Promise { return this.streamSimple(model, context, options).result(); } } @@ -311,6 +462,7 @@ export interface CreateProviderOptions { * `Models.refresh(provider)`), and a later call retries. */ refreshModels?: () => Promise[]>; + filterModels?: (models: readonly Model[], credential: Credential | undefined) => readonly Model[]; /** Single implementation, or map keyed by `model.api` for mixed-API providers. */ api: ProviderStreams | Partial>; } @@ -363,6 +515,7 @@ export function createProvider(input: CreateProviderOpti return inflightRefresh; } : undefined, + filterModels: input.filterModels, stream: (model, context, options) => dispatch(model, (streams) => streams.stream(model, context, options)), streamSimple: (model, context, options) => dispatch(model, (streams) => streams.streamSimple(model, context, options)), diff --git a/packages/ai/src/oauth.ts b/packages/ai/src/oauth.ts index 487816d8..039a587f 100644 --- a/packages/ai/src/oauth.ts +++ b/packages/ai/src/oauth.ts @@ -1 +1,10 @@ -export * from "./utils/oauth/index.ts"; +/** Type-only compatibility entry point for coding-agent extension OAuth declarations. */ +export type { + OAuthAuthInfo, + OAuthCredentials, + OAuthDeviceCodeInfo, + OAuthLoginCallbacks, + OAuthPrompt, + OAuthSelectOption, + OAuthSelectPrompt, +} from "./compat/extension-oauth-types.ts"; diff --git a/packages/ai/src/providers/amazon-bedrock.models.ts b/packages/ai/src/providers/amazon-bedrock.models.ts index aaddca5f..a3bc9162 100644 --- a/packages/ai/src/providers/amazon-bedrock.models.ts +++ b/packages/ai/src/providers/amazon-bedrock.models.ts @@ -1290,6 +1290,60 @@ export const AMAZON_BEDROCK_MODELS = { contextWindow: 272000, maxTokens: 128000, } satisfies Model<"bedrock-converse-stream">, + "openai.gpt-5.6-luna": { + id: "openai.gpt-5.6-luna", + name: "GPT-5.6 Luna", + api: "bedrock-converse-stream", + provider: "amazon-bedrock", + baseUrl: "https://bedrock-runtime.us-east-1.amazonaws.com", + reasoning: true, + thinkingLevelMap: {"xhigh":"xhigh"}, + input: ["text", "image"], + cost: { + input: 1, + output: 6, + cacheRead: 0.1, + cacheWrite: 1.25, + }, + contextWindow: 272000, + maxTokens: 128000, + } satisfies Model<"bedrock-converse-stream">, + "openai.gpt-5.6-sol": { + id: "openai.gpt-5.6-sol", + name: "GPT-5.6 Sol", + api: "bedrock-converse-stream", + provider: "amazon-bedrock", + baseUrl: "https://bedrock-runtime.us-east-1.amazonaws.com", + reasoning: true, + thinkingLevelMap: {"xhigh":"xhigh"}, + input: ["text", "image"], + cost: { + input: 5, + output: 30, + cacheRead: 0.5, + cacheWrite: 6.25, + }, + contextWindow: 272000, + maxTokens: 128000, + } satisfies Model<"bedrock-converse-stream">, + "openai.gpt-5.6-terra": { + id: "openai.gpt-5.6-terra", + name: "GPT-5.6 Terra", + api: "bedrock-converse-stream", + provider: "amazon-bedrock", + baseUrl: "https://bedrock-runtime.us-east-1.amazonaws.com", + reasoning: true, + thinkingLevelMap: {"xhigh":"xhigh"}, + input: ["text", "image"], + cost: { + input: 2.5, + output: 15, + cacheRead: 0.25, + cacheWrite: 3.125, + }, + contextWindow: 272000, + maxTokens: 128000, + } satisfies Model<"bedrock-converse-stream">, "openai.gpt-oss-120b": { id: "openai.gpt-oss-120b", name: "gpt-oss-120b", diff --git a/packages/ai/src/providers/amazon-bedrock.ts b/packages/ai/src/providers/amazon-bedrock.ts index d839ab6a..83a7050f 100644 --- a/packages/ai/src/providers/amazon-bedrock.ts +++ b/packages/ai/src/providers/amazon-bedrock.ts @@ -4,16 +4,61 @@ import { createProvider, type Provider } from "../models.ts"; import { AMAZON_BEDROCK_MODELS } from "./amazon-bedrock.models.ts"; /** - * Bedrock auth is ambient: the AWS SDK's default credential chain handles the - * actual signing, so `resolve` only reports whether the provider is - * configured. A stored credential key is surfaced as the bearer token. + * Bedrock accepts a bearer token or the AWS SDK's default credential chain. + * The login flow can store a token/profile choice; resolve also detects ambient + * AWS credentials without copying them into pi's credential store. */ const bedrockAuth: ApiKeyAuth = { - name: "AWS credentials", + name: "AWS credentials or bearer token", + login: async (interaction) => { + const method = await interaction.prompt({ + type: "select", + message: "Select Amazon Bedrock authentication method:", + options: [ + { id: "bearer-token", label: "Bearer token" }, + { id: "aws-profile", label: "AWS profile" }, + { id: "credential-chain", label: "Existing AWS credential chain" }, + ], + }); + if (method === "bearer-token") { + return { + type: "api_key", + key: await interaction.prompt({ type: "secret", message: "Enter Amazon Bedrock bearer token" }), + }; + } + interaction.notify({ + type: "info", + message: "Amazon Bedrock supports AWS profiles, IAM credentials, and role-based credentials.", + links: [ + { + label: "AWS credential provider chain", + url: "https://docs.aws.amazon.com/sdkref/latest/guide/standardized-credentials.html", + }, + ], + }); + if (method === "aws-profile") { + return { + type: "api_key", + env: { AWS_PROFILE: await interaction.prompt({ type: "text", message: "Enter AWS profile name" }) }, + }; + } + if (method !== "credential-chain") throw new Error(`Unknown Amazon Bedrock auth method: ${method}`); + await interaction.prompt({ + type: "text", + message: "Configure AWS credentials, then press Enter to continue", + }); + return { type: "api_key" }; + }, resolve: async ({ ctx, credential }) => { if (credential?.key) return { auth: { apiKey: credential.key }, source: "stored credential" }; if (await ctx.env("AWS_BEARER_TOKEN_BEDROCK")) return { auth: {}, source: "AWS_BEARER_TOKEN_BEDROCK" }; - if (await ctx.env("AWS_PROFILE")) return { auth: {}, source: "AWS_PROFILE" }; + if (credential?.env?.AWS_PROFILE ?? (await ctx.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"))) { return { auth: {}, source: "AWS access keys" }; } diff --git a/packages/ai/src/providers/anthropic.ts b/packages/ai/src/providers/anthropic.ts index 6570fc38..ce0b5b9d 100644 --- a/packages/ai/src/providers/anthropic.ts +++ b/packages/ai/src/providers/anthropic.ts @@ -1,7 +1,7 @@ import { anthropicMessagesApi } from "../api/anthropic-messages.lazy.ts"; import { envApiKeyAuth, lazyOAuth } from "../auth/helpers.ts"; +import { loadAnthropicOAuth } from "../auth/oauth/load.ts"; import { createProvider, type Provider } from "../models.ts"; -import { loadAnthropicOAuth } from "../utils/oauth/load.ts"; import { ANTHROPIC_MODELS } from "./anthropic.models.ts"; export function anthropicProvider(): Provider<"anthropic-messages"> { diff --git a/packages/ai/src/providers/azure-openai-responses.models.ts b/packages/ai/src/providers/azure-openai-responses.models.ts index 7ec9ef68..96f00b24 100644 --- a/packages/ai/src/providers/azure-openai-responses.models.ts +++ b/packages/ai/src/providers/azure-openai-responses.models.ts @@ -660,6 +660,23 @@ export const AZURE_OPENAI_RESPONSES_MODELS = { contextWindow: 1050000, maxTokens: 128000, } satisfies Model<"azure-openai-responses">, + "gpt-realtime-2.1": { + id: "gpt-realtime-2.1", + name: "GPT-Realtime-2.1", + api: "azure-openai-responses", + provider: "azure-openai-responses", + baseUrl: "", + reasoning: true, + input: ["text", "image"], + cost: { + input: 4, + output: 24, + cacheRead: 0.4, + cacheWrite: 0, + }, + contextWindow: 128000, + maxTokens: 32000, + } satisfies Model<"azure-openai-responses">, "o1": { id: "o1", name: "o1", diff --git a/packages/ai/src/providers/cerebras.models.ts b/packages/ai/src/providers/cerebras.models.ts index 24c35bae..7428d537 100644 --- a/packages/ai/src/providers/cerebras.models.ts +++ b/packages/ai/src/providers/cerebras.models.ts @@ -52,7 +52,7 @@ export const CEREBRAS_MODELS = { cost: { input: 2.25, output: 2.75, - cacheRead: 0, + cacheRead: 2.25, cacheWrite: 0, }, contextWindow: 131072, diff --git a/packages/ai/src/providers/cloudflare-ai-gateway.models.ts b/packages/ai/src/providers/cloudflare-ai-gateway.models.ts index c0e1033c..dd049567 100644 --- a/packages/ai/src/providers/cloudflare-ai-gateway.models.ts +++ b/packages/ai/src/providers/cloudflare-ai-gateway.models.ts @@ -528,6 +528,60 @@ export const CLOUDFLARE_AI_GATEWAY_MODELS = { contextWindow: 1050000, maxTokens: 128000, } satisfies Model<"openai-responses">, + "gpt-5.6-luna": { + id: "gpt-5.6-luna", + name: "GPT-5.6 Luna", + api: "openai-responses", + provider: "cloudflare-ai-gateway", + baseUrl: "https://gateway.ai.cloudflare.com/v1/{CLOUDFLARE_ACCOUNT_ID}/{CLOUDFLARE_GATEWAY_ID}/openai", + reasoning: true, + thinkingLevelMap: {"off":null,"xhigh":"xhigh","max":"max"}, + input: ["text", "image"], + cost: { + input: 1, + output: 6, + cacheRead: 0.1, + cacheWrite: 0, + }, + contextWindow: 1050000, + maxTokens: 128000, + } satisfies Model<"openai-responses">, + "gpt-5.6-sol": { + id: "gpt-5.6-sol", + name: "GPT-5.6 Sol", + api: "openai-responses", + provider: "cloudflare-ai-gateway", + baseUrl: "https://gateway.ai.cloudflare.com/v1/{CLOUDFLARE_ACCOUNT_ID}/{CLOUDFLARE_GATEWAY_ID}/openai", + reasoning: true, + thinkingLevelMap: {"off":null,"xhigh":"xhigh","max":"max"}, + input: ["text", "image"], + cost: { + input: 5, + output: 30, + cacheRead: 0.5, + cacheWrite: 0, + }, + contextWindow: 1050000, + maxTokens: 128000, + } satisfies Model<"openai-responses">, + "gpt-5.6-terra": { + id: "gpt-5.6-terra", + name: "GPT-5.6 Terra", + api: "openai-responses", + provider: "cloudflare-ai-gateway", + baseUrl: "https://gateway.ai.cloudflare.com/v1/{CLOUDFLARE_ACCOUNT_ID}/{CLOUDFLARE_GATEWAY_ID}/openai", + reasoning: true, + thinkingLevelMap: {"off":null,"xhigh":"xhigh","max":"max"}, + input: ["text", "image"], + cost: { + input: 2.5, + output: 15, + cacheRead: 0.25, + cacheWrite: 0, + }, + contextWindow: 1050000, + maxTokens: 128000, + } satisfies Model<"openai-responses">, "o1": { id: "o1", name: "o1", @@ -685,4 +739,22 @@ export const CLOUDFLARE_AI_GATEWAY_MODELS = { contextWindow: 131072, maxTokens: 131072, } satisfies Model<"openai-completions">, + "workers-ai/@cf/zai-org/glm-5.2": { + id: "workers-ai/@cf/zai-org/glm-5.2", + name: "Glm 5.2", + api: "openai-completions", + provider: "cloudflare-ai-gateway", + baseUrl: "https://gateway.ai.cloudflare.com/v1/{CLOUDFLARE_ACCOUNT_ID}/{CLOUDFLARE_GATEWAY_ID}/compat", + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false,"sendSessionAffinityHeaders":true}, + reasoning: true, + input: ["text"], + cost: { + input: 1.4, + output: 4.4, + cacheRead: 0.26, + cacheWrite: 0, + }, + contextWindow: 262144, + maxTokens: 262144, + } satisfies Model<"openai-completions">, } as const; diff --git a/packages/ai/src/providers/cloudflare-ai-gateway.ts b/packages/ai/src/providers/cloudflare-ai-gateway.ts index 9f6ff5f6..50c10569 100644 --- a/packages/ai/src/providers/cloudflare-ai-gateway.ts +++ b/packages/ai/src/providers/cloudflare-ai-gateway.ts @@ -4,6 +4,7 @@ import { openAIResponsesApi } from "../api/openai-responses.lazy.ts"; import { createProvider, type Provider } from "../models.ts"; import { CLOUDFLARE_AI_GATEWAY_MODELS } from "./cloudflare-ai-gateway.models.ts"; import { cloudflareAIGatewayAuth } from "./cloudflare-auth.ts"; +import { cloudflareStreams } from "./cloudflare-stream.ts"; export function cloudflareAIGatewayProvider(): Provider< "anthropic-messages" | "openai-completions" | "openai-responses" @@ -14,9 +15,9 @@ export function cloudflareAIGatewayProvider(): Provider< auth: { apiKey: cloudflareAIGatewayAuth() }, models: Object.values(CLOUDFLARE_AI_GATEWAY_MODELS), api: { - "anthropic-messages": anthropicMessagesApi(), - "openai-completions": openAICompletionsApi(), - "openai-responses": openAIResponsesApi(), + "anthropic-messages": cloudflareStreams(anthropicMessagesApi()), + "openai-completions": cloudflareStreams(openAICompletionsApi()), + "openai-responses": cloudflareStreams(openAIResponsesApi()), }, }); } diff --git a/packages/ai/src/providers/cloudflare-auth.ts b/packages/ai/src/providers/cloudflare-auth.ts index 511e8d73..0d220ec6 100644 --- a/packages/ai/src/providers/cloudflare-auth.ts +++ b/packages/ai/src/providers/cloudflare-auth.ts @@ -1,5 +1,5 @@ import type { ApiKeyAuth, ApiKeyCredential, AuthContext } from "../auth/types.ts"; -import type { Api, ImagesApi, ImagesModel, Model, ProviderEnv } from "../types.ts"; +import type { ProviderEnv } from "../types.ts"; const CLOUDFLARE_API_KEY = "CLOUDFLARE_API_KEY"; const CLOUDFLARE_ACCOUNT_ID = "CLOUDFLARE_ACCOUNT_ID"; @@ -19,22 +19,11 @@ async function resolveValue( return ctx.env(name); } -function resolveCloudflareBaseUrl( - model: Model | ImagesModel, - accountId: string, - gatewayId: string | undefined, -): string { - return model.baseUrl - .replaceAll(`{${CLOUDFLARE_ACCOUNT_ID}}`, accountId) - .replaceAll(`{${CLOUDFLARE_GATEWAY_ID}}`, gatewayId ?? ""); -} - async function resolveCloudflareEnv( kind: CloudflareAuthKind, - model: Model | ImagesModel, ctx: AuthContext, credential: ApiKeyCredential | undefined, -): Promise<{ apiKey: string; env: ProviderEnv; baseUrl: string; source: string } | undefined> { +): 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; @@ -47,7 +36,6 @@ async function resolveCloudflareEnv( CLOUDFLARE_ACCOUNT_ID: accountId, ...(gatewayId ? { CLOUDFLARE_GATEWAY_ID: gatewayId } : {}), }, - baseUrl: resolveCloudflareBaseUrl(model, accountId, gatewayId), source: credential ? "stored credential" : CLOUDFLARE_API_KEY, }; } @@ -55,16 +43,16 @@ async function resolveCloudflareEnv( export function cloudflareWorkersAIAuth(): ApiKeyAuth { return { name: "Cloudflare API key", - login: async (callbacks) => { - const key = await callbacks.prompt({ type: "secret", message: "Enter Cloudflare API key" }); - const accountId = await callbacks.prompt({ type: "text", message: "Enter Cloudflare account ID" }); + login: async (interaction) => { + const key = await interaction.prompt({ type: "secret", message: "Enter Cloudflare API key" }); + const accountId = await interaction.prompt({ type: "text", message: "Enter Cloudflare account ID" }); return { type: "api_key", key, env: { CLOUDFLARE_ACCOUNT_ID: accountId } }; }, - resolve: async ({ model, ctx, credential }) => { - const resolved = await resolveCloudflareEnv("workers-ai", model, ctx, credential); + resolve: async ({ ctx, credential }) => { + const resolved = await resolveCloudflareEnv("workers-ai", ctx, credential); if (!resolved) return undefined; return { - auth: { apiKey: resolved.apiKey, baseUrl: resolved.baseUrl }, + auth: { apiKey: resolved.apiKey }, env: resolved.env, source: resolved.source, }; @@ -75,18 +63,18 @@ export function cloudflareWorkersAIAuth(): ApiKeyAuth { export function cloudflareAIGatewayAuth(): ApiKeyAuth { return { name: "Cloudflare API key", - login: async (callbacks) => { - const key = await callbacks.prompt({ type: "secret", message: "Enter Cloudflare API key" }); - const accountId = await callbacks.prompt({ type: "text", message: "Enter Cloudflare account ID" }); - const gatewayId = await callbacks.prompt({ type: "text", message: "Enter Cloudflare AI Gateway ID" }); + login: async (interaction) => { + const key = await interaction.prompt({ type: "secret", message: "Enter Cloudflare API key" }); + const accountId = await interaction.prompt({ type: "text", message: "Enter Cloudflare account ID" }); + const gatewayId = await interaction.prompt({ type: "text", message: "Enter Cloudflare AI Gateway ID" }); return { type: "api_key", key, env: { CLOUDFLARE_ACCOUNT_ID: accountId, CLOUDFLARE_GATEWAY_ID: gatewayId }, }; }, - resolve: async ({ model, ctx, credential }) => { - const resolved = await resolveCloudflareEnv("ai-gateway", model, ctx, credential); + resolve: async ({ ctx, credential }) => { + const resolved = await resolveCloudflareEnv("ai-gateway", ctx, credential); if (!resolved) return undefined; return { auth: { @@ -95,7 +83,6 @@ export function cloudflareAIGatewayAuth(): ApiKeyAuth { Authorization: null, "x-api-key": null, }, - baseUrl: resolved.baseUrl, }, env: resolved.env, source: resolved.source, diff --git a/packages/ai/src/providers/cloudflare-stream.ts b/packages/ai/src/providers/cloudflare-stream.ts new file mode 100644 index 00000000..7284f656 --- /dev/null +++ b/packages/ai/src/providers/cloudflare-stream.ts @@ -0,0 +1,28 @@ +import type { Api, Model, ProviderEnv, ProviderStreams } from "../types.ts"; + +const CLOUDFLARE_ACCOUNT_ID = "CLOUDFLARE_ACCOUNT_ID"; +const CLOUDFLARE_GATEWAY_ID = "CLOUDFLARE_GATEWAY_ID"; + +export function resolveCloudflareModel( + model: Model, + env: ProviderEnv | undefined, +): Model { + if (!env) return model; + const baseUrl = model.baseUrl + .replaceAll(`{${CLOUDFLARE_ACCOUNT_ID}}`, env[CLOUDFLARE_ACCOUNT_ID] ?? `{${CLOUDFLARE_ACCOUNT_ID}}`) + .replaceAll(`{${CLOUDFLARE_GATEWAY_ID}}`, env[CLOUDFLARE_GATEWAY_ID] ?? `{${CLOUDFLARE_GATEWAY_ID}}`); + return baseUrl === model.baseUrl ? model : { ...model, baseUrl }; +} + +/** + * Wrap an API implementation so Cloudflare account/gateway endpoint + * placeholders materialize from the resolved provider env before dispatch. + */ +export function cloudflareStreams(streams: ProviderStreams): ProviderStreams { + return { + stream: (model, context, options) => + streams.stream(resolveCloudflareModel(model, options?.env), context, options), + streamSimple: (model, context, options) => + streams.streamSimple(resolveCloudflareModel(model, options?.env), context, options), + }; +} diff --git a/packages/ai/src/providers/cloudflare-workers-ai.ts b/packages/ai/src/providers/cloudflare-workers-ai.ts index 9e376a5c..89216b5e 100644 --- a/packages/ai/src/providers/cloudflare-workers-ai.ts +++ b/packages/ai/src/providers/cloudflare-workers-ai.ts @@ -1,6 +1,7 @@ import { openAICompletionsApi } from "../api/openai-completions.lazy.ts"; import { createProvider, type Provider } from "../models.ts"; import { cloudflareWorkersAIAuth } from "./cloudflare-auth.ts"; +import { cloudflareStreams } from "./cloudflare-stream.ts"; import { CLOUDFLARE_WORKERS_AI_MODELS } from "./cloudflare-workers-ai.models.ts"; export function cloudflareWorkersAIProvider(): Provider<"openai-completions"> { @@ -9,6 +10,6 @@ export function cloudflareWorkersAIProvider(): Provider<"openai-completions"> { name: "Cloudflare Workers AI", auth: { apiKey: cloudflareWorkersAIAuth() }, models: Object.values(CLOUDFLARE_WORKERS_AI_MODELS), - api: openAICompletionsApi(), + api: cloudflareStreams(openAICompletionsApi()), }); } diff --git a/packages/ai/src/providers/github-copilot.models.ts b/packages/ai/src/providers/github-copilot.models.ts index 47a2f3d9..c6b82331 100644 --- a/packages/ai/src/providers/github-copilot.models.ts +++ b/packages/ai/src/providers/github-copilot.models.ts @@ -252,7 +252,7 @@ export const GITHUB_COPILOT_MODELS = { cacheRead: 0.2, cacheWrite: 0, }, - contextWindow: 200000, + contextWindow: 1000000, maxTokens: 64000, } satisfies Model<"openai-completions">, "gemini-3.5-flash": { @@ -445,6 +445,63 @@ export const GITHUB_COPILOT_MODELS = { contextWindow: 1000000, maxTokens: 128000, } satisfies Model<"openai-responses">, + "gpt-5.6-luna": { + id: "gpt-5.6-luna", + name: "GPT-5.6 Luna", + api: "openai-responses", + provider: "github-copilot", + baseUrl: "https://api.individual.githubcopilot.com", + headers: {"User-Agent":"GitHubCopilotChat/0.35.0","Editor-Version":"vscode/1.107.0","Editor-Plugin-Version":"copilot-chat/0.35.0","Copilot-Integration-Id":"vscode-chat"}, + reasoning: true, + thinkingLevelMap: {"off":null,"minimal":"low","xhigh":"xhigh","max":"max"}, + input: ["text", "image"], + cost: { + input: 1, + output: 6, + cacheRead: 0.1, + cacheWrite: 0, + }, + contextWindow: 1050000, + maxTokens: 128000, + } satisfies Model<"openai-responses">, + "gpt-5.6-sol": { + id: "gpt-5.6-sol", + name: "GPT-5.6 Sol", + api: "openai-responses", + provider: "github-copilot", + baseUrl: "https://api.individual.githubcopilot.com", + headers: {"User-Agent":"GitHubCopilotChat/0.35.0","Editor-Version":"vscode/1.107.0","Editor-Plugin-Version":"copilot-chat/0.35.0","Copilot-Integration-Id":"vscode-chat"}, + reasoning: true, + thinkingLevelMap: {"off":null,"minimal":"low","xhigh":"xhigh","max":"max"}, + input: ["text", "image"], + cost: { + input: 5, + output: 30, + cacheRead: 0.5, + cacheWrite: 0, + }, + contextWindow: 1050000, + maxTokens: 128000, + } satisfies Model<"openai-responses">, + "gpt-5.6-terra": { + id: "gpt-5.6-terra", + name: "GPT-5.6 Terra", + api: "openai-responses", + provider: "github-copilot", + baseUrl: "https://api.individual.githubcopilot.com", + headers: {"User-Agent":"GitHubCopilotChat/0.35.0","Editor-Version":"vscode/1.107.0","Editor-Plugin-Version":"copilot-chat/0.35.0","Copilot-Integration-Id":"vscode-chat"}, + reasoning: true, + thinkingLevelMap: {"off":null,"minimal":"low","xhigh":"xhigh","max":"max"}, + input: ["text", "image"], + cost: { + input: 2.5, + output: 15, + cacheRead: 0.25, + cacheWrite: 0, + }, + contextWindow: 1050000, + maxTokens: 128000, + } satisfies Model<"openai-responses">, "kimi-k2.7-code": { id: "kimi-k2.7-code", name: "Kimi K2.7 Code", diff --git a/packages/ai/src/providers/github-copilot.ts b/packages/ai/src/providers/github-copilot.ts index c935ad5d..f5c5d7c5 100644 --- a/packages/ai/src/providers/github-copilot.ts +++ b/packages/ai/src/providers/github-copilot.ts @@ -2,8 +2,8 @@ import { anthropicMessagesApi } from "../api/anthropic-messages.lazy.ts"; import { openAICompletionsApi } from "../api/openai-completions.lazy.ts"; import { openAIResponsesApi } from "../api/openai-responses.lazy.ts"; import { envApiKeyAuth, lazyOAuth } from "../auth/helpers.ts"; +import { loadGitHubCopilotOAuth } from "../auth/oauth/load.ts"; import { createProvider, type Provider } from "../models.ts"; -import { loadGitHubCopilotOAuth } from "../utils/oauth/load.ts"; import { GITHUB_COPILOT_MODELS } from "./github-copilot.models.ts"; export function githubCopilotProvider(): Provider<"anthropic-messages" | "openai-completions" | "openai-responses"> { @@ -16,6 +16,15 @@ export function githubCopilotProvider(): Provider<"anthropic-messages" | "openai oauth: lazyOAuth({ name: "GitHub Copilot", load: loadGitHubCopilotOAuth }), }, models: Object.values(GITHUB_COPILOT_MODELS), + filterModels: (models, credential) => { + if (credential?.type !== "oauth") return models; + const availableModelIds = credential.availableModelIds; + if (!Array.isArray(availableModelIds) || !availableModelIds.every((id) => typeof id === "string")) { + return models; + } + const available = new Set(availableModelIds); + return models.filter((model) => available.has(model.id)); + }, api: { "anthropic-messages": anthropicMessagesApi(), "openai-completions": openAICompletionsApi(), diff --git a/packages/ai/src/providers/google-vertex.ts b/packages/ai/src/providers/google-vertex.ts index af84fc70..e66aa6ef 100644 --- a/packages/ai/src/providers/google-vertex.ts +++ b/packages/ai/src/providers/google-vertex.ts @@ -12,16 +12,71 @@ const VERTEX_ADC_PATH = "~/.config/gcloud/application_default_credentials.json"; */ const vertexAuth: ApiKeyAuth = { name: "Google Cloud credentials", + login: async (interaction) => { + const method = await interaction.prompt({ + type: "select", + message: "Select Google Vertex AI authentication method:", + options: [ + { id: "api-key", label: "Google Cloud API key" }, + { id: "adc", label: "Application Default Credentials" }, + { id: "service-account", label: "Service account credentials file" }, + ], + }); + if (method === "api-key") { + return { + type: "api_key", + key: await interaction.prompt({ type: "secret", message: "Enter Google Cloud API key" }), + }; + } + if (method !== "adc" && method !== "service-account") { + throw new Error(`Unknown Google Vertex AI auth method: ${method}`); + } + interaction.notify({ + type: "info", + message: + method === "adc" + ? "Run `gcloud auth application-default login`, then provide the project and location." + : "Provide a service account credentials file, project, and location.", + links: [ + { + label: "Application Default Credentials", + url: "https://cloud.google.com/docs/authentication/provide-credentials-adc", + }, + ], + }); + const project = await interaction.prompt({ type: "text", message: "Enter Google Cloud project ID" }); + const location = await interaction.prompt({ type: "text", message: "Enter Google Cloud location" }); + const credentialsPath = + method === "service-account" + ? await interaction.prompt({ type: "text", message: "Enter service account credentials file path" }) + : undefined; + return { + type: "api_key", + env: { + GOOGLE_CLOUD_PROJECT: project, + GOOGLE_CLOUD_LOCATION: location, + ...(credentialsPath ? { GOOGLE_APPLICATION_CREDENTIALS: credentialsPath } : {}), + }, + }; + }, resolve: async ({ ctx, credential }) => { const key = credential?.key ?? (await ctx.env("GOOGLE_CLOUD_API_KEY")); if (key) return { auth: { apiKey: key }, source: credential?.key ? "stored credential" : "GOOGLE_CLOUD_API_KEY" }; - const adcPath = await ctx.env("GOOGLE_APPLICATION_CREDENTIALS"); + const adcPath = + credential?.env?.GOOGLE_APPLICATION_CREDENTIALS ?? (await ctx.env("GOOGLE_APPLICATION_CREDENTIALS")); const hasCredentials = await ctx.fileExists(adcPath ?? VERTEX_ADC_PATH); - const hasProject = Boolean((await ctx.env("GOOGLE_CLOUD_PROJECT")) ?? (await ctx.env("GCLOUD_PROJECT"))); - const hasLocation = Boolean(await ctx.env("GOOGLE_CLOUD_LOCATION")); - if (hasCredentials && hasProject && hasLocation) { - return { auth: {}, source: "gcloud application default credentials" }; + 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")); + if (hasCredentials && project && location) { + return { + auth: {}, + env: credential?.env, + source: credential ? "stored credential" : "gcloud application default credentials", + }; } return undefined; }, diff --git a/packages/ai/src/providers/openai-codex.ts b/packages/ai/src/providers/openai-codex.ts index 6ccdb6ef..cbf9d792 100644 --- a/packages/ai/src/providers/openai-codex.ts +++ b/packages/ai/src/providers/openai-codex.ts @@ -1,7 +1,7 @@ import { openAICodexResponsesApi } from "../api/openai-codex-responses.lazy.ts"; import { lazyOAuth } from "../auth/helpers.ts"; +import { loadOpenAICodexOAuth } from "../auth/oauth/load.ts"; import { createProvider, type Provider } from "../models.ts"; -import { loadOpenAICodexOAuth } from "../utils/oauth/load.ts"; import { OPENAI_CODEX_MODELS } from "./openai-codex.models.ts"; export function openaiCodexProvider(): Provider<"openai-codex-responses"> { diff --git a/packages/ai/src/providers/openai.models.ts b/packages/ai/src/providers/openai.models.ts index 1ea1d15b..5f5846f6 100644 --- a/packages/ai/src/providers/openai.models.ts +++ b/packages/ai/src/providers/openai.models.ts @@ -667,6 +667,23 @@ export const OPENAI_MODELS = { contextWindow: 272000, maxTokens: 128000, } satisfies Model<"openai-responses">, + "gpt-realtime-2.1": { + id: "gpt-realtime-2.1", + name: "GPT-Realtime-2.1", + api: "openai-responses", + provider: "openai", + baseUrl: "https://api.openai.com/v1", + reasoning: true, + input: ["text", "image"], + cost: { + input: 4, + output: 24, + cacheRead: 0.4, + cacheWrite: 0, + }, + contextWindow: 128000, + maxTokens: 32000, + } satisfies Model<"openai-responses">, "o1": { id: "o1", name: "o1", diff --git a/packages/ai/src/providers/opencode.models.ts b/packages/ai/src/providers/opencode.models.ts index ad18727a..77b22513 100644 --- a/packages/ai/src/providers/opencode.models.ts +++ b/packages/ai/src/providers/opencode.models.ts @@ -674,6 +674,60 @@ export const OPENCODE_MODELS = { contextWindow: 1050000, maxTokens: 128000, } satisfies Model<"openai-responses">, + "gpt-5.6-luna": { + id: "gpt-5.6-luna", + name: "GPT-5.6 Luna", + api: "openai-responses", + provider: "opencode", + baseUrl: "https://opencode.ai/zen/v1", + reasoning: true, + thinkingLevelMap: {"off":null,"xhigh":"xhigh","max":"max"}, + input: ["text", "image"], + cost: { + input: 1, + output: 6, + cacheRead: 0.1, + cacheWrite: 1.25, + }, + contextWindow: 1050000, + maxTokens: 128000, + } satisfies Model<"openai-responses">, + "gpt-5.6-sol": { + id: "gpt-5.6-sol", + name: "GPT-5.6 Sol", + api: "openai-responses", + provider: "opencode", + baseUrl: "https://opencode.ai/zen/v1", + reasoning: true, + thinkingLevelMap: {"off":null,"xhigh":"xhigh","max":"max"}, + input: ["text", "image"], + cost: { + input: 5, + output: 30, + cacheRead: 0.5, + cacheWrite: 6.25, + }, + contextWindow: 1050000, + maxTokens: 128000, + } satisfies Model<"openai-responses">, + "gpt-5.6-terra": { + id: "gpt-5.6-terra", + name: "GPT-5.6 Terra", + api: "openai-responses", + provider: "opencode", + baseUrl: "https://opencode.ai/zen/v1", + reasoning: true, + thinkingLevelMap: {"off":null,"xhigh":"xhigh","max":"max"}, + input: ["text", "image"], + cost: { + input: 2.5, + output: 15, + cacheRead: 0.25, + cacheWrite: 3.125, + }, + contextWindow: 1050000, + maxTokens: 128000, + } satisfies Model<"openai-responses">, "grok-4.5": { id: "grok-4.5", name: "Grok 4.5", @@ -726,7 +780,7 @@ export const OPENCODE_MODELS = { cacheRead: 0, cacheWrite: 0, }, - contextWindow: 256000, + contextWindow: 190000, maxTokens: 64000, } satisfies Model<"openai-completions">, "kimi-k2.5": { diff --git a/packages/ai/src/providers/openrouter.models.ts b/packages/ai/src/providers/openrouter.models.ts index f2b0eebb..5f7db58f 100644 --- a/packages/ai/src/providers/openrouter.models.ts +++ b/packages/ai/src/providers/openrouter.models.ts @@ -461,24 +461,6 @@ export const OPENROUTER_MODELS = { contextWindow: 262144, maxTokens: 80000, } satisfies Model<"openai-completions">, - "arcee-ai/trinity-mini": { - id: "arcee-ai/trinity-mini", - name: "Arcee AI: Trinity Mini", - api: "openai-completions", - provider: "openrouter", - baseUrl: "https://openrouter.ai/api/v1", - compat: {"supportsDeveloperRole":false,"thinkingFormat":"openrouter"}, - reasoning: true, - input: ["text"], - cost: { - input: 0.045, - output: 0.15, - cacheRead: 0, - cacheWrite: 0, - }, - contextWindow: 131072, - maxTokens: 131072, - } satisfies Model<"openai-completions">, "arcee-ai/virtuoso-large": { id: "arcee-ai/virtuoso-large", name: "Arcee AI: Virtuoso Large", @@ -687,8 +669,8 @@ export const OPENROUTER_MODELS = { reasoning: true, input: ["text"], cost: { - input: 0.21, - output: 0.79, + input: 0.25, + output: 0.95, cacheRead: 0.13, cacheWrite: 0, }, @@ -759,9 +741,9 @@ export const OPENROUTER_MODELS = { reasoning: true, input: ["text"], cost: { - input: 0.2288, - output: 0.3432, - cacheRead: 0.02288, + input: 0.2145, + output: 0.32175, + cacheRead: 0.02145, cacheWrite: 0, }, contextWindow: 131072, @@ -1121,13 +1103,13 @@ export const OPENROUTER_MODELS = { reasoning: true, input: ["text", "image"], cost: { - input: 0.12, + input: 0.06, output: 0.35, - cacheRead: 0.09, + cacheRead: 0, cacheWrite: 0, }, contextWindow: 262144, - maxTokens: 262144, + maxTokens: 8192, } satisfies Model<"openai-completions">, "google/gemma-4-31b-it:free": { id: "google/gemma-4-31b-it:free", @@ -1145,7 +1127,7 @@ export const OPENROUTER_MODELS = { cacheWrite: 0, }, contextWindow: 262144, - maxTokens: 8192, + maxTokens: 32768, } satisfies Model<"openai-completions">, "ibm-granite/granite-4.1-8b": { id: "ibm-granite/granite-4.1-8b", @@ -1238,6 +1220,24 @@ export const OPENROUTER_MODELS = { contextWindow: 262144, maxTokens: 65536, } satisfies Model<"openai-completions">, + "kwaipilot/kat-coder-air-v2.5": { + id: "kwaipilot/kat-coder-air-v2.5", + name: "Kwaipilot: KAT-Coder-Air V2.5", + api: "openai-completions", + provider: "openrouter", + baseUrl: "https://openrouter.ai/api/v1", + compat: {"supportsDeveloperRole":false,"thinkingFormat":"openrouter"}, + reasoning: false, + input: ["text"], + cost: { + input: 0.15, + output: 0.6, + cacheRead: 0.03, + cacheWrite: 0, + }, + contextWindow: 256000, + maxTokens: 80000, + } satisfies Model<"openai-completions">, "kwaipilot/kat-coder-pro-v2": { id: "kwaipilot/kat-coder-pro-v2", name: "Kwaipilot: KAT-Coder-Pro V2", @@ -1256,23 +1256,23 @@ export const OPENROUTER_MODELS = { contextWindow: 256000, maxTokens: 80000, } satisfies Model<"openai-completions">, - "liquid/lfm-2.5-1.2b-thinking:free": { - id: "liquid/lfm-2.5-1.2b-thinking:free", - name: "LiquidAI: LFM2.5-1.2B-Thinking (free)", + "kwaipilot/kat-coder-pro-v2.5": { + id: "kwaipilot/kat-coder-pro-v2.5", + name: "Kwaipilot: KAT-Coder-Pro V2.5", api: "openai-completions", provider: "openrouter", baseUrl: "https://openrouter.ai/api/v1", compat: {"supportsDeveloperRole":false,"thinkingFormat":"openrouter"}, - reasoning: true, + reasoning: false, input: ["text"], cost: { - input: 0, - output: 0, - cacheRead: 0, + input: 0.74, + output: 2.96, + cacheRead: 0.15, cacheWrite: 0, }, - contextWindow: 32768, - maxTokens: 4096, + contextWindow: 256000, + maxTokens: 80000, } satisfies Model<"openai-completions">, "meta-llama/llama-3.1-70b-instruct": { id: "meta-llama/llama-3.1-70b-instruct", @@ -1356,8 +1356,8 @@ export const OPENROUTER_MODELS = { reasoning: false, input: ["text", "image"], cost: { - input: 0.15, - output: 0.6, + input: 0.2, + output: 0.8, cacheRead: 0, cacheWrite: 0, }, @@ -1878,9 +1878,9 @@ export const OPENROUTER_MODELS = { reasoning: true, input: ["text", "image"], cost: { - input: 0.65, + input: 0.66, output: 3.41, - cacheRead: 0.14, + cacheRead: 0.15, cacheWrite: 0, }, contextWindow: 262144, @@ -1896,9 +1896,9 @@ export const OPENROUTER_MODELS = { reasoning: true, input: ["text", "image"], cost: { - input: 0.72, + input: 0.719, output: 3.49, - cacheRead: 0.159, + cacheRead: 0.149, cacheWrite: 0, }, contextWindow: 262144, @@ -2456,11 +2456,11 @@ export const OPENROUTER_MODELS = { cost: { input: 0.05, output: 0.4, - cacheRead: 0.01, + cacheRead: 0.005, cacheWrite: 0, }, contextWindow: 400000, - maxTokens: 4096, + maxTokens: 128000, } satisfies Model<"openai-completions">, "openai/gpt-5-pro": { id: "openai/gpt-5-pro", @@ -2492,7 +2492,7 @@ export const OPENROUTER_MODELS = { cost: { input: 1.25, output: 10, - cacheRead: 0.13, + cacheRead: 0.125, cacheWrite: 0, }, contextWindow: 400000, @@ -2976,26 +2976,8 @@ export const OPENROUTER_MODELS = { reasoning: true, input: ["text"], cost: { - input: 0.036, - output: 0.18, - cacheRead: 0, - cacheWrite: 0, - }, - contextWindow: 131072, - maxTokens: 4096, - } satisfies Model<"openai-completions">, - "openai/gpt-oss-120b:free": { - id: "openai/gpt-oss-120b:free", - name: "OpenAI: gpt-oss-120b (free)", - api: "openai-completions", - provider: "openrouter", - baseUrl: "https://openrouter.ai/api/v1", - compat: {"thinkingFormat":"openrouter"}, - reasoning: true, - input: ["text"], - cost: { - input: 0, - output: 0, + input: 0.03, + output: 0.15, cacheRead: 0, cacheWrite: 0, }, @@ -3481,7 +3463,7 @@ export const OPENROUTER_MODELS = { input: ["text"], cost: { input: 0.09, - output: 0.1, + output: 0.55, cacheRead: 0, cacheWrite: 0, }, @@ -4074,13 +4056,13 @@ export const OPENROUTER_MODELS = { reasoning: true, input: ["text", "image"], cost: { - input: 0.285, + input: 0.289, output: 2.4, - cacheRead: 0.15, + cacheRead: 0, cacheWrite: 0, }, contextWindow: 262144, - maxTokens: 262140, + maxTokens: 131072, } satisfies Model<"openai-completions">, "qwen/qwen3.6-35b-a3b": { id: "qwen/qwen3.6-35b-a3b", @@ -4561,12 +4543,12 @@ export const OPENROUTER_MODELS = { input: ["text"], cost: { input: 0.43, - output: 1.74, + output: 1.75, cacheRead: 0.08, cacheWrite: 0, }, - contextWindow: 202752, - maxTokens: 131072, + contextWindow: 200000, + maxTokens: 16384, } satisfies Model<"openai-completions">, "z-ai/glm-4.6v": { id: "z-ai/glm-4.6v", @@ -4638,7 +4620,7 @@ export const OPENROUTER_MODELS = { cacheWrite: 0, }, contextWindow: 202752, - maxTokens: 4096, + maxTokens: 128000, } satisfies Model<"openai-completions">, "z-ai/glm-5-turbo": { id: "z-ai/glm-5-turbo", @@ -4687,13 +4669,13 @@ export const OPENROUTER_MODELS = { thinkingLevelMap: {"xhigh":"xhigh"}, input: ["text"], cost: { - input: 0.54, - output: 1.76, - cacheRead: 0.1, + input: 0.924, + output: 2.904, + cacheRead: 0.1716, cacheWrite: 0, }, contextWindow: 1048576, - maxTokens: 101376, + maxTokens: 131072, } satisfies Model<"openai-completions">, "z-ai/glm-5v-turbo": { id: "z-ai/glm-5v-turbo", @@ -4831,9 +4813,9 @@ export const OPENROUTER_MODELS = { reasoning: true, input: ["text", "image"], cost: { - input: 0.65, + input: 0.66, output: 3.41, - cacheRead: 0.14, + cacheRead: 0.15, cacheWrite: 0, }, contextWindow: 262144, diff --git a/packages/ai/src/providers/vercel-ai-gateway.models.ts b/packages/ai/src/providers/vercel-ai-gateway.models.ts index 48027249..00dc85eb 100644 --- a/packages/ai/src/providers/vercel-ai-gateway.models.ts +++ b/packages/ai/src/providers/vercel-ai-gateway.models.ts @@ -497,23 +497,6 @@ export const VERCEL_AI_GATEWAY_MODELS = { contextWindow: 200000, maxTokens: 4096, } satisfies Model<"anthropic-messages">, - "anthropic/claude-3.5-haiku": { - id: "anthropic/claude-3.5-haiku", - name: "Claude 3.5 Haiku", - api: "anthropic-messages", - provider: "vercel-ai-gateway", - baseUrl: "https://ai-gateway.vercel.sh", - reasoning: false, - input: ["text", "image"], - cost: { - input: 0.8, - output: 4, - cacheRead: 0.08, - cacheWrite: 1, - }, - contextWindow: 200000, - maxTokens: 8192, - } satisfies Model<"anthropic-messages">, "anthropic/claude-fable-5": { id: "anthropic/claude-fable-5", name: "Claude Fable 5", @@ -565,7 +548,7 @@ export const VERCEL_AI_GATEWAY_MODELS = { cacheWrite: 18.75, }, contextWindow: 200000, - maxTokens: 32000, + maxTokens: 8192, } satisfies Model<"anthropic-messages">, "anthropic/claude-opus-4.1": { id: "anthropic/claude-opus-4.1", @@ -673,7 +656,7 @@ export const VERCEL_AI_GATEWAY_MODELS = { cacheWrite: 3.75, }, contextWindow: 1000000, - maxTokens: 64000, + maxTokens: 8192, } satisfies Model<"anthropic-messages">, "anthropic/claude-sonnet-4.5": { id: "anthropic/claude-sonnet-4.5", @@ -730,23 +713,6 @@ export const VERCEL_AI_GATEWAY_MODELS = { contextWindow: 1000000, maxTokens: 128000, } satisfies Model<"anthropic-messages">, - "arcee-ai/trinity-large-preview": { - id: "arcee-ai/trinity-large-preview", - name: "Trinity Large Preview", - api: "anthropic-messages", - provider: "vercel-ai-gateway", - baseUrl: "https://ai-gateway.vercel.sh", - reasoning: false, - input: ["text"], - cost: { - input: 0.25, - output: 1, - cacheRead: 0, - cacheWrite: 0, - }, - contextWindow: 131000, - maxTokens: 131000, - } satisfies Model<"anthropic-messages">, "arcee-ai/trinity-large-thinking": { id: "arcee-ai/trinity-large-thinking", name: "Trinity Large Thinking", @@ -875,12 +841,12 @@ export const VERCEL_AI_GATEWAY_MODELS = { reasoning: true, input: ["text"], cost: { - input: 0.6, - output: 1.7, - cacheRead: 0, + input: 0.21, + output: 0.79, + cacheRead: 0.13, cacheWrite: 0, }, - contextWindow: 128000, + contextWindow: 163840, maxTokens: 128000, } satisfies Model<"anthropic-messages">, "deepseek/deepseek-v3.1-terminus": { @@ -945,7 +911,7 @@ export const VERCEL_AI_GATEWAY_MODELS = { cost: { input: 0.14, output: 0.28, - cacheRead: 0.0028, + cacheRead: 0.028, cacheWrite: 0, }, contextWindow: 1000000, @@ -1206,6 +1172,23 @@ export const VERCEL_AI_GATEWAY_MODELS = { contextWindow: 1000000, maxTokens: 32000, } satisfies Model<"anthropic-messages">, + "kwaipilot/kat-coder-air-v2.5": { + id: "kwaipilot/kat-coder-air-v2.5", + name: "Kat Coder Air V2.5", + api: "anthropic-messages", + provider: "vercel-ai-gateway", + baseUrl: "https://ai-gateway.vercel.sh", + reasoning: true, + input: ["text"], + cost: { + input: 0.15, + output: 0.6, + cacheRead: 0.03, + cacheWrite: 0, + }, + contextWindow: 256000, + maxTokens: 80000, + } satisfies Model<"anthropic-messages">, "kwaipilot/kat-coder-pro-v1": { id: "kwaipilot/kat-coder-pro-v1", name: "KAT-Coder-Pro V1", @@ -1240,39 +1223,22 @@ export const VERCEL_AI_GATEWAY_MODELS = { contextWindow: 256000, maxTokens: 256000, } satisfies Model<"anthropic-messages">, - "meituan/longcat-flash-chat": { - id: "meituan/longcat-flash-chat", - name: "LongCat Flash Chat", - api: "anthropic-messages", - provider: "vercel-ai-gateway", - baseUrl: "https://ai-gateway.vercel.sh", - reasoning: false, - input: ["text"], - cost: { - input: 0, - output: 0, - cacheRead: 0, - cacheWrite: 0, - }, - contextWindow: 128000, - maxTokens: 100000, - } satisfies Model<"anthropic-messages">, - "meituan/longcat-flash-thinking-2601": { - id: "meituan/longcat-flash-thinking-2601", - name: "LongCat Flash Thinking 2601", + "kwaipilot/kat-coder-pro-v2.5": { + id: "kwaipilot/kat-coder-pro-v2.5", + name: "Kat Coder Pro V2.5", api: "anthropic-messages", provider: "vercel-ai-gateway", baseUrl: "https://ai-gateway.vercel.sh", reasoning: true, input: ["text"], cost: { - input: 0, - output: 0, - cacheRead: 0, + input: 0.74, + output: 2.96, + cacheRead: 0.15, cacheWrite: 0, }, - contextWindow: 32768, - maxTokens: 32768, + contextWindow: 256000, + maxTokens: 80000, } satisfies Model<"anthropic-messages">, "meta/llama-3.1-70b": { id: "meta/llama-3.1-70b", @@ -1400,7 +1366,7 @@ export const VERCEL_AI_GATEWAY_MODELS = { provider: "vercel-ai-gateway", baseUrl: "https://ai-gateway.vercel.sh", reasoning: true, - input: ["text"], + input: ["text", "image"], cost: { input: 1.25, output: 4.25, @@ -1580,23 +1546,6 @@ export const VERCEL_AI_GATEWAY_MODELS = { contextWindow: 256000, maxTokens: 256000, } satisfies Model<"anthropic-messages">, - "mistral/devstral-small": { - id: "mistral/devstral-small", - name: "Devstral Small 1.1", - api: "anthropic-messages", - provider: "vercel-ai-gateway", - baseUrl: "https://ai-gateway.vercel.sh", - reasoning: false, - input: ["text"], - cost: { - input: 0.1, - output: 0.3, - cacheRead: 0, - cacheWrite: 0, - }, - contextWindow: 128000, - maxTokens: 64000, - } satisfies Model<"anthropic-messages">, "mistral/devstral-small-2": { id: "mistral/devstral-small-2", name: "Devstral Small 2", @@ -1801,23 +1750,6 @@ export const VERCEL_AI_GATEWAY_MODELS = { contextWindow: 128000, maxTokens: 4000, } satisfies Model<"anthropic-messages">, - "mistral/pixtral-large": { - id: "mistral/pixtral-large", - name: "Pixtral Large", - api: "anthropic-messages", - provider: "vercel-ai-gateway", - baseUrl: "https://ai-gateway.vercel.sh", - reasoning: false, - input: ["text", "image"], - cost: { - input: 2, - output: 6, - cacheRead: 0, - cacheWrite: 0, - }, - contextWindow: 128000, - maxTokens: 4000, - } satisfies Model<"anthropic-messages">, "moonshotai/kimi-k2": { id: "moonshotai/kimi-k2", name: "Kimi K2 Instruct", @@ -2972,40 +2904,6 @@ export const VERCEL_AI_GATEWAY_MODELS = { contextWindow: 256000, maxTokens: 256000, } satisfies Model<"anthropic-messages">, - "xiaomi/mimo-v2-flash": { - id: "xiaomi/mimo-v2-flash", - name: "MiMo V2 Flash", - api: "anthropic-messages", - provider: "vercel-ai-gateway", - baseUrl: "https://ai-gateway.vercel.sh", - reasoning: true, - input: ["text"], - cost: { - input: 0.1, - output: 0.3, - cacheRead: 0.01, - cacheWrite: 0, - }, - contextWindow: 262144, - maxTokens: 32000, - } satisfies Model<"anthropic-messages">, - "xiaomi/mimo-v2-pro": { - id: "xiaomi/mimo-v2-pro", - name: "MiMo V2 Pro", - api: "anthropic-messages", - provider: "vercel-ai-gateway", - baseUrl: "https://ai-gateway.vercel.sh", - reasoning: true, - input: ["text"], - cost: { - input: 1, - output: 3, - cacheRead: 0.2, - cacheWrite: 0, - }, - contextWindow: 1000000, - maxTokens: 128000, - } satisfies Model<"anthropic-messages">, "xiaomi/mimo-v2.5": { id: "xiaomi/mimo-v2.5", name: "MiMo M2.5", @@ -3270,9 +3168,9 @@ export const VERCEL_AI_GATEWAY_MODELS = { reasoning: true, input: ["text"], cost: { - input: 3, - output: 10.25, - cacheRead: 0.5, + input: 2.1, + output: 6.6, + cacheRead: 0.21, cacheWrite: 0, }, contextWindow: 1000000, diff --git a/packages/ai/src/utils/oauth/index.ts b/packages/ai/src/utils/oauth/index.ts deleted file mode 100644 index a57badda..00000000 --- a/packages/ai/src/utils/oauth/index.ts +++ /dev/null @@ -1,160 +0,0 @@ -/** - * OAuth credential management for AI providers. - * - * This module handles login, token refresh, and credential storage - * for OAuth-based providers: - * - Anthropic (Claude Pro/Max) - * - GitHub Copilot - */ - -// Anthropic -export { anthropicOAuthProvider, loginAnthropic, refreshAnthropicToken } from "./anthropic.ts"; -export * from "./device-code.ts"; -// GitHub Copilot -export { - getGitHubCopilotBaseUrl, - githubCopilotOAuthProvider, - loginGitHubCopilot, - normalizeDomain, - refreshGitHubCopilotToken, -} from "./github-copilot.ts"; -// OpenAI Codex (ChatGPT OAuth) -export { - loginOpenAICodex, - loginOpenAICodexDeviceCode, - OPENAI_CODEX_BROWSER_LOGIN_METHOD, - OPENAI_CODEX_DEVICE_CODE_LOGIN_METHOD, - openaiCodexOAuthProvider, - refreshOpenAICodexToken, -} from "./openai-codex.ts"; - -export * from "./types.ts"; - -// ============================================================================ -// Provider Registry -// ============================================================================ - -import { anthropicOAuthProvider } from "./anthropic.ts"; -import { githubCopilotOAuthProvider } from "./github-copilot.ts"; -import { openaiCodexOAuthProvider } from "./openai-codex.ts"; -import type { OAuthCredentials, OAuthProviderId, OAuthProviderInfo, OAuthProviderInterface } from "./types.ts"; - -const BUILT_IN_OAUTH_PROVIDERS: OAuthProviderInterface[] = [ - anthropicOAuthProvider, - githubCopilotOAuthProvider, - openaiCodexOAuthProvider, -]; - -const oauthProviderRegistry = new Map( - BUILT_IN_OAUTH_PROVIDERS.map((provider) => [provider.id, provider]), -); - -/** - * Get an OAuth provider by ID - */ -export function getOAuthProvider(id: OAuthProviderId): OAuthProviderInterface | undefined { - return oauthProviderRegistry.get(id); -} - -/** - * Register a custom OAuth provider - */ -export function registerOAuthProvider(provider: OAuthProviderInterface): void { - oauthProviderRegistry.set(provider.id, provider); -} - -/** - * Unregister an OAuth provider. - * - * If the provider is built-in, restores the built-in implementation. - * Custom providers are removed completely. - */ -export function unregisterOAuthProvider(id: string): void { - const builtInProvider = BUILT_IN_OAUTH_PROVIDERS.find((provider) => provider.id === id); - if (builtInProvider) { - oauthProviderRegistry.set(id, builtInProvider); - return; - } - oauthProviderRegistry.delete(id); -} - -/** - * Reset OAuth providers to built-ins. - */ -export function resetOAuthProviders(): void { - oauthProviderRegistry.clear(); - for (const provider of BUILT_IN_OAUTH_PROVIDERS) { - oauthProviderRegistry.set(provider.id, provider); - } -} - -/** - * Get all registered OAuth providers - */ -export function getOAuthProviders(): OAuthProviderInterface[] { - return Array.from(oauthProviderRegistry.values()); -} - -/** - * @deprecated Use getOAuthProviders() which returns OAuthProviderInterface[] - */ -export function getOAuthProviderInfoList(): OAuthProviderInfo[] { - return getOAuthProviders().map((p) => ({ - id: p.id, - name: p.name, - available: true, - })); -} - -// ============================================================================ -// High-level API (uses provider registry) -// ============================================================================ - -/** - * Refresh token for any OAuth provider. - * @deprecated Use getOAuthProvider(id).refreshToken() instead - */ -export async function refreshOAuthToken( - providerId: OAuthProviderId, - credentials: OAuthCredentials, -): Promise { - const provider = getOAuthProvider(providerId); - if (!provider) { - throw new Error(`Unknown OAuth provider: ${providerId}`); - } - return provider.refreshToken(credentials); -} - -/** - * Get API key for a provider from OAuth credentials. - * Automatically refreshes expired tokens. - * - * @returns API key string and updated credentials, or null if no credentials - * @throws Error if refresh fails - */ -export async function getOAuthApiKey( - providerId: OAuthProviderId, - credentials: Record, -): Promise<{ newCredentials: OAuthCredentials; apiKey: string } | null> { - const provider = getOAuthProvider(providerId); - if (!provider) { - throw new Error(`Unknown OAuth provider: ${providerId}`); - } - - let creds = credentials[providerId]; - if (!creds) { - return null; - } - - // Refresh if expired - if (Date.now() >= creds.expires) { - try { - creds = await provider.refreshToken(creds); - } catch (_error) { - throw new Error(`Failed to refresh OAuth token for ${providerId}`); - } - } - - const apiKey = provider.getApiKey(creds); - return { newCredentials: creds, apiKey }; -} diff --git a/packages/ai/src/utils/oauth/types.ts b/packages/ai/src/utils/oauth/types.ts deleted file mode 100644 index 008be405..00000000 --- a/packages/ai/src/utils/oauth/types.ts +++ /dev/null @@ -1,79 +0,0 @@ -import type { Api, Model } from "../../types.ts"; - -export type OAuthCredentials = { - refresh: string; - access: string; - expires: number; - [key: string]: unknown; -}; - -export type OAuthProviderId = string; - -/** @deprecated Use OAuthProviderId instead */ -export type OAuthProvider = OAuthProviderId; - -export type OAuthPrompt = { - message: string; - placeholder?: string; - allowEmpty?: boolean; -}; - -export type OAuthAuthInfo = { - url: string; - instructions?: string; -}; - -export type OAuthDeviceCodeInfo = { - userCode: string; - verificationUri: string; - intervalSeconds?: number; - expiresInSeconds?: number; -}; - -export type OAuthSelectOption = { - id: string; - label: string; -}; - -export type OAuthSelectPrompt = { - message: string; - options: OAuthSelectOption[]; -}; - -export interface OAuthLoginCallbacks { - onAuth: (info: OAuthAuthInfo) => void; - onDeviceCode: (info: OAuthDeviceCodeInfo) => void; - onPrompt: (prompt: OAuthPrompt) => Promise; - onProgress?: (message: string) => void; - onManualCodeInput?: () => Promise; - /** Show an interactive selector and return the selected option id, or undefined on cancel. */ - onSelect: (prompt: OAuthSelectPrompt) => Promise; - signal?: AbortSignal; -} - -export interface OAuthProviderInterface { - readonly id: OAuthProviderId; - readonly name: string; - - /** Run the login flow, return credentials to persist */ - login(callbacks: OAuthLoginCallbacks): Promise; - - /** Whether login uses a local callback server and supports manual code input. */ - usesCallbackServer?: boolean; - - /** Refresh expired credentials, return updated credentials to persist */ - refreshToken(credentials: OAuthCredentials): Promise; - - /** Convert credentials to API key string for the provider */ - getApiKey(credentials: OAuthCredentials): string; - - /** Optional: modify models for this provider (e.g., update baseUrl) */ - modifyModels?(models: Model[], credentials: OAuthCredentials): Model[]; -} - -/** @deprecated Use OAuthProviderInterface instead */ -export interface OAuthProviderInfo { - id: OAuthProviderId; - name: string; - available: boolean; -} diff --git a/packages/ai/test/anthropic-oauth.test.ts b/packages/ai/test/anthropic-oauth.test.ts index ae3ae093..679dd7d5 100644 --- a/packages/ai/test/anthropic-oauth.test.ts +++ b/packages/ai/test/anthropic-oauth.test.ts @@ -1,6 +1,6 @@ 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"; -import { anthropicOAuth, loginAnthropic, refreshAnthropicToken } from "../src/utils/oauth/anthropic.ts"; function jsonResponse(body: unknown, status: number = 200): Response { return new Response(JSON.stringify(body), { @@ -53,18 +53,16 @@ describe.sequential("Anthropic OAuth", () => { }); vi.stubGlobal("fetch", fetchMock); - const credentials = await loginAnthropic({ - onAuth: (info) => { - authUrl = info.url; + const credentials = await anthropicOAuth.login({ + notify: (event) => { + if (event.type === "auth_url") authUrl = event.url; }, - onPrompt: async () => "", - onManualCodeInput: async () => { + prompt: async (prompt) => { + if (prompt.type !== "manual_code") throw new Error(`Unexpected prompt: ${prompt.type}`); const url = new URL(authUrl); const state = url.searchParams.get("state"); const redirectUri = url.searchParams.get("redirect_uri"); - if (!state || !redirectUri) { - throw new Error("Missing OAuth state or redirect_uri in auth URL"); - } + if (!state || !redirectUri) throw new Error("Missing OAuth state or redirect_uri in auth URL"); return `${redirectUri}?code=manual-code&state=${state}`; }, }); @@ -91,7 +89,12 @@ describe.sequential("Anthropic OAuth", () => { }); vi.stubGlobal("fetch", fetchMock); - const credentials = await refreshAnthropicToken("refresh-token"); + const credentials = await anthropicOAuth.refresh({ + type: "oauth", + access: "old-access-token", + refresh: "refresh-token", + expires: 0, + }); expect(credentials.access).toBe("new-access-token"); expect(credentials.refresh).toBe("new-refresh-token"); diff --git a/packages/ai/test/cloudflare-stream.test.ts b/packages/ai/test/cloudflare-stream.test.ts new file mode 100644 index 00000000..34469bae --- /dev/null +++ b/packages/ai/test/cloudflare-stream.test.ts @@ -0,0 +1,65 @@ +import { describe, expect, it } from "vitest"; +import { cloudflareStreams } from "../src/providers/cloudflare-stream.ts"; +import type { Api, Context, Model } from "../src/types.ts"; +import { AssistantMessageEventStream } from "../src/utils/event-stream.ts"; + +const model: Model = { + id: "model", + name: "model", + api: "openai-completions", + provider: "cloudflare-ai-gateway", + baseUrl: "https://gateway.ai.cloudflare.com/v1/{CLOUDFLARE_ACCOUNT_ID}/{CLOUDFLARE_GATEWAY_ID}/openai", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1000, + maxTokens: 100, +}; + +const context: Context = { messages: [] }; + +describe("Cloudflare provider streams", () => { + it("materializes the model endpoint before dispatch", () => { + const captured: string[] = []; + const streams = cloudflareStreams({ + stream: (requestModel) => { + captured.push(requestModel.baseUrl); + return new AssistantMessageEventStream(); + }, + streamSimple: (requestModel) => { + captured.push(requestModel.baseUrl); + return new AssistantMessageEventStream(); + }, + }); + const env = { + CLOUDFLARE_ACCOUNT_ID: "account", + CLOUDFLARE_GATEWAY_ID: "gateway", + }; + + streams.stream(model, context, { env }); + streams.streamSimple(model, context, { env }); + + expect(captured).toEqual([ + "https://gateway.ai.cloudflare.com/v1/account/gateway/openai", + "https://gateway.ai.cloudflare.com/v1/account/gateway/openai", + ]); + }); + + it("keeps placeholders when the provider env does not resolve them", () => { + let captured: string | undefined; + const streams = cloudflareStreams({ + stream: (requestModel) => { + captured = requestModel.baseUrl; + return new AssistantMessageEventStream(); + }, + streamSimple: (requestModel) => { + captured = requestModel.baseUrl; + return new AssistantMessageEventStream(); + }, + }); + + streams.streamSimple(model, context, {}); + + expect(captured).toBe(model.baseUrl); + }); +}); diff --git a/packages/ai/test/codex-websocket-cached-probe.ts b/packages/ai/test/codex-websocket-cached-probe.ts index b5104080..74f8f57e 100644 --- a/packages/ai/test/codex-websocket-cached-probe.ts +++ b/packages/ai/test/codex-websocket-cached-probe.ts @@ -9,7 +9,7 @@ import { tmpdir } from "node:os"; import { join, resolve } from "node:path"; import { Type } from "typebox"; -import { AuthStorage } from "../../coding-agent/src/core/auth-storage.ts"; +import { ModelRuntime } from "../../coding-agent/src/core/model-runtime.ts"; import { closeOpenAICodexWebSocketSessions, getOpenAICodexWebSocketDebugStats, @@ -166,8 +166,9 @@ async function main(): Promise { const model = getModel("openai-codex", "gpt-5.5") as Model<"openai-codex-responses"> | undefined; if (!model) throw new Error("Model openai-codex/gpt-5.5 not found"); const modelWithMaxTokens = { ...model, maxTokens: args.maxTokens }; - const authStorage = AuthStorage.create(); - const apiKey = (await authStorage.getApiKey("openai-codex")) ?? (await authStorage.getApiKey("openai")); + const modelRuntime = await ModelRuntime.create(); + const apiKey = + (await modelRuntime.getAuth("openai-codex"))?.auth.apiKey ?? (await modelRuntime.getAuth("openai"))?.auth.apiKey; if (!apiKey) { throw new Error("No OpenAI Codex API key found in coding-agent auth storage."); } diff --git a/packages/ai/test/github-copilot-oauth.test.ts b/packages/ai/test/github-copilot-oauth.test.ts index 963eef7d..26738f29 100644 --- a/packages/ai/test/github-copilot-oauth.test.ts +++ b/packages/ai/test/github-copilot-oauth.test.ts @@ -1,10 +1,8 @@ import { afterEach, describe, expect, it, vi } from "vitest"; -import { getModels } from "../src/compat.ts"; -import { - githubCopilotOAuthProvider, - loginGitHubCopilot, - refreshGitHubCopilotToken, -} from "../src/utils/oauth/github-copilot.ts"; +import { InMemoryCredentialStore } from "../src/auth/credential-store.ts"; +import { githubCopilotOAuth } from "../src/auth/oauth/github-copilot.ts"; +import { createModels } from "../src/models.ts"; +import { githubCopilotProvider } from "../src/providers/github-copilot.ts"; function jsonResponse(body: unknown, status: number = 200): Response { return new Response(JSON.stringify(body), { @@ -28,6 +26,33 @@ function getUrl(input: unknown): string { throw new Error(`Unsupported fetch input: ${String(input)}`); } +function loginGitHubCopilotForTest(options: { + onDeviceCode(info: { + userCode: string; + verificationUri: string; + intervalSeconds?: number; + expiresInSeconds?: number; + }): void; + onPrompt(prompt: { message: string; placeholder?: string; allowEmpty?: boolean }): Promise; + onProgress?(message: string): void; + signal?: AbortSignal; +}) { + return githubCopilotOAuth.login({ + signal: options.signal, + prompt: (prompt) => { + if (prompt.type !== "text") throw new Error(`Unexpected prompt: ${prompt.type}`); + return options.onPrompt({ message: prompt.message, placeholder: prompt.placeholder, allowEmpty: true }); + }, + notify: (event) => { + if (event.type === "device_code") { + const { type: _, ...info } = event; + options.onDeviceCode(info); + } + if (event.type === "progress") options.onProgress?.(event.message); + }, + }); +} + describe("GitHub Copilot OAuth device flow", () => { afterEach(() => { vi.unstubAllGlobals(); @@ -76,13 +101,19 @@ describe("GitHub Copilot OAuth device flow", () => { vi.stubGlobal("fetch", fetchMock); - const credentials = await refreshGitHubCopilotToken("ghu_refresh_token"); + const credentials = await githubCopilotOAuth.refresh({ + type: "oauth", + access: "old-access-token", + refresh: "ghu_refresh_token", + expires: 0, + }); expect(credentials.availableModelIds).toEqual(["gpt-4.1"]); - const modifiedModels = githubCopilotOAuthProvider.modifyModels?.(getModels("github-copilot"), credentials) ?? []; - expect(modifiedModels.filter((model) => model.provider === "github-copilot").map((model) => model.id)).toEqual([ - "gpt-4.1", - ]); + const store = new InMemoryCredentialStore(); + await store.modify("github-copilot", async () => ({ ...credentials, type: "oauth" })); + const models = createModels({ credentials: store }); + models.setProvider(githubCopilotProvider()); + expect((await models.getAvailable("github-copilot")).map((model) => model.id)).toEqual(["gpt-4.1"]); }); it("reports device-code details through onDeviceCode", async () => { @@ -127,7 +158,7 @@ describe("GitHub Copilot OAuth device flow", () => { vi.stubGlobal("fetch", fetchMock); const onDeviceCode = vi.fn(); - const loginPromise = loginGitHubCopilot({ + const loginPromise = loginGitHubCopilotForTest({ onDeviceCode, onPrompt: async () => "", }); @@ -166,7 +197,7 @@ describe("GitHub Copilot OAuth device flow", () => { const onDeviceCode = vi.fn(); await expect( - loginGitHubCopilot({ + loginGitHubCopilotForTest({ onDeviceCode, onPrompt: async () => "", }), @@ -220,7 +251,7 @@ describe("GitHub Copilot OAuth device flow", () => { vi.stubGlobal("fetch", fetchMock); const onDeviceCode = vi.fn(); - const loginPromise = loginGitHubCopilot({ + const loginPromise = loginGitHubCopilotForTest({ onDeviceCode, onPrompt: async () => "", }); @@ -308,7 +339,7 @@ describe("GitHub Copilot OAuth device flow", () => { vi.stubGlobal("fetch", fetchMock); - const loginPromise = loginGitHubCopilot({ + const loginPromise = loginGitHubCopilotForTest({ onDeviceCode: () => {}, onPrompt: async () => "", onProgress: () => {}, @@ -382,7 +413,7 @@ describe("GitHub Copilot OAuth device flow", () => { vi.stubGlobal("fetch", fetchMock); - const loginPromise = loginGitHubCopilot({ + const loginPromise = loginGitHubCopilotForTest({ onDeviceCode: () => {}, onPrompt: async () => "", }); diff --git a/packages/ai/test/images-models.test.ts b/packages/ai/test/images-models.test.ts index 0a2a9e56..a9dd0407 100644 --- a/packages/ai/test/images-models.test.ts +++ b/packages/ai/test/images-models.test.ts @@ -51,10 +51,10 @@ function testProvider(input: { auth: { apiKey: { name: "Test key", - resolve: async ({ ctx }) => { + resolve: async ({ ctx, credential }) => { if (!input.envVar) return { auth: {} }; - const key = await ctx.env(input.envVar); - return key ? { auth: { apiKey: key }, source: input.envVar } : undefined; + const key = credential?.key ?? (await ctx.env(input.envVar)); + return key ? { auth: { apiKey: key }, source: credential ? "stored" : input.envVar } : undefined; }, }, }, @@ -93,6 +93,8 @@ describe("ImagesModels", () => { const model = models.getModel("p1", "model-a")!; expect((await models.getAuth(model))?.auth.apiKey).toBe("env-key"); + expect((await models.getAuth(model.provider))?.auth.apiKey).toBe("env-key"); + expect((await models.getAuth(model, { apiKey: "explicit-key" }))?.auth.apiKey).toBe("explicit-key"); const result = await models.generateImages(model, context); expect(result.stopReason).toBe("stop"); diff --git a/packages/ai/test/models-runtime.test.ts b/packages/ai/test/models-runtime.test.ts index 19f3f316..00e96d83 100644 --- a/packages/ai/test/models-runtime.test.ts +++ b/packages/ai/test/models-runtime.test.ts @@ -106,6 +106,22 @@ function testOAuth(overrides?: Partial): OAuthAuth { } describe("Models runtime", () => { + it("enumerates credential metadata without exposing secrets", async () => { + const credentials = new InMemoryCredentialStore(); + await credentials.modify("api-provider", async () => ({ type: "api_key", key: "secret" })); + await credentials.modify("oauth-provider", async () => ({ + type: "oauth", + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + })); + + expect(await credentials.list()).toEqual([ + { providerId: "api-provider", type: "api_key" }, + { providerId: "oauth-provider", type: "oauth" }, + ]); + }); + it("applies request-wide pricing tiers above the configured input threshold", () => { const model = testModel("openai", "gpt-5.6-sol"); model.cost = { @@ -246,8 +262,10 @@ describe("Models runtime", () => { models.setProvider(testProvider({ id: "p1", auth: { apiKey: envKeyAuth("env-key"), oauth: testOAuth() } })); const model = testModel("p1", "model-a"); - // nothing stored: ambient env resolves + // model and provider-id overloads resolve the same provider-scoped auth expect((await models.getAuth(model))?.auth.apiKey).toBe("env-key"); + expect((await models.getAuth(model.provider))?.auth.apiKey).toBe("env-key"); + expect((await models.getAuth(model, { apiKey: "explicit-key" }))?.auth.apiKey).toBe("explicit-key"); // stored oauth credential (persisted via the single write path): beats ambient env await credentials.modify("p1", async () => ({ @@ -256,17 +274,69 @@ describe("Models runtime", () => { refresh: "r", expires: Date.now() + 100000, })); - const resolution = await models.getAuth(model); + const resolution = await models.getAuth(model.provider); expect(resolution?.auth.apiKey).toBe("oauth-token"); expect(resolution?.source).toBe("OAuth"); // stored api-key credential resolves through apiKey auth, beats env await credentials.modify("p1", async () => ({ type: "api_key", key: "stored-key" })); - const apiKeyResolution = await models.getAuth(model); + const apiKeyResolution = await models.getAuth(model.provider); expect(apiKeyResolution?.auth.apiKey).toBe("stored-key"); expect(apiKeyResolution?.source).toBe("stored"); }); + it("checks provider auth without refreshing OAuth and filters available models", async () => { + const credentials = new InMemoryCredentialStore(); + let refreshes = 0; + const models = createModels({ credentials }); + models.setProvider(testProvider({ id: "ambient", auth: { apiKey: envKeyAuth("env-key") } })); + models.setProvider(testProvider({ id: "missing", auth: { apiKey: envKeyAuth(undefined) } })); + models.setProvider( + testProvider({ + id: "oauth", + auth: { + oauth: testOAuth({ + refresh: async (credential) => { + refreshes++; + return credential; + }, + }), + }, + }), + ); + await credentials.modify("oauth", async () => ({ + type: "oauth", + access: "expired", + refresh: "refresh", + expires: 0, + })); + + expect(await models.checkAuth("ambient")).toEqual({ source: "env", type: "api_key" }); + expect(await models.checkAuth("missing")).toBeUndefined(); + expect(await models.checkAuth("oauth")).toEqual({ source: "OAuth", type: "oauth" }); + expect(refreshes).toBe(0); + expect((await models.getAvailable()).map((model) => model.provider)).toEqual(["ambient", "oauth"]); + expect((await models.getAvailable("ambient")).map((model) => model.provider)).toEqual(["ambient"]); + }); + + it("runs provider login and logout through the credential store", async () => { + const credentials = new InMemoryCredentialStore(); + const apiKey = envKeyAuth(undefined); + apiKey.login = async () => ({ type: "api_key", key: "logged-in" }); + const models = createModels({ credentials }); + models.setProvider(testProvider({ id: "p1", auth: { apiKey } })); + + const credential = await models.login("p1", "api_key", { + prompt: async () => "unused", + notify: () => {}, + }); + expect(credential).toEqual({ type: "api_key", key: "logged-in" }); + expect(await credentials.read("p1")).toEqual(credential); + + await models.logout("p1"); + expect(await credentials.read("p1")).toBeUndefined(); + }); + it("a stored credential without a matching handler blocks ambient fallback", async () => { const credentials = new InMemoryCredentialStore(); const models = createModels({ credentials }); @@ -274,7 +344,7 @@ describe("Models runtime", () => { models.setProvider(testProvider({ id: "p1", auth: { apiKey: envKeyAuth("env-key") } })); await credentials.modify("p1", async () => ({ type: "oauth", access: "a", refresh: "r", expires: 0 })); - expect(await models.getAuth(testModel("p1", "model-a"))).toBeUndefined(); + expect(await models.getAuth("p1")).toBeUndefined(); }); it("refreshes expired oauth credentials and persists the rotated credential", async () => { @@ -291,7 +361,7 @@ describe("Models runtime", () => { expires: 0, })); - const resolution = await models.getAuth(testModel("p1", "model-a")); + const resolution = await models.getAuth("p1"); expect(resolution?.auth.apiKey).toBe("new-token"); expect(((await credentials.read("p1")) as { access: string }).access).toBe("new-token"); }); @@ -307,7 +377,7 @@ describe("Models runtime", () => { models.setProvider(testProvider({ id: "p1", auth: { oauth } })); await credentials.modify("p1", async () => ({ type: "oauth", access: "old", refresh: "r", expires: 0 })); - await expect(models.getAuth(testModel("p1", "model-a"))).rejects.toMatchObject({ code: "oauth" }); + await expect(models.getAuth("p1")).rejects.toMatchObject({ code: "oauth" }); // credential preserved for retry / re-login expect(((await credentials.read("p1")) as { access: string }).access).toBe("old"); }); @@ -328,7 +398,7 @@ describe("Models runtime", () => { models.setProvider(testProvider({ id: "p1", auth: { oauth } })); const model = testModel("p1", "model-a"); - const [a, b] = await Promise.all([models.getAuth(model), models.getAuth(model)]); + const [a, b] = await Promise.all([models.getAuth(model.provider), models.getAuth(model.provider)]); expect(refreshes).toBe(1); expect(a?.auth.apiKey).toBe("new-1"); expect(b?.auth.apiKey).toBe("new-1"); @@ -339,6 +409,7 @@ describe("Models runtime", () => { const base = new InMemoryCredentialStore(); const credentials: CredentialStore = { read: (pid) => base.read(pid), + list: () => base.list(), modify: (pid, fn) => { modifies++; return base.modify(pid, fn); @@ -354,7 +425,7 @@ describe("Models runtime", () => { const models = createModels({ credentials }); models.setProvider(testProvider({ id: "p1", auth: { oauth: testOAuth() } })); - expect((await models.getAuth(testModel("p1", "model-a")))?.auth.apiKey).toBe("valid"); + expect((await models.getAuth("p1"))?.auth.apiKey).toBe("valid"); expect(modifies).toBe(0); }); @@ -364,16 +435,18 @@ describe("Models runtime", () => { read: async () => { throw new Error("disk on fire"); }, + list: async () => [], modify: async () => undefined, delete: async () => {}, }; const models = createModels({ credentials: readFailing }); models.setProvider(testProvider({ id: "p1", auth: { apiKey: envKeyAuth("env-key") } })); - await expect(models.getAuth(testModel("p1", "model-a"))).rejects.toMatchObject({ code: "auth" }); + await expect(models.getAuth("p1")).rejects.toMatchObject({ code: "auth" }); // modify failure during refresh const modifyFailing: CredentialStore = { read: async () => ({ type: "oauth", access: "old", refresh: "r", expires: 0 }), + list: async () => [{ providerId: "p1", type: "oauth" }], modify: async () => { throw new Error("disk on fire"); }, @@ -381,7 +454,7 @@ describe("Models runtime", () => { }; const oauthModels = createModels({ credentials: modifyFailing }); oauthModels.setProvider(testProvider({ id: "p1", auth: { oauth: testOAuth() } })); - await expect(oauthModels.getAuth(testModel("p1", "model-a"))).rejects.toMatchObject({ code: "auth" }); + await expect(oauthModels.getAuth("p1")).rejects.toMatchObject({ code: "auth" }); }); it("wraps api-key auth failures in ModelsError", async () => { @@ -393,7 +466,7 @@ describe("Models runtime", () => { }; const models = createModels(); models.setProvider(testProvider({ id: "p1", auth: { apiKey: failing } })); - await expect(models.getAuth(testModel("p1", "model-a"))).rejects.toMatchObject({ code: "auth" }); + await expect(models.getAuth("p1")).rejects.toMatchObject({ code: "auth" }); }); it("uses explicit request api key and env during provider auth resolution", async () => { @@ -427,7 +500,7 @@ describe("Models runtime", () => { resolve: async () => ({ auth: { apiKey: "resolved-key", - headers: { "x-a": "auth", "x-b": "auth" }, + headers: { Authorization: "Bearer resolved-key", "x-a": "auth", "x-b": "auth" }, baseUrl: "https://auth.test/v1", }, }), @@ -438,12 +511,12 @@ describe("Models runtime", () => { const result = await models.completeSimple(model, context, { apiKey: "explicit-key", - headers: { "x-b": "explicit" }, + headers: { authorization: "Explicit token", "x-b": "explicit" }, }); expect(result.stopReason).toBe("stop"); expect(calls).toHaveLength(1); expect(calls[0].options?.apiKey).toBe("explicit-key"); - expect(calls[0].options?.headers).toEqual({ "x-a": "auth", "x-b": "explicit" }); + expect(calls[0].options?.headers).toEqual({ authorization: "Explicit token", "x-a": "auth", "x-b": "explicit" }); expect(calls[0].model.baseUrl).toBe("https://auth.test/v1"); // without explicit options, resolved auth applies @@ -452,6 +525,36 @@ describe("Models runtime", () => { expect(calls[1].options?.apiKey).toBe("resolved-key"); }); + it("adds model headers only for model auth and transforms assembled headers once", async () => { + const calls: ProviderCall[] = []; + const models = createModels(); + models.setProvider(testProvider({ id: "p1", auth: { apiKey: envKeyAuth("key") }, calls })); + const model = testModel("p1", "model-a"); + model.headers = { "x-model": "model", "x-shared": "model" }; + + expect((await models.getAuth("p1"))?.auth.headers).toBeUndefined(); + expect((await models.getAuth(model))?.auth.headers).toEqual({ "x-model": "model", "x-shared": "model" }); + + let transforms = 0; + await models.completeSimple(model, context, { + headers: { "x-explicit": "explicit", "X-Shared": "explicit" }, + transformHeaders: async (headers) => { + transforms++; + expect(headers).toEqual({ "x-model": "model", "x-explicit": "explicit", "X-Shared": "explicit" }); + return { ...headers, "x-transformed": "yes" }; + }, + }); + + expect(transforms).toBe(1); + expect(calls[0].options?.headers).toEqual({ + "x-model": "model", + "x-explicit": "explicit", + "X-Shared": "explicit", + "x-transformed": "yes", + }); + expect(calls[0].options).not.toHaveProperty("transformHeaders"); + }); + it("produces an error stream for unknown providers instead of throwing", async () => { const models = createModels(); const result = await models.completeSimple(testModel("ghost", "model-a"), context); diff --git a/packages/ai/test/oauth-auth.test.ts b/packages/ai/test/oauth-auth.test.ts index 43be6f72..68534901 100644 --- a/packages/ai/test/oauth-auth.test.ts +++ b/packages/ai/test/oauth-auth.test.ts @@ -1,17 +1,23 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { InMemoryCredentialStore } from "../src/auth/credential-store.ts"; +import { anthropicOAuth } from "../src/auth/oauth/anthropic.ts"; +import { githubCopilotOAuth } from "../src/auth/oauth/github-copilot.ts"; +import { openaiCodexOAuth } from "../src/auth/oauth/openai-codex.ts"; import { createModels } from "../src/models.ts"; +import * as extensionOAuthCompatibility from "../src/oauth.ts"; import { anthropicProvider } from "../src/providers/anthropic.ts"; import { githubCopilotProvider } from "../src/providers/github-copilot.ts"; -import { anthropicOAuth } from "../src/utils/oauth/anthropic.ts"; -import { githubCopilotOAuth } from "../src/utils/oauth/github-copilot.ts"; -import { openaiCodexOAuth } from "../src/utils/oauth/openai-codex.ts"; function jsonResponse(body: unknown, status = 200): Response { return new Response(JSON.stringify(body), { status, headers: { "Content-Type": "application/json" } }); } describe.sequential("OAuthAuth adapters", () => { + it("keeps the extension OAuth barrel free of built-in flow implementations", () => { + expect(extensionOAuthCompatibility).not.toHaveProperty("loginAnthropic"); + expect(extensionOAuthCompatibility).not.toHaveProperty("anthropicOAuth"); + }); + afterEach(() => { vi.unstubAllGlobals(); }); @@ -104,7 +110,7 @@ describe("OAuth through Models.getAuth (lazy load chain)", () => { models.setProvider(anthropicProvider()); const model = models.getModels("anthropic")[0]; - const result = await models.getAuth(model); + const result = await models.getAuth(model.provider); expect(result?.auth.apiKey).toBe("oauth-access-token"); expect(result?.source).toBe("OAuth"); }); @@ -122,7 +128,7 @@ describe("OAuth through Models.getAuth (lazy load chain)", () => { models.setProvider(githubCopilotProvider()); const model = models.getModels("github-copilot")[0]; - const result = await models.getAuth(model); + const result = await models.getAuth(model.provider); expect(result?.auth.apiKey).toBe(access); expect(result?.auth.baseUrl).toBe("https://api.business.githubcopilot.com"); }); diff --git a/packages/ai/test/oauth-device-code.test.ts b/packages/ai/test/oauth-device-code.test.ts index 23f97042..fe146f58 100644 --- a/packages/ai/test/oauth-device-code.test.ts +++ b/packages/ai/test/oauth-device-code.test.ts @@ -1,5 +1,5 @@ import { afterEach, describe, expect, it, vi } from "vitest"; -import { pollOAuthDeviceCodeFlow } from "../src/utils/oauth/device-code.ts"; +import { pollOAuthDeviceCodeFlow } from "../src/auth/oauth/device-code.ts"; describe("OAuth device-code polling", () => { afterEach(() => { diff --git a/packages/ai/test/oauth.ts b/packages/ai/test/oauth.ts index 8ee7c94d..57060b89 100644 --- a/packages/ai/test/oauth.ts +++ b/packages/ai/test/oauth.ts @@ -8,8 +8,8 @@ import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs"; import { homedir } from "os"; import { dirname, join } from "path"; -import { getOAuthApiKey } from "../src/utils/oauth/index.ts"; -import type { OAuthCredentials, OAuthProvider } from "../src/utils/oauth/types.ts"; +import type { OAuthCredentials } from "../src/auth/types.ts"; +import { builtinProviders } from "../src/providers/all.ts"; const AUTH_PATH = join(homedir(), ".pi", "agent", "auth.json"); @@ -65,28 +65,18 @@ export async function resolveApiKey(provider: string): Promise = {}; - for (const [key, value] of Object.entries(storage)) { - if (value.type === "oauth") { - const { type: _, ...creds } = value; - oauthCredentials[key] = creds; - } - } - - let result: { newCredentials: OAuthCredentials; apiKey: string } | null = null; + const oauth = builtinProviders().find((candidate) => candidate.id === provider)?.auth.oauth; + if (!oauth) return undefined; + let credential = entry; try { - result = await getOAuthApiKey(provider as OAuthProvider, oauthCredentials); - } catch (e) { - console.log(JSON.stringify(e)); + if (Date.now() >= credential.expires) credential = await oauth.refresh(credential); + } catch (error) { + console.log(JSON.stringify(error)); + return undefined; } - if (!result) return undefined; - - // Save refreshed credentials back to auth.json - storage[provider] = { type: "oauth", ...result.newCredentials }; + storage[provider] = credential; saveAuthStorage(storage); - - return result.apiKey; + return (await oauth.toAuth(credential)).apiKey; } return undefined; diff --git a/packages/ai/test/openai-codex-oauth.test.ts b/packages/ai/test/openai-codex-oauth.test.ts index 820fbe6b..64c0e8c7 100644 --- a/packages/ai/test/openai-codex-oauth.test.ts +++ b/packages/ai/test/openai-codex-oauth.test.ts @@ -1,9 +1,5 @@ import { afterEach, describe, expect, it, vi } from "vitest"; -import { - loginOpenAICodexDeviceCode, - openaiCodexOAuthProvider, - refreshOpenAICodexToken, -} from "../src/utils/oauth/openai-codex.ts"; +import { openaiCodexOAuth } from "../src/auth/oauth/openai-codex.ts"; function jsonResponse(body: unknown, status: number = 200): Response { return new Response(JSON.stringify(body), { @@ -45,6 +41,30 @@ function deviceAuthPendingResponse(): Response { ); } +function loginOpenAICodexDeviceCodeForTest(options: { + onDeviceCode(info: { + userCode: string; + verificationUri: string; + intervalSeconds?: number; + expiresInSeconds?: number; + }): void; + signal?: AbortSignal; +}) { + return openaiCodexOAuth.login({ + signal: options.signal, + prompt: async (prompt) => { + if (prompt.type !== "select") throw new Error(`Unexpected prompt: ${prompt.type}`); + return "device_code"; + }, + notify: (event) => { + if (event.type === "device_code") { + const { type: _, ...info } = event; + options.onDeviceCode(info); + } + }, + }); +} + describe("OpenAI Codex OAuth", () => { afterEach(() => { vi.restoreAllMocks(); @@ -125,7 +145,7 @@ describe("OpenAI Codex OAuth", () => { vi.stubGlobal("fetch", fetchMock); - const credentialsPromise = loginOpenAICodexDeviceCode({ + const credentialsPromise = loginOpenAICodexDeviceCodeForTest({ onDeviceCode: (info) => deviceInfos.push(info), }); @@ -159,7 +179,7 @@ describe("OpenAI Codex OAuth", () => { const accessToken = createAccessToken("account-456"); const selectPrompts: Array<{ message: string; - options: Array<{ id: string; label: string }>; + options: readonly { id: string; label: string }[]; }> = []; const deviceInfos: Array<{ userCode: string; @@ -199,20 +219,22 @@ describe("OpenAI Codex OAuth", () => { ); await expect( - openaiCodexOAuthProvider.login({ - onAuth: () => { - throw new Error("Browser login should not start"); - }, - onDeviceCode: (info) => deviceInfos.push(info), - onPrompt: async () => { - throw new Error("Prompt should not be used"); - }, - onSelect: async (prompt) => { + openaiCodexOAuth.login({ + prompt: async (prompt) => { + if (prompt.type !== "select") throw new Error("Text prompt should not be used"); selectPrompts.push(prompt); return "device_code"; }, + notify: (event) => { + if (event.type === "auth_url") throw new Error("Browser login should not start"); + if (event.type === "device_code") { + const { type: _, ...info } = event; + deviceInfos.push(info); + } + }, }), ).resolves.toMatchObject({ + type: "oauth", access: accessToken, refresh: "refresh-token", accountId: "account-456", @@ -220,6 +242,7 @@ describe("OpenAI Codex OAuth", () => { expect(selectPrompts).toEqual([ { + type: "select", message: "Select OpenAI Codex login method:", options: [ { id: "browser", label: "Browser login (default)" }, @@ -239,11 +262,11 @@ describe("OpenAI Codex OAuth", () => { it("cancels when OpenAI Codex login method selection is cancelled", async () => { await expect( - openaiCodexOAuthProvider.login({ - onAuth: () => {}, - onDeviceCode: () => {}, - onPrompt: async () => "", - onSelect: async () => undefined, + openaiCodexOAuth.login({ + prompt: async () => { + throw new Error("Login cancelled"); + }, + notify: () => {}, }), ).rejects.toThrow("Login cancelled"); }); @@ -273,7 +296,7 @@ describe("OpenAI Codex OAuth", () => { }), ); - const credentialsPromise = loginOpenAICodexDeviceCode({ + const credentialsPromise = loginOpenAICodexDeviceCodeForTest({ onDeviceCode: () => {}, signal: controller.signal, }); @@ -317,7 +340,7 @@ describe("OpenAI Codex OAuth", () => { }), ); - const credentialsPromise = loginOpenAICodexDeviceCode({ + const credentialsPromise = loginOpenAICodexDeviceCodeForTest({ onDeviceCode: () => {}, }); const rejectionPromise = credentialsPromise.then( @@ -380,7 +403,7 @@ describe("OpenAI Codex OAuth", () => { }), ); - const credentialsPromise = loginOpenAICodexDeviceCode({ + const credentialsPromise = loginOpenAICodexDeviceCodeForTest({ onDeviceCode: () => {}, }); @@ -418,7 +441,7 @@ describe("OpenAI Codex OAuth", () => { ); await expect( - loginOpenAICodexDeviceCode({ + loginOpenAICodexDeviceCodeForTest({ onDeviceCode: () => {}, }), ).rejects.toThrow( @@ -443,9 +466,14 @@ describe("OpenAI Codex OAuth", () => { }), ); - await expect(refreshOpenAICodexToken("invalid-refresh-token")).rejects.toThrow( - /OpenAI Codex token refresh failed \(401\).*Could not validate your token/, - ); + await expect( + openaiCodexOAuth.refresh({ + type: "oauth", + access: "invalid-access-token", + refresh: "invalid-refresh-token", + expires: 0, + }), + ).rejects.toThrow(/OpenAI Codex token refresh failed \(401\).*Could not validate your token/); expect(consoleError).not.toHaveBeenCalled(); }); }); diff --git a/packages/ai/test/openai-responses-copilot-provider.test.ts b/packages/ai/test/openai-responses-copilot-provider.test.ts index 3a62ce58..00b3a369 100644 --- a/packages/ai/test/openai-responses-copilot-provider.test.ts +++ b/packages/ai/test/openai-responses-copilot-provider.test.ts @@ -252,11 +252,12 @@ describe("openai-responses provider defaults", () => { expect(captured).toEqual({ sessionId: null, clientRequestId: null }); }); - it.each([ - ["gpt-5.4", "priority", 2], - ["gpt-5.5", "priority", 2.5], - ["gpt-5.5", "flex", 0.5], - ] as const)("applies %s %s service-tier cost multiplier", async (modelId, serviceTier, multiplier) => { + async function streamServiceTierUsage( + modelId: "gpt-5.4" | "gpt-5.5", + serviceTier: "priority" | "flex", + inputTokens: number, + outputTokens: number, + ) { const model = getModel("openai", modelId); const sse = `${[ `data: ${JSON.stringify({ @@ -265,9 +266,9 @@ describe("openai-responses provider defaults", () => { status: "completed", service_tier: serviceTier, usage: { - input_tokens: 1000000, - output_tokens: 1000000, - total_tokens: 2000000, + input_tokens: inputTokens, + output_tokens: outputTokens, + total_tokens: inputTokens + outputTokens, input_tokens_details: { cached_tokens: 0 }, }, }, @@ -290,10 +291,39 @@ describe("openai-responses provider defaults", () => { { apiKey: "test-key", serviceTier }, ); - const result = await stream.result(); + return { model, result: await stream.result() }; + } - expect(result.usage.cost.input).toBe(model.cost.input * multiplier); - expect(result.usage.cost.output).toBe(model.cost.output * multiplier); - expect(result.usage.cost.total).toBe((model.cost.input + model.cost.output) * multiplier); + it.each([ + ["gpt-5.4", "priority", 2], + ["gpt-5.5", "priority", 2.5], + ["gpt-5.5", "flex", 0.5], + ] as const)("applies %s %s service-tier cost multiplier", async (modelId, serviceTier, multiplier) => { + // Stay below the 272K long-context tier threshold so base rates apply. + const inputTokens = 200000; + const outputTokens = 100000; + const { model, result } = await streamServiceTierUsage(modelId, serviceTier, inputTokens, outputTokens); + + const expectedInput = (model.cost.input / 1_000_000) * inputTokens * multiplier; + const expectedOutput = (model.cost.output / 1_000_000) * outputTokens * multiplier; + expect(result.usage.cost.input).toBe(expectedInput); + expect(result.usage.cost.output).toBe(expectedOutput); + expect(result.usage.cost.total).toBe(expectedInput + expectedOutput); + }); + + it("applies the service-tier multiplier on top of long-context tier pricing", async () => { + // Above the 272K input threshold the long-context tier rates apply, then the multiplier. + const inputTokens = 1000000; + const outputTokens = 100000; + const multiplier = 2; + const { model, result } = await streamServiceTierUsage("gpt-5.4", "priority", inputTokens, outputTokens); + + const tier = model.cost.tiers?.find((entry) => inputTokens > entry.inputTokensAbove); + if (!tier) throw new Error("expected gpt-5.4 to define a long-context pricing tier"); + const expectedInput = (tier.input / 1_000_000) * inputTokens * multiplier; + const expectedOutput = (tier.output / 1_000_000) * outputTokens * multiplier; + expect(result.usage.cost.input).toBe(expectedInput); + expect(result.usage.cost.output).toBe(expectedOutput); + expect(result.usage.cost.total).toBe(expectedInput + expectedOutput); }); }); diff --git a/packages/ai/test/providers.test.ts b/packages/ai/test/providers.test.ts index 0f6e25bd..c02fb6b7 100644 --- a/packages/ai/test/providers.test.ts +++ b/packages/ai/test/providers.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "vitest"; import { envApiKeyAuth } from "../src/auth/helpers.ts"; -import type { AuthContext } from "../src/auth/types.ts"; +import type { AuthContext, AuthEvent } from "../src/auth/types.ts"; import { createModels, createProvider } from "../src/models.ts"; import { builtinModels, builtinProviders } from "../src/providers/all.ts"; import { amazonBedrockProvider } from "../src/providers/amazon-bedrock.ts"; @@ -49,40 +49,69 @@ describe("builtin providers", () => { models.setProvider(anthropicProvider()); const model = models.getModel("anthropic", "claude-haiku-4-5")!; - const result = await models.getAuth(model); + const result = await models.getAuth(model.provider); expect(result?.auth.apiKey).toBe("oauth-token"); expect(result?.source).toBe("ANTHROPIC_OAUTH_TOKEN"); }); + it("runs provider-owned Bedrock bearer token and AWS profile login flows", async () => { + const auth = amazonBedrockProvider().auth.apiKey!; + const bearerAnswers = ["bearer-token", "bedrock-token"]; + expect( + await auth.login?.({ + prompt: async () => bearerAnswers.shift()!, + notify: () => {}, + }), + ).toEqual({ type: "api_key", key: "bedrock-token" }); + + const profileAnswers = ["aws-profile", "work"]; + const events: AuthEvent[] = []; + expect( + await auth.login?.({ + prompt: async () => profileAnswers.shift()!, + notify: (event) => events.push(event), + }), + ).toEqual({ type: "api_key", env: { AWS_PROFILE: "work" } }); + expect(events).toEqual([ + expect.objectContaining({ + type: "info", + links: [expect.objectContaining({ label: "AWS credential provider chain" })], + }), + ]); + expect( + await auth.resolve({ + ctx: fakeAuthContext({}), + credential: { type: "api_key", env: { AWS_PROFILE: "work" } }, + }), + ).toMatchObject({ auth: {}, env: { AWS_PROFILE: "work" } }); + }); + it("reports bedrock as configured from ambient AWS credentials without an api key", async () => { const models = createModels({ authContext: fakeAuthContext({ AWS_PROFILE: "dev" }) }); models.setProvider(amazonBedrockProvider()); const model = models.getModels("amazon-bedrock")[0]; - const result = await models.getAuth(model); + const result = await models.getAuth(model.provider); expect(result?.auth).toEqual({}); expect(result?.source).toBe("AWS_PROFILE"); const unconfigured = createModels({ authContext: fakeAuthContext({}) }); unconfigured.setProvider(amazonBedrockProvider()); - expect(await unconfigured.getAuth(model)).toBeUndefined(); + expect(await unconfigured.getAuth(model.provider)).toBeUndefined(); }); it("requires Cloudflare Workers AI account config and returns scoped env", async () => { const missingAccount = createModels({ authContext: fakeAuthContext({ CLOUDFLARE_API_KEY: "cf-key" }) }); missingAccount.setProvider(cloudflareWorkersAIProvider()); const model = missingAccount.getModels("cloudflare-workers-ai")[0]; - expect(await missingAccount.getAuth(model)).toBeUndefined(); + expect(await missingAccount.getAuth(model.provider)).toBeUndefined(); const configured = createModels({ authContext: fakeAuthContext({ CLOUDFLARE_API_KEY: "cf-key", CLOUDFLARE_ACCOUNT_ID: "account-id" }), }); configured.setProvider(cloudflareWorkersAIProvider()); - const result = await configured.getAuth(model); - expect(result?.auth).toEqual({ - apiKey: "cf-key", - baseUrl: "https://api.cloudflare.com/client/v4/accounts/account-id/ai/v1", - }); + const result = await configured.getAuth(model.provider); + expect(result?.auth).toEqual({ apiKey: "cf-key" }); expect(result?.env).toEqual({ CLOUDFLARE_ACCOUNT_ID: "account-id" }); }); @@ -92,7 +121,7 @@ describe("builtin providers", () => { }); missingGateway.setProvider(cloudflareAIGatewayProvider()); const model = missingGateway.getModels("cloudflare-ai-gateway")[0]; - expect(await missingGateway.getAuth(model)).toBeUndefined(); + expect(await missingGateway.getAuth(model.provider)).toBeUndefined(); const configured = createModels({ authContext: fakeAuthContext({ @@ -102,14 +131,13 @@ describe("builtin providers", () => { }), }); configured.setProvider(cloudflareAIGatewayProvider()); - const result = await configured.getAuth(model); + const result = await configured.getAuth(model.provider); expect(result?.auth).toEqual({ headers: { "cf-aig-authorization": "Bearer cf-key", Authorization: null, "x-api-key": null, }, - baseUrl: "https://gateway.ai.cloudflare.com/v1/account-id/gateway-id/anthropic", }); expect(result?.env).toEqual({ CLOUDFLARE_ACCOUNT_ID: "account-id", @@ -117,6 +145,47 @@ describe("builtin providers", () => { }); }); + it("runs provider-owned Vertex API key and ADC login flows", async () => { + const auth = googleVertexProvider().auth.apiKey!; + const keyAnswers = ["api-key", "vertex-key"]; + expect( + await auth.login?.({ + prompt: async () => keyAnswers.shift()!, + notify: () => {}, + }), + ).toEqual({ type: "api_key", key: "vertex-key" }); + + const adcAnswers = ["adc", "project-id", "us-central1"]; + const events: AuthEvent[] = []; + expect( + await auth.login?.({ + prompt: async () => adcAnswers.shift()!, + notify: (event) => events.push(event), + }), + ).toEqual({ + type: "api_key", + env: { GOOGLE_CLOUD_PROJECT: "project-id", GOOGLE_CLOUD_LOCATION: "us-central1" }, + }); + expect(events).toEqual([ + expect.objectContaining({ + type: "info", + links: [expect.objectContaining({ label: "Application Default Credentials" })], + }), + ]); + expect( + await auth.resolve({ + ctx: fakeAuthContext({}, ["~/.config/gcloud/application_default_credentials.json"]), + credential: { + type: "api_key", + env: { GOOGLE_CLOUD_PROJECT: "project-id", GOOGLE_CLOUD_LOCATION: "us-central1" }, + }, + }), + ).toMatchObject({ + auth: {}, + env: { GOOGLE_CLOUD_PROJECT: "project-id", GOOGLE_CLOUD_LOCATION: "us-central1" }, + }); + }); + it("resolves vertex via ADC file plus project and location", async () => { const adc = "~/.config/gcloud/application_default_credentials.json"; const configured = createModels({ @@ -125,40 +194,38 @@ describe("builtin providers", () => { configured.setProvider(googleVertexProvider()); const model = configured.getModels("google-vertex")[0]; - const result = await configured.getAuth(model); + const result = await configured.getAuth(model.provider); expect(result?.auth).toEqual({}); expect(result?.source).toContain("application default"); // ADC without project/location is not configured const partial = createModels({ authContext: fakeAuthContext({ GOOGLE_CLOUD_PROJECT: "proj" }, [adc]) }); partial.setProvider(googleVertexProvider()); - expect(await partial.getAuth(model)).toBeUndefined(); + expect(await partial.getAuth(model.provider)).toBeUndefined(); // explicit key wins over ADC const keyed = createModels({ authContext: fakeAuthContext({ GOOGLE_CLOUD_API_KEY: "vertex-key" }) }); keyed.setProvider(googleVertexProvider()); - expect((await keyed.getAuth(model))?.auth.apiKey).toBe("vertex-key"); + expect((await keyed.getAuth(model.provider))?.auth.apiKey).toBe("vertex-key"); }); }); describe("envApiKeyAuth", () => { it("prefers the stored credential key and falls back through env vars in order", async () => { const auth = envApiKeyAuth("Test key", ["FIRST_KEY", "SECOND_KEY"]); - const model = { provider: "p1" } as Model; const stored = await auth.resolve({ - model, ctx: fakeAuthContext({ FIRST_KEY: "env" }), credential: { type: "api_key", key: "stored" }, }); expect(stored?.auth.apiKey).toBe("stored"); expect(stored?.source).toBe("stored credential"); - const second = await auth.resolve({ model, ctx: fakeAuthContext({ SECOND_KEY: "second" }) }); + const second = await auth.resolve({ ctx: fakeAuthContext({ SECOND_KEY: "second" }) }); expect(second?.auth.apiKey).toBe("second"); expect(second?.source).toBe("SECOND_KEY"); - expect(await auth.resolve({ model, ctx: fakeAuthContext({}) })).toBeUndefined(); + expect(await auth.resolve({ ctx: fakeAuthContext({}) })).toBeUndefined(); }); it("login prompts for a secret and returns an api-key credential", async () => { diff --git a/packages/ai/test/scratch.ts b/packages/ai/test/scratch.ts index c2d83649..fec7b31a 100644 --- a/packages/ai/test/scratch.ts +++ b/packages/ai/test/scratch.ts @@ -21,7 +21,7 @@ models.setProvider(anthropicProvider()); const model = models.getModel("anthropic", "claude-haiku-4-5"); if (!model) throw new Error("model not found"); -const auth = await models.getAuth(model); +const auth = await models.getAuth(model.provider); console.log(`model: ${model.provider}/${model.id}`); console.log(`auth: ${auth ? `configured via ${auth.source}` : "not configured"}\n`); if (!auth) process.exit(1); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 2d810f29..2b13e87b 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,11 +2,76 @@ ## [Unreleased] +### Breaking Changes + +- Replaced the SDK's `CreateAgentSessionOptions.authStorage` and `modelRegistry` options with the async `modelRuntime` option. `AuthStorage` and its storage backends are no longer exported; use `ModelRuntime` (or a custom pi-ai `CredentialStore`), or `readStoredCredential()` for one-off reads of auth.json. +- Removed redundant `ModelRuntime.getAll()`, `find()`, `getSnapshot()`, and `getAuthOptions()` projections. Use the pi-ai `Models` methods `getModels()`, `getModel()`, `getProviders()`, and `checkAuth()` directly. +- Replaced SDK request-auth assembly through `ModelRegistry.getApiKeyAndHeaders()` with `ModelRuntime.getAuth()`. Passing a provider ID returns provider-scoped auth; passing a model also resolves built-in, `models.json`, and extension model headers. +- Changed extension-facing `ModelRegistry.refresh()` from synchronous `void` to `Promise` because `models.json` loading is asynchronous. Extensions must await it before making synchronous registry reads. +- Removed extension OAuth `modifyModels`. Provider catalogs are now composed independently of credentials; credential-specific availability belongs to canonical provider filtering. The legacy extension OAuth callback and credential types remain available from pi-ai's root and `oauth` subpath. + +#### SDK migration + +Construct one `ModelRuntime` and pass it to `createAgentSession()`: + +```typescript +// Before +const authStorage = AuthStorage.create(authPath); +const modelRegistry = await ModelRegistry.create(authStorage, modelsPath); +authStorage.setRuntimeApiKey("anthropic", apiKey); +const { session } = await createAgentSession({ authStorage, modelRegistry }); + +// After +const modelRuntime = await ModelRuntime.create({ authPath, modelsPath }); +// Or: ModelRuntime.create({ credentials: myCredentialStore, modelsPath }) +modelRuntime.setRuntimeApiKey("anthropic", apiKey); +const { session } = await createAgentSession({ modelRuntime }); +``` + +Replace `ModelRegistry` projections with the corresponding `ModelRuntime`/pi-ai `Models` methods: + +```typescript +const allModels = modelRuntime.getModels(); +const model = modelRuntime.getModel(providerId, modelId); +const availableModels = await modelRuntime.getAvailable(); +const authStatus = await modelRuntime.checkAuth(providerId); +const requestAuth = await modelRuntime.getAuth(model); // Includes model headers + +modelRuntime.registerProvider(providerId, providerConfig); // Still synchronous +await modelRuntime.reloadConfig(); +``` + +`ModelRuntime.stream*()` resolves auth and configured headers itself. Do not call `getAuth(model)` before streaming merely to reconstruct request options. For SDK-level header interception, use the Models-only transform so auth is resolved once: + +```typescript +modelRuntime.streamSimple(model, context, { + transformHeaders: async (headers) => ({ + ...headers, + "X-Request-ID": requestId, + }), +}); +``` + +Use `ModelRuntime` for model lookup, availability, provider auth, login/logout, runtime API-key overrides, provider registration, and config refresh. `ModelRegistry` remains a synchronous-read compatibility facade for extensions; SDK code should use `ModelRuntime`. Extensions that explicitly refresh it must await completion: + +```typescript +await ctx.modelRegistry.refresh(); +const models = ctx.modelRegistry.getAll(); +``` + + ### Added +- Added `ModelRuntime` as the canonical async SDK and internal model/auth facade while preserving the synchronous extension-facing `ModelRegistry` API. `ModelRuntime.create()` accepts any pi-ai `CredentialStore` through its `credentials` option. +- Added provider-owned `/login` discovery directly from registered pi-ai providers, including ambient auth status and informational links. - Added the opt-in `max` thinking level across CLI, SDK, RPC, model selection, and themes. Custom themes can define `thinkingMax`; existing themes fall back to `thinkingXhigh`. - Added request-wide input-token pricing tiers to custom model costs in `models.json`, `modelOverrides`, and extension-registered providers. +### Changed + +- Changed `ModelRuntime` to compose built-in providers, immutable `models.json` configuration, and extension overlays through ad-hoc pi-ai provider methods. +- Changed `ModelRuntime` to own final request assembly: `getAuth(model)` includes configured model headers, stream methods resolve auth once, and `before_provider_headers` runs as the Models-only header transform before provider dispatch. + ## [0.80.5] - 2026-07-09 ## [0.80.4] - 2026-07-09 diff --git a/packages/coding-agent/README.md b/packages/coding-agent/README.md index dc69b4d2..53267570 100644 --- a/packages/coding-agent/README.md +++ b/packages/coding-agent/README.md @@ -453,14 +453,12 @@ See [docs/packages.md](docs/packages.md). ### SDK ```typescript -import { AuthStorage, createAgentSession, ModelRegistry, SessionManager } from "@earendil-works/pi-coding-agent"; +import { createAgentSession, ModelRuntime, SessionManager } from "@earendil-works/pi-coding-agent"; -const authStorage = AuthStorage.create(); -const modelRegistry = ModelRegistry.create(authStorage); +const modelRuntime = await ModelRuntime.create(); const { session } = await createAgentSession({ sessionManager: SessionManager.inMemory(), - authStorage, - modelRegistry, + modelRuntime, }); await session.prompt("What files are in the current directory?"); diff --git a/packages/coding-agent/docs/custom-provider.md b/packages/coding-agent/docs/custom-provider.md index 6a8a38ed..f2138c6a 100644 --- a/packages/coding-agent/docs/custom-provider.md +++ b/packages/coding-agent/docs/custom-provider.md @@ -253,6 +253,8 @@ pi.registerProvider("custom-api", { }); ``` +The key is resolved for each request. An explicit request `Authorization` header takes precedence over the generated value. + ## OAuth Support Add OAuth/SSO authentication that integrates with `/login`: @@ -312,15 +314,6 @@ pi.registerProvider("corporate-ai", { getApiKey(credentials: OAuthCredentials): string { return credentials.access; - }, - - // Optional: modify models based on user's subscription - modifyModels(models, credentials) { - const region = decodeRegionFromToken(credentials.access); - return models.map(m => ({ - ...m, - baseUrl: `https://${region}.ai.corp.com/v1` - })); } } }); @@ -330,7 +323,7 @@ After registration, users can authenticate via `/login corporate-ai`. ### OAuthLoginCallbacks -The `callbacks` object provides three ways to authenticate: +The `callbacks` object provides UI-neutral interactions for the provider-owned flow: ```typescript interface OAuthLoginCallbacks { @@ -345,6 +338,9 @@ interface OAuthLoginCallbacks { expiresInSeconds?: number; }): void; + // Show transient progress + onProgress?(message: string): void; + // Prompt user for input (for manual token entry) onPrompt(params: { message: string }): Promise; @@ -660,7 +656,6 @@ interface ProviderConfig { login(callbacks: OAuthLoginCallbacks): Promise; refreshToken(credentials: OAuthCredentials): Promise; getApiKey(credentials: OAuthCredentials): string; - modifyModels?(models: Model[], credentials: OAuthCredentials): Model[]; }; } ``` diff --git a/packages/coding-agent/docs/sdk.md b/packages/coding-agent/docs/sdk.md index c84d6afc..9fcfdeec 100644 --- a/packages/coding-agent/docs/sdk.md +++ b/packages/coding-agent/docs/sdk.md @@ -16,16 +16,12 @@ See [examples/sdk/](../examples/sdk/) for working examples from minimal to full ## Quick Start ```typescript -import { AuthStorage, createAgentSession, ModelRegistry, SessionManager } from "@earendil-works/pi-coding-agent"; - -// Set up credential storage and model registry -const authStorage = AuthStorage.create(); -const modelRegistry = ModelRegistry.create(authStorage); +import { createAgentSession, ModelRuntime, SessionManager } from "@earendil-works/pi-coding-agent"; +const modelRuntime = await ModelRuntime.create(); const { session } = await createAgentSession({ sessionManager: SessionManager.inMemory(), - authStorage, - modelRegistry, + modelRuntime, }); session.subscribe((event) => { @@ -369,10 +365,9 @@ When you pass a custom `ResourceLoader`, `cwd` and `agentDir` no longer control ```typescript import { getModel } from "@earendil-works/pi-ai"; -import { AuthStorage, ModelRegistry } from "@earendil-works/pi-coding-agent"; +import { ModelRuntime } from "@earendil-works/pi-coding-agent"; -const authStorage = AuthStorage.create(); -const modelRegistry = ModelRegistry.create(authStorage); +const modelRuntime = await ModelRuntime.create(); // Find specific built-in model (doesn't check if API key exists) const opus = getModel("anthropic", "claude-opus-4-5"); @@ -380,10 +375,10 @@ if (!opus) throw new Error("Model not found"); // Find any model by provider/id, including custom models from models.json // (doesn't check if API key exists) -const customModel = modelRegistry.find("my-provider", "my-model"); +const customModel = modelRuntime.getModel("my-provider", "my-model"); -// Get only models that have valid API keys configured -const available = await modelRegistry.getAvailable(); +// Get only models that have valid authentication configured +const available = await modelRuntime.getAvailable(); const { session } = await createAgentSession({ model: opus, @@ -395,8 +390,7 @@ const { session } = await createAgentSession({ { model: haiku, thinkingLevel: "off" }, ], - authStorage, - modelRegistry, + modelRuntime, }); ``` @@ -415,14 +409,14 @@ import { const cliModel = resolveCliModel({ cliModel: "anthropic/claude-opus-4-5:high", - modelRegistry, + modelRuntime, }); if (cliModel.error) throw new Error(cliModel.error); if (cliModel.warning) console.warn(cliModel.warning); const { scopedModels, diagnostics } = await resolveModelScopeWithDiagnostics( ["anthropic/*:high", "gpt-5"], - modelRegistry, + modelRuntime, ); for (const diagnostic of diagnostics) { console.warn(diagnostic.message); @@ -435,40 +429,41 @@ for (const diagnostic of diagnostics) { ### API Keys and OAuth -API key resolution priority (handled by AuthStorage): +Authentication resolution priority (handled by `ModelRuntime`): 1. Runtime overrides (via `setRuntimeApiKey`, not persisted) 2. Stored credentials in `auth.json` (API keys or OAuth tokens) 3. Environment variables (`ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, etc.) 4. Fallback resolver (for custom provider keys from `models.json`) ```typescript -import { AuthStorage, ModelRegistry } from "@earendil-works/pi-coding-agent"; +import { InMemoryCredentialStore } from "@earendil-works/pi-ai"; +import { createAgentSession, ModelRuntime } from "@earendil-works/pi-coding-agent"; // Default: uses ~/.pi/agent/auth.json and ~/.pi/agent/models.json -const authStorage = AuthStorage.create(); -const modelRegistry = ModelRegistry.create(authStorage); +const modelRuntime = await ModelRuntime.create(); -const { session } = await createAgentSession({ - sessionManager: SessionManager.inMemory(), - authStorage, - modelRegistry, -}); +// Provider-owned auth methods and current status +for (const provider of modelRuntime.getProviders()) { + const status = await modelRuntime.checkAuth(provider.id); + console.log(provider.name, provider.auth, status); +} // Runtime API key override (not persisted to disk) -authStorage.setRuntimeApiKey("anthropic", "sk-my-temp-key"); +modelRuntime.setRuntimeApiKey("anthropic", "sk-my-temp-key"); -// Custom auth storage location -const customAuth = AuthStorage.create("/my/app/auth.json"); -const customRegistry = ModelRegistry.create(customAuth, "/my/app/models.json"); - -const { session } = await createAgentSession({ - sessionManager: SessionManager.inMemory(), - authStorage: customAuth, - modelRegistry: customRegistry, +// Custom credential and model locations +const customRuntime = await ModelRuntime.create({ + authPath: "/my/app/auth.json", + modelsPath: "/my/app/models.json", }); -// No custom models.json (built-in models only) -const simpleRegistry = ModelRegistry.inMemory(authStorage); +// Or inject any pi-ai CredentialStore +const credentials = new InMemoryCredentialStore(); +const inMemoryRuntime = await ModelRuntime.create({ credentials }); + +const { session } = await createAgentSession({ + modelRuntime: customRuntime, +}); ``` > See [examples/sdk/09-api-keys-and-oauth.ts](../examples/sdk/09-api-keys-and-oauth.ts) @@ -927,26 +922,22 @@ interface LoadExtensionsResult { import { getModel } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { - AuthStorage, createAgentSession, DefaultResourceLoader, defineTool, - ModelRegistry, + ModelRuntime, SessionManager, SettingsManager, } from "@earendil-works/pi-coding-agent"; -// Set up auth storage (custom location) -const authStorage = AuthStorage.create("/custom/agent/auth.json"); - -// Runtime API key override (not persisted) +const modelRuntime = await ModelRuntime.create({ + authPath: "/custom/agent/auth.json", + modelsPath: "/custom/agent/models.json", +}); if (process.env.MY_KEY) { - authStorage.setRuntimeApiKey("anthropic", process.env.MY_KEY); + modelRuntime.setRuntimeApiKey("anthropic", process.env.MY_KEY); } -// Model registry (no custom models.json) -const modelRegistry = ModelRegistry.create(authStorage); - // Inline tool const statusTool = defineTool({ name: "status", @@ -982,8 +973,7 @@ const { session } = await createAgentSession({ model, thinkingLevel: "off", - authStorage, - modelRegistry, + modelRuntime, tools: ["read", "bash", "status"], customTools: [statusTool], @@ -1149,8 +1139,8 @@ createAgentSessionRuntime AgentSessionRuntime // Auth and Models -AuthStorage -ModelRegistry +ModelRuntime // implements pi-ai Models and owns credential storage +ModelRegistry // synchronous extension compatibility facade 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 cfa80dbe..c4bebc32 100644 --- a/packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts +++ b/packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts @@ -46,7 +46,7 @@ import { import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; // ============================================================================= -// OAuth Implementation (copied from packages/ai/src/utils/oauth/anthropic.ts) +// OAuth implementation adapted for the legacy extension compatibility interface. // ============================================================================= const decode = (s: string) => atob(s); diff --git a/packages/coding-agent/examples/sdk/02-custom-model.ts b/packages/coding-agent/examples/sdk/02-custom-model.ts index 641d553e..ab947dc1 100644 --- a/packages/coding-agent/examples/sdk/02-custom-model.ts +++ b/packages/coding-agent/examples/sdk/02-custom-model.ts @@ -5,11 +5,9 @@ */ import { getModel } from "@earendil-works/pi-ai/compat"; -import { AuthStorage, createAgentSession, ModelRegistry } from "@earendil-works/pi-coding-agent"; +import { createAgentSession, ModelRuntime } from "@earendil-works/pi-coding-agent"; -// Set up auth storage and model registry -const authStorage = AuthStorage.create(); -const modelRegistry = ModelRegistry.create(authStorage); +const modelRuntime = await ModelRuntime.create(); // Option 1: Find a specific built-in model by provider/id const opus = getModel("anthropic", "claude-opus-4-5"); @@ -18,13 +16,13 @@ if (opus) { } // Option 2: Find model via registry (includes custom models from models.json) -const customModel = modelRegistry.find("my-provider", "my-model"); +const customModel = modelRuntime.getModel("my-provider", "my-model"); if (customModel) { console.log(`Found custom model: ${customModel.provider}/${customModel.id}`); } // Option 3: Pick from available models (have valid API keys) -const available = await modelRegistry.getAvailable(); +const available = await modelRuntime.getAvailable(); console.log( "Available models:", available.map((m) => `${m.provider}/${m.id}`), @@ -34,8 +32,7 @@ if (available.length > 0) { const { session } = await createAgentSession({ model: available[0], thinkingLevel: "medium", // off, low, medium, high - authStorage, - modelRegistry, + modelRuntime, }); try { diff --git a/packages/coding-agent/examples/sdk/09-api-keys-and-oauth.ts b/packages/coding-agent/examples/sdk/09-api-keys-and-oauth.ts index 120662e8..179b1d22 100644 --- a/packages/coding-agent/examples/sdk/09-api-keys-and-oauth.ts +++ b/packages/coding-agent/examples/sdk/09-api-keys-and-oauth.ts @@ -1,52 +1,34 @@ /** * API Keys and OAuth * - * Configure API key resolution via AuthStorage and ModelRegistry. + * Configure provider auth through ModelRuntime. */ -import { AuthStorage, createAgentSession, ModelRegistry, SessionManager } from "@earendil-works/pi-coding-agent"; - -// Default: AuthStorage uses ~/.pi/agent/auth.json -// ModelRegistry loads built-in + custom models from ~/.pi/agent/models.json -const authStorage = AuthStorage.create(); -const modelRegistry = ModelRegistry.create(authStorage); +import { createAgentSession, ModelRuntime, SessionManager } from "@earendil-works/pi-coding-agent"; +const modelRuntime = await ModelRuntime.create(); const { session: defaultAuthSession } = await createAgentSession({ sessionManager: SessionManager.inMemory(), - authStorage, - modelRegistry, + modelRuntime, }); -console.log("Session with default auth storage and model registry"); +console.log("Session with default model runtime"); defaultAuthSession.dispose(); -// Custom auth storage location -const customAuthStorage = AuthStorage.create("/tmp/my-app/auth.json"); -const customModelRegistry = ModelRegistry.create(customAuthStorage, "/tmp/my-app/models.json"); - +const customRuntime = await ModelRuntime.create({ + authPath: "/tmp/my-app/auth.json", + modelsPath: "/tmp/my-app/models.json", +}); const { session: customAuthSession } = await createAgentSession({ sessionManager: SessionManager.inMemory(), - authStorage: customAuthStorage, - modelRegistry: customModelRegistry, + modelRuntime: customRuntime, }); -console.log("Session with custom auth storage location"); +console.log("Session with custom auth and models locations"); customAuthSession.dispose(); -// Runtime API key override (not persisted to disk) -authStorage.setRuntimeApiKey("anthropic", "sk-my-temp-key"); +modelRuntime.setRuntimeApiKey("anthropic", "sk-my-temp-key"); const { session: runtimeKeySession } = await createAgentSession({ sessionManager: SessionManager.inMemory(), - authStorage, - modelRegistry, + modelRuntime, }); console.log("Session with runtime API key override"); runtimeKeySession.dispose(); - -// No models.json - only built-in models -const simpleRegistry = ModelRegistry.inMemory(authStorage); -const { session: builtInModelsSession } = await createAgentSession({ - sessionManager: SessionManager.inMemory(), - authStorage, - modelRegistry: simpleRegistry, -}); -console.log("Session with only built-in models"); -builtInModelsSession.dispose(); diff --git a/packages/coding-agent/examples/sdk/12-full-control.ts b/packages/coding-agent/examples/sdk/12-full-control.ts index cd23108b..066a5bbd 100644 --- a/packages/coding-agent/examples/sdk/12-full-control.ts +++ b/packages/coding-agent/examples/sdk/12-full-control.ts @@ -6,26 +6,22 @@ import { getModel } from "@earendil-works/pi-ai/compat"; import { - AuthStorage, createAgentSession, createExtensionRuntime, - ModelRegistry, + ModelRuntime, type ResourceLoader, SessionManager, SettingsManager, } from "@earendil-works/pi-coding-agent"; -// Custom auth storage location -const authStorage = AuthStorage.create("/tmp/my-agent/auth.json"); - -// Runtime API key override (not persisted) +const modelRuntime = await ModelRuntime.create({ + authPath: "/tmp/my-agent/auth.json", + modelsPath: "/tmp/my-agent/models.json", +}); if (process.env.MY_ANTHROPIC_KEY) { - authStorage.setRuntimeApiKey("anthropic", process.env.MY_ANTHROPIC_KEY); + modelRuntime.setRuntimeApiKey("anthropic", process.env.MY_ANTHROPIC_KEY); } -// Model registry with no custom models.json -const modelRegistry = ModelRegistry.inMemory(authStorage); - const model = getModel("anthropic", "claude-sonnet-4-5"); if (!model) throw new Error("Model not found"); @@ -55,8 +51,7 @@ const { session } = await createAgentSession({ agentDir: "/tmp/my-agent", model, thinkingLevel: "off", - authStorage, - modelRegistry, + modelRuntime, resourceLoader, tools: ["read", "bash"], sessionManager: SessionManager.inMemory(cwd), diff --git a/packages/coding-agent/examples/sdk/README.md b/packages/coding-agent/examples/sdk/README.md index f79157de..71689b01 100644 --- a/packages/coding-agent/examples/sdk/README.md +++ b/packages/coding-agent/examples/sdk/README.md @@ -34,46 +34,44 @@ npx tsx examples/sdk/01-minimal.ts ```typescript import { getModel } from "@earendil-works/pi-ai"; import { - AuthStorage, createAgentSession, DefaultResourceLoader, - ModelRegistry, + ModelRuntime, SessionManager, SettingsManager, } from "@earendil-works/pi-coding-agent"; -// Auth and models setup -const authStorage = AuthStorage.create(); -const modelRegistry = ModelRegistry.create(authStorage); +const modelRuntime = await ModelRuntime.create(); // Minimal -const { session } = await createAgentSession({ authStorage, modelRegistry }); +const { session } = await createAgentSession({ modelRuntime }); // Custom model const model = getModel("anthropic", "claude-opus-4-5"); -const { session } = await createAgentSession({ model, thinkingLevel: "high", authStorage, modelRegistry }); +const { session } = await createAgentSession({ model, thinkingLevel: "high", modelRuntime }); // Modify prompt const loader = new DefaultResourceLoader({ systemPromptOverride: (base) => `${base}\n\nBe concise.`, }); await loader.reload(); -const { session } = await createAgentSession({ resourceLoader: loader, authStorage, modelRegistry }); +const { session } = await createAgentSession({ resourceLoader: loader, modelRuntime }); // Read-only -const { session } = await createAgentSession({ tools: ["read", "grep", "find", "ls"], authStorage, modelRegistry }); +const { session } = await createAgentSession({ tools: ["read", "grep", "find", "ls"], modelRuntime }); // In-memory const { session } = await createAgentSession({ sessionManager: SessionManager.inMemory(), - authStorage, - modelRegistry, + modelRuntime, }); // Full control -const customAuth = AuthStorage.create("/my/app/auth.json"); -customAuth.setRuntimeApiKey("anthropic", process.env.MY_KEY!); -const customRegistry = ModelRegistry.create(customAuth); +const customRuntime = await ModelRuntime.create({ + authPath: "/my/app/auth.json", + modelsPath: "/my/app/models.json", +}); +customRuntime.setRuntimeApiKey("anthropic", process.env.MY_KEY!); const resourceLoader = new DefaultResourceLoader({ systemPromptOverride: () => "You are helpful.", @@ -86,8 +84,7 @@ await resourceLoader.reload(); const { session } = await createAgentSession({ model, - authStorage: customAuth, - modelRegistry: customRegistry, + modelRuntime: customRuntime, resourceLoader, tools: ["read", "bash", "my_tool"], customTools: [myTool], @@ -108,8 +105,7 @@ await session.prompt("Hello"); | Option | Default | Description | |--------|---------|-------------| -| `authStorage` | `AuthStorage.create()` | Credential storage | -| `modelRegistry` | `ModelRegistry.create(authStorage)` | Model registry | +| `modelRuntime` | Runtime using `agentDir/auth.json` and `models.json` | Canonical model and authentication runtime | | `cwd` | `process.cwd()` | Working directory | | `agentDir` | `~/.pi/agent` | Config directory | | `model` | From settings/first available | Model to use | diff --git a/packages/coding-agent/src/cli/list-models.ts b/packages/coding-agent/src/cli/list-models.ts index b648fb96..10c1cf36 100644 --- a/packages/coding-agent/src/cli/list-models.ts +++ b/packages/coding-agent/src/cli/list-models.ts @@ -6,7 +6,7 @@ import type { Api, Model } from "@earendil-works/pi-ai"; import { fuzzyFilter } from "@earendil-works/pi-tui"; import chalk from "chalk"; import { formatNoModelsAvailableMessage } from "../core/auth-guidance.ts"; -import type { ModelRegistry } from "../core/model-registry.ts"; +import type { ModelRuntime } from "../core/model-runtime.ts"; /** * Format a number as human-readable (e.g., 200000 -> "200K", 1000000 -> "1M") @@ -26,13 +26,13 @@ function formatTokenCount(count: number): string { /** * List available models, optionally filtered by search pattern */ -export async function listModels(modelRegistry: ModelRegistry, searchPattern?: string): Promise { - const loadError = modelRegistry.getError(); +export async function listModels(modelRuntime: ModelRuntime, searchPattern?: string): Promise { + const loadError = modelRuntime.getError(); if (loadError) { console.error(chalk.yellow(`Warning: errors loading models.json:\n${loadError}`)); } - const models = modelRegistry.getAvailable(); + const models = [...(await modelRuntime.getAvailable())]; 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 3a6035cb..7514c9d1 100644 --- a/packages/coding-agent/src/core/agent-session-services.ts +++ b/packages/coding-agent/src/core/agent-session-services.ts @@ -3,9 +3,8 @@ import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { Model } from "@earendil-works/pi-ai"; import { getAgentDir } from "../config.ts"; import { resolvePath } from "../utils/paths.ts"; -import { AuthStorage } from "./auth-storage.ts"; import type { SessionStartEvent, ToolDefinition } from "./extensions/index.ts"; -import { ModelRegistry } from "./model-registry.ts"; +import { ModelRuntime } from "./model-runtime.ts"; import { DefaultResourceLoader, type DefaultResourceLoaderOptions, @@ -38,9 +37,8 @@ export interface AgentSessionRuntimeDiagnostic { export interface CreateAgentSessionServicesOptions { cwd: string; agentDir?: string; - authStorage?: AuthStorage; settingsManager?: SettingsManager; - modelRegistry?: ModelRegistry; + modelRuntime?: ModelRuntime; extensionFlagValues?: Map; resourceLoaderOptions?: Omit; resourceLoaderReloadOptions?: ResourceLoaderReloadOptions; @@ -74,9 +72,8 @@ export interface CreateAgentSessionFromServicesOptions { export interface AgentSessionServices { cwd: string; agentDir: string; - authStorage: AuthStorage; + modelRuntime: ModelRuntime; settingsManager: SettingsManager; - modelRegistry: ModelRegistry; resourceLoader: ResourceLoader; diagnostics: AgentSessionRuntimeDiagnostic[]; } @@ -139,9 +136,13 @@ export async function createAgentSessionServices( ): Promise { const cwd = resolvePath(options.cwd); const agentDir = options.agentDir ? resolvePath(options.agentDir) : getAgentDir(); - const authStorage = options.authStorage ?? AuthStorage.create(join(agentDir, "auth.json")); + const modelRuntime = + options.modelRuntime ?? + (await ModelRuntime.create({ + authPath: join(agentDir, "auth.json"), + modelsPath: join(agentDir, "models.json"), + })); const settingsManager = options.settingsManager ?? SettingsManager.create(cwd, agentDir); - const modelRegistry = options.modelRegistry ?? ModelRegistry.create(authStorage, join(agentDir, "models.json")); const resourceLoader = new DefaultResourceLoader({ ...(options.resourceLoaderOptions ?? {}), cwd, @@ -154,7 +155,7 @@ export async function createAgentSessionServices( const extensionsResult = resourceLoader.getExtensions(); for (const { name, config, extensionPath } of extensionsResult.runtime.pendingProviderRegistrations) { try { - modelRegistry.registerProvider(name, config); + modelRuntime.registerProvider(name, config); } catch (error) { const message = error instanceof Error ? error.message : String(error); diagnostics.push({ @@ -169,9 +170,8 @@ export async function createAgentSessionServices( return { cwd, agentDir, - authStorage, + modelRuntime, settingsManager, - modelRegistry, resourceLoader, diagnostics, }; @@ -190,9 +190,8 @@ export async function createAgentSessionFromServices( return createAgentSession({ cwd: options.services.cwd, agentDir: options.services.agentDir, - authStorage: options.services.authStorage, + modelRuntime: options.services.modelRuntime, settingsManager: options.services.settingsManager, - modelRegistry: options.services.modelRegistry, resourceLoader: options.services.resourceLoader, sessionManager: options.sessionManager, model: options.model, diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index c4c2d366..6fb23906 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -24,7 +24,15 @@ import type { PrepareNextTurnContext, ThinkingLevel, } from "@earendil-works/pi-agent-core"; -import type { AssistantMessage, ImageContent, Message, Model, TextContent } from "@earendil-works/pi-ai/compat"; +import type { + AssistantMessage, + AuthResult, + ImageContent, + Message, + Model, + ProviderHeaders, + TextContent, +} from "@earendil-works/pi-ai/compat"; import { clampThinkingLevel, cleanupSessionResources, @@ -83,7 +91,8 @@ import { } from "./extensions/index.ts"; import { emitSessionShutdownEvent } from "./extensions/runner.ts"; import type { BashExecutionMessage, CustomMessage } from "./messages.ts"; -import type { ModelRegistry } from "./model-registry.ts"; +import { ModelRegistry } from "./model-registry.ts"; +import type { ModelRuntime } from "./model-runtime.ts"; import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.ts"; import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.ts"; import type { BranchSummaryEntry, CompactionEntry, SessionEntry, SessionManager } from "./session-manager.ts"; @@ -159,6 +168,12 @@ export type AgentSessionEventListener = (event: AgentSessionEvent) => void; // Types // ============================================================================ +function withoutDeletedHeaders(headers: ProviderHeaders | undefined): Record | undefined { + return headers + ? Object.fromEntries(Object.entries(headers).filter((entry): entry is [string, string] => entry[1] !== null)) + : undefined; +} + export interface AgentSessionConfig { agent: Agent; sessionManager: SessionManager; @@ -170,8 +185,8 @@ export interface AgentSessionConfig { resourceLoader: ResourceLoader; /** SDK custom tools registered outside extensions */ customTools?: ToolDefinition[]; - /** Model registry for API key resolution and model discovery */ - modelRegistry: ModelRegistry; + /** Canonical model/auth runtime used by coding-agent internals. */ + modelRuntime: ModelRuntime; /** Initial active built-in tool names. Default: [read, bash, edit, write] */ initialActiveToolNames?: string[]; /** Optional allowlist of tool names. When provided, only these tool names are exposed. */ @@ -325,8 +340,7 @@ export class AgentSession { private _extensionErrorListener?: ExtensionErrorListener; private _extensionErrorUnsubscriber?: () => void; - // Model registry for API key resolution - private _modelRegistry: ModelRegistry; + private _modelRuntime: ModelRuntime; // Tool registry for extension getTools/setTools private _toolRegistry: Map = new Map(); @@ -347,7 +361,7 @@ export class AgentSession { this._resourceLoader = config.resourceLoader; this._customTools = config.customTools ?? []; this._cwd = config.cwd; - this._modelRegistry = config.modelRegistry; + this._modelRuntime = config.modelRuntime; this._extensionRunnerRef = config.extensionRunnerRef; this._initialActiveToolNames = config.initialActiveToolNames; this._allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined; @@ -367,9 +381,8 @@ export class AgentSession { }); } - /** Model registry for API key resolution and model discovery */ - get modelRegistry(): ModelRegistry { - return this._modelRegistry; + get modelRuntime(): ModelRuntime { + return this._modelRuntime; } private async _getRequiredRequestAuth(model: Model): Promise<{ @@ -377,18 +390,25 @@ export class AgentSession { headers?: Record; env?: Record; }> { - const result = await this._modelRegistry.getApiKeyAndHeaders(model); - if (!result.ok) { - if (result.error.startsWith("No API key found")) { + let result: AuthResult | undefined; + try { + result = await this._modelRuntime.getAuth(model); + } catch (error) { + const cause = error instanceof Error ? error.cause : undefined; + if (cause instanceof Error && cause.message === "authHeader requires a resolved API key") { throw new Error(formatNoApiKeyFoundMessage(model.provider)); } - throw new Error(result.error); + throw error; } - if (result.apiKey) { - return { apiKey: result.apiKey, headers: result.headers, env: result.env }; + if (result?.auth.apiKey) { + return { + apiKey: result.auth.apiKey, + headers: withoutDeletedHeaders(result.auth.headers), + env: result.env, + }; } - const isOAuth = this._modelRegistry.isUsingOAuth(model); + const isOAuth = this._modelRuntime.isUsingOAuth(model.provider); if (isOAuth) { throw new Error( `Authentication failed for "${model.provider}". ` + @@ -408,8 +428,14 @@ export class AgentSession { return this._getRequiredRequestAuth(model); } - const result = await this._modelRegistry.getApiKeyAndHeaders(model); - return result.ok ? { apiKey: result.apiKey, headers: result.headers, env: result.env } : {}; + try { + const result = await this._modelRuntime.getAuth(model); + return result + ? { apiKey: result.auth.apiKey, headers: withoutDeletedHeaders(result.auth.headers), env: result.env } + : {}; + } catch { + return {}; + } } /** @@ -1141,8 +1167,11 @@ export class AgentSession { throw new Error(formatNoModelSelectedMessage()); } - if (!this._modelRegistry.hasConfiguredAuth(this.model)) { - const isOAuth = this._modelRegistry.isUsingOAuth(this.model); + const hasConfiguredAuth = + this._modelRuntime.hasConfiguredAuth(this.model.provider) || + (await this._modelRuntime.checkAuth(this.model.provider)) !== undefined; + if (!hasConfiguredAuth) { + const isOAuth = this._modelRuntime.isUsingOAuth(this.model.provider); if (isOAuth) { throw new Error( `Authentication failed for "${this.model.provider}". ` + @@ -1535,7 +1564,7 @@ export class AgentSession { * @throws Error if no auth is configured for the model */ async setModel(model: Model): Promise { - if (!this._modelRegistry.hasConfiguredAuth(model)) { + if (!(await this._modelRuntime.checkAuth(model.provider))) { throw new Error(`No API key for ${model.provider}/${model.id}`); } @@ -1565,7 +1594,13 @@ export class AgentSession { } private async _cycleScopedModel(direction: "forward" | "backward"): Promise { - const scopedModels = this._scopedModels.filter((scoped) => this._modelRegistry.hasConfiguredAuth(scoped.model)); + const checks = await Promise.all( + this._scopedModels.map(async (scoped) => ({ + scoped, + auth: await this._modelRuntime.checkAuth(scoped.model.provider), + })), + ); + const scopedModels = checks.filter(({ auth }) => auth !== undefined).map(({ scoped }) => scoped); if (scopedModels.length <= 1) return undefined; const currentModel = this.model; @@ -1594,7 +1629,7 @@ export class AgentSession { } private async _cycleAvailableModel(direction: "forward" | "backward"): Promise { - const availableModels = await this._modelRegistry.getAvailable(); + const availableModels = await this._modelRuntime.getAvailable(); if (availableModels.length <= 1) return undefined; const currentModel = this.model; @@ -2004,12 +2039,10 @@ export class AgentSession { let headers: Record | undefined; let env: Record | undefined; if (this.agent.streamFn === streamSimple) { - const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model); - if (!authResult.ok || !authResult.apiKey) { - return false; - } - apiKey = authResult.apiKey; - headers = authResult.headers; + const authResult = await this._modelRuntime.getAuth(this.model); + if (!authResult?.auth.apiKey) return false; + apiKey = authResult.auth.apiKey; + headers = withoutDeletedHeaders(authResult.auth.headers); env = authResult.env; } else { ({ apiKey, headers, env } = await this._getCompactionRequestAuth(this.model)); @@ -2267,7 +2300,7 @@ export class AgentSession { return; } - const refreshedModel = this._modelRegistry.find(currentModel.provider, currentModel.id); + const refreshedModel = this._modelRuntime.getModel(currentModel.provider, currentModel.id); if (!refreshedModel || refreshedModel === currentModel) { return; } @@ -2343,7 +2376,7 @@ export class AgentSession { refreshTools: () => this._refreshToolRegistry(), getCommands, setModel: async (model) => { - if (!this.modelRegistry.hasConfiguredAuth(model)) return false; + if (!this._modelRuntime.hasConfiguredAuth(model.provider)) return false; await this.setModel(model); return true; }, @@ -2383,11 +2416,11 @@ export class AgentSession { }, { registerProvider: (name, config) => { - this._modelRegistry.registerProvider(name, config); + this._modelRuntime.registerProvider(name, config); this._refreshCurrentModelFromRegistry(); }, unregisterProvider: (name) => { - this._modelRegistry.unregisterProvider(name); + this._modelRuntime.unregisterProvider(name); this._refreshCurrentModelFromRegistry(); }, }, @@ -2523,7 +2556,7 @@ export class AgentSession { extensionsResult.runtime, this._cwd, this.sessionManager, - this._modelRegistry, + new ModelRegistry(this._modelRuntime), ); if (this._extensionRunnerRef) { this._extensionRunnerRef.current = this._extensionRunner; diff --git a/packages/coding-agent/src/core/auth-storage.ts b/packages/coding-agent/src/core/auth-storage.ts index 725b6f3d..8991dc41 100644 --- a/packages/coding-agent/src/core/auth-storage.ts +++ b/packages/coding-agent/src/core/auth-storage.ts @@ -1,19 +1,9 @@ /** - * Credential storage for API keys and OAuth tokens. - * Handles loading, saving, and refreshing credentials from auth.json. - * - * Uses file locking to prevent race conditions when multiple pi instances - * try to refresh tokens simultaneously. + * CredentialStore implementation backed by auth.json. + * Provider auth orchestration belongs to ModelRuntime and pi-ai Models. */ -import { - findEnvKeys, - getEnvApiKey, - type OAuthCredentials, - type OAuthLoginCallbacks, - type OAuthProviderId, -} from "@earendil-works/pi-ai/compat"; -import { getOAuthApiKey, getOAuthProvider, getOAuthProviders } from "@earendil-works/pi-ai/oauth"; +import type { Credential, CredentialInfo, CredentialStore } from "@earendil-works/pi-ai"; import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs"; import { dirname, join } from "path"; import lockfile from "proper-lockfile"; @@ -21,29 +11,7 @@ import { getAgentDir } from "../config.ts"; import { normalizePath } from "../utils/paths.ts"; import { resolveConfigValue } from "./resolve-config-value.ts"; -export type ApiKeyCredential = { - type: "api_key"; - key: string; - env?: Record; -}; - -export type OAuthCredential = { - type: "oauth"; -} & OAuthCredentials; - -export type AuthCredential = ApiKeyCredential | OAuthCredential; - -export type AuthStorageData = Record; - -export type AuthStatus = { - configured: boolean; - source?: "stored" | "runtime" | "environment" | "fallback" | "models_json_key" | "models_json_command"; - label?: string; -}; - -export interface GetApiKeyOptions { - includeFallback?: boolean; -} +type AuthStorageData = Record; type LockResult = { result: T; @@ -200,11 +168,8 @@ export class InMemoryAuthStorageBackend implements AuthStorageBackend { /** * Credential storage backed by a JSON file. */ -export class AuthStorage { +export class AuthStorage implements CredentialStore { private data: AuthStorageData = {}; - private runtimeOverrides: Map = new Map(); - private loadError: Error | null = null; - private errors: Error[] = []; private storage: AuthStorageBackend; private constructor(storage: AuthStorageBackend) { @@ -226,26 +191,6 @@ export class AuthStorage { return AuthStorage.fromStorage(storage); } - /** - * Set a runtime API key override (not persisted to disk). - * Used for CLI --api-key flag. - */ - setRuntimeApiKey(provider: string, apiKey: string): void { - this.runtimeOverrides.set(provider, apiKey); - } - - /** - * Remove a runtime API key override. - */ - removeRuntimeApiKey(provider: string): void { - this.runtimeOverrides.delete(provider); - } - - private recordError(error: unknown): void { - const normalizedError = error instanceof Error ? error : new Error(String(error)); - this.errors.push(normalizedError); - } - private parseStorageData(content: string | undefined): AuthStorageData { if (!content) { return {}; @@ -264,276 +209,63 @@ export class AuthStorage { return { result: undefined }; }); this.data = this.parseStorageData(content); - this.loadError = null; - } catch (error) { - this.loadError = error as Error; - this.recordError(error); + } catch { + // Preserve the last valid in-memory snapshot. } } - private persistProviderChange(provider: string, credential: AuthCredential | undefined): AuthStorageData { - if (this.loadError) { - this.reload(); - } - - if (this.loadError) { - const error = new Error( - `Cannot update auth storage because it could not be loaded: ${this.loadError.message}`, - ); - this.recordError(error); - throw error; - } - - try { - let persistedData: AuthStorageData = {}; - this.storage.withLock((current) => { - const currentData = this.parseStorageData(current); - const merged: AuthStorageData = { ...currentData }; - if (credential) { - merged[provider] = credential; - } else { - delete merged[provider]; - } - persistedData = merged; - return { result: undefined, next: JSON.stringify(merged, null, 2) }; - }); - this.loadError = null; - return persistedData; - } catch (error) { - this.recordError(error); - throw error; - } + async read(provider: string): Promise { + const credential = this.data[provider]; + if (credential?.type !== "api_key") return credential; + if (credential.key === undefined) return credential; + return { ...credential, key: resolveConfigValue(credential.key, credential.env) }; } - /** - * Get credential for a provider. - */ - get(provider: string): AuthCredential | undefined { - return this.data[provider] ?? undefined; - } - - /** - * Get provider-scoped environment values for an API key credential. - */ - getProviderEnv(provider: string): Record | undefined { - const cred = this.data[provider]; - return cred?.type === "api_key" && cred.env ? { ...cred.env } : undefined; - } - - /** - * Set credential for a provider. - */ - set(provider: string, credential: AuthCredential): void { - this.data = this.persistProviderChange(provider, credential); - } - - /** - * Remove credential for a provider. - */ - remove(provider: string): void { - this.data = this.persistProviderChange(provider, undefined); - } - - /** - * List all providers with credentials. - */ - list(): string[] { - return Object.keys(this.data); - } - - /** - * Check if credentials exist for a provider in auth.json. - */ - has(provider: string): boolean { - return provider in this.data; - } - - /** - * Check if any form of auth is configured for a provider. - * Unlike getApiKey(), this doesn't refresh OAuth tokens. - */ - hasAuth(provider: string): boolean { - if (this.runtimeOverrides.has(provider)) return true; - if (this.data[provider]) return true; - if (getEnvApiKey(provider)) return true; - return false; - } - - /** - * Return auth status without exposing credential values or refreshing tokens. - */ - getAuthStatus(provider: string): AuthStatus { - if (this.data[provider]) { - return { configured: true, source: "stored" }; - } - - if (this.runtimeOverrides.has(provider)) { - return { configured: false, source: "runtime", label: "--api-key" }; - } - - const envKeys = findEnvKeys(provider); - if (envKeys?.[0]) { - return { configured: false, source: "environment", label: envKeys[0] }; - } - - return { configured: false }; - } - - /** - * Get all credentials (for passing to getOAuthApiKey). - */ - getAll(): AuthStorageData { - return { ...this.data }; - } - - drainErrors(): Error[] { - const drained = [...this.errors]; - this.errors = []; - return drained; - } - - /** - * Login to an OAuth provider. - */ - async login(providerId: OAuthProviderId, callbacks: OAuthLoginCallbacks): Promise { - const provider = getOAuthProvider(providerId); - if (!provider) { - throw new Error(`Unknown OAuth provider: ${providerId}`); - } - - const credentials = await provider.login(callbacks); - this.set(providerId, { type: "oauth", ...credentials }); - } - - /** - * Logout from a provider. - */ - logout(provider: string): void { - this.remove(provider); - } - - /** - * Refresh OAuth token with backend locking to prevent race conditions. - * Multiple pi instances may try to refresh simultaneously when tokens expire. - */ - private async refreshOAuthTokenWithLock( - providerId: OAuthProviderId, - ): Promise<{ apiKey: string; newCredentials: OAuthCredentials } | null> { - const provider = getOAuthProvider(providerId); - if (!provider) { - return null; - } - - const result = await this.storage.withLockAsync(async (current) => { - const currentData = this.parseStorageData(current); - this.data = currentData; - this.loadError = null; - - const cred = currentData[providerId]; - if (cred?.type !== "oauth") { - return { result: null }; + async modify( + provider: string, + fn: (current: Credential | undefined) => Promise, + ): Promise { + return this.storage.withLockAsync(async (content) => { + const currentData = this.parseStorageData(content); + const next = await fn(currentData[provider]); + if (next === undefined) { + this.data = currentData; + return { result: currentData[provider] }; } - if (Date.now() < cred.expires) { - return { result: { apiKey: provider.getApiKey(cred), newCredentials: cred } }; - } - - const oauthCreds: Record = {}; - for (const [key, value] of Object.entries(currentData)) { - if (value.type === "oauth") { - oauthCreds[key] = value; - } - } - - const refreshed = await getOAuthApiKey(providerId, oauthCreds); - if (!refreshed) { - return { result: null }; - } - - const merged: AuthStorageData = { - ...currentData, - [providerId]: { type: "oauth", ...refreshed.newCredentials }, - }; + const merged: AuthStorageData = { ...currentData, [provider]: next }; this.data = merged; - this.loadError = null; - return { result: refreshed, next: JSON.stringify(merged, null, 2) }; + return { result: next, next: JSON.stringify(merged, null, 2) }; }); - - return result; } - /** - * Get API key for a provider. - * Priority: - * 1. Runtime override (CLI --api-key) - * 2. API key from auth.json - * 3. OAuth token from auth.json (auto-refreshed with locking) - * 4. Environment variable - */ - async getApiKey(providerId: string, options: GetApiKeyOptions = {}): Promise { - // Runtime override takes highest priority - const runtimeKey = this.runtimeOverrides.get(providerId); - if (runtimeKey) { - return runtimeKey; - } - - const cred = this.data[providerId]; - - if (cred?.type === "api_key") { - return resolveConfigValue(cred.key, cred.env); - } - - if (cred?.type === "oauth") { - const provider = getOAuthProvider(providerId); - if (!provider) { - // Unknown OAuth provider, can't get API key - return undefined; - } - - // Check if token needs refresh - const needsRefresh = Date.now() >= cred.expires; - - if (needsRefresh) { - // Use locked refresh to prevent race conditions - try { - const result = await this.refreshOAuthTokenWithLock(providerId); - if (result) { - return result.apiKey; - } - } catch (error) { - this.recordError(error); - // Refresh failed - re-read file to check if another instance succeeded - this.reload(); - const updatedCred = this.data[providerId]; - - if (updatedCred?.type === "oauth" && Date.now() < updatedCred.expires) { - // Another instance refreshed successfully, use those credentials - return provider.getApiKey(updatedCred); - } - - // Refresh truly failed - return undefined so model discovery skips this provider - // User can /login to re-authenticate (credentials preserved for retry) - return undefined; - } - } else { - // Token not expired, use current access token - return provider.getApiKey(cred); - } - } - - if (options.includeFallback === false) return undefined; - - // Fall back to environment variable - const envKey = getEnvApiKey(providerId); - if (envKey) return envKey; - - return undefined; + async delete(provider: string): Promise { + await this.storage.withLockAsync(async (content) => { + const currentData = this.parseStorageData(content); + delete currentData[provider]; + this.data = currentData; + return { result: undefined, next: JSON.stringify(currentData, null, 2) }; + }); } - /** - * Get all registered OAuth providers - */ - getOAuthProviders() { - return getOAuthProviders(); + /** List credential metadata without resolving configured key values. */ + async list(): Promise { + return Object.entries(this.data).map(([providerId, credential]) => ({ providerId, type: credential.type })); + } +} + +/** + * One-off synchronous read of a stored credential from an auth.json file, + * without instantiating a store or resolving configured key values. + */ +export function readStoredCredential( + providerId: string, + authPath: string = join(getAgentDir(), "auth.json"), +): Credential | undefined { + try { + const data = JSON.parse(readFileSync(normalizePath(authPath), "utf-8")) as AuthStorageData; + return data[providerId]; + } catch { + return undefined; } } diff --git a/packages/coding-agent/src/core/cache-stats.ts b/packages/coding-agent/src/core/cache-stats.ts index 8c9a4943..f6054a2b 100644 --- a/packages/coding-agent/src/core/cache-stats.ts +++ b/packages/coding-agent/src/core/cache-stats.ts @@ -29,9 +29,9 @@ export interface CacheWasteTotals { missCount: number; } -/** Minimal pricing lookup, satisfied by ModelRegistry. Cost is $/million tokens. */ +/** Minimal pricing lookup, satisfied by ModelRuntime. Cost is $/million tokens. */ export interface ModelPriceSource { - find(provider: string, modelId: string): { cost: { cacheRead: number } } | undefined; + getModel(provider: string, modelId: string): { cost: { cacheRead: number } } | undefined; } /** The last request seen by the scan; everything in its prompt should be cached. */ @@ -79,7 +79,7 @@ function detectMiss( const readPerToken = usage.cacheRead > 0 ? usage.cost.cacheRead / usage.cacheRead - : (models.find(message.provider, message.model)?.cost.cacheRead ?? 0) / 1_000_000; + : (models.getModel(message.provider, message.model)?.cost.cacheRead ?? 0) / 1_000_000; return { missedTokens, diff --git a/packages/coding-agent/src/core/extensions/loader.ts b/packages/coding-agent/src/core/extensions/loader.ts index 0daeb262..01524d2b 100644 --- a/packages/coding-agent/src/core/extensions/loader.ts +++ b/packages/coding-agent/src/core/extensions/loader.ts @@ -10,6 +10,7 @@ import { fileURLToPath } from "node:url"; import * as _bundledPiAgentCore from "@earendil-works/pi-agent-core"; import * as _bundledPiAiCompat from "@earendil-works/pi-ai/compat"; import * as _bundledPiAiOauth from "@earendil-works/pi-ai/oauth"; +import * as _bundledPiAiProviders from "@earendil-works/pi-ai/providers/all"; import type { KeyId } from "@earendil-works/pi-tui"; import * as _bundledPiTui from "@earendil-works/pi-tui"; import { createJiti } from "jiti/static"; @@ -58,12 +59,14 @@ const VIRTUAL_MODULES: Record = { "@earendil-works/pi-ai": _bundledPiAiCompat, "@earendil-works/pi-ai/compat": _bundledPiAiCompat, "@earendil-works/pi-ai/oauth": _bundledPiAiOauth, + "@earendil-works/pi-ai/providers/all": _bundledPiAiProviders, "@earendil-works/pi-coding-agent": _bundledPiCodingAgent, "@mariozechner/pi-agent-core": _bundledPiAgentCore, "@mariozechner/pi-tui": _bundledPiTui, "@mariozechner/pi-ai": _bundledPiAiCompat, "@mariozechner/pi-ai/compat": _bundledPiAiCompat, "@mariozechner/pi-ai/oauth": _bundledPiAiOauth, + "@mariozechner/pi-ai/providers/all": _bundledPiAiProviders, "@mariozechner/pi-coding-agent": _bundledPiCodingAgent, }; @@ -102,20 +105,26 @@ function getAliases(): Record { // global API keep working at runtime until compat is removed. const piAiCompatEntry = resolveWorkspaceOrImport("ai/dist/compat.js", "@earendil-works/pi-ai/compat"); const piAiOauthEntry = resolveWorkspaceOrImport("ai/dist/oauth.js", "@earendil-works/pi-ai/oauth"); + const piAiProvidersEntry = resolveWorkspaceOrImport( + "ai/dist/providers/all.js", + "@earendil-works/pi-ai/providers/all", + ); _aliases = { "@earendil-works/pi-coding-agent": piCodingAgentEntry, "@earendil-works/pi-agent-core": piAgentCoreEntry, "@earendil-works/pi-tui": piTuiEntry, - "@earendil-works/pi-ai": piAiCompatEntry, + "@earendil-works/pi-ai/providers/all": piAiProvidersEntry, "@earendil-works/pi-ai/compat": piAiCompatEntry, "@earendil-works/pi-ai/oauth": piAiOauthEntry, + "@earendil-works/pi-ai": piAiCompatEntry, "@mariozechner/pi-coding-agent": piCodingAgentEntry, "@mariozechner/pi-agent-core": piAgentCoreEntry, "@mariozechner/pi-tui": piTuiEntry, - "@mariozechner/pi-ai": piAiCompatEntry, + "@mariozechner/pi-ai/providers/all": piAiProvidersEntry, "@mariozechner/pi-ai/compat": piAiCompatEntry, "@mariozechner/pi-ai/oauth": piAiOauthEntry, + "@mariozechner/pi-ai": piAiCompatEntry, typebox: typeboxEntry, "typebox/compile": typeboxCompileEntry, "typebox/value": typeboxValueEntry, diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 2a699a7d..b853ec49 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -602,6 +602,10 @@ export class ExtensionRunner { }); } + getModelRegistry(): ModelRegistry { + return this.modelRegistry; + } + getRegisteredCommands(): ResolvedCommand[] { this.commandDiagnostics = []; return this.resolveRegisteredCommands(); diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index a3c9f735..93195cbe 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -1424,14 +1424,14 @@ export interface ProviderConfig { oauth?: { /** Display name for the provider in login UI. */ name: string; + /** @deprecated Retained for source compatibility; canonical auth flows ignore it. */ + usesCallbackServer?: boolean; /** Run the login flow, return credentials to persist. */ login(callbacks: OAuthLoginCallbacks): Promise; /** Refresh expired credentials, return updated credentials to persist. */ refreshToken(credentials: OAuthCredentials): Promise; /** Convert credentials to API key string for the provider. */ getApiKey(credentials: OAuthCredentials): string; - /** Optional: modify models for this provider (e.g., update baseUrl based on credentials). */ - modifyModels?(models: Model[], credentials: OAuthCredentials): Model[]; }; } diff --git a/packages/coding-agent/src/core/model-config.ts b/packages/coding-agent/src/core/model-config.ts new file mode 100644 index 00000000..9b54be81 --- /dev/null +++ b/packages/coding-agent/src/core/model-config.ts @@ -0,0 +1,277 @@ +/** Immutable, credential-blind models.json snapshot. */ + +import { readFile } from "node:fs/promises"; +import { type Static, Type } from "typebox"; +import { Compile } from "typebox/compile"; +import type { TLocalizedValidationError } from "typebox/error"; +import { stripJsonComments } from "../utils/json.ts"; +import { normalizePath } from "../utils/paths.ts"; + +const PercentileCutoffsSchema = Type.Object({ + p50: Type.Optional(Type.Number()), + p75: Type.Optional(Type.Number()), + p90: Type.Optional(Type.Number()), + p99: Type.Optional(Type.Number()), +}); + +const OpenRouterRoutingSchema = Type.Object({ + allow_fallbacks: Type.Optional(Type.Boolean()), + require_parameters: Type.Optional(Type.Boolean()), + data_collection: Type.Optional(Type.Union([Type.Literal("deny"), Type.Literal("allow")])), + zdr: Type.Optional(Type.Boolean()), + enforce_distillable_text: Type.Optional(Type.Boolean()), + order: Type.Optional(Type.Array(Type.String())), + only: Type.Optional(Type.Array(Type.String())), + ignore: Type.Optional(Type.Array(Type.String())), + quantizations: Type.Optional(Type.Array(Type.String())), + sort: Type.Optional( + Type.Union([ + Type.String(), + Type.Object({ + by: Type.Optional(Type.String()), + partition: Type.Optional(Type.Union([Type.String(), Type.Null()])), + }), + ]), + ), + max_price: Type.Optional( + Type.Object({ + prompt: Type.Optional(Type.Union([Type.Number(), Type.String()])), + completion: Type.Optional(Type.Union([Type.Number(), Type.String()])), + image: Type.Optional(Type.Union([Type.Number(), Type.String()])), + audio: Type.Optional(Type.Union([Type.Number(), Type.String()])), + request: Type.Optional(Type.Union([Type.Number(), Type.String()])), + }), + ), + preferred_min_throughput: Type.Optional(Type.Union([Type.Number(), PercentileCutoffsSchema])), + preferred_max_latency: Type.Optional(Type.Union([Type.Number(), PercentileCutoffsSchema])), +}); + +const VercelGatewayRoutingSchema = Type.Object({ + only: Type.Optional(Type.Array(Type.String())), + order: Type.Optional(Type.Array(Type.String())), +}); + +const ThinkingLevelMapValueSchema = Type.Union([Type.String(), Type.Null()]); +const ThinkingLevelMapSchema = Type.Object({ + off: Type.Optional(ThinkingLevelMapValueSchema), + minimal: Type.Optional(ThinkingLevelMapValueSchema), + low: Type.Optional(ThinkingLevelMapValueSchema), + medium: Type.Optional(ThinkingLevelMapValueSchema), + high: Type.Optional(ThinkingLevelMapValueSchema), + xhigh: Type.Optional(ThinkingLevelMapValueSchema), + max: Type.Optional(ThinkingLevelMapValueSchema), +}); + +const ChatTemplateKwargScalarSchema = Type.Union([Type.String(), Type.Number(), Type.Boolean(), Type.Null()]); +const ChatTemplateKwargVariableSchema = Type.Object({ + $var: Type.Union([Type.Literal("thinking.enabled"), Type.Literal("thinking.effort")]), + omitWhenOff: Type.Optional(Type.Boolean()), +}); +const ChatTemplateKwargSchema = Type.Union([ChatTemplateKwargScalarSchema, ChatTemplateKwargVariableSchema]); + +const OpenAICompletionsCompatSchema = Type.Object({ + supportsStore: Type.Optional(Type.Boolean()), + supportsDeveloperRole: Type.Optional(Type.Boolean()), + supportsReasoningEffort: Type.Optional(Type.Boolean()), + supportsUsageInStreaming: Type.Optional(Type.Boolean()), + maxTokensField: Type.Optional(Type.Union([Type.Literal("max_completion_tokens"), Type.Literal("max_tokens")])), + requiresToolResultName: Type.Optional(Type.Boolean()), + requiresAssistantAfterToolResult: Type.Optional(Type.Boolean()), + requiresThinkingAsText: Type.Optional(Type.Boolean()), + requiresReasoningContentOnAssistantMessages: Type.Optional(Type.Boolean()), + thinkingFormat: Type.Optional( + Type.Union([ + Type.Literal("openai"), + Type.Literal("openrouter"), + Type.Literal("together"), + Type.Literal("deepseek"), + Type.Literal("zai"), + Type.Literal("qwen"), + Type.Literal("chat-template"), + Type.Literal("qwen-chat-template"), + Type.Literal("string-thinking"), + Type.Literal("ant-ling"), + ]), + ), + chatTemplateKwargs: Type.Optional(Type.Record(Type.String(), ChatTemplateKwargSchema)), + cacheControlFormat: Type.Optional(Type.Literal("anthropic")), + openRouterRouting: Type.Optional(OpenRouterRoutingSchema), + vercelGatewayRouting: Type.Optional(VercelGatewayRoutingSchema), + supportsStrictMode: Type.Optional(Type.Boolean()), + supportsLongCacheRetention: Type.Optional(Type.Boolean()), +}); + +const OpenAIResponsesCompatSchema = Type.Object({ + supportsDeveloperRole: Type.Optional(Type.Boolean()), + sendSessionIdHeader: Type.Optional(Type.Boolean()), + supportsLongCacheRetention: Type.Optional(Type.Boolean()), +}); + +const AnthropicMessagesCompatSchema = Type.Object({ + supportsEagerToolInputStreaming: Type.Optional(Type.Boolean()), + supportsLongCacheRetention: Type.Optional(Type.Boolean()), + sendSessionAffinityHeaders: Type.Optional(Type.Boolean()), + supportsCacheControlOnTools: Type.Optional(Type.Boolean()), + forceAdaptiveThinking: Type.Optional(Type.Boolean()), +}); + +const ProviderCompatSchema = Type.Union([ + OpenAICompletionsCompatSchema, + OpenAIResponsesCompatSchema, + AnthropicMessagesCompatSchema, +]); + +const ModelCostRatesSchema = { + input: Type.Number(), + output: Type.Number(), + cacheRead: Type.Number(), + cacheWrite: Type.Number(), +}; +const ModelCostTierSchema = Type.Object({ + inputTokensAbove: Type.Number(), + ...ModelCostRatesSchema, +}); +const ModelCostSchema = Type.Object({ + ...ModelCostRatesSchema, + tiers: Type.Optional(Type.Array(ModelCostTierSchema)), +}); + +const ModelDefinitionSchema = Type.Object({ + id: Type.String({ minLength: 1 }), + name: Type.Optional(Type.String({ minLength: 1 })), + api: Type.Optional(Type.String({ minLength: 1 })), + baseUrl: Type.Optional(Type.String({ minLength: 1 })), + reasoning: Type.Optional(Type.Boolean()), + thinkingLevelMap: Type.Optional(ThinkingLevelMapSchema), + input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))), + cost: Type.Optional(ModelCostSchema), + contextWindow: Type.Optional(Type.Number()), + maxTokens: Type.Optional(Type.Number()), + headers: Type.Optional(Type.Record(Type.String(), Type.String())), + compat: Type.Optional(ProviderCompatSchema), +}); + +const ModelOverrideSchema = Type.Object({ + name: Type.Optional(Type.String({ minLength: 1 })), + reasoning: Type.Optional(Type.Boolean()), + thinkingLevelMap: Type.Optional(ThinkingLevelMapSchema), + input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))), + cost: Type.Optional( + Type.Object({ + input: Type.Optional(Type.Number()), + output: Type.Optional(Type.Number()), + cacheRead: Type.Optional(Type.Number()), + cacheWrite: Type.Optional(Type.Number()), + tiers: Type.Optional(Type.Array(ModelCostTierSchema)), + }), + ), + contextWindow: Type.Optional(Type.Number()), + maxTokens: Type.Optional(Type.Number()), + headers: Type.Optional(Type.Record(Type.String(), Type.String())), + compat: Type.Optional(ProviderCompatSchema), +}); + +const ProviderConfigSchema = Type.Object({ + name: Type.Optional(Type.String({ minLength: 1 })), + baseUrl: Type.Optional(Type.String({ minLength: 1 })), + apiKey: Type.Optional(Type.String({ minLength: 1 })), + api: Type.Optional(Type.String({ minLength: 1 })), + headers: Type.Optional(Type.Record(Type.String(), Type.String())), + compat: Type.Optional(ProviderCompatSchema), + authHeader: Type.Optional(Type.Boolean()), + models: Type.Optional(Type.Array(ModelDefinitionSchema)), + modelOverrides: Type.Optional(Type.Record(Type.String(), ModelOverrideSchema)), +}); + +const ModelsConfigSchema = Type.Object({ + providers: Type.Record(Type.String(), ProviderConfigSchema), +}); +const validateModelsConfig = Compile(ModelsConfigSchema); + +export type ModelsJsonModel = Static; +export type ModelsJsonModelOverride = Static; +export type ModelsJsonProvider = Static; +type ModelsJson = Static; + +function formatValidationPath(error: TLocalizedValidationError): string { + if (error.keyword === "required") { + const requiredProperties = (error.params as { requiredProperties?: string[] }).requiredProperties; + const requiredProperty = requiredProperties?.[0]; + if (requiredProperty) { + const basePath = error.instancePath.replace(/^\//, "").replace(/\//g, "."); + return basePath ? `${basePath}.${requiredProperty}` : requiredProperty; + } + } + const path = error.instancePath.replace(/^\//, "").replace(/\//g, "."); + return path || "root"; +} + +function deepFreeze(value: T): T { + if (typeof value !== "object" || value === null || Object.isFrozen(value)) return value; + for (const child of Object.values(value)) deepFreeze(child); + return Object.freeze(value); +} + +/** One immutable load of models.json. */ +export class ModelConfig { + private readonly providers: ReadonlyMap; + private readonly error: string | undefined; + + private constructor(providers: ReadonlyMap, error?: string) { + this.providers = providers; + this.error = error; + } + + static async load(modelsJsonPath: string | undefined): Promise { + if (!modelsJsonPath) return new ModelConfig(new Map()); + const path = normalizePath(modelsJsonPath); + let content: string; + try { + content = await readFile(path, "utf-8"); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") return new ModelConfig(new Map()); + return new ModelConfig( + new Map(), + `Failed to load models.json: ${error instanceof Error ? error.message : error}\n\nFile: ${path}`, + ); + } + + let parsed: unknown; + try { + parsed = JSON.parse(stripJsonComments(content)); + } catch (error) { + return new ModelConfig( + new Map(), + `Failed to parse models.json: ${error instanceof Error ? error.message : error}\n\nFile: ${path}`, + ); + } + + if (!validateModelsConfig.Check(parsed)) { + const errors = + validateModelsConfig + .Errors(parsed) + .map((error) => ` - ${formatValidationPath(error)}: ${error.message}`) + .join("\n") || "Unknown schema error"; + return new ModelConfig(new Map(), `Invalid models.json schema:\n${errors}\n\nFile: ${path}`); + } + + const config = parsed as ModelsJson; + const providers = new Map(); + for (const [providerId, provider] of Object.entries(config.providers)) { + providers.set(providerId, deepFreeze(structuredClone(provider))); + } + return new ModelConfig(providers); + } + + getProvider(providerId: string): ModelsJsonProvider | undefined { + return this.providers.get(providerId); + } + + getProviderIds(): readonly string[] { + return [...this.providers.keys()]; + } + + getError(): string | undefined { + return this.error; + } +} diff --git a/packages/coding-agent/src/core/model-registry.ts b/packages/coding-agent/src/core/model-registry.ts index ccf15430..a3ef09fa 100644 --- a/packages/coding-agent/src/core/model-registry.ts +++ b/packages/coding-agent/src/core/model-registry.ts @@ -1,261 +1,8 @@ -/** - * Model registry - manages built-in and custom models, provides API key resolution. - */ - -import { - type AnthropicMessagesCompat, - type Api, - type AssistantMessageEventStream, - type Context, - getModels, - getProviders, - type KnownProvider, - type Model, - type OAuthProviderInterface, - type OpenAICompletionsCompat, - type OpenAIResponsesCompat, - registerApiProvider, - resetApiProviders, - type SimpleStreamOptions, -} from "@earendil-works/pi-ai/compat"; -import { registerOAuthProvider, resetOAuthProviders } from "@earendil-works/pi-ai/oauth"; -import { existsSync, readFileSync } from "fs"; -import { join } from "path"; -import { type Static, Type } from "typebox"; -import { Compile } from "typebox/compile"; -import type { TLocalizedValidationError } from "typebox/error"; -import { getAgentDir } from "../config.ts"; -import { stripJsonComments } from "../utils/json.ts"; -import { normalizePath } from "../utils/paths.ts"; -import type { AuthStatus, AuthStorage } from "./auth-storage.ts"; -import { BUILT_IN_PROVIDER_DISPLAY_NAMES } from "./provider-display-names.ts"; -import { - clearConfigValueCache, - getConfigValueEnvVarNames, - isCommandConfigValue, - isConfigValueConfigured, - resolveConfigValueOrThrow, - resolveConfigValueUncached, - resolveHeadersOrThrow, -} from "./resolve-config-value.ts"; - -// Schema for OpenRouter routing preferences -const PercentileCutoffsSchema = Type.Object({ - p50: Type.Optional(Type.Number()), - p75: Type.Optional(Type.Number()), - p90: Type.Optional(Type.Number()), - p99: Type.Optional(Type.Number()), -}); - -const OpenRouterRoutingSchema = Type.Object({ - allow_fallbacks: Type.Optional(Type.Boolean()), - require_parameters: Type.Optional(Type.Boolean()), - data_collection: Type.Optional(Type.Union([Type.Literal("deny"), Type.Literal("allow")])), - zdr: Type.Optional(Type.Boolean()), - enforce_distillable_text: Type.Optional(Type.Boolean()), - order: Type.Optional(Type.Array(Type.String())), - only: Type.Optional(Type.Array(Type.String())), - ignore: Type.Optional(Type.Array(Type.String())), - quantizations: Type.Optional(Type.Array(Type.String())), - sort: Type.Optional( - Type.Union([ - Type.String(), - Type.Object({ - by: Type.Optional(Type.String()), - partition: Type.Optional(Type.Union([Type.String(), Type.Null()])), - }), - ]), - ), - max_price: Type.Optional( - Type.Object({ - prompt: Type.Optional(Type.Union([Type.Number(), Type.String()])), - completion: Type.Optional(Type.Union([Type.Number(), Type.String()])), - image: Type.Optional(Type.Union([Type.Number(), Type.String()])), - audio: Type.Optional(Type.Union([Type.Number(), Type.String()])), - request: Type.Optional(Type.Union([Type.Number(), Type.String()])), - }), - ), - preferred_min_throughput: Type.Optional(Type.Union([Type.Number(), PercentileCutoffsSchema])), - preferred_max_latency: Type.Optional(Type.Union([Type.Number(), PercentileCutoffsSchema])), -}); - -// Schema for Vercel AI Gateway routing preferences -const VercelGatewayRoutingSchema = Type.Object({ - only: Type.Optional(Type.Array(Type.String())), - order: Type.Optional(Type.Array(Type.String())), -}); - -// Schema for thinking level support and provider-specific values -const ThinkingLevelMapValueSchema = Type.Union([Type.String(), Type.Null()]); -const ThinkingLevelMapSchema = Type.Object({ - off: Type.Optional(ThinkingLevelMapValueSchema), - minimal: Type.Optional(ThinkingLevelMapValueSchema), - low: Type.Optional(ThinkingLevelMapValueSchema), - medium: Type.Optional(ThinkingLevelMapValueSchema), - high: Type.Optional(ThinkingLevelMapValueSchema), - xhigh: Type.Optional(ThinkingLevelMapValueSchema), - max: Type.Optional(ThinkingLevelMapValueSchema), -}); - -const ChatTemplateKwargScalarSchema = Type.Union([Type.String(), Type.Number(), Type.Boolean(), Type.Null()]); -const ChatTemplateKwargVariableSchema = Type.Object({ - $var: Type.Union([Type.Literal("thinking.enabled"), Type.Literal("thinking.effort")]), - omitWhenOff: Type.Optional(Type.Boolean()), -}); -const ChatTemplateKwargSchema = Type.Union([ChatTemplateKwargScalarSchema, ChatTemplateKwargVariableSchema]); - -const OpenAICompletionsCompatSchema = Type.Object({ - supportsStore: Type.Optional(Type.Boolean()), - supportsDeveloperRole: Type.Optional(Type.Boolean()), - supportsReasoningEffort: Type.Optional(Type.Boolean()), - supportsUsageInStreaming: Type.Optional(Type.Boolean()), - maxTokensField: Type.Optional(Type.Union([Type.Literal("max_completion_tokens"), Type.Literal("max_tokens")])), - requiresToolResultName: Type.Optional(Type.Boolean()), - requiresAssistantAfterToolResult: Type.Optional(Type.Boolean()), - requiresThinkingAsText: Type.Optional(Type.Boolean()), - requiresReasoningContentOnAssistantMessages: Type.Optional(Type.Boolean()), - thinkingFormat: Type.Optional( - Type.Union([ - Type.Literal("openai"), - Type.Literal("openrouter"), - Type.Literal("together"), - Type.Literal("deepseek"), - Type.Literal("zai"), - Type.Literal("qwen"), - Type.Literal("chat-template"), - Type.Literal("qwen-chat-template"), - Type.Literal("string-thinking"), - Type.Literal("ant-ling"), - ]), - ), - chatTemplateKwargs: Type.Optional(Type.Record(Type.String(), ChatTemplateKwargSchema)), - cacheControlFormat: Type.Optional(Type.Literal("anthropic")), - openRouterRouting: Type.Optional(OpenRouterRoutingSchema), - vercelGatewayRouting: Type.Optional(VercelGatewayRoutingSchema), - supportsStrictMode: Type.Optional(Type.Boolean()), - supportsLongCacheRetention: Type.Optional(Type.Boolean()), -}); - -const OpenAIResponsesCompatSchema = Type.Object({ - supportsDeveloperRole: Type.Optional(Type.Boolean()), - sendSessionIdHeader: Type.Optional(Type.Boolean()), - supportsLongCacheRetention: Type.Optional(Type.Boolean()), -}); - -const AnthropicMessagesCompatSchema = Type.Object({ - supportsEagerToolInputStreaming: Type.Optional(Type.Boolean()), - supportsLongCacheRetention: Type.Optional(Type.Boolean()), - sendSessionAffinityHeaders: Type.Optional(Type.Boolean()), - supportsCacheControlOnTools: Type.Optional(Type.Boolean()), - forceAdaptiveThinking: Type.Optional(Type.Boolean()), -}); - -const ProviderCompatSchema = Type.Union([ - OpenAICompletionsCompatSchema, - OpenAIResponsesCompatSchema, - AnthropicMessagesCompatSchema, -]); - -const ModelCostRatesSchema = { - input: Type.Number(), - output: Type.Number(), - cacheRead: Type.Number(), - cacheWrite: Type.Number(), -}; -const ModelCostTierSchema = Type.Object({ - inputTokensAbove: Type.Number(), - ...ModelCostRatesSchema, -}); -const ModelCostSchema = Type.Object({ - ...ModelCostRatesSchema, - tiers: Type.Optional(Type.Array(ModelCostTierSchema)), -}); - -// Schema for custom model definition -// Most fields are optional with sensible defaults for local models (Ollama, LM Studio, etc.) -const ModelDefinitionSchema = Type.Object({ - id: Type.String({ minLength: 1 }), - name: Type.Optional(Type.String({ minLength: 1 })), - api: Type.Optional(Type.String({ minLength: 1 })), - baseUrl: Type.Optional(Type.String({ minLength: 1 })), - reasoning: Type.Optional(Type.Boolean()), - thinkingLevelMap: Type.Optional(ThinkingLevelMapSchema), - input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))), - cost: Type.Optional(ModelCostSchema), - contextWindow: Type.Optional(Type.Number()), - maxTokens: Type.Optional(Type.Number()), - headers: Type.Optional(Type.Record(Type.String(), Type.String())), - compat: Type.Optional(ProviderCompatSchema), -}); - -// Schema for per-model overrides (all fields optional, merged with built-in model) -const ModelOverrideSchema = Type.Object({ - name: Type.Optional(Type.String({ minLength: 1 })), - reasoning: Type.Optional(Type.Boolean()), - thinkingLevelMap: Type.Optional(ThinkingLevelMapSchema), - input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))), - cost: Type.Optional( - Type.Object({ - input: Type.Optional(Type.Number()), - output: Type.Optional(Type.Number()), - cacheRead: Type.Optional(Type.Number()), - cacheWrite: Type.Optional(Type.Number()), - tiers: Type.Optional(Type.Array(ModelCostTierSchema)), - }), - ), - contextWindow: Type.Optional(Type.Number()), - maxTokens: Type.Optional(Type.Number()), - headers: Type.Optional(Type.Record(Type.String(), Type.String())), - compat: Type.Optional(ProviderCompatSchema), -}); - -type ModelOverride = Static; - -const ProviderConfigSchema = Type.Object({ - name: Type.Optional(Type.String({ minLength: 1 })), - baseUrl: Type.Optional(Type.String({ minLength: 1 })), - apiKey: Type.Optional(Type.String({ minLength: 1 })), - api: Type.Optional(Type.String({ minLength: 1 })), - headers: Type.Optional(Type.Record(Type.String(), Type.String())), - compat: Type.Optional(ProviderCompatSchema), - authHeader: Type.Optional(Type.Boolean()), - models: Type.Optional(Type.Array(ModelDefinitionSchema)), - modelOverrides: Type.Optional(Type.Record(Type.String(), ModelOverrideSchema)), -}); - -const ModelsConfigSchema = Type.Object({ - providers: Type.Record(Type.String(), ProviderConfigSchema), -}); - -const validateModelsConfig = Compile(ModelsConfigSchema); - -type ModelsConfig = Static; - -function formatValidationPath(error: TLocalizedValidationError): string { - if (error.keyword === "required") { - const requiredProperties = (error.params as { requiredProperties?: string[] }).requiredProperties; - const requiredProperty = requiredProperties?.[0]; - if (requiredProperty) { - const basePath = error.instancePath.replace(/^\//, "").replace(/\//g, "."); - return basePath ? `${basePath}.${requiredProperty}` : requiredProperty; - } - } - const path = error.instancePath.replace(/^\//, "").replace(/\//g, "."); - return path || "root"; -} - -/** Provider override config (baseUrl, compat) without request auth/headers */ -interface ProviderOverride { - baseUrl?: string; - compat?: Model["compat"]; -} - -interface ProviderRequestConfig { - apiKey?: string; - headers?: Record; - authHeader?: boolean; -} +import type { Api, Model } from "@earendil-works/pi-ai"; +import type { ModelRuntime } from "./model-runtime.ts"; +import type { AuthStatus, ProviderConfigInput } from "./provider-composer.ts"; +export type { ProviderConfigInput } from "./provider-composer.ts"; export type ResolvedRequestAuth = | { ok: true; @@ -263,756 +10,117 @@ export type ResolvedRequestAuth = headers?: Record; env?: Record; } - | { - ok: false; - error: string; - }; - -/** Result of loading custom models from models.json */ -interface CustomModelsResult { - models: Model[]; - /** Providers with baseUrl/headers/apiKey overrides for built-in models */ - overrides: Map; - /** Per-model overrides: provider -> modelId -> override */ - modelOverrides: Map>; - error: string | undefined; -} - -function emptyCustomModelsResult(error?: string): CustomModelsResult { - return { models: [], overrides: new Map(), modelOverrides: new Map(), error }; -} - -function mergeCompat( - baseCompat: Model["compat"], - overrideCompat: ModelOverride["compat"], -): Model["compat"] | undefined { - if (!overrideCompat) return baseCompat; - - const base = baseCompat as OpenAICompletionsCompat | OpenAIResponsesCompat | AnthropicMessagesCompat | undefined; - const override = overrideCompat as OpenAICompletionsCompat | OpenAIResponsesCompat | AnthropicMessagesCompat; - const merged = { ...base, ...override } as OpenAICompletionsCompat | OpenAIResponsesCompat | AnthropicMessagesCompat; - - const baseCompletions = base as OpenAICompletionsCompat | undefined; - const overrideCompletions = override as OpenAICompletionsCompat; - const mergedCompletions = merged as OpenAICompletionsCompat; - - if (baseCompletions?.openRouterRouting || overrideCompletions.openRouterRouting) { - mergedCompletions.openRouterRouting = { - ...baseCompletions?.openRouterRouting, - ...overrideCompletions.openRouterRouting, - }; - } - - if (baseCompletions?.vercelGatewayRouting || overrideCompletions.vercelGatewayRouting) { - mergedCompletions.vercelGatewayRouting = { - ...baseCompletions?.vercelGatewayRouting, - ...overrideCompletions.vercelGatewayRouting, - }; - } - - if (baseCompletions?.chatTemplateKwargs || overrideCompletions.chatTemplateKwargs) { - mergedCompletions.chatTemplateKwargs = { - ...baseCompletions?.chatTemplateKwargs, - ...overrideCompletions.chatTemplateKwargs, - }; - } - - return merged as Model["compat"]; -} + | { ok: false; error: string }; +export { clearApiKeyCache } from "./provider-composer.ts"; /** - * Deep merge a model override into a model. - * Handles nested objects (cost, compat) by merging rather than replacing. - */ -function applyModelOverride(model: Model, override: ModelOverride): Model { - const result = { ...model }; - - // Simple field overrides - if (override.name !== undefined) result.name = override.name; - if (override.reasoning !== undefined) result.reasoning = override.reasoning; - if (override.thinkingLevelMap !== undefined) { - result.thinkingLevelMap = { ...model.thinkingLevelMap, ...override.thinkingLevelMap }; - } - if (override.input !== undefined) result.input = override.input as ("text" | "image")[]; - if (override.contextWindow !== undefined) result.contextWindow = override.contextWindow; - if (override.maxTokens !== undefined) result.maxTokens = override.maxTokens; - - // Merge cost (partial override) - if (override.cost) { - result.cost = { - input: override.cost.input ?? model.cost.input, - output: override.cost.output ?? model.cost.output, - cacheRead: override.cost.cacheRead ?? model.cost.cacheRead, - cacheWrite: override.cost.cacheWrite ?? model.cost.cacheWrite, - tiers: override.cost.tiers ?? model.cost.tiers, - }; - } - - // Deep merge compat - result.compat = mergeCompat(model.compat, override.compat); - - return result; -} - -/** Clear the config value command cache. Exported for testing. */ -export const clearApiKeyCache = clearConfigValueCache; - -/** - * Model registry - loads and manages models, resolves API keys via AuthStorage. + * Synchronous compatibility facade exposed to extensions. + * Coding-agent internals use ModelRuntime directly. */ export class ModelRegistry { - private models: Model[] = []; - private providerRequestConfigs: Map = new Map(); - private modelRequestHeaders: Map> = new Map(); - private configModelOverrides: Map> = new Map(); - private registeredProviders: Map = new Map(); - private loadError: string | undefined = undefined; - readonly authStorage: AuthStorage; - private modelsJsonPath: string | undefined; + private readonly runtime: ModelRuntime; - private constructor(authStorage: AuthStorage, modelsJsonPath: string | undefined) { - this.authStorage = authStorage; - this.modelsJsonPath = modelsJsonPath ? normalizePath(modelsJsonPath) : undefined; - this.loadModels(); + constructor(runtime: ModelRuntime) { + this.runtime = runtime; } - static create(authStorage: AuthStorage, modelsJsonPath: string = join(getAgentDir(), "models.json")): ModelRegistry { - return new ModelRegistry(authStorage, modelsJsonPath); + /** Reload models.json asynchronously. Await before making synchronous registry reads. */ + refresh(): Promise { + return this.runtime.reloadConfig(); } - static inMemory(authStorage: AuthStorage): ModelRegistry { - return new ModelRegistry(authStorage, undefined); - } - - /** - * Reload models from disk (built-in + custom from models.json). - */ - refresh(): void { - this.providerRequestConfigs.clear(); - this.modelRequestHeaders.clear(); - this.loadError = undefined; - - // Ensure dynamic API/OAuth registrations are rebuilt from current provider state. - resetApiProviders(); - resetOAuthProviders(); - - this.loadModels(); - - for (const [providerName, config] of this.registeredProviders.entries()) { - this.applyProviderConfig(providerName, config); - } - } - - /** - * Get any error from loading models.json (undefined if no error). - */ getError(): string | undefined { - return this.loadError; + return this.runtime.getError(); } - private loadModels(): void { - // Load custom models and overrides from models.json - const { - models: customModels, - overrides, - modelOverrides, - error, - } = this.modelsJsonPath ? this.loadCustomModels(this.modelsJsonPath) : emptyCustomModelsResult(); - this.configModelOverrides = modelOverrides; - - if (error) { - this.loadError = error; - // Keep built-in models even if custom models failed to load - } - - const builtInModels = this.loadBuiltInModels(overrides, modelOverrides); - let combined = this.mergeCustomModels(builtInModels, customModels); - - // Let OAuth providers modify their models (e.g., update baseUrl) - for (const oauthProvider of this.authStorage.getOAuthProviders()) { - const cred = this.authStorage.get(oauthProvider.id); - if (cred?.type === "oauth" && oauthProvider.modifyModels) { - combined = oauthProvider.modifyModels(combined, cred); - } - } - - this.models = combined; - } - - /** Load built-in models and apply provider/model overrides */ - private loadBuiltInModels( - overrides: Map, - modelOverrides: Map>, - ): Model[] { - return getProviders().flatMap((provider) => { - const models = getModels(provider as KnownProvider) as Model[]; - const providerOverride = overrides.get(provider); - const perModelOverrides = modelOverrides.get(provider); - - return models.map((m) => { - let model = m; - - // Apply provider-level baseUrl/headers/compat override - if (providerOverride) { - model = { - ...model, - baseUrl: providerOverride.baseUrl ?? model.baseUrl, - compat: mergeCompat(model.compat, providerOverride.compat), - }; - } - - // Apply per-model override - const modelOverride = perModelOverrides?.get(m.id); - if (modelOverride) { - model = applyModelOverride(model, modelOverride); - } - - return model; - }); - }); - } - - private getConfiguredModelOverride(providerName: string, modelId: string): ModelOverride | undefined { - return this.configModelOverrides.get(providerName)?.get(modelId); - } - - private applyConfiguredModelOverride(providerName: string, model: Model): Model { - const modelOverride = this.getConfiguredModelOverride(providerName, model.id); - return modelOverride ? applyModelOverride(model, modelOverride) : model; - } - - /** Merge custom models into built-in list by provider+id (custom wins on conflicts). */ - private mergeCustomModels(builtInModels: Model[], customModels: Model[]): Model[] { - const merged = [...builtInModels]; - for (const customModel of customModels) { - const existingIndex = merged.findIndex((m) => m.provider === customModel.provider && m.id === customModel.id); - if (existingIndex >= 0) { - merged[existingIndex] = customModel; - } else { - merged.push(customModel); - } - } - return merged; - } - - private loadCustomModels(modelsJsonPath: string): CustomModelsResult { - if (!existsSync(modelsJsonPath)) { - return emptyCustomModelsResult(); - } - - try { - const content = readFileSync(modelsJsonPath, "utf-8"); - const parsed = JSON.parse(stripJsonComments(content)) as unknown; - - if (!validateModelsConfig.Check(parsed)) { - const errors = - validateModelsConfig - .Errors(parsed) - .map((error) => ` - ${formatValidationPath(error)}: ${error.message}`) - .join("\n") || "Unknown schema error"; - return emptyCustomModelsResult(`Invalid models.json schema:\n${errors}\n\nFile: ${modelsJsonPath}`); - } - - const config = parsed as ModelsConfig; - - // Additional validation - this.validateConfig(config); - - const overrides = new Map(); - const modelOverrides = new Map>(); - - for (const [providerName, providerConfig] of Object.entries(config.providers)) { - if (providerConfig.baseUrl || providerConfig.compat) { - overrides.set(providerName, { - baseUrl: providerConfig.baseUrl, - compat: providerConfig.compat, - }); - } - - this.storeProviderRequestConfig(providerName, providerConfig); - - if (providerConfig.modelOverrides) { - modelOverrides.set(providerName, new Map(Object.entries(providerConfig.modelOverrides))); - for (const [modelId, modelOverride] of Object.entries(providerConfig.modelOverrides)) { - this.storeModelHeaders(providerName, modelId, modelOverride.headers); - } - } - } - - return { models: this.parseModels(config), overrides, modelOverrides, error: undefined }; - } catch (error) { - if (error instanceof SyntaxError) { - return emptyCustomModelsResult(`Failed to parse models.json: ${error.message}\n\nFile: ${modelsJsonPath}`); - } - return emptyCustomModelsResult( - `Failed to load models.json: ${error instanceof Error ? error.message : error}\n\nFile: ${modelsJsonPath}`, - ); - } - } - - private validateConfig(config: ModelsConfig): void { - const builtInProviders = new Set(getProviders()); - - for (const [providerName, providerConfig] of Object.entries(config.providers)) { - const isBuiltIn = builtInProviders.has(providerName); - const hasProviderApi = !!providerConfig.api; - const models = providerConfig.models ?? []; - const hasModelOverrides = - providerConfig.modelOverrides && Object.keys(providerConfig.modelOverrides).length > 0; - - if (models.length === 0) { - // Override-only config: needs baseUrl, headers, compat, modelOverrides, or some combination. - if (!providerConfig.baseUrl && !providerConfig.headers && !providerConfig.compat && !hasModelOverrides) { - throw new Error( - `Provider ${providerName}: must specify "baseUrl", "headers", "compat", "modelOverrides", or "models".`, - ); - } - } else if (!isBuiltIn) { - // Non-built-in providers with custom models require an endpoint. - // Auth can come from auth.json, --api-key, or provider request config. - if (!providerConfig.baseUrl) { - throw new Error(`Provider ${providerName}: "baseUrl" is required when defining custom models.`); - } - } - // Built-in providers with custom models: baseUrl/apiKey/api are optional, - // inherited from built-in models. Auth comes from env vars / auth storage. - - for (const modelDef of models) { - const hasModelApi = !!modelDef.api; - - if (!hasProviderApi && !hasModelApi && !isBuiltIn) { - throw new Error( - `Provider ${providerName}, model ${modelDef.id}: no "api" specified. Set at provider or model level.`, - ); - } - // For built-in providers, api is optional — inherited from built-in models. - - if (!modelDef.id) throw new Error(`Provider ${providerName}: model missing "id"`); - // Validate contextWindow/maxTokens only if provided (they have defaults) - if (modelDef.contextWindow !== undefined && modelDef.contextWindow <= 0) - throw new Error(`Provider ${providerName}, model ${modelDef.id}: invalid contextWindow`); - if (modelDef.maxTokens !== undefined && modelDef.maxTokens <= 0) - throw new Error(`Provider ${providerName}, model ${modelDef.id}: invalid maxTokens`); - } - } - } - - private parseModels(config: ModelsConfig): Model[] { - const models: Model[] = []; - const builtInProviders = new Set(getProviders()); - - // Cache built-in defaults (api, baseUrl) per provider, extracted from first model. - const builtInDefaultsCache = new Map(); - const getBuiltInDefaults = (providerName: string): { api: string; baseUrl: string } | undefined => { - if (!builtInProviders.has(providerName)) return undefined; - if (builtInDefaultsCache.has(providerName)) return builtInDefaultsCache.get(providerName); - const builtIn = getModels(providerName as KnownProvider) as Model[]; - if (builtIn.length === 0) return undefined; - const defaults = { api: builtIn[0].api, baseUrl: builtIn[0].baseUrl }; - builtInDefaultsCache.set(providerName, defaults); - return defaults; - }; - - for (const [providerName, providerConfig] of Object.entries(config.providers)) { - const modelDefs = providerConfig.models ?? []; - if (modelDefs.length === 0) continue; // Override-only, no custom models - - const builtInDefaults = getBuiltInDefaults(providerName); - - for (const modelDef of modelDefs) { - const api = modelDef.api ?? providerConfig.api ?? builtInDefaults?.api; - if (!api) continue; - - const baseUrl = modelDef.baseUrl ?? providerConfig.baseUrl ?? builtInDefaults?.baseUrl; - if (!baseUrl) continue; - - const compat = mergeCompat(providerConfig.compat, modelDef.compat); - this.storeModelHeaders(providerName, modelDef.id, modelDef.headers); - - const defaultCost = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; - models.push({ - id: modelDef.id, - name: modelDef.name ?? modelDef.id, - api: api as Api, - provider: providerName, - baseUrl, - reasoning: modelDef.reasoning ?? false, - thinkingLevelMap: modelDef.thinkingLevelMap, - input: (modelDef.input ?? ["text"]) as ("text" | "image")[], - cost: modelDef.cost ?? defaultCost, - contextWindow: modelDef.contextWindow ?? 128000, - maxTokens: modelDef.maxTokens ?? 16384, - headers: undefined, - compat, - } as Model); - } - } - - return models; - } - - /** - * Get all models (built-in + custom). - * If models.json had errors, returns only built-in models. - */ getAll(): Model[] { - return this.models; + return [...this.runtime.getModels()]; } - /** - * Get only models that have auth configured. - * This is a fast check that doesn't refresh OAuth tokens. - */ getAvailable(): Model[] { - return this.models.filter((m) => this.hasConfiguredAuth(m)); + return [...this.runtime.getAvailableSnapshot()]; } - /** - * Find a model by provider and ID. - */ find(provider: string, modelId: string): Model | undefined { - return this.models.find((m) => m.provider === provider && m.id === modelId); + return this.runtime.getModel(provider, modelId); } - /** - * Get API key for a model. - */ hasConfiguredAuth(model: Model): boolean { - const providerApiKey = this.providerRequestConfigs.get(model.provider)?.apiKey; - return ( - this.authStorage.hasAuth(model.provider) || - (providerApiKey !== undefined && isConfigValueConfigured(providerApiKey)) - ); + return this.runtime.hasConfiguredAuth(model.provider); } - private getModelRequestKey(provider: string, modelId: string): string { - return `${provider}:${modelId}`; - } - - private storeProviderRequestConfig( - providerName: string, - config: { - apiKey?: string; - headers?: Record; - authHeader?: boolean; - }, - ): void { - if (!config.apiKey && !config.headers && !config.authHeader) { - return; - } - - this.providerRequestConfigs.set(providerName, { - apiKey: config.apiKey, - headers: config.headers, - authHeader: config.authHeader, - }); - } - - private storeModelHeaders(providerName: string, modelId: string, headers?: Record): void { - const key = this.getModelRequestKey(providerName, modelId); - if (!headers || Object.keys(headers).length === 0) { - this.modelRequestHeaders.delete(key); - return; - } - this.modelRequestHeaders.set(key, headers); - } - - /** - * Get API key and request headers for a model. - */ async getApiKeyAndHeaders(model: Model): Promise { try { - const providerConfig = this.providerRequestConfigs.get(model.provider); - const providerEnv = this.authStorage.getProviderEnv(model.provider); - const apiKeyFromAuthStorage = await this.authStorage.getApiKey(model.provider, { includeFallback: false }); - const apiKey = - apiKeyFromAuthStorage ?? - (providerConfig?.apiKey - ? resolveConfigValueOrThrow( - providerConfig.apiKey, - `API key for provider "${model.provider}"`, - providerEnv, - ) - : undefined); - - const providerHeaders = resolveHeadersOrThrow( - providerConfig?.headers, - `provider "${model.provider}"`, - providerEnv, - ); - const modelHeaders = resolveHeadersOrThrow( - this.modelRequestHeaders.get(this.getModelRequestKey(model.provider, model.id)), - `model "${model.provider}/${model.id}"`, - providerEnv, - ); - - let headers = - model.headers || providerHeaders || modelHeaders - ? { ...model.headers, ...providerHeaders, ...modelHeaders } - : undefined; - - if (providerConfig?.authHeader) { - if (!apiKey) { + const resolution = await this.runtime.getAuth(model); + if (!resolution) { + const compatibility = this.runtime.getCompatibilityRequestConfig(model); + if (compatibility.authHeader) { return { ok: false, error: `No API key found for "${model.provider}"` }; } - headers = { ...headers, Authorization: `Bearer ${apiKey}` }; + const headers = compatibility.headers + ? Object.fromEntries( + Object.entries(compatibility.headers).filter( + (entry): entry is [string, string] => entry[1] !== null, + ), + ) + : undefined; + return { ok: true, headers }; } - - return { - ok: true, - apiKey, - headers: headers && Object.keys(headers).length > 0 ? headers : undefined, - env: providerEnv && Object.keys(providerEnv).length > 0 ? providerEnv : undefined, - }; + const headers = resolution.auth.headers + ? Object.fromEntries( + Object.entries(resolution.auth.headers).filter( + (entry): entry is [string, string] => entry[1] !== null, + ), + ) + : undefined; + return { ok: true, apiKey: resolution.auth.apiKey, headers, env: resolution.env }; } catch (error) { + const cause = error instanceof Error ? error.cause : undefined; + const message = + cause instanceof Error ? cause.message : error instanceof Error ? error.message : String(error); return { ok: false, - error: error instanceof Error ? error.message : String(error), + error: + message === "authHeader requires a resolved API key" + ? `No API key found for "${model.provider}"` + : message, }; } } - /** - * Return auth status for a provider, including request auth configured in models.json. - * This intentionally does not execute command-backed config values. - */ getProviderAuthStatus(provider: string): AuthStatus { - const authStatus = this.authStorage.getAuthStatus(provider); - if (authStatus.source) { - return authStatus; - } - - const providerApiKey = this.providerRequestConfigs.get(provider)?.apiKey; - if (!providerApiKey) { - return authStatus; - } - - if (isCommandConfigValue(providerApiKey)) { - return { configured: true, source: "models_json_command" }; - } - - const envVarNames = getConfigValueEnvVarNames(providerApiKey); - if (envVarNames.length > 0) { - return isConfigValueConfigured(providerApiKey) - ? { configured: true, source: "environment", label: envVarNames.join(", ") } - : { configured: false }; - } - - return { configured: true, source: "models_json_key" }; + return this.runtime.getProviderAuthStatus(provider); } - /** - * Get display name for a provider. - */ getProviderDisplayName(provider: string): string { - const registeredProvider = this.registeredProviders.get(provider); - const oauthProvider = this.authStorage.getOAuthProviders().find((p) => p.id === provider); - - return ( - registeredProvider?.name ?? - registeredProvider?.oauth?.name ?? - oauthProvider?.name ?? - BUILT_IN_PROVIDER_DISPLAY_NAMES[provider] ?? - provider - ); + return this.runtime.getProvider(provider)?.name ?? provider; } - /** - * Get API key for a provider. - */ async getApiKeyForProvider(provider: string): Promise { - const apiKey = await this.authStorage.getApiKey(provider); - if (apiKey !== undefined) { - return apiKey; + try { + return (await this.runtime.getAuth(provider))?.auth.apiKey; + } catch { + return undefined; } - - const providerApiKey = this.providerRequestConfigs.get(provider)?.apiKey; - return providerApiKey - ? resolveConfigValueUncached(providerApiKey, this.authStorage.getProviderEnv(provider)) - : undefined; } - /** - * Check if a model is using OAuth credentials (subscription). - */ isUsingOAuth(model: Model): boolean { - const cred = this.authStorage.get(model.provider); - return cred?.type === "oauth"; + return this.runtime.isUsingOAuth(model.provider); } - /** - * Register a provider dynamically (from extensions). - * - * If provider has models: replaces all existing models for this provider. - * If provider has only baseUrl/headers: overrides existing models' URLs. - * If provider has oauth: registers OAuth provider for /login support. - */ registerProvider(providerName: string, config: ProviderConfigInput): void { - this.validateProviderConfig(providerName, config); - this.applyProviderConfig(providerName, config); - this.upsertRegisteredProvider(providerName, config); + this.runtime.registerProvider(providerName, config); } - /** - * Unregister a previously registered provider. - * - * Removes the provider from the registry and reloads models from disk so that - * built-in models overridden by this provider are restored to their original state. - * Also resets dynamic OAuth and API stream registrations before reapplying - * remaining dynamic providers. - * Has no effect if the provider was never registered. - */ unregisterProvider(providerName: string): void { - if (!this.registeredProviders.has(providerName)) return; - this.registeredProviders.delete(providerName); - this.refresh(); + this.runtime.unregisterProvider(providerName); } - /** - * Upsert a provider config into registeredProviders. - * If the provider is already registered, defined values in the incoming config - * override existing ones; undefined values are preserved from the stored config. - * If the provider is not registered, the incoming config is stored as-is. - */ - private upsertRegisteredProvider(providerName: string, config: ProviderConfigInput): void { - const existing = this.registeredProviders.get(providerName); - if (!existing) { - this.registeredProviders.set(providerName, config); - return; - } - for (const k of Object.keys(config) as (keyof ProviderConfigInput)[]) { - if (config[k] !== undefined) { - (existing as Record)[k] = config[k]; - } - } + getRegisteredProviderConfig(providerName: string): ProviderConfigInput | undefined { + return this.runtime.getRegisteredProviderConfig(providerName); } - private validateProviderConfig(providerName: string, config: ProviderConfigInput): void { - if (config.streamSimple && !config.api) { - throw new Error(`Provider ${providerName}: "api" is required when registering streamSimple.`); - } - - if (!config.models || config.models.length === 0) { - return; - } - - if (!config.baseUrl) { - throw new Error(`Provider ${providerName}: "baseUrl" is required when defining models.`); - } - if (!config.apiKey && !config.oauth) { - throw new Error(`Provider ${providerName}: "apiKey" or "oauth" is required when defining models.`); - } - - for (const modelDef of config.models) { - const api = modelDef.api || config.api; - if (!api) { - throw new Error(`Provider ${providerName}, model ${modelDef.id}: no "api" specified.`); - } - } - } - - private applyProviderConfig(providerName: string, config: ProviderConfigInput): void { - // Register OAuth provider if provided - if (config.oauth) { - // Ensure the OAuth provider ID matches the provider name - const oauthProvider: OAuthProviderInterface = { - ...config.oauth, - id: providerName, - }; - registerOAuthProvider(oauthProvider); - } - - if (config.streamSimple) { - const streamSimple = config.streamSimple; - registerApiProvider( - { - api: config.api!, - stream: (model, context, options) => streamSimple(model, context, options as SimpleStreamOptions), - streamSimple, - }, - `provider:${providerName}`, - ); - } - - this.storeProviderRequestConfig(providerName, config); - - if (config.models && config.models.length > 0) { - // Full replacement: remove existing models for this provider - this.models = this.models.filter((m) => m.provider !== providerName); - - // Parse and add new models - for (const modelDef of config.models) { - const api = modelDef.api || config.api; - const modelOverride = this.getConfiguredModelOverride(providerName, modelDef.id); - const headers = - modelDef.headers || modelOverride?.headers - ? { ...modelDef.headers, ...modelOverride?.headers } - : undefined; - this.storeModelHeaders(providerName, modelDef.id, headers); - - const model = this.applyConfiguredModelOverride(providerName, { - id: modelDef.id, - name: modelDef.name, - api: api as Api, - provider: providerName, - baseUrl: modelDef.baseUrl ?? config.baseUrl!, - reasoning: modelDef.reasoning, - thinkingLevelMap: modelDef.thinkingLevelMap, - input: modelDef.input as ("text" | "image")[], - cost: modelDef.cost, - contextWindow: modelDef.contextWindow, - maxTokens: modelDef.maxTokens, - headers: undefined, - compat: modelDef.compat, - } as Model); - this.models.push(model); - } - - // Apply OAuth modifyModels if credentials exist (e.g., to update baseUrl) - if (config.oauth?.modifyModels) { - const cred = this.authStorage.get(providerName); - if (cred?.type === "oauth") { - this.models = config.oauth.modifyModels(this.models, cred); - } - } - } else if (config.baseUrl || config.headers) { - // Override-only: update baseUrl for existing models. Request headers are resolved per request. - this.models = this.models.map((m) => { - if (m.provider !== providerName) return m; - return { - ...m, - baseUrl: config.baseUrl ?? m.baseUrl, - }; - }); - } + getRegisteredProviderIds(): readonly string[] { + return this.runtime.getRegisteredProviderIds(); } } - -/** - * Input type for registerProvider API. - */ -export interface ProviderConfigInput { - name?: string; - baseUrl?: string; - apiKey?: string; - api?: Api; - streamSimple?: (model: Model, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream; - headers?: Record; - authHeader?: boolean; - /** OAuth provider for /login support */ - oauth?: Omit; - models?: Array<{ - id: string; - name: string; - api?: Api; - baseUrl?: string; - reasoning: boolean; - thinkingLevelMap?: Model["thinkingLevelMap"]; - input: ("text" | "image")[]; - cost: Model["cost"]; - contextWindow: number; - maxTokens: number; - headers?: Record; - compat?: Model["compat"]; - }>; -} diff --git a/packages/coding-agent/src/core/model-resolver.ts b/packages/coding-agent/src/core/model-resolver.ts index 3a3341df..23c5801e 100644 --- a/packages/coding-agent/src/core/model-resolver.ts +++ b/packages/coding-agent/src/core/model-resolver.ts @@ -8,7 +8,7 @@ import chalk from "chalk"; import { minimatch } from "minimatch"; import { isValidThinkingLevel } from "../cli/args.ts"; import { DEFAULT_THINKING_LEVEL } from "./defaults.ts"; -import type { ModelRegistry } from "./model-registry.ts"; +import type { ModelRuntime } from "./model-runtime.ts"; /** Default model IDs for each known provider */ export const defaultModelPerProvider: Record = { @@ -268,9 +268,9 @@ export interface ResolveModelScopeResult { export async function resolveModelScopeWithDiagnostics( patterns: string[], - modelRegistry: ModelRegistry, + modelRuntime: ModelRuntime, ): Promise { - const availableModels = await modelRegistry.getAvailable(); + const availableModels = [...(await modelRuntime.getAvailable())]; const scopedModels: ScopedModel[] = []; const diagnostics: ModelScopeDiagnostic[] = []; @@ -330,8 +330,8 @@ export async function resolveModelScopeWithDiagnostics( return { scopedModels, diagnostics }; } -export async function resolveModelScope(patterns: string[], modelRegistry: ModelRegistry): Promise { - const { scopedModels, diagnostics } = await resolveModelScopeWithDiagnostics(patterns, modelRegistry); +export async function resolveModelScope(patterns: string[], modelRuntime: ModelRuntime): Promise { + const { scopedModels, diagnostics } = await resolveModelScopeWithDiagnostics(patterns, modelRuntime); for (const diagnostic of diagnostics) { console.warn(chalk.yellow(`Warning: ${diagnostic.message}`)); } @@ -364,9 +364,9 @@ export function resolveCliModel(options: { cliProvider?: string; cliModel?: string; cliThinking?: ThinkingLevel; - modelRegistry: ModelRegistry; + modelRuntime: ModelRuntime; }): ResolveCliModelResult { - const { cliProvider, cliModel, cliThinking, modelRegistry } = options; + const { cliProvider, cliModel, cliThinking, modelRuntime } = options; if (!cliModel) { return { model: undefined, warning: undefined, error: undefined }; @@ -374,7 +374,7 @@ export function resolveCliModel(options: { // Important: use *all* models here, not just models with pre-configured auth. // This allows "--api-key" to be used for first-time setup. - const availableModels = modelRegistry.getAll(); + const availableModels = [...modelRuntime.getModels()]; if (availableModels.length === 0) { return { model: undefined, @@ -454,8 +454,8 @@ export function resolveCliModel(options: { const rawExactMatches = availableModels.filter( (m) => m.id.toLowerCase() === cliModel.toLowerCase() && !modelsAreEqual(m, model), ); - if (rawExactMatches.length > 0 && !modelRegistry.hasConfiguredAuth(model)) { - const authenticatedRawMatches = rawExactMatches.filter((m) => modelRegistry.hasConfiguredAuth(m)); + if (rawExactMatches.length > 0 && !modelRuntime.hasConfiguredAuth(model.provider)) { + const authenticatedRawMatches = rawExactMatches.filter((m) => modelRuntime.hasConfiguredAuth(m.provider)); if (authenticatedRawMatches.length === 1) { return { model: authenticatedRawMatches[0], @@ -555,7 +555,7 @@ export async function findInitialModel(options: { defaultProvider?: string; defaultModelId?: string; defaultThinkingLevel?: ThinkingLevel; - modelRegistry: ModelRegistry; + modelRuntime: ModelRuntime; }): Promise { const { cliProvider, @@ -565,7 +565,7 @@ export async function findInitialModel(options: { defaultProvider, defaultModelId, defaultThinkingLevel, - modelRegistry, + modelRuntime, } = options; let model: Model | undefined; @@ -576,7 +576,7 @@ export async function findInitialModel(options: { const resolved = resolveCliModel({ cliProvider, cliModel, - modelRegistry, + modelRuntime, }); if (resolved.error) { console.error(chalk.red(resolved.error)); @@ -598,8 +598,8 @@ export async function findInitialModel(options: { // 3. Try saved default from settings if auth is configured. if (defaultProvider && defaultModelId) { - const found = modelRegistry.find(defaultProvider, defaultModelId); - if (found && modelRegistry.hasConfiguredAuth(found)) { + const found = modelRuntime.getModel(defaultProvider, defaultModelId); + if (found && modelRuntime.hasConfiguredAuth(found.provider)) { model = found; if (defaultThinkingLevel) { thinkingLevel = defaultThinkingLevel; @@ -609,7 +609,7 @@ export async function findInitialModel(options: { } // 4. Try first available model with valid API key - const availableModels = await modelRegistry.getAvailable(); + const availableModels = [...(await modelRuntime.getAvailable())]; if (availableModels.length > 0) { // Try to find a default model from known providers @@ -637,12 +637,12 @@ export async function restoreModelFromSession( savedModelId: string, currentModel: Model | undefined, shouldPrintMessages: boolean, - modelRegistry: ModelRegistry, + modelRuntime: ModelRuntime, ): Promise<{ model: Model | undefined; fallbackMessage: string | undefined }> { - const restoredModel = modelRegistry.find(savedProvider, savedModelId); + const restoredModel = modelRuntime.getModel(savedProvider, savedModelId); // Check if restored model exists and still has auth configured - const hasConfiguredAuth = restoredModel ? modelRegistry.hasConfiguredAuth(restoredModel) : false; + const hasConfiguredAuth = restoredModel ? modelRuntime.hasConfiguredAuth(restoredModel.provider) : false; if (restoredModel && hasConfiguredAuth) { if (shouldPrintMessages) { @@ -670,7 +670,7 @@ export async function restoreModelFromSession( } // Try to find any available model - const availableModels = await modelRegistry.getAvailable(); + const availableModels = [...(await modelRuntime.getAvailable())]; 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 new file mode 100644 index 00000000..b8c24f00 --- /dev/null +++ b/packages/coding-agent/src/core/model-runtime.ts @@ -0,0 +1,489 @@ +import { join } from "node:path"; +import { + type Api, + type ApiStreamOptions, + type AssistantMessage, + type AssistantMessageEventStream, + type AuthCheck, + type AuthInteraction, + type AuthResult, + type AuthType, + type Context, + type Credential, + type CredentialInfo, + type CredentialStore, + createModels, + lazyStream, + type Model, + type Models, + type ModelsApiStreamOptions, + ModelsError, + type ModelsSimpleStreamOptions, + type ModelsStreamTransforms, + type MutableModels, + type Provider, + type ProviderHeaders, + type SimpleStreamOptions, + type StreamOptions, +} from "@earendil-works/pi-ai"; +import { builtinProviders } from "@earendil-works/pi-ai/providers/all"; +import { getAgentDir } from "../config.ts"; +import { AuthStorage as DefaultAuthStorage } from "./auth-storage.ts"; +import { ModelConfig } from "./model-config.ts"; +import { + type AuthStatus, + type CompatibilityRequestConfig, + composeModelProvider, + configuredRequestAuthStatus, + type ProviderConfigInput, + resolveCompatibilityRequestConfig, + resolveConfiguredModelHeaders, + validateExtensionProvider, +} from "./provider-composer.ts"; +import { RuntimeCredentials } from "./runtime-credentials.ts"; + +interface ModelRuntimeSnapshot { + all: readonly Model[]; + available: readonly Model[]; + configuredProviders: ReadonlySet; + storedProviders: ReadonlySet; + auth: ReadonlyMap; +} + +export interface CreateModelRuntimeOptions { + /** Credential storage. Defaults to the file at authPath. */ + credentials?: CredentialStore; + authPath?: string; + modelsPath?: string | null; +} + +export interface ModelRuntimeAuthOverrides { + apiKey?: string; + env?: Record; +} + +function mergeHeaders( + base: ProviderHeaders | undefined, + override: ProviderHeaders | undefined, +): ProviderHeaders | undefined { + if (!base && !override) return undefined; + const merged = { ...base }; + for (const [name, value] of Object.entries(override ?? {})) { + const lowerName = name.toLowerCase(); + for (const existingName of Object.keys(merged)) { + if (existingName.toLowerCase() === lowerName) delete merged[existingName]; + } + merged[name] = value; + } + return merged; +} + +/** Configured pi-ai Models collection used by coding-agent and SDK consumers. */ +export class ModelRuntime implements Models { + private readonly models: MutableModels; + private readonly credentials: RuntimeCredentials; + private readonly builtins: ReadonlyMap; + private readonly extensionProviders = new Map(); + private readonly compositionErrors = new Map(); + private readonly modelsPath: string | undefined; + private config: ModelConfig; + private snapshot: ModelRuntimeSnapshot = { + all: [], + available: [], + configuredProviders: new Set(), + storedProviders: new Set(), + auth: new Map(), + }; + private availabilityRefresh: Promise | undefined; + private availabilityError: string | undefined; + + private constructor( + credentials: RuntimeCredentials, + config: ModelConfig, + modelsPath: string | undefined, + providers: readonly Provider[], + ) { + this.credentials = credentials; + this.config = config; + this.modelsPath = modelsPath; + this.builtins = new Map(providers.map((provider) => [provider.id, provider])); + this.models = createModels({ credentials }); + this.rebuildProviders(); + } + + static async create(options: CreateModelRuntimeOptions = {}): Promise { + const credentials = new RuntimeCredentials(options.credentials ?? DefaultAuthStorage.create(options.authPath)); + const modelsPath = + options.modelsPath === null ? undefined : (options.modelsPath ?? join(getAgentDir(), "models.json")); + const config = await ModelConfig.load(modelsPath); + const runtime = new ModelRuntime(credentials, config, modelsPath, builtinProviders()); + await runtime.refreshAvailability(); + return runtime; + } + + private providerIds(): Set { + return new Set([...this.builtins.keys(), ...this.config.getProviderIds(), ...this.extensionProviders.keys()]); + } + + private recomposeProvider(providerId: string): void { + const base = this.builtins.get(providerId); + const extension = this.extensionProviders.get(providerId); + if (!base && !this.config.getProvider(providerId) && !extension) { + this.models.deleteProvider(providerId); + this.compositionErrors.delete(providerId); + return; + } + if (base && !this.config.getProvider(providerId) && !extension) { + // No overlays: use the builtin untouched so its auth/login/stream behavior is exact. + this.models.setProvider(base); + this.compositionErrors.delete(providerId); + return; + } + try { + this.models.setProvider(composeModelProvider(providerId, base, this.config, extension)); + this.compositionErrors.delete(providerId); + } catch (error) { + this.compositionErrors.set(providerId, error instanceof Error ? error.message : String(error)); + if (base) this.models.setProvider(base); + else this.models.deleteProvider(providerId); + } + } + + private rebuildProviders(): void { + this.models.clearProviders(); + this.compositionErrors.clear(); + for (const providerId of this.providerIds()) this.recomposeProvider(providerId); + this.updateModelSnapshot(); + } + + private updateModelSnapshot(): void { + const all = [...this.models.getModels()]; + this.snapshot = { + ...this.snapshot, + all, + available: all.filter((model) => this.snapshot.configuredProviders.has(model.provider)), + }; + } + + private async runAvailabilityRefresh(): Promise { + const providers = this.models.getProviders(); + const [available, checks, credentials] = await Promise.all([ + this.models.getAvailable(), + Promise.all( + providers.map( + async (provider): Promise<[string, AuthCheck | undefined]> => [ + provider.id, + await this.models.checkAuth(provider.id), + ], + ), + ), + this.credentials.list(), + ]); + const auth = new Map(checks); + const configuredProviders = new Set( + checks + .filter((entry): entry is [string, AuthCheck] => entry[1] !== undefined) + .map(([providerId]) => providerId), + ); + this.snapshot = { + all: [...this.models.getModels()], + available: [...available], + configuredProviders, + storedProviders: new Set(credentials.map((entry) => entry.providerId)), + auth, + }; + this.availabilityError = undefined; + } + + private queueAvailabilityRefresh(after: Promise | undefined): Promise { + const refresh = (after ?? Promise.resolve()).catch(() => {}).then(() => this.runAvailabilityRefresh()); + const recorded = refresh.catch((error) => { + this.availabilityError = error instanceof Error ? error.message : String(error); + throw error; + }); + const tracked = recorded.finally(() => { + if (this.availabilityRefresh === tracked) this.availabilityRefresh = undefined; + }); + this.availabilityRefresh = tracked; + return tracked; + } + + /** Coalesce concurrent readers onto the pending refresh. */ + private refreshAvailability(): Promise { + return this.availabilityRefresh ?? this.queueAvailabilityRefresh(undefined); + } + + /** Mutations must not observe an in-flight refresh started before them. */ + private forceRefreshAvailability(): Promise { + return this.queueAvailabilityRefresh(this.availabilityRefresh); + } + + getProviders(): readonly Provider[] { + return this.models.getProviders(); + } + + getProvider(providerId: string): Provider | undefined { + return this.models.getProvider(providerId); + } + + getModels(providerId?: string): readonly Model[] { + return this.models.getModels(providerId); + } + + getModel(providerId: string, modelId: string): Model | undefined { + return this.models.getModel(providerId, modelId); + } + + async checkAuth(providerId: string): Promise { + return this.models.checkAuth(providerId); + } + + async getAvailable(providerId?: string): Promise[]> { + if (providerId) { + if (this.availabilityRefresh) { + await this.availabilityRefresh; + return this.snapshot.available.filter((model) => model.provider === providerId); + } + try { + return await this.models.getAvailable(providerId); + } catch (error) { + this.availabilityError = error instanceof Error ? error.message : String(error); + throw error; + } + } + await this.refreshAvailability(); + return this.snapshot.available; + } + + getAvailableSnapshot(): readonly Model[] { + return this.snapshot.available; + } + + getError(): string | undefined { + const errors: string[] = []; + const configError = this.config.getError(); + if (configError) errors.push(configError); + for (const [providerId, error] of this.compositionErrors) { + errors.push(`Provider "${providerId}": ${error}`); + } + if (this.availabilityError) errors.push(`Availability refresh: ${this.availabilityError}`); + return errors.length > 0 ? errors.join("\n\n") : undefined; + } + + getRegisteredProviderConfig(providerId: string): ProviderConfigInput | undefined { + return this.extensionProviders.get(providerId); + } + + getRegisteredProviderIds(): readonly string[] { + return [...this.extensionProviders.keys()]; + } + + /** @internal Compatibility fallback for ModelRegistry when provider auth is unconfigured. */ + getCompatibilityRequestConfig(model: Model): CompatibilityRequestConfig { + return resolveCompatibilityRequestConfig( + model, + this.config.getProvider(model.provider), + this.extensionProviders.get(model.provider), + ); + } + + isUsingOAuth(providerId: string): boolean { + return this.snapshot.auth.get(providerId)?.type === "oauth"; + } + + hasConfiguredAuth(providerId: string): boolean { + return this.snapshot.configuredProviders.has(providerId); + } + + getAuth(providerId: string, overrides?: ModelRuntimeAuthOverrides): Promise; + getAuth(model: Model, overrides?: ModelRuntimeAuthOverrides): Promise; + async getAuth( + providerOrModel: string | Model, + overrides: ModelRuntimeAuthOverrides = {}, + ): Promise { + if (typeof providerOrModel === "string") return this.models.getAuth(providerOrModel, overrides); + const resolution = await this.models.getAuth(providerOrModel, overrides); + if (!resolution) return undefined; + const configuredHeaders = resolveConfiguredModelHeaders( + providerOrModel, + this.config.getProvider(providerOrModel.provider), + this.extensionProviders.get(providerOrModel.provider), + { ...(resolution.env ?? {}), ...(overrides.env ?? {}) }, + ); + return { + ...resolution, + auth: { + ...resolution.auth, + headers: mergeHeaders(resolution.auth.headers, configuredHeaders), + }, + }; + } + + setRuntimeApiKey(providerId: string, apiKey: string): void { + this.credentials.setRuntimeApiKey(providerId, apiKey); + const auth = new Map(this.snapshot.auth).set(providerId, { type: "api_key", source: "runtime API key" }); + const configuredProviders = new Set(this.snapshot.configuredProviders).add(providerId); + const storedProviders = new Set(this.snapshot.storedProviders).add(providerId); + this.snapshot = { + ...this.snapshot, + auth, + configuredProviders, + storedProviders, + available: this.snapshot.all.filter((model) => configuredProviders.has(model.provider)), + }; + void this.forceRefreshAvailability().catch(() => {}); + } + + removeRuntimeApiKey(providerId: string): void { + this.credentials.removeRuntimeApiKey(providerId); + void this.forceRefreshAvailability().catch(() => {}); + } + + listCredentials(): Promise { + return this.credentials.list(); + } + + getProviderAuthStatus(providerId: string): AuthStatus { + if (this.credentials.hasRuntimeApiKey(providerId)) return { configured: true, source: "runtime" }; + if (this.snapshot.storedProviders.has(providerId)) return { configured: true, source: "stored" }; + const configured = configuredRequestAuthStatus( + this.config.getProvider(providerId), + this.extensionProviders.get(providerId), + ); + if (configured) return configured; + const check = this.snapshot.auth.get(providerId); + return check ? { configured: true, source: "environment", label: check.source } : { configured: false }; + } + + private async prepareRequest( + model: Model, + options: (StreamOptions & ModelsStreamTransforms) | undefined, + ): 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 }); + if (!resolution) throw new ModelsError("auth", `Provider is not configured: ${model.provider}`); + + const { transformHeaders, ...providerOptions } = options ?? {}; + let headers = mergeHeaders(resolution.auth.headers, providerOptions.headers); + if (transformHeaders) headers = await transformHeaders(headers ?? {}); + const env = + resolution.env || providerOptions.env + ? { ...(resolution.env ?? {}), ...(providerOptions.env ?? {}) } + : undefined; + return { + provider, + model: resolution.auth.baseUrl ? { ...model, baseUrl: resolution.auth.baseUrl } : model, + options: { + ...providerOptions, + apiKey: providerOptions.apiKey ?? resolution.auth.apiKey, + headers, + env, + }, + }; + } + + stream( + model: Model, + context: Context, + options?: ModelsApiStreamOptions, + ): AssistantMessageEventStream { + return lazyStream(model, async () => { + const prepared = await this.prepareRequest( + model, + options as (StreamOptions & ModelsStreamTransforms) | undefined, + ); + return prepared.provider.stream( + prepared.model as Model, + context, + prepared.options as ApiStreamOptions, + ); + }); + } + + complete( + model: Model, + context: Context, + options?: ModelsApiStreamOptions, + ): Promise { + return this.stream(model, context, options).result(); + } + + streamSimple(model: Model, context: Context, options?: ModelsSimpleStreamOptions): AssistantMessageEventStream { + return lazyStream(model, async () => { + const prepared = await this.prepareRequest(model, options); + return prepared.provider.streamSimple(prepared.model, context, prepared.options as SimpleStreamOptions); + }); + } + + completeSimple(model: Model, context: Context, options?: ModelsSimpleStreamOptions): Promise { + 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); + await this.forceRefreshAvailability(); + return credential; + } + + async logout(providerId: string): Promise { + await this.models.logout(providerId); + await this.forceRefreshAvailability(); + } + + async reloadConfig(): Promise { + this.config = await ModelConfig.load(this.modelsPath); + this.rebuildProviders(); + await this.forceRefreshAvailability(); + } + + async refresh(providerId?: string): Promise { + await this.models.refresh(providerId); + this.updateModelSnapshot(); + await this.forceRefreshAvailability(); + } + + registerProvider(providerId: string, config: ProviderConfigInput): void { + // Validate the incoming registration on its own, like the legacy registry: + // a broken re-registration must throw without touching the stored config. + validateExtensionProvider(providerId, this.builtins.get(providerId), this.config.getProvider(providerId), config); + // Re-registration merges defined values over the previous registration and + // preserves undefined ones, matching the legacy ModelRegistry contract. + const previous = this.extensionProviders.get(providerId); + const effective: ProviderConfigInput = { ...previous }; + for (const [key, value] of Object.entries(config)) { + if (value !== undefined) (effective as Record)[key] = value; + } + this.extensionProviders.set(providerId, effective); + this.recomposeProvider(providerId); + this.updateModelSnapshot(); + if ( + this.snapshot.storedProviders.has(providerId) || + configuredRequestAuthStatus(this.config.getProvider(providerId), effective)?.configured + ) { + const configuredProviders = new Set(this.snapshot.configuredProviders).add(providerId); + const auth = new Map(this.snapshot.auth); + // Provisional entry until the async refresh lands; never clobber a real check result. + if (!auth.get(providerId)) { + auth.set(providerId, { + type: effective.oauth && !effective.apiKey ? "oauth" : "api_key", + source: "configured provider", + }); + } + this.snapshot = { + ...this.snapshot, + auth, + configuredProviders, + available: this.snapshot.all.filter((model) => configuredProviders.has(model.provider)), + }; + } + void this.forceRefreshAvailability().catch(() => {}); + } + + unregisterProvider(providerId: string): void { + this.extensionProviders.delete(providerId); + this.recomposeProvider(providerId); + this.updateModelSnapshot(); + void this.forceRefreshAvailability().catch(() => {}); + } +} diff --git a/packages/coding-agent/src/core/provider-composer.ts b/packages/coding-agent/src/core/provider-composer.ts new file mode 100644 index 00000000..d1b21c1c --- /dev/null +++ b/packages/coding-agent/src/core/provider-composer.ts @@ -0,0 +1,513 @@ +import { + type Api, + type ApiKeyAuth, + type AssistantMessageEventStream, + type AuthContext, + type AuthInteraction, + type AuthResult, + type Context, + type Credential, + lazyStream, + type Model, + type ModelAuth, + type OAuthAuth, + type OAuthCredentials, + type OAuthLoginCallbacks, + type Provider, + type ProviderHeaders, + type SimpleStreamOptions, + type StreamOptions, +} from "@earendil-works/pi-ai"; +import { getApiProvider } from "@earendil-works/pi-ai/compat"; +import type { ModelConfig, ModelsJsonModel, ModelsJsonModelOverride, ModelsJsonProvider } from "./model-config.ts"; +import { + clearConfigValueCache, + getConfigValueEnvVarNames, + isCommandConfigValue, + isConfigValueConfigured, + resolveConfigValueOrThrow, + resolveHeadersOrThrow, +} from "./resolve-config-value.ts"; + +export interface ExtensionOAuthConfig { + name: string; + /** @deprecated Retained for extension source compatibility; ignored by canonical auth flows. */ + usesCallbackServer?: boolean; + login(callbacks: OAuthLoginCallbacks): Promise; + refreshToken(credentials: OAuthCredentials): Promise; + getApiKey(credentials: OAuthCredentials): string; +} + +/** Input type for the extension registerProvider API. */ +export interface ProviderConfigInput { + name?: string; + baseUrl?: string; + apiKey?: string; + api?: Api; + streamSimple?: (model: Model, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream; + headers?: Record; + authHeader?: boolean; + oauth?: ExtensionOAuthConfig; + models?: Array<{ + id: string; + name: string; + api?: Api; + baseUrl?: string; + reasoning: boolean; + thinkingLevelMap?: Model["thinkingLevelMap"]; + input: ("text" | "image")[]; + cost: Model["cost"]; + contextWindow: number; + maxTokens: number; + headers?: Record; + compat?: Model["compat"]; + }>; +} + +export type AuthStatus = { + configured: boolean; + source?: "stored" | "runtime" | "environment" | "fallback" | "models_json_key" | "models_json_command"; + label?: string; +}; + +export const clearApiKeyCache = clearConfigValueCache; + +function mergeCompat( + base: Model["compat"], + override: Model["compat"] | ModelsJsonModelOverride["compat"], +): Model["compat"] { + if (!override) return base; + const merged = { ...base, ...override } as NonNullable["compat"]>; + const baseNested = base as Record | undefined; + const overrideNested = override as Record; + const mergedNested = merged as Record; + for (const key of ["openRouterRouting", "vercelGatewayRouting", "chatTemplateKwargs"] as const) { + const baseValue = baseNested?.[key]; + const overrideValue = overrideNested[key]; + if ( + (typeof baseValue === "object" && baseValue !== null) || + (typeof overrideValue === "object" && overrideValue !== null) + ) { + mergedNested[key] = { ...(baseValue as object | undefined), ...(overrideValue as object | undefined) }; + } + } + return merged; +} + +function applyModelOverride(model: Model, override: ModelsJsonModelOverride): Model { + return { + ...model, + name: override.name ?? model.name, + reasoning: override.reasoning ?? model.reasoning, + thinkingLevelMap: override.thinkingLevelMap + ? { ...model.thinkingLevelMap, ...override.thinkingLevelMap } + : model.thinkingLevelMap, + input: (override.input as ("text" | "image")[] | undefined) ?? model.input, + cost: override.cost + ? { + input: override.cost.input ?? model.cost.input, + output: override.cost.output ?? model.cost.output, + cacheRead: override.cost.cacheRead ?? model.cost.cacheRead, + cacheWrite: override.cost.cacheWrite ?? model.cost.cacheWrite, + tiers: override.cost.tiers ?? model.cost.tiers, + } + : model.cost, + contextWindow: override.contextWindow ?? model.contextWindow, + maxTokens: override.maxTokens ?? model.maxTokens, + compat: mergeCompat(model.compat, override.compat), + }; +} + +function modelFromJson( + providerId: string, + definition: ModelsJsonModel, + providerConfig: ModelsJsonProvider, + defaults: Model | undefined, +): Model { + const api = definition.api ?? providerConfig.api ?? defaults?.api; + if (!api) { + throw new Error( + `Provider ${providerId}, model ${definition.id}: no "api" specified. Set at provider or model level.`, + ); + } + const baseUrl = definition.baseUrl ?? providerConfig.baseUrl ?? defaults?.baseUrl; + if (!baseUrl) throw new Error(`Provider ${providerId}: "baseUrl" is required when defining custom models.`); + if (definition.contextWindow !== undefined && definition.contextWindow <= 0) { + throw new Error(`Provider ${providerId}, model ${definition.id}: invalid contextWindow`); + } + if (definition.maxTokens !== undefined && definition.maxTokens <= 0) { + throw new Error(`Provider ${providerId}, model ${definition.id}: invalid maxTokens`); + } + return { + id: definition.id, + name: definition.name ?? definition.id, + api: api as Api, + provider: providerId, + baseUrl, + reasoning: definition.reasoning ?? false, + thinkingLevelMap: definition.thinkingLevelMap, + input: (definition.input ?? ["text"]) as ("text" | "image")[], + cost: definition.cost ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: definition.contextWindow ?? 128000, + maxTokens: definition.maxTokens ?? 16384, + headers: undefined, + compat: mergeCompat(providerConfig.compat, definition.compat), + }; +} + +function applyModelsJson( + providerId: string, + baseModels: readonly Model[], + config: ModelsJsonProvider | undefined, +): Model[] { + if (!config) return [...baseModels]; + const hasOverrides = config.modelOverrides && Object.keys(config.modelOverrides).length > 0; + if ( + !config.models?.length && + !config.baseUrl && + !config.headers && + !config.compat && + !hasOverrides && + !config.apiKey && + config.authHeader === undefined + ) { + throw new Error( + `Provider ${providerId}: must specify "baseUrl", "headers", "compat", "modelOverrides", or "models".`, + ); + } + + const models: Model[] = baseModels.map((model) => ({ + ...model, + baseUrl: config.baseUrl ?? model.baseUrl, + compat: mergeCompat(model.compat, config.compat), + })); + for (const definition of config.models ?? []) { + const existingIndex = models.findIndex((model) => model.id === definition.id); + const defaults = existingIndex >= 0 ? models[existingIndex] : models[0]; + const model = modelFromJson(providerId, definition, config, defaults); + if (existingIndex >= 0) models[existingIndex] = model; + else models.push(model); + } + return models; +} + +function applyExtension( + providerId: string, + models: readonly Model[], + config: ProviderConfigInput | undefined, +): Model[] { + if (!config) return [...models]; + if (!config.models) { + return config.baseUrl ? models.map((model) => ({ ...model, baseUrl: config.baseUrl! })) : [...models]; + } + return config.models.map((definition) => { + const defaults = models.find((model) => model.id === definition.id) ?? models[0]; + const api = definition.api ?? config.api ?? defaults?.api; + if (!api) { + throw new Error( + `Provider ${providerId}, model ${definition.id}: no "api" specified. Set at provider or model level.`, + ); + } + const baseUrl = definition.baseUrl ?? config.baseUrl ?? defaults?.baseUrl; + if (!baseUrl) throw new Error(`Provider ${providerId}: "baseUrl" is required when defining custom models.`); + return { + ...definition, + api, + provider: providerId, + baseUrl, + headers: undefined, + }; + }); +} + +function adaptOAuth(config: ExtensionOAuthConfig): OAuthAuth { + return { + name: config.name, + login: async (callbacks) => { + const credential = await config.login({ + onAuth: (info) => callbacks.notify({ type: "auth_url", ...info }), + onDeviceCode: (info) => callbacks.notify({ type: "device_code", ...info }), + onPrompt: (prompt) => callbacks.prompt({ type: "text", ...prompt }), + onProgress: (message) => callbacks.notify({ type: "progress", message }), + onManualCodeInput: () => callbacks.prompt({ type: "manual_code", message: "Paste the authorization code" }), + onSelect: (prompt) => callbacks.prompt({ type: "select", ...prompt }), + signal: callbacks.signal, + }); + return { ...credential, type: "oauth" }; + }, + refresh: async (credential) => ({ ...(await config.refreshToken(credential)), type: "oauth" }), + toAuth: async (credential) => ({ apiKey: config.getApiKey(credential) }), + }; +} + +function withConfiguredAuth( + auth: ModelAuth, + headers: Record | undefined, + authHeader: boolean, +): ModelAuth { + let mergedHeaders: ProviderHeaders | undefined = + auth.headers || headers ? { ...auth.headers, ...headers } : undefined; + if (authHeader) { + if (!auth.apiKey) throw new Error("authHeader requires a resolved API key"); + mergedHeaders = { ...mergedHeaders, Authorization: `Bearer ${auth.apiKey}` }; + } + return { ...auth, headers: mergedHeaders }; +} + +function configuredApiKey( + config: ModelsJsonProvider | undefined, + extension: ProviderConfigInput | undefined, +): string | undefined { + return extension?.apiKey ?? config?.apiKey; +} + +function configuredHeaders( + config: ModelsJsonProvider | undefined, + extension: ProviderConfigInput | undefined, +): Record | undefined { + if (!config?.headers && !extension?.headers) return undefined; + return { ...config?.headers, ...extension?.headers }; +} + +async function configContextEnv( + values: readonly string[], + ctx: AuthContext, + explicit?: Record, +): Promise | undefined> { + const env = { ...explicit }; + for (const name of new Set(values.flatMap(getConfigValueEnvVarNames))) { + if (env[name] !== undefined) continue; + const value = await ctx.env(name); + if (value !== undefined) env[name] = value; + } + return Object.keys(env).length > 0 ? env : undefined; +} + +function composeApiKeyAuth( + providerId: string, + base: Provider | undefined, + config: ModelsJsonProvider | undefined, + extension: ProviderConfigInput | undefined, +): ApiKeyAuth | undefined { + const inherited = base?.auth.apiKey; + const rawKey = configuredApiKey(config, extension); + const oauth = extension?.oauth ?? base?.auth.oauth; + // OAuth-only providers get no fabricated API-key login method. + if (!inherited && rawKey === undefined && oauth) return undefined; + const rawHeaders = configuredHeaders(config, extension); + const authHeader = extension?.authHeader ?? config?.authHeader ?? false; + return { + name: inherited?.name ?? "API key", + login: + inherited?.login ?? + (async (interaction: AuthInteraction) => ({ + type: "api_key", + key: await interaction.prompt({ type: "secret", message: "Enter API key" }), + })), + check: async (input) => { + if (input.credential) { + if (inherited?.check) return inherited.check(input); + if (input.credential.key) return { type: "api_key", source: "stored credential" }; + const resolved = await inherited?.resolve(input); + return resolved ? { type: "api_key", source: resolved.source } : undefined; + } + if (rawKey !== undefined) { + if (isCommandConfigValue(rawKey)) return { type: "api_key", source: "configured API key" }; + const envNames = getConfigValueEnvVarNames(rawKey); + for (const name of envNames) { + if ((await input.ctx.env(name)) === undefined) return undefined; + } + return { type: "api_key", source: "configured API key" }; + } + if (inherited?.check) return inherited.check(input); + const resolved = await inherited?.resolve(input); + return resolved ? { type: "api_key", source: resolved.source } : undefined; + }, + resolve: async (input) => { + let result: AuthResult | undefined; + if (input.credential) { + result = inherited + ? await inherited.resolve(input) + : input.credential.key + ? { auth: { apiKey: input.credential.key }, env: input.credential.env, source: "stored credential" } + : undefined; + } else if (rawKey !== undefined) { + const env = await configContextEnv([rawKey], input.ctx); + const key = resolveConfigValueOrThrow(rawKey, `API key for provider "${providerId}"`, env); + result = inherited + ? await inherited.resolve({ ...input, credential: { type: "api_key", key } }) + : { auth: { apiKey: key }, source: "configured API key" }; + } else { + result = await inherited?.resolve(input); + } + if (!result) return undefined; + const explicitEnv = { ...(input.credential?.env ?? {}), ...(result.env ?? {}) }; + const headerEnv = await configContextEnv(Object.values(rawHeaders ?? {}), input.ctx, explicitEnv); + const headers = resolveHeadersOrThrow(rawHeaders, `provider "${providerId}"`, headerEnv); + return { ...result, auth: withConfiguredAuth(result.auth, headers, authHeader) }; + }, + }; +} + +function composeOAuthAuth( + providerId: string, + base: Provider | undefined, + config: ModelsJsonProvider | undefined, + extension: ProviderConfigInput | undefined, +): OAuthAuth | undefined { + const oauth = extension?.oauth ? adaptOAuth(extension.oauth) : base?.auth.oauth; + if (!oauth) return undefined; + const rawHeaders = configuredHeaders(config, extension); + const authHeader = extension?.authHeader ?? config?.authHeader ?? false; + return { + ...oauth, + toAuth: async (credential) => { + const auth = await oauth.toAuth(credential); + const env = credential.env; + const headers = resolveHeadersOrThrow( + rawHeaders, + `provider "${providerId}"`, + typeof env === "object" && env !== null ? (env as Record) : undefined, + ); + return withConfiguredAuth(auth, headers, authHeader); + }, + }; +} + +function rawModelHeaders( + model: Model, + config: ModelsJsonProvider | undefined, + extension: ProviderConfigInput | undefined, +): Record | undefined { + const definition = config?.models?.find((entry) => entry.id === model.id); + const extensionModel = extension?.models?.find((entry) => entry.id === model.id); + const headers = { + ...config?.modelOverrides?.[model.id]?.headers, + ...definition?.headers, + ...extensionModel?.headers, + }; + return Object.keys(headers).length > 0 ? headers : undefined; +} + +export function validateExtensionProvider( + providerId: string, + base: Provider | undefined, + modelsConfig: ModelsJsonProvider | undefined, + extension: ProviderConfigInput, +): void { + if (extension.streamSimple && !extension.api) { + throw new Error(`Provider ${providerId}: "api" is required when registering streamSimple.`); + } + applyExtension(providerId, applyModelsJson(providerId, base?.getModels() ?? [], modelsConfig), extension); +} + +/** Compose built-in, models.json, and extension layers without reading credentials. */ +export function composeModelProvider( + providerId: string, + base: Provider | undefined, + modelConfig: ModelConfig, + extension: ProviderConfigInput | undefined, +): Provider { + const config = modelConfig.getProvider(providerId); + // models.json modelOverrides are the topmost user-config layer: they apply once, + // after custom-model upserts and extension model replacement. + const getModels = () => + applyExtension(providerId, applyModelsJson(providerId, base?.getModels() ?? [], config), extension).map( + (model) => { + const override = config?.modelOverrides?.[model.id]; + return override ? applyModelOverride(model, override) : model; + }, + ); + // Validate eagerly so registration/reload reports structural errors immediately. + getModels(); + const apiKey = composeApiKeyAuth(providerId, base, config, extension); + const oauth = composeOAuthAuth(providerId, base, config, extension); + if (!apiKey && !oauth) throw new Error(`Provider ${providerId}: no authentication method configured.`); + + const supportsBaseApi = (model: Model) => base?.getModels().some((entry) => entry.api === model.api) ?? false; + const streamWith = ( + model: Model, + context: Context, + options: StreamOptions | undefined, + simple: boolean, + ): AssistantMessageEventStream => + lazyStream(model, async () => { + if (extension?.streamSimple && model.api === extension.api) { + return extension.streamSimple(model, context, options as SimpleStreamOptions); + } + if (base && supportsBaseApi(model)) { + return simple + ? base.streamSimple(model, context, options as SimpleStreamOptions) + : base.stream(model, context, options); + } + const api = getApiProvider(model.api); + if (!api) throw new Error(`No API provider registered for api: ${model.api}`); + return simple + ? api.streamSimple(model, context, options as SimpleStreamOptions) + : api.stream(model, context, options); + }); + + return { + id: providerId, + name: extension?.name ?? config?.name ?? base?.name ?? extension?.oauth?.name ?? providerId, + baseUrl: extension?.baseUrl ?? config?.baseUrl ?? base?.baseUrl, + headers: base?.headers, + auth: { ...(apiKey ? { apiKey } : {}), ...(oauth ? { oauth } : {}) }, + getModels, + refreshModels: base?.refreshModels ? () => base.refreshModels!() : undefined, + filterModels: base?.filterModels + ? (models, credential: Credential | undefined) => base.filterModels!(models, credential) + : undefined, + stream: (model, context, options) => streamWith(model, context, options, false), + streamSimple: (model, context, options) => streamWith(model, context, options, true), + }; +} + +export function resolveConfiguredModelHeaders( + model: Model, + config: ModelsJsonProvider | undefined, + extension: ProviderConfigInput | undefined, + env?: Record, +): Record | undefined { + return resolveHeadersOrThrow( + rawModelHeaders(model, config, extension), + `model "${model.provider}/${model.id}"`, + env, + ); +} + +export interface CompatibilityRequestConfig { + headers?: ProviderHeaders; + authHeader: boolean; +} + +export function resolveCompatibilityRequestConfig( + model: Model, + config: ModelsJsonProvider | undefined, + extension: ProviderConfigInput | undefined, +): CompatibilityRequestConfig { + const configured = resolveHeadersOrThrow( + { ...configuredHeaders(config, extension), ...rawModelHeaders(model, config, extension) }, + `model "${model.provider}/${model.id}"`, + ); + return { + headers: model.headers || configured ? { ...model.headers, ...configured } : undefined, + authHeader: extension?.authHeader ?? config?.authHeader ?? false, + }; +} + +export function configuredRequestAuthStatus( + config: ModelsJsonProvider | undefined, + extension: ProviderConfigInput | undefined, +): AuthStatus | undefined { + const value = configuredApiKey(config, extension); + if (value === undefined) return undefined; + if (isCommandConfigValue(value)) return { configured: true, source: "models_json_command" }; + const names = getConfigValueEnvVarNames(value); + if (names.length > 0) { + return isConfigValueConfigured(value) + ? { configured: true, source: "environment", label: names.join(", ") } + : { configured: false }; + } + return { configured: true, source: extension?.apiKey !== undefined ? "fallback" : "models_json_key" }; +} diff --git a/packages/coding-agent/src/core/provider-display-names.ts b/packages/coding-agent/src/core/provider-display-names.ts deleted file mode 100644 index d33c3d7d..00000000 --- a/packages/coding-agent/src/core/provider-display-names.ts +++ /dev/null @@ -1,35 +0,0 @@ -export const BUILT_IN_PROVIDER_DISPLAY_NAMES: Record = { - anthropic: "Anthropic", - "amazon-bedrock": "Amazon Bedrock", - "ant-ling": "Ant Ling", - "azure-openai-responses": "Azure OpenAI Responses", - cerebras: "Cerebras", - "cloudflare-ai-gateway": "Cloudflare AI Gateway", - "cloudflare-workers-ai": "Cloudflare Workers AI", - deepseek: "DeepSeek", - fireworks: "Fireworks", - google: "Google Gemini", - "google-vertex": "Google Vertex AI", - groq: "Groq", - huggingface: "Hugging Face", - "kimi-coding": "Kimi For Coding", - mistral: "Mistral", - minimax: "MiniMax", - "minimax-cn": "MiniMax (China)", - moonshotai: "Moonshot AI", - "moonshotai-cn": "Moonshot AI (China)", - nvidia: "NVIDIA NIM", - opencode: "OpenCode Zen", - "opencode-go": "OpenCode Go", - openai: "OpenAI", - openrouter: "OpenRouter", - together: "Together AI", - "vercel-ai-gateway": "Vercel AI Gateway", - xai: "xAI", - zai: "ZAI Coding Plan (Global)", - "zai-coding-cn": "ZAI Coding Plan (China)", - xiaomi: "Xiaomi MiMo", - "xiaomi-token-plan-cn": "Xiaomi MiMo Token Plan (China)", - "xiaomi-token-plan-ams": "Xiaomi MiMo Token Plan (Amsterdam)", - "xiaomi-token-plan-sgp": "Xiaomi MiMo Token Plan (Singapore)", -}; diff --git a/packages/coding-agent/src/core/runtime-credentials.ts b/packages/coding-agent/src/core/runtime-credentials.ts new file mode 100644 index 00000000..022b145e --- /dev/null +++ b/packages/coding-agent/src/core/runtime-credentials.ts @@ -0,0 +1,48 @@ +import type { Credential, CredentialInfo, CredentialStore } from "@earendil-works/pi-ai"; + +/** Async credential store overlay for non-persistent runtime API keys. */ +export class RuntimeCredentials implements CredentialStore { + private readonly store: CredentialStore; + private readonly overrides = new Map(); + + constructor(store: CredentialStore) { + this.store = store; + } + + setRuntimeApiKey(providerId: string, apiKey: string): void { + this.overrides.set(providerId, apiKey); + } + + removeRuntimeApiKey(providerId: string): void { + this.overrides.delete(providerId); + } + + hasRuntimeApiKey(providerId: string): boolean { + return this.overrides.has(providerId); + } + + async read(providerId: string): Promise { + const override = this.overrides.get(providerId); + return override ? { type: "api_key", key: override } : this.store.read(providerId); + } + + async list(): Promise { + const entries = new Map((await this.store.list()).map((entry) => [entry.providerId, entry])); + for (const providerId of this.overrides.keys()) { + entries.set(providerId, { providerId, type: "api_key" }); + } + return [...entries.values()]; + } + + modify( + providerId: string, + fn: (current: Credential | undefined) => Promise, + ): Promise { + return this.store.modify(providerId, fn); + } + + async delete(providerId: string): Promise { + this.overrides.delete(providerId); + await this.store.delete(providerId); + } +} diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 4fc98670..9ac4bc76 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -1,16 +1,15 @@ import { join } from "node:path"; import { Agent, type AgentMessage, type ThinkingLevel } from "@earendil-works/pi-agent-core"; -import { clampThinkingLevel, type Message, type Model, streamSimple } from "@earendil-works/pi-ai/compat"; +import { clampThinkingLevel, type Message, type Model } from "@earendil-works/pi-ai/compat"; import { getAgentDir } from "../config.ts"; import { resolvePath } from "../utils/paths.ts"; import { AgentSession } from "./agent-session.ts"; import { formatNoModelsAvailableMessage } from "./auth-guidance.ts"; -import { AuthStorage } from "./auth-storage.ts"; import { DEFAULT_THINKING_LEVEL } from "./defaults.ts"; import type { ExtensionRunner, LoadExtensionsResult, SessionStartEvent, ToolDefinition } from "./extensions/index.ts"; import { convertToLlm } from "./messages.ts"; -import { ModelRegistry } from "./model-registry.ts"; import { findInitialModel } from "./model-resolver.ts"; +import { ModelRuntime } from "./model-runtime.ts"; import { mergeProviderAttributionHeaders } from "./provider-attribution.ts"; import type { ResourceLoader } from "./resource-loader.ts"; import { DefaultResourceLoader } from "./resource-loader.ts"; @@ -37,10 +36,8 @@ export interface CreateAgentSessionOptions { /** Global config directory. Default: ~/.pi/agent */ agentDir?: string; - /** Auth storage for credentials. Default: AuthStorage.create(agentDir/auth.json) */ - authStorage?: AuthStorage; - /** Model registry. Default: ModelRegistry.create(authStorage, agentDir/models.json) */ - modelRegistry?: ModelRegistry; + /** Canonical model/auth runtime. Defaults to a runtime using agentDir/auth.json and models.json. */ + modelRuntime?: ModelRuntime; /** Model to use. Default: from settings, else first available */ model?: Model; @@ -169,11 +166,9 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const agentDir = options.agentDir ? resolvePath(options.agentDir) : getDefaultAgentDir(); let resourceLoader = options.resourceLoader; - // Use provided or create AuthStorage and ModelRegistry const authPath = options.agentDir ? join(agentDir, "auth.json") : undefined; const modelsPath = options.agentDir ? join(agentDir, "models.json") : undefined; - const authStorage = options.authStorage ?? AuthStorage.create(authPath); - const modelRegistry = options.modelRegistry ?? ModelRegistry.create(authStorage, modelsPath); + const modelRuntime = options.modelRuntime ?? (await ModelRuntime.create({ authPath, modelsPath })); const settingsManager = options.settingsManager ?? SettingsManager.create(cwd, agentDir); const sessionManager = options.sessionManager ?? SessionManager.create(cwd, getDefaultSessionDir(cwd, agentDir)); @@ -194,8 +189,8 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} // If session has data, try to restore model from it if (!model && hasExistingSession && existingSession.model) { - const restoredModel = modelRegistry.find(existingSession.model.provider, existingSession.model.modelId); - if (restoredModel && modelRegistry.hasConfiguredAuth(restoredModel)) { + const restoredModel = modelRuntime.getModel(existingSession.model.provider, existingSession.model.modelId); + if (restoredModel && modelRuntime.hasConfiguredAuth(restoredModel.provider)) { model = restoredModel; } if (!model) { @@ -211,7 +206,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} defaultProvider: settingsManager.getDefaultProvider(), defaultModelId: settingsManager.getDefaultModel(), defaultThinkingLevel: settingsManager.getDefaultThinkingLevel(), - modelRegistry, + modelRuntime, }); model = result.model; if (!model) { @@ -300,11 +295,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} }, convertToLlm: convertToLlmWithBlockImages, streamFn: async (model, context, options) => { - const auth = await modelRegistry.getApiKeyAndHeaders(model); - if (!auth.ok) { - throw new Error(auth.error); - } - const env = auth.env || options?.env ? { ...(auth.env ?? {}), ...(options?.env ?? {}) } : undefined; const providerRetrySettings = settingsManager.getProviderRetrySettings(); const httpIdleTimeoutMs = settingsManager.getHttpIdleTimeoutMs(); // SDKs treat timeout=0 as 0ms (immediate timeout), not "no timeout". @@ -313,28 +303,24 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const timeoutMs = options?.timeoutMs ?? providerRetrySettings.timeoutMs ?? effectiveTimeoutMs; const websocketConnectTimeoutMs = options?.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs(); - let headers = mergeProviderAttributionHeaders( - model, - settingsManager, - options?.sessionId, - auth.headers, - options?.headers, - ); - // Let extensions inject/adjust per-request headers (e.g. tracing, session correlation) - // after static assembly, before the provider HTTP call. const headerRunner = extensionRunnerRef.current; - if (headerRunner?.hasHandlers("before_provider_headers")) { - headers = await headerRunner.emitBeforeProviderHeaders(headers ?? {}); - } - return streamSimple(model, context, { + return modelRuntime.streamSimple(model, context, { ...options, - apiKey: auth.apiKey, - env, timeoutMs, websocketConnectTimeoutMs, maxRetries: options?.maxRetries ?? providerRetrySettings.maxRetries, maxRetryDelayMs: options?.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs, - headers, + transformHeaders: async (requestHeaders) => { + const headers = mergeProviderAttributionHeaders( + model, + settingsManager, + options?.sessionId, + requestHeaders, + ); + return headerRunner?.hasHandlers("before_provider_headers") + ? headerRunner.emitBeforeProviderHeaders(headers ?? {}) + : (headers ?? {}); + }, }); }, onPayload: async (payload, _model) => { @@ -390,7 +376,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} scopedModels: options.scopedModels, resourceLoader, customTools: options.customTools, - modelRegistry, + modelRuntime, initialActiveToolNames, allowedToolNames, excludedToolNames, diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 7c6de94b..26757fb8 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -23,17 +23,7 @@ export { parseSkillBlock, type SessionStats, } from "./core/agent-session.ts"; -// Auth and model registry -export { - type ApiKeyCredential, - type AuthCredential, - type AuthStatus, - AuthStorage, - type AuthStorageBackend, - FileAuthStorageBackend, - InMemoryAuthStorageBackend, - type OAuthCredential, -} from "./core/auth-storage.ts"; +export { readStoredCredential } from "./core/auth-storage.ts"; // Compaction export { type BranchPreparation, @@ -178,6 +168,11 @@ export { resolveModelScopeWithDiagnostics, type ScopedModel, } from "./core/model-resolver.ts"; +export { + type CreateModelRuntimeOptions, + ModelRuntime, + type ModelRuntimeAuthOverrides, +} from "./core/model-runtime.ts"; export type { PackageManager, PathMetadata, diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index 11c3f120..e4da50a1 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -23,12 +23,11 @@ import { createAgentSessionServices, } from "./core/agent-session-services.ts"; import { formatNoModelsAvailableMessage } from "./core/auth-guidance.ts"; -import { AuthStorage } from "./core/auth-storage.ts"; import { exportFromFile } from "./core/export-html/index.ts"; import type { InlineExtension } from "./core/extensions/types.ts"; import { applyHttpProxySettings, configureHttpDispatcher } from "./core/http-dispatcher.ts"; -import type { ModelRegistry } from "./core/model-registry.ts"; import { resolveCliModel, resolveModelScope, type ScopedModel } from "./core/model-resolver.ts"; +import type { ModelRuntime } from "./core/model-runtime.ts"; import { restoreStdout, takeOverStdout } from "./core/output-guard.ts"; import { type AppMode, resolveProjectTrusted } from "./core/project-trust.ts"; import type { CreateAgentSessionOptions } from "./core/sdk.ts"; @@ -358,7 +357,7 @@ function buildSessionOptions( parsed: Args, scopedModels: ScopedModel[], hasExistingSession: boolean, - modelRegistry: ModelRegistry, + modelRuntime: ModelRuntime, settingsManager: SettingsManager, ): { options: CreateAgentSessionOptions; @@ -377,7 +376,7 @@ function buildSessionOptions( cliProvider: parsed.provider, cliModel: parsed.model, cliThinking: parsed.thinking, - modelRegistry, + modelRuntime, }); if (resolved.warning) { diagnostics.push({ type: "warning", message: resolved.warning }); @@ -400,7 +399,7 @@ function buildSessionOptions( // Check if saved default is in scoped models - use it if so, otherwise first scoped model const savedProvider = settingsManager.getDefaultProvider(); const savedModelId = settingsManager.getDefaultModel(); - const savedModel = savedProvider && savedModelId ? modelRegistry.find(savedProvider, savedModelId) : undefined; + const savedModel = savedProvider && savedModelId ? modelRuntime.getModel(savedProvider, savedModelId) : undefined; const savedInScope = savedModel ? scopedModels.find((sm) => modelsAreEqual(sm.model, savedModel)) : undefined; if (savedInScope) { @@ -433,7 +432,7 @@ function buildSessionOptions( })); } - // API key from CLI - set in authStorage + // API key from CLI - set as a non-persistent runtime override // (handled by caller before createAgentSession) // Tools @@ -611,7 +610,6 @@ export async function main(args: string[], options?: MainOptions) { const resolvedSkillPaths = resolveCliPaths(cwd, parsed.skills); const resolvedPromptTemplatePaths = resolveCliPaths(cwd, parsed.promptTemplates); const resolvedThemePaths = resolveCliPaths(cwd, parsed.themes); - const authStorage = AuthStorage.create(); const createRuntime: CreateAgentSessionRuntimeFactory = async ({ cwd, agentDir, @@ -634,7 +632,6 @@ export async function main(args: string[], options?: MainOptions) { const services = await createAgentSessionServices({ cwd, agentDir, - authStorage, settingsManager: runtimeSettingsManager, extensionFlagValues: parsed.unknownFlags, resourceLoaderReloadOptions: shouldResolveProjectTrust @@ -676,7 +673,7 @@ export async function main(args: string[], options?: MainOptions) { extensionFactories: options?.extensionFactories, }, }); - const { settingsManager, modelRegistry, resourceLoader } = services; + const { settingsManager, modelRuntime, resourceLoader } = services; const diagnostics: AgentSessionRuntimeDiagnostic[] = [ ...projectTrustDiagnostics, ...services.diagnostics, @@ -689,7 +686,7 @@ export async function main(args: string[], options?: MainOptions) { const modelPatterns = parsed.models ?? settingsManager.getEnabledModels(); const scopedModels = - modelPatterns && modelPatterns.length > 0 ? await resolveModelScope(modelPatterns, modelRegistry) : []; + modelPatterns && modelPatterns.length > 0 ? await resolveModelScope(modelPatterns, modelRuntime) : []; const { options: sessionOptions, cliThinkingFromModel, @@ -698,7 +695,7 @@ export async function main(args: string[], options?: MainOptions) { parsed, scopedModels, sessionManager.buildSessionContext().messages.length > 0, - modelRegistry, + modelRuntime, settingsManager, ); diagnostics.push(...sessionOptionDiagnostics); @@ -710,7 +707,8 @@ export async function main(args: string[], options?: MainOptions) { message: "--api-key requires a model to be specified via --model, --provider/--model, or --models", }); } else { - authStorage.setRuntimeApiKey(sessionOptions.model.provider, parsed.apiKey); + modelRuntime.setRuntimeApiKey(sessionOptions.model.provider, parsed.apiKey); + await services.modelRuntime.getAvailable(); } } @@ -745,7 +743,7 @@ export async function main(args: string[], options?: MainOptions) { }); time("createAgentSessionRuntime"); const { services, session, modelFallbackMessage } = runtime; - const { settingsManager, modelRegistry, resourceLoader } = services; + const { settingsManager, modelRuntime, resourceLoader } = services; applyHttpProxySettings(settingsManager.getGlobalSettings().httpProxy); configureHttpDispatcher(settingsManager.getHttpIdleTimeoutMs()); @@ -759,7 +757,7 @@ export async function main(args: string[], options?: MainOptions) { if (parsed.listModels !== undefined) { const searchPattern = typeof parsed.listModels === "string" ? parsed.listModels : undefined; - await listModels(modelRegistry, searchPattern); + await listModels(modelRuntime, searchPattern); process.exit(0); } diff --git a/packages/coding-agent/src/modes/interactive/components/footer.ts b/packages/coding-agent/src/modes/interactive/components/footer.ts index 4b2914fd..a781f893 100644 --- a/packages/coding-agent/src/modes/interactive/components/footer.ts +++ b/packages/coding-agent/src/modes/interactive/components/footer.ts @@ -138,7 +138,7 @@ export class FooterComponent implements Component { statsParts.push(`CH${latestCacheHitRate.toFixed(1)}%`); } // Show cost with "(sub)" indicator if using OAuth subscription - const usingSubscription = state.model ? this.session.modelRegistry.isUsingOAuth(state.model) : false; + const usingSubscription = state.model ? this.session.modelRuntime.isUsingOAuth(state.model.provider) : false; if (totalCost || usingSubscription) { const costStr = `$${totalCost.toFixed(3)}${usingSubscription ? " (sub)" : ""}`; statsParts.push(costStr); diff --git a/packages/coding-agent/src/modes/interactive/components/login-dialog.ts b/packages/coding-agent/src/modes/interactive/components/login-dialog.ts index 3fc4b516..010262cd 100644 --- a/packages/coding-agent/src/modes/interactive/components/login-dialog.ts +++ b/packages/coding-agent/src/modes/interactive/components/login-dialog.ts @@ -1,4 +1,4 @@ -import { getOAuthProviders, type OAuthDeviceCodeInfo } from "@earendil-works/pi-ai/oauth"; +import type { AuthInfoLink, OAuthDeviceCodeInfo } from "@earendil-works/pi-ai"; import { Container, type Focusable, getKeybindings, Input, Spacer, Text, type TUI } from "@earendil-works/pi-tui"; import { openBrowser } from "../../../utils/open-browser.ts"; import { theme } from "../theme/theme.ts"; @@ -38,8 +38,7 @@ export class LoginDialogComponent extends Container implements Focusable { this.tui = tui; this.onComplete = onComplete; - const providerInfo = getOAuthProviders().find((p) => p.id === providerId); - const providerName = providerNameOverride || providerInfo?.name || providerId; + const providerName = providerNameOverride || providerId; const title = titleOverride ?? `Login to ${providerName}`; // Top border @@ -176,17 +175,19 @@ export class LoginDialogComponent extends Container implements Focusable { }); } - /** - * Show informational text without prompting for input. - */ - showInfo(lines: string[]): void { - this.contentContainer.clear(); + /** Show provider-owned information and links without starting an auth callback flow. */ + showInfo(message: string, links: readonly AuthInfoLink[] = [], showCloseHint = false): void { this.contentContainer.addChild(new Spacer(1)); - for (const line of lines) { - this.contentContainer.addChild(new Text(line, 1, 0)); + this.contentContainer.addChild(new Text(theme.fg("text", message), 1, 0)); + for (const link of links) { + const text = link.label ? `${link.label}: ${link.url}` : link.url; + const hyperlink = `\x1b]8;;${link.url}\x07${text}\x1b]8;;\x07`; + this.contentContainer.addChild(new Text(theme.fg("accent", hyperlink), 1, 0)); + } + if (showCloseHint) { + this.contentContainer.addChild(new Spacer(1)); + this.contentContainer.addChild(new Text(`(${keyHint("tui.select.cancel", "to close")})`, 1, 0)); } - this.contentContainer.addChild(new Spacer(1)); - this.contentContainer.addChild(new Text(`(${keyHint("tui.select.cancel", "to close")})`, 1, 0)); this.tui.requestRender(); } 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 32711929..8a8c722a 100644 --- a/packages/coding-agent/src/modes/interactive/components/model-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/model-selector.ts @@ -9,7 +9,7 @@ import { Text, type TUI, } from "@earendil-works/pi-tui"; -import type { ModelRegistry } from "../../../core/model-registry.ts"; +import type { ModelRuntime } from "../../../core/model-runtime.ts"; import type { SettingsManager } from "../../../core/settings-manager.ts"; import { getModelSelectorSearchText } from "../model-search.ts"; import { theme } from "../theme/theme.ts"; @@ -52,7 +52,7 @@ export class ModelSelectorComponent extends Container implements Focusable { private selectedIndex: number = 0; private currentModel?: Model; private settingsManager: SettingsManager; - private modelRegistry: ModelRegistry; + private modelRuntime: ModelRuntime; private onSelectCallback: (model: Model) => void; private onCancelCallback: () => void; private errorMessage?: string; @@ -66,7 +66,7 @@ export class ModelSelectorComponent extends Container implements Focusable { tui: TUI, currentModel: Model | undefined, settingsManager: SettingsManager, - modelRegistry: ModelRegistry, + modelRuntime: ModelRuntime, scopedModels: ReadonlyArray, onSelect: (model: Model) => void, onCancel: () => void, @@ -77,7 +77,7 @@ export class ModelSelectorComponent extends Container implements Focusable { this.tui = tui; this.currentModel = currentModel; this.settingsManager = settingsManager; - this.modelRegistry = modelRegistry; + this.modelRuntime = modelRuntime; this.scopedModels = scopedModels; this.scope = scopedModels.length > 0 ? "scoped" : "all"; this.onSelectCallback = onSelect; @@ -139,17 +139,17 @@ export class ModelSelectorComponent extends Container implements Focusable { let models: ModelItem[]; // Refresh to pick up any changes to models.json - this.modelRegistry.refresh(); + await this.modelRuntime.refresh(); // Check for models.json errors - const loadError = this.modelRegistry.getError(); + const loadError = this.modelRuntime.getError(); if (loadError) { this.errorMessage = loadError; } // Load available models (built-in models still work even if models.json failed) try { - const availableModels = await this.modelRegistry.getAvailable(); + const availableModels = await this.modelRuntime.getAvailable(); models = availableModels.map((model: Model) => ({ provider: model.provider, id: model.id, @@ -166,7 +166,7 @@ export class ModelSelectorComponent extends Container implements Focusable { this.allModels = this.sortModels(models); this.scopedModels = this.scopedModels.map((scoped) => { - const refreshed = this.modelRegistry.find(scoped.model.provider, scoped.model.id); + const refreshed = this.modelRuntime.getModel(scoped.model.provider, scoped.model.id); return refreshed ? { ...scoped, model: refreshed } : scoped; }); this.scopedModelItems = this.scopedModels.map((scoped) => ({ diff --git a/packages/coding-agent/src/modes/interactive/components/oauth-selector.ts b/packages/coding-agent/src/modes/interactive/components/oauth-selector.ts index 514f189f..d599168b 100644 --- a/packages/coding-agent/src/modes/interactive/components/oauth-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/oauth-selector.ts @@ -1,3 +1,4 @@ +import type { ApiKeyAuth, AuthCheck, OAuthAuth } from "@earendil-works/pi-ai"; import { Container, type Focusable, @@ -7,7 +8,6 @@ import { Spacer, TruncatedText, } from "@earendil-works/pi-tui"; -import type { AuthStatus, AuthStorage } from "../../../core/auth-storage.ts"; import { theme } from "../theme/theme.ts"; import { DynamicBorder } from "./dynamic-border.ts"; @@ -15,6 +15,8 @@ export type AuthSelectorProvider = { id: string; name: string; authType: "oauth" | "api_key"; + method?: ApiKeyAuth | OAuthAuth; + status?: AuthCheck; }; export function formatAuthSelectorProviderType(authType: AuthSelectorProvider["authType"]): string { @@ -42,26 +44,20 @@ export class OAuthSelectorComponent extends Container implements Focusable { private filteredProviders: AuthSelectorProvider[]; private selectedIndex: number = 0; private mode: "login" | "logout"; - private authStorage: AuthStorage; - private getAuthStatus: (providerId: string) => AuthStatus; private onSelectCallback: (providerId: string, authType: AuthSelectorProvider["authType"]) => void; private onCancelCallback: () => void; private showAuthTypeLabels: boolean; constructor( mode: "login" | "logout", - authStorage: AuthStorage, providers: AuthSelectorProvider[], onSelect: (providerId: string, authType: AuthSelectorProvider["authType"]) => void, onCancel: () => void, - getAuthStatus?: (providerId: string) => AuthStatus, initialSearchInput?: string, ) { super(); this.mode = mode; - this.authStorage = authStorage; - this.getAuthStatus = getAuthStatus ?? ((providerId) => this.authStorage.getAuthStatus(providerId)); this.allProviders = providers; this.filteredProviders = providers; this.showAuthTypeLabels = new Set(providers.map((provider) => provider.authType)).size > 1; @@ -105,7 +101,11 @@ export class OAuthSelectorComponent extends Container implements Focusable { private filterProviders(query: string): void { this.filteredProviders = query - ? fuzzyFilter(this.allProviders, query, (provider) => `${provider.name} ${provider.id} ${provider.authType}`) + ? fuzzyFilter( + this.allProviders, + query, + (provider) => `${provider.name} ${provider.id} ${provider.authType} ${provider.method?.name ?? ""}`, + ) : this.allProviders; this.selectedIndex = Math.max(0, Math.min(this.selectedIndex, Math.max(0, this.filteredProviders.length - 1))); this.updateList(); @@ -162,29 +162,22 @@ export class OAuthSelectorComponent extends Container implements Focusable { } private formatStatusIndicator(provider: AuthSelectorProvider): string { - const credential = this.authStorage.get(provider.id); - if (credential?.type === provider.authType) return theme.fg("success", " ✓ configured"); - if (credential) { - const label = credential.type === "oauth" ? "subscription configured" : "API key configured"; + if (!provider.status) return theme.fg("muted", " • unconfigured"); + if (provider.status.type !== provider.authType) { + const label = provider.status.type === "oauth" ? "subscription configured" : "API key configured"; return theme.fg("muted", " • ") + theme.fg("warning", label); } - if (provider.authType !== "api_key") return theme.fg("muted", " • unconfigured"); - - const status = this.getAuthStatus(provider.id); - switch (status.source) { - case "environment": - return theme.fg("success", ` ✓ env: ${status.label ?? "API key"}`); - case "runtime": - return theme.fg("success", " ✓ runtime API key"); - case "fallback": - return theme.fg("success", " ✓ custom API key"); - case "models_json_key": - return theme.fg("success", " ✓ key in models.json"); - case "models_json_command": - return theme.fg("success", " ✓ command in models.json"); - default: - return theme.fg("muted", " • unconfigured"); + if ( + !provider.status.source || + provider.status.source === "OAuth" || + provider.status.source === "stored credential" + ) { + return theme.fg("success", " ✓ configured"); } + const source = /^[A-Z][A-Z0-9_]*(?:, [A-Z][A-Z0-9_]*)*$/.test(provider.status.source) + ? `env: ${provider.status.source}` + : provider.status.source; + return theme.fg("success", ` ✓ ${source}`); } handleInput(keyData: string): void { diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 8f492a18..5decd8ad 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -8,15 +8,8 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; import type { AgentMessage } from "@earendil-works/pi-agent-core"; -import { - type AssistantMessage, - getProviders, - type ImageContent, - type Message, - type Model, - type OAuthProviderId, - type OAuthSelectPrompt, -} from "@earendil-works/pi-ai/compat"; +import type { AuthEvent, AuthPrompt } from "@earendil-works/pi-ai"; +import type { AssistantMessage, ImageContent, Message, Model } from "@earendil-works/pi-ai/compat"; import type { AutocompleteItem, AutocompleteProvider, @@ -54,7 +47,6 @@ import { getAgentDir, getAuthPath, getDebugLogPath, - getDocsPath, getShareViewerUrl, VERSION, } from "../../config.ts"; @@ -85,7 +77,6 @@ import { type AppKeybinding, KeybindingsManager } from "../../core/keybindings.t import { createCompactionSummaryMessage } from "../../core/messages.ts"; import { defaultModelPerProvider, findExactModelReferenceMatch, resolveModelScope } from "../../core/model-resolver.ts"; import { DefaultPackageManager } from "../../core/package-manager.ts"; -import { BUILT_IN_PROVIDER_DISPLAY_NAMES } from "../../core/provider-display-names.ts"; import type { ResourceDiagnostic } from "../../core/resource-loader.ts"; import { formatMissingSessionCwdPrompt, MissingSessionCwdError } from "../../core/session-cwd.ts"; import { type SessionEntry, SessionManager, sessionEntryToContextMessages } from "../../core/session-manager.ts"; @@ -212,7 +203,7 @@ function isDeadTerminalError(error: unknown): boolean { } const ANTHROPIC_SUBSCRIPTION_AUTH_WARNING = - "Anthropic subscription auth is active. Third-party harness usage draws from extra usage and is billed per token, not your Claude plan limits. Manage extra usage at https://claude.ai/settings/usage."; + "Anthropic subscription auth is active. Third-party harness usage draws from extra usage and is billed per token, not your Claude plan limits. Manage extra usage at https://claude.ai/settings/usage. Disable this warning in /settings."; function isAnthropicSubscriptionAuthKey(apiKey: string | undefined): boolean { return typeof apiKey === "string" && apiKey.startsWith("sk-ant-oat"); @@ -248,24 +239,6 @@ function hasDefaultModelProvider(providerId: string): providerId is keyof typeof return providerId in defaultModelPerProvider; } -const BEDROCK_PROVIDER_ID = "amazon-bedrock"; - -const BUILT_IN_MODEL_PROVIDERS = new Set(getProviders()); - -export function isApiKeyLoginProvider( - providerId: string, - oauthProviderIds: ReadonlySet, - builtInProviderIds: ReadonlySet = BUILT_IN_MODEL_PROVIDERS, -): boolean { - if (BUILT_IN_PROVIDER_DISPLAY_NAMES[providerId]) { - return true; - } - if (builtInProviderIds.has(providerId)) { - return false; - } - return !oauthProviderIds.has(providerId); -} - type LoginProviderCompletionOption = { id: string; name: string; @@ -571,12 +544,12 @@ export class InteractiveMode { const modelCommand = slashCommands.find((command) => command.name === "model"); if (modelCommand) { - modelCommand.getArgumentCompletions = (prefix: string): AutocompleteItem[] | null => { + modelCommand.getArgumentCompletions = async (prefix: string): Promise => { // Get available models (scoped or from registry) const models = this.session.scopedModels.length > 0 ? this.session.scopedModels.map((s) => s.model) - : this.session.modelRegistry.getAvailable(); + : await this.session.modelRuntime.getAvailable(); if (models.length === 0) return null; @@ -879,7 +852,7 @@ export class InteractiveMode { this.showWarning(`Migrated credentials to auth.json: ${migratedProviders.join(", ")}`); } - const modelsJsonError = this.session.modelRegistry.getError(); + const modelsJsonError = this.session.modelRuntime.getError(); if (modelsJsonError) { this.showError(`models.json error: ${modelsJsonError}`); } @@ -1779,7 +1752,7 @@ export class InteractiveMode { hasUI: true, cwd: this.sessionManager.getCwd(), sessionManager: this.sessionManager, - modelRegistry: this.session.modelRegistry, + modelRegistry: extensionRunner.getModelRegistry(), model: this.session.model, isIdle: () => this.session.isIdle, isProjectTrusted: () => this.settingsManager.isProjectTrusted(), @@ -3288,7 +3261,7 @@ export class InteractiveMode { // Cache-miss notices are not persisted; re-derive them from the full entry // list and re-inject them after the assistant messages that paid for them. const cacheMisses = this.settingsManager.getShowCacheMissNotices() - ? collectCacheMisses(this.sessionManager.getEntries(), this.session.modelRegistry) + ? collectCacheMisses(this.sessionManager.getEntries(), this.session.modelRuntime) : new Map(); if (options.updateFooter) { @@ -3392,7 +3365,7 @@ export class InteractiveMode { if (!this.settingsManager.getShowCacheMissNotices()) return; // Entries don't contain `message` yet: message_end fires before persistence. - const miss = detectCacheMiss(this.sessionManager.getEntries(), message, this.session.modelRegistry); + const miss = detectCacheMiss(this.sessionManager.getEntries(), message, this.session.modelRuntime); if (miss) this.addCacheMissNotice(miss); } @@ -4325,9 +4298,9 @@ export class InteractiveMode { return this.session.scopedModels.map((scoped) => scoped.model); } - this.session.modelRegistry.refresh(); try { - return await this.session.modelRegistry.getAvailable(); + await this.session.modelRuntime.refresh(); + return [...(await this.session.modelRuntime.getAvailable())]; } catch { return []; } @@ -4353,15 +4326,13 @@ export class InteractiveMode { return; } - const storedCredential = this.session.modelRegistry.authStorage.get("anthropic"); - if (storedCredential?.type === "oauth") { - this.anthropicSubscriptionWarningShown = true; - this.showWarning(ANTHROPIC_SUBSCRIPTION_AUTH_WARNING); - return; - } - try { - const apiKey = await this.session.modelRegistry.getApiKeyForProvider(model.provider); + if ((await this.session.modelRuntime.checkAuth("anthropic"))?.type === "oauth") { + this.anthropicSubscriptionWarningShown = true; + this.showWarning(ANTHROPIC_SUBSCRIPTION_AUTH_WARNING); + return; + } + const apiKey = (await this.session.modelRuntime.getAuth(model.provider))?.auth.apiKey; if (!isAnthropicSubscriptionAuthKey(apiKey)) { return; } @@ -4429,7 +4400,7 @@ export class InteractiveMode { this.ui, this.session.model, this.settingsManager, - this.session.modelRegistry, + this.session.modelRuntime, this.session.scopedModels, async (model) => { try { @@ -4457,8 +4428,8 @@ export class InteractiveMode { private async showModelsSelector(): Promise { // Get all available models - this.session.modelRegistry.refresh(); - const allModels = this.session.modelRegistry.getAvailable(); + await this.session.modelRuntime.refresh(); + const allModels = [...(await this.session.modelRuntime.getAvailable())]; if (allModels.length === 0) { this.showStatus("No models available"); @@ -4479,7 +4450,7 @@ export class InteractiveMode { // Fall back to settings const patterns = this.settingsManager.getEnabledModels(); if (patterns !== undefined && patterns.length > 0) { - const scopedModels = await resolveModelScope(patterns, this.session.modelRegistry); + const scopedModels = await resolveModelScope(patterns, this.session.modelRuntime); currentEnabledIds = scopedModels.map((scoped) => `${scoped.model.provider}/${scoped.model.id}`); } } @@ -4488,7 +4459,7 @@ export class InteractiveMode { const updateSessionModels = async (enabledIds: string[] | null) => { currentEnabledIds = enabledIds === null ? null : [...enabledIds]; if (enabledIds && enabledIds.length > 0 && enabledIds.length < allModels.length) { - const newScopedModels = await resolveModelScope(enabledIds, this.session.modelRegistry); + const newScopedModels = await resolveModelScope(enabledIds, this.session.modelRuntime); this.session.setScopedModels( newScopedModels.map((sm) => ({ model: sm.model, @@ -4790,48 +4761,46 @@ export class InteractiveMode { } private getLoginProviderOptions(authType?: "oauth" | "api_key"): AuthSelectorProvider[] { - const authStorage = this.session.modelRegistry.authStorage; - const oauthProviders = authStorage.getOAuthProviders(); - const oauthProviderIds = new Set(oauthProviders.map((provider) => provider.id)); - const options: AuthSelectorProvider[] = oauthProviders.map((provider) => ({ - id: provider.id, - name: provider.name, - authType: "oauth", - })); - - const modelProviders = new Set(this.session.modelRegistry.getAll().map((model) => model.provider)); - for (const providerId of modelProviders) { - if (!isApiKeyLoginProvider(providerId, oauthProviderIds)) { - continue; + const options: AuthSelectorProvider[] = []; + for (const provider of this.session.modelRuntime.getProviders()) { + const authStatus = this.session.modelRuntime.getProviderAuthStatus(provider.id); + const status = authStatus.configured + ? { + type: this.session.modelRuntime.isUsingOAuth(provider.id) ? ("oauth" as const) : ("api_key" as const), + source: authStatus.label ?? authStatus.source, + } + : undefined; + if ((!authType || authType === "oauth") && provider.auth.oauth) { + options.push({ + id: provider.id, + name: provider.name, + authType: "oauth", + method: provider.auth.oauth, + status, + }); + } + if ((!authType || authType === "api_key") && provider.auth.apiKey) { + options.push({ + id: provider.id, + name: provider.name, + authType: "api_key", + method: provider.auth.apiKey, + status, + }); } - options.push({ - id: providerId, - name: this.session.modelRegistry.getProviderDisplayName(providerId), - authType: "api_key", - }); } - - const filteredOptions = authType ? options.filter((option) => option.authType === authType) : options; - return filteredOptions.sort((a, b) => a.name.localeCompare(b.name)); + return options.sort((a, b) => a.name.localeCompare(b.name)); } - private getLogoutProviderOptions(): AuthSelectorProvider[] { - const authStorage = this.session.modelRegistry.authStorage; - const options: AuthSelectorProvider[] = []; - - for (const providerId of authStorage.list()) { - const credential = authStorage.get(providerId); - if (!credential) { - continue; - } - options.push({ + private async getLogoutProviderOptions(): Promise { + return (await this.session.modelRuntime.listCredentials()) + .map(({ providerId, type }) => ({ id: providerId, - name: this.session.modelRegistry.getProviderDisplayName(providerId), - authType: credential.type, - }); - } - - return options.sort((a, b) => a.name.localeCompare(b.name)); + name: this.session.modelRuntime.getProvider(providerId)?.name ?? providerId, + authType: type, + status: { type, source: "stored credential" }, + })) + .sort((a, b) => a.name.localeCompare(b.name)); } private findLoginProviderOptions(providerRef: string): AuthSelectorProvider[] { @@ -4848,6 +4817,7 @@ export class InteractiveMode { } private async handleLoginCommand(providerRef?: string): Promise { + await this.session.modelRuntime.getAvailable(); if (!providerRef) { this.showLoginAuthTypeSelector(); return; @@ -4873,10 +4843,10 @@ export class InteractiveMode { private async startProviderLogin(providerOption: AuthSelectorProvider): Promise { if (providerOption.authType === "oauth") { await this.showLoginDialog(providerOption.id, providerOption.name); - } else if (providerOption.id === BEDROCK_PROVIDER_ID) { - this.showBedrockSetupDialog(providerOption.id, providerOption.name); - } else { + } else if (providerOption.method?.login) { await this.showApiKeyLoginDialog(providerOption.id, providerOption.name); + } else { + this.showAmbientAuthDialog(providerOption); } } @@ -4951,7 +4921,6 @@ export class InteractiveMode { this.showSelector((done) => { const selector = new OAuthSelectorComponent( "login", - this.session.modelRegistry.authStorage, providerOptions, async (providerId, selectedAuthType) => { done(); @@ -4973,7 +4942,6 @@ export class InteractiveMode { this.ui.requestRender(); } }, - (providerId) => this.session.modelRegistry.getProviderAuthStatus(providerId), initialSearchInput, ); return { component: selector, focus: selector }; @@ -4986,7 +4954,7 @@ export class InteractiveMode { return; } - const providerOptions = this.getLogoutProviderOptions(); + const providerOptions = await this.getLogoutProviderOptions(); 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.", @@ -4997,7 +4965,6 @@ export class InteractiveMode { this.showSelector((done) => { const selector = new OAuthSelectorComponent( mode, - this.session.modelRegistry.authStorage, providerOptions, async (providerId: string) => { done(); @@ -5008,8 +4975,7 @@ export class InteractiveMode { } try { - this.session.modelRegistry.authStorage.logout(providerOption.id); - this.session.modelRegistry.refresh(); + await this.session.modelRuntime.logout(providerOption.id); await this.updateAvailableProviderCount(); const message = providerOption.authType === "oauth" @@ -5035,14 +5001,14 @@ export class InteractiveMode { authType: "oauth" | "api_key", previousModel: Model | undefined, ): Promise { - this.session.modelRegistry.refresh(); + 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 = this.session.modelRegistry.getAvailable(); + const availableModels = await this.session.modelRuntime.getAvailable(); 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.`; @@ -5082,7 +5048,7 @@ export class InteractiveMode { } } - private showBedrockSetupDialog(providerId: string, providerName: string): void { + private showAmbientAuthDialog(providerOption: AuthSelectorProvider): void { const restoreEditor = () => { this.editorContainer.clear(); this.editorContainer.addChild(this.editor); @@ -5092,17 +5058,12 @@ export class InteractiveMode { const dialog = new LoginDialogComponent( this.ui, - providerId, + providerOption.id, () => restoreEditor(), - providerName, - "Amazon Bedrock setup", + providerOption.name, + `${providerOption.name} setup`, ); - dialog.showInfo([ - theme.fg("text", "Amazon Bedrock uses AWS credentials instead of a single API key."), - theme.fg("text", "Configure an AWS profile, IAM keys, bearer token, or role-based credentials."), - theme.fg("muted", "See:"), - theme.fg("accent", ` ${path.join(getDocsPath(), "providers.md")}`), - ]); + dialog.showInfo(`${providerOption.method?.name ?? "Authentication"} is configured outside pi.`, [], true); this.editorContainer.clear(); this.editorContainer.addChild(dialog); @@ -5135,13 +5096,7 @@ export class InteractiveMode { }; try { - const apiKey = (await dialog.showPrompt("Enter API key:")).trim(); - if (!apiKey) { - throw new Error("API key cannot be empty."); - } - - this.session.modelRegistry.authStorage.set(providerId, { type: "api_key", key: apiKey }); - + await this.loginProvider(dialog, providerId, "api_key"); restoreEditor(); await this.completeProviderAuthentication(providerId, providerName, "api_key", previousModel); } catch (error: unknown) { @@ -5153,8 +5108,11 @@ export class InteractiveMode { } } - private showOAuthLoginSelect(dialog: LoginDialogComponent, prompt: OAuthSelectPrompt): Promise { - return new Promise((resolve) => { + private showAuthSelect( + dialog: LoginDialogComponent, + prompt: Extract, + ): Promise { + return new Promise((resolve, reject) => { const restoreDialog = () => { this.editorContainer.clear(); this.editorContainer.addChild(dialog); @@ -5167,11 +5125,13 @@ export class InteractiveMode { labels, (optionLabel) => { restoreDialog(); - resolve(prompt.options.find((option) => option.label === optionLabel)?.id); + const id = prompt.options.find((option) => option.label === optionLabel)?.id; + if (id) resolve(id); + else reject(new Error("Login cancelled")); }, () => { restoreDialog(); - resolve(undefined); + reject(new Error("Login cancelled")); }, ); this.editorContainer.clear(); @@ -5181,40 +5141,63 @@ export class InteractiveMode { }); } + private async showAuthPrompt(dialog: LoginDialogComponent, prompt: AuthPrompt): Promise { + let response: Promise; + if (prompt.type === "select") { + response = this.showAuthSelect(dialog, prompt); + } else if (prompt.type === "manual_code") { + response = dialog.showManualInput(prompt.message); + } else { + response = dialog.showPrompt(prompt.message, prompt.placeholder); + } + if (!prompt.signal) return response; + if (prompt.signal.aborted) throw new Error("Login cancelled"); + const signal = prompt.signal; + let onAbort: (() => void) | undefined; + const aborted = new Promise((_resolve, reject) => { + onAbort = () => reject(new Error("Login cancelled")); + signal.addEventListener("abort", onAbort, { once: true }); + }); + try { + return await Promise.race([response, aborted]); + } finally { + if (onAbort) signal.removeEventListener("abort", onAbort); + } + } + + private notifyAuthDialog(dialog: LoginDialogComponent, event: AuthEvent): void { + if (event.type === "auth_url") { + dialog.showAuth(event.url, event.instructions); + } else if (event.type === "device_code") { + dialog.showDeviceCode(event); + dialog.showWaiting("Waiting for authentication..."); + } else if (event.type === "info") { + dialog.showInfo(event.message, event.links); + } else { + dialog.showProgress(event.message); + } + } + + private async loginProvider( + dialog: LoginDialogComponent, + providerId: string, + method: "api_key" | "oauth", + ): Promise { + await this.session.modelRuntime.login(providerId, method, { + signal: dialog.signal, + prompt: (prompt) => this.showAuthPrompt(dialog, prompt), + notify: (event) => this.notifyAuthDialog(dialog, event), + }); + } + private async showLoginDialog(providerId: string, providerName: string): Promise { - const providerInfo = this.session.modelRegistry.authStorage - .getOAuthProviders() - .find((provider) => provider.id === providerId); const previousModel = this.session.model; - - // Providers that use callback servers (can paste redirect URL) - const usesCallbackServer = providerInfo?.usesCallbackServer ?? false; - - // Create login dialog component - const dialog = new LoginDialogComponent( - this.ui, - providerId, - (_success, _message) => { - // Completion handled below - }, - providerName, - ); - - // Show dialog in editor container + const dialog = new LoginDialogComponent(this.ui, providerId, (_success, _message) => {}, providerName); this.editorContainer.clear(); this.editorContainer.addChild(dialog); this.ui.setFocus(dialog); this.ui.requestRender(); - // Promise for manual code input (racing with callback server) - let manualCodeResolve: ((code: string) => void) | undefined; - let manualCodeReject: ((err: Error) => void) | undefined; - const manualCodePromise = new Promise((resolve, reject) => { - manualCodeResolve = resolve; - manualCodeReject = reject; - }); - - // Restore editor helper const restoreEditor = () => { this.editorContainer.clear(); this.editorContainer.addChild(this.editor); @@ -5223,51 +5206,7 @@ export class InteractiveMode { }; try { - await this.session.modelRegistry.authStorage.login(providerId as OAuthProviderId, { - onAuth: (info: { url: string; instructions?: string }) => { - dialog.showAuth(info.url, info.instructions); - - if (usesCallbackServer) { - // Show input for manual paste, racing with callback - dialog - .showManualInput("Paste redirect URL below, or complete login in browser:") - .then((value) => { - if (value && manualCodeResolve) { - manualCodeResolve(value); - manualCodeResolve = undefined; - } - }) - .catch(() => { - if (manualCodeReject) { - manualCodeReject(new Error("Login cancelled")); - manualCodeReject = undefined; - } - }); - } - // For Anthropic: onPrompt is called immediately after - }, - - onDeviceCode: (info) => { - dialog.showDeviceCode(info); - dialog.showWaiting("Waiting for authentication..."); - }, - - onPrompt: async (prompt: { message: string; placeholder?: string }) => { - return dialog.showPrompt(prompt.message, prompt.placeholder); - }, - - onProgress: (message: string) => { - dialog.showProgress(message); - }, - - onSelect: (prompt: OAuthSelectPrompt) => this.showOAuthLoginSelect(dialog, prompt), - - onManualCodeInput: () => manualCodePromise, - - signal: dialog.signal, - }); - - // Success + await this.loginProvider(dialog, providerId, "oauth"); restoreEditor(); await this.completeProviderAuthentication(providerId, providerName, "oauth", previousModel); } catch (error: unknown) { @@ -5368,7 +5307,7 @@ export class InteractiveMode { showDiagnosticsWhenQuiet: true, }); const savedImplicitProjectTrust = this.maybeSaveImplicitProjectTrustAfterReload(); - const modelsJsonError = this.session.modelRegistry.getError(); + const modelsJsonError = this.session.modelRuntime.getError(); if (modelsJsonError) { this.showError(`models.json error: ${modelsJsonError}`); } @@ -5613,7 +5552,7 @@ export class InteractiveMode { const stats = this.session.getSessionStats(); const sessionName = this.sessionManager.getSessionName(); const entries = this.sessionManager.getEntries(); - const cacheWaste = computeCacheWaste(entries, this.session.modelRegistry); + const cacheWaste = computeCacheWaste(entries, this.session.modelRuntime); // Cost/token totals per provider/model actually used (e.g. OpenRouter `auto` // resolves to a concrete responseModel), sorted by cost descending. diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index b099f874..58b64531 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -465,7 +465,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}`); @@ -483,7 +483,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { @@ -18,7 +18,7 @@ describe("AgentSession auto-compaction queue resume", () => { let settingsManager: SettingsManager; let tempDir: string; - beforeEach(() => { + beforeEach(async () => { tempDir = join(tmpdir(), `pi-auto-compaction-queue-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); vi.useFakeTimers(); @@ -35,15 +35,15 @@ describe("AgentSession auto-compaction queue resume", () => { sessionManager = SessionManager.inMemory(); settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - authStorage.setRuntimeApiKey("anthropic", "test-key"); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); + const modelRegistry = await createModelRegistry(authStorage, tempDir); session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), }); }); diff --git a/packages/coding-agent/test/agent-session-branching.test.ts b/packages/coding-agent/test/agent-session-branching.test.ts index 76d355e1..7016ecca 100644 --- a/packages/coding-agent/test/agent-session-branching.test.ts +++ b/packages/coding-agent/test/agent-session-branching.test.ts @@ -48,7 +48,7 @@ describe.skipIf(!API_KEY)("AgentSession forking", () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; sessionManager = noSession ? SessionManager.inMemory(tempDir) : SessionManager.create(tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - authStorage.setRuntimeApiKey("anthropic", API_KEY!); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: API_KEY! })); const servicesOptions = { agentDir: tempDir, diff --git a/packages/coding-agent/test/agent-session-compaction.test.ts b/packages/coding-agent/test/agent-session-compaction.test.ts index 516c3b45..61ad0a2b 100644 --- a/packages/coding-agent/test/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/agent-session-compaction.test.ts @@ -1,3 +1,4 @@ +import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; /** * E2E tests for AgentSession compaction behavior. * @@ -15,7 +16,6 @@ import { getModel } from "@earendil-works/pi-ai/compat"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AgentSession, type AgentSessionEvent } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; import { createCodingTools } from "../src/index.ts"; @@ -27,7 +27,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => { let sessionManager: SessionManager; let events: AgentSessionEvent[]; - beforeEach(() => { + beforeEach(async () => { // Create temp directory for session files tempDir = join(tmpdir(), `pi-compaction-test-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); @@ -45,7 +45,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => { } }); - function createSession(inMemory = false) { + async function createSession(inMemory = false) { const model = getModel("anthropic", "claude-sonnet-4-5")!; const agent = new Agent({ getApiKey: () => API_KEY, @@ -61,14 +61,14 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => { // Use minimal keepRecentTokens so small test conversations have something to summarize settingsManager.applyOverrides({ compaction: { keepRecentTokens: 1 } }); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage); + const modelRegistry = await createModelRegistry(authStorage); session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), }); @@ -81,7 +81,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => { } it("should trigger manual compaction via compact()", async () => { - createSession(); + await createSession(); // Send a few prompts to build up history await session.prompt("What is 2+2? Reply with just the number."); @@ -107,7 +107,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => { }, 120000); it("should maintain valid session state after compaction", async () => { - createSession(); + await createSession(); // Build up history await session.prompt("What is the capital of France? One word answer."); @@ -132,7 +132,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => { }, 180000); it("should persist compaction to session file", async () => { - createSession(); + await createSession(); await session.prompt("Say hello"); await session.agent.waitForIdle(); @@ -160,7 +160,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => { }, 120000); it("should work with --no-session mode (in-memory only)", async () => { - createSession(true); // in-memory mode + await createSession(true); // in-memory mode // Send prompts await session.prompt("What is 2+2? Reply with just the number."); @@ -182,7 +182,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => { }, 120000); it("should emit compaction events during manual compaction", async () => { - createSession(); + await createSession(); // Build some history await session.prompt("Say hello"); diff --git a/packages/coding-agent/test/agent-session-concurrent.test.ts b/packages/coding-agent/test/agent-session-concurrent.test.ts index 46599068..bd825ec7 100644 --- a/packages/coding-agent/test/agent-session-concurrent.test.ts +++ b/packages/coding-agent/test/agent-session-concurrent.test.ts @@ -1,3 +1,4 @@ +import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; /** * Tests for AgentSession concurrent prompt guard. */ @@ -18,7 +19,6 @@ import { Type } from "typebox"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; import type { BuildSystemPromptOptions } from "../src/core/system-prompt.ts"; @@ -62,7 +62,7 @@ describe("AgentSession concurrent prompt guard", () => { let session: AgentSession; let tempDir: string; - beforeEach(() => { + beforeEach(async () => { tempDir = join(tmpdir(), `pi-concurrent-test-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); }); @@ -78,7 +78,7 @@ describe("AgentSession concurrent prompt guard", () => { } }); - function createSession() { + async function createSession() { const model = getModel("anthropic", "claude-sonnet-4-5")!; let abortSignal: AbortSignal | undefined; @@ -111,16 +111,16 @@ describe("AgentSession concurrent prompt guard", () => { const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); + const modelRegistry = await createModelRegistry(authStorage, tempDir); // Set a runtime API key so validation passes - authStorage.setRuntimeApiKey("anthropic", "test-key"); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), }); @@ -128,7 +128,7 @@ describe("AgentSession concurrent prompt guard", () => { } it("should throw when prompt() called while streaming", async () => { - createSession(); + await createSession(); // Start first prompt (don't await, it will block until abort) const firstPrompt = session.prompt("First message"); @@ -150,7 +150,7 @@ describe("AgentSession concurrent prompt guard", () => { }); it("should allow steer() while streaming", async () => { - createSession(); + await createSession(); // Start first prompt const firstPrompt = session.prompt("First message"); @@ -166,7 +166,7 @@ describe("AgentSession concurrent prompt guard", () => { }); it("should allow followUp() while streaming", async () => { - createSession(); + await createSession(); // Start first prompt const firstPrompt = session.prompt("First message"); @@ -236,8 +236,8 @@ describe("AgentSession concurrent prompt guard", () => { const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); - authStorage.setRuntimeApiKey("anthropic", "test-key"); + const modelRegistry = await createModelRegistry(authStorage, tempDir); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); const extensionsResult = await createTestExtensionsResult([ (pi) => { @@ -255,7 +255,7 @@ describe("AgentSession concurrent prompt guard", () => { sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader({ extensionsResult }), }); session.subscribe((event) => { @@ -314,15 +314,15 @@ describe("AgentSession concurrent prompt guard", () => { const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); - authStorage.setRuntimeApiKey("anthropic", "test-key"); + const modelRegistry = await createModelRegistry(authStorage, tempDir); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), }); @@ -420,15 +420,15 @@ describe("AgentSession concurrent prompt guard", () => { const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); - authStorage.setRuntimeApiKey("anthropic", "test-key"); + const modelRegistry = await createModelRegistry(authStorage, tempDir); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), baseToolsOverride: { dummy: tool }, }); @@ -567,15 +567,15 @@ describe("AgentSession concurrent prompt guard", () => { const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); - authStorage.setRuntimeApiKey("anthropic", "test-key"); + const modelRegistry = await createModelRegistry(authStorage, tempDir); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), baseToolsOverride: { dummy: tool }, }); diff --git a/packages/coding-agent/test/agent-session-dynamic-provider.test.ts b/packages/coding-agent/test/agent-session-dynamic-provider.test.ts index eee5583d..87014322 100644 --- a/packages/coding-agent/test/agent-session-dynamic-provider.test.ts +++ b/packages/coding-agent/test/agent-session-dynamic-provider.test.ts @@ -4,6 +4,7 @@ import { join } from "node:path"; import { getModel } from "@earendil-works/pi-ai/compat"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; +import { ModelRuntime } from "../src/core/model-runtime.ts"; import { DefaultResourceLoader } from "../src/core/resource-loader.ts"; import type { ExtensionFactory } from "../src/core/sdk.ts"; import { createAgentSession } from "../src/core/sdk.ts"; @@ -30,7 +31,11 @@ describe("AgentSession dynamic provider registration", () => { const settingsManager = SettingsManager.create(tempDir, agentDir); const sessionManager = SessionManager.inMemory(); const authStorage = AuthStorage.create(join(agentDir, "auth.json")); - authStorage.setRuntimeApiKey("anthropic", "test-key"); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); + const modelRuntime = await ModelRuntime.create({ + credentials: authStorage, + modelsPath: join(agentDir, "models.json"), + }); const resourceLoader = new DefaultResourceLoader({ cwd: tempDir, agentDir, @@ -45,7 +50,7 @@ describe("AgentSession dynamic provider registration", () => { model: getModel("anthropic", "claude-sonnet-4-5")!, settingsManager, sessionManager, - authStorage, + modelRuntime, resourceLoader, }); diff --git a/packages/coding-agent/test/agent-session-retry.test.ts b/packages/coding-agent/test/agent-session-retry.test.ts index 6e3d0582..f7d5eced 100644 --- a/packages/coding-agent/test/agent-session-retry.test.ts +++ b/packages/coding-agent/test/agent-session-retry.test.ts @@ -7,9 +7,9 @@ import { Type } from "typebox"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; +import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; import { createTestResourceLoader } from "./utilities.ts"; class MockAssistantStream extends EventStream { @@ -54,7 +54,7 @@ describe("AgentSession retry", () => { let session: AgentSession; let tempDir: string; - beforeEach(() => { + beforeEach(async () => { tempDir = join(tmpdir(), `pi-retry-test-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); }); @@ -68,7 +68,11 @@ describe("AgentSession retry", () => { } }); - function createSession(options?: { failCount?: number; maxRetries?: number; delayAssistantMessageEndMs?: number }) { + async function createSession(options?: { + failCount?: number; + maxRetries?: number; + delayAssistantMessageEndMs?: number; + }) { const failCount = options?.failCount ?? 1; const maxRetries = options?.maxRetries ?? 3; const delayAssistantMessageEndMs = options?.delayAssistantMessageEndMs ?? 0; @@ -102,8 +106,8 @@ describe("AgentSession retry", () => { const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); - authStorage.setRuntimeApiKey("anthropic", "test-key"); + const modelRegistry = await createModelRegistry(authStorage, tempDir); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); settingsManager.applyOverrides({ retry: { enabled: true, maxRetries, baseDelayMs: 1 } }); session = new AgentSession({ @@ -111,7 +115,7 @@ describe("AgentSession retry", () => { sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), }); @@ -130,7 +134,7 @@ describe("AgentSession retry", () => { } it("retries after a transient error and succeeds", async () => { - const created = createSession({ failCount: 1 }); + const created = await createSession({ failCount: 1 }); const events: string[] = []; created.session.subscribe((event) => { if (event.type === "auto_retry_start") events.push(`start:${event.attempt}`); @@ -145,7 +149,7 @@ describe("AgentSession retry", () => { }); it("exhausts max retries and emits failure", async () => { - const created = createSession({ failCount: 99, maxRetries: 2 }); + const created = await createSession({ failCount: 99, maxRetries: 2 }); const events: string[] = []; created.session.subscribe((event) => { if (event.type === "auto_retry_start") events.push(`start:${event.attempt}`); @@ -162,7 +166,7 @@ describe("AgentSession retry", () => { }); it("prompt waits for retry completion even when assistant message_end handling is delayed", async () => { - const created = createSession({ failCount: 1, delayAssistantMessageEndMs: 40 }); + const created = await createSession({ failCount: 1, delayAssistantMessageEndMs: 40 }); await created.session.prompt("Test"); @@ -171,7 +175,7 @@ describe("AgentSession retry", () => { }); it("retries provider network_error failures", async () => { - const created = createSession({ failCount: 0 }); + const created = await createSession({ failCount: 0 }); let callCount = 0; const streamFn = () => { callCount++; @@ -204,15 +208,15 @@ describe("AgentSession retry", () => { const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); - authStorage.setRuntimeApiKey("anthropic", "test-key"); + const modelRegistry = await createModelRegistry(authStorage, tempDir); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } }); session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), }); @@ -289,8 +293,8 @@ describe("AgentSession retry", () => { const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); - authStorage.setRuntimeApiKey("anthropic", "test-key"); + const modelRegistry = await createModelRegistry(authStorage, tempDir); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } }); session = new AgentSession({ @@ -298,7 +302,7 @@ describe("AgentSession retry", () => { sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), baseToolsOverride: { echo: echoTool }, }); diff --git a/packages/coding-agent/test/agent-session-runtime-events.test.ts b/packages/coding-agent/test/agent-session-runtime-events.test.ts index 348c0451..32bb828f 100644 --- a/packages/coding-agent/test/agent-session-runtime-events.test.ts +++ b/packages/coding-agent/test/agent-session-runtime-events.test.ts @@ -10,6 +10,7 @@ import { createAgentSessionServices, } from "../src/core/agent-session-runtime.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; +import { ModelRuntime } from "../src/core/model-runtime.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import type { ExtensionFactory, @@ -42,11 +43,33 @@ describe("AgentSessionRuntime session lifecycle events", () => { faux.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two"), fauxAssistantMessage("three")]); const authStorage = AuthStorage.inMemory(); - authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + await authStorage.modify(faux.getModel().provider, async () => ({ type: "api_key", key: "faux-key" })); + const modelRuntime = await ModelRuntime.create({ + credentials: authStorage, + modelsPath: join(tempDir, "models.json"), + }); + const model = faux.getModel(); + modelRuntime.registerProvider(model.provider, { + baseUrl: model.baseUrl, + api: model.api, + models: [ + { + id: model.id, + name: model.name, + api: model.api, + reasoning: model.reasoning, + input: model.input, + cost: model.cost, + contextWindow: model.contextWindow, + maxTokens: model.maxTokens, + baseUrl: model.baseUrl, + }, + ], + }); const runtimeOptions = { agentDir: tempDir, - authStorage, + modelRuntime, model: faux.getModel(), resourceLoaderOptions: { extensionFactories: [extensionFactory], diff --git a/packages/coding-agent/test/agent-session-stats.test.ts b/packages/coding-agent/test/agent-session-stats.test.ts index 0b6c4853..9a442c73 100644 --- a/packages/coding-agent/test/agent-session-stats.test.ts +++ b/packages/coding-agent/test/agent-session-stats.test.ts @@ -3,9 +3,9 @@ import { type AssistantMessage, getModel, type Usage } from "@earendil-works/pi- import { describe, expect, it } from "vitest"; import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; +import { createInMemoryModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; import { createTestResourceLoader } from "./utilities.ts"; const model = getModel("anthropic", "claude-sonnet-4-5")!; @@ -48,11 +48,11 @@ function createUserMessage(text: string, timestamp: number) { }; } -function createSession() { +async function createSession() { const settingsManager = SettingsManager.inMemory(); const sessionManager = SessionManager.inMemory(); const authStorage = AuthStorage.inMemory(); - authStorage.setRuntimeApiKey("anthropic", "test-key"); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); const session = new AgentSession({ agent: new Agent({ getApiKey: () => "test-key", @@ -66,7 +66,7 @@ function createSession() { sessionManager, settingsManager, cwd: process.cwd(), - modelRegistry: ModelRegistry.inMemory(authStorage), + modelRuntime: getModelRuntime(await createInMemoryModelRegistry(authStorage)), resourceLoader: createTestResourceLoader(), }); @@ -78,8 +78,8 @@ function syncAgentMessages(session: AgentSession, sessionManager: SessionManager } describe("AgentSession.getSessionStats", () => { - it("exposes the current context usage alongside token totals", () => { - const { session, sessionManager } = createSession(); + it("exposes the current context usage alongside token totals", async () => { + const { session, sessionManager } = await createSession(); try { sessionManager.appendMessage(createUserMessage("hello", 1)); @@ -96,8 +96,8 @@ describe("AgentSession.getSessionStats", () => { } }); - it("reports unknown current context usage immediately after compaction", () => { - const { session, sessionManager } = createSession(); + it("reports unknown current context usage immediately after compaction", async () => { + const { session, sessionManager } = await createSession(); try { sessionManager.appendMessage(createUserMessage("first", 1)); @@ -119,8 +119,8 @@ describe("AgentSession.getSessionStats", () => { } }); - it("uses post-compaction usage for current context instead of stale kept usage", () => { - const { session, sessionManager } = createSession(); + it("uses post-compaction usage for current context instead of stale kept usage", async () => { + const { session, sessionManager } = await createSession(); try { sessionManager.appendMessage(createUserMessage("first", 1)); @@ -143,8 +143,8 @@ describe("AgentSession.getSessionStats", () => { } }); - it("ignores zero-usage messages when checking for post-compaction context usage", () => { - const { session, sessionManager } = createSession(); + it("ignores zero-usage messages when checking for post-compaction context usage", async () => { + const { session, sessionManager } = await createSession(); try { sessionManager.appendMessage(createUserMessage("first", 1)); diff --git a/packages/coding-agent/test/agent-session-tree-navigation.test.ts b/packages/coding-agent/test/agent-session-tree-navigation.test.ts index 361e5bff..5cd20450 100644 --- a/packages/coding-agent/test/agent-session-tree-navigation.test.ts +++ b/packages/coding-agent/test/agent-session-tree-navigation.test.ts @@ -15,8 +15,8 @@ import { API_KEY, createTestSession, type TestSessionContext } from "./utilities describe.skipIf(!API_KEY)("AgentSession tree navigation e2e", () => { let ctx: TestSessionContext; - beforeEach(() => { - ctx = createTestSession({ + beforeEach(async () => { + ctx = await createTestSession({ systemPrompt: "You are a helpful assistant. Reply with just a few words.", settingsOverrides: { compaction: { keepRecentTokens: 1 } }, }); @@ -279,8 +279,8 @@ describe.skipIf(!API_KEY)("AgentSession tree navigation e2e", () => { describe.skipIf(!API_KEY)("AgentSession tree navigation - branch scenarios", () => { let ctx: TestSessionContext; - beforeEach(() => { - ctx = createTestSession({ + beforeEach(async () => { + ctx = await createTestSession({ systemPrompt: "You are a helpful assistant. Reply with just a few words.", }); }); diff --git a/packages/coding-agent/test/auth-storage.test.ts b/packages/coding-agent/test/auth-storage.test.ts index 62ecd81a..28c49ba9 100644 --- a/packages/coding-agent/test/auth-storage.test.ts +++ b/packages/coding-agent/test/auth-storage.test.ts @@ -1,17 +1,14 @@ import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import { registerOAuthProvider } from "@earendil-works/pi-ai/oauth"; +import { createModels, type Provider } from "@earendil-works/pi-ai"; import lockfile from "proper-lockfile"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { clearConfigValueCache, resolveConfigValueUncached } from "../src/core/resolve-config-value.ts"; -import * as shellModule from "../src/utils/shell.ts"; describe("AuthStorage", () => { let tempDir: string; let authJsonPath: string; - let authStorage: AuthStorage; beforeEach(() => { tempDir = join(tmpdir(), `pi-test-auth-storage-${Date.now()}-${Math.random().toString(36).slice(2)}`); @@ -20,680 +17,201 @@ describe("AuthStorage", () => { }); afterEach(() => { - if (tempDir && existsSync(tempDir)) { - rmSync(tempDir, { recursive: true }); - } - clearConfigValueCache(); + if (existsSync(tempDir)) rmSync(tempDir, { recursive: true }); vi.restoreAllMocks(); }); - function writeAuthJson(data: Record) { + function writeAuthJson(data: Record): void { writeFileSync(authJsonPath, JSON.stringify(data)); } - function toShPath(value: string): string { - return value.replace(/\\/g, "/").replace(/"/g, '\\"'); - } + test("reads and resolves stored API-key credentials", async () => { + const original = process.env.TEST_AUTH_STORAGE_KEY; + process.env.TEST_AUTH_STORAGE_KEY = "environment-key"; + try { + writeAuthJson({ anthropic: { type: "api_key", key: "$TEST_AUTH_STORAGE_KEY" } }); + const storage = AuthStorage.create(authJsonPath); + expect(await storage.read("anthropic")).toEqual({ type: "api_key", key: "environment-key" }); + } finally { + if (original === undefined) delete process.env.TEST_AUTH_STORAGE_KEY; + else process.env.TEST_AUTH_STORAGE_KEY = original; + } + }); - describe("API key resolution", () => { - test("literal API key is returned directly", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "sk-ant-literal-key" }, - }); + test("resolves command-backed API-key credentials", async () => { + writeAuthJson({ anthropic: { type: "api_key", key: "!printf 'command-key'" } }); + const storage = AuthStorage.create(authJsonPath); + expect(await storage.read("anthropic")).toEqual({ type: "api_key", key: "command-key" }); + }); - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); + test("returns OAuth credentials unchanged", async () => { + const credential = { + type: "oauth" as const, + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + }; + const storage = AuthStorage.inMemory({ anthropic: credential }); + expect(await storage.read("anthropic")).toEqual(credential); + }); - expect(apiKey).toBe("sk-ant-literal-key"); + test("credential-scoped env takes precedence and remains inspectable", async () => { + writeAuthJson({ + anthropic: { + type: "api_key", + key: "$SCOPED_KEY", + env: { SCOPED_KEY: "scoped-value", REGION: "test-region" }, + }, + }); + const storage = AuthStorage.create(authJsonPath); + expect(await storage.read("anthropic")).toMatchObject({ + key: "scoped-value", + env: { SCOPED_KEY: "scoped-value", REGION: "test-region" }, + }); + }); + + test("modify persists a credential while preserving unrelated external edits", async () => { + writeAuthJson({ anthropic: { type: "api_key", key: "old" } }); + const storage = AuthStorage.create(authJsonPath); + writeAuthJson({ + anthropic: { type: "api_key", key: "old" }, + openai: { type: "api_key", key: "external" }, }); - test("apiKey with ! prefix executes command and uses stdout", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "!echo test-api-key-from-command" }, - }); + await storage.modify("anthropic", async () => ({ type: "api_key", key: "new" })); - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); + expect(JSON.parse(readFileSync(authJsonPath, "utf8"))).toEqual({ + anthropic: { type: "api_key", key: "new" }, + openai: { type: "api_key", key: "external" }, + }); + }); - expect(apiKey).toBe("test-api-key-from-command"); + test("modify with undefined leaves the current credential unchanged", async () => { + writeAuthJson({ anthropic: { type: "api_key", key: "stored" } }); + const storage = AuthStorage.create(authJsonPath); + expect(await storage.modify("anthropic", async () => undefined)).toEqual({ type: "api_key", key: "stored" }); + expect(await storage.read("anthropic")).toEqual({ type: "api_key", key: "stored" }); + }); + + test("serializes concurrent modifications", async () => { + writeAuthJson({}); + const first = AuthStorage.create(authJsonPath); + const second = AuthStorage.create(authJsonPath); + await Promise.all([ + first.modify("anthropic", async () => ({ type: "api_key", key: "anthropic-key" })), + second.modify("openai", async () => ({ type: "api_key", key: "openai-key" })), + ]); + expect(JSON.parse(readFileSync(authJsonPath, "utf8"))).toEqual({ + anthropic: { type: "api_key", key: "anthropic-key" }, + openai: { type: "api_key", key: "openai-key" }, + }); + }); + + test("delete removes one credential while preserving others", async () => { + writeAuthJson({ + anthropic: { type: "api_key", key: "anthropic-key" }, + openai: { type: "api_key", key: "openai-key" }, + }); + const storage = AuthStorage.create(authJsonPath); + writeAuthJson({ + anthropic: { type: "api_key", key: "anthropic-key" }, + openai: { type: "api_key", key: "openai-key" }, + google: { type: "api_key", key: "external-key" }, + }); + await storage.delete("anthropic"); + await expect(storage.list()).resolves.toEqual([ + { providerId: "openai", type: "api_key" }, + { providerId: "google", type: "api_key" }, + ]); + expect(await storage.read("anthropic")).toBeUndefined(); + expect(await storage.read("openai")).toEqual({ type: "api_key", key: "openai-key" }); + expect(await storage.read("google")).toEqual({ type: "api_key", key: "external-key" }); + }); + + test("in-memory storage implements the same credential-store behavior", async () => { + const storage = AuthStorage.inMemory({ anthropic: { type: "api_key", key: "initial" } }); + expect(await storage.read("anthropic")).toEqual({ type: "api_key", key: "initial" }); + await storage.modify("anthropic", async () => ({ type: "api_key", key: "updated" })); + expect(await storage.read("anthropic")).toEqual({ type: "api_key", key: "updated" }); + await storage.delete("anthropic"); + await expect(storage.list()).resolves.toEqual([]); + }); + + test("does not write after lock acquisition failure and recovers on retry", async () => { + writeAuthJson({ anthropic: { type: "api_key", key: "stored" } }); + const storage = AuthStorage.create(authJsonPath); + const lockSpy = vi.spyOn(lockfile, "lock").mockRejectedValueOnce(new Error("lock unavailable")); + + await expect(storage.modify("openai", async () => ({ type: "api_key", key: "new" }))).rejects.toThrow( + "lock unavailable", + ); + expect(JSON.parse(readFileSync(authJsonPath, "utf8"))).toEqual({ + anthropic: { type: "api_key", key: "stored" }, }); - test("apiKey with ! prefix trims whitespace from command output", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "!echo ' spaced-key '" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("spaced-key"); + lockSpy.mockRestore(); + await storage.modify("openai", async () => ({ type: "api_key", key: "new" })); + expect(JSON.parse(readFileSync(authJsonPath, "utf8"))).toEqual({ + anthropic: { type: "api_key", key: "stored" }, + openai: { type: "api_key", key: "new" }, }); + }); - test("apiKey with ! prefix handles multiline output (uses trimmed result)", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "!printf 'line1\\nline2'" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("line1\nline2"); + test("surfaces a compromised OAuth refresh lock and allows a later retry", async () => { + const providerId = "oauth-provider"; + writeAuthJson({ + [providerId]: { + type: "oauth", + access: "expired-access", + refresh: "refresh-token", + expires: 0, + }, }); - - test("apiKey with ! prefix returns undefined on command failure", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "!exit 1" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBeUndefined(); - }); - - test("apiKey with ! prefix returns undefined on nonexistent command", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "!nonexistent-command-12345" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBeUndefined(); - }); - - test("apiKey with ! prefix returns undefined on empty output", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "!printf ''" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBeUndefined(); - }); - - test("apiKey with $ prefix resolves to env value", async () => { - const originalEnv = process.env.TEST_AUTH_API_KEY_12345; - process.env.TEST_AUTH_API_KEY_12345 = "env-api-key-value"; - - try { - writeAuthJson({ - anthropic: { type: "api_key", key: "$TEST_AUTH_API_KEY_12345" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("env-api-key-value"); - } finally { - if (originalEnv === undefined) { - delete process.env.TEST_AUTH_API_KEY_12345; - } else { - process.env.TEST_AUTH_API_KEY_12345 = originalEnv; - } - } - }); - - test("apiKey env bag takes precedence over process.env", async () => { - const originalEnv = process.env.TEST_AUTH_SCOPED_API_KEY_12345; - process.env.TEST_AUTH_SCOPED_API_KEY_12345 = "process-env-value"; - - try { - writeAuthJson({ - anthropic: { - type: "api_key", - key: "$TEST_AUTH_SCOPED_API_KEY_12345", - env: { TEST_AUTH_SCOPED_API_KEY_12345: "credential-env-value" }, + const storage = AuthStorage.create(authJsonPath); + const provider: Provider = { + id: providerId, + name: "OAuth Provider", + auth: { + oauth: { + name: "OAuth", + login: async () => { + throw new Error("not used"); }, - }); - - authStorage = AuthStorage.create(authJsonPath); - - expect(await authStorage.getApiKey("anthropic")).toBe("credential-env-value"); - expect(authStorage.getProviderEnv("anthropic")).toEqual({ - TEST_AUTH_SCOPED_API_KEY_12345: "credential-env-value", - }); - } finally { - if (originalEnv === undefined) { - delete process.env.TEST_AUTH_SCOPED_API_KEY_12345; - } else { - process.env.TEST_AUTH_SCOPED_API_KEY_12345 = originalEnv; - } - } - }); - - test("apiKey with braced env syntax resolves to env value", async () => { - const originalEnv = process.env.TEST_AUTH_BRACED_API_KEY_12345; - process.env.TEST_AUTH_BRACED_API_KEY_12345 = "braced-env-api-key-value"; - const bracedKey = "$" + "{TEST_AUTH_BRACED_API_KEY_12345}"; - - try { - writeAuthJson({ - anthropic: { type: "api_key", key: bracedKey }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("braced-env-api-key-value"); - } finally { - if (originalEnv === undefined) { - delete process.env.TEST_AUTH_BRACED_API_KEY_12345; - } else { - process.env.TEST_AUTH_BRACED_API_KEY_12345 = originalEnv; - } - } - }); - - test("apiKey interpolates braced env references inside literals", async () => { - const originalPartA = process.env.TEST_AUTH_INTERPOLATED_PART_A_12345; - const originalPartB = process.env.TEST_AUTH_INTERPOLATED_PART_B_12345; - process.env.TEST_AUTH_INTERPOLATED_PART_A_12345 = "left"; - process.env.TEST_AUTH_INTERPOLATED_PART_B_12345 = "right"; - const interpolatedKey = [ - "$", - "{TEST_AUTH_INTERPOLATED_PART_A_12345}_$", - "{TEST_AUTH_INTERPOLATED_PART_B_12345}", - ].join(""); - - try { - writeAuthJson({ - anthropic: { type: "api_key", key: interpolatedKey }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("left_right"); - } finally { - if (originalPartA === undefined) { - delete process.env.TEST_AUTH_INTERPOLATED_PART_A_12345; - } else { - process.env.TEST_AUTH_INTERPOLATED_PART_A_12345 = originalPartA; - } - if (originalPartB === undefined) { - delete process.env.TEST_AUTH_INTERPOLATED_PART_B_12345; - } else { - process.env.TEST_AUTH_INTERPOLATED_PART_B_12345 = originalPartB; - } - } - }); - - test("apiKey with $$ prefix escapes a leading dollar", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "$$TEST_AUTH_API_KEY_12345" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("$TEST_AUTH_API_KEY_12345"); - }); - - test("apiKey with $! escapes a literal bang and still interpolates later env refs", async () => { - const originalEnv = process.env.TEST_AUTH_API_KEY_12345; - process.env.TEST_AUTH_API_KEY_12345 = "env-api-key-value"; - - try { - writeAuthJson({ - anthropic: { type: "api_key", key: "$!literal-$TEST_AUTH_API_KEY_12345" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("!literal-env-api-key-value"); - } finally { - if (originalEnv === undefined) { - delete process.env.TEST_AUTH_API_KEY_12345; - } else { - process.env.TEST_AUTH_API_KEY_12345 = originalEnv; - } - } - }); - - test("plain API key is used directly even when it matches an env var", async () => { - const originalEnv = process.env.TEST_AUTH_API_KEY_12345; - process.env.TEST_AUTH_API_KEY_12345 = "env-api-key-value"; - - try { - writeAuthJson({ - anthropic: { type: "api_key", key: "TEST_AUTH_API_KEY_12345" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("TEST_AUTH_API_KEY_12345"); - } finally { - if (originalEnv === undefined) { - delete process.env.TEST_AUTH_API_KEY_12345; - } else { - process.env.TEST_AUTH_API_KEY_12345 = originalEnv; - } - } - }); - - test("literal public API key is not corrupted by the Windows PUBLIC env var", async () => { - const originalPublic = process.env.PUBLIC; - process.env.PUBLIC = "C:\\Users\\Public"; - - try { - writeAuthJson({ - opencode: { type: "api_key", key: "public" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("opencode"); - - expect(apiKey).toBe("public"); - } finally { - if (originalPublic === undefined) { - delete process.env.PUBLIC; - } else { - process.env.PUBLIC = originalPublic; - } - } - }); - - test("apiKey as literal value is used directly when not an env var", async () => { - // Make sure this isn't an env var - delete process.env.literal_api_key_value; - - writeAuthJson({ - anthropic: { type: "api_key", key: "literal_api_key_value" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("literal_api_key_value"); - }); - - test("apiKey command can use shell features like pipes", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "!echo 'hello world' | tr ' ' '-'" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("hello-world"); - }); - - test("command config uses stdin when configured shell requires it", () => { - if (process.platform === "win32") return; - const platformDescriptor = Object.getOwnPropertyDescriptor(process, "platform"); - vi.spyOn(shellModule, "getShellConfig").mockReturnValue({ - shell: "/bin/bash", - args: ["-s"], - commandTransport: "stdin", - }); - - try { - Object.defineProperty(process, "platform", { - configurable: true, - value: "win32", - }); - const nameExpansion = "$" + "{name}"; - - expect(resolveConfigValueUncached(`!name='World'; echo "Hello, ${nameExpansion}!"`)).toBe("Hello, World!"); - } finally { - if (platformDescriptor) { - Object.defineProperty(process, "platform", platformDescriptor); - } - } - }); - - describe("caching", () => { - test("command is only executed once per process", async () => { - // Use a command that writes to a file to count invocations - const counterFile = join(tempDir, "counter"); - writeFileSync(counterFile, "0"); - - const counterPath = toShPath(counterFile); - const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; echo "key-value"'`; - writeAuthJson({ - anthropic: { type: "api_key", key: command }, - }); - - authStorage = AuthStorage.create(authJsonPath); - - // Call multiple times - await authStorage.getApiKey("anthropic"); - await authStorage.getApiKey("anthropic"); - await authStorage.getApiKey("anthropic"); - - // Command should have only run once - const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); - expect(count).toBe(1); - }); - - test("cache persists across AuthStorage instances", async () => { - const counterFile = join(tempDir, "counter"); - writeFileSync(counterFile, "0"); - - const counterPath = toShPath(counterFile); - const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; echo "key-value"'`; - writeAuthJson({ - anthropic: { type: "api_key", key: command }, - }); - - // Create multiple AuthStorage instances - const storage1 = AuthStorage.create(authJsonPath); - await storage1.getApiKey("anthropic"); - - const storage2 = AuthStorage.create(authJsonPath); - await storage2.getApiKey("anthropic"); - - // Command should still have only run once - const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); - expect(count).toBe(1); - }); - - test("clearConfigValueCache allows command to run again", async () => { - const counterFile = join(tempDir, "counter"); - writeFileSync(counterFile, "0"); - - const counterPath = toShPath(counterFile); - const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; echo "key-value"'`; - writeAuthJson({ - anthropic: { type: "api_key", key: command }, - }); - - authStorage = AuthStorage.create(authJsonPath); - await authStorage.getApiKey("anthropic"); - - // Clear cache and call again - clearConfigValueCache(); - await authStorage.getApiKey("anthropic"); - - // Command should have run twice - const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); - expect(count).toBe(2); - }); - - test("different commands are cached separately", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "!echo key-anthropic" }, - openai: { type: "api_key", key: "!echo key-openai" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - - const keyA = await authStorage.getApiKey("anthropic"); - const keyB = await authStorage.getApiKey("openai"); - - expect(keyA).toBe("key-anthropic"); - expect(keyB).toBe("key-openai"); - }); - - test("failed commands are cached (not retried)", async () => { - const counterFile = join(tempDir, "counter"); - writeFileSync(counterFile, "0"); - - const counterPath = toShPath(counterFile); - const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; exit 1'`; - writeAuthJson({ - anthropic: { type: "api_key", key: command }, - }); - - authStorage = AuthStorage.create(authJsonPath); - - // Call multiple times - all should return undefined - const key1 = await authStorage.getApiKey("anthropic"); - const key2 = await authStorage.getApiKey("anthropic"); - - expect(key1).toBeUndefined(); - expect(key2).toBeUndefined(); - - // Command should have only run once despite failures - const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); - expect(count).toBe(1); - }); - - test("environment variables are not cached (changes are picked up)", async () => { - const envVarName = "TEST_AUTH_KEY_CACHE_TEST_98765"; - const originalEnv = process.env[envVarName]; - - try { - process.env[envVarName] = "first-value"; - - writeAuthJson({ - anthropic: { type: "api_key", key: `$${envVarName}` }, - }); - - authStorage = AuthStorage.create(authJsonPath); - - const key1 = await authStorage.getApiKey("anthropic"); - expect(key1).toBe("first-value"); - - // Change env var - process.env[envVarName] = "second-value"; - - const key2 = await authStorage.getApiKey("anthropic"); - expect(key2).toBe("second-value"); - } finally { - if (originalEnv === undefined) { - delete process.env[envVarName]; - } else { - process.env[envVarName] = originalEnv; - } - } - }); - }); - }); - - describe("oauth lock compromise handling", () => { - test("returns undefined on compromised lock and allows a later retry", async () => { - const providerId = `test-oauth-provider-${Date.now()}-${Math.random().toString(36).slice(2)}`; - registerOAuthProvider({ - id: providerId, - name: "Test OAuth Provider", - async login() { - throw new Error("Not used in this test"); - }, - async refreshToken(credentials) { - return { - ...credentials, - access: "refreshed-access-token", + refresh: async (credential) => ({ + ...credential, + access: "refreshed-access", expires: Date.now() + 60_000, - }; + }), + toAuth: async (credential) => ({ apiKey: credential.access }), }, - getApiKey(credentials) { - return `Bearer ${credentials.access}`; - }, - }); + }, + getModels: () => [], + stream: () => { + throw new Error("not used"); + }, + streamSimple: () => { + throw new Error("not used"); + }, + }; + const models = createModels({ credentials: storage }); + models.setProvider(provider); - writeAuthJson({ - [providerId]: { - type: "oauth", - refresh: "refresh-token", - access: "expired-access-token", - expires: Date.now() - 10_000, - }, - }); - - authStorage = AuthStorage.create(authJsonPath); - - const realLock = lockfile.lock.bind(lockfile); - const lockSpy = vi.spyOn(lockfile, "lock"); - lockSpy.mockImplementationOnce(async (file, options) => { - options?.onCompromised?.(new Error("Unable to update lock within the stale threshold")); - return realLock(file, options); - }); - - const firstTry = await authStorage.getApiKey(providerId); - expect(firstTry).toBeUndefined(); - - lockSpy.mockRestore(); - - const secondTry = await authStorage.getApiKey(providerId); - expect(secondTry).toBe("Bearer refreshed-access-token"); + const realLock = lockfile.lock.bind(lockfile); + const lockSpy = vi.spyOn(lockfile, "lock").mockImplementationOnce(async (file, options) => { + options?.onCompromised?.(new Error("lock compromised")); + return realLock(file, options); }); + await expect(models.getAuth(providerId)).rejects.toMatchObject({ code: "auth" }); + + lockSpy.mockRestore(); + await expect(models.getAuth(providerId)).resolves.toMatchObject({ auth: { apiKey: "refreshed-access" } }); }); - describe("persistence semantics", () => { - test("set preserves unrelated external edits", () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "old-anthropic" }, - openai: { type: "api_key", key: "openai-key" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - - // Simulate external edit while process is running - writeAuthJson({ - anthropic: { type: "api_key", key: "old-anthropic" }, - openai: { type: "api_key", key: "openai-key" }, - google: { type: "api_key", key: "google-key" }, - }); - - authStorage.set("anthropic", { type: "api_key", key: "new-anthropic" }); - - const updated = JSON.parse(readFileSync(authJsonPath, "utf-8")) as Record; - expect(updated.anthropic.key).toBe("new-anthropic"); - expect(updated.openai.key).toBe("openai-key"); - expect(updated.google.key).toBe("google-key"); - }); - - test("remove preserves unrelated external edits", () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "anthropic-key" }, - openai: { type: "api_key", key: "openai-key" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - - // Simulate external edit while process is running - writeAuthJson({ - anthropic: { type: "api_key", key: "anthropic-key" }, - openai: { type: "api_key", key: "openai-key" }, - google: { type: "api_key", key: "google-key" }, - }); - - authStorage.remove("anthropic"); - - const updated = JSON.parse(readFileSync(authJsonPath, "utf-8")) as Record; - expect(updated.anthropic).toBeUndefined(); - expect(updated.openai.key).toBe("openai-key"); - expect(updated.google.key).toBe("google-key"); - }); - - test("throws and does not overwrite malformed auth file after load error", () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "anthropic-key" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - writeFileSync(authJsonPath, "{invalid-json", "utf-8"); - - authStorage.reload(); - expect(() => authStorage.set("openai", { type: "api_key", key: "openai-key" })).toThrow( - "Cannot update auth storage because it could not be loaded", - ); - - const raw = readFileSync(authJsonPath, "utf-8"); - expect(raw).toBe("{invalid-json"); - expect(authStorage.has("openai")).toBe(false); - }); - - test("throws when a stale auth lock prevents persistence", () => { - writeAuthJson({}); - writeFileSync(`${authJsonPath}.lock`, "", "utf-8"); - - authStorage = AuthStorage.create(authJsonPath); - expect(() => authStorage.set("github-copilot", { type: "api_key", key: "copilot-key" })).toThrow( - "Cannot update auth storage because it could not be loaded", - ); - - expect(readFileSync(authJsonPath, "utf-8")).toBe("{}"); - expect(authStorage.has("github-copilot")).toBe(false); - }); - - test("recovers from an earlier load error before persisting", () => { - writeAuthJson({}); - const lockPath = `${authJsonPath}.lock`; - writeFileSync(lockPath, "", "utf-8"); - - authStorage = AuthStorage.create(authJsonPath); - rmSync(lockPath); - authStorage.set("github-copilot", { type: "api_key", key: "copilot-key" }); - - const updated = JSON.parse(readFileSync(authJsonPath, "utf-8")) as Record; - expect(updated["github-copilot"].key).toBe("copilot-key"); - expect(authStorage.has("github-copilot")).toBe(true); - }); - - test("reload records parse errors and drainErrors clears buffer", () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "anthropic-key" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - writeFileSync(authJsonPath, "{invalid-json", "utf-8"); - - authStorage.reload(); - - // Keeps previous in-memory data on reload failure - expect(authStorage.get("anthropic")).toEqual({ type: "api_key", key: "anthropic-key" }); - - const firstDrain = authStorage.drainErrors(); - expect(firstDrain.length).toBeGreaterThan(0); - expect(firstDrain[0]).toBeInstanceOf(Error); - - const secondDrain = authStorage.drainErrors(); - expect(secondDrain).toHaveLength(0); - }); - }); - - describe("auth status", () => { - test("does not expose stored API keys or OAuth tokens", () => { - authStorage = AuthStorage.inMemory({ - anthropic: { type: "api_key", key: "secret-api-key" }, - openai: { - type: "oauth", - access: "secret-access-token", - refresh: "secret-refresh-token", - expires: Date.now() + 1000, - }, - }); - - expect(authStorage.getAuthStatus("anthropic")).toEqual({ configured: true, source: "stored" }); - expect(authStorage.getAuthStatus("openai")).toEqual({ configured: true, source: "stored" }); - expect(JSON.stringify(authStorage.getAuthStatus("anthropic"))).not.toContain("secret-api-key"); - expect(JSON.stringify(authStorage.getAuthStatus("openai"))).not.toContain("secret-access-token"); - expect(JSON.stringify(authStorage.getAuthStatus("openai"))).not.toContain("secret-refresh-token"); - }); - }); - - describe("runtime overrides", () => { - test("runtime override takes priority over auth.json", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "!echo stored-key" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - authStorage.setRuntimeApiKey("anthropic", "runtime-key"); - - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("runtime-key"); - }); - - test("removing runtime override falls back to auth.json", async () => { - writeAuthJson({ - anthropic: { type: "api_key", key: "!echo stored-key" }, - }); - - authStorage = AuthStorage.create(authJsonPath); - authStorage.setRuntimeApiKey("anthropic", "runtime-key"); - authStorage.removeRuntimeApiKey("anthropic"); - - const apiKey = await authStorage.getApiKey("anthropic"); - - expect(apiKey).toBe("stored-key"); - }); + test("does not overwrite malformed auth files", async () => { + writeAuthJson({ anthropic: { type: "api_key", key: "stored" } }); + const storage = AuthStorage.create(authJsonPath); + writeFileSync(authJsonPath, "{invalid-json", "utf8"); + await expect(storage.modify("openai", async () => ({ type: "api_key", key: "new" }))).rejects.toThrow(); + expect(readFileSync(authJsonPath, "utf8")).toBe("{invalid-json"); }); }); diff --git a/packages/coding-agent/test/cache-stats.test.ts b/packages/coding-agent/test/cache-stats.test.ts index 1d43bfed..a7c45e69 100644 --- a/packages/coding-agent/test/cache-stats.test.ts +++ b/packages/coding-agent/test/cache-stats.test.ts @@ -12,7 +12,7 @@ const zeroCost = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }; const models: ModelPriceSource = { // $/million tokens; used as cache-read price fallback on full-miss turns - find: () => ({ cost: { cacheRead: 0.3 } }), + getModel: () => ({ cost: { cacheRead: 0.3 } }), }; function assistant(options: { diff --git a/packages/coding-agent/test/compaction-extensions.test.ts b/packages/coding-agent/test/compaction-extensions.test.ts index 62ca9780..fabcb38d 100644 --- a/packages/coding-agent/test/compaction-extensions.test.ts +++ b/packages/coding-agent/test/compaction-extensions.test.ts @@ -1,3 +1,4 @@ +import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; /** * Tests for compaction extension events (before_compact / compact). */ @@ -17,7 +18,6 @@ import { type SessionCompactEvent, type SessionEvent, } from "../src/core/extensions/index.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; import { createSyntheticSourceInfo } from "../src/core/source-info.ts"; @@ -31,7 +31,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { let tempDir: string; let capturedEvents: SessionEvent[]; - beforeEach(() => { + beforeEach(async () => { tempDir = join(tmpdir(), `pi-compaction-extensions-test-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); capturedEvents = []; @@ -85,7 +85,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { }; } - function createSession(extensions: Extension[]) { + async function createSession(extensions: Extension[]) { const model = getModel("anthropic", "claude-sonnet-4-5")!; const agent = new Agent({ getApiKey: () => API_KEY, @@ -100,7 +100,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { const settingsManager = SettingsManager.create(tempDir, tempDir); settingsManager.applyOverrides({ compaction: { keepRecentTokens: 1 } }); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage); + const modelRegistry = await createModelRegistry(authStorage); const runtime = createExtensionRuntime(); const resourceLoader = { @@ -113,7 +113,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader, }); @@ -122,7 +122,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { it("should emit before_compact and compact events", async () => { const extension = createExtension(); - createSession([extension]); + await createSession([extension]); await session.prompt("What is 2+2? Reply with just the number."); await session.agent.waitForIdle(); @@ -158,7 +158,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { it("should allow extensions to cancel compaction", async () => { const extension = createExtension(() => ({ cancel: true })); - createSession([extension]); + await createSession([extension]); await session.prompt("What is 2+2? Reply with just the number."); await session.agent.waitForIdle(); @@ -184,7 +184,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { } return undefined; }); - createSession([extension]); + await createSession([extension]); await session.prompt("What is 2+2? Reply with just the number."); await session.agent.waitForIdle(); @@ -208,7 +208,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { it("should include entries in compact event after compaction is saved", async () => { const extension = createExtension(); - createSession([extension]); + await createSession([extension]); await session.prompt("What is 2+2? Reply with just the number."); await session.agent.waitForIdle(); @@ -259,7 +259,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { shortcuts: new Map(), }; - createSession([throwingExtension]); + await createSession([throwingExtension]); await session.prompt("What is 2+2? Reply with just the number."); await session.agent.waitForIdle(); @@ -339,7 +339,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { shortcuts: new Map(), }; - createSession([extension1, extension2]); + await createSession([extension1, extension2]); await session.prompt("What is 2+2? Reply with just the number."); await session.agent.waitForIdle(); @@ -356,7 +356,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { capturedBeforeEvent = event; return undefined; }); - createSession([extension]); + await createSession([extension]); await session.prompt("What is 2+2? Reply with just the number."); await session.agent.waitForIdle(); @@ -378,10 +378,9 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { expect(Array.isArray(event.branchEntries)).toBe(true); - // sessionManager, modelRegistry, and model are now on ctx, not event - // Verify they're accessible via session + // sessionManager and model runtime remain available on the session. expect(typeof session.sessionManager.getEntries).toBe("function"); - expect(typeof session.modelRegistry.getApiKeyAndHeaders).toBe("function"); + expect(typeof session.modelRuntime.getAuth).toBe("function"); const entries = session.sessionManager.getEntries(); expect(Array.isArray(entries)).toBe(true); @@ -403,7 +402,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { } return undefined; }); - createSession([extension]); + await createSession([extension]); await session.prompt("What is 2+2? Reply with just the number."); await session.agent.waitForIdle(); diff --git a/packages/coding-agent/test/config-value-migration.test.ts b/packages/coding-agent/test/config-value-migration.test.ts index 8a273b53..8bad16cf 100644 --- a/packages/coding-agent/test/config-value-migration.test.ts +++ b/packages/coding-agent/test/config-value-migration.test.ts @@ -4,9 +4,10 @@ import * as path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { ENV_AGENT_DIR } from "../src/config.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { runMigrations } from "../src/migrations.ts"; +import { createModelRegistry } from "./model-runtime-test-utils.ts"; + describe("config value env var syntax migration", () => { const tempDirs: string[] = []; @@ -71,7 +72,7 @@ describe("config value env var syntax migration", () => { it.each([ ["malformed", '{\n "providers": {\n'], ["blank", ""], - ])("does not throw on %s models.json during migrations", (_name, content) => { + ])("does not throw on %s models.json during migrations", async (_name, content) => { const agentDir = createAgentDir(); const modelsPath = path.join(agentDir, "models.json"); fs.writeFileSync(modelsPath, content, "utf-8"); @@ -79,7 +80,7 @@ describe("config value env var syntax migration", () => { withAgentDir(agentDir, () => expect(() => runMigrations(agentDir)).not.toThrow()); expect(fs.readFileSync(modelsPath, "utf-8")).toBe(content); - const registry = ModelRegistry.create(AuthStorage.create(path.join(agentDir, "auth.json")), modelsPath); + const registry = await createModelRegistry(AuthStorage.create(path.join(agentDir, "auth.json")), modelsPath); const loadError = registry.getError(); expect(loadError).toContain("Failed to parse models.json"); expect(loadError).toContain(`File: ${modelsPath}`); @@ -148,7 +149,7 @@ describe("config value env var syntax migration", () => { expect(provider.modelOverrides?.["model-b"]?.headers?.["x-override-key"]).toBe("OVERRIDE_API_KEY"); expect(logSpy).not.toHaveBeenCalled(); - const registry = ModelRegistry.create( + const registry = await createModelRegistry( AuthStorage.create(path.join(agentDir, "auth.json")), path.join(agentDir, "models.json"), ); diff --git a/packages/coding-agent/test/extensions-discovery.test.ts b/packages/coding-agent/test/extensions-discovery.test.ts index af1aba52..6456a50d 100644 --- a/packages/coding-agent/test/extensions-discovery.test.ts +++ b/packages/coding-agent/test/extensions-discovery.test.ts @@ -51,6 +51,42 @@ describe("extensions discovery", () => { expect(result.extensions.map((e) => path.basename(e.path)).sort()).toEqual(["bar.ts", "foo.ts"]); }); + it("loads the coding-agent entrypoint without rewriting pi-ai provider subpaths", async () => { + fs.writeFileSync( + path.join(extensionsDir, "coding-agent-import.ts"), + ` + import { getAgentDir } from "@earendil-works/pi-coding-agent"; + void getAgentDir; + export default function(pi) { + pi.registerCommand("test", { handler: async () => {} }); + } + `, + ); + + const result = await discoverAndLoadExtensions([], tempDir, tempDir); + + expect(result.errors).toHaveLength(0); + expect(result.extensions).toHaveLength(1); + }); + + it("keeps the type-only pi-ai OAuth compatibility barrel resolvable", async () => { + fs.writeFileSync( + path.join(extensionsDir, "oauth-import.ts"), + ` + import * as oauth from "@earendil-works/pi-ai/oauth"; + void oauth; + export default function(pi) { + pi.registerCommand("test", { handler: async () => {} }); + } + `, + ); + + const result = await discoverAndLoadExtensions([], tempDir, tempDir); + + expect(result.errors).toEqual([]); + expect(result.extensions).toHaveLength(1); + }); + it("discovers direct .js files in extensions/", async () => { fs.writeFileSync(path.join(extensionsDir, "foo.js"), extensionCode); diff --git a/packages/coding-agent/test/extensions-input-event.test.ts b/packages/coding-agent/test/extensions-input-event.test.ts index 357fb6b8..f90f8490 100644 --- a/packages/coding-agent/test/extensions-input-event.test.ts +++ b/packages/coding-agent/test/extensions-input-event.test.ts @@ -5,9 +5,10 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; import { discoverAndLoadExtensions } from "../src/core/extensions/loader.ts"; import { ExtensionRunner } from "../src/core/extensions/runner.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { SessionManager } from "../src/core/session-manager.ts"; +import { createModelRegistry } from "./model-runtime-test-utils.ts"; + describe("Input Event", () => { let tempDir: string; let extensionsDir: string; @@ -29,7 +30,7 @@ describe("Input Event", () => { for (let i = 0; i < extensions.length; i++) fs.writeFileSync(path.join(extensionsDir, `e${i}.ts`), extensions[i]); const result = await discoverAndLoadExtensions([], tempDir, tempDir); const sm = SessionManager.inMemory(); - const mr = ModelRegistry.create(AuthStorage.create(path.join(tempDir, "auth.json"))); + const mr = await createModelRegistry(AuthStorage.create(path.join(tempDir, "auth.json"))); return new ExtensionRunner(result.extensions, result.runtime, tempDir, sm, mr); } diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index d4115a35..69f20be9 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -1,3 +1,4 @@ +import { createModelRegistry } from "./model-runtime-test-utils.ts"; /** * Tests for ExtensionRunner - conflict detection, error handling, tool wrapping. */ @@ -16,7 +17,7 @@ import type { ProviderConfig, } from "../src/core/extensions/types.ts"; import { KeybindingsManager, type KeyId } from "../src/core/keybindings.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; +import type { ModelRegistry } from "../src/core/model-registry.ts"; import { SessionManager } from "../src/core/session-manager.ts"; describe("ExtensionRunner", () => { @@ -26,13 +27,13 @@ describe("ExtensionRunner", () => { let modelRegistry: ModelRegistry; const defaultKeybindings = new KeybindingsManager().getEffectiveConfig(); - beforeEach(() => { + beforeEach(async () => { tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "pi-runner-test-")); extensionsDir = path.join(tempDir, "extensions"); fs.mkdirSync(extensionsDir); sessionManager = SessionManager.inMemory(); const authStorage = AuthStorage.create(path.join(tempDir, "auth.json")); - modelRegistry = ModelRegistry.create(authStorage); + modelRegistry = await createModelRegistry(authStorage); }); afterEach(() => { @@ -817,7 +818,7 @@ describe("ExtensionRunner", () => { }); describe("provider registration", () => { - it("bindCore ignores invalid queued registrations and reports extension error", () => { + it("bindCore ignores invalid queued registrations and reports extension error", async () => { const runtime = createExtensionRuntime(); runtime.registerProvider( "broken-provider", @@ -837,7 +838,7 @@ describe("ExtensionRunner", () => { expect(errors).toEqual([ '/tmp/broken-extension.ts: Provider broken-provider: "api" is required when registering streamSimple.', ]); - expect(() => modelRegistry.refresh()).not.toThrow(); + await expect(modelRegistry.refresh()).resolves.toBeUndefined(); }); it("pre-bind unregister removes all queued registrations for a provider", () => { diff --git a/packages/coding-agent/test/footer-width.test.ts b/packages/coding-agent/test/footer-width.test.ts index 51df1eda..c3d707a9 100644 --- a/packages/coding-agent/test/footer-width.test.ts +++ b/packages/coding-agent/test/footer-width.test.ts @@ -52,7 +52,7 @@ function createSession(options: { getCwd: () => "/tmp/project", }, getContextUsage: () => ({ contextWindow: 200_000, percent: 12.3 }), - modelRegistry: { + modelRuntime: { isUsingOAuth: () => false, }, }; diff --git a/packages/coding-agent/test/interactive-mode-anthropic-warning.test.ts b/packages/coding-agent/test/interactive-mode-anthropic-warning.test.ts index 10e35192..bcda13ae 100644 --- a/packages/coding-agent/test/interactive-mode-anthropic-warning.test.ts +++ b/packages/coding-agent/test/interactive-mode-anthropic-warning.test.ts @@ -7,19 +7,20 @@ function createSettingsManager(warnings: { anthropicExtraUsage?: boolean } = {}) }; } +function createModelRuntime(credential: { type: "oauth" } | undefined, apiKey?: string) { + return { + checkAuth: vi.fn().mockResolvedValue(credential), + getAuth: vi.fn().mockResolvedValue(apiKey ? { auth: { apiKey } } : undefined), + }; +} + describe("InteractiveMode.maybeWarnAboutAnthropicSubscriptionAuth", () => { test("warns once when Anthropic subscription auth is detected", async () => { + const modelRuntime = createModelRuntime(undefined, "sk-ant-oat01-test"); const fakeThis: any = { anthropicSubscriptionWarningShown: false, settingsManager: createSettingsManager(), - session: { - modelRegistry: { - authStorage: { - get: vi.fn().mockReturnValue(undefined), - }, - getApiKeyForProvider: vi.fn().mockResolvedValue("sk-ant-oat01-test"), - }, - }, + session: { modelRuntime }, showWarning: vi.fn(), }; @@ -31,21 +32,15 @@ describe("InteractiveMode.maybeWarnAboutAnthropicSubscriptionAuth", () => { }); expect(fakeThis.showWarning).toHaveBeenCalledTimes(1); - expect(fakeThis.session.modelRegistry.getApiKeyForProvider).toHaveBeenCalledTimes(1); + expect(modelRuntime.getAuth).toHaveBeenCalledTimes(1); }); test("warns when Anthropic OAuth is stored even if token refresh lookup would fail", async () => { + const modelRuntime = createModelRuntime({ type: "oauth" }); const fakeThis: any = { anthropicSubscriptionWarningShown: false, settingsManager: createSettingsManager(), - session: { - modelRegistry: { - authStorage: { - get: vi.fn().mockReturnValue({ type: "oauth" }), - }, - getApiKeyForProvider: vi.fn().mockResolvedValue(undefined), - }, - }, + session: { modelRuntime }, showWarning: vi.fn(), }; @@ -54,21 +49,15 @@ describe("InteractiveMode.maybeWarnAboutAnthropicSubscriptionAuth", () => { }); expect(fakeThis.showWarning).toHaveBeenCalledTimes(1); - expect(fakeThis.session.modelRegistry.getApiKeyForProvider).not.toHaveBeenCalled(); + expect(modelRuntime.getAuth).not.toHaveBeenCalled(); }); test("does not warn for non-Anthropic models", async () => { + const modelRuntime = createModelRuntime(undefined); const fakeThis: any = { anthropicSubscriptionWarningShown: false, settingsManager: createSettingsManager(), - session: { - modelRegistry: { - authStorage: { - get: vi.fn(), - }, - getApiKeyForProvider: vi.fn(), - }, - }, + session: { modelRuntime }, showWarning: vi.fn(), }; @@ -77,21 +66,15 @@ describe("InteractiveMode.maybeWarnAboutAnthropicSubscriptionAuth", () => { }); expect(fakeThis.showWarning).not.toHaveBeenCalled(); - expect(fakeThis.session.modelRegistry.getApiKeyForProvider).not.toHaveBeenCalled(); + expect(modelRuntime.getAuth).not.toHaveBeenCalled(); }); test("does not warn when Anthropic extra usage warning is disabled", async () => { + const modelRuntime = createModelRuntime(undefined); const fakeThis: any = { anthropicSubscriptionWarningShown: false, settingsManager: createSettingsManager({ anthropicExtraUsage: false }), - session: { - modelRegistry: { - authStorage: { - get: vi.fn(), - }, - getApiKeyForProvider: vi.fn(), - }, - }, + session: { modelRuntime }, showWarning: vi.fn(), }; @@ -100,7 +83,7 @@ describe("InteractiveMode.maybeWarnAboutAnthropicSubscriptionAuth", () => { }); expect(fakeThis.showWarning).not.toHaveBeenCalled(); - expect(fakeThis.session.modelRegistry.authStorage.get).not.toHaveBeenCalled(); - expect(fakeThis.session.modelRegistry.getApiKeyForProvider).not.toHaveBeenCalled(); + expect(modelRuntime.checkAuth).not.toHaveBeenCalled(); + expect(modelRuntime.getAuth).not.toHaveBeenCalled(); }); }); diff --git a/packages/coding-agent/test/interactive-mode-status.test.ts b/packages/coding-agent/test/interactive-mode-status.test.ts index c705d739..54c615bc 100644 --- a/packages/coding-agent/test/interactive-mode-status.test.ts +++ b/packages/coding-agent/test/interactive-mode-status.test.ts @@ -379,7 +379,7 @@ describe("InteractiveMode.createBaseAutocompleteProvider", () => { type FakeInteractiveMode = { session: { scopedModels: Array<{ model: TestModel }>; - modelRegistry: { getAvailable: () => TestModel[] }; + modelRuntime: { getAvailable: () => TestModel[] }; promptTemplates: []; extensionRunner: { getRegisteredCommands: () => [] }; resourceLoader: { getSkills: () => { skills: [] } }; @@ -402,7 +402,7 @@ describe("InteractiveMode.createBaseAutocompleteProvider", () => { const fakeThis: FakeInteractiveMode = { session: { scopedModels: [], - modelRegistry: { getAvailable: () => models }, + modelRuntime: { getAvailable: () => models }, promptTemplates: [], extensionRunner: { getRegisteredCommands: () => [] }, resourceLoader: { getSkills: () => ({ skills: [] }) }, @@ -429,7 +429,7 @@ describe("InteractiveMode.createBaseAutocompleteProvider", () => { type FakeInteractiveMode = { session: { scopedModels: []; - modelRegistry: { getAvailable: () => [] }; + modelRuntime: { getAvailable: () => [] }; promptTemplates: []; extensionRunner: { getRegisteredCommands: () => [] }; resourceLoader: { getSkills: () => { skills: [] } }; @@ -449,7 +449,7 @@ describe("InteractiveMode.createBaseAutocompleteProvider", () => { const fakeThis: FakeInteractiveMode = { session: { scopedModels: [], - modelRegistry: { getAvailable: () => [] }, + modelRuntime: { getAvailable: () => [] }, promptTemplates: [], extensionRunner: { getRegisteredCommands: () => [] }, resourceLoader: { getSkills: () => ({ skills: [] }) }, diff --git a/packages/coding-agent/test/model-registry.test.ts b/packages/coding-agent/test/model-registry.test.ts index 33eb2e00..ef4b133b 100644 --- a/packages/coding-agent/test/model-registry.test.ts +++ b/packages/coding-agent/test/model-registry.test.ts @@ -9,10 +9,11 @@ import type { OpenAICompletionsCompat, } from "@earendil-works/pi-ai/compat"; import { getApiProvider, getSupportedThinkingLevels } from "@earendil-works/pi-ai/compat"; -import { getOAuthProvider } from "@earendil-works/pi-ai/oauth"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { clearApiKeyCache, ModelRegistry, type ProviderConfigInput } from "../src/core/model-registry.ts"; +import { clearApiKeyCache, type ModelRegistry, type ProviderConfigInput } from "../src/core/model-registry.ts"; + +import { createModelRegistry } from "./model-runtime-test-utils.ts"; describe("ModelRegistry", () => { let tempDir: string; @@ -96,12 +97,12 @@ describe("ModelRegistry", () => { }; describe("baseUrl override (no custom models)", () => { - test("overriding baseUrl keeps all built-in models", () => { + test("overriding baseUrl keeps all built-in models", async () => { writeRawModelsJson({ anthropic: overrideConfig("https://my-proxy.example.com/v1"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const anthropicModels = getModelsForProvider(registry, "anthropic"); // Should have multiple built-in models, not just one @@ -109,12 +110,12 @@ describe("ModelRegistry", () => { expect(anthropicModels.some((m) => m.id.includes("claude"))).toBe(true); }); - test("overriding baseUrl changes URL on all built-in models", () => { + test("overriding baseUrl changes URL on all built-in models", async () => { writeRawModelsJson({ anthropic: overrideConfig("https://my-proxy.example.com/v1"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const anthropicModels = getModelsForProvider(registry, "anthropic"); // All models should have the new baseUrl @@ -130,7 +131,7 @@ describe("ModelRegistry", () => { }), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const anthropicModels = getModelsForProvider(registry, "anthropic"); for (const model of anthropicModels) { @@ -151,7 +152,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(registry.getError()).toBeUndefined(); const anthropicModels = getModelsForProvider(registry, "anthropic"); @@ -164,12 +165,26 @@ describe("ModelRegistry", () => { } }); - test("baseUrl-only override does not affect other providers", () => { + test("unconfigured compatibility auth includes static model headers", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); + const base = registry.getAll()[0]; + const model = { + ...base, + provider: "missing-provider", + headers: { "X-Static-Model": "static-value" }, + }; + + const auth = await registry.getApiKeyAndHeaders(model); + + expect(auth).toEqual({ ok: true, headers: { "X-Static-Model": "static-value" } }); + }); + + test("baseUrl-only override does not affect other providers", async () => { writeRawModelsJson({ anthropic: overrideConfig("https://my-proxy.example.com/v1"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const googleModels = getModelsForProvider(registry, "google"); // Google models should still have their original baseUrl @@ -177,7 +192,7 @@ describe("ModelRegistry", () => { expect(googleModels[0].baseUrl).not.toBe("https://my-proxy.example.com/v1"); }); - test("can mix baseUrl override and models merge", () => { + test("can mix baseUrl override and models merge", async () => { writeRawModelsJson({ // baseUrl-only for anthropic anthropic: overrideConfig("https://anthropic-proxy.example.com/v1"), @@ -189,7 +204,7 @@ describe("ModelRegistry", () => { ), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); // Anthropic: multiple built-in models with new baseUrl const anthropicModels = getModelsForProvider(registry, "anthropic"); @@ -202,11 +217,11 @@ describe("ModelRegistry", () => { expect(googleModels.some((m) => m.id === "gemini-custom")).toBe(true); }); - test("refresh() picks up baseUrl override changes", () => { + test("refresh() picks up baseUrl override changes", async () => { writeRawModelsJson({ anthropic: overrideConfig("https://first-proxy.example.com/v1"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(getModelsForProvider(registry, "anthropic")[0].baseUrl).toBe("https://first-proxy.example.com/v1"); @@ -214,14 +229,14 @@ describe("ModelRegistry", () => { writeRawModelsJson({ anthropic: overrideConfig("https://second-proxy.example.com/v1"), }); - registry.refresh(); + await registry.refresh(); expect(getModelsForProvider(registry, "anthropic")[0].baseUrl).toBe("https://second-proxy.example.com/v1"); }); }); describe("custom models merge behavior", () => { - test("built-in provider custom models inherit api and baseUrl without explicit fields", () => { + test("built-in provider custom models inherit api and baseUrl without explicit fields", async () => { // Built-in providers already have api/baseUrl on every model, and auth // comes from env vars / auth storage. No need to specify them. writeRawModelsJson({ @@ -237,7 +252,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(registry.getError()).toBeUndefined(); const model = registry.find("openrouter", "fake-provider/fake-model"); @@ -246,7 +261,7 @@ describe("ModelRegistry", () => { expect(model?.baseUrl).toBe("https://openrouter.ai/api/v1"); }); - test("non-built-in provider custom models still require baseUrl", () => { + test("non-built-in provider custom models still require baseUrl", async () => { writeRawModelsJson({ "my-custom-provider": { apiKey: "test-key", @@ -261,16 +276,29 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(registry.getError()).toContain("baseUrl"); }); - test("custom provider with same name as built-in merges with built-in models", () => { + test("reports every provider composition error", async () => { + writeRawModelsJson({ + "broken-one": { api: "openai-completions", models: [{ id: "one" }] }, + "broken-two": { api: "openai-completions", models: [{ id: "two" }] }, + }); + + const registry = await createModelRegistry(authStorage, modelsJsonPath); + const error = registry.getError(); + + expect(error).toContain('Provider "broken-one"'); + expect(error).toContain('Provider "broken-two"'); + }); + + test("custom provider with same name as built-in merges with built-in models", async () => { writeModelsJson({ anthropic: providerConfig("https://my-proxy.example.com/v1", [{ id: "claude-custom" }]), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const anthropicModels = getModelsForProvider(registry, "anthropic"); expect(anthropicModels.length).toBeGreaterThan(1); @@ -278,7 +306,7 @@ describe("ModelRegistry", () => { expect(anthropicModels.some((m) => m.id.includes("claude"))).toBe(true); }); - test("custom model with same id replaces built-in model by id", () => { + test("custom model with same id replaces built-in model by id", async () => { writeModelsJson({ openrouter: providerConfig( "https://my-proxy.example.com/v1", @@ -287,7 +315,7 @@ describe("ModelRegistry", () => { ), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnetModels = models.filter((m) => m.id === "anthropic/claude-sonnet-4"); @@ -295,23 +323,23 @@ describe("ModelRegistry", () => { expect(sonnetModels[0].baseUrl).toBe("https://my-proxy.example.com/v1"); }); - test("custom provider with same name as built-in does not affect other built-in providers", () => { + test("custom provider with same name as built-in does not affect other built-in providers", async () => { writeModelsJson({ anthropic: providerConfig("https://my-proxy.example.com/v1", [{ id: "claude-custom" }]), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(getModelsForProvider(registry, "google").length).toBeGreaterThan(0); expect(getModelsForProvider(registry, "openai").length).toBeGreaterThan(0); }); - test("provider-level baseUrl applies to both built-in and custom models", () => { + test("provider-level baseUrl applies to both built-in and custom models", async () => { writeModelsJson({ anthropic: providerConfig("https://merged-proxy.example.com/v1", [{ id: "claude-custom" }]), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const anthropicModels = getModelsForProvider(registry, "anthropic"); for (const model of anthropicModels) { @@ -319,7 +347,7 @@ describe("ModelRegistry", () => { } }); - test("provider-level compat applies to custom models", () => { + test("provider-level compat applies to custom models", async () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com/v1", @@ -342,14 +370,14 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const compat = registry.find("demo", "demo-model")?.compat as OpenAICompletionsCompat | undefined; expect(compat?.supportsUsageInStreaming).toBe(false); expect(compat?.maxTokensField).toBe("max_tokens"); }); - test("model-level compat overrides provider-level compat for custom models", () => { + test("model-level compat overrides provider-level compat for custom models", async () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com/v1", @@ -376,14 +404,14 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const compat = registry.find("demo", "demo-model")?.compat as OpenAICompletionsCompat | undefined; expect(compat?.supportsUsageInStreaming).toBe(true); expect(compat?.maxTokensField).toBe("max_completion_tokens"); }); - test("provider-level compat applies to built-in models", () => { + test("provider-level compat applies to built-in models", async () => { writeRawModelsJson({ openrouter: { compat: { @@ -393,7 +421,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); expect(models.length).toBeGreaterThan(0); @@ -404,7 +432,7 @@ describe("ModelRegistry", () => { } }); - test("model schema accepts thinkingLevelMap and compat schema accepts supportsStrictMode and cacheControlFormat", () => { + test("model schema accepts thinkingLevelMap and compat schema accepts supportsStrictMode and cacheControlFormat", async () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com/v1", @@ -431,7 +459,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const model = registry.find("demo", "demo-model"); const compat = model?.compat as OpenAICompletionsCompat | undefined; @@ -441,7 +469,7 @@ describe("ModelRegistry", () => { expect(compat?.cacheControlFormat).toBe("anthropic"); }); - test("compat schema accepts chat template thinking configuration", () => { + test("compat schema accepts chat template thinking configuration", async () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com/v1", @@ -467,7 +495,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const compat = registry.find("demo", "demo-model")?.compat as OpenAICompletionsCompat | undefined; expect(registry.getError()).toBeUndefined(); @@ -478,7 +506,7 @@ describe("ModelRegistry", () => { }); }); - test("compat schema accepts Anthropic eager tool input streaming flag", () => { + test("compat schema accepts Anthropic eager tool input streaming flag", async () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com", @@ -500,14 +528,14 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const compat = registry.find("demo", "demo-model")?.compat as AnthropicMessagesCompat | undefined; expect(registry.getError()).toBeUndefined(); expect(compat?.supportsEagerToolInputStreaming).toBe(false); }); - test("compat schema accepts long cache retention flag", () => { + test("compat schema accepts long cache retention flag", async () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com", @@ -529,14 +557,14 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const compat = registry.find("demo", "demo-model")?.compat as AnthropicMessagesCompat | undefined; expect(registry.getError()).toBeUndefined(); expect(compat?.supportsLongCacheRetention).toBe(false); }); - test("model-level baseUrl overrides provider-level baseUrl for custom models", () => { + test("model-level baseUrl overrides provider-level baseUrl for custom models", async () => { writeRawModelsJson({ "opencode-go": { baseUrl: "https://opencode.ai/zen/go/v1", @@ -565,7 +593,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const m25 = registry.find("opencode-go", "minimax-m2.5"); const glm5 = registry.find("opencode-go", "glm-5"); @@ -573,7 +601,7 @@ describe("ModelRegistry", () => { expect(glm5?.baseUrl).toBe("https://opencode.ai/zen/go/v1"); }); - test("modelOverrides still apply when provider also defines models", () => { + test("modelOverrides still apply when provider also defines models", async () => { writeRawModelsJson({ openrouter: { baseUrl: "https://my-proxy.example.com/v1", @@ -598,7 +626,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); expect(models.some((m) => m.id === "custom/openrouter-model")).toBe(true); @@ -607,18 +635,18 @@ describe("ModelRegistry", () => { ).toBe(true); }); - test("refresh() reloads merged custom models from disk", () => { + test("refresh() reloads merged custom models from disk", async () => { writeModelsJson({ anthropic: providerConfig("https://first-proxy.example.com/v1", [{ id: "claude-custom" }]), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(getModelsForProvider(registry, "anthropic").some((m) => m.id === "claude-custom")).toBe(true); // Update and refresh writeModelsJson({ anthropic: providerConfig("https://second-proxy.example.com/v1", [{ id: "claude-custom-2" }]), }); - registry.refresh(); + await registry.refresh(); const anthropicModels = getModelsForProvider(registry, "anthropic"); expect(anthropicModels.some((m) => m.id === "claude-custom")).toBe(false); @@ -626,16 +654,16 @@ describe("ModelRegistry", () => { expect(anthropicModels.some((m) => m.id.includes("claude"))).toBe(true); }); - test("removing custom models from models.json keeps built-in provider models", () => { + test("removing custom models from models.json keeps built-in provider models", async () => { writeModelsJson({ anthropic: providerConfig("https://proxy.example.com/v1", [{ id: "claude-custom" }]), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(getModelsForProvider(registry, "anthropic").some((m) => m.id === "claude-custom")).toBe(true); // Remove custom models and refresh writeModelsJson({}); - registry.refresh(); + await registry.refresh(); const anthropicModels = getModelsForProvider(registry, "anthropic"); expect(anthropicModels.some((m) => m.id === "claude-custom")).toBe(false); @@ -644,7 +672,7 @@ describe("ModelRegistry", () => { }); describe("modelOverrides (per-model customization)", () => { - test("model override applies to a single built-in model", () => { + test("model override applies to a single built-in model", async () => { writeRawModelsJson({ openrouter: { modelOverrides: { @@ -655,7 +683,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); @@ -666,7 +694,7 @@ describe("ModelRegistry", () => { expect(opus?.name).not.toBe("Custom Sonnet Name"); }); - test("model override with compat.openRouterRouting", () => { + test("model override with compat.openRouterRouting", async () => { writeRawModelsJson({ openrouter: { modelOverrides: { @@ -679,7 +707,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); @@ -687,7 +715,7 @@ describe("ModelRegistry", () => { expect(compat?.openRouterRouting).toEqual({ only: ["amazon-bedrock"] }); }); - test("model override deep merges compat settings", () => { + test("model override deep merges compat settings", async () => { writeRawModelsJson({ openrouter: { modelOverrides: { @@ -700,7 +728,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); @@ -709,7 +737,7 @@ describe("ModelRegistry", () => { expect(compat?.openRouterRouting).toEqual({ order: ["anthropic", "together"] }); }); - test("multiple model overrides on same provider", () => { + test("multiple model overrides on same provider", async () => { writeRawModelsJson({ openrouter: { modelOverrides: { @@ -723,7 +751,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); @@ -735,7 +763,7 @@ describe("ModelRegistry", () => { expect(opusCompat?.openRouterRouting).toEqual({ only: ["anthropic"] }); }); - test("model override combined with baseUrl override", () => { + test("model override combined with baseUrl override", async () => { writeRawModelsJson({ openrouter: { baseUrl: "https://my-proxy.example.com/v1", @@ -747,7 +775,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); @@ -761,7 +789,7 @@ describe("ModelRegistry", () => { expect(opus?.name).not.toBe("Proxied Sonnet"); }); - test("model override for non-existent model ID is ignored", () => { + test("model override for non-existent model ID is ignored", async () => { writeRawModelsJson({ openrouter: { modelOverrides: { @@ -772,7 +800,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); // Should not create a new model @@ -781,7 +809,7 @@ describe("ModelRegistry", () => { expect(registry.getError()).toBeUndefined(); }); - test("model override can change cost fields partially", () => { + test("model override can change cost fields partially", async () => { writeRawModelsJson({ openrouter: { modelOverrides: { @@ -792,7 +820,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); @@ -813,7 +841,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); expect(sonnet).toBeDefined(); @@ -825,7 +853,7 @@ describe("ModelRegistry", () => { } }); - test("refresh() picks up model override changes", () => { + test("refresh() picks up model override changes", async () => { writeRawModelsJson({ openrouter: { modelOverrides: { @@ -836,7 +864,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect( getModelsForProvider(registry, "openrouter").find((m) => m.id === "anthropic/claude-sonnet-4")?.name, ).toBe("First Name"); @@ -851,14 +879,14 @@ describe("ModelRegistry", () => { }, }, }); - registry.refresh(); + await registry.refresh(); expect( getModelsForProvider(registry, "openrouter").find((m) => m.id === "anthropic/claude-sonnet-4")?.name, ).toBe("Second Name"); }); - test("removing model override restores built-in values", () => { + test("removing model override restores built-in values", async () => { writeRawModelsJson({ openrouter: { modelOverrides: { @@ -869,7 +897,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const customName = getModelsForProvider(registry, "openrouter").find( (m) => m.id === "anthropic/claude-sonnet-4", )?.name; @@ -877,7 +905,7 @@ describe("ModelRegistry", () => { // Remove override and refresh writeRawModelsJson({}); - registry.refresh(); + await registry.refresh(); const restoredName = getModelsForProvider(registry, "openrouter").find( (m) => m.id === "anthropic/claude-sonnet-4", @@ -887,12 +915,12 @@ describe("ModelRegistry", () => { }); describe("dynamic provider lifecycle", () => { - test("getProviderDisplayName resolves registered, OAuth, built-in, and fallback names", () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + test("getProviderDisplayName resolves registered, OAuth, built-in, and fallback names", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(registry.getProviderDisplayName("openai")).toBe("OpenAI"); expect(registry.getProviderDisplayName("github-copilot")).toBe("GitHub Copilot"); - expect(registry.getProviderDisplayName("zai")).toBe("ZAI Coding Plan (Global)"); + expect(registry.getProviderDisplayName("zai")).toBe("Z.AI"); expect(registry.getProviderDisplayName("unknown-provider")).toBe("unknown-provider"); registry.registerProvider("named-provider", { @@ -957,7 +985,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider("extension-provider", { baseUrl: "https://provider.test/v1", apiKey: "test-key", @@ -996,21 +1024,22 @@ describe("ModelRegistry", () => { }); test("stored API key env propagates to request auth and resolves headers", async () => { - authStorage.set("cloudflare-ai-gateway", { + await authStorage.modify("cloudflare-ai-gateway", async () => ({ type: "api_key", key: "$CLOUDFLARE_API_KEY", env: { CLOUDFLARE_API_KEY: "stored-cf-token", CLOUDFLARE_ACCOUNT_ID: "stored-account", + CLOUDFLARE_GATEWAY_ID: "stored-gateway", }, - }); + })); writeRawModelsJson({ "cloudflare-ai-gateway": { headers: { "x-account": "$CLOUDFLARE_ACCOUNT_ID" }, }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const model = registry.getAll().find((m) => m.provider === "cloudflare-ai-gateway"); expect(model).toBeDefined(); @@ -1018,11 +1047,14 @@ describe("ModelRegistry", () => { expect(auth).toEqual({ ok: true, - apiKey: "stored-cf-token", - headers: { "x-account": "stored-account" }, + apiKey: undefined, + headers: { + "cf-aig-authorization": "Bearer stored-cf-token", + "x-account": "stored-account", + }, env: { - CLOUDFLARE_API_KEY: "stored-cf-token", CLOUDFLARE_ACCOUNT_ID: "stored-account", + CLOUDFLARE_GATEWAY_ID: "stored-gateway", }, }); }); @@ -1037,7 +1069,7 @@ describe("ModelRegistry", () => { const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); try { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider("literal-provider", { ...providerConfig("https://provider.test/v1", [{ id: "demo-model" }], "openai-completions"), @@ -1080,8 +1112,8 @@ describe("ModelRegistry", () => { } }); - test("failed registerProvider does not persist invalid streamSimple config", () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + test("failed registerProvider does not persist invalid streamSimple config", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(() => registry.registerProvider("broken-provider", { @@ -1091,11 +1123,11 @@ describe("ModelRegistry", () => { }), ).toThrow('Provider broken-provider: "api" is required when registering streamSimple.'); - expect(() => registry.refresh()).not.toThrow(); + await expect(registry.refresh()).resolves.toBeUndefined(); }); - test("failed registerProvider does not remove existing provider models", () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + test("failed registerProvider does not remove existing provider models", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider("demo-provider", { baseUrl: "https://provider.test/v1", @@ -1135,12 +1167,12 @@ describe("ModelRegistry", () => { ).toThrow('Provider demo-provider, model broken-model: no "api" specified.'); expect(registry.find("demo-provider", "demo-model")).toBeDefined(); - expect(() => registry.refresh()).not.toThrow(); + await expect(registry.refresh()).resolves.toBeUndefined(); expect(registry.find("demo-provider", "demo-model")).toBeDefined(); }); - test("unregisterProvider removes custom OAuth provider and restores built-in OAuth provider", () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + test("unregisterProvider removes the runtime OAuth overlay without mutating global state", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider("anthropic", { oauth: { @@ -1155,15 +1187,15 @@ describe("ModelRegistry", () => { }, }); - expect(getOAuthProvider("anthropic")?.name).toBe("Custom Anthropic OAuth"); + expect(registry.getRegisteredProviderConfig("anthropic")?.oauth?.name).toBe("Custom Anthropic OAuth"); registry.unregisterProvider("anthropic"); - expect(getOAuthProvider("anthropic")?.name).not.toBe("Custom Anthropic OAuth"); + expect(registry.getRegisteredProviderConfig("anthropic")).toBeUndefined(); }); - test("unregisterProvider removes custom streamSimple override and restores built-in API stream handler", () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + test("streamSimple overlays do not mutate the global compat API registry", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider("stream-override-provider", { api: "openai-completions", @@ -1178,7 +1210,7 @@ describe("ModelRegistry", () => { } catch (error) { threwCustomOverride = error instanceof Error && error.message === "custom streamSimple override"; } - expect(threwCustomOverride).toBe(true); + expect(threwCustomOverride).toBe(false); registry.unregisterProvider("stream-override-provider"); @@ -1193,52 +1225,52 @@ describe("ModelRegistry", () => { }); describe("dynamic provider override persistence", () => { - test("baseUrl-only override keeps built-in provider models after refresh", () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + test("baseUrl-only override keeps built-in provider models after refresh", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider("anthropic", { baseUrl: "https://proxy.test/anthropic" }); - registry.refresh(); + await registry.refresh(); const anthropicModels = getModelsForProvider(registry, "anthropic"); expect(anthropicModels.length).toBeGreaterThan(1); expect(anthropicModels.every((m) => m.baseUrl === "https://proxy.test/anthropic")).toBe(true); }); - test("models-only override replaces built-in provider models after refresh", () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + test("models-only override replaces built-in provider models after refresh", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider("anthropic", { ...providerConfig("https://custom.test/anthropic", [{ id: "custom-claude" }], "anthropic-messages"), baseUrl: "https://custom.test/anthropic", }); - registry.refresh(); + await registry.refresh(); expect(getModelsForProvider(registry, "anthropic").map((m) => m.id)).toEqual(["custom-claude"]); expect(registry.find("anthropic", "custom-claude")?.baseUrl).toBe("https://custom.test/anthropic"); }); - test("models plus baseUrl override replaces built-in provider models after refresh", () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + test("models plus baseUrl override replaces built-in provider models after refresh", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider("anthropic", { ...providerConfig("https://custom.test/anthropic", [{ id: "custom-claude" }], "anthropic-messages"), baseUrl: "https://custom.test/anthropic", }); registry.registerProvider("anthropic", { baseUrl: "https://proxy.test/anthropic" }); - registry.refresh(); + await registry.refresh(); expect(getModelsForProvider(registry, "anthropic").map((m) => m.id)).toEqual(["custom-claude"]); expect(registry.find("anthropic", "custom-claude")?.baseUrl).toBe("https://proxy.test/anthropic"); }); - test("models-only custom provider registration survives refresh", () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + test("models-only custom provider registration survives refresh", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider( "custom-provider", providerConfig("https://custom.test/v1", [{ id: "custom-a" }, { id: "custom-b" }], "openai-completions"), ); - registry.refresh(); + await registry.refresh(); expect(getModelsForProvider(registry, "custom-provider").map((m) => m.id)).toEqual([ "custom-a", @@ -1246,15 +1278,15 @@ describe("ModelRegistry", () => { ]); }); - test("baseUrl-only override keeps custom provider models after refresh", () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + test("baseUrl-only override keeps custom provider models after refresh", async () => { + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider( "custom-provider", providerConfig("https://custom.test/v1", [{ id: "custom-a" }, { id: "custom-b" }], "openai-completions"), ); registry.registerProvider("custom-provider", { baseUrl: "https://proxy.test/custom" }); - registry.refresh(); + await registry.refresh(); expect(getModelsForProvider(registry, "custom-provider").map((m) => m.id)).toEqual([ "custom-a", @@ -1268,14 +1300,14 @@ describe("ModelRegistry", () => { }); test("headers-only override keeps custom provider models after refresh", async () => { - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); registry.registerProvider( "custom-provider", providerConfig("https://custom.test/v1", [{ id: "custom-a" }, { id: "custom-b" }], "openai-completions"), ); registry.registerProvider("custom-provider", { headers: { "x-proxy": "enabled" } }); - registry.refresh(); + await registry.refresh(); const models = getModelsForProvider(registry, "custom-provider"); expect(models.map((m) => m.id)).toEqual(["custom-a", "custom-b"]); @@ -1314,7 +1346,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("!echo test-api-key-from-command"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("test-api-key-from-command"); @@ -1325,7 +1357,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("!echo ' spaced-key '"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("spaced-key"); @@ -1336,7 +1368,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("!printf 'line1\\nline2'"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("line1\nline2"); @@ -1347,7 +1379,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("!exit 1"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBeUndefined(); @@ -1358,7 +1390,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("!nonexistent-command-12345"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBeUndefined(); @@ -1369,7 +1401,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("!printf ''"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBeUndefined(); @@ -1384,7 +1416,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("$TEST_API_KEY_12345"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("env-api-key-value"); @@ -1407,7 +1439,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(bracedKey), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("braced-env-api-key-value"); @@ -1434,7 +1466,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(interpolatedKey), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("left_right"); @@ -1457,7 +1489,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("$$TEST_API_KEY_12345"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("$TEST_API_KEY_12345"); @@ -1472,7 +1504,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("$!literal-$TEST_API_KEY_12345"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("!literal-env-api-key-value"); @@ -1494,7 +1526,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("TEST_API_KEY_12345"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("TEST_API_KEY_12345"); @@ -1515,7 +1547,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("literal_api_key_value"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("literal_api_key_value"); @@ -1526,7 +1558,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey("!echo 'hello world' | tr ' ' '-'"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider"); expect(apiKey).toBe("hello-world"); @@ -1543,7 +1575,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(command), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); await registry.getApiKeyForProvider("custom-provider"); await registry.getApiKeyForProvider("custom-provider"); await registry.getApiKeyForProvider("custom-provider"); @@ -1562,10 +1594,10 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(command), }); - const registry1 = ModelRegistry.create(authStorage, modelsJsonPath); + const registry1 = await createModelRegistry(authStorage, modelsJsonPath); await registry1.getApiKeyForProvider("custom-provider"); - const registry2 = ModelRegistry.create(authStorage, modelsJsonPath); + const registry2 = await createModelRegistry(authStorage, modelsJsonPath); await registry2.getApiKeyForProvider("custom-provider"); const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); @@ -1578,7 +1610,7 @@ describe("ModelRegistry", () => { "provider-b": providerWithApiKey("!echo key-b"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const keyA = await registry.getApiKeyForProvider("provider-a"); const keyB = await registry.getApiKeyForProvider("provider-b"); @@ -1597,7 +1629,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(command), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const key1 = await registry.getApiKeyForProvider("custom-provider"); const key2 = await registry.getApiKeyForProvider("custom-provider"); @@ -1608,7 +1640,7 @@ describe("ModelRegistry", () => { expect(count).toBe(2); }); - test("provider auth status reports apiKey environment variables from models.json", () => { + test("provider auth status reports apiKey environment variables from models.json", async () => { const envVarName = "TEST_API_KEY_STATUS_TEST_98765"; const originalEnv = process.env[envVarName]; @@ -1619,7 +1651,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(`$${envVarName}`), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: true, @@ -1635,7 +1667,7 @@ describe("ModelRegistry", () => { } }); - test("provider auth status reports interpolated apiKey environment variables", () => { + test("provider auth status reports interpolated apiKey environment variables", async () => { const envVarNameA = "TEST_API_KEY_STATUS_PART_A_98765"; const envVarNameB = "TEST_API_KEY_STATUS_PART_B_98765"; const originalEnvA = process.env[envVarNameA]; @@ -1649,7 +1681,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(interpolatedKey), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: true, @@ -1670,12 +1702,12 @@ describe("ModelRegistry", () => { } }); - test("provider auth status reports non-env apiKey values from models.json as a config key", () => { + test("provider auth status reports non-env apiKey values from models.json as a config key", async () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("literal_api_key_value"), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: true, @@ -1683,7 +1715,7 @@ describe("ModelRegistry", () => { }); }); - test("missing explicit env apiKey keeps provider unavailable", () => { + test("missing explicit env apiKey keeps provider unavailable", async () => { const envVarName = "TEST_API_KEY_MISSING_TEST_98765"; const originalEnv = process.env[envVarName]; delete process.env[envVarName]; @@ -1693,7 +1725,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(`$${envVarName}`), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: false }); expect(registry.getAvailable().some((model) => model.provider === "custom-provider")).toBe(false); @@ -1706,7 +1738,7 @@ describe("ModelRegistry", () => { } }); - test("provider auth status reports command apiKey values from models.json without executing them", () => { + test("provider auth status reports command apiKey values from models.json without executing them", async () => { const counterFile = join(tempDir, "status-counter"); writeFileSync(counterFile, "0"); const counterPath = toShPath(counterFile); @@ -1715,7 +1747,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(command), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: true, @@ -1735,7 +1767,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(`$${envVarName}`), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const key1 = await registry.getApiKeyForProvider("custom-provider"); expect(key1).toBe("first-value"); @@ -1763,7 +1795,7 @@ describe("ModelRegistry", () => { "custom-provider": providerWithApiKey(command), }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const available = registry.getAvailable(); expect(available.some((m) => m.provider === "custom-provider")).toBe(true); @@ -1771,16 +1803,16 @@ describe("ModelRegistry", () => { expect(count).toBe(0); }); - test("getAvailable filters GitHub Copilot OAuth models to account picker availability", () => { - authStorage.set("github-copilot", { + test("getAvailable filters GitHub Copilot OAuth models to account picker availability", async () => { + await authStorage.modify("github-copilot", async () => ({ type: "oauth", refresh: "github-access-token", access: "tid=test;exp=9999999999;proxy-ep=proxy.individual.githubcopilot.com;", expires: Date.now() + 60_000, availableModelIds: ["gpt-4.1"], - }); + })); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); expect( registry @@ -1802,7 +1834,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const model = registry.find("custom-provider", "test-model"); expect(model).toBeDefined(); @@ -1823,6 +1855,62 @@ describe("ModelRegistry", () => { }); }); + test("getApiKeyAndHeaders resolves configured auth exactly once", async () => { + const counterFile = join(tempDir, "auth-counter"); + writeFileSync(counterFile, "0"); + const counterPath = toShPath(counterFile); + writeRawModelsJson({ + "custom-provider": { + ...providerWithApiKey( + `!sh -c 'count=$(cat "${counterPath}"); count=$((count + 1)); echo "$count" > "${counterPath}"; echo "token-$count"'`, + ), + authHeader: true, + }, + }); + + const registry = await createModelRegistry(authStorage, modelsJsonPath); + const auth = await registry.getApiKeyAndHeaders(registry.find("custom-provider", "test-model")!); + + expect(auth).toEqual({ + ok: true, + apiKey: "token-1", + headers: { Authorization: "Bearer token-1" }, + }); + expect(readFileSync(counterFile, "utf-8").trim()).toBe("1"); + }); + + test("stored credentials bypass lower-priority configured auth commands", async () => { + const counterFile = join(tempDir, "fallback-counter"); + writeFileSync(counterFile, "0"); + const counterPath = toShPath(counterFile); + writeRawModelsJson({ + "custom-provider": providerWithApiKey(`!sh -c 'echo 1 > "${counterPath}"; echo fallback-key'`), + }); + await authStorage.modify("custom-provider", async () => ({ type: "api_key", key: "stored-key" })); + + const registry = await createModelRegistry(authStorage, modelsJsonPath); + const auth = await registry.getApiKeyAndHeaders(registry.find("custom-provider", "test-model")!); + + expect(auth).toMatchObject({ ok: true, apiKey: "stored-key" }); + expect(readFileSync(counterFile, "utf-8").trim()).toBe("0"); + }); + + test("getApiKeyAndHeaders preserves the legacy missing-key authHeader error", async () => { + writeRawModelsJson({ + "custom-provider": { + baseUrl: "https://example.test/v1", + api: "openai-completions", + authHeader: true, + models: [{ id: "test-model" }], + }, + }); + + const registry = await createModelRegistry(authStorage, modelsJsonPath); + const auth = await registry.getApiKeyAndHeaders(registry.find("custom-provider", "test-model")!); + + expect(auth).toEqual({ ok: false, error: 'No API key found for "custom-provider"' }); + }); + test("getApiKeyAndHeaders returns an error for failed authHeader resolution", async () => { writeRawModelsJson({ "custom-provider": { @@ -1831,7 +1919,7 @@ describe("ModelRegistry", () => { }, }); - const registry = ModelRegistry.create(authStorage, modelsJsonPath); + const registry = await createModelRegistry(authStorage, modelsJsonPath); const model = registry.find("custom-provider", "test-model"); expect(model).toBeDefined(); diff --git a/packages/coding-agent/test/model-resolver.test.ts b/packages/coding-agent/test/model-resolver.test.ts index 24575829..22dca860 100644 --- a/packages/coding-agent/test/model-resolver.test.ts +++ b/packages/coding-agent/test/model-resolver.test.ts @@ -260,12 +260,12 @@ describe("resolveModelScopeWithDiagnostics", () => { describe("resolveCliModel", () => { test("resolves --model provider/id without --provider", () => { const registry = { - getAll: () => allModels, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => allModels, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliModel: "openai/gpt-4o", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -275,13 +275,13 @@ describe("resolveCliModel", () => { test("resolves fuzzy patterns within an explicit provider", () => { const registry = { - getAll: () => allModels, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => allModels, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliProvider: "openai", cliModel: "4o", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -291,12 +291,12 @@ describe("resolveCliModel", () => { test("supports --model : (without explicit --thinking)", () => { const registry = { - getAll: () => allModels, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => allModels, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliModel: "sonnet:high", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -306,12 +306,12 @@ describe("resolveCliModel", () => { test("prefers exact model id match over provider inference (OpenRouter-style ids)", () => { const registry = { - getAll: () => allModels, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => allModels, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliModel: "openai/gpt-4o:extended", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -321,13 +321,13 @@ describe("resolveCliModel", () => { test("does not strip invalid :suffix as thinking level in --model (treat as raw id)", () => { const registry = { - getAll: () => allModels, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => allModels, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliProvider: "openai", cliModel: "gpt-4o:extended", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -337,13 +337,13 @@ describe("resolveCliModel", () => { test("allows custom model ids for explicit providers without double prefixing", () => { const registry = { - getAll: () => allModels, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => allModels, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliProvider: "openrouter", cliModel: "openrouter/openai/ghost-model", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -353,13 +353,13 @@ describe("resolveCliModel", () => { test("returns a clear error when there are no models", () => { const registry = { - getAll: () => [], - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => [], + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliProvider: "openai", cliModel: "gpt-4o", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.model).toBeUndefined(); @@ -394,13 +394,13 @@ describe("resolveCliModel", () => { maxTokens: 8192, }; const registry = { - getAll: () => [...allModels, zaiModel, gatewayModel], + getModels: () => [...allModels, zaiModel, gatewayModel], hasConfiguredAuth: () => true, - } as unknown as Parameters[0]["modelRegistry"]; + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliModel: "zai/glm-5", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -434,13 +434,13 @@ describe("resolveCliModel", () => { maxTokens: 8192, }; const registry = { - getAll: () => [...allModels, commandcodeModel, xiaomiModel], - hasConfiguredAuth: (model: Model<"anthropic-messages">) => model.provider === "commandcode", - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => [...allModels, commandcodeModel, xiaomiModel], + hasConfiguredAuth: (provider: string) => provider === "commandcode", + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliModel: "xiaomi/mimo-v2.5-pro", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -450,12 +450,12 @@ describe("resolveCliModel", () => { test("resolves provider-prefixed fuzzy patterns (openrouter/qwen -> openrouter model)", () => { const registry = { - getAll: () => allModels, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => allModels, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliModel: "openrouter/qwen", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -483,12 +483,12 @@ describe("resolveCliModel", () => { test("strips :thinking suffix from custom model id in fallback path", () => { const registry = { - getAll: () => modelsWithNeuralwatt, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => modelsWithNeuralwatt, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliModel: "neuralwatt/zai-org/GLM-5.1-FP8:high", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -501,12 +501,12 @@ describe("resolveCliModel", () => { test("custom model without thinking suffix works normally in fallback path", () => { const registry = { - getAll: () => modelsWithNeuralwatt, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => modelsWithNeuralwatt, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliModel: "neuralwatt/zai-org/GLM-5.1-FP8", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -517,13 +517,13 @@ describe("resolveCliModel", () => { test("all valid thinking levels work in fallback path", () => { const registry = { - getAll: () => modelsWithNeuralwatt, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => modelsWithNeuralwatt, + } as unknown as Parameters[0]["modelRuntime"]; for (const level of ["off", "minimal", "low", "medium", "high", "xhigh", "max"]) { const result = resolveCliModel({ cliModel: `neuralwatt/zai-org/GLM-5.1-FP8:${level}`, - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -534,12 +534,12 @@ describe("resolveCliModel", () => { test("invalid thinking suffix on custom model is treated as part of model id", () => { const registry = { - getAll: () => modelsWithNeuralwatt, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => modelsWithNeuralwatt, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliModel: "neuralwatt/zai-org/GLM-5.1-FP8:banana", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -551,13 +551,13 @@ describe("resolveCliModel", () => { test("explicit --provider with custom model:thinking strips suffix correctly", () => { const registry = { - getAll: () => modelsWithNeuralwatt, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => modelsWithNeuralwatt, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliProvider: "neuralwatt", cliModel: "zai-org/GLM-5.1-FP8:high", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -568,13 +568,13 @@ describe("resolveCliModel", () => { test("with explicit --thinking, :suffix is kept as part of model id", () => { const registry = { - getAll: () => modelsWithNeuralwatt, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => modelsWithNeuralwatt, + } as unknown as Parameters[0]["modelRuntime"]; const result = resolveCliModel({ cliModel: "neuralwatt/zai-org/GLM-5.1-FP8:high", cliThinking: "medium", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.error).toBeUndefined(); @@ -606,15 +606,15 @@ describe("default model selection", () => { test("findInitialModel accepts explicit provider custom model ids", async () => { const registry = { - getAll: () => allModels, - } as unknown as Parameters[0]["modelRegistry"]; + getModels: () => allModels, + } as unknown as Parameters[0]["modelRuntime"]; const result = await findInitialModel({ cliProvider: "openrouter", cliModel: "openrouter/openai/ghost-model", scopedModels: [], isContinuing: false, - modelRegistry: registry, + modelRuntime: registry, }); expect(result.model?.provider).toBe("openrouter"); @@ -637,12 +637,12 @@ describe("default model selection", () => { const registry = { getAvailable: async () => [aiGatewayModel], - } as unknown as Parameters[0]["modelRegistry"]; + } as unknown as Parameters[0]["modelRuntime"]; const result = await findInitialModel({ scopedModels: [], isContinuing: false, - modelRegistry: registry, + modelRuntime: registry, }); expect(result.model?.provider).toBe("vercel-ai-gateway"); @@ -668,20 +668,20 @@ describe("default model selection", () => { baseUrl: "http://spark-two:8000/v1", }; const registry = { - find: (provider: string, modelId: string) => + getModel: (provider: string, modelId: string) => provider === savedDeepSeekModel.provider && modelId === savedDeepSeekModel.id ? savedDeepSeekModel : undefined, - hasConfiguredAuth: (model: Model<"anthropic-messages">) => model.provider === "spark-two", + hasConfiguredAuth: (provider: string) => provider === "spark-two", getAvailable: async () => [localDeepSeekModel], - } as unknown as Parameters[0]["modelRegistry"]; + } as unknown as Parameters[0]["modelRuntime"]; const result = await findInitialModel({ scopedModels: [], isContinuing: false, defaultProvider: "deepseek", defaultModelId: "deepseek-v4-flash", - modelRegistry: registry, + modelRuntime: registry, }); expect(result.model?.provider).toBe("spark-two"); diff --git a/packages/coding-agent/test/model-runtime-auth-options.test.ts b/packages/coding-agent/test/model-runtime-auth-options.test.ts new file mode 100644 index 00000000..74f65be4 --- /dev/null +++ b/packages/coding-agent/test/model-runtime-auth-options.test.ts @@ -0,0 +1,256 @@ +import { type AuthType, type CredentialStore, InMemoryCredentialStore } from "@earendil-works/pi-ai"; +import { describe, expect, it } from "vitest"; +import { AuthStorage } from "../src/core/auth-storage.ts"; +import { ModelRuntime } from "../src/core/model-runtime.ts"; + +function authOptions(runtime: ModelRuntime, type?: AuthType) { + return runtime + .getProviders() + .flatMap((provider) => [ + ...(!type || type === "oauth" + ? provider.auth.oauth + ? [{ type: "oauth" as const, provider, method: provider.auth.oauth }] + : [] + : []), + ...(!type || type === "api_key" + ? provider.auth.apiKey + ? [{ type: "api_key" as const, provider, method: provider.auth.apiKey }] + : [] + : []), + ]); +} + +function testModel(id: string) { + return { + id, + name: id, + reasoning: false, + input: ["text"] as ("text" | "image")[], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 10000, + maxTokens: 1000, + }; +} + +describe("ModelRuntime auth options", () => { + it("accepts a pi-ai CredentialStore", async () => { + const credentials = new InMemoryCredentialStore(); + await credentials.modify("anthropic", async () => ({ type: "api_key", key: "stored-key" })); + const runtime = await ModelRuntime.create({ credentials, modelsPath: null }); + + expect((await runtime.getAuth("anthropic"))?.auth.apiKey).toBe("stored-key"); + }); + + it("scopes provider availability reads and records refresh failures", async () => { + const base = new InMemoryCredentialStore(); + const reads: string[] = []; + let failReads = false; + const credentials: CredentialStore = { + read: async (providerId) => { + reads.push(providerId); + if (failReads) throw new Error(`read failed for ${providerId}`); + return base.read(providerId); + }, + list: () => base.list(), + modify: (providerId, fn) => base.modify(providerId, fn), + delete: (providerId) => base.delete(providerId), + }; + const runtime = await ModelRuntime.create({ credentials, modelsPath: null }); + + reads.length = 0; + await runtime.getAvailable("anthropic"); + expect(new Set(reads)).toEqual(new Set(["anthropic"])); + + failReads = true; + await expect(runtime.getAvailable("anthropic")).rejects.toThrow("Credential store read failed for anthropic"); + expect(runtime.getError()).toContain("Availability refresh: Credential store read failed for anthropic"); + + failReads = false; + await runtime.getAvailable(); + expect(runtime.getError()).toBeUndefined(); + }); + + it("projects provider-owned methods, names, and status", async () => { + const runtime = await ModelRuntime.create({ credentials: AuthStorage.inMemory(), modelsPath: null }); + const options = authOptions(runtime); + + expect(options).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + type: "api_key", + provider: expect.objectContaining({ id: "amazon-bedrock", name: "Amazon Bedrock" }), + method: expect.objectContaining({ name: "AWS credentials or bearer token" }), + }), + expect.objectContaining({ + type: "api_key", + provider: expect.objectContaining({ id: "google-vertex", name: "Google Vertex AI" }), + method: expect.objectContaining({ name: "Google Cloud credentials" }), + }), + expect.objectContaining({ + type: "oauth", + provider: expect.objectContaining({ id: "anthropic", name: "Anthropic" }), + }), + expect.objectContaining({ + type: "api_key", + provider: expect.objectContaining({ id: "cloudflare-ai-gateway", name: "Cloudflare AI Gateway" }), + }), + expect.objectContaining({ + type: "api_key", + provider: expect.objectContaining({ id: "cloudflare-workers-ai", name: "Cloudflare Workers AI" }), + }), + ]), + ); + expect(authOptions(runtime, "api_key").every((option) => option.type === "api_key")).toBe(true); + expect(authOptions(runtime, "oauth").every((option) => option.type === "oauth")).toBe(true); + expect(options.some((option) => option.provider.id === "openai-codex" && option.type === "api_key")).toBe(false); + }); + + it("attaches the provider's active auth status to every method option", async () => { + const runtime = await ModelRuntime.create({ + credentials: AuthStorage.inMemory({ + anthropic: { + type: "oauth", + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }, + }), + modelsPath: null, + }); + + const options = authOptions(runtime).filter((option) => option.provider.id === "anthropic"); + expect(options).toHaveLength(2); + expect(await runtime.checkAuth("anthropic")).toMatchObject({ type: "oauth" }); + }); + + it("constructs an API key method for an extension API-key provider", async () => { + const runtime = await ModelRuntime.create({ credentials: AuthStorage.inMemory(), modelsPath: null }); + runtime.registerProvider("extension-api-key", { + name: "Extension API Key", + baseUrl: "https://example.test/v1", + apiKey: "$EXTENSION_TEST_API_KEY", + api: "openai-completions", + models: [testModel("extension-model")], + }); + + const options = authOptions(runtime).filter((option) => option.provider.id === "extension-api-key"); + expect(options).toHaveLength(1); + expect(options[0]).toMatchObject({ + type: "api_key", + provider: { id: "extension-api-key", name: "Extension API Key" }, + method: { name: "API key" }, + }); + expect(options[0]?.method.login).toBeTypeOf("function"); + }); + + it("resolves configured auth from request-scoped environment overrides", async () => { + const runtime = await ModelRuntime.create({ credentials: AuthStorage.inMemory(), modelsPath: null }); + runtime.registerProvider("request-env-provider", { + baseUrl: "https://example.test/v1", + apiKey: "$REQUEST_SCOPED_API_KEY", + headers: { "x-request-value": "$REQUEST_SCOPED_HEADER" }, + api: "openai-completions", + models: [testModel("request-env-model")], + }); + + const auth = await runtime.getAuth("request-env-provider", { + env: { REQUEST_SCOPED_API_KEY: "request-key", REQUEST_SCOPED_HEADER: "request-header" }, + }); + + expect(auth?.auth).toEqual({ apiKey: "request-key", headers: { "x-request-value": "request-header" } }); + }); + + it("lets an explicit Authorization header override authHeader case-insensitively", async () => { + const runtime = await ModelRuntime.create({ credentials: AuthStorage.inMemory(), modelsPath: null }); + let capturedHeaders: Record | undefined; + runtime.registerProvider("auth-header-provider", { + baseUrl: "https://example.test/v1", + apiKey: "generated-key", + authHeader: true, + api: "openai-completions", + streamSimple: (_model, _context, options) => { + capturedHeaders = options?.headers; + throw new Error("captured"); + }, + models: [testModel("auth-header-model")], + }); + const model = runtime.getModel("auth-header-provider", "auth-header-model"); + expect(model).toBeDefined(); + + await runtime.completeSimple(model!, { messages: [] }, { headers: { authorization: "Explicit token" } }); + + expect(capturedHeaders).toEqual({ authorization: "Explicit token" }); + }); + + it("transforms fully assembled headers once without forwarding the transform", async () => { + const runtime = await ModelRuntime.create({ credentials: AuthStorage.inMemory(), modelsPath: null }); + let capturedHeaders: Record | undefined; + let transforms = 0; + runtime.registerProvider("header-provider", { + baseUrl: "https://example.test/v1", + apiKey: "generated-key", + authHeader: true, + headers: { "x-provider": "provider" }, + api: "openai-completions", + streamSimple: (_model, _context, options) => { + expect(options).not.toHaveProperty("transformHeaders"); + capturedHeaders = options?.headers; + throw new Error("captured"); + }, + models: [{ ...testModel("header-model"), headers: { "x-model": "model" } }], + }); + const model = runtime.getModel("header-provider", "header-model"); + expect(model).toBeDefined(); + + await runtime.completeSimple( + model!, + { messages: [] }, + { + headers: { "x-explicit": "explicit" }, + transformHeaders: async (headers) => { + transforms++; + expect(headers).toEqual({ + Authorization: "Bearer generated-key", + "x-provider": "provider", + "x-model": "model", + "x-explicit": "explicit", + }); + return { ...headers, "x-transformed": "yes" }; + }, + }, + ); + + expect(transforms).toBe(1); + expect(capturedHeaders).toEqual({ + Authorization: "Bearer generated-key", + "x-provider": "provider", + "x-model": "model", + "x-explicit": "explicit", + "x-transformed": "yes", + }); + }); + + 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", { + 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 (credentials) => credentials, + getApiKey: (credentials) => credentials.access, + }, + models: [testModel("extension-model")], + }); + + const options = authOptions(runtime).filter((option) => option.provider.id === "extension-oauth"); + expect(options).toHaveLength(1); + expect(options[0]).toMatchObject({ + type: "oauth", + provider: { id: "extension-oauth", name: "Extension OAuth" }, + method: { name: "Extension subscription" }, + }); + }); +}); diff --git a/packages/coding-agent/test/model-runtime-cloudflare-compat.test.ts b/packages/coding-agent/test/model-runtime-cloudflare-compat.test.ts new file mode 100644 index 00000000..cb91b3c4 --- /dev/null +++ b/packages/coding-agent/test/model-runtime-cloudflare-compat.test.ts @@ -0,0 +1,95 @@ +import { complete, resetApiProviders } from "@earendil-works/pi-ai/compat"; +import { describe, expect, it, vi } from "vitest"; +import { AuthStorage } from "../src/core/auth-storage.ts"; +import { ModelRegistry } from "../src/core/model-registry.ts"; +import { ModelRuntime } from "../src/core/model-runtime.ts"; + +const openAIState = vi.hoisted(() => ({ clientOptions: undefined as unknown })); + +vi.mock("openai", () => { + class FakeOpenAI { + constructor(options: unknown) { + openAIState.clientOptions = options; + } + + chat = { + completions: { + create: () => { + const stream = { + async *[Symbol.asyncIterator]() { + yield { + choices: [{ delta: {}, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1 }, + }; + }, + }; + const promise = Promise.resolve(stream) as Promise & { + withResponse(): Promise<{ + data: typeof stream; + response: { status: number; headers: Headers }; + }>; + }; + promise.withResponse = async () => ({ + data: stream, + response: { status: 200, headers: new Headers() }, + }); + return promise; + }, + }, + }; + } + + return { default: FakeOpenAI }; +}); + +async function createCloudflareRuntime(): Promise<{ modelRuntime: ModelRuntime; modelRegistry: ModelRegistry }> { + const authStorage = AuthStorage.inMemory(); + await authStorage.modify("cloudflare-ai-gateway", async () => ({ + type: "api_key", + key: "test-token", + env: { + CLOUDFLARE_ACCOUNT_ID: "test-account", + CLOUDFLARE_GATEWAY_ID: "test-gateway", + }, + })); + const modelRuntime = await ModelRuntime.create({ credentials: authStorage, modelsPath: null }); + return { modelRuntime, modelRegistry: new ModelRegistry(modelRuntime) }; +} + +describe("ModelRegistry Cloudflare compat streaming", () => { + it("materializes the Cloudflare endpoint through ModelRuntime streaming", async () => { + const { modelRuntime } = await createCloudflareRuntime(); + const model = modelRuntime.getModel("cloudflare-ai-gateway", "workers-ai/@cf/moonshotai/kimi-k2.5"); + expect(model).toBeDefined(); + + resetApiProviders(); + await modelRuntime.completeSimple(model!, { messages: [] }); + + const clientOptions = openAIState.clientOptions as { + baseURL?: string; + defaultHeaders?: Record; + }; + expect(clientOptions.baseURL).toBe("https://gateway.ai.cloudflare.com/v1/test-account/test-gateway/compat"); + expect(clientOptions.defaultHeaders?.["cf-aig-authorization"]).toBe("Bearer test-token"); + }); + + it("materializes the Cloudflare endpoint after extension-style auth resolution", async () => { + const { modelRegistry } = await createCloudflareRuntime(); + const model = modelRegistry.find("cloudflare-ai-gateway", "workers-ai/@cf/moonshotai/kimi-k2.5"); + expect(model).toBeDefined(); + + resetApiProviders(); + const auth = await modelRegistry.getApiKeyAndHeaders(model!); + expect(auth.ok).toBe(true); + if (!auth.ok) throw new Error(auth.error); + + await complete(model!, { messages: [] }, auth); + + const clientOptions = openAIState.clientOptions as { + baseURL?: string; + defaultHeaders?: Record; + }; + expect(clientOptions.baseURL).toBe("https://gateway.ai.cloudflare.com/v1/test-account/test-gateway/compat"); + expect(clientOptions.defaultHeaders?.["cf-aig-authorization"]).toBe("Bearer test-token"); + }); +}); diff --git a/packages/coding-agent/test/model-runtime-test-utils.ts b/packages/coding-agent/test/model-runtime-test-utils.ts new file mode 100644 index 00000000..3b180e72 --- /dev/null +++ b/packages/coding-agent/test/model-runtime-test-utils.ts @@ -0,0 +1,25 @@ +import type { CredentialStore } from "@earendil-works/pi-ai"; +import { ModelRegistry } from "../src/core/model-registry.ts"; +import { ModelRuntime } from "../src/core/model-runtime.ts"; + +const runtimes = new WeakMap(); + +function wrap(runtime: ModelRuntime): ModelRegistry { + const registry = new ModelRegistry(runtime); + runtimes.set(registry, runtime); + return registry; +} + +export async function createModelRegistry(credentials: CredentialStore, modelsPath?: string): Promise { + return wrap(await ModelRuntime.create({ credentials, modelsPath })); +} + +export async function createInMemoryModelRegistry(credentials: CredentialStore): Promise { + return wrap(await ModelRuntime.create({ credentials, modelsPath: null })); +} + +export function getModelRuntime(modelRegistry: ModelRegistry): ModelRuntime { + const runtime = runtimes.get(modelRegistry); + if (!runtime) throw new Error("ModelRegistry was not created by the test helper"); + return runtime; +} diff --git a/packages/coding-agent/test/oauth-selector.test.ts b/packages/coding-agent/test/oauth-selector.test.ts index 7164db54..71cfb2f5 100644 --- a/packages/coding-agent/test/oauth-selector.test.ts +++ b/packages/coding-agent/test/oauth-selector.test.ts @@ -1,15 +1,11 @@ import { setKeybindings } from "@earendil-works/pi-tui"; -import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; -import { AuthStorage } from "../src/core/auth-storage.ts"; +import { beforeAll, beforeEach, describe, expect, it } from "vitest"; import { KeybindingsManager } from "../src/core/keybindings.ts"; -import { BUILT_IN_PROVIDER_DISPLAY_NAMES } from "../src/core/provider-display-names.ts"; import { OAuthSelectorComponent } from "../src/modes/interactive/components/oauth-selector.ts"; -import { isApiKeyLoginProvider } from "../src/modes/interactive/interactive-mode.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"; -const originalOpenAiApiKey = process.env.OPENAI_API_KEY; - describe("OAuthSelectorComponent", () => { beforeAll(() => { initTheme("dark"); @@ -19,119 +15,133 @@ describe("OAuthSelectorComponent", () => { setKeybindings(new KeybindingsManager()); }); - afterEach(() => { - if (originalOpenAiApiKey === undefined) { - delete process.env.OPENAI_API_KEY; - } else { - process.env.OPENAI_API_KEY = originalOpenAiApiKey; - } - }); - - it("keeps built-in API key providers separate from OAuth-only providers", () => { - const oauthProviderIds = new Set(["anthropic", "github-copilot", "custom-oauth"]); - const builtInProviderIds = new Set(["anthropic", "github-copilot", "amazon-bedrock", "openai"]); - - expect(isApiKeyLoginProvider("anthropic", oauthProviderIds, builtInProviderIds)).toBe(true); - expect(BUILT_IN_PROVIDER_DISPLAY_NAMES.anthropic).toBe("Anthropic"); - expect(isApiKeyLoginProvider("openai", oauthProviderIds, builtInProviderIds)).toBe(true); - expect(isApiKeyLoginProvider("github-copilot", oauthProviderIds, builtInProviderIds)).toBe(false); - expect(isApiKeyLoginProvider("amazon-bedrock", oauthProviderIds, builtInProviderIds)).toBe(true); - expect(isApiKeyLoginProvider("custom-oauth", oauthProviderIds, builtInProviderIds)).toBe(false); - expect(isApiKeyLoginProvider("custom-api", oauthProviderIds, builtInProviderIds)).toBe(true); - }); - - it("shows stored OAuth auth distinctly in the API key selector", () => { - const authStorage = AuthStorage.inMemory({ - anthropic: { - type: "oauth", - access: "access-token", - refresh: "refresh-token", - expires: Date.now() + 60_000, + it("projects provider-owned auth options without provider-specific filtering", () => { + const getLoginProviderOptions = ( + InteractiveMode as unknown as { + prototype: { + getLoginProviderOptions( + this: object, + authType?: "oauth" | "api_key", + ): Array<{ id: string; name: string; authType: string; method?: { name: string; login?: unknown } }>; + }; + } + ).prototype.getLoginProviderOptions; + const providers = [ + { + id: "anthropic", + name: "Anthropic", + auth: { + oauth: { name: "Anthropic (Claude Pro/Max)", login: async () => ({}) }, + apiKey: { name: "Anthropic API key", login: async () => ({}) }, + }, }, - }); + { + id: "google-vertex", + name: "Google Vertex AI", + auth: { apiKey: { name: "Google Cloud credentials" } }, + }, + ]; + const fakeThis = { + session: { + modelRuntime: { + getProviders: () => providers, + getProviderAuthStatus: () => ({ configured: false }), + isUsingOAuth: () => false, + }, + }, + }; + + const apiKeyOptions = getLoginProviderOptions.call(fakeThis, "api_key"); + expect(apiKeyOptions).toMatchObject([ + { + id: "anthropic", + name: "Anthropic", + authType: "api_key", + method: { name: "Anthropic API key" }, + }, + { + id: "google-vertex", + name: "Google Vertex AI", + authType: "api_key", + method: { name: "Google Cloud credentials" }, + }, + ]); + expect(getLoginProviderOptions.call(fakeThis, "oauth")).toMatchObject([ + { id: "anthropic", name: "Anthropic", authType: "oauth" }, + ]); + }); + + it("renders an option without compiled auth status as unconfigured", () => { const selector = new OAuthSelectorComponent( "login", - authStorage, - [{ id: "anthropic", name: "Anthropic", authType: "api_key" }], + [{ id: "google", name: "Google", authType: "api_key", status: undefined }], () => {}, () => {}, ); const output = stripAnsi(selector.render(120).join("\n")); + expect(output).toContain("unconfigured"); + expect(output).not.toContain("✓ configured"); + }); - expect(output).toContain("Anthropic"); + it("shows OAuth auth distinctly in the API key selector", () => { + const selector = new OAuthSelectorComponent( + "login", + [{ id: "anthropic", name: "Anthropic", authType: "api_key", status: { type: "oauth", source: "OAuth" } }], + () => {}, + () => {}, + ); + + const output = stripAnsi(selector.render(120).join("\n")); expect(output).toContain("subscription configured"); }); it("shows environment API key auth as configured", () => { - process.env.OPENAI_API_KEY = "test-openai-key"; - const authStorage = AuthStorage.inMemory(); const selector = new OAuthSelectorComponent( "login", - authStorage, - [{ id: "openai", name: "OpenAI", authType: "api_key" }], + [{ id: "openai", name: "OpenAI", authType: "api_key", status: { type: "api_key", source: "OPENAI_API_KEY" } }], () => {}, () => {}, ); const output = stripAnsi(selector.render(120).join("\n")); - - expect(output).toContain("OpenAI"); expect(output).toContain("✓ env: OPENAI_API_KEY"); expect(output).not.toContain("unconfigured"); }); - it("shows custom provider environment API key auth from status resolver", () => { - const authStorage = AuthStorage.inMemory(); - const selector = new OAuthSelectorComponent( - "login", - authStorage, - [{ id: "ollama", name: "ollama", authType: "api_key" }], - () => {}, - () => {}, - () => ({ configured: true, source: "environment", label: "OLLAMA_API_KEY" }), - ); - - const output = stripAnsi(selector.render(120).join("\n")); - - expect(output).toContain("ollama"); - expect(output).toContain("✓ env: OLLAMA_API_KEY"); - expect(output).not.toContain("unconfigured"); - }); - it("shows models.json API key auth as configured", () => { - const authStorage = AuthStorage.inMemory(); const selector = new OAuthSelectorComponent( "login", - authStorage, - [{ id: "local-proxy", name: "local-proxy", authType: "api_key" }], + [ + { + id: "local-proxy", + name: "local-proxy", + authType: "api_key", + status: { type: "api_key", source: "key in models.json" }, + }, + ], () => {}, () => {}, - () => ({ configured: true, source: "models_json_key" }), ); - const output = stripAnsi(selector.render(120).join("\n")); - - expect(output).toContain("local-proxy"); - expect(output).toContain("✓ key in models.json"); - expect(output).not.toContain("unconfigured"); + expect(stripAnsi(selector.render(120).join("\n"))).toContain("✓ key in models.json"); }); it("shows models.json command auth as configured", () => { - const authStorage = AuthStorage.inMemory(); const selector = new OAuthSelectorComponent( "login", - authStorage, - [{ id: "op-proxy", name: "op-proxy", authType: "api_key" }], + [ + { + id: "op-proxy", + name: "op-proxy", + authType: "api_key", + status: { type: "api_key", source: "command in models.json" }, + }, + ], () => {}, () => {}, - () => ({ configured: true, source: "models_json_command" }), ); - const output = stripAnsi(selector.render(120).join("\n")); - - expect(output).toContain("op-proxy"); - expect(output).toContain("✓ command in models.json"); - expect(output).not.toContain("unconfigured"); + expect(stripAnsi(selector.render(120).join("\n"))).toContain("✓ command in models.json"); }); }); diff --git a/packages/coding-agent/test/resolve-config-value.test.ts b/packages/coding-agent/test/resolve-config-value.test.ts new file mode 100644 index 00000000..b202f1c8 --- /dev/null +++ b/packages/coding-agent/test/resolve-config-value.test.ts @@ -0,0 +1,121 @@ +import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { + clearConfigValueCache, + resolveConfigValue, + resolveConfigValueUncached, +} from "../src/core/resolve-config-value.ts"; +import * as shellModule from "../src/utils/shell.ts"; + +describe("resolveConfigValue", () => { + let tempDir: string; + + beforeEach(() => { + tempDir = join(tmpdir(), `pi-config-value-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(tempDir, { recursive: true }); + clearConfigValueCache(); + }); + + afterEach(() => { + if (existsSync(tempDir)) rmSync(tempDir, { recursive: true }); + clearConfigValueCache(); + vi.restoreAllMocks(); + }); + + test("resolves literals, environment templates, and escapes", () => { + process.env.TEST_CONFIG_LEFT = "left"; + process.env.TEST_CONFIG_RIGHT = "right"; + try { + expect(resolveConfigValue("literal-key")).toBe("literal-key"); + expect(resolveConfigValue("$TEST_CONFIG_LEFT")).toBe("left"); + expect(resolveConfigValue("$" + "{TEST_CONFIG_LEFT}_$TEST_CONFIG_RIGHT")).toBe("left_right"); + expect(resolveConfigValue("$$TEST_CONFIG_LEFT")).toBe("$TEST_CONFIG_LEFT"); + expect(resolveConfigValue("$!literal-$TEST_CONFIG_RIGHT")).toBe("!literal-right"); + } finally { + delete process.env.TEST_CONFIG_LEFT; + delete process.env.TEST_CONFIG_RIGHT; + } + }); + + test("uses credential-scoped environment before process.env", () => { + process.env.TEST_CONFIG_SCOPED = "process"; + try { + expect(resolveConfigValue("$TEST_CONFIG_SCOPED", { TEST_CONFIG_SCOPED: "credential" })).toBe("credential"); + } finally { + delete process.env.TEST_CONFIG_SCOPED; + } + }); + + test("executes shell commands and trims their output", () => { + expect(resolveConfigValue("!echo ' spaced-key '")).toBe("spaced-key"); + expect(resolveConfigValue("!printf 'line1\\nline2'")).toBe("line1\nline2"); + expect(resolveConfigValue("!echo 'hello world' | tr ' ' '-'")).toBe("hello-world"); + }); + + test.each(["!exit 1", "!nonexistent-command-12345", "!printf ''"])( + "returns undefined when command resolution fails: %s", + (command) => { + expect(resolveConfigValue(command)).toBeUndefined(); + }, + ); + + test("caches successful and failed commands until explicitly cleared", () => { + const counterFile = join(tempDir, "counter"); + writeFileSync(counterFile, "0"); + const escapedPath = counterFile.replace(/\\/g, "/").replace(/"/g, '\\"'); + const success = `!sh -c 'count=$(cat "${escapedPath}"); echo $((count + 1)) > "${escapedPath}"; echo value'`; + + expect(resolveConfigValue(success)).toBe("value"); + expect(resolveConfigValue(success)).toBe("value"); + expect(readFileSync(counterFile, "utf-8").trim()).toBe("1"); + + clearConfigValueCache(); + expect(resolveConfigValue(success)).toBe("value"); + expect(readFileSync(counterFile, "utf-8").trim()).toBe("2"); + + const failure = `!sh -c 'count=$(cat "${escapedPath}"); echo $((count + 1)) > "${escapedPath}"; exit 1'`; + expect(resolveConfigValue(failure)).toBeUndefined(); + expect(resolveConfigValue(failure)).toBeUndefined(); + expect(readFileSync(counterFile, "utf-8").trim()).toBe("3"); + }); + + test("does not cache environment values", () => { + process.env.TEST_CONFIG_DYNAMIC = "first"; + try { + expect(resolveConfigValue("$TEST_CONFIG_DYNAMIC")).toBe("first"); + process.env.TEST_CONFIG_DYNAMIC = "second"; + expect(resolveConfigValue("$TEST_CONFIG_DYNAMIC")).toBe("second"); + } finally { + delete process.env.TEST_CONFIG_DYNAMIC; + } + }); + + test("uncached resolution executes a command on every call", () => { + const counterFile = join(tempDir, "uncached-counter"); + writeFileSync(counterFile, "0"); + const escapedPath = counterFile.replace(/\\/g, "/").replace(/"/g, '\\"'); + const command = `!sh -c 'count=$(cat "${escapedPath}"); echo $((count + 1)) > "${escapedPath}"; echo value'`; + expect(resolveConfigValueUncached(command)).toBe("value"); + expect(resolveConfigValueUncached(command)).toBe("value"); + expect(readFileSync(counterFile, "utf-8").trim()).toBe("2"); + }); + + test("uses stdin when the configured Windows shell requires it", () => { + if (process.platform === "win32") return; + const platformDescriptor = Object.getOwnPropertyDescriptor(process, "platform"); + vi.spyOn(shellModule, "getShellConfig").mockReturnValue({ + shell: "/bin/bash", + args: ["-s"], + commandTransport: "stdin", + }); + try { + Object.defineProperty(process, "platform", { configurable: true, value: "win32" }); + const expansion = "$" + "{name}"; + expect(resolveConfigValueUncached(`!name='World'; echo "Hello, ${expansion}!"`)).toBe("Hello, World!"); + } finally { + if (platformDescriptor) Object.defineProperty(process, "platform", platformDescriptor); + } + }); +}); diff --git a/packages/coding-agent/test/resource-loader.test.ts b/packages/coding-agent/test/resource-loader.test.ts index 72cee79a..e458ec7b 100644 --- a/packages/coding-agent/test/resource-loader.test.ts +++ b/packages/coding-agent/test/resource-loader.test.ts @@ -5,13 +5,14 @@ import { pathToFileURL } from "node:url"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; import { ExtensionRunner } from "../src/core/extensions/runner.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { DefaultResourceLoader } from "../src/core/resource-loader.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; import type { Skill } from "../src/core/skills.ts"; import { createSyntheticSourceInfo } from "../src/core/source-info.ts"; +import { createModelRegistry } from "./model-runtime-test-utils.ts"; + describe("DefaultResourceLoader", () => { let tempDir: string; let agentDir: string; @@ -277,7 +278,7 @@ export default function(pi) { const sessionManager = SessionManager.inMemory(); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage); + const modelRegistry = await createModelRegistry(authStorage); const runner = new ExtensionRunner( extensionsResult.extensions, extensionsResult.runtime, @@ -721,7 +722,7 @@ export default function(pi: ExtensionAPI) { const sessionManager = SessionManager.inMemory(); const authStorage = AuthStorage.create(join(tempDir, "auth-explicit.json")); - const modelRegistry = ModelRegistry.create(authStorage); + const modelRegistry = await createModelRegistry(authStorage); const runner = new ExtensionRunner( extensionsResult.extensions, extensionsResult.runtime, diff --git a/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts b/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts index 5e4d5b08..155d2d38 100644 --- a/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts +++ b/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts @@ -13,10 +13,10 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { AgentSession } from "../src/core/agent-session.ts"; import type { AgentSessionRuntime } from "../src/core/agent-session-runtime.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; import { runRpcMode } from "../src/modes/rpc/rpc-mode.ts"; +import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; import { createTestResourceLoader } from "./utilities.ts"; const rpcIo = vi.hoisted(() => ({ @@ -95,10 +95,10 @@ function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } -function createRuntimeHost(options: { withAuth: boolean; responseDelayMs: number; model?: Model }): { +async function createRuntimeHost(options: { withAuth: boolean; responseDelayMs: number; model?: Model }): Promise<{ runtimeHost: AgentSessionRuntime; cleanup: () => Promise; -} { +}> { const tempDir = join(tmpdir(), `pi-rpc-prompt-${Date.now()}-${Math.random().toString(36).slice(2)}`); mkdirSync(tempDir, { recursive: true }); @@ -129,9 +129,9 @@ function createRuntimeHost(options: { withAuth: boolean; responseDelayMs: number const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); + const modelRegistry = await createModelRegistry(authStorage, tempDir); if (options.withAuth) { - authStorage.setRuntimeApiKey("anthropic", "test-key"); + await authStorage.modify("anthropic", async () => ({ type: "api_key", key: "test-key" })); } const session = new AgentSession({ @@ -139,7 +139,7 @@ function createRuntimeHost(options: { withAuth: boolean; responseDelayMs: number sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), }); @@ -177,7 +177,7 @@ async function startRpcMode(options: { withAuth: boolean; responseDelayMs: numbe rpcIo.outputLines = []; rpcIo.lineHandler = undefined; - const { runtimeHost, cleanup } = createRuntimeHost(options); + const { runtimeHost, cleanup } = await createRuntimeHost(options); void runRpcMode(runtimeHost); await vi.waitFor(() => expect(rpcIo.lineHandler).toBeDefined()); diff --git a/packages/coding-agent/test/runtime-credentials.test.ts b/packages/coding-agent/test/runtime-credentials.test.ts new file mode 100644 index 00000000..a044f715 --- /dev/null +++ b/packages/coding-agent/test/runtime-credentials.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, test } from "vitest"; +import { AuthStorage } from "../src/core/auth-storage.ts"; +import { RuntimeCredentials } from "../src/core/runtime-credentials.ts"; + +describe("RuntimeCredentials", () => { + test("runtime overrides mask stored credentials without persisting", async () => { + const storage = AuthStorage.inMemory({ anthropic: { type: "api_key", key: "stored-key" } }); + const credentials = new RuntimeCredentials(storage); + + credentials.setRuntimeApiKey("anthropic", "runtime-key"); + expect(await credentials.read("anthropic")).toEqual({ type: "api_key", key: "runtime-key" }); + expect(await storage.read("anthropic")).toEqual({ type: "api_key", key: "stored-key" }); + + credentials.removeRuntimeApiKey("anthropic"); + expect(await credentials.read("anthropic")).toEqual({ type: "api_key", key: "stored-key" }); + }); + + test("enumeration merges overrides without exposing keys", async () => { + const storage = AuthStorage.inMemory({ + anthropic: { type: "oauth", access: "access", refresh: "refresh", expires: Date.now() + 60_000 }, + }); + const credentials = new RuntimeCredentials(storage); + credentials.setRuntimeApiKey("anthropic", "runtime-key"); + credentials.setRuntimeApiKey("openai", "other-runtime-key"); + + expect(await credentials.list()).toEqual([ + { providerId: "anthropic", type: "api_key" }, + { providerId: "openai", type: "api_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); + credentials.setRuntimeApiKey("anthropic", "runtime-key"); + + await credentials.delete("anthropic"); + + expect(await credentials.read("anthropic")).toBeUndefined(); + expect(await credentials.list()).toEqual([]); + }); +}); diff --git a/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts b/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts index ff6a9798..a6d3566d 100644 --- a/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts +++ b/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts @@ -28,11 +28,11 @@ import { import { AuthStorage } from "../src/core/auth-storage.ts"; import { createExtensionRuntime } from "../src/core/extensions/loader.ts"; import type { ToolDefinition } from "../src/core/extensions/types.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import type { ResourceLoader } from "../src/core/resource-loader.ts"; import { createAgentSession } from "../src/core/sdk.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; +import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; type Transport = "sse" | "websocket" | "websocket-cached" | "auto"; @@ -275,7 +275,7 @@ async function main(): Promise { mkdirSync(dirname(args.sessionPath), { recursive: true }); const authStorage = AuthStorage.create(); - const modelRegistry = ModelRegistry.create(authStorage); + const modelRegistry = await createModelRegistry(authStorage); const model = getModel("openai-codex", "gpt-5.5"); if (!model) { @@ -296,6 +296,7 @@ async function main(): Promise { models: [baseModel], }); + const modelRuntime = getModelRuntime(modelRegistry); const settingsManager = SettingsManager.inMemory({ compaction: { enabled: false }, retry: { enabled: false }, @@ -315,8 +316,7 @@ async function main(): Promise { resourceLoader, sessionManager: SessionManager.open(args.sessionPath), settingsManager, - authStorage, - modelRegistry, + modelRuntime, }); session.setActiveToolsByName(["deterministic_probe"]); diff --git a/packages/coding-agent/test/sdk-openrouter-attribution.test.ts b/packages/coding-agent/test/sdk-openrouter-attribution.test.ts index baaa5746..53f7e942 100644 --- a/packages/coding-agent/test/sdk-openrouter-attribution.test.ts +++ b/packages/coding-agent/test/sdk-openrouter-attribution.test.ts @@ -11,11 +11,12 @@ import { } from "@earendil-works/pi-ai"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { createAgentSession } from "../src/core/sdk.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; +import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; + describe("createAgentSession provider attribution headers", () => { let tempDir: string; let cwd: string; @@ -96,24 +97,20 @@ describe("createAgentSession provider attribution headers", () => { } const authStorage = AuthStorage.create(join(agentDir, "auth.json")); - authStorage.setRuntimeApiKey(model.provider, "test-api-key"); - const modelRegistry = ModelRegistry.create(authStorage, join(agentDir, "models.json")); - const registeredProviders = ["capture-provider"]; + await authStorage.modify(model.provider, async () => ({ type: "api_key", key: "test-api-key" })); + const modelRegistry = await createModelRegistry(authStorage, join(agentDir, "models.json")); let capturedOptions: SimpleStreamOptions | undefined; - modelRegistry.registerProvider("capture-provider", { - api: "openai-completions", + modelRegistry.registerProvider(model.provider, { + api: model.api, + headers: options.providerHeaders, streamSimple: (_model, _context, providerOptions) => { capturedOptions = providerOptions; return createDoneStream(); }, }); - if (options.providerHeaders) { - modelRegistry.registerProvider(model.provider, { headers: options.providerHeaders }); - registeredProviders.push(model.provider); - } - + const modelRuntime = getModelRuntime(modelRegistry); const sessionManager = SessionManager.inMemory(cwd); if (options.sessionId) { sessionManager.newSession({ id: options.sessionId }); @@ -123,14 +120,13 @@ describe("createAgentSession provider attribution headers", () => { cwd, agentDir, model, - authStorage, - modelRegistry, + modelRuntime, settingsManager, sessionManager, }); try { - await session.agent.streamFn( + const stream = await session.agent.streamFn( model, { messages: [] }, { @@ -138,12 +134,11 @@ describe("createAgentSession provider attribution headers", () => { ...(options.requestHeaders ? { headers: options.requestHeaders } : {}), }, ); + await stream.result(); return capturedOptions?.headers; } finally { session.dispose(); - for (const provider of registeredProviders.reverse()) { - modelRegistry.unregisterProvider(provider); - } + modelRegistry.unregisterProvider(model.provider); } } diff --git a/packages/coding-agent/test/sdk-stream-options.test.ts b/packages/coding-agent/test/sdk-stream-options.test.ts index 1d565c7e..f61cade0 100644 --- a/packages/coding-agent/test/sdk-stream-options.test.ts +++ b/packages/coding-agent/test/sdk-stream-options.test.ts @@ -1,4 +1,4 @@ -import { mkdirSync, mkdtempSync, rmSync } from "node:fs"; +import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { @@ -10,11 +10,12 @@ import { } from "@earendil-works/pi-ai"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { createAgentSession } from "../src/core/sdk.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; +import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; + describe("createAgentSession stream options", () => { let tempDir: string; let cwd: string; @@ -46,6 +47,7 @@ describe("createAgentSession stream options", () => { cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 128000, maxTokens: 4096, + headers: { "x-model": "model" }, }; } @@ -76,36 +78,44 @@ describe("createAgentSession stream options", () => { api: Api, settings: { httpIdleTimeoutMs?: number; websocketConnectTimeoutMs?: number }, requestOptions: SimpleStreamOptions = {}, + extensionSource?: string, ): Promise { const model = createModel(api); const settingsManager = SettingsManager.inMemory(settings); + if (extensionSource) { + const extensionsDir = join(agentDir, "extensions"); + mkdirSync(extensionsDir, { recursive: true }); + writeFileSync(join(extensionsDir, "headers.ts"), extensionSource); + } const authStorage = AuthStorage.create(join(agentDir, "auth.json")); - authStorage.setRuntimeApiKey(model.provider, "test-api-key"); - const modelRegistry = ModelRegistry.create(authStorage, join(agentDir, "models.json")); + await authStorage.modify(model.provider, async () => ({ type: "api_key", key: "test-api-key" })); + const modelRegistry = await createModelRegistry(authStorage, join(agentDir, "models.json")); let capturedOptions: SimpleStreamOptions | undefined; modelRegistry.registerProvider(model.provider, { api, + headers: { "x-provider": "provider" }, streamSimple: (_model, _context, providerOptions) => { capturedOptions = providerOptions; return createDoneStream(api); }, }); + const modelRuntime = getModelRuntime(modelRegistry); const sessionManager = SessionManager.inMemory(cwd); const { session } = await createAgentSession({ cwd, agentDir, model, - authStorage, - modelRegistry, + modelRuntime, settingsManager, sessionManager, }); try { - await session.agent.streamFn(model, { messages: [] }, requestOptions); + const stream = await session.agent.streamFn(model, { messages: [] }, requestOptions); + await stream.result(); return capturedOptions; } finally { session.dispose(); @@ -150,4 +160,29 @@ describe("createAgentSession stream options", () => { expect(options?.websocketConnectTimeoutMs).toBe(0); }); + + it("runs before_provider_headers on assembled headers without forwarding the transform", async () => { + const options = await captureStreamOptions( + "openai-completions", + {}, + { headers: { "x-explicit": "explicit" } }, + `export default function (pi) { + pi.on("before_provider_headers", (event) => { + event.headers["x-hook"] = [ + event.headers["x-provider"], + event.headers["x-model"], + event.headers["x-explicit"], + ].join(":"); + }); + }`, + ); + + expect(options?.headers).toMatchObject({ + "x-provider": "provider", + "x-model": "model", + "x-explicit": "explicit", + "x-hook": "provider:model:explicit", + }); + expect(options).not.toHaveProperty("transformHeaders"); + }); }); diff --git a/packages/coding-agent/test/suite/agent-session-runtime.test.ts b/packages/coding-agent/test/suite/agent-session-runtime.test.ts index cf1a81c2..9f6c1388 100644 --- a/packages/coding-agent/test/suite/agent-session-runtime.test.ts +++ b/packages/coding-agent/test/suite/agent-session-runtime.test.ts @@ -52,7 +52,7 @@ describe("AgentSessionRuntime characterization", () => { faux.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two"), fauxAssistantMessage("three")]); const authStorage = AuthStorage.inMemory(); - authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + await authStorage.modify(faux.getModel().provider, async () => ({ type: "api_key", key: "faux-key" })); const runtimeOptions = { agentDir: tempDir, @@ -343,7 +343,7 @@ describe("AgentSessionRuntime characterization", () => { faux.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two"), fauxAssistantMessage("three")]); const authStorage = AuthStorage.inMemory(); - authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + await authStorage.modify(faux.getModel().provider, async () => ({ type: "api_key", key: "faux-key" })); const runtimeOptions = { agentDir: tempDir, @@ -454,7 +454,7 @@ describe("AgentSessionRuntime characterization", () => { mkdirSync(secondDir, { recursive: true }); const { runtime, faux, tempDir } = await createRuntimeForTest(() => {}, { cwd: firstDir }); const otherAuthStorage = AuthStorage.inMemory(); - otherAuthStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + await otherAuthStorage.modify(faux.getModel().provider, async () => ({ type: "api_key", key: "faux-key" })); const otherRuntimeOptions = { agentDir: tempDir, authStorage: otherAuthStorage, @@ -527,7 +527,7 @@ describe("AgentSessionRuntime characterization", () => { const otherDir = join(tempDir, "other"); mkdirSync(otherDir, { recursive: true }); const otherAuthStorage = AuthStorage.inMemory(); - otherAuthStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + await otherAuthStorage.modify(faux.getModel().provider, async () => ({ type: "api_key", key: "faux-key" })); const otherRuntimeOptions = { agentDir: tempDir, authStorage: otherAuthStorage, diff --git a/packages/coding-agent/test/suite/harness.ts b/packages/coding-agent/test/suite/harness.ts index 1f01cad0..b7b27971 100644 --- a/packages/coding-agent/test/suite/harness.ts +++ b/packages/coding-agent/test/suite/harness.ts @@ -1,3 +1,4 @@ +import { createInMemoryModelRegistry, getModelRuntime } from "../model-runtime-test-utils.ts"; /** * Local test harness for the new coding-agent test suite. */ @@ -18,7 +19,6 @@ import { AgentSession, type AgentSessionEvent } from "../../src/core/agent-sessi import { AuthStorage } from "../../src/core/auth-storage.ts"; import type { ExtensionRunner } from "../../src/core/extensions/index.ts"; import { convertToLlm } from "../../src/core/messages.ts"; -import { ModelRegistry } from "../../src/core/model-registry.ts"; import { SessionManager } from "../../src/core/session-manager.ts"; import type { Settings } from "../../src/core/settings-manager.ts"; import { SettingsManager } from "../../src/core/settings-manager.ts"; @@ -113,9 +113,9 @@ export async function createHarness(options: HarnessOptions = {}): Promise ({ type: "api_key", key: "faux-key" })); } - const modelRegistry = ModelRegistry.inMemory(authStorage); + const modelRegistry = await createInMemoryModelRegistry(authStorage); if (withConfiguredAuth) { modelRegistry.registerProvider(model.provider, { baseUrl: model.baseUrl, @@ -178,7 +178,7 @@ export async function createHarness(options: HarnessOptions = {}): Promise { @@ -32,13 +33,17 @@ describe("issue #2753 reload stale resource settings", () => { models: [{ id: "faux-1", reasoning: false }], }); const authStorage = AuthStorage.inMemory(); - authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + await authStorage.modify(faux.getModel().provider, async () => ({ type: "api_key", key: "faux-key" })); + const modelRuntime = await ModelRuntime.create({ + credentials: authStorage, + modelsPath: join(agentDir, "models.json"), + }); const createRuntime: CreateAgentSessionRuntimeFactory = async ({ cwd, sessionManager, sessionStartEvent }) => { const services = await createAgentSessionServices({ cwd, agentDir, - authStorage, + modelRuntime, resourceLoaderOptions: { extensionFactories: [ (pi) => { diff --git a/packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts b/packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts index b6f7002e..9923cf68 100644 --- a/packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts +++ b/packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts @@ -11,6 +11,7 @@ import { createAgentSessionServices, } from "../../../src/core/agent-session-runtime.ts"; import { AuthStorage } from "../../../src/core/auth-storage.ts"; +import { ModelRuntime } from "../../../src/core/model-runtime.ts"; import { SessionManager } from "../../../src/core/session-manager.ts"; import type { ExtensionAPI, ExtensionCommandContext, ExtensionFactory } from "../../../src/index.ts"; @@ -45,13 +46,17 @@ describe("regression #2860: replaced session callbacks", () => { faux.setResponses(responses.map((response) => fauxAssistantMessage(response))); const authStorage = AuthStorage.inMemory(); - authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + await authStorage.modify(faux.getModel().provider, async () => ({ type: "api_key", key: "faux-key" })); + const modelRuntime = await ModelRuntime.create({ + credentials: authStorage, + modelsPath: join(tempDir, "models.json"), + }); const createRuntime: CreateAgentSessionRuntimeFactory = async ({ cwd, sessionManager, sessionStartEvent }) => { const services = await createAgentSessionServices({ cwd, agentDir: tempDir, - authStorage, + modelRuntime, resourceLoaderOptions: { extensionFactories: [ (pi: ExtensionAPI) => { diff --git a/packages/coding-agent/test/suite/regressions/3217-scoped-model-order.test.ts b/packages/coding-agent/test/suite/regressions/3217-scoped-model-order.test.ts index b32535ba..28fe7ada 100644 --- a/packages/coding-agent/test/suite/regressions/3217-scoped-model-order.test.ts +++ b/packages/coding-agent/test/suite/regressions/3217-scoped-model-order.test.ts @@ -1,5 +1,5 @@ import { setKeybindings, type TUI } from "@earendil-works/pi-tui"; -import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; +import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import { KeybindingsManager } from "../../../src/core/keybindings.ts"; import { ModelSelectorComponent } from "../../../src/modes/interactive/components/model-selector.ts"; import { ScopedModelsSelectorComponent } from "../../../src/modes/interactive/components/scoped-models-selector.ts"; @@ -13,10 +13,6 @@ function createFakeTui(): TUI { } as unknown as TUI; } -async function waitForAsyncRender(): Promise { - await new Promise((resolve) => setTimeout(resolve, 0)); -} - describe("issue #3217 scoped model ordering", () => { const harnesses: Harness[] = []; @@ -83,13 +79,15 @@ describe("issue #3217 scoped model ordering", () => { createFakeTui(), modelOne, harness.settingsManager, - harness.session.modelRegistry, + harness.session.modelRuntime, [{ model: modelTwo }, { model: modelOne }, { model: modelThree }], () => {}, () => {}, ); - await waitForAsyncRender(); + await vi.waitFor(() => { + expect(stripAnsi(selector.render(120).join("\n"))).toContain(`[${modelOne.provider}]`); + }); const renderedLines = stripAnsi(selector.render(120).join("\n")) .split("\n") diff --git a/packages/coding-agent/test/suite/regressions/5433-extension-oauth-prompt-input.test.ts b/packages/coding-agent/test/suite/regressions/5433-extension-oauth-prompt-input.test.ts index 06562b3d..5d8448e6 100644 --- a/packages/coding-agent/test/suite/regressions/5433-extension-oauth-prompt-input.test.ts +++ b/packages/coding-agent/test/suite/regressions/5433-extension-oauth-prompt-input.test.ts @@ -70,6 +70,20 @@ describe("LoginDialogComponent OAuth prompts", () => { expect(output).toContain("First prompt:"); }); + test("preserves neutral information and links when showing a prompt", () => { + const dialog = createDialog(); + + dialog.showInfo("Configure credentials outside pi.", [ + { label: "Provider documentation", url: "https://example.invalid/docs" }, + ]); + dialog.showPrompt("Press Enter to continue:"); + + const output = renderDialog(dialog).join("\n"); + expect(output).toContain("Configure credentials outside pi."); + expect(output).toContain("Provider documentation: https://example.invalid/docs"); + expect(output).toContain("Press Enter to continue:"); + }); + test("keeps previous manual input stable when a later prompt is active", async () => { const dialog = createDialog(); diff --git a/packages/coding-agent/test/suite/regressions/5596-missing-theme-export.test.ts b/packages/coding-agent/test/suite/regressions/5596-missing-theme-export.test.ts index 296993e0..e8ef3b35 100644 --- a/packages/coding-agent/test/suite/regressions/5596-missing-theme-export.test.ts +++ b/packages/coding-agent/test/suite/regressions/5596-missing-theme-export.test.ts @@ -7,10 +7,10 @@ import { afterEach, describe, expect, it } from "vitest"; import { AgentSession } from "../../../src/core/agent-session.ts"; import { AuthStorage } from "../../../src/core/auth-storage.ts"; import { convertToLlm } from "../../../src/core/messages.ts"; -import { ModelRegistry } from "../../../src/core/model-registry.ts"; import { SessionManager } from "../../../src/core/session-manager.ts"; import { SettingsManager } from "../../../src/core/settings-manager.ts"; import { initTheme } from "../../../src/modes/interactive/theme/theme.ts"; +import { createInMemoryModelRegistry, getModelRuntime } from "../../model-runtime-test-utils.ts"; import { createTestResourceLoader } from "../../utilities.ts"; describe("regression #5596: missing configured theme export", () => { @@ -32,8 +32,8 @@ describe("regression #5596: missing configured theme export", () => { const model = faux.getModel(); const authStorage = AuthStorage.inMemory(); - authStorage.setRuntimeApiKey(model.provider, "faux-key"); - const modelRegistry = ModelRegistry.inMemory(authStorage); + await authStorage.modify(model.provider, async () => ({ type: "api_key", key: "faux-key" })); + const modelRegistry = await createInMemoryModelRegistry(authStorage); modelRegistry.registerProvider(model.provider, { baseUrl: model.baseUrl, apiKey: "faux-key", @@ -67,7 +67,7 @@ describe("regression #5596: missing configured theme export", () => { sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), }); cleanups.push(() => { diff --git a/packages/coding-agent/test/suite/regressions/5661-uppercase-header-values.test.ts b/packages/coding-agent/test/suite/regressions/5661-uppercase-header-values.test.ts index e625a3c6..40671d96 100644 --- a/packages/coding-agent/test/suite/regressions/5661-uppercase-header-values.test.ts +++ b/packages/coding-agent/test/suite/regressions/5661-uppercase-header-values.test.ts @@ -3,8 +3,8 @@ import { join } from "node:path"; import { afterEach, describe, expect, it } from "vitest"; import { ENV_AGENT_DIR } from "../../../src/config.ts"; import { AuthStorage } from "../../../src/core/auth-storage.ts"; -import { ModelRegistry } from "../../../src/core/model-registry.ts"; import { runMigrations } from "../../../src/migrations.ts"; +import { createModelRegistry } from "../../model-runtime-test-utils.ts"; import { createHarness } from "../harness.ts"; describe("regression #5661: uppercase models.json header values", () => { @@ -79,7 +79,7 @@ describe("regression #5661: uppercase models.json header values", () => { expect(migrated.providers["my-provider"]?.apiKey).toBe("CUSTOM_API_KEY"); expect(migrated.providers["my-provider"]?.headers?.Authorization).toBe("BEARER"); - const registry = ModelRegistry.create(AuthStorage.create(join(harness.tempDir, "auth.json")), modelsPath); + const registry = await createModelRegistry(AuthStorage.create(join(harness.tempDir, "auth.json")), modelsPath); const model = registry.find("my-provider", "my-model"); expect(model).toBeDefined(); expect(await registry.getApiKeyAndHeaders(model!)).toMatchObject({ diff --git a/packages/coding-agent/test/test-harness.test.ts b/packages/coding-agent/test/test-harness.test.ts index 2d6e7866..0c4a151f 100644 --- a/packages/coding-agent/test/test-harness.test.ts +++ b/packages/coding-agent/test/test-harness.test.ts @@ -17,7 +17,7 @@ describe("test harness", () => { }); it("simple text response", async () => { - harness = createHarness({ responses: ["hello world"] }); + harness = await createHarness({ responses: ["hello world"] }); await harness.session.prompt("hi"); @@ -32,7 +32,7 @@ describe("test harness", () => { }); it("response sequence", async () => { - harness = createHarness({ responses: ["first", "second", "third"] }); + harness = await createHarness({ responses: ["first", "second", "third"] }); await harness.session.prompt("a"); await harness.session.prompt("b"); @@ -60,7 +60,7 @@ describe("test harness", () => { }, }; - harness = createHarness({ + harness = await createHarness({ responses: [{ toolCalls: [{ name: "echo", args: { text: "hi" } }] }, "done after tool"], tools: [echoTool], baseToolsOverride: { echo: echoTool }, @@ -76,7 +76,7 @@ describe("test harness", () => { }); it("error response", async () => { - harness = createHarness({ + harness = await createHarness({ responses: [{ error: "something broke" }], }); @@ -89,7 +89,7 @@ describe("test harness", () => { }); it("retry on transient error", async () => { - harness = createHarness({ + harness = await createHarness({ responses: [{ error: "overloaded_error" }, "recovered"], settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } }, }); @@ -107,7 +107,7 @@ describe("test harness", () => { }); it("custom usage numbers", async () => { - harness = createHarness({ + harness = await createHarness({ responses: [{ text: "big response", usage: { input: 100000, output: 5000 } }], }); @@ -119,7 +119,7 @@ describe("test harness", () => { }); it("event capture", async () => { - harness = createHarness({ responses: ["hello"] }); + harness = await createHarness({ responses: ["hello"] }); await harness.session.prompt("hi"); @@ -134,7 +134,7 @@ describe("test harness", () => { }); it("context capture", async () => { - harness = createHarness({ responses: ["reply"] }); + harness = await createHarness({ responses: ["reply"] }); await harness.session.prompt("my question"); @@ -145,7 +145,7 @@ describe("test harness", () => { }); it("wraps around when more calls than responses", async () => { - harness = createHarness({ responses: ["a", "b"] }); + harness = await createHarness({ responses: ["a", "b"] }); await harness.session.prompt("1"); await harness.session.prompt("2"); @@ -161,7 +161,7 @@ describe("test harness", () => { }); it("streams text deltas", async () => { - harness = createHarness({ responses: ["hello world"] }); + harness = await createHarness({ responses: ["hello world"] }); await harness.session.prompt("hi"); @@ -175,7 +175,7 @@ describe("test harness", () => { }); it("streams thinking deltas", async () => { - harness = createHarness({ + harness = await createHarness({ responses: [{ thinking: "let me think about this", text: "answer" }], }); @@ -203,7 +203,7 @@ describe("test harness", () => { execute: async () => ({ content: [{ type: "text", text: "echoed" }], details: {} }), }; - harness = createHarness({ + harness = await createHarness({ responses: [{ toolCalls: [{ name: "echo", args: { text: "hi" } }] }, "done"], tools: [echoTool], baseToolsOverride: { echo: echoTool }, @@ -230,7 +230,7 @@ describe("test harness", () => { execute: async () => ({ content: [{ type: "text", text: "echoed" }], details: {} }), }; - harness = createHarness({ + harness = await createHarness({ responses: [ { thinking: "hmm", @@ -310,7 +310,7 @@ describe("test harness", () => { }); it("session persistence works", async () => { - harness = createHarness({ responses: ["persisted"] }); + harness = await createHarness({ responses: ["persisted"] }); await harness.session.prompt("hi"); diff --git a/packages/coding-agent/test/test-harness.ts b/packages/coding-agent/test/test-harness.ts index ea4fe246..3dc9f628 100644 --- a/packages/coding-agent/test/test-harness.ts +++ b/packages/coding-agent/test/test-harness.ts @@ -1,3 +1,4 @@ +import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; /** * Test harness for AgentSession runtime testing. * @@ -28,7 +29,6 @@ import type { import { createAssistantMessageEventStream } from "@earendil-works/pi-ai"; import { AgentSession, type AgentSessionEvent } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import type { Settings } from "../src/core/settings-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; @@ -361,11 +361,11 @@ function createTempDir(): string { return tempDir; } -function createHarnessWithResourceLoader( +async function createHarnessWithResourceLoader( options: HarnessOptions, resourceLoader: ResourceLoader, tempDir: string, -): Harness { +): Promise { const baseModel = options.model ?? fauxModel; const model: Model = options.contextWindow ? { ...baseModel, contextWindow: options.contextWindow } : baseModel; @@ -389,15 +389,32 @@ function createHarnessWithResourceLoader( } const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - authStorage.setRuntimeApiKey(model.provider, "faux-key"); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); + await authStorage.modify(model.provider, async () => ({ type: "api_key", key: "faux-key" })); + const modelRegistry = await createModelRegistry(authStorage, tempDir); + modelRegistry.registerProvider(model.provider, { + baseUrl: model.baseUrl, + api: model.api, + models: [ + { + id: model.id, + name: model.name, + api: model.api, + reasoning: model.reasoning, + input: model.input, + cost: model.cost, + contextWindow: model.contextWindow, + maxTokens: model.maxTokens, + baseUrl: model.baseUrl, + }, + ], + }); const session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader, baseToolsOverride: options.baseToolsOverride, }); @@ -429,18 +446,18 @@ function createHarnessWithResourceLoader( }; } -export function createHarness(options: HarnessOptions = {}): Harness { +export async function createHarness(options: HarnessOptions = {}): Promise { if (options.extensionFactories?.length) { throw new Error("createHarness does not support extensionFactories. Use createHarnessWithExtensions()."); } const tempDir = createTempDir(); - return createHarnessWithResourceLoader(options, options.resourceLoader ?? createTestResourceLoader(), tempDir); + return await createHarnessWithResourceLoader(options, options.resourceLoader ?? createTestResourceLoader(), tempDir); } export async function createHarnessWithExtensions(options: HarnessOptions = {}): Promise { const tempDir = createTempDir(); const extensionsResult = await createTestExtensionsResult(options.extensionFactories ?? [], tempDir); const resourceLoader = options.resourceLoader ?? createTestResourceLoader({ extensionsResult }); - return createHarnessWithResourceLoader(options, resourceLoader, tempDir); + return await createHarnessWithResourceLoader(options, resourceLoader, tempDir); } diff --git a/packages/coding-agent/test/utilities.ts b/packages/coding-agent/test/utilities.ts index c76e89eb..f7568e9f 100644 --- a/packages/coding-agent/test/utilities.ts +++ b/packages/coding-agent/test/utilities.ts @@ -1,3 +1,4 @@ +import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; /** * Shared test utilities for coding-agent tests. */ @@ -6,8 +7,9 @@ import { chmodSync, existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } import { homedir, tmpdir } from "node:os"; import { dirname, join } from "node:path"; import { Agent } from "@earendil-works/pi-agent-core"; -import { getModel, type OAuthCredentials, type OAuthProvider } from "@earendil-works/pi-ai/compat"; -import { getOAuthApiKey } from "@earendil-works/pi-ai/oauth"; +import type { OAuthCredentials } from "@earendil-works/pi-ai"; +import { getModel } from "@earendil-works/pi-ai/compat"; +import { builtinProviders } from "@earendil-works/pi-ai/providers/all"; import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; import { createEventBus } from "../src/core/event-bus.ts"; @@ -18,7 +20,6 @@ import type { LoadExtensionsResult, } from "../src/core/extensions/index.ts"; import { createExtensionRuntime, loadExtensionFromFactory } from "../src/core/extensions/loader.ts"; -import { ModelRegistry } from "../src/core/model-registry.ts"; import type { ResourceLoader } from "../src/core/resource-loader.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; @@ -88,23 +89,15 @@ export async function resolveApiKey(provider: string): Promise = {}; - for (const [key, value] of Object.entries(storage)) { - if (value.type === "oauth") { - const { type: _, ...creds } = value; - oauthCredentials[key] = creds; - } + const oauth = builtinProviders().find((candidate) => candidate.id === provider)?.auth.oauth; + if (!oauth) return undefined; + let credential = entry; + if (Date.now() >= credential.expires) { + credential = await oauth.refresh(credential); + storage[provider] = credential; + saveAuthStorage(storage); } - - const result = await getOAuthApiKey(provider as OAuthProvider, oauthCredentials); - if (!result) return undefined; - - // Save refreshed credentials back to auth.json - storage[provider] = { type: "oauth", ...result.newCredentials }; - saveAuthStorage(storage); - - return result.apiKey; + return (await oauth.toAuth(credential)).apiKey; } return undefined; @@ -241,7 +234,7 @@ export function createTestResourceLoader(options: CreateTestResourceLoaderOption * Create an AgentSession for testing with proper setup and cleanup. * Use this for e2e tests that need real LLM calls. */ -export function createTestSession(options: TestSessionOptions = {}): TestSessionContext { +export async function createTestSession(options: TestSessionOptions = {}): Promise { const tempDir = join(tmpdir(), `pi-test-${Date.now()}-${Math.random().toString(36).slice(2)}`); mkdirSync(tempDir, { recursive: true }); @@ -263,14 +256,14 @@ export function createTestSession(options: TestSessionOptions = {}): TestSession } const authStorage = AuthStorage.create(join(tempDir, "auth.json")); - const modelRegistry = ModelRegistry.create(authStorage, tempDir); + const modelRegistry = await createModelRegistry(authStorage, tempDir); const session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, - modelRegistry, + modelRuntime: getModelRuntime(modelRegistry), resourceLoader: createTestResourceLoader(), }); diff --git a/packages/coding-agent/vitest.config.ts b/packages/coding-agent/vitest.config.ts index 0582b7df..2c6b8cf3 100644 --- a/packages/coding-agent/vitest.config.ts +++ b/packages/coding-agent/vitest.config.ts @@ -4,6 +4,7 @@ import { defineConfig } from "vitest/config"; const aiSrcIndex = fileURLToPath(new URL("../ai/src/index.ts", import.meta.url)); const aiSrcCompat = fileURLToPath(new URL("../ai/src/compat.ts", import.meta.url)); const aiSrcOAuth = fileURLToPath(new URL("../ai/src/oauth.ts", import.meta.url)); +const aiSrcProviders = fileURLToPath(new URL("../ai/src/providers", import.meta.url)); const agentSrcIndex = fileURLToPath(new URL("../agent/src/index.ts", import.meta.url)); const tuiSrcIndex = fileURLToPath(new URL("../tui/src/index.ts", import.meta.url)); @@ -25,6 +26,7 @@ export default defineConfig({ { find: /^@earendil-works\/pi-ai$/, replacement: aiSrcIndex }, { find: /^@earendil-works\/pi-ai\/compat$/, replacement: aiSrcCompat }, { find: /^@earendil-works\/pi-ai\/oauth$/, replacement: aiSrcOAuth }, + { find: /^@earendil-works\/pi-ai\/providers\/(.+)$/, replacement: `${aiSrcProviders}/$1.ts` }, { find: /^@earendil-works\/pi-agent-core$/, replacement: agentSrcIndex }, { find: /^@earendil-works\/pi-tui$/, replacement: tuiSrcIndex }, { find: /^@mariozechner\/pi-ai$/, replacement: aiSrcIndex }, diff --git a/packages/orchestrator/src/radius.ts b/packages/orchestrator/src/radius.ts index 2d6347f9..c3ac35f6 100644 --- a/packages/orchestrator/src/radius.ts +++ b/packages/orchestrator/src/radius.ts @@ -1,5 +1,6 @@ import { hostname, platform } from "node:os"; -import { AuthStorage, type OAuthCredential } from "@earendil-works/pi-coding-agent"; +import type { OAuthCredential } from "@earendil-works/pi-ai"; +import { readStoredCredential } from "@earendil-works/pi-coding-agent"; import { getOrchestratorDir, getSocketPath, VERSION } from "./config.ts"; import { loadMachine, saveMachine } from "./storage.ts"; import type { InstanceRecord, MachineRecord, RadiusRegistration } from "./types.ts"; @@ -116,15 +117,9 @@ export function getRadiusOrchestratorBaseUrl(): string { return new URL(DEFAULT_ORCHESTRATOR_BASE_PATH, getRadiusUrl()).toString(); } -const radiusAuthStorage = AuthStorage.create(); - function getStoredRadiusCredential(): OAuthCredential | undefined { - radiusAuthStorage.reload(); - const credential = radiusAuthStorage.get(RADIUS_PROVIDER); - if (!credential || credential.type !== "oauth") { - return undefined; - } - return credential; + const credential = readStoredCredential(RADIUS_PROVIDER); + return credential?.type === "oauth" ? credential : undefined; } export function getRadiusAccessToken(): string {