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/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(),
});