Archived
Merge pull request #64 from jmfederico/fix/issue-62-authstorage
fix: migrate auth/model plumbing to ModelRuntime (fixes #62)
This commit is contained in:
Federico Jaramillo Martinez
@@ -1,52 +1,70 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getLoginProviderOptions, getLogoutProviderOptions, isApiKeyLoginProvider, type AuthProviderModelRegistry } from "./authProviderOptions";
|
||||
import { getLoginProviderOptions, getLogoutProviderOptions, type AuthProviderRuntime } from "./authProviderOptions";
|
||||
|
||||
function registry(): AuthProviderModelRegistry {
|
||||
const credentials = new Map<string, { type: "oauth" | "api_key" }>();
|
||||
credentials.set("openai", { type: "api_key" });
|
||||
function runtime(configuredProviders: ReadonlySet<string> = new Set(["openai"])): AuthProviderRuntime {
|
||||
const credentials = [{ providerId: "openai", type: "api_key" as const }];
|
||||
// Auth shapes mirror what the Pi SDK actually reports for these providers:
|
||||
// github-copilot supports both methods, openai-codex is OAuth-only, and
|
||||
// ambient providers resolve credentials without offering interactive login.
|
||||
const providers = [
|
||||
{ id: "anthropic", name: "Anthropic", auth: { oauth: {}, apiKey: { login: () => undefined } } },
|
||||
{ id: "github-copilot", name: "GitHub Copilot", auth: { oauth: {}, apiKey: { login: () => undefined } } },
|
||||
{ id: "openai-codex", name: "ChatGPT Plus/Pro (Codex Subscription)", auth: { oauth: {} } },
|
||||
{ id: "openai", name: "OpenAI", auth: { apiKey: { login: () => undefined } } },
|
||||
{ id: "custom", name: "Custom", auth: { apiKey: { login: () => undefined } } },
|
||||
{ id: "cloudflare-ai-gateway", name: "Cloudflare AI Gateway", auth: { apiKey: { login: () => undefined } } },
|
||||
{ id: "cloudflare-workers-ai", name: "Cloudflare Workers AI", auth: { apiKey: { login: () => undefined } } },
|
||||
{ id: "amazon-bedrock", name: "Amazon Bedrock", auth: { apiKey: { login: () => undefined } } },
|
||||
{ id: "google-vertex", name: "Google Vertex AI", auth: { apiKey: { login: () => undefined } } },
|
||||
{ id: "ambient", name: "Ambient credentials", auth: { apiKey: {} } },
|
||||
];
|
||||
return {
|
||||
authStorage: {
|
||||
getOAuthProviders: () => [
|
||||
{ id: "anthropic", name: "Anthropic (Claude Pro/Max)" },
|
||||
{ id: "github-copilot", name: "GitHub Copilot" },
|
||||
{ id: "openai-codex", name: "ChatGPT Plus/Pro (Codex Subscription)" },
|
||||
],
|
||||
list: () => Array.from(credentials.keys()),
|
||||
get: (provider: string) => credentials.get(provider),
|
||||
},
|
||||
getAll: () => [
|
||||
{ provider: "anthropic" },
|
||||
{ provider: "openai" },
|
||||
{ provider: "openai-codex" },
|
||||
{ provider: "github-copilot" },
|
||||
{ provider: "custom" },
|
||||
],
|
||||
getProviderDisplayName: (provider: string) => ({ anthropic: "Anthropic", openai: "OpenAI", custom: "Custom" }[provider] ?? provider),
|
||||
getProviders: () => providers,
|
||||
listCredentials: () => Promise.resolve(credentials),
|
||||
getProviderAuthStatus: (provider: string) => (provider === "openai" ? { configured: true, source: "stored" } : { configured: false }),
|
||||
hasConfiguredAuth: (provider: string) => configuredProviders.has(provider),
|
||||
};
|
||||
}
|
||||
|
||||
describe("auth provider options", () => {
|
||||
it("keeps OAuth-only providers out of API key login options", () => {
|
||||
expect(isApiKeyLoginProvider("openai-codex", new Set(["openai-codex"]))).toBe(false);
|
||||
expect(isApiKeyLoginProvider("github-copilot", new Set(["github-copilot"]))).toBe(false);
|
||||
expect(isApiKeyLoginProvider("openai", new Set(["openai-codex"]))).toBe(true);
|
||||
});
|
||||
|
||||
it("builds login options for OAuth-only, dual-auth, and API-key providers", () => {
|
||||
const options = getLoginProviderOptions(registry());
|
||||
it("offers each interactive login method reported by the backend", () => {
|
||||
const options = getLoginProviderOptions(runtime());
|
||||
expect(options).toEqual(expect.arrayContaining([
|
||||
// Dual-capable providers surface both login methods, driven purely by SDK data.
|
||||
expect.objectContaining({ id: "anthropic", authType: "oauth" }),
|
||||
expect.objectContaining({ id: "anthropic", authType: "api_key" }),
|
||||
expect.objectContaining({ id: "openai", authType: "api_key", status: { configured: true, source: "stored" } }),
|
||||
expect.objectContaining({ id: "github-copilot", authType: "oauth" }),
|
||||
expect.objectContaining({ id: "github-copilot", authType: "api_key" }),
|
||||
// OAuth-only provider surfaces only oauth.
|
||||
expect.objectContaining({ id: "openai-codex", authType: "oauth" }),
|
||||
// API-key options use the generic AuthInteraction flow, including
|
||||
// multi-field and select-first providers the legacy form cannot execute.
|
||||
expect.objectContaining({ id: "openai", authType: "api_key", loginFlow: "interactive", status: { configured: true, source: "stored" } }),
|
||||
expect.objectContaining({ id: "custom", authType: "api_key", loginFlow: "interactive" }),
|
||||
expect.objectContaining({ id: "cloudflare-ai-gateway", authType: "api_key", loginFlow: "interactive" }),
|
||||
expect.objectContaining({ id: "cloudflare-workers-ai", authType: "api_key", loginFlow: "interactive" }),
|
||||
expect.objectContaining({ id: "amazon-bedrock", authType: "api_key", loginFlow: "interactive" }),
|
||||
expect.objectContaining({ id: "google-vertex", authType: "api_key", loginFlow: "interactive" }),
|
||||
]));
|
||||
expect(options).not.toEqual(expect.arrayContaining([expect.objectContaining({ id: "openai-codex", authType: "api_key" })]));
|
||||
expect(options).not.toEqual(expect.arrayContaining([expect.objectContaining({ id: "openai", authType: "oauth" })]));
|
||||
expect(options).not.toEqual(expect.arrayContaining([expect.objectContaining({ id: "ambient", authType: "api_key" })]));
|
||||
});
|
||||
|
||||
it("returns only currently stored credentials for logout", () => {
|
||||
expect(getLogoutProviderOptions(registry())).toEqual([
|
||||
expect.objectContaining({ id: "openai", authType: "api_key" }),
|
||||
it("does not report a stored credential as configured when provider resolution is incomplete", async () => {
|
||||
const unresolvedRuntime = runtime(new Set());
|
||||
|
||||
expect(getLoginProviderOptions(unresolvedRuntime, "api_key")).toEqual(expect.arrayContaining([
|
||||
expect.objectContaining({ id: "openai", status: { configured: false } }),
|
||||
]));
|
||||
expect(await getLogoutProviderOptions(unresolvedRuntime)).toEqual([
|
||||
expect.objectContaining({ id: "openai", authType: "api_key", status: { configured: false } }),
|
||||
]);
|
||||
});
|
||||
|
||||
it("returns only currently stored credentials for logout", async () => {
|
||||
expect(await getLogoutProviderOptions(runtime())).toEqual([
|
||||
expect.objectContaining({ id: "openai", authType: "api_key", status: { configured: true, source: "stored" } }),
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,62 +1,78 @@
|
||||
import type { AuthProviderOption, AuthProviderStatus, AuthType } from "../../shared/apiTypes.js";
|
||||
|
||||
const OAUTH_ONLY_PROVIDERS = new Set(["github-copilot", "openai-codex"]);
|
||||
|
||||
export interface AuthProviderModelRegistry {
|
||||
authStorage: {
|
||||
getOAuthProviders(): { id: string; name: string }[];
|
||||
list(): string[];
|
||||
get(provider: string): { type: AuthType } | undefined;
|
||||
};
|
||||
getAll(): { provider: string }[];
|
||||
getProviderDisplayName(provider: string): string;
|
||||
getProviderAuthStatus(provider: string): AuthProviderStatus;
|
||||
/** Minimal provider shape needed to enumerate login/logout options. */
|
||||
interface AuthProviderInfo {
|
||||
id: string;
|
||||
name: string;
|
||||
auth: { apiKey?: { login?: unknown }; oauth?: unknown };
|
||||
}
|
||||
|
||||
export function getLoginProviderOptions(modelRegistry: AuthProviderModelRegistry, authType?: AuthType): AuthProviderOption[] {
|
||||
const oauthProviders = modelRegistry.authStorage.getOAuthProviders();
|
||||
const oauthProviderIds = new Set(oauthProviders.map((provider) => provider.id));
|
||||
const options: AuthProviderOption[] = oauthProviders.map((provider) => ({
|
||||
id: provider.id,
|
||||
name: provider.name,
|
||||
authType: "oauth",
|
||||
status: modelRegistry.getProviderAuthStatus(provider.id),
|
||||
}));
|
||||
/** Non-secret stored-credential metadata, keyed by provider id. */
|
||||
interface AuthProviderCredentialInfo {
|
||||
providerId: string;
|
||||
type: AuthType;
|
||||
}
|
||||
|
||||
const modelProviders = new Set(modelRegistry.getAll().map((model) => model.provider));
|
||||
for (const providerId of modelProviders) {
|
||||
if (!isApiKeyLoginProvider(providerId, oauthProviderIds)) continue;
|
||||
/**
|
||||
* Structural slice of the SDK `ModelRuntime` used to derive auth provider
|
||||
* options. Kept structural (rather than `Pick<ModelRuntime, ...>`) so tests can
|
||||
* supply a lightweight double without constructing a full runtime; a real
|
||||
* `ModelRuntime` satisfies it.
|
||||
*/
|
||||
export interface AuthProviderRuntime {
|
||||
getProviders(): readonly AuthProviderInfo[];
|
||||
listCredentials(): Promise<readonly AuthProviderCredentialInfo[]>;
|
||||
getProviderAuthStatus(providerId: string): AuthProviderStatus;
|
||||
hasConfiguredAuth(providerId: string): boolean;
|
||||
}
|
||||
|
||||
export function getLoginProviderOptions(runtime: AuthProviderRuntime, authType?: AuthType): AuthProviderOption[] {
|
||||
const providers = runtime.getProviders();
|
||||
|
||||
const options: AuthProviderOption[] = [];
|
||||
for (const provider of providers) {
|
||||
if (provider.auth.oauth === undefined) continue;
|
||||
options.push({
|
||||
id: providerId,
|
||||
name: modelRegistry.getProviderDisplayName(providerId),
|
||||
id: provider.id,
|
||||
name: provider.name,
|
||||
authType: "oauth",
|
||||
status: truthfulProviderStatus(runtime, provider.id),
|
||||
});
|
||||
}
|
||||
|
||||
for (const provider of providers) {
|
||||
if (provider.auth.apiKey?.login === undefined) continue;
|
||||
options.push({
|
||||
id: provider.id,
|
||||
name: provider.name,
|
||||
authType: "api_key",
|
||||
status: modelRegistry.getProviderAuthStatus(providerId),
|
||||
status: truthfulProviderStatus(runtime, provider.id),
|
||||
loginFlow: "interactive",
|
||||
});
|
||||
}
|
||||
|
||||
return filterAndSort(options, authType);
|
||||
}
|
||||
|
||||
export function getLogoutProviderOptions(modelRegistry: AuthProviderModelRegistry): AuthProviderOption[] {
|
||||
export async function getLogoutProviderOptions(runtime: AuthProviderRuntime): Promise<AuthProviderOption[]> {
|
||||
const providerNames = new Map(runtime.getProviders().map((provider) => [provider.id, provider.name]));
|
||||
const options: AuthProviderOption[] = [];
|
||||
for (const providerId of modelRegistry.authStorage.list()) {
|
||||
const credential = modelRegistry.authStorage.get(providerId);
|
||||
if (credential === undefined) continue;
|
||||
for (const credential of await runtime.listCredentials()) {
|
||||
options.push({
|
||||
id: providerId,
|
||||
name: modelRegistry.getProviderDisplayName(providerId),
|
||||
id: credential.providerId,
|
||||
name: providerNames.get(credential.providerId) ?? credential.providerId,
|
||||
authType: credential.type,
|
||||
status: modelRegistry.getProviderAuthStatus(providerId),
|
||||
status: truthfulProviderStatus(runtime, credential.providerId),
|
||||
});
|
||||
}
|
||||
return filterAndSort(options);
|
||||
}
|
||||
|
||||
export function isApiKeyLoginProvider(providerId: string, oauthProviderIds: ReadonlySet<string>): boolean {
|
||||
if (OAUTH_ONLY_PROVIDERS.has(providerId)) return false;
|
||||
if (providerId === "anthropic") return true;
|
||||
if (oauthProviderIds.has(providerId)) return false;
|
||||
return true;
|
||||
function truthfulProviderStatus(runtime: AuthProviderRuntime, providerId: string): AuthProviderStatus {
|
||||
const reported = runtime.getProviderAuthStatus(providerId);
|
||||
// ModelRuntime reports any stored entry as configured before checking whether
|
||||
// the provider can resolve all required credential and ambient fields.
|
||||
return reported.configured && !runtime.hasConfiguredAuth(providerId) ? { configured: false } : reported;
|
||||
}
|
||||
|
||||
function filterAndSort(options: AuthProviderOption[], authType?: AuthType): AuthProviderOption[] {
|
||||
|
||||
@@ -4,7 +4,7 @@ import type { AuthService } from "./authService.js";
|
||||
export function registerAuthRoutes(app: FastifyInstance, auth: AuthService, prefix = ""): void {
|
||||
app.get<{ Querystring: { mode?: "login" | "logout"; authType?: "oauth" | "api_key" } }>(`${prefix}/auth/providers`, async (request, reply) => {
|
||||
try {
|
||||
return auth.authProviders(request.query.mode ?? "login", request.query.authType);
|
||||
return await auth.authProviders(request.query.mode ?? "login", request.query.authType);
|
||||
} catch (error) {
|
||||
return reply.code(404).send({ error: error instanceof Error ? error.message : String(error) });
|
||||
}
|
||||
@@ -12,7 +12,17 @@ export function registerAuthRoutes(app: FastifyInstance, auth: AuthService, pref
|
||||
|
||||
app.post<{ Body: { providerId: string; key: string } }>(`${prefix}/auth/api-key`, async (request, reply) => {
|
||||
try {
|
||||
return auth.saveApiKey(request.body.providerId, request.body.key);
|
||||
return await auth.saveApiKey(request.body.providerId, request.body.key);
|
||||
} catch (error) {
|
||||
return reply.code(400).send({ error: error instanceof Error ? error.message : String(error) });
|
||||
}
|
||||
});
|
||||
|
||||
// Additive endpoint for newer browsers; the one-secret route remains for
|
||||
// rolling compatibility with older browser bundles.
|
||||
app.post<{ Body: { providerId: string } }>(`${prefix}/auth/api-key/interactive`, async (request, reply) => {
|
||||
try {
|
||||
return await auth.startApiKeyLogin(request.body.providerId);
|
||||
} catch (error) {
|
||||
return reply.code(400).send({ error: error instanceof Error ? error.message : String(error) });
|
||||
}
|
||||
@@ -20,7 +30,7 @@ export function registerAuthRoutes(app: FastifyInstance, auth: AuthService, pref
|
||||
|
||||
app.post<{ Body: { providerId: string } }>(`${prefix}/auth/logout`, async (request, reply) => {
|
||||
try {
|
||||
return auth.logoutProvider(request.body.providerId);
|
||||
return await auth.logoutProvider(request.body.providerId);
|
||||
} catch (error) {
|
||||
return reply.code(400).send({ error: error instanceof Error ? error.message : String(error) });
|
||||
}
|
||||
@@ -28,7 +38,7 @@ export function registerAuthRoutes(app: FastifyInstance, auth: AuthService, pref
|
||||
|
||||
app.post<{ Body: { providerId: string } }>(`${prefix}/auth/oauth`, async (request, reply) => {
|
||||
try {
|
||||
return auth.startOAuthLogin(request.body.providerId);
|
||||
return await auth.startOAuthLogin(request.body.providerId);
|
||||
} catch (error) {
|
||||
return reply.code(400).send({ error: error instanceof Error ? error.message : String(error) });
|
||||
}
|
||||
|
||||
@@ -1,98 +1,505 @@
|
||||
import { mkdtemp, readFile, rm } from "node:fs/promises";
|
||||
import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { AuthStorage, ModelRegistry } from "@earendil-works/pi-coding-agent";
|
||||
import { ModelRuntime } from "@earendil-works/pi-coding-agent";
|
||||
import { InMemoryCredentialStore, type AuthPrompt, type Credential } from "@earendil-works/pi-ai";
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import type { OAuthFlowState } from "../../shared/apiTypes.js";
|
||||
import { AuthService, type AuthChange } from "./authService.js";
|
||||
import { AuthService, createModelRuntimeForAgentDir, type AuthChange, type AuthServiceLogger } from "./authService.js";
|
||||
import { OAuthLoginFlowService } from "./oauthLoginFlowService.js";
|
||||
|
||||
const tempDirs: string[] = [];
|
||||
|
||||
afterEach(async () => {
|
||||
vi.unstubAllEnvs();
|
||||
await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
|
||||
});
|
||||
|
||||
describe("AuthService", () => {
|
||||
it("saves API keys and emits a global auth change", () => {
|
||||
const { auth, authStorage, changes } = createAuthService();
|
||||
it("saves API keys and emits a global auth change after the runtime refreshes", async () => {
|
||||
const { auth, runtime, credentials, changes } = await createAuthService();
|
||||
const reloadConfig = vi.spyOn(runtime, "reloadConfig").mockResolvedValue(undefined);
|
||||
const refresh = vi.spyOn(runtime, "refresh");
|
||||
|
||||
expect(auth.saveApiKey("anthropic", "sk-test")).toEqual({ accepted: true });
|
||||
await expect(auth.saveApiKey("anthropic", "sk-test")).resolves.toEqual({ accepted: true });
|
||||
|
||||
expect(authStorage.get("anthropic")).toEqual({ type: "api_key", key: "sk-test" });
|
||||
await expect(credentials.read("anthropic")).resolves.toEqual({ type: "api_key", key: "sk-test" });
|
||||
expect(reloadConfig).toHaveBeenCalledOnce();
|
||||
expect(refresh).toHaveBeenCalledOnce();
|
||||
expect(changes).toEqual([{}]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("logs out providers and emits the removed provider id", () => {
|
||||
const { auth, authStorage, changes } = createAuthService({ anthropic: { type: "api_key", key: "sk-test" } });
|
||||
it("logs out providers and emits the removed provider id after the runtime refreshes", async () => {
|
||||
const { auth, runtime, credentials, changes } = await createAuthService({ anthropic: { type: "api_key", key: "sk-test" } });
|
||||
const refresh = vi.spyOn(runtime, "refresh");
|
||||
|
||||
expect(auth.logoutProvider("anthropic")).toEqual({ accepted: true });
|
||||
await expect(auth.logoutProvider("anthropic")).resolves.toEqual({ accepted: true });
|
||||
|
||||
expect(authStorage.get("anthropic")).toBeUndefined();
|
||||
await expect(credentials.read("anthropic")).resolves.toBeUndefined();
|
||||
expect(refresh).toHaveBeenCalledOnce();
|
||||
expect(changes).toEqual([{ removedProviderId: "anthropic" }]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("rejects blank API keys", () => {
|
||||
const { auth, changes } = createAuthService();
|
||||
it("persists an API key and attempts every listener when failure logging throws", async () => {
|
||||
const loggingFailure = new Error("auth logger failed");
|
||||
const error = vi.fn(() => { throw loggingFailure; });
|
||||
const logger: AuthServiceLogger = { error };
|
||||
const { auth, credentials, changes } = await createAuthService({}, logger);
|
||||
const failure = new Error("session auth refresh failed");
|
||||
const attempts: string[] = [];
|
||||
auth.subscribe(() => {
|
||||
attempts.push("throwing");
|
||||
throw failure;
|
||||
});
|
||||
auth.subscribe(async () => {
|
||||
await Promise.resolve();
|
||||
attempts.push("healthy");
|
||||
});
|
||||
|
||||
expect(() => { auth.saveApiKey("anthropic", " "); }).toThrow("API key is required");
|
||||
await expect(auth.saveApiKey("anthropic", "sk-test")).resolves.toEqual({ accepted: true });
|
||||
|
||||
await expect(credentials.read("anthropic")).resolves.toEqual({ type: "api_key", key: "sk-test" });
|
||||
expect(changes).toEqual([{}]);
|
||||
expect(attempts).toEqual(["throwing", "healthy"]);
|
||||
expect(error).toHaveBeenCalledWith(
|
||||
{ err: failure, operation: "login", providerId: "anthropic", authType: "api_key" },
|
||||
"auth-change listener failed",
|
||||
);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("removes a credential when auth-change propagation rejects", async () => {
|
||||
const error = vi.fn();
|
||||
const logger: AuthServiceLogger = { error };
|
||||
const { auth, credentials, changes } = await createAuthService(
|
||||
{ anthropic: { type: "api_key", key: "sk-test" } },
|
||||
logger,
|
||||
);
|
||||
const failure = new Error("session logout refresh failed");
|
||||
auth.subscribe(() => Promise.reject(failure));
|
||||
|
||||
await expect(auth.logoutProvider("anthropic")).resolves.toEqual({ accepted: true });
|
||||
|
||||
await expect(credentials.read("anthropic")).resolves.toBeUndefined();
|
||||
expect(changes).toEqual([{ removedProviderId: "anthropic" }]);
|
||||
expect(error).toHaveBeenCalledWith(
|
||||
{ err: failure, operation: "logout", providerId: "anthropic" },
|
||||
"auth-change listener failed",
|
||||
);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("rejects blank API keys", async () => {
|
||||
const { auth, changes } = await createAuthService();
|
||||
|
||||
await expect(auth.saveApiKey("anthropic", " ")).rejects.toThrow("API key is required");
|
||||
expect(changes).toEqual([]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("keeps existing file-backed credentials unchanged when legacy Cloudflare setup cannot finish", async () => {
|
||||
const seed = {
|
||||
"cloudflare-ai-gateway": {
|
||||
type: "api_key" as const,
|
||||
key: "existing-secret",
|
||||
env: { CLOUDFLARE_ACCOUNT_ID: "existing-account", CLOUDFLARE_GATEWAY_ID: "existing-gateway" },
|
||||
},
|
||||
};
|
||||
const { auth, authPath, changes } = await createFileBackedAuthService(seed);
|
||||
const before = await readFile(authPath, "utf8");
|
||||
|
||||
await expect(auth.saveApiKey("cloudflare-ai-gateway", "new-secret")).rejects.toThrow(
|
||||
"Cloudflare AI Gateway requires interactive setup; use Pi's generic /login flow",
|
||||
);
|
||||
|
||||
await expect(readFile(authPath, "utf8")).resolves.toBe(before);
|
||||
expect(changes).toEqual([]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it.each([
|
||||
{ providerId: "amazon-bedrock", providerName: "Amazon Bedrock" },
|
||||
{ providerId: "google-vertex", providerName: "Google Vertex AI" },
|
||||
])("keeps an empty file-backed store unchanged when legacy $providerName setup starts with a selection", async ({ providerId, providerName }) => {
|
||||
const { auth, authPath, changes } = await createFileBackedAuthService({});
|
||||
const before = await readFile(authPath, "utf8");
|
||||
|
||||
await expect(auth.saveApiKey(providerId, "submitted-secret")).rejects.toThrow(
|
||||
`${providerName} requires interactive setup; use Pi's generic /login flow`,
|
||||
);
|
||||
|
||||
await expect(readFile(authPath, "utf8")).resolves.toBe(before);
|
||||
expect(changes).toEqual([]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("executes Cloudflare multi-field API-key setup through the interactive flow", async () => {
|
||||
const { auth, credentials, changes } = await createAuthService();
|
||||
|
||||
const state = await auth.startApiKeyLogin("cloudflare-ai-gateway");
|
||||
expect(state.prompt).toMatchObject({ message: "Enter Cloudflare API key", promptType: "secret" });
|
||||
if (state.prompt === undefined) throw new Error("Expected Cloudflare key prompt");
|
||||
auth.respondToOAuthFlow(state.flowId, state.prompt.requestId, "cf-secret");
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(auth.oauthFlow(state.flowId).prompt).toMatchObject({ message: "Enter Cloudflare account ID", promptType: "text" });
|
||||
});
|
||||
const accountPrompt = auth.oauthFlow(state.flowId).prompt;
|
||||
if (accountPrompt === undefined) throw new Error("Expected Cloudflare account prompt");
|
||||
auth.respondToOAuthFlow(state.flowId, accountPrompt.requestId, "account-1");
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(auth.oauthFlow(state.flowId).prompt).toMatchObject({ message: "Enter Cloudflare AI Gateway ID", promptType: "text" });
|
||||
});
|
||||
const gatewayPrompt = auth.oauthFlow(state.flowId).prompt;
|
||||
if (gatewayPrompt === undefined) throw new Error("Expected Cloudflare gateway prompt");
|
||||
auth.respondToOAuthFlow(state.flowId, gatewayPrompt.requestId, "gateway-1");
|
||||
|
||||
await vi.waitFor(() => { expect(auth.oauthFlow(state.flowId).status).toBe("complete"); });
|
||||
await expect(credentials.read("cloudflare-ai-gateway")).resolves.toEqual({
|
||||
type: "api_key",
|
||||
key: "cf-secret",
|
||||
env: { CLOUDFLARE_ACCOUNT_ID: "account-1", CLOUDFLARE_GATEWAY_ID: "gateway-1" },
|
||||
});
|
||||
expect(changes).toEqual([{}]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it.each([
|
||||
{ providerId: "amazon-bedrock", selection: "bearer-token", secretPrompt: "Enter Amazon Bedrock bearer token" },
|
||||
{ providerId: "google-vertex", selection: "api-key", secretPrompt: "Enter Google Cloud API key" },
|
||||
])("executes $providerId select-first API-key setup through the interactive flow", async ({ providerId, selection, secretPrompt }) => {
|
||||
const { auth, credentials, changes } = await createAuthService();
|
||||
|
||||
const state = await auth.startApiKeyLogin(providerId);
|
||||
expect(state.select).toBeDefined();
|
||||
if (state.select === undefined) throw new Error("Expected auth method selection");
|
||||
auth.respondToOAuthFlow(state.flowId, state.select.requestId, selection);
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(auth.oauthFlow(state.flowId).prompt).toMatchObject({ message: secretPrompt, promptType: "secret" });
|
||||
});
|
||||
const prompt = auth.oauthFlow(state.flowId).prompt;
|
||||
if (prompt === undefined) throw new Error("Expected provider secret prompt");
|
||||
auth.respondToOAuthFlow(state.flowId, prompt.requestId, "provider-secret");
|
||||
|
||||
await vi.waitFor(() => { expect(auth.oauthFlow(state.flowId).status).toBe("complete"); });
|
||||
await expect(credentials.read(providerId)).resolves.toEqual({ type: "api_key", key: "provider-secret" });
|
||||
expect(changes).toEqual([{}]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("reports a key-only legacy Cloudflare credential as unconfigured", async () => {
|
||||
vi.stubEnv("CLOUDFLARE_ACCOUNT_ID", "");
|
||||
vi.stubEnv("CLOUDFLARE_GATEWAY_ID", "");
|
||||
const { auth } = await createFileBackedAuthService({
|
||||
"cloudflare-ai-gateway": { type: "api_key", key: "legacy-secret" },
|
||||
});
|
||||
|
||||
const response = await auth.authProviders("login", "api_key");
|
||||
|
||||
expect(response.providers).toEqual(expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
id: "cloudflare-ai-gateway",
|
||||
loginFlow: "interactive",
|
||||
status: { configured: false },
|
||||
}),
|
||||
]));
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("reports a stored Cloudflare key as configured when ambient fields complete it", async () => {
|
||||
vi.stubEnv("CLOUDFLARE_ACCOUNT_ID", "ambient-account");
|
||||
vi.stubEnv("CLOUDFLARE_GATEWAY_ID", "ambient-gateway");
|
||||
const { auth } = await createFileBackedAuthService({
|
||||
"cloudflare-ai-gateway": { type: "api_key", key: "legacy-secret" },
|
||||
});
|
||||
|
||||
const response = await auth.authProviders("login", "api_key");
|
||||
|
||||
expect(response.providers).toEqual(expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
id: "cloudflare-ai-gateway",
|
||||
loginFlow: "interactive",
|
||||
status: { configured: true, source: "stored" },
|
||||
}),
|
||||
]));
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it.each([
|
||||
{ label: "text", prompt: { type: "text", message: "Account" } satisfies AuthPrompt },
|
||||
{
|
||||
label: "select",
|
||||
prompt: { type: "select", message: "Region", options: [{ id: "us", label: "US" }] } satisfies AuthPrompt,
|
||||
},
|
||||
{ label: "manual-code", prompt: { type: "manual_code", message: "Code" } satisfies AuthPrompt },
|
||||
])("rejects a first $label prompt before credential persistence", async ({ prompt }) => {
|
||||
const { auth, runtime, credentials, changes } = await createAuthService();
|
||||
const login = mockLoginPromptsBeforePersistence(runtime, credentials, [prompt]);
|
||||
|
||||
await expect(auth.saveApiKey("anthropic", "sk-test")).rejects.toThrow(
|
||||
"Anthropic requires interactive setup; use Pi's generic /login flow",
|
||||
);
|
||||
|
||||
expect(login).toHaveBeenCalledOnce();
|
||||
await expect(credentials.read("anthropic")).resolves.toBeUndefined();
|
||||
expect(changes).toEqual([]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("rejects a repeated secret prompt before credential persistence", async () => {
|
||||
const { auth, runtime, credentials, changes } = await createAuthService();
|
||||
const login = mockLoginPromptsBeforePersistence(runtime, credentials, [
|
||||
{ type: "secret", message: "API key" },
|
||||
{ type: "secret", message: "API key again" },
|
||||
]);
|
||||
|
||||
await expect(auth.saveApiKey("anthropic", "sk-test")).rejects.toThrow(
|
||||
"Anthropic requires interactive setup; use Pi's generic /login flow",
|
||||
);
|
||||
|
||||
expect(login).toHaveBeenCalledOnce();
|
||||
await expect(credentials.read("anthropic")).resolves.toBeUndefined();
|
||||
expect(changes).toEqual([]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("rejects an aborted secret prompt before credential persistence", async () => {
|
||||
const { auth, runtime, credentials, changes } = await createAuthService();
|
||||
const abort = new AbortController();
|
||||
abort.abort();
|
||||
const login = mockLoginPromptsBeforePersistence(runtime, credentials, [
|
||||
{ type: "secret", message: "API key", signal: abort.signal },
|
||||
]);
|
||||
|
||||
await expect(auth.saveApiKey("anthropic", "sk-test")).rejects.toThrow("Login cancelled");
|
||||
|
||||
expect(login).toHaveBeenCalledOnce();
|
||||
await expect(credentials.read("anthropic")).resolves.toBeUndefined();
|
||||
expect(changes).toEqual([]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("rejects unknown providers before starting API-key login", async () => {
|
||||
const { auth, runtime, credentials, changes } = await createAuthService();
|
||||
const login = vi.spyOn(runtime, "login");
|
||||
|
||||
await expect(auth.saveApiKey("unknown-provider", "sk-test")).rejects.toThrow(
|
||||
"API key provider not found: unknown-provider",
|
||||
);
|
||||
|
||||
expect(login).not.toHaveBeenCalled();
|
||||
await expect(credentials.read("unknown-provider")).resolves.toBeUndefined();
|
||||
expect(changes).toEqual([]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("rejects ambient-only providers before starting API-key login", async () => {
|
||||
const { auth, runtime, credentials, changes } = await createAuthService();
|
||||
const providers = [...runtime.getProviders()];
|
||||
const interactiveProvider = providers.find((provider) => provider.auth.apiKey?.login !== undefined);
|
||||
if (interactiveProvider?.auth.apiKey === undefined) throw new Error("Expected an interactive API-key provider");
|
||||
const ambientApiKey = { ...interactiveProvider.auth.apiKey };
|
||||
delete ambientApiKey.login;
|
||||
const ambientProvider = {
|
||||
...interactiveProvider,
|
||||
id: "ambient-only",
|
||||
name: "Ambient Only",
|
||||
auth: { apiKey: ambientApiKey },
|
||||
};
|
||||
vi.spyOn(runtime, "getProviders").mockReturnValue([...providers, ambientProvider]);
|
||||
const login = vi.spyOn(runtime, "login");
|
||||
|
||||
await expect(auth.saveApiKey("ambient-only", "sk-test")).rejects.toThrow(
|
||||
"Ambient Only does not support interactive API-key setup",
|
||||
);
|
||||
|
||||
expect(login).not.toHaveBeenCalled();
|
||||
await expect(credentials.read("ambient-only")).resolves.toBeUndefined();
|
||||
expect(changes).toEqual([]);
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("reloads models.json before enumerating and validating OAuth providers", async () => {
|
||||
const agentDir = await tempAgentDir();
|
||||
const modelsPath = join(agentDir, "models.json");
|
||||
const runtime = await ModelRuntime.create({
|
||||
credentials: new InMemoryCredentialStore(),
|
||||
modelsPath,
|
||||
allowModelNetwork: false,
|
||||
});
|
||||
const authFlows = new CapturingOAuthLoginFlowService();
|
||||
const auth = await AuthService.create({ runtime, authFlows });
|
||||
|
||||
await writeFile(modelsPath, radiusModelsConfig("First Radius"));
|
||||
const response = await auth.authProviders("login", "oauth");
|
||||
expect(response.providers).toEqual(expect.arrayContaining([
|
||||
expect.objectContaining({ id: "test-radius", name: "First Radius", authType: "oauth" }),
|
||||
]));
|
||||
|
||||
await writeFile(modelsPath, radiusModelsConfig("Updated Radius"));
|
||||
await expect(auth.startOAuthLogin("test-radius")).resolves.toMatchObject({
|
||||
providerId: "test-radius",
|
||||
providerName: "Updated Radius",
|
||||
status: "running",
|
||||
});
|
||||
expect(authFlows.startCalls.at(0)).toMatchObject({
|
||||
providerId: "test-radius",
|
||||
providerName: "Updated Radius",
|
||||
runtime,
|
||||
});
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("stores credentials in the configured agent directory", async () => {
|
||||
const agentDir = await tempAgentDir();
|
||||
const auth = new AuthService({ agentDir });
|
||||
const runtime = await createModelRuntimeForAgentDir(agentDir, false);
|
||||
const auth = await AuthService.create({ runtime });
|
||||
|
||||
auth.saveApiKey("anthropic", "sk-test");
|
||||
await auth.saveApiKey("anthropic", "sk-test");
|
||||
|
||||
await expect(readFile(join(agentDir, "auth.json"), "utf8")).resolves.toContain("sk-test");
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("refreshes auth state after OAuth login completes", () => {
|
||||
const authStorage = AuthStorage.inMemory();
|
||||
const modelRegistry = ModelRegistry.create(authStorage);
|
||||
it("reconciles cancellation after ModelRuntime persists OAuth but before its refresh completes", async () => {
|
||||
const { auth, runtime, credentials, changes } = await createAuthService();
|
||||
const provider = runtime.getProviders().find((option) => option.id === "anthropic" && option.auth.oauth !== undefined);
|
||||
if (provider?.auth.oauth === undefined) throw new Error("Expected built-in OAuth provider");
|
||||
const credential: Credential = {
|
||||
type: "oauth",
|
||||
refresh: "refresh-token",
|
||||
access: "access-token",
|
||||
expires: Date.now() + 60_000,
|
||||
};
|
||||
vi.spyOn(provider.auth.oauth, "login").mockResolvedValue(credential);
|
||||
vi.spyOn(runtime, "reloadConfig").mockResolvedValue(undefined);
|
||||
const refreshStarted = deferred<undefined>();
|
||||
const finishRefresh = deferred<undefined>();
|
||||
const refresh = vi.spyOn(runtime, "refresh").mockImplementation(async () => {
|
||||
refreshStarted.resolve(undefined);
|
||||
await finishRefresh.promise;
|
||||
return { aborted: false, errors: new Map() };
|
||||
});
|
||||
|
||||
const state = await auth.startOAuthLogin(provider.id);
|
||||
await refreshStarted.promise;
|
||||
|
||||
await expect(credentials.read(provider.id)).resolves.toEqual(credential);
|
||||
expect(auth.cancelOAuthFlow(state.flowId)).toMatchObject({ status: "cancelled", error: "Login cancelled" });
|
||||
expect(changes).toEqual([]);
|
||||
|
||||
finishRefresh.resolve(undefined);
|
||||
await vi.waitFor(() => { expect(auth.oauthFlow(state.flowId).status).toBe("complete"); });
|
||||
|
||||
expect(auth.oauthFlow(state.flowId)).toMatchObject({ status: "complete", progress: ["Login complete"] });
|
||||
expect(auth.oauthFlow(state.flowId)).not.toHaveProperty("error");
|
||||
await expect(credentials.read(provider.id)).resolves.toEqual(credential);
|
||||
expect(changes).toEqual([{}]);
|
||||
expect(refresh).toHaveBeenCalledOnce();
|
||||
auth.dispose();
|
||||
});
|
||||
|
||||
it("emits an auth change after OAuth login completes without refreshing twice", async () => {
|
||||
const runtime = await ModelRuntime.create({
|
||||
credentials: new InMemoryCredentialStore(),
|
||||
modelsPath: null,
|
||||
allowModelNetwork: false,
|
||||
});
|
||||
const authFlows = new CapturingOAuthLoginFlowService();
|
||||
const auth = new AuthService({ modelRegistry, authFlows });
|
||||
const auth = await AuthService.create({ runtime, authFlows });
|
||||
const changes: AuthChange[] = [];
|
||||
auth.subscribe((change) => { changes.push(change); });
|
||||
const reload = vi.spyOn(authStorage, "reload");
|
||||
const refresh = vi.spyOn(modelRegistry, "refresh");
|
||||
const provider = authStorage.getOAuthProviders().find((option) => option.id === "anthropic");
|
||||
const refresh = vi.spyOn(runtime, "refresh");
|
||||
const provider = runtime.getProviders().find((option) => option.id === "anthropic" && option.auth.oauth !== undefined);
|
||||
if (provider === undefined) throw new Error("Expected built-in OAuth provider");
|
||||
|
||||
expect(auth.startOAuthLogin(provider.id)).toMatchObject({ providerId: provider.id, providerName: provider.name, status: "running" });
|
||||
await expect(auth.startOAuthLogin(provider.id)).resolves.toMatchObject({ providerId: provider.id, providerName: provider.name, status: "running" });
|
||||
|
||||
const startOptions = authFlows.startCalls.at(0);
|
||||
if (startOptions === undefined) throw new Error("Expected OAuth flow to start");
|
||||
expect(startOptions.providerId).toBe(provider.id);
|
||||
expect(startOptions.providerName).toBe(provider.name);
|
||||
expect(startOptions.authStorage).toBe(authStorage);
|
||||
expect(startOptions.runtime).toBe(runtime);
|
||||
expect(changes).toEqual([]);
|
||||
|
||||
reload.mockClear();
|
||||
refresh.mockClear();
|
||||
if (startOptions.onComplete === undefined) throw new Error("Expected OAuth completion callback");
|
||||
startOptions.onComplete();
|
||||
|
||||
expect(reload).toHaveBeenCalledOnce();
|
||||
expect(refresh).toHaveBeenCalledOnce();
|
||||
await startOptions.onComplete();
|
||||
expect(changes).toEqual([{}]);
|
||||
|
||||
expect(refresh).not.toHaveBeenCalled();
|
||||
auth.dispose();
|
||||
expect(authFlows.disposed).toBe(true);
|
||||
});
|
||||
|
||||
it("completes OAuth when an auth-change listener and failure logging throw", async () => {
|
||||
const loggingFailure = new Error("auth logger failed");
|
||||
const error = vi.fn(() => { throw loggingFailure; });
|
||||
const logger: AuthServiceLogger = { error };
|
||||
const { auth, runtime, changes } = await createAuthService({}, logger);
|
||||
const provider = runtime.getProviders().find((option) => option.id === "anthropic" && option.auth.oauth !== undefined);
|
||||
if (provider === undefined) throw new Error("Expected built-in OAuth provider");
|
||||
vi.spyOn(runtime, "login").mockResolvedValue({
|
||||
type: "oauth",
|
||||
refresh: "refresh-token",
|
||||
access: "access-token",
|
||||
expires: Date.now() + 60_000,
|
||||
});
|
||||
const failure = new Error("session OAuth refresh failed");
|
||||
auth.subscribe(() => Promise.reject(failure));
|
||||
|
||||
const state = await auth.startOAuthLogin(provider.id);
|
||||
await vi.waitFor(() => { expect(auth.oauthFlow(state.flowId).status).toBe("complete"); });
|
||||
|
||||
expect(changes).toEqual([{}]);
|
||||
expect(error).toHaveBeenCalledWith(
|
||||
{ err: failure, operation: "login", providerId: provider.id, authType: "oauth" },
|
||||
"auth-change listener failed",
|
||||
);
|
||||
auth.dispose();
|
||||
});
|
||||
});
|
||||
|
||||
function createAuthService(data: Parameters<typeof AuthStorage.inMemory>[0] = {}) {
|
||||
const authStorage = AuthStorage.inMemory(data);
|
||||
const modelRegistry = ModelRegistry.create(authStorage);
|
||||
const auth = new AuthService({ modelRegistry });
|
||||
async function createAuthService(seed: Record<string, Credential> = {}, logger?: AuthServiceLogger) {
|
||||
const credentials = new InMemoryCredentialStore();
|
||||
for (const [providerId, credential] of Object.entries(seed)) {
|
||||
await credentials.modify(providerId, () => Promise.resolve(credential));
|
||||
}
|
||||
const runtime = await ModelRuntime.create({ credentials, modelsPath: null, allowModelNetwork: false });
|
||||
const auth = await AuthService.create({ runtime, ...(logger === undefined ? {} : { logger }) });
|
||||
const changes: AuthChange[] = [];
|
||||
auth.subscribe((change) => { changes.push(change); });
|
||||
return { auth, authStorage, changes };
|
||||
return { auth, runtime, credentials, changes };
|
||||
}
|
||||
|
||||
async function createFileBackedAuthService(seed: Record<string, Credential>) {
|
||||
const agentDir = await tempAgentDir();
|
||||
const authPath = join(agentDir, "auth.json");
|
||||
await writeFile(authPath, JSON.stringify(seed, null, 2));
|
||||
const runtime = await createModelRuntimeForAgentDir(agentDir, false);
|
||||
const auth = await AuthService.create({ runtime });
|
||||
const changes: AuthChange[] = [];
|
||||
auth.subscribe((change) => { changes.push(change); });
|
||||
return { auth, runtime, authPath, changes };
|
||||
}
|
||||
|
||||
function mockLoginPromptsBeforePersistence(
|
||||
runtime: ModelRuntime,
|
||||
credentials: InMemoryCredentialStore,
|
||||
prompts: readonly AuthPrompt[],
|
||||
) {
|
||||
return vi.spyOn(runtime, "login").mockImplementation(async (providerId, _authType, interaction) => {
|
||||
let key: string | undefined;
|
||||
for (const prompt of prompts) key = await interaction.prompt(prompt);
|
||||
if (key === undefined) throw new Error("Expected at least one login prompt");
|
||||
const credential: Credential = { type: "api_key", key };
|
||||
await credentials.modify(providerId, () => Promise.resolve(credential));
|
||||
return credential;
|
||||
});
|
||||
}
|
||||
|
||||
async function tempAgentDir(): Promise<string> {
|
||||
@@ -101,6 +508,28 @@ async function tempAgentDir(): Promise<string> {
|
||||
return dir;
|
||||
}
|
||||
|
||||
function deferred<T>() {
|
||||
let resolveValue: (value: T) => void = () => undefined;
|
||||
let rejectValue: (reason?: unknown) => void = () => undefined;
|
||||
const promise = new Promise<T>((resolve, reject) => {
|
||||
resolveValue = resolve;
|
||||
rejectValue = reject;
|
||||
});
|
||||
return { promise, resolve: resolveValue, reject: rejectValue };
|
||||
}
|
||||
|
||||
function radiusModelsConfig(name: string): string {
|
||||
return JSON.stringify({
|
||||
providers: {
|
||||
"test-radius": {
|
||||
name,
|
||||
baseUrl: "https://radius.example.test/v1",
|
||||
oauth: "radius",
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
class CapturingOAuthLoginFlowService extends OAuthLoginFlowService {
|
||||
readonly startCalls: Parameters<OAuthLoginFlowService["start"]>[0][] = [];
|
||||
disposed = false;
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import { join } from "node:path";
|
||||
import { AuthStorage, ModelRegistry } from "@earendil-works/pi-coding-agent";
|
||||
import { ModelRuntime } from "@earendil-works/pi-coding-agent";
|
||||
import type { AuthInteraction } from "@earendil-works/pi-ai";
|
||||
import type { AuthProvidersResponse, AuthType, OAuthFlowState } from "../../shared/apiTypes.js";
|
||||
import { getLoginProviderOptions, getLogoutProviderOptions } from "./authProviderOptions.js";
|
||||
import { OAuthLoginFlowService } from "./oauthLoginFlowService.js";
|
||||
@@ -8,28 +9,53 @@ export interface AuthChange {
|
||||
removedProviderId?: string;
|
||||
}
|
||||
|
||||
type AuthChangeListener = (change: AuthChange) => void;
|
||||
type ModelRegistryInstance = ReturnType<typeof ModelRegistry.create>;
|
||||
type AuthChangeListener = (change: AuthChange) => void | Promise<void>;
|
||||
|
||||
export interface AuthServiceDependencies {
|
||||
agentDir?: string;
|
||||
modelRegistry?: ModelRegistryInstance;
|
||||
runtime?: ModelRuntime;
|
||||
authFlows?: OAuthLoginFlowService;
|
||||
logger?: AuthServiceLogger;
|
||||
}
|
||||
|
||||
export function createModelRegistryForAgentDir(agentDir: string): ModelRegistryInstance {
|
||||
const authStorage = AuthStorage.create(join(agentDir, "auth.json"));
|
||||
return ModelRegistry.create(authStorage, join(agentDir, "models.json"));
|
||||
/** Minimal structured-logging seam for non-fatal auth propagation failures. */
|
||||
export interface AuthServiceLogger {
|
||||
error(details: Record<string, unknown>, message: string): void;
|
||||
}
|
||||
|
||||
interface AuthChangeContext {
|
||||
operation: "login" | "logout";
|
||||
providerId: string;
|
||||
authType?: AuthType;
|
||||
}
|
||||
|
||||
const noopLogger: AuthServiceLogger = { error() { /* no-op */ } };
|
||||
|
||||
export function createModelRuntimeForAgentDir(agentDir: string, allowModelNetwork?: boolean): Promise<ModelRuntime> {
|
||||
return ModelRuntime.create({
|
||||
authPath: join(agentDir, "auth.json"),
|
||||
modelsPath: join(agentDir, "models.json"),
|
||||
...(allowModelNetwork === undefined ? {} : { allowModelNetwork }),
|
||||
});
|
||||
}
|
||||
|
||||
export class AuthService {
|
||||
readonly modelRegistry: ModelRegistryInstance;
|
||||
readonly runtime: ModelRuntime;
|
||||
private readonly authFlows: OAuthLoginFlowService;
|
||||
private readonly logger: AuthServiceLogger;
|
||||
private readonly listeners = new Set<AuthChangeListener>();
|
||||
|
||||
constructor(deps: AuthServiceDependencies = {}) {
|
||||
this.modelRegistry = deps.modelRegistry ?? (deps.agentDir === undefined ? ModelRegistry.create(AuthStorage.create()) : createModelRegistryForAgentDir(deps.agentDir));
|
||||
this.authFlows = deps.authFlows ?? new OAuthLoginFlowService();
|
||||
private constructor(runtime: ModelRuntime, authFlows: OAuthLoginFlowService, logger: AuthServiceLogger) {
|
||||
this.runtime = runtime;
|
||||
this.authFlows = authFlows;
|
||||
this.logger = logger;
|
||||
}
|
||||
|
||||
static async create(deps: AuthServiceDependencies = {}): Promise<AuthService> {
|
||||
const runtime = deps.runtime ?? (deps.agentDir === undefined ? await ModelRuntime.create({}) : await createModelRuntimeForAgentDir(deps.agentDir));
|
||||
const logger = deps.logger ?? noopLogger;
|
||||
const authFlows = deps.authFlows ?? new OAuthLoginFlowService({ logger });
|
||||
return new AuthService(runtime, authFlows, logger);
|
||||
}
|
||||
|
||||
subscribe(listener: AuthChangeListener): () => void {
|
||||
@@ -44,34 +70,60 @@ export class AuthService {
|
||||
this.listeners.clear();
|
||||
}
|
||||
|
||||
authProviders(mode: "login" | "logout", authType?: AuthType): AuthProvidersResponse {
|
||||
this.modelRegistry.refresh();
|
||||
const providers = mode === "logout" ? getLogoutProviderOptions(this.modelRegistry) : getLoginProviderOptions(this.modelRegistry, authType);
|
||||
async authProviders(mode: "login" | "logout", authType?: AuthType): Promise<AuthProvidersResponse> {
|
||||
await this.runtime.reloadConfig();
|
||||
const providers = mode === "logout" ? await getLogoutProviderOptions(this.runtime) : getLoginProviderOptions(this.runtime, authType);
|
||||
return { providers };
|
||||
}
|
||||
|
||||
saveApiKey(providerId: string, key: string): { accepted: true } {
|
||||
async saveApiKey(providerId: string, key: string): Promise<{ accepted: true }> {
|
||||
if (key.trim() === "") throw new Error("API key is required");
|
||||
this.modelRegistry.authStorage.set(providerId, { type: "api_key", key });
|
||||
this.refreshAuthState();
|
||||
const provider = await this.requireApiKeyLoginProvider(providerId);
|
||||
let promptAttempted = false;
|
||||
const interaction: AuthInteraction = {
|
||||
prompt: (prompt) => {
|
||||
if (promptAttempted) {
|
||||
throw new Error(`${provider.name} requires interactive setup; use Pi's generic /login flow`);
|
||||
}
|
||||
promptAttempted = true;
|
||||
if (prompt.signal?.aborted === true) throw new Error("Login cancelled");
|
||||
if (prompt.type !== "secret") {
|
||||
throw new Error(`${provider.name} requires interactive setup; use Pi's generic /login flow`);
|
||||
}
|
||||
return Promise.resolve(key);
|
||||
},
|
||||
notify: () => undefined,
|
||||
};
|
||||
await this.runtime.login(providerId, "api_key", interaction);
|
||||
await this.emit({}, { operation: "login", providerId, authType: "api_key" });
|
||||
return { accepted: true };
|
||||
}
|
||||
|
||||
logoutProvider(providerId: string): { accepted: true } {
|
||||
this.modelRegistry.authStorage.logout(providerId);
|
||||
this.refreshAuthState({ removedProviderId: providerId });
|
||||
async logoutProvider(providerId: string): Promise<{ accepted: true }> {
|
||||
await this.runtime.logout(providerId);
|
||||
await this.emit({ removedProviderId: providerId }, { operation: "logout", providerId });
|
||||
return { accepted: true };
|
||||
}
|
||||
|
||||
startOAuthLogin(providerId: string): OAuthFlowState {
|
||||
const provider = this.requireOAuthLoginProvider(providerId);
|
||||
async startApiKeyLogin(providerId: string): Promise<OAuthFlowState> {
|
||||
const provider = await this.requireApiKeyLoginProvider(providerId);
|
||||
return this.authFlows.start({
|
||||
providerId,
|
||||
providerName: provider.name,
|
||||
authStorage: this.modelRegistry.authStorage,
|
||||
onComplete: () => {
|
||||
this.refreshAuthState();
|
||||
},
|
||||
runtime: this.runtime,
|
||||
authType: "api_key",
|
||||
onComplete: () => this.emit({}, { operation: "login", providerId, authType: "api_key" }),
|
||||
});
|
||||
}
|
||||
|
||||
async startOAuthLogin(providerId: string): Promise<OAuthFlowState> {
|
||||
const provider = await this.requireOAuthLoginProvider(providerId);
|
||||
return this.authFlows.start({
|
||||
providerId,
|
||||
providerName: provider.name,
|
||||
runtime: this.runtime,
|
||||
authType: "oauth",
|
||||
onComplete: () => this.emit({}, { operation: "login", providerId, authType: "oauth" }),
|
||||
});
|
||||
}
|
||||
|
||||
@@ -87,19 +139,38 @@ export class AuthService {
|
||||
return this.authFlows.cancel(flowId);
|
||||
}
|
||||
|
||||
private refreshAuthState(change: AuthChange = {}): void {
|
||||
this.modelRegistry.authStorage.reload();
|
||||
this.modelRegistry.refresh();
|
||||
this.emit(change);
|
||||
private async emit(change: AuthChange, context: AuthChangeContext): Promise<void> {
|
||||
const results = await Promise.allSettled([...this.listeners].map(async (listener) => listener(change)));
|
||||
for (const result of results) {
|
||||
if (result.status === "rejected") {
|
||||
this.logErrorNoThrow({ err: result.reason, ...context }, "auth-change listener failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private emit(change: AuthChange): void {
|
||||
for (const listener of this.listeners) listener(change);
|
||||
private logErrorNoThrow(details: Record<string, unknown>, message: string): void {
|
||||
try {
|
||||
this.logger.error(details, message);
|
||||
} catch {
|
||||
// A diagnostic failure cannot turn an already-committed auth mutation into an API failure.
|
||||
}
|
||||
}
|
||||
|
||||
private requireOAuthLoginProvider(providerId: string) {
|
||||
this.modelRegistry.refresh();
|
||||
const provider = getLoginProviderOptions(this.modelRegistry, "oauth").find((option) => option.id === providerId);
|
||||
private async requireApiKeyLoginProvider(providerId: string) {
|
||||
await this.runtime.reloadConfig();
|
||||
const provider = getLoginProviderOptions(this.runtime, "api_key").find((option) => option.id === providerId);
|
||||
if (provider !== undefined) return provider;
|
||||
|
||||
const knownProvider = this.runtime.getProviders().find((option) => option.id === providerId);
|
||||
if (knownProvider !== undefined) {
|
||||
throw new Error(`${knownProvider.name} does not support interactive API-key setup`);
|
||||
}
|
||||
throw new Error(`API key provider not found: ${providerId}`);
|
||||
}
|
||||
|
||||
private async requireOAuthLoginProvider(providerId: string) {
|
||||
await this.runtime.reloadConfig();
|
||||
const provider = getLoginProviderOptions(this.runtime, "oauth").find((option) => option.id === providerId);
|
||||
if (provider === undefined) throw new Error(`OAuth provider not found: ${providerId}`);
|
||||
return provider;
|
||||
}
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
import type { OAuthLoginCallbacks } from "@earendil-works/pi-ai";
|
||||
import type { AuthStorage } from "@earendil-works/pi-coding-agent";
|
||||
import type { AuthInteraction, AuthType } from "@earendil-works/pi-ai";
|
||||
import type { ModelRuntime } from "@earendil-works/pi-coding-agent";
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { OAuthLoginFlowService } from "./oauthLoginFlowService.js";
|
||||
|
||||
type LoginHandler = (providerId: string, callbacks: OAuthLoginCallbacks) => Promise<void>;
|
||||
type LoginHandler = (providerId: string, interaction: AuthInteraction) => Promise<void>;
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
@@ -17,19 +17,19 @@ describe("OAuthLoginFlowService", () => {
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
authStorage: fakeAuthStorage(async (_providerId, callbacks) => {
|
||||
callbacks.onAuth({ url: "https://example.test/auth", instructions: "Open it" });
|
||||
callbacks.onProgress?.("Waiting for code");
|
||||
promptValue = await callbacks.onPrompt({ message: "Paste code", placeholder: "code" });
|
||||
callbacks.onProgress?.(`Got ${promptValue}`);
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
interaction.notify({ type: "auth_url", url: "https://example.test/auth", instructions: "Open it" });
|
||||
interaction.notify({ type: "progress", message: "Waiting for code" });
|
||||
promptValue = await interaction.prompt({ type: "text", message: "Paste code", placeholder: "code" });
|
||||
interaction.notify({ type: "progress", message: `Got ${promptValue}` });
|
||||
}),
|
||||
onComplete,
|
||||
});
|
||||
|
||||
const prompt = state.prompt;
|
||||
if (prompt === undefined) throw new Error("Expected prompt");
|
||||
expect(state).toMatchObject({ auth: { url: "https://example.test/auth" }, progress: ["Waiting for code"] });
|
||||
expect(prompt).toMatchObject({ message: "Paste code", placeholder: "code", kind: "prompt" });
|
||||
expect(state).toMatchObject({ auth: { url: "https://example.test/auth", instructions: "Open it" }, progress: ["Waiting for code"] });
|
||||
expect(prompt).toMatchObject({ message: "Paste code", placeholder: "code", kind: "prompt", promptType: "text", allowEmpty: true });
|
||||
|
||||
const afterRespond = service.respond(state.flowId, prompt.requestId, "abc123");
|
||||
expect(afterRespond.prompt).toBeUndefined();
|
||||
@@ -41,23 +41,191 @@ describe("OAuthLoginFlowService", () => {
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("runs API-key login through the same AuthInteraction transport", async () => {
|
||||
const authTypes: AuthType[] = [];
|
||||
let key: string | undefined;
|
||||
const service = new OAuthLoginFlowService();
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
key = await interaction.prompt({ type: "secret", message: "Enter API key" });
|
||||
}, authTypes),
|
||||
authType: "api_key",
|
||||
});
|
||||
|
||||
const prompt = state.prompt;
|
||||
if (prompt === undefined) throw new Error("Expected API-key prompt");
|
||||
service.respond(state.flowId, prompt.requestId, "sk-test");
|
||||
await flushAsyncLogin();
|
||||
|
||||
expect(authTypes).toEqual(["api_key"]);
|
||||
expect(key).toBe("sk-test");
|
||||
expect(service.get(state.flowId).status).toBe("complete");
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("awaits async completion propagation before marking the flow complete", async () => {
|
||||
const completion = deferred<undefined>();
|
||||
const service = new OAuthLoginFlowService();
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
runtime: fakeRuntime(() => Promise.resolve()),
|
||||
onComplete: () => completion.promise,
|
||||
});
|
||||
|
||||
await flushAsyncLogin();
|
||||
expect(service.get(state.flowId).status).toBe("running");
|
||||
|
||||
completion.resolve(undefined);
|
||||
await flushAsyncLogin();
|
||||
expect(service.get(state.flowId).status).toBe("complete");
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("keeps a committed login complete when its completion callback and logger throw", async () => {
|
||||
const completionFailure = new Error("completion propagation failed");
|
||||
const loggingFailure = new Error("OAuth logger failed");
|
||||
const error = vi.fn(() => { throw loggingFailure; });
|
||||
const onComplete = vi.fn(() => { throw completionFailure; });
|
||||
const service = new OAuthLoginFlowService({ logger: { error } });
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
runtime: fakeRuntime(() => Promise.resolve()),
|
||||
onComplete,
|
||||
});
|
||||
|
||||
await vi.waitFor(() => { expect(service.get(state.flowId).status).toBe("complete"); });
|
||||
|
||||
expect(service.get(state.flowId)).toMatchObject({ status: "complete", progress: ["Login complete"] });
|
||||
expect(onComplete).toHaveBeenCalledOnce();
|
||||
expect(error).toHaveBeenCalledWith(
|
||||
{ err: completionFailure, flowId: state.flowId, providerId: "test-provider" },
|
||||
"login completion callback failed",
|
||||
);
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("allows blank text responses for providers that use blank as a default", async () => {
|
||||
let domain: string | undefined;
|
||||
const service = new OAuthLoginFlowService();
|
||||
const state = service.start({
|
||||
providerId: "github-copilot",
|
||||
providerName: "GitHub Copilot",
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
domain = await interaction.prompt({
|
||||
type: "text",
|
||||
message: "GitHub Enterprise URL/domain (blank for github.com)",
|
||||
});
|
||||
}),
|
||||
});
|
||||
|
||||
const prompt = state.prompt;
|
||||
if (prompt === undefined) throw new Error("Expected text prompt");
|
||||
expect(prompt).toMatchObject({ kind: "prompt", promptType: "text", allowEmpty: true });
|
||||
|
||||
service.respond(state.flowId, prompt.requestId, "");
|
||||
await flushAsyncLogin();
|
||||
|
||||
expect(domain).toBe("");
|
||||
expect(service.get(state.flowId).status).toBe("complete");
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("preserves secret prompt semantics behind the legacy prompt kind", () => {
|
||||
const service = new OAuthLoginFlowService();
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
await interaction.prompt({ type: "secret", message: "Enter secret", placeholder: "token" });
|
||||
}),
|
||||
});
|
||||
|
||||
const prompt = state.prompt;
|
||||
if (prompt === undefined) throw new Error("Expected secret prompt");
|
||||
expect(prompt).toMatchObject({
|
||||
kind: "prompt",
|
||||
promptType: "secret",
|
||||
message: "Enter secret",
|
||||
placeholder: "token",
|
||||
});
|
||||
expect(prompt).not.toHaveProperty("allowEmpty");
|
||||
expect(() => { service.respond(state.flowId, prompt.requestId, ""); }).toThrow("A value is required");
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("preserves info-event links without replacing the authorization URL", () => {
|
||||
const service = new OAuthLoginFlowService();
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
interaction.notify({ type: "auth_url", url: "https://example.test/login" });
|
||||
interaction.notify({
|
||||
type: "info",
|
||||
message: "Review the provider setup guide",
|
||||
links: [{ url: "https://example.test/docs", label: "Setup guide" }],
|
||||
});
|
||||
await interaction.prompt({ type: "text", message: "Continue" });
|
||||
}),
|
||||
});
|
||||
|
||||
expect(state).toMatchObject({
|
||||
auth: { url: "https://example.test/login" },
|
||||
progress: ["Review the provider setup guide"],
|
||||
info: [{ message: "Review the provider setup guide", links: [{ url: "https://example.test/docs", label: "Setup guide" }] }],
|
||||
});
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("surfaces device-code events through the auth field", () => {
|
||||
const service = new OAuthLoginFlowService();
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
interaction.notify({
|
||||
type: "device_code",
|
||||
userCode: "WXYZ-1234",
|
||||
verificationUri: "https://example.test/device",
|
||||
intervalSeconds: 5,
|
||||
expiresInSeconds: 900,
|
||||
});
|
||||
await interaction.prompt({ type: "text", message: "Waiting" });
|
||||
}),
|
||||
});
|
||||
|
||||
expect(service.get(state.flowId)).toMatchObject({
|
||||
auth: {
|
||||
url: "https://example.test/device",
|
||||
instructions: "Enter code: WXYZ-1234",
|
||||
deviceCode: { userCode: "WXYZ-1234", intervalSeconds: 5, expiresInSeconds: 900 },
|
||||
},
|
||||
});
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("round-trips select responses", async () => {
|
||||
let selectedValue: string | undefined;
|
||||
const service = new OAuthLoginFlowService();
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
authStorage: fakeAuthStorage(async (_providerId, callbacks) => {
|
||||
selectedValue = await callbacks.onSelect({
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
selectedValue = await interaction.prompt({
|
||||
type: "select",
|
||||
message: "Choose account",
|
||||
options: [{ id: "work", label: "Work" }, { id: "personal", label: "Personal" }],
|
||||
options: [{ id: "work", label: "Work", description: "Company account" }, { id: "personal", label: "Personal" }],
|
||||
});
|
||||
}),
|
||||
});
|
||||
|
||||
const select = state.select;
|
||||
if (select === undefined) throw new Error("Expected select prompt");
|
||||
expect(select).toMatchObject({ message: "Choose account", options: [{ value: "work", label: "Work" }, { value: "personal", label: "Personal" }] });
|
||||
expect(select).toMatchObject({ message: "Choose account", options: [{ value: "work", label: "Work", description: "Company account" }, { value: "personal", label: "Personal" }] });
|
||||
|
||||
service.respond(state.flowId, select.requestId, "personal");
|
||||
await flushAsyncLogin();
|
||||
@@ -67,40 +235,97 @@ describe("OAuthLoginFlowService", () => {
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("uses a manual-code prompt for callback-server flows", async () => {
|
||||
let manualValue: string | undefined;
|
||||
it("rejects responses outside the pending select options", () => {
|
||||
const service = new OAuthLoginFlowService();
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
authStorage: fakeAuthStorage(async (_providerId, callbacks) => {
|
||||
const manualCodeInput = callbacks.onManualCodeInput;
|
||||
if (manualCodeInput === undefined) throw new Error("Expected manual-code callback");
|
||||
manualValue = await manualCodeInput();
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
await interaction.prompt({
|
||||
type: "select",
|
||||
message: "Choose account",
|
||||
options: [{ id: "work", label: "Work" }],
|
||||
});
|
||||
}),
|
||||
});
|
||||
|
||||
const select = state.select;
|
||||
if (select === undefined) throw new Error("Expected select prompt");
|
||||
expect(() => { service.respond(state.flowId, select.requestId, "personal"); }).toThrow("Invalid login selection");
|
||||
expect(service.get(state.flowId).select).toEqual(select);
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("uses a manual-code prompt for callback-server flows and cleans up its abort listener", async () => {
|
||||
let manualValue: string | undefined;
|
||||
const service = new OAuthLoginFlowService();
|
||||
const controller = new AbortController();
|
||||
const removeAbortListener = vi.spyOn(controller.signal, "removeEventListener");
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
manualValue = await interaction.prompt({
|
||||
type: "manual_code",
|
||||
message: "Paste the callback URL or authorization code",
|
||||
signal: controller.signal,
|
||||
});
|
||||
}),
|
||||
});
|
||||
|
||||
const prompt = state.prompt;
|
||||
if (prompt === undefined) throw new Error("Expected manual prompt");
|
||||
expect(prompt).toMatchObject({ kind: "manual", message: "Paste the callback URL or authorization code" });
|
||||
expect(prompt).toMatchObject({ kind: "manual", promptType: "manual_code", message: "Paste the callback URL or authorization code" });
|
||||
|
||||
service.respond(state.flowId, prompt.requestId, "https://localhost/callback?code=abc");
|
||||
await flushAsyncLogin();
|
||||
|
||||
expect(manualValue).toBe("https://localhost/callback?code=abc");
|
||||
expect(removeAbortListener).toHaveBeenCalledWith("abort", expect.any(Function));
|
||||
expect(service.get(state.flowId).status).toBe("complete");
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("rejects a pending prompt when its own signal aborts without ending the flow", async () => {
|
||||
const promptRejected = deferred<Error>();
|
||||
const service = new OAuthLoginFlowService();
|
||||
const controller = new AbortController();
|
||||
const removeAbortListener = vi.spyOn(controller.signal, "removeEventListener");
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
try {
|
||||
await interaction.prompt({ type: "manual_code", message: "Paste code", signal: controller.signal });
|
||||
} catch (error) {
|
||||
promptRejected.resolve(toError(error));
|
||||
}
|
||||
// The flow keeps running (e.g. the callback server resolves it) until we
|
||||
// resolve the follow-up prompt below.
|
||||
await interaction.prompt({ type: "text", message: "Waiting for callback" });
|
||||
}),
|
||||
});
|
||||
|
||||
expect(state.prompt).toMatchObject({ kind: "manual" });
|
||||
controller.abort();
|
||||
await expect(promptRejected.promise).resolves.toMatchObject({ message: "Prompt cancelled" });
|
||||
expect(removeAbortListener).toHaveBeenCalledWith("abort", expect.any(Function));
|
||||
|
||||
const afterAbort = service.get(state.flowId);
|
||||
expect(afterAbort.status).toBe("running");
|
||||
expect(afterAbort.prompt).toMatchObject({ kind: "prompt", message: "Waiting for callback" });
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
it("rejects pending prompts when cancelled", async () => {
|
||||
const promptRejected = deferred<Error>();
|
||||
const service = new OAuthLoginFlowService();
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
authStorage: fakeAuthStorage(async (_providerId, callbacks) => {
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
try {
|
||||
await callbacks.onPrompt({ message: "Paste code" });
|
||||
await interaction.prompt({ type: "text", message: "Paste code" });
|
||||
} catch (error) {
|
||||
promptRejected.resolve(toError(error));
|
||||
throw error;
|
||||
@@ -122,9 +347,9 @@ describe("OAuthLoginFlowService", () => {
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
authStorage: fakeAuthStorage(async (_providerId, callbacks) => {
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
try {
|
||||
await callbacks.onPrompt({ message: "Paste code" });
|
||||
await interaction.prompt({ type: "text", message: "Paste code" });
|
||||
} catch (error) {
|
||||
promptRejected.resolve(toError(error));
|
||||
throw error;
|
||||
@@ -137,7 +362,7 @@ describe("OAuthLoginFlowService", () => {
|
||||
service.dispose();
|
||||
|
||||
await expect(promptRejected.promise).resolves.toMatchObject({ message: "Login cancelled" });
|
||||
expect(() => { service.get(state.flowId); }).toThrow("OAuth login flow not found");
|
||||
expect(() => { service.get(state.flowId); }).toThrow("Login flow not found");
|
||||
});
|
||||
|
||||
it("rejects stale or duplicate responses", () => {
|
||||
@@ -145,8 +370,8 @@ describe("OAuthLoginFlowService", () => {
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
authStorage: fakeAuthStorage(async (_providerId, callbacks) => {
|
||||
await callbacks.onPrompt({ message: "Paste code" });
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
await interaction.prompt({ type: "text", message: "Paste code" });
|
||||
}),
|
||||
});
|
||||
|
||||
@@ -154,7 +379,7 @@ describe("OAuthLoginFlowService", () => {
|
||||
if (prompt === undefined) throw new Error("Expected prompt");
|
||||
|
||||
service.respond(state.flowId, prompt.requestId, "abc123");
|
||||
expect(() => { service.respond(state.flowId, prompt.requestId, "abc123"); }).toThrow("OAuth login request expired");
|
||||
expect(() => { service.respond(state.flowId, prompt.requestId, "abc123"); }).toThrow("Login request expired");
|
||||
service.dispose();
|
||||
});
|
||||
|
||||
@@ -165,9 +390,9 @@ describe("OAuthLoginFlowService", () => {
|
||||
const state = service.start({
|
||||
providerId: "test-provider",
|
||||
providerName: "Test Provider",
|
||||
authStorage: fakeAuthStorage(async (_providerId, callbacks) => {
|
||||
runtime: fakeRuntime(async (_providerId, interaction) => {
|
||||
try {
|
||||
await callbacks.onPrompt({ message: "Paste code" });
|
||||
await interaction.prompt({ type: "text", message: "Paste code" });
|
||||
} catch (error) {
|
||||
promptRejected.resolve(toError(error));
|
||||
throw error;
|
||||
@@ -177,18 +402,25 @@ describe("OAuthLoginFlowService", () => {
|
||||
|
||||
await vi.advanceTimersByTimeAsync(1000);
|
||||
|
||||
expect(service.get(state.flowId)).toMatchObject({ status: "error", error: "OAuth login flow expired" });
|
||||
await expect(promptRejected.promise).resolves.toMatchObject({ message: "OAuth login flow expired" });
|
||||
expect(service.get(state.flowId)).toMatchObject({ status: "error", error: "Login flow expired" });
|
||||
await expect(promptRejected.promise).resolves.toMatchObject({ message: "Login flow expired" });
|
||||
|
||||
await vi.advanceTimersByTimeAsync(1000);
|
||||
|
||||
expect(() => { service.get(state.flowId); }).toThrow("OAuth login flow not found");
|
||||
expect(() => { service.get(state.flowId); }).toThrow("Login flow not found");
|
||||
service.dispose();
|
||||
});
|
||||
});
|
||||
|
||||
function fakeAuthStorage(login: LoginHandler): Pick<AuthStorage, "login"> {
|
||||
return { login };
|
||||
function fakeRuntime(login: LoginHandler, authTypes?: AuthType[]): Pick<ModelRuntime, "login"> {
|
||||
return {
|
||||
login: (providerId, type, interaction) => {
|
||||
authTypes?.push(type);
|
||||
return login(providerId, interaction).then(() => type === "api_key"
|
||||
? { type: "api_key", key: "test" }
|
||||
: { type: "oauth", refresh: "r", access: "a", expires: 0 });
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async function flushAsyncLogin(): Promise<void> {
|
||||
|
||||
@@ -1,16 +1,21 @@
|
||||
import crypto from "node:crypto";
|
||||
import type { OAuthLoginCallbacks, OAuthSelectPrompt, OAuthPrompt } from "@earendil-works/pi-ai";
|
||||
import type { AuthStorage } from "@earendil-works/pi-coding-agent";
|
||||
import type { AuthEvent, AuthInteraction, AuthPrompt, AuthType } from "@earendil-works/pi-ai";
|
||||
import type { ModelRuntime } from "@earendil-works/pi-coding-agent";
|
||||
import type { CommandOption, OAuthFlowState } from "../../shared/apiTypes.js";
|
||||
|
||||
type OAuthLoginStorage = Pick<AuthStorage, "login">;
|
||||
/** The single runtime capability this service drives — narrowed for testable DI. */
|
||||
type OAuthLoginRuntime = Pick<ModelRuntime, "login">;
|
||||
type TimerHandle = ReturnType<typeof setTimeout>;
|
||||
type SelectPrompt = Extract<AuthPrompt, { type: "select" }>;
|
||||
type ValuePrompt = Exclude<AuthPrompt, { type: "select" }>;
|
||||
|
||||
interface PendingOAuthRequest {
|
||||
requestId: string;
|
||||
allowEmpty: boolean;
|
||||
resolve: (value: string | undefined) => void;
|
||||
resolve: (value: string) => void;
|
||||
reject: (error: Error) => void;
|
||||
allowedValues?: ReadonlySet<string>;
|
||||
cleanup?: () => void;
|
||||
}
|
||||
|
||||
interface OAuthFlowRecord {
|
||||
@@ -22,32 +27,46 @@ interface OAuthFlowRecord {
|
||||
cleanupTimer?: TimerHandle;
|
||||
}
|
||||
|
||||
export interface OAuthLoginFlowLogger {
|
||||
error(details: Record<string, unknown>, message: string): void;
|
||||
}
|
||||
|
||||
export interface OAuthLoginFlowServiceOptions {
|
||||
terminalTtlMs?: number;
|
||||
runningTtlMs?: number;
|
||||
now?: () => number;
|
||||
logger?: OAuthLoginFlowLogger;
|
||||
}
|
||||
|
||||
const DEFAULT_TERMINAL_TTL_MS = 5 * 60 * 1000;
|
||||
const DEFAULT_RUNNING_TTL_MS = 30 * 60 * 1000;
|
||||
const noopLogger: OAuthLoginFlowLogger = { error() { /* no-op */ } };
|
||||
|
||||
/**
|
||||
* AuthInteraction transport shared by OAuth and provider-driven API-key login.
|
||||
* The historical class and wire names remain for rolling browser/sessiond compatibility.
|
||||
*/
|
||||
export class OAuthLoginFlowService {
|
||||
private readonly flows = new Map<string, OAuthFlowRecord>();
|
||||
private readonly terminalTtlMs: number;
|
||||
private readonly runningTtlMs: number;
|
||||
private readonly now: () => number;
|
||||
private readonly logger: OAuthLoginFlowLogger;
|
||||
|
||||
constructor(options: OAuthLoginFlowServiceOptions = {}) {
|
||||
this.terminalTtlMs = options.terminalTtlMs ?? DEFAULT_TERMINAL_TTL_MS;
|
||||
this.runningTtlMs = options.runningTtlMs ?? DEFAULT_RUNNING_TTL_MS;
|
||||
this.now = options.now ?? (() => Date.now());
|
||||
this.logger = options.logger ?? noopLogger;
|
||||
}
|
||||
|
||||
start(options: {
|
||||
providerId: string;
|
||||
providerName: string;
|
||||
authStorage: OAuthLoginStorage;
|
||||
onComplete?: () => void;
|
||||
runtime: OAuthLoginRuntime;
|
||||
/** Defaults to OAuth so established callers retain their existing behavior. */
|
||||
authType?: AuthType;
|
||||
onComplete?: () => void | Promise<void>;
|
||||
}): OAuthFlowState {
|
||||
const flowId = crypto.randomUUID();
|
||||
const abort = new AbortController();
|
||||
@@ -66,58 +85,43 @@ export class OAuthLoginFlowService {
|
||||
this.flows.set(flowId, record);
|
||||
this.scheduleRunningExpiry(record);
|
||||
|
||||
const callbacks: OAuthLoginCallbacks = {
|
||||
// Adapt the pi-ai AuthInteraction contract onto the web-UI flow state:
|
||||
// `prompt()` returns the entered/selected string; `notify()` surfaces
|
||||
// out-of-band login events (auth URL, device code, progress).
|
||||
const interaction: AuthInteraction = {
|
||||
signal: abort.signal,
|
||||
onAuth: (info) => {
|
||||
if (!this.isCurrentRunning(record)) return;
|
||||
this.updateState(record, { ...record.state, auth: info });
|
||||
},
|
||||
// Device-code flows have no redirect URL; reuse the auth field so the web UI
|
||||
// shows the verification link and user code without a dedicated API shape.
|
||||
onDeviceCode: (info) => {
|
||||
if (!this.isCurrentRunning(record)) return;
|
||||
this.updateState(record, { ...record.state, auth: { url: info.verificationUri, instructions: `Enter code: ${info.userCode}` } });
|
||||
},
|
||||
onPrompt: (prompt) => this.waitForPrompt(record, prompt, "prompt"),
|
||||
onManualCodeInput: () => this.waitForPrompt(record, { message: "Paste the callback URL or authorization code", allowEmpty: false }, "manual"),
|
||||
onSelect: (prompt) => this.waitForSelect(record, prompt),
|
||||
onProgress: (message) => {
|
||||
if (!this.isCurrentRunning(record)) return;
|
||||
this.updateState(record, { ...record.state, progress: [...record.state.progress, message] });
|
||||
},
|
||||
prompt: (prompt) => this.handlePrompt(record, prompt),
|
||||
notify: (event) => { this.handleEvent(record, event); },
|
||||
};
|
||||
|
||||
void options.authStorage.login(options.providerId, callbacks)
|
||||
.then(() => {
|
||||
if (!this.isCurrentRunning(record)) return;
|
||||
record.pending = undefined;
|
||||
this.markTerminal(record, { ...withoutInteraction(record.state), status: "complete", progress: [...record.state.progress, "Login complete"] });
|
||||
options.onComplete?.();
|
||||
})
|
||||
.catch((error: unknown) => {
|
||||
if (this.flows.get(record.flowId) !== record) return;
|
||||
record.pending = undefined;
|
||||
void options.runtime.login(options.providerId, options.authType ?? "oauth", interaction).then(
|
||||
() => this.reconcileCommittedLogin(record, options.onComplete),
|
||||
(error: unknown) => {
|
||||
if (!this.isCurrent(record)) return;
|
||||
this.clearPending(record);
|
||||
if (record.state.status !== "running") return;
|
||||
this.markTerminal(record, { ...withoutInteraction(record.state), status: "error", error: error instanceof Error ? error.message : String(error) });
|
||||
});
|
||||
},
|
||||
);
|
||||
|
||||
return this.get(flowId);
|
||||
}
|
||||
|
||||
get(flowId: string): OAuthFlowState {
|
||||
const record = this.flows.get(flowId);
|
||||
if (record === undefined) throw new Error("OAuth login flow not found");
|
||||
if (record === undefined) throw new Error("Login flow not found");
|
||||
return cloneState(record.state);
|
||||
}
|
||||
|
||||
respond(flowId: string, requestId: string, value: string): OAuthFlowState {
|
||||
const record = this.flows.get(flowId);
|
||||
if (record === undefined) throw new Error("OAuth login flow not found");
|
||||
if (record === undefined) throw new Error("Login flow not found");
|
||||
if (record.state.status !== "running") return cloneState(record.state);
|
||||
const pending = record.pending;
|
||||
if (pending?.requestId !== requestId) throw new Error("OAuth login request expired");
|
||||
if (pending?.requestId !== requestId) throw new Error("Login request expired");
|
||||
if (!pending.allowEmpty && value.trim() === "") throw new Error("A value is required");
|
||||
record.pending = undefined;
|
||||
if (pending.allowedValues !== undefined && !pending.allowedValues.has(value)) throw new Error("Invalid login selection");
|
||||
this.clearPending(record);
|
||||
this.updateState(record, withoutInteraction(record.state));
|
||||
pending.resolve(value);
|
||||
return cloneState(record.state);
|
||||
@@ -125,11 +129,10 @@ export class OAuthLoginFlowService {
|
||||
|
||||
cancel(flowId: string): OAuthFlowState {
|
||||
const record = this.flows.get(flowId);
|
||||
if (record === undefined) throw new Error("OAuth login flow not found");
|
||||
if (record === undefined) throw new Error("Login flow not found");
|
||||
if (record.state.status === "running") {
|
||||
record.abort.abort();
|
||||
const pending = record.pending;
|
||||
record.pending = undefined;
|
||||
const pending = this.clearPending(record);
|
||||
this.markTerminal(record, { ...withoutInteraction(record.state), status: "cancelled", error: "Login cancelled" });
|
||||
pending?.reject(new Error("Login cancelled"));
|
||||
}
|
||||
@@ -140,51 +143,179 @@ export class OAuthLoginFlowService {
|
||||
for (const record of this.flows.values()) {
|
||||
this.clearTimer(record);
|
||||
record.abort.abort();
|
||||
const pending = record.pending;
|
||||
record.pending = undefined;
|
||||
const pending = this.clearPending(record);
|
||||
pending?.reject(new Error("Login cancelled"));
|
||||
}
|
||||
this.flows.clear();
|
||||
}
|
||||
|
||||
private waitForPrompt(record: OAuthFlowRecord, prompt: OAuthPrompt, kind: "prompt" | "manual"): Promise<string> {
|
||||
private handlePrompt(record: OAuthFlowRecord, prompt: AuthPrompt): Promise<string> {
|
||||
if (prompt.type === "select") return this.waitForSelect(record, prompt);
|
||||
return this.waitForPrompt(record, prompt);
|
||||
}
|
||||
|
||||
private handleEvent(record: OAuthFlowRecord, event: AuthEvent): void {
|
||||
if (!this.isCurrentRunning(record)) return;
|
||||
switch (event.type) {
|
||||
case "auth_url":
|
||||
this.updateState(record, { ...record.state, auth: { url: event.url, ...(event.instructions === undefined ? {} : { instructions: event.instructions }) } });
|
||||
return;
|
||||
// Keep the legacy auth URL/instructions while adding structured metadata
|
||||
// that newer browsers can use during rolling sessiond upgrades.
|
||||
case "device_code":
|
||||
this.updateState(record, {
|
||||
...record.state,
|
||||
auth: {
|
||||
url: event.verificationUri,
|
||||
instructions: `Enter code: ${event.userCode}`,
|
||||
deviceCode: {
|
||||
userCode: event.userCode,
|
||||
...(event.intervalSeconds === undefined ? {} : { intervalSeconds: event.intervalSeconds }),
|
||||
...(event.expiresInSeconds === undefined ? {} : { expiresInSeconds: event.expiresInSeconds }),
|
||||
},
|
||||
},
|
||||
});
|
||||
return;
|
||||
case "info":
|
||||
this.updateState(record, {
|
||||
...record.state,
|
||||
progress: [...record.state.progress, event.message],
|
||||
info: [
|
||||
...(record.state.info ?? []),
|
||||
{
|
||||
message: event.message,
|
||||
...(event.links === undefined ? {} : {
|
||||
links: event.links.map((link) => ({
|
||||
url: link.url,
|
||||
...(link.label === undefined ? {} : { label: link.label }),
|
||||
})),
|
||||
}),
|
||||
},
|
||||
],
|
||||
});
|
||||
return;
|
||||
case "progress":
|
||||
this.updateState(record, { ...record.state, progress: [...record.state.progress, event.message] });
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
private waitForPrompt(record: OAuthFlowRecord, prompt: ValuePrompt): Promise<string> {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (!this.isCurrentRunning(record)) {
|
||||
reject(new Error("Login cancelled"));
|
||||
return;
|
||||
}
|
||||
const requestId = crypto.randomUUID();
|
||||
record.pending = { requestId, allowEmpty: prompt.allowEmpty === true, resolve: (value) => { resolve(value ?? ""); }, reject };
|
||||
const pending: PendingOAuthRequest = {
|
||||
requestId,
|
||||
allowEmpty: prompt.type === "text",
|
||||
resolve,
|
||||
reject,
|
||||
};
|
||||
record.pending = pending;
|
||||
if (!this.bindPromptSignal(record, pending, prompt.signal)) return;
|
||||
const base = withoutInteraction(record.state);
|
||||
this.updateState(record, {
|
||||
...base,
|
||||
prompt: {
|
||||
requestId,
|
||||
message: prompt.message,
|
||||
kind,
|
||||
kind: prompt.type === "manual_code" ? "manual" : "prompt",
|
||||
promptType: prompt.type,
|
||||
...(prompt.type === "text" ? { allowEmpty: true } : {}),
|
||||
...(prompt.placeholder === undefined ? {} : { placeholder: prompt.placeholder }),
|
||||
...(prompt.allowEmpty === true ? { allowEmpty: true } : {}),
|
||||
},
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
private waitForSelect(record: OAuthFlowRecord, prompt: OAuthSelectPrompt): Promise<string | undefined> {
|
||||
private waitForSelect(record: OAuthFlowRecord, prompt: SelectPrompt): Promise<string> {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (!this.isCurrentRunning(record)) {
|
||||
reject(new Error("Login cancelled"));
|
||||
return;
|
||||
}
|
||||
const requestId = crypto.randomUUID();
|
||||
const options: CommandOption[] = prompt.options.map((option) => ({ value: option.id, label: option.label }));
|
||||
record.pending = { requestId, allowEmpty: true, resolve, reject };
|
||||
const options: CommandOption[] = prompt.options.map((option) => ({
|
||||
value: option.id,
|
||||
label: option.label,
|
||||
...(option.description === undefined ? {} : { description: option.description }),
|
||||
}));
|
||||
const pending: PendingOAuthRequest = {
|
||||
requestId,
|
||||
allowEmpty: false,
|
||||
resolve,
|
||||
reject,
|
||||
allowedValues: new Set(options.map((option) => option.value)),
|
||||
};
|
||||
record.pending = pending;
|
||||
if (!this.bindPromptSignal(record, pending, prompt.signal)) return;
|
||||
const base = withoutInteraction(record.state);
|
||||
this.updateState(record, { ...base, select: { requestId, message: prompt.message, options } });
|
||||
});
|
||||
}
|
||||
|
||||
// A prompt may carry its own AbortSignal (e.g. a manual_code prompt raced
|
||||
// against a callback server). When it fires, drop just that pending request
|
||||
// and clear the interaction from state — the overall login keeps running.
|
||||
private bindPromptSignal(record: OAuthFlowRecord, pending: PendingOAuthRequest, signal?: AbortSignal): boolean {
|
||||
if (signal === undefined) return true;
|
||||
const onAbort = () => {
|
||||
if (record.pending !== pending) return;
|
||||
this.clearPending(record);
|
||||
if (this.isCurrentRunning(record)) this.updateState(record, withoutInteraction(record.state));
|
||||
pending.reject(new Error("Prompt cancelled"));
|
||||
};
|
||||
pending.cleanup = () => { signal.removeEventListener("abort", onAbort); };
|
||||
signal.addEventListener("abort", onAbort, { once: true });
|
||||
if (signal.aborted) {
|
||||
onAbort();
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
private clearPending(record: OAuthFlowRecord): PendingOAuthRequest | undefined {
|
||||
const pending = record.pending;
|
||||
record.pending = undefined;
|
||||
pending?.cleanup?.();
|
||||
return pending;
|
||||
}
|
||||
|
||||
// ModelRuntime persists the credential before its post-login refresh. If a
|
||||
// cancellation lands during that refresh, the resolved login is committed
|
||||
// truth and must supersede the transient cancelled state.
|
||||
private async reconcileCommittedLogin(record: OAuthFlowRecord, onComplete?: () => void | Promise<void>): Promise<void> {
|
||||
if (this.isCurrent(record)) this.clearPending(record);
|
||||
try {
|
||||
await onComplete?.();
|
||||
} catch (error) {
|
||||
this.logErrorNoThrow(
|
||||
{ err: error, flowId: record.flowId, providerId: record.state.providerId },
|
||||
"login completion callback failed",
|
||||
);
|
||||
}
|
||||
if (!this.isCurrent(record)) return;
|
||||
const completed = withoutInteraction(record.state);
|
||||
delete completed.error;
|
||||
this.markTerminal(record, { ...completed, status: "complete", progress: [...record.state.progress, "Login complete"] });
|
||||
}
|
||||
|
||||
private isCurrent(record: OAuthFlowRecord): boolean {
|
||||
return this.flows.get(record.flowId) === record;
|
||||
}
|
||||
|
||||
private isCurrentRunning(record: OAuthFlowRecord): boolean {
|
||||
return this.flows.get(record.flowId) === record && record.state.status === "running";
|
||||
return this.isCurrent(record) && record.state.status === "running";
|
||||
}
|
||||
|
||||
private logErrorNoThrow(details: Record<string, unknown>, message: string): void {
|
||||
try {
|
||||
this.logger.error(details, message);
|
||||
} catch {
|
||||
// Logging is post-commit diagnostics and must never change auth truth.
|
||||
}
|
||||
}
|
||||
|
||||
private updateState(record: OAuthFlowRecord, state: OAuthFlowState): void {
|
||||
@@ -226,10 +357,9 @@ export class OAuthLoginFlowService {
|
||||
private expireRunningFlow(record: OAuthFlowRecord): void {
|
||||
if (!this.isCurrentRunning(record)) return;
|
||||
record.abort.abort();
|
||||
const pending = record.pending;
|
||||
record.pending = undefined;
|
||||
this.markTerminal(record, { ...withoutInteraction(record.state), status: "error", error: "OAuth login flow expired" });
|
||||
pending?.reject(new Error("OAuth login flow expired"));
|
||||
const pending = this.clearPending(record);
|
||||
this.markTerminal(record, { ...withoutInteraction(record.state), status: "error", error: "Login flow expired" });
|
||||
pending?.reject(new Error("Login flow expired"));
|
||||
}
|
||||
|
||||
private setTimer(record: OAuthFlowRecord, delayMs: number, callback: () => void): void {
|
||||
@@ -256,9 +386,20 @@ function cloneState(state: OAuthFlowState): OAuthFlowState {
|
||||
return {
|
||||
...state,
|
||||
progress: [...state.progress],
|
||||
...(state.auth === undefined ? {} : { auth: { ...state.auth } }),
|
||||
...(state.auth === undefined ? {} : {
|
||||
auth: {
|
||||
...state.auth,
|
||||
...(state.auth.deviceCode === undefined ? {} : { deviceCode: { ...state.auth.deviceCode } }),
|
||||
},
|
||||
}),
|
||||
...(state.prompt === undefined ? {} : { prompt: { ...state.prompt } }),
|
||||
...(state.select === undefined ? {} : { select: { ...state.select, options: state.select.options.map((option) => ({ ...option })) } }),
|
||||
...(state.info === undefined ? {} : {
|
||||
info: state.info.map((item) => ({
|
||||
...item,
|
||||
...(item.links === undefined ? {} : { links: item.links.map((link) => ({ ...link })) }),
|
||||
})),
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { PiSessionService } from "./piSessionService.js";
|
||||
import { CapturingSessionEventHub, fakeRuntime, fakeSessionManager, runtimeCreator, sessionGateway, sessionRecord, sessionRef } from "./piSessionService.testSupport.js";
|
||||
import { CapturingSessionEventHub, fakeRuntime, fakeSessionManager, runtimeCreator, sessionGateway, sessionRecord, sessionRef, testModelRuntime } from "./piSessionService.testSupport.js";
|
||||
|
||||
const TEST_AGENT_DIR = "/tmp/pi-web-test-agent";
|
||||
|
||||
@@ -15,6 +15,7 @@ describe("PiSessionService archive and cleanup", () => {
|
||||
const fake = fakeRuntime("root", { sessionFile: root.path });
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([{ sessionId: "archived-child", cwd: "/workspace", archivedAt: "2026-01-02T00:00:00.000Z", originalPath: archivedChild.path, archivePath: "/archive/archived-child.jsonl", created: "2026-01-01T00:00:00.000Z", modified: "2026-01-01T00:01:00.000Z", messageCount: 1, firstMessage: "archived", parentSessionPath: root.path }]),
|
||||
@@ -49,6 +50,7 @@ describe("PiSessionService archive and cleanup", () => {
|
||||
const deletedSessionIds: string[] = [];
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([]),
|
||||
get: (sessionId) => Promise.resolve(sessionId === "archived" || "archived".startsWith(sessionId)
|
||||
@@ -83,6 +85,7 @@ describe("PiSessionService archive and cleanup", () => {
|
||||
const archiveMany = vi.fn((inputs: readonly { sessionId: string; cwd: string }[]) => Promise.resolve(inputs.map((input) => ({ sessionId: input.sessionId, cwd: input.cwd, archivedAt: "2026-01-03T00:00:00.000Z" }))));
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([]),
|
||||
get: () => Promise.resolve(undefined),
|
||||
@@ -118,6 +121,7 @@ describe("PiSessionService archive and cleanup", () => {
|
||||
const archiveMany = vi.fn((inputs: readonly { sessionId: string; cwd: string }[]) => Promise.resolve(inputs.map((input) => ({ sessionId: input.sessionId, cwd: input.cwd, archivedAt: "2026-01-03T00:00:00.000Z" }))));
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: () => {
|
||||
createCalls += 1;
|
||||
return Promise.resolve(busy.runtime);
|
||||
@@ -159,6 +163,7 @@ describe("PiSessionService archive and cleanup", () => {
|
||||
const deleteArchivedMany = vi.fn((sessionIds: readonly string[]) => Promise.resolve([...sessionIds]));
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(busy.runtime),
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([busyRecord, idleRecord]),
|
||||
@@ -196,6 +201,7 @@ describe("PiSessionService archive and cleanup", () => {
|
||||
const listCalls: string[] = [];
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([
|
||||
{ sessionId: "legacy-a", cwd: "/workspace", archivedAt: "2026-01-02T00:00:00.000Z" },
|
||||
@@ -239,6 +245,7 @@ describe("PiSessionService archive and cleanup", () => {
|
||||
const otherArchived = { sessionId: "archived-other", cwd: "/other-project", archivedAt: "2026-04-01T00:00:00.000Z", archivePath: "/archive/archived-other.jsonl" };
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
now: () => new Date("2026-06-25T00:00:00.000Z"),
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([archived, otherArchived]),
|
||||
@@ -292,6 +299,7 @@ describe("PiSessionService archive and cleanup", () => {
|
||||
const deleteArchivedMany = vi.fn((sessionIds: readonly string[]) => Promise.resolve([...sessionIds]));
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
now: () => new Date("2026-06-25T00:00:00.000Z"),
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([
|
||||
@@ -334,6 +342,7 @@ describe("PiSessionService archive and cleanup", () => {
|
||||
const archivedInputs: string[] = [];
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
now: () => new Date("2026-06-25T00:00:00.000Z"),
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
archiveStore: {
|
||||
|
||||
@@ -3,7 +3,7 @@ import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { PiSessionService, type PiAgentSession, type PiSessionRuntime } from "./piSessionService.js";
|
||||
import { CapturingSessionEventHub, emptyArchiveStore, fakeRuntime, fakeSessionManager, runtimeCreator, sessionGateway, sessionRecord, sessionRef, type RuntimeCreator } from "./piSessionService.testSupport.js";
|
||||
import { CapturingSessionEventHub, emptyArchiveStore, fakeRuntime, fakeSessionManager, runtimeCreator, sessionGateway, sessionRecord, sessionRef, testModelRuntime, type RuntimeCreator } from "./piSessionService.testSupport.js";
|
||||
|
||||
const TEST_AGENT_DIR = "/tmp/pi-web-test-agent";
|
||||
|
||||
@@ -31,6 +31,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
};
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime,
|
||||
sessionManager: sessionGateway([]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -60,6 +61,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
try {
|
||||
service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -87,6 +89,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
const open = vi.fn(() => fakeSessionManager());
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: {
|
||||
create: () => fakeSessionManager(),
|
||||
@@ -136,6 +139,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
const open = vi.spyOn(gateway, "open");
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: emptyArchiveStore(),
|
||||
createAgentRuntime,
|
||||
sessionManager: gateway,
|
||||
@@ -190,6 +194,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
};
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: emptyArchiveStore(),
|
||||
createAgentRuntime,
|
||||
sessionManager: sessionGateway([sessionRecord(sessionId)]),
|
||||
@@ -230,6 +235,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
const fake = fakeRuntime(sessionId);
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: emptyArchiveStore(),
|
||||
createAgentRuntime: () => {
|
||||
createStarted.resolve();
|
||||
@@ -264,6 +270,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
fake.runtime.setRebindSession = (callback) => { rebindSession = callback; };
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -291,6 +298,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
});
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -326,6 +334,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
});
|
||||
service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("idle-session")]),
|
||||
heartbeatIntervalMs: 1_000,
|
||||
@@ -362,6 +371,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
});
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("completion-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -381,6 +391,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
it("uses injected archive and session-manager gateways for listing", async () => {
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([{ sessionId: "archived", cwd: "/workspace", archivedAt: "2026-01-01T00:00:00.000Z" }]),
|
||||
get: () => Promise.resolve(undefined),
|
||||
@@ -411,6 +422,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
it("lists archived records that have been moved out of the active session directory", async () => {
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([{ sessionId: "archived", cwd: "/workspace", archivedAt: "2026-01-02T00:00:00.000Z", originalPath: "/sessions/archived.jsonl", archivePath: "/archive/archived.jsonl", created: "2026-01-01T00:00:00.000Z", modified: "2026-01-01T00:01:00.000Z", messageCount: 2, firstMessage: "bye" }]),
|
||||
get: () => Promise.resolve(undefined),
|
||||
@@ -442,6 +454,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
const fake = fakeRuntime("runtime-reload-session");
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("runtime-reload-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -475,6 +488,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
};
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime,
|
||||
sessionManager: sessionGateway([sessionRecord("reload-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -499,6 +513,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
const fake = fakeRuntime("busy-session", { isStreaming: true });
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("busy-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -514,6 +529,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
it("refuses to reload an archived session", async () => {
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([]),
|
||||
get: (sessionId) => Promise.resolve(sessionId === "archived" || "archived".startsWith(sessionId)
|
||||
@@ -536,6 +552,7 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
|
||||
const reconciliations: { cwd: string; sessionIds: string[] }[] = [];
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: {
|
||||
list: () => Promise.resolve([{ sessionId: "archived", cwd: "/workspace", archivedAt: "2026-01-02T00:00:00.000Z", originalPath: "/sessions/archived.jsonl", archivePath: "/archive/archived.jsonl", created: "2026-01-01T00:00:00.000Z", modified: "2026-01-01T00:01:00.000Z", messageCount: 2, firstMessage: "bye" }]),
|
||||
get: () => Promise.resolve(undefined),
|
||||
@@ -573,6 +590,7 @@ describe("PiSessionService.streamSnapshot", () => {
|
||||
const fake = fakeRuntime("snap-idle");
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -601,6 +619,7 @@ describe("PiSessionService.streamSnapshot", () => {
|
||||
const fake = fakeRuntime("snap-live", { state: { streamingMessage } });
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
import { createAssistantMessageEventStream, type AssistantMessage } from "@earendil-works/pi-ai";
|
||||
import { mkdtemp, rm, writeFile } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { createAssistantMessageEventStream, InMemoryCredentialStore, type AssistantMessage } from "@earendil-works/pi-ai";
|
||||
import type { StreamFn } from "@earendil-works/pi-agent-core";
|
||||
import { AuthStorage, ModelRegistry } from "@earendil-works/pi-coding-agent";
|
||||
import { ModelRuntime } from "@earendil-works/pi-coding-agent";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { PiSessionService } from "./piSessionService.js";
|
||||
import { CapturingSessionEventHub, fakeRuntime, runtimeCreator, sessionGateway, sessionRecord, sessionRef, TEST_MODEL_ID, TEST_MODEL_PROVIDER, testModel, type RuntimeCreator } from "./piSessionService.testSupport.js";
|
||||
import { CapturingSessionEventHub, createTestModelRuntime, fakeRuntime, runtimeCreator, seedCredential, sessionGateway, sessionRecord, sessionRef, TEST_MODEL_ID, TEST_MODEL_PROVIDER, testModel, testModelRuntime, type RuntimeCreator } from "./piSessionService.testSupport.js";
|
||||
|
||||
const TEST_AGENT_DIR = "/tmp/pi-web-test-agent";
|
||||
|
||||
@@ -12,6 +15,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
const fake = fakeRuntime("prompt-session");
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("prompt-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -30,6 +34,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
const hub = new CapturingSessionEventHub();
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("echo-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -60,6 +65,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
};
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime,
|
||||
sessionManager: sessionGateway([sessionRecord("prompt-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -96,6 +102,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
const fake = fakeRuntime("name-session", { model, agent: { streamFn } });
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("name-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -118,6 +125,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
});
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("status-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -139,6 +147,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
});
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("dedupe-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -155,6 +164,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
const fake = fakeRuntime("queued-session", { isStreaming: true });
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("queued-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -181,6 +191,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
};
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("compacting-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -250,6 +261,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
fake.session.clearQueue = clearRuntimeQueue;
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("clear-queue-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -290,6 +302,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
const fake = fakeRuntime("clear-empty-queue-session");
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("clear-empty-queue-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -309,6 +322,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
const fake = fakeRuntime("abort-session");
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("abort-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -326,6 +340,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
const fake = fakeRuntime("abort-compaction-session", { isCompacting: true });
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("abort-compaction-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -341,17 +356,66 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("reloads models.json before listing and selecting models", async () => {
|
||||
const agentDir = await mkdtemp(join(tmpdir(), "pi-web-model-runtime-"));
|
||||
try {
|
||||
const modelsPath = join(agentDir, "models.json");
|
||||
await writeLocalModelsConfig(modelsPath, "initial-model");
|
||||
const modelRuntime = await ModelRuntime.create({
|
||||
credentials: new InMemoryCredentialStore(),
|
||||
modelsPath,
|
||||
allowModelNetwork: false,
|
||||
});
|
||||
const setSessionModel = vi.fn(() => Promise.resolve());
|
||||
const fake = fakeRuntime("models-session", { modelRuntime, setModel: setSessionModel });
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir,
|
||||
modelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("models-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
});
|
||||
|
||||
try {
|
||||
await writeLocalModelsConfig(modelsPath, "listed-model");
|
||||
const listed = await service.availableModels(sessionRef("models-session"));
|
||||
expect(listed).toEqual(expect.arrayContaining([
|
||||
expect.objectContaining({ provider: "test-local", id: "listed-model" }),
|
||||
]));
|
||||
expect(listed).not.toEqual(expect.arrayContaining([
|
||||
expect.objectContaining({ provider: "test-local", id: "initial-model" }),
|
||||
]));
|
||||
|
||||
await writeLocalModelsConfig(modelsPath, "selected-model");
|
||||
await expect(service.setModel(sessionRef("models-session"), "test-local", "selected-model")).resolves.toBeDefined();
|
||||
expect(setSessionModel).toHaveBeenCalledWith(expect.objectContaining({
|
||||
provider: "test-local",
|
||||
id: "selected-model",
|
||||
}));
|
||||
} finally {
|
||||
await service.dispose();
|
||||
}
|
||||
} finally {
|
||||
await rm(agentDir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("refreshes auth state and dedupes warnings when logout removes the current model's credentials", async () => {
|
||||
const hub = new CapturingSessionEventHub();
|
||||
const authStorage = AuthStorage.inMemory({ anthropic: { type: "api_key", key: "sk-test" } });
|
||||
const modelRegistry = ModelRegistry.inMemory(authStorage);
|
||||
const model = modelRegistry.find(TEST_MODEL_PROVIDER, TEST_MODEL_ID);
|
||||
// The shared model runtime reads a live credential store. Mutating the store
|
||||
// and refreshing here simulates the committed snapshot that
|
||||
// ModelRuntime.login()/logout() establishes before AuthService emits.
|
||||
// applyAuthChange then only needs to notify active sessions.
|
||||
const credentials = new InMemoryCredentialStore();
|
||||
await seedCredential(credentials, "anthropic", { type: "api_key", key: "sk-test" });
|
||||
const modelRuntime = await createTestModelRuntime(credentials);
|
||||
const model = modelRuntime.getModel(TEST_MODEL_PROVIDER, TEST_MODEL_ID);
|
||||
if (model === undefined) throw new Error("Expected Anthropic model fixture");
|
||||
const fake = fakeRuntime("auth-session", { model, modelRegistry });
|
||||
const fake = fakeRuntime("auth-session", { model, modelRuntime });
|
||||
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRegistry,
|
||||
modelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("auth-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -361,7 +425,8 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
hub.sessionEvents.length = 0;
|
||||
hub.globalEvents.length = 0;
|
||||
|
||||
authStorage.logout("anthropic");
|
||||
await credentials.delete("anthropic");
|
||||
await modelRuntime.refresh();
|
||||
service.applyAuthChange({ removedProviderId: "anthropic" });
|
||||
service.applyAuthChange({ removedProviderId: "anthropic" });
|
||||
|
||||
@@ -369,9 +434,11 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
expect(warningCount()).toBe(1);
|
||||
expect(hub.globalEvents.some((event) => event.type === "status.update" && event.status.sessionId === "auth-session")).toBe(true);
|
||||
|
||||
authStorage.set("anthropic", { type: "api_key", key: "sk-new" });
|
||||
await seedCredential(credentials, "anthropic", { type: "api_key", key: "sk-new" });
|
||||
await modelRuntime.refresh();
|
||||
service.applyAuthChange();
|
||||
authStorage.logout("anthropic");
|
||||
await credentials.delete("anthropic");
|
||||
await modelRuntime.refresh();
|
||||
service.applyAuthChange({ removedProviderId: "anthropic" });
|
||||
expect(warningCount()).toBe(2);
|
||||
|
||||
@@ -382,6 +449,7 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
const fake = fakeRuntime("stop-session");
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord("stop-session")]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
@@ -394,3 +462,25 @@ describe("PiSessionService prompt, queue, and auth warnings", () => {
|
||||
await service.dispose();
|
||||
});
|
||||
});
|
||||
|
||||
async function writeLocalModelsConfig(path: string, modelId: string): Promise<void> {
|
||||
await writeFile(path, JSON.stringify({
|
||||
providers: {
|
||||
"test-local": {
|
||||
name: "Test Local",
|
||||
baseUrl: "http://127.0.0.1:1234/v1",
|
||||
apiKey: "offline-test-key",
|
||||
api: "openai-completions",
|
||||
models: [{
|
||||
id: modelId,
|
||||
name: modelId,
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 1_000,
|
||||
maxTokens: 100,
|
||||
}],
|
||||
},
|
||||
},
|
||||
}));
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { PiSessionService, type PiAgentSession } from "./piSessionService.js";
|
||||
import type { SpawnTargetDecision } from "./spawnTargetResolver.js";
|
||||
import { CapturingSessionEventHub, fakeRuntime, runtimeCreator, sessionGateway, testModel, type RuntimeCreator } from "./piSessionService.testSupport.js";
|
||||
import { CapturingSessionEventHub, fakeRuntime, runtimeCreator, sessionGateway, testModel, testModelRuntime, type RuntimeCreator } from "./piSessionService.testSupport.js";
|
||||
|
||||
const TEST_AGENT_DIR = "/tmp/pi-web-test-agent";
|
||||
|
||||
@@ -12,6 +12,7 @@ describe("PiSessionService", () => {
|
||||
const log: { details: Record<string, unknown>; message: string }[] = [];
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([]),
|
||||
spawnTargets: { resolveSpawnTarget: () => Promise.resolve(decision) },
|
||||
@@ -45,6 +46,7 @@ describe("PiSessionService", () => {
|
||||
};
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime,
|
||||
sessionManager: sessionGateway([]),
|
||||
spawnTargets: { resolveSpawnTarget: () => Promise.resolve({ allowed: true, cwd: "/workspace-feature" }) },
|
||||
@@ -80,6 +82,7 @@ describe("PiSessionService", () => {
|
||||
const fake = fakeRuntime("spawned-x");
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
|
||||
@@ -4,7 +4,7 @@ import { join } from "node:path";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { PiSessionService, type PiAgentSession } from "./piSessionService.js";
|
||||
import type { SpawnTargetDecision } from "./spawnTargetResolver.js";
|
||||
import { CapturingSessionEventHub, emptyArchiveStore, fakeRuntime, fakeSessionManager, runtimeCreator, sessionGateway, sessionRecord, sessionRef, testModel, type RuntimeCreator } from "./piSessionService.testSupport.js";
|
||||
import { CapturingSessionEventHub, emptyArchiveStore, fakeRuntime, fakeSessionManager, runtimeCreator, sessionGateway, sessionRecord, sessionRef, testModel, testModelRuntime, type RuntimeCreator } from "./piSessionService.testSupport.js";
|
||||
|
||||
const TEST_AGENT_DIR = "/tmp/pi-web-test-agent";
|
||||
|
||||
@@ -40,6 +40,7 @@ describe("PiSessionService", () => {
|
||||
};
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime,
|
||||
sessionManager: sessionGateway([]),
|
||||
archiveStore,
|
||||
@@ -81,6 +82,7 @@ describe("PiSessionService", () => {
|
||||
};
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime,
|
||||
sessionManager: sessionGateway([]),
|
||||
archiveStore: emptyArchiveStore(),
|
||||
@@ -121,6 +123,7 @@ describe("PiSessionService", () => {
|
||||
let index = 0;
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: () => {
|
||||
const runtime = runtimes[index] ?? child.runtime;
|
||||
index += 1;
|
||||
@@ -173,6 +176,7 @@ describe("PiSessionService", () => {
|
||||
const open = vi.fn(() => childManager);
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: () => {
|
||||
const runtime = runtimes[index] ?? child.runtime;
|
||||
index += 1;
|
||||
@@ -215,6 +219,7 @@ describe("PiSessionService", () => {
|
||||
});
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(parent.runtime),
|
||||
sessionManager: { create: () => parent.session.sessionManager, list: () => Promise.resolve([]), listAll: () => Promise.resolve([]), open: () => fakeSessionManager() },
|
||||
archiveStore: emptyArchiveStore(),
|
||||
@@ -240,6 +245,7 @@ describe("PiSessionService", () => {
|
||||
});
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(parent.runtime),
|
||||
sessionManager: { create: () => parent.session.sessionManager, list: () => Promise.resolve([]), listAll: () => Promise.resolve([]), open: () => fakeSessionManager() },
|
||||
archiveStore: emptyArchiveStore(),
|
||||
@@ -262,6 +268,7 @@ describe("PiSessionService", () => {
|
||||
});
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(parent.runtime),
|
||||
sessionManager: { create: () => parent.session.sessionManager, list: () => Promise.resolve([]), listAll: () => Promise.resolve([]), open: () => fakeSessionManager() },
|
||||
archiveStore: emptyArchiveStore(),
|
||||
@@ -283,6 +290,7 @@ describe("PiSessionService", () => {
|
||||
});
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(parent.runtime),
|
||||
sessionManager: { create: () => parent.session.sessionManager, list: () => Promise.resolve([]), listAll: () => Promise.resolve([childRecord]), open: () => fakeSessionManager() },
|
||||
archiveStore: emptyArchiveStore(),
|
||||
@@ -304,6 +312,7 @@ describe("PiSessionService", () => {
|
||||
});
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(forkedParent.runtime),
|
||||
sessionManager: { create: () => forkedParent.session.sessionManager, list: () => Promise.resolve([]), listAll: () => Promise.resolve([]), open: () => fakeSessionManager() },
|
||||
archiveStore: emptyArchiveStore(),
|
||||
@@ -343,6 +352,7 @@ describe("PiSessionService", () => {
|
||||
const open = vi.fn((path: string) => path === parentFile ? parentManager : childManager);
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: (_createRuntime, options) => {
|
||||
delegationCapabilities.push(options.delegationToolsEnabled);
|
||||
const runtime = runtimes[index] ?? parent.runtime;
|
||||
@@ -409,6 +419,7 @@ describe("PiSessionService", () => {
|
||||
});
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: () => {
|
||||
const runtime = runtimes[index] ?? parent.runtime;
|
||||
index += 1;
|
||||
@@ -469,6 +480,7 @@ describe("PiSessionService", () => {
|
||||
const open = vi.fn((path: string) => path === parentFile ? parentManager : childManager);
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: () => {
|
||||
const runtime = runtimes[index] ?? parent.runtime;
|
||||
index += 1;
|
||||
@@ -536,6 +548,7 @@ describe("PiSessionService", () => {
|
||||
});
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime,
|
||||
sessionManager: {
|
||||
create: () => parentManager,
|
||||
@@ -614,6 +627,7 @@ describe("PiSessionService", () => {
|
||||
});
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime,
|
||||
sessionManager: {
|
||||
create: () => copiedParentManager,
|
||||
@@ -674,6 +688,7 @@ describe("PiSessionService", () => {
|
||||
const open = vi.fn((path: string) => path === parentFile ? parentManager : childManager);
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: () => {
|
||||
const runtime = runtimes[index] ?? parent.runtime;
|
||||
index += 1;
|
||||
@@ -728,6 +743,7 @@ describe("PiSessionService", () => {
|
||||
const open = vi.fn((path: string) => path === actualParentFile ? parent.session.sessionManager : childManager);
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: () => {
|
||||
const runtime = runtimes[index] ?? parent.runtime;
|
||||
index += 1;
|
||||
@@ -770,6 +786,7 @@ describe("PiSessionService", () => {
|
||||
const open = vi.fn(() => childManager);
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(child.runtime),
|
||||
sessionManager: {
|
||||
create: () => childManager,
|
||||
@@ -959,6 +976,7 @@ describe("PiSessionService", () => {
|
||||
const fake = fakeRuntime("nope");
|
||||
const service = new PiSessionService(new CapturingSessionEventHub(), {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { AuthStorage, ModelRegistry } from "@earendil-works/pi-coding-agent";
|
||||
import { ModelRuntime } from "@earendil-works/pi-coding-agent";
|
||||
import { InMemoryCredentialStore, type Credential, type CredentialStore } from "@earendil-works/pi-ai";
|
||||
import type { GlobalSessionEvent, SessionUiEvent } from "../../shared/apiTypes.js";
|
||||
import { SessionEventHub } from "../realtime/sessionEventHub.js";
|
||||
import type { PiAgentSession, PiSessionManager, PiSessionRuntime, PiSessionServiceDependencies } from "./piSessionService.js";
|
||||
@@ -62,8 +63,34 @@ export function sessionRef(id: string, cwd = "/workspace") {
|
||||
export const TEST_MODEL_PROVIDER = "anthropic";
|
||||
export const TEST_MODEL_ID = "claude-sonnet-4-5-20250929";
|
||||
|
||||
/**
|
||||
* Seed a credential into an {@link InMemoryCredentialStore}. `modify` is the
|
||||
* only write path on the pi-ai `CredentialStore` contract, so tests that need a
|
||||
* pre-populated store go through it rather than mutating internals.
|
||||
*/
|
||||
export async function seedCredential(store: InMemoryCredentialStore, providerId: string, credential: Credential): Promise<void> {
|
||||
await store.modify(providerId, () => Promise.resolve(credential));
|
||||
}
|
||||
|
||||
/**
|
||||
* Build a real {@link ModelRuntime} over an in-memory credential store — the
|
||||
* async test seam that replaces the removed `ModelRegistry.create(AuthStorage
|
||||
* .inMemory())`. Pass a pre-seeded store to exercise credential-dependent
|
||||
* behavior (e.g. auth-loss warnings).
|
||||
*/
|
||||
export function createTestModelRuntime(credentials: CredentialStore = new InMemoryCredentialStore()): Promise<ModelRuntime> {
|
||||
return ModelRuntime.create({ credentials, modelsPath: null, allowModelNetwork: false });
|
||||
}
|
||||
|
||||
/**
|
||||
* Shared runtime for the common case where a test only needs model catalog
|
||||
* reads and no configured auth. Built once so the many `fakeRuntime` sessions
|
||||
* and `PiSessionService` constructions can inject it synchronously.
|
||||
*/
|
||||
export const testModelRuntime = await createTestModelRuntime();
|
||||
|
||||
export function testModel(): NonNullable<PiAgentSession["model"]> {
|
||||
const model = ModelRegistry.inMemory(AuthStorage.inMemory()).find(TEST_MODEL_PROVIDER, TEST_MODEL_ID);
|
||||
const model = testModelRuntime.getModel(TEST_MODEL_PROVIDER, TEST_MODEL_ID);
|
||||
if (model === undefined) throw new Error("test model not found");
|
||||
return model;
|
||||
}
|
||||
@@ -88,7 +115,7 @@ export function fakeRuntime(sessionId = "session-1", patch: Partial<TestSession>
|
||||
pendingMessageCount: 0,
|
||||
sessionManager: fakeSessionManager(),
|
||||
settingsManager: { getWarnings: () => ({}), setWarnings: () => undefined },
|
||||
modelRegistry: ModelRegistry.create(AuthStorage.inMemory()),
|
||||
modelRuntime: testModelRuntime,
|
||||
scopedModels: [],
|
||||
extensionRunner: { getRegisteredCommands: () => [] },
|
||||
promptTemplates: [],
|
||||
|
||||
@@ -1,20 +1,21 @@
|
||||
import { statSync } from "node:fs";
|
||||
import { join } from "node:path";
|
||||
import { open, readFile, writeFile } from "node:fs/promises";
|
||||
import type { ImageContent } from "@earendil-works/pi-ai";
|
||||
import type { StreamFn } from "@earendil-works/pi-agent-core";
|
||||
import {
|
||||
AuthStorage,
|
||||
createAgentSessionFromServices,
|
||||
createAgentSessionRuntime,
|
||||
createAgentSessionServices,
|
||||
createEditToolDefinition,
|
||||
defineTool,
|
||||
ModelRegistry,
|
||||
readStoredCredential,
|
||||
SessionManager,
|
||||
type AgentSessionRuntimeDiagnostic,
|
||||
type AgentSessionServices,
|
||||
type CreateAgentSessionRuntimeFactory,
|
||||
type EditToolDetails,
|
||||
type ModelRuntime,
|
||||
type ResourceDiagnostic,
|
||||
} from "@earendil-works/pi-coding-agent";
|
||||
import type { ClientArchiveSessionsResponse, ClientCommand, ClientCommandResult, ClientMessagePage, ClientSession, ClientSessionCleanupExecuteResponse, ClientSessionCleanupPreviewResponse, ClientSessionModel, ClientSessionStatus, ClientThinkingLevel, SessionStreamSnapshot, SessionUiEvent } from "../types.js";
|
||||
@@ -26,7 +27,6 @@ import { SessionCommandService } from "./sessionCommandService.js";
|
||||
import { SessionArchiveStore, type ArchivedSessionRecord, type ArchiveSessionInput } from "./sessionArchiveStore.js";
|
||||
import { findArchiveCandidateByIdOrPrefix, planSessionArchiveTree, type SessionArchiveTreeCandidate } from "./sessionArchiveTree.js";
|
||||
import type { ActiveSession } from "./sessionRuntimeStore.js";
|
||||
import { createModelRegistryForAgentDir, type AuthChange } from "./authService.js";
|
||||
import { deterministicSessionName, fallbackSessionName, generateShortSessionName } from "./sessionNameGenerator.js";
|
||||
import { computeEditPreview, type EditPreviewResult } from "./editPreview.js";
|
||||
import { attachmentsToInlineImages, saveAttachmentsToWorkspace } from "./attachmentService.js";
|
||||
@@ -34,6 +34,7 @@ import { parsePromptAttachments } from "../../shared/promptAttachments.js";
|
||||
import type { SavedPromptAttachment, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkFailure, SessionBulkMutationRef, SessionWarning } from "../../shared/apiTypes.js";
|
||||
import type { SessionRouteLookup, SessionRouteRef, SessionRouteService } from "./sessionService.js";
|
||||
|
||||
import { type AuthChange } from "./authService.js";
|
||||
import { canonicalizeStoredCwd, cwdPathsEqual } from "../workingDirectory.js";
|
||||
import type { WorkspaceActivityService } from "../activity/workspaceActivityService.js";
|
||||
import { createSpawnSessionToolDefinition, type SpawnSessionInvocation, type SpawnSessionResult } from "./spawnSessionTool.js";
|
||||
@@ -171,7 +172,6 @@ interface BulkDeletePlanItem {
|
||||
}
|
||||
|
||||
type AgentModel = NonNullable<SpawnSessionInvocation["model"]>;
|
||||
type ModelRegistryInstance = ReturnType<typeof ModelRegistry.create>;
|
||||
|
||||
export interface PiSessionManager {
|
||||
getCwd(): string;
|
||||
@@ -208,7 +208,7 @@ interface PiExtensionBindings {
|
||||
}
|
||||
|
||||
export interface PiAgentSession {
|
||||
modelRegistry: ModelRegistryInstance;
|
||||
modelRuntime: ModelRuntime;
|
||||
/**
|
||||
* Narrow read/write of the SDK `SettingsManager`, exposing only the warning
|
||||
* suppression flags consumed here (e.g. `anthropicExtraUsage`). Used to gate
|
||||
@@ -383,11 +383,12 @@ const ANTHROPIC_EXTRA_USAGE_DISMISS_ID = "anthropicExtraUsage";
|
||||
* synchronous live status computation.
|
||||
*/
|
||||
export function anthropicSubscriptionWarning(
|
||||
session: Pick<PiAgentSession, "model" | "modelRegistry" | "settingsManager">,
|
||||
session: Pick<PiAgentSession, "model" | "settingsManager">,
|
||||
authPath?: string,
|
||||
): SessionWarning | undefined {
|
||||
if (session.settingsManager.getWarnings().anthropicExtraUsage === false) return undefined;
|
||||
if (session.model?.provider !== "anthropic") return undefined;
|
||||
const credential = session.modelRegistry.authStorage.get("anthropic");
|
||||
const credential = readStoredCredential("anthropic", authPath);
|
||||
if (credential === undefined) return undefined;
|
||||
const isSubscriptionAuth = credential.type === "oauth"
|
||||
? true
|
||||
@@ -487,14 +488,13 @@ export function createPiWebCustomToolDefinitions(
|
||||
}
|
||||
|
||||
function createDefaultRuntimeFactory(
|
||||
authStorage: AuthStorage,
|
||||
modelRegistry: ModelRegistryInstance,
|
||||
modelRuntime: ModelRuntime,
|
||||
sessionManagers: Pick<PiSessionManagerGateway, "open">,
|
||||
spawn?: SpawnSessionFn,
|
||||
subsessions?: SubsessionToolDeps,
|
||||
): PiWebCreateAgentSessionRuntimeFactory {
|
||||
return async ({ cwd, agentDir, sessionManager, sessionStartEvent, initialModel, delegationToolsEnabled }) => {
|
||||
const services = await createAgentSessionServices({ cwd, agentDir, authStorage, modelRegistry });
|
||||
const services = await createAgentSessionServices({ cwd, agentDir, modelRuntime });
|
||||
const resolvedDelegationToolsEnabled = delegationToolsEnabled
|
||||
?? await sessionAllowsDelegationTools(sessionManager, sessionManagers);
|
||||
const customTools = createPiWebCustomToolDefinitions(cwd, resolvedDelegationToolsEnabled, spawn, subsessions);
|
||||
@@ -539,7 +539,7 @@ export interface PiSessionServiceDependencies {
|
||||
archiveStore?: SessionArchiveRepository;
|
||||
createRuntime?: PiWebCreateAgentSessionRuntimeFactory;
|
||||
createAgentRuntime?: CreateAgentRuntime;
|
||||
modelRegistry?: ModelRegistryInstance;
|
||||
modelRuntime: ModelRuntime;
|
||||
heartbeatIntervalMs?: number;
|
||||
workspaceActivity?: Pick<WorkspaceActivityService, "applySessionStatus" | "applySessionActivity" | "removeSession" | "reconcileSessionActivity">;
|
||||
/**
|
||||
@@ -589,7 +589,7 @@ export class PiSessionService implements SessionRouteService {
|
||||
private readonly sessionManager: PiSessionManagerGateway;
|
||||
private readonly createRuntime: PiWebCreateAgentSessionRuntimeFactory;
|
||||
private readonly createAgentRuntime: CreateAgentRuntime;
|
||||
private readonly modelRegistry: ModelRegistryInstance;
|
||||
private readonly modelRuntime: ModelRuntime;
|
||||
private readonly workspaceActivity: Pick<WorkspaceActivityService, "applySessionStatus" | "applySessionActivity" | "removeSession" | "reconcileSessionActivity"> | undefined;
|
||||
private readonly spawnTargets: SpawnTargetResolver | undefined;
|
||||
private readonly logger: PiSessionLogger;
|
||||
@@ -599,7 +599,7 @@ export class PiSessionService implements SessionRouteService {
|
||||
this.archiveStore = deps.archiveStore ?? new SessionArchiveStore();
|
||||
this.agentDir = deps.agentDir;
|
||||
this.sessionManager = deps.sessionManager;
|
||||
this.modelRegistry = deps.modelRegistry ?? createModelRegistryForAgentDir(this.agentDir);
|
||||
this.modelRuntime = deps.modelRuntime;
|
||||
this.spawnTargets = deps.spawnTargets;
|
||||
this.logger = deps.logger ?? noopLogger;
|
||||
this.now = deps.now ?? (() => new Date());
|
||||
@@ -607,8 +607,7 @@ export class PiSessionService implements SessionRouteService {
|
||||
// also require the spawn capability (they share its project-scope resolver).
|
||||
const subsessionsActive = this.spawnTargets !== undefined && deps.subsessionsEnabled === true;
|
||||
this.createRuntime = deps.createRuntime ?? createDefaultRuntimeFactory(
|
||||
this.modelRegistry.authStorage,
|
||||
this.modelRegistry,
|
||||
this.modelRuntime,
|
||||
this.sessionManager,
|
||||
this.spawnTargets === undefined ? undefined : (input) => this.spawnSession(input),
|
||||
!subsessionsActive ? undefined : {
|
||||
@@ -1159,22 +1158,22 @@ export class PiSessionService implements SessionRouteService {
|
||||
|
||||
async availableModels(ref: PiSessionLookup): Promise<ClientSessionModel[]> {
|
||||
const session = await this.getOrOpen(ref);
|
||||
session.modelRegistry.refresh();
|
||||
await session.modelRuntime.reloadConfig();
|
||||
const models = session.scopedModels.length > 0
|
||||
? session.scopedModels.map((scoped) => scoped.model)
|
||||
: session.modelRegistry.getAvailable();
|
||||
: session.modelRuntime.getAvailableSnapshot();
|
||||
return models.map(modelToClientModel);
|
||||
}
|
||||
|
||||
async setModel(ref: PiSessionLookup, provider: string, modelId: string): Promise<ClientSessionStatus> {
|
||||
await this.assertWritable(ref);
|
||||
const session = await this.getOrOpen(ref);
|
||||
session.modelRegistry.refresh();
|
||||
await session.modelRuntime.reloadConfig();
|
||||
const candidates = session.scopedModels.length > 0
|
||||
? session.scopedModels.map((scoped) => scoped.model)
|
||||
: session.modelRegistry.getAvailable();
|
||||
: session.modelRuntime.getAvailableSnapshot();
|
||||
const model = candidates.find((candidate) => candidate.provider === provider && candidate.id === modelId)
|
||||
?? session.modelRegistry.find(provider, modelId);
|
||||
?? session.modelRuntime.getModel(provider, modelId);
|
||||
if (model === undefined) throw new Error(`Model not found: ${provider}/${modelId}`);
|
||||
await session.setModel(model);
|
||||
this.publishActivity(session, `model: ${model.id}`, "idle", model.provider);
|
||||
@@ -2020,10 +2019,11 @@ export class PiSessionService implements SessionRouteService {
|
||||
}
|
||||
|
||||
applyAuthChange(change: AuthChange = {}): void {
|
||||
this.modelRegistry.refresh();
|
||||
// ModelRuntime.login()/logout() refresh the shared runtime before AuthService
|
||||
// emits the change, so no refresh is needed here. Keeping this synchronous
|
||||
// also lets every active session observe the same committed auth snapshot.
|
||||
for (const active of this.active.values()) {
|
||||
const { session } = active.runtime;
|
||||
session.modelRegistry.refresh();
|
||||
this.syncCurrentModelAuthWarning(session, change.removedProviderId);
|
||||
this.publishStatus(session);
|
||||
}
|
||||
@@ -2034,9 +2034,9 @@ export class PiSessionService implements SessionRouteService {
|
||||
if (model === undefined) return;
|
||||
if (model.provider === "unknown" && model.id === "unknown") return;
|
||||
const warningKey = authLossWarningKey(session.sessionId, model.provider, model.id);
|
||||
const registered = session.modelRegistry.find(model.provider, model.id);
|
||||
const registered = session.modelRuntime.getModel(model.provider, model.id);
|
||||
if (registered === undefined) return;
|
||||
if (session.modelRegistry.hasConfiguredAuth(registered)) {
|
||||
if (session.modelRuntime.hasConfiguredAuth(model.provider)) {
|
||||
this.authLossWarnings.delete(warningKey);
|
||||
return;
|
||||
}
|
||||
@@ -2182,7 +2182,7 @@ export class PiSessionService implements SessionRouteService {
|
||||
private warningsForSession(session: PiAgentSession): SessionWarning[] {
|
||||
const runtime = this.active.get(session.sessionId)?.runtime;
|
||||
const warnings = runtime === undefined ? [] : collectRuntimeWarnings(runtime);
|
||||
const anthropic = anthropicSubscriptionWarning(session);
|
||||
const anthropic = anthropicSubscriptionWarning(session, join(this.agentDir, "auth.json"));
|
||||
if (anthropic !== undefined) warnings.push(anthropic);
|
||||
return warnings;
|
||||
}
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { AuthStorage, ModelRegistry, type AgentSessionRuntimeDiagnostic, type ResourceDiagnostic } from "@earendil-works/pi-coding-agent";
|
||||
import { mkdtemp, rm, writeFile } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import { type AgentSessionRuntimeDiagnostic, type ResourceDiagnostic } from "@earendil-works/pi-coding-agent";
|
||||
import { anthropicSubscriptionWarning, collectRuntimeWarnings, dismissSessionWarning, type RuntimeWarningSources } from "./piSessionService.js";
|
||||
import { testModel } from "./piSessionService.testSupport.js";
|
||||
import type { PiAgentSession } from "./piSessionService.js";
|
||||
import type { SessionWarning } from "../../shared/apiTypes.js";
|
||||
|
||||
@@ -83,47 +87,55 @@ describe("collectRuntimeWarnings", () => {
|
||||
const ANTHROPIC_SUBSCRIPTION_AUTH_WARNING =
|
||||
"Anthropic subscription auth is active. Third-party harness usage draws from extra usage and is billed per token, not your Claude plan limits. Manage extra usage at https://claude.ai/settings/usage.";
|
||||
|
||||
type SubscriptionSession = Pick<PiAgentSession, "model" | "modelRegistry" | "settingsManager">;
|
||||
type SubscriptionSession = Pick<PiAgentSession, "model" | "settingsManager">;
|
||||
|
||||
function anthropicModel(provider: string): PiAgentSession["model"] {
|
||||
const registry = ModelRegistry.inMemory(AuthStorage.inMemory());
|
||||
const model = registry.getAll().find((candidate) => candidate.provider === provider) ?? registry.getAll()[0];
|
||||
if (model === undefined) throw new Error("expected at least one built-in model");
|
||||
return { ...model, provider };
|
||||
// anthropicSubscriptionWarning only reads `model.provider`, so any built-in
|
||||
// model re-tagged with the desired provider is a sufficient fixture.
|
||||
return { ...testModel(), provider };
|
||||
}
|
||||
|
||||
function subscriptionSession(options: {
|
||||
provider?: string;
|
||||
anthropicExtraUsage?: boolean;
|
||||
credential?: AuthStorage;
|
||||
}): SubscriptionSession {
|
||||
const authStorage = options.credential ?? AuthStorage.inMemory();
|
||||
return {
|
||||
model: options.provider === undefined ? undefined : anthropicModel(options.provider),
|
||||
settingsManager: {
|
||||
getWarnings: () => (options.anthropicExtraUsage === undefined ? {} : { anthropicExtraUsage: options.anthropicExtraUsage }),
|
||||
setWarnings: () => undefined,
|
||||
},
|
||||
modelRegistry: ModelRegistry.create(authStorage),
|
||||
};
|
||||
}
|
||||
|
||||
function anthropicAuth(credential: { type: "oauth" } | { type: "api_key"; key: string }): AuthStorage {
|
||||
const authStorage = AuthStorage.inMemory();
|
||||
if (credential.type === "oauth") {
|
||||
authStorage.set("anthropic", { type: "oauth", access: "a", refresh: "r", expires: Date.now() + 3_600_000 });
|
||||
} else {
|
||||
authStorage.set("anthropic", { type: "api_key", key: credential.key });
|
||||
}
|
||||
return authStorage;
|
||||
const tempDirs: string[] = [];
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
|
||||
});
|
||||
|
||||
/**
|
||||
* Write an `auth.json` holding a single anthropic credential and return its
|
||||
* path. `anthropicSubscriptionWarning` reads it via `readStoredCredential`, so
|
||||
* the credential seam is the on-disk auth file rather than an in-memory store.
|
||||
*/
|
||||
async function anthropicAuthPath(credential: { type: "oauth" } | { type: "api_key"; key: string }): Promise<string> {
|
||||
const dir = await mkdtemp(join(tmpdir(), "pi-web-warnings-"));
|
||||
tempDirs.push(dir);
|
||||
const authPath = join(dir, "auth.json");
|
||||
const stored = credential.type === "oauth"
|
||||
? { type: "oauth", access: "a", refresh: "r", expires: Date.now() + 3_600_000 }
|
||||
: { type: "api_key", key: credential.key };
|
||||
await writeFile(authPath, JSON.stringify({ anthropic: stored }));
|
||||
return authPath;
|
||||
}
|
||||
|
||||
describe("anthropicSubscriptionWarning", () => {
|
||||
it("warns with the verbatim SDK wording for a stored oauth credential", () => {
|
||||
expect(anthropicSubscriptionWarning(subscriptionSession({
|
||||
provider: "anthropic",
|
||||
credential: anthropicAuth({ type: "oauth" }),
|
||||
}))).toEqual({
|
||||
it("warns with the verbatim SDK wording for a stored oauth credential", async () => {
|
||||
expect(anthropicSubscriptionWarning(
|
||||
subscriptionSession({ provider: "anthropic" }),
|
||||
await anthropicAuthPath({ type: "oauth" }),
|
||||
)).toEqual({
|
||||
severity: "warning",
|
||||
message: ANTHROPIC_SUBSCRIPTION_AUTH_WARNING,
|
||||
source: "anthropic",
|
||||
@@ -131,37 +143,41 @@ describe("anthropicSubscriptionWarning", () => {
|
||||
} satisfies SessionWarning);
|
||||
});
|
||||
|
||||
it("warns for an sk-ant-oat subscription API key", () => {
|
||||
expect(anthropicSubscriptionWarning(subscriptionSession({
|
||||
provider: "anthropic",
|
||||
credential: anthropicAuth({ type: "api_key", key: "sk-ant-oat-abc123" }),
|
||||
}))?.message).toBe(ANTHROPIC_SUBSCRIPTION_AUTH_WARNING);
|
||||
it("warns for an sk-ant-oat subscription API key", async () => {
|
||||
expect(anthropicSubscriptionWarning(
|
||||
subscriptionSession({ provider: "anthropic" }),
|
||||
await anthropicAuthPath({ type: "api_key", key: "sk-ant-oat-abc123" }),
|
||||
)?.message).toBe(ANTHROPIC_SUBSCRIPTION_AUTH_WARNING);
|
||||
});
|
||||
|
||||
it("does not warn for a standard anthropic API key", () => {
|
||||
expect(anthropicSubscriptionWarning(subscriptionSession({
|
||||
provider: "anthropic",
|
||||
credential: anthropicAuth({ type: "api_key", key: "sk-ant-api-abc123" }),
|
||||
}))).toBeUndefined();
|
||||
it("does not warn for a standard anthropic API key", async () => {
|
||||
expect(anthropicSubscriptionWarning(
|
||||
subscriptionSession({ provider: "anthropic" }),
|
||||
await anthropicAuthPath({ type: "api_key", key: "sk-ant-api-abc123" }),
|
||||
)).toBeUndefined();
|
||||
});
|
||||
|
||||
it("respects the anthropicExtraUsage suppression gate", () => {
|
||||
expect(anthropicSubscriptionWarning(subscriptionSession({
|
||||
provider: "anthropic",
|
||||
anthropicExtraUsage: false,
|
||||
credential: anthropicAuth({ type: "oauth" }),
|
||||
}))).toBeUndefined();
|
||||
it("respects the anthropicExtraUsage suppression gate", async () => {
|
||||
expect(anthropicSubscriptionWarning(
|
||||
subscriptionSession({ provider: "anthropic", anthropicExtraUsage: false }),
|
||||
await anthropicAuthPath({ type: "oauth" }),
|
||||
)).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not warn when the active provider is not anthropic", () => {
|
||||
expect(anthropicSubscriptionWarning(subscriptionSession({
|
||||
provider: "openai",
|
||||
credential: anthropicAuth({ type: "oauth" }),
|
||||
}))).toBeUndefined();
|
||||
it("does not warn when the active provider is not anthropic", async () => {
|
||||
expect(anthropicSubscriptionWarning(
|
||||
subscriptionSession({ provider: "openai" }),
|
||||
await anthropicAuthPath({ type: "oauth" }),
|
||||
)).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not warn when no anthropic credential is stored", () => {
|
||||
expect(anthropicSubscriptionWarning(subscriptionSession({ provider: "anthropic" }))).toBeUndefined();
|
||||
it("does not warn when no anthropic credential is stored", async () => {
|
||||
const dir = await mkdtemp(join(tmpdir(), "pi-web-warnings-"));
|
||||
tempDirs.push(dir);
|
||||
expect(anthropicSubscriptionWarning(
|
||||
subscriptionSession({ provider: "anthropic" }),
|
||||
join(dir, "auth.json"),
|
||||
)).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
import type { MessagePage, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkMutationRef, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse, SessionStatus, SessionStreamSnapshot } from "../../shared/apiTypes.js";
|
||||
import { SessionEventHub } from "../realtime/sessionEventHub.js";
|
||||
import { PiSessionService, type PiSessionManagerGateway } from "./piSessionService.js";
|
||||
import { testModelRuntime } from "./piSessionService.testSupport.js";
|
||||
import type { SessionRouteLookup, SessionRouteService } from "./sessionService.js";
|
||||
import { registerSessionRoutes } from "./sessionRoutes.js";
|
||||
import type { NormalizedSessionCleanupRequest } from "./sessionCleanup.js";
|
||||
@@ -20,7 +21,7 @@ beforeEach(async () => {
|
||||
await app.register(fastifyWebsocket);
|
||||
sessionManager = new RejectingSessionManager();
|
||||
const eventHub = new SessionEventHub();
|
||||
service = new PiSessionService(eventHub, { agentDir: TEST_AGENT_DIR, sessionManager, heartbeatIntervalMs: 60_000 });
|
||||
service = new PiSessionService(eventHub, { agentDir: TEST_AGENT_DIR, modelRuntime: testModelRuntime, sessionManager, heartbeatIntervalMs: 60_000 });
|
||||
registerSessionRoutes(app, service, eventHub);
|
||||
});
|
||||
|
||||
|
||||
Reference in New Issue
Block a user