ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(coding-agent): replace model registry with model runtime

Browse Source 
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.
This commit is contained in:
Mario Zechner
2026-07-14 17:48:45 +02:00
parent 6731a0ba9e
commit 9993c96907
133 changed files with 5103 additions and 4340 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/CHANGELOG.md
+65
View File
@@ -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<void>` 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
packages/coding-agent/README.md
+3 -5
View File
@@ -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?");
packages/coding-agent/docs/custom-provider.md
+6 -11
View File
@@ -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<string>;
@@ -660,7 +656,6 @@ interface ProviderConfig {
login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials>;
refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials>;
getApiKey(credentials: OAuthCredentials): string;
modifyModels?(models: Model<Api>[], credentials: OAuthCredentials): Model<Api>[];
};
}
```
packages/coding-agent/docs/sdk.md
+41 -51
View File
@@ -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
packages/coding-agent/examples/extensions/custom-provider-anthropic/index.ts
+1 -1
View File
@@ -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);
packages/coding-agent/examples/sdk/02-custom-model.ts
+5 -8
View File
@@ -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 {
packages/coding-agent/examples/sdk/09-api-keys-and-oauth.ts
+13 -31
View File
@@ -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();
packages/coding-agent/examples/sdk/12-full-control.ts
+7 -12
View File
@@ -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),
packages/coding-agent/examples/sdk/README.md
+14 -18
View File
@@ -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 |
packages/coding-agent/src/cli/list-models.ts
+4 -4
View File
@@ -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<void> {
const loadError = modelRegistry.getError();
export async function listModels(modelRuntime: ModelRuntime, searchPattern?: string): Promise<void> {
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());
packages/coding-agent/src/core/agent-session-services.ts
+12 -13
View File
@@ -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<string, boolean | string>;
resourceLoaderOptions?: Omit<DefaultResourceLoaderOptions, "cwd" | "agentDir" | "settingsManager">;
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<AgentSessionServices> {
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,
packages/coding-agent/src/core/agent-session.ts
+68 -35
View File
@@ -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<string, string> | 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<string, AgentTool> = 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<any>): Promise<{
@@ -377,18 +390,25 @@ export class AgentSession {
headers?: Record<string, string>;
env?: Record<string, string>;
}> {
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<any>): Promise<void> {
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<ModelCycleResult | undefined> {
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<ModelCycleResult | undefined> {
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<string, string> | undefined;
let env: Record<string, string> | 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;
packages/coding-agent/src/core/auth-storage.ts
+50 -318
View File
@@ -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<string, string>;
};
export type OAuthCredential = {
type: "oauth";
} & OAuthCredentials;
export type AuthCredential = ApiKeyCredential | OAuthCredential;
export type AuthStorageData = Record<string, AuthCredential>;
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<string, Credential>;
type LockResult<T> = {
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<string, string> = 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<Credential | undefined> {
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<string, string> | 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<void> {
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<Credential | undefined>,
): Promise<Credential | undefined> {
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<string, OAuthCredentials> = {};
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<string | undefined> {
// 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<void> {
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<readonly CredentialInfo[]> {
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;
}
}
packages/coding-agent/src/core/cache-stats.ts
+3 -3
View File
@@ -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,
packages/coding-agent/src/core/extensions/loader.ts
+11 -2
View File
@@ -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<string, unknown> = {
"@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<string, string> {
// 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,
packages/coding-agent/src/core/extensions/runner.ts
+4
View File
@@ -602,6 +602,10 @@ export class ExtensionRunner {
});
}
getModelRegistry(): ModelRegistry {
return this.modelRegistry;
}
getRegisteredCommands(): ResolvedCommand[] {
this.commandDiagnostics = [];
return this.resolveRegisteredCommands();
packages/coding-agent/src/core/extensions/types.ts
+2 -2
View File
@@ -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<OAuthCredentials>;
/** Refresh expired credentials, return updated credentials to persist. */
refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials>;
/** 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<Api>[], credentials: OAuthCredentials): Model<Api>[];
};
}
packages/coding-agent/src/core/model-config.ts
+277
View File
@@ -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<typeof ModelDefinitionSchema>;
export type ModelsJsonModelOverride = Static<typeof ModelOverrideSchema>;
export type ModelsJsonProvider = Static<typeof ProviderConfigSchema>;
type ModelsJson = Static<typeof ModelsConfigSchema>;
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<T>(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<string, ModelsJsonProvider>;
private readonly error: string | undefined;
private constructor(providers: ReadonlyMap<string, ModelsJsonProvider>, error?: string) {
this.providers = providers;
this.error = error;
}
static async load(modelsJsonPath: string | undefined): Promise<ModelConfig> {
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<string, ModelsJsonProvider>();
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;
}
}
packages/coding-agent/src/core/model-registry.ts
+59 -951
View File
File diff suppressed because it is too large Load Diff
packages/coding-agent/src/core/model-resolver.ts
+20 -20
View File
@@ -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<KnownProvider, string> = {
@@ -268,9 +268,9 @@ export interface ResolveModelScopeResult {
export async function resolveModelScopeWithDiagnostics(
patterns: string[],
modelRegistry: ModelRegistry,
modelRuntime: ModelRuntime,
): Promise<ResolveModelScopeResult> {
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<ScopedModel[]> {
const { scopedModels, diagnostics } = await resolveModelScopeWithDiagnostics(patterns, modelRegistry);
export async function resolveModelScope(patterns: string[], modelRuntime: ModelRuntime): Promise<ScopedModel[]> {
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<InitialModelResult> {
const {
cliProvider,
@@ -565,7 +565,7 @@ export async function findInitialModel(options: {
defaultProvider,
defaultModelId,
defaultThinkingLevel,
modelRegistry,
modelRuntime,
} = options;
let model: Model<Api> | 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<Api> | undefined,
shouldPrintMessages: boolean,
modelRegistry: ModelRegistry,
modelRuntime: ModelRuntime,
): Promise<{ model: Model<Api> | 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
packages/coding-agent/src/core/model-runtime.ts
+489
View File
@@ -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<Api>[];
available: readonly Model<Api>[];
configuredProviders: ReadonlySet<string>;
storedProviders: ReadonlySet<string>;
auth: ReadonlyMap<string, AuthCheck | undefined>;
}
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<string, string>;
}
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<string, Provider>;
private readonly extensionProviders = new Map<string, ProviderConfigInput>();
private readonly compositionErrors = new Map<string, string>();
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<void> | 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<ModelRuntime> {
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<string> {
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<void> {
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<void> | undefined): Promise<void> {
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<void> {
return this.availabilityRefresh ?? this.queueAvailabilityRefresh(undefined);
}
/** Mutations must not observe an in-flight refresh started before them. */
private forceRefreshAvailability(): Promise<void> {
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<Api>[] {
return this.models.getModels(providerId);
}
getModel(providerId: string, modelId: string): Model<Api> | undefined {
return this.models.getModel(providerId, modelId);
}
async checkAuth(providerId: string): Promise<AuthCheck | undefined> {
return this.models.checkAuth(providerId);
}
async getAvailable(providerId?: string): Promise<readonly Model<Api>[]> {
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<Api>[] {
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<Api>): 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<AuthResult | undefined>;
getAuth(model: Model<Api>, overrides?: ModelRuntimeAuthOverrides): Promise<AuthResult | undefined>;
async getAuth(
providerOrModel: string | Model<Api>,
overrides: ModelRuntimeAuthOverrides = {},
): Promise<AuthResult | undefined> {
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<readonly CredentialInfo[]> {
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<Api>,
options: (StreamOptions & ModelsStreamTransforms) | undefined,
): Promise<{ provider: Provider; model: Model<Api>; options: StreamOptions }> {
const provider = this.models.getProvider(model.provider);
if (!provider) throw new ModelsError("provider", `Unknown provider: ${model.provider}`);
const resolution = await this.getAuth(model, { apiKey: options?.apiKey, env: options?.env });
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<TApi extends Api>(
model: Model<TApi>,
context: Context,
options?: ModelsApiStreamOptions<TApi>,
): AssistantMessageEventStream {
return lazyStream(model, async () => {
const prepared = await this.prepareRequest(
model,
options as (StreamOptions & ModelsStreamTransforms) | undefined,
);
return prepared.provider.stream(
prepared.model as Model<TApi>,
context,
prepared.options as ApiStreamOptions<TApi>,
);
});
}
complete<TApi extends Api>(
model: Model<TApi>,
context: Context,
options?: ModelsApiStreamOptions<TApi>,
): Promise<AssistantMessage> {
return this.stream(model, context, options).result();
}
streamSimple(model: Model<Api>, 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<Api>, context: Context, options?: ModelsSimpleStreamOptions): Promise<AssistantMessage> {
return this.streamSimple(model, context, options).result();
}
async login(providerId: string, type: AuthType, interaction: AuthInteraction): Promise<Credential> {
const credential = await this.models.login(providerId, type, interaction);
await this.forceRefreshAvailability();
return credential;
}
async logout(providerId: string): Promise<void> {
await this.models.logout(providerId);
await this.forceRefreshAvailability();
}
async reloadConfig(): Promise<void> {
this.config = await ModelConfig.load(this.modelsPath);
this.rebuildProviders();
await this.forceRefreshAvailability();
}
async refresh(providerId?: string): Promise<void> {
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<string, unknown>)[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(() => {});
}
}
packages/coding-agent/src/core/provider-composer.ts
+513
View File
@@ -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<OAuthCredentials>;
refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials>;
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<Api>, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream;
headers?: Record<string, string>;
authHeader?: boolean;
oauth?: ExtensionOAuthConfig;
models?: Array<{
id: string;
name: string;
api?: Api;
baseUrl?: string;
reasoning: boolean;
thinkingLevelMap?: Model<Api>["thinkingLevelMap"];
input: ("text" | "image")[];
cost: Model<Api>["cost"];
contextWindow: number;
maxTokens: number;
headers?: Record<string, string>;
compat?: Model<Api>["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<Api>["compat"],
override: Model<Api>["compat"] | ModelsJsonModelOverride["compat"],
): Model<Api>["compat"] {
if (!override) return base;
const merged = { ...base, ...override } as NonNullable<Model<Api>["compat"]>;
const baseNested = base as Record<string, unknown> | undefined;
const overrideNested = override as Record<string, unknown>;
const mergedNested = merged as Record<string, unknown>;
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<Api>, override: ModelsJsonModelOverride): Model<Api> {
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<Api> | undefined,
): Model<Api> {
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<Api>[],
config: ModelsJsonProvider | undefined,
): Model<Api>[] {
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<Api>[] = 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<Api>[],
config: ProviderConfigInput | undefined,
): Model<Api>[] {
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<string, string> | 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<string, string> | undefined {
if (!config?.headers && !extension?.headers) return undefined;
return { ...config?.headers, ...extension?.headers };
}
async function configContextEnv(
values: readonly string[],
ctx: AuthContext,
explicit?: Record<string, string>,
): Promise<Record<string, string> | 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<string, string>) : undefined,
);
return withConfiguredAuth(auth, headers, authHeader);
},
};
}
function rawModelHeaders(
model: Model<Api>,
config: ModelsJsonProvider | undefined,
extension: ProviderConfigInput | undefined,
): Record<string, string> | 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<Api>) => base?.getModels().some((entry) => entry.api === model.api) ?? false;
const streamWith = (
model: Model<Api>,
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<Api>,
config: ModelsJsonProvider | undefined,
extension: ProviderConfigInput | undefined,
env?: Record<string, string>,
): Record<string, string> | 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<Api>,
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" };
}
packages/coding-agent/src/core/provider-display-names.ts
-35
View File
@@ -1,35 +0,0 @@
export const BUILT_IN_PROVIDER_DISPLAY_NAMES: Record<string, string> = {
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)",
};
packages/coding-agent/src/core/runtime-credentials.ts
+48
View File
@@ -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<string, string>();
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<Credential | undefined> {
const override = this.overrides.get(providerId);
return override ? { type: "api_key", key: override } : this.store.read(providerId);
}
async list(): Promise<readonly CredentialInfo[]> {
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<Credential | undefined>,
): Promise<Credential | undefined> {
return this.store.modify(providerId, fn);
}
async delete(providerId: string): Promise<void> {
this.overrides.delete(providerId);
await this.store.delete(providerId);
}
}
packages/coding-agent/src/core/sdk.ts
+21 -35
View File
@@ -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<any>;
@@ -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,
packages/coding-agent/src/index.ts
+6 -11
View File
@@ -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,
packages/coding-agent/src/main.ts
+12 -14
View File
@@ -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);
}
packages/coding-agent/src/modes/interactive/components/footer.ts
+1 -1
View File
@@ -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);
packages/coding-agent/src/modes/interactive/components/login-dialog.ts
+13 -12
View File
@@ -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();
}
packages/coding-agent/src/modes/interactive/components/model-selector.ts
+8 -8
View File
@@ -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<any>;
private settingsManager: SettingsManager;
private modelRegistry: ModelRegistry;
private modelRuntime: ModelRuntime;
private onSelectCallback: (model: Model<any>) => void;
private onCancelCallback: () => void;
private errorMessage?: string;
@@ -66,7 +66,7 @@ export class ModelSelectorComponent extends Container implements Focusable {
tui: TUI,
currentModel: Model<any> | undefined,
settingsManager: SettingsManager,
modelRegistry: ModelRegistry,
modelRuntime: ModelRuntime,
scopedModels: ReadonlyArray<ScopedModelItem>,
onSelect: (model: Model<any>) => 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<any>) => ({
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) => ({
packages/coding-agent/src/modes/interactive/components/oauth-selector.ts
+21 -28
View File
@@ -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 {
packages/coding-agent/src/modes/interactive/interactive-mode.ts
+133 -194
View File
@@ -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<string>(getProviders());
export function isApiKeyLoginProvider(
providerId: string,
oauthProviderIds: ReadonlySet<string>,
builtInProviderIds: ReadonlySet<string> = 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<AutocompleteItem[] | null> => {
// 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<AssistantMessage, CacheMiss>();
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<void> {
// 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<AuthSelectorProvider[]> {
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<void> {
await this.session.modelRuntime.getAvailable();
if (!providerRef) {
this.showLoginAuthTypeSelector();
return;
@@ -4873,10 +4843,10 @@ export class InteractiveMode {
private async startProviderLogin(providerOption: AuthSelectorProvider): Promise<void> {
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<any> | undefined,
): Promise<void> {
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<any> | 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<string | undefined> {
return new Promise((resolve) => {
private showAuthSelect(
dialog: LoginDialogComponent,
prompt: Extract<AuthPrompt, { type: "select" }>,
): Promise<string> {
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<string> {
let response: Promise<string>;
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<string>((_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<void> {
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<void> {
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<string>((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.
packages/coding-agent/src/modes/rpc/rpc-mode.ts
+2 -2
View File
@@ -465,7 +465,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
// =================================================================
case "set_model": {
const models = await session.modelRegistry.getAvailable();
const models = await session.modelRuntime.getAvailable();
const model = models.find((m) => m.provider === command.provider && m.id === command.modelId);
if (!model) {
return error(id, "set_model", `Model not found: ${command.provider}/${command.modelId}`);
@@ -483,7 +483,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
}
case "get_available_models": {
const models = await session.modelRegistry.getAvailable();
const models = await session.modelRuntime.getAvailable();
return success(id, "get_available_models", { models });
}
packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts
+5 -5
View File
@@ -7,9 +7,9 @@ import { getModel } from "@earendil-works/pi-ai/compat";
import { afterEach, beforeEach, describe, expect, it, vi } 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";
describe("AgentSession auto-compaction queue resume", () => {
@@ -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(),
});
});
packages/coding-agent/test/agent-session-branching.test.ts
+1 -1
View File
@@ -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,
packages/coding-agent/test/agent-session-compaction.test.ts
+10 -10
View File
@@ -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");
packages/coding-agent/test/agent-session-concurrent.test.ts
+21 -21
View File
@@ -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 },
});
packages/coding-agent/test/agent-session-dynamic-provider.test.ts
+7 -2
View File
@@ -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,
});
packages/coding-agent/test/agent-session-retry.test.ts
+20 -16
View File
@@ -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<AssistantMessageEvent, AssistantMessage> {
@@ -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 },
});
packages/coding-agent/test/agent-session-runtime-events.test.ts
+25 -2
View File
@@ -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],
packages/coding-agent/test/agent-session-stats.test.ts
+12 -12
View File
@@ -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));
packages/coding-agent/test/agent-session-tree-navigation.test.ts
+4 -4
View File
@@ -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.",
});
});
packages/coding-agent/test/auth-storage.test.ts
+172 -654
View File
@@ -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<string, unknown>) {
function writeAuthJson(data: Record<string, unknown>): 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<string, { key: string }>;
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<string, { key: string }>;
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<string, { key: string }>;
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");
});
});
packages/coding-agent/test/cache-stats.test.ts
+1 -1
View File
@@ -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: {
packages/coding-agent/test/compaction-extensions.test.ts
+15 -16
View File
@@ -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();
packages/coding-agent/test/config-value-migration.test.ts
+5 -4
View File
@@ -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"),
);
packages/coding-agent/test/extensions-discovery.test.ts
+36
View File
@@ -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);
packages/coding-agent/test/extensions-input-event.test.ts
+3 -2
View File
@@ -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);
}
packages/coding-agent/test/extensions-runner.test.ts
+6 -5
View File
@@ -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", () => {
packages/coding-agent/test/footer-width.test.ts
+1 -1
View File
@@ -52,7 +52,7 @@ function createSession(options: {
getCwd: () => "/tmp/project",
},
getContextUsage: () => ({ contextWindow: 200_000, percent: 12.3 }),
modelRegistry: {
modelRuntime: {
isUsingOAuth: () => false,
},
};
packages/coding-agent/test/interactive-mode-anthropic-warning.test.ts
+20 -37
View File
@@ -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();
});
});
packages/coding-agent/test/interactive-mode-status.test.ts
+4 -4
View File
@@ -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: [] }) },
packages/coding-agent/test/model-registry.test.ts
+238 -150
View File
File diff suppressed because it is too large Load Diff
packages/coding-agent/test/model-resolver.test.ts
+58 -58
View File
@@ -260,12 +260,12 @@ describe("resolveModelScopeWithDiagnostics", () => {
describe("resolveCliModel", () => {
test("resolves --model provider/id without --provider", () => {
const registry = {
getAll: () => allModels,
} as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => allModels,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => allModels,
} as unknown as Parameters<typeof resolveCliModel>[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 <pattern>:<thinking> (without explicit --thinking)", () => {
const registry = {
getAll: () => allModels,
} as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => allModels,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => allModels,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => allModels,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => allModels,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => [],
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => [...allModels, commandcodeModel, xiaomiModel],
hasConfiguredAuth: (provider: string) => provider === "commandcode",
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => allModels,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => modelsWithNeuralwatt,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => modelsWithNeuralwatt,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => modelsWithNeuralwatt,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => modelsWithNeuralwatt,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => modelsWithNeuralwatt,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof resolveCliModel>[0]["modelRegistry"];
getModels: () => modelsWithNeuralwatt,
} as unknown as Parameters<typeof resolveCliModel>[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<typeof findInitialModel>[0]["modelRegistry"];
getModels: () => allModels,
} as unknown as Parameters<typeof findInitialModel>[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<typeof findInitialModel>[0]["modelRegistry"];
} as unknown as Parameters<typeof findInitialModel>[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<typeof findInitialModel>[0]["modelRegistry"];
} as unknown as Parameters<typeof findInitialModel>[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");
packages/coding-agent/test/model-runtime-auth-options.test.ts
+256
View File
@@ -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<string, string | null> | 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<string, string | null> | 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" },
});
});
});
packages/coding-agent/test/model-runtime-cloudflare-compat.test.ts
+95
View File
@@ -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<typeof stream> & {
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<string, unknown>;
};
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<string, unknown>;
};
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");
});
});
packages/coding-agent/test/model-runtime-test-utils.ts
+25
View File
@@ -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<ModelRegistry, ModelRuntime>();
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<ModelRegistry> {
return wrap(await ModelRuntime.create({ credentials, modelsPath }));
}
export async function createInMemoryModelRegistry(credentials: CredentialStore): Promise<ModelRegistry> {
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;
}
packages/coding-agent/test/oauth-selector.test.ts
+90 -80
View File
@@ -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");
});
});
packages/coding-agent/test/resolve-config-value.test.ts
+121
View File
@@ -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);
}
});
});
packages/coding-agent/test/resource-loader.test.ts
+4 -3
View File
@@ -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,
packages/coding-agent/test/rpc-prompt-response-semantics.test.ts
+7 -7
View File
@@ -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<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
function createRuntimeHost(options: { withAuth: boolean; responseDelayMs: number; model?: Model<any> }): {
async function createRuntimeHost(options: { withAuth: boolean; responseDelayMs: number; model?: Model<any> }): Promise<{
runtimeHost: AgentSessionRuntime;
cleanup: () => Promise<void>;
} {
}> {
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());
packages/coding-agent/test/runtime-credentials.test.ts
+42
View File
@@ -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([]);
});
});
packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts
+4 -4
View File
@@ -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<void> {
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<void> {
models: [baseModel],
});
const modelRuntime = getModelRuntime(modelRegistry);
const settingsManager = SettingsManager.inMemory({
compaction: { enabled: false },
retry: { enabled: false },
@@ -315,8 +316,7 @@ async function main(): Promise<void> {
resourceLoader,
sessionManager: SessionManager.open(args.sessionPath),
settingsManager,
authStorage,
modelRegistry,
modelRuntime,
});
session.setActiveToolsByName(["deterministic_probe"]);
packages/coding-agent/test/sdk-openrouter-attribution.test.ts
+12 -17
View File
@@ -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);
}
}
packages/coding-agent/test/sdk-stream-options.test.ts
+42 -7
View File
@@ -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<SimpleStreamOptions | undefined> {
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");
});
});
packages/coding-agent/test/suite/agent-session-runtime.test.ts
+4 -4
View File
@@ -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,
packages/coding-agent/test/suite/harness.ts
+4 -4
View File
@@ -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<Harne
const authStorage = AuthStorage.inMemory();
if (withConfiguredAuth) {
authStorage.setRuntimeApiKey(model.provider, "faux-key");
await authStorage.modify(model.provider, async () => ({ 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<Harne
sessionManager,
settingsManager,
cwd: tempDir,
modelRegistry,
modelRuntime: getModelRuntime(modelRegistry),
resourceLoader,
baseToolsOverride: toolMap,
initialActiveToolNames: options.initialActiveToolNames,
packages/coding-agent/test/suite/regressions/2753-reload-stale-resource-settings.test.ts
+7 -2
View File
@@ -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";
describe("issue #2753 reload stale resource settings", () => {
@@ -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) => {
packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts
+7 -2
View File
@@ -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) => {
packages/coding-agent/test/suite/regressions/3217-scoped-model-order.test.ts
+5 -7
View File
@@ -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<void> {
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")
packages/coding-agent/test/suite/regressions/5433-extension-oauth-prompt-input.test.ts
+14
View File
@@ -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();
packages/coding-agent/test/suite/regressions/5596-missing-theme-export.test.ts
+4 -4
View File
@@ -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(() => {
packages/coding-agent/test/suite/regressions/5661-uppercase-header-values.test.ts
+2 -2
View File
@@ -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({
packages/coding-agent/test/test-harness.test.ts
+14 -14
View File
@@ -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");
packages/coding-agent/test/test-harness.ts
+26 -9
View File
@@ -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<Harness> {
const baseModel = options.model ?? fauxModel;
const model: Model<any> = 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<Harness> {
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<Harness> {
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);
}
packages/coding-agent/test/utilities.ts
+15 -22
View File
@@ -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<string | undefine
}
if (entry.type === "oauth") {
// Build OAuthCredentials record for getOAuthApiKey
const oauthCredentials: Record<string, OAuthCredentials> = {};
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<TestSessionContext> {
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(),
});
packages/coding-agent/vitest.config.ts
+2
View File
@@ -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 },