Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
18 kB · 455 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { resetOAuthProviders } from "@earendil-works/pi-ai/oauth";import type { Component, OverlayHandle, TUI } from "@earendil-works/pi-tui";import stripAnsi from "strip-ansi";import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";import { AuthStorage } from "../src/core/auth-storage.js";import type { McpRemoveAccountResult } from "../src/core/mcp/connection-store.js";import { ModelRegistry } from "../src/core/model-registry.js";import { PRIME_INFERENCE_PROVIDER_ID } from "../src/core/prime-inference-auth.js";import type { AuthenticationResult } from "../src/modes/interactive/auth-flows.js";import { ProviderAuthFlows, type ProviderAuthFlowsHost } from "../src/modes/interactive/auth-flows.js";import { ExtensionSelectorComponent } from "../src/modes/interactive/components/extension-selector.js";import { LoginDialogComponent } from "../src/modes/interactive/components/login-dialog.js";import { OAuthSelectorComponent } from "../src/modes/interactive/components/oauth-selector.js";import { initTheme } from "../src/modes/interactive/theme/theme.js";
function createOverlayHandle(): OverlayHandle { return { hide: vi.fn(), setHidden: vi.fn(), isHidden: () => false, focus: vi.fn(), unfocus: vi.fn(), isFocused: () => true, };}
function createFakeTui(overlays: Component[] = []): TUI { return { terminal: { columns: 80, rows: 24 }, requestRender: vi.fn(), showOverlay: vi.fn((component: Component) => { overlays.push(component); return createOverlayHandle(); }), } as unknown as TUI;}
function createHost(authStorage: AuthStorage): { host: ProviderAuthFlowsHost; statusMessages: string[]; errorMessages: string[]; overlays: Component[]; panels: Component[];} { const statusMessages: string[] = []; const errorMessages: string[] = []; const overlays: Component[] = []; const panels: Component[] = []; const modelRegistry = { authStorage, refresh: vi.fn(), getAll: () => [], getProviderDisplayName: (providerId: string) => providerId, getProviderAuthStatus: (providerId: string) => authStorage.getAuthStatus(providerId), } as unknown as ModelRegistry;
return { host: { ui: createFakeTui(overlays), modelRegistry, showStatus: (message) => statusMessages.push(message), showError: (message) => errorMessages.push(message), showAuthPanel: (component) => { panels.push(component); return () => { const index = panels.lastIndexOf(component); if (index !== -1) { panels.splice(index, 1); } }; }, getAuthPanelRows: () => 24, getAvailableModels: async () => [], }, statusMessages, errorMessages, overlays, panels, };}
describe("Prime login completion guard", () => { it("does not complete the Prime login after the dialog was aborted", async () => { const onAuthChanged = vi.fn(); const host = { modelRegistry: { refresh: vi.fn(), authStorage: { setPrimeInferenceApiKey: vi.fn() } }, onAuthChanged, isOnboardingSurface: () => true, }; const flows = new ProviderAuthFlows(host as never); const complete = ( flows as unknown as { completePrimeInferenceLogin: ( apiKey: string, dialog: unknown, close: () => void, team?: unknown, ) => Promise<AuthenticationResult>; } ).completePrimeInferenceLogin; const controller = new AbortController(); controller.abort(); const dialog = { signal: controller.signal } as never; const select = vi .spyOn( flows as unknown as { selectPrimeInferenceTeam: () => Promise<string | undefined> }, "selectPrimeInferenceTeam", ) .mockResolvedValue(undefined);
const result = await complete.call(flows, "key", dialog, () => {}, undefined);
expect(result).toEqual({ status: "cancelled" }); expect(onAuthChanged).not.toHaveBeenCalled(); expect(host.modelRegistry.refresh).not.toHaveBeenCalled(); expect(select).toHaveBeenCalledOnce(); });});
describe("ProviderAuthFlows", () => { let tempDir: string; let authJsonPath: string; let primeConfigPath: string; let originalHome: string | undefined; let originalPrimeTeamId: string | undefined;
beforeAll(() => { initTheme("dark"); });
beforeEach(() => { resetOAuthProviders(); vi.stubEnv("PRIME_AGENT_INFERENCE_API_BASE_URL", ""); vi.stubEnv("PRIME_AGENT_INFERENCE_FRONTEND_URL", ""); tempDir = join(tmpdir(), `pi-auth-flows-${Date.now()}-${Math.random().toString(36).slice(2)}`); mkdirSync(tempDir, { recursive: true }); authJsonPath = join(tempDir, "auth.json"); primeConfigPath = join(tempDir, "prime-config.json"); writeFileSync(authJsonPath, "{}"); originalHome = process.env.HOME; originalPrimeTeamId = process.env.PRIME_TEAM_ID; });
afterEach(() => { if (originalHome === undefined) { delete process.env.HOME; } else { process.env.HOME = originalHome; } if (originalPrimeTeamId === undefined) { delete process.env.PRIME_TEAM_ID; } else { process.env.PRIME_TEAM_ID = originalPrimeTeamId; } if (existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); } resetOAuthProviders(); vi.restoreAllMocks(); vi.unstubAllEnvs(); });
it("does not import default CLI credentials when CLI reuse is disabled", async () => { process.env.HOME = tempDir; const defaultPrimeDir = join(tempDir, ".prime"); mkdirSync(defaultPrimeDir, { recursive: true }); writeFileSync(join(defaultPrimeDir, "config.json"), JSON.stringify({ api_key: "prime-cli-key" })); const authStorage = AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }); const fetchMock = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("offline")); const { host, overlays, panels } = createHost(authStorage); const result = new ProviderAuthFlows(host).runPrimeInferenceLogin(); await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledOnce()); expect(String(fetchMock.mock.calls[0]?.[0])).toBe("https://api.primeintellect.ai/api/v1/auth_challenge/generate"); panels[0]?.handleInput?.("\x1b"); await expect(result).resolves.toEqual({ status: "cancelled" }); expect(authStorage.has(PRIME_INFERENCE_PROVIDER_ID)).toBe(false); // The login dialog mounts inline in the host, never as an overlay. expect(overlays).toHaveLength(0); });
it("the generic /logout route delegates MCP logouts whole before touching auth", async () => { const authStorage = AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }); authStorage.set("mcp:acme-2", { type: "oauth", access: "real-for-acme-2", refresh: "r", expires: Date.now() + 3600_000, endpoint: "https://mcp.acme.test/mcp", }); const { host, panels, statusMessages, errorMessages } = createHost(authStorage); const logout = vi.spyOn(authStorage, "logout"); const delegated = vi.fn(async () => "removed" as McpRemoveAccountResult); (host as { onMcpAccountLogout?: unknown }).onMcpAccountLogout = delegated;
const logoutResult = new ProviderAuthFlows(host).runLogout(); expect(panels).toHaveLength(1); // #2331: the route selector mounts inline, not as an overlay for (const char of "acme-2") { panels[0]?.handleInput?.(char); } panels[0]?.handleInput?.("\r"); await expect(logoutResult).resolves.toBe("mcp:acme-2"); // The route never touched auth directly for the MCP id: the host-owned // critical section (store->auth) did everything. expect(logout).not.toHaveBeenCalled(); expect(delegated).toHaveBeenCalledWith("mcp:acme-2"); expect(errorMessages).toEqual([]); expect(statusMessages.join("\n")).toContain("Logged out of acme-2"); });
it("a failed MCP logout outcome reports the failure honestly instead of success", async () => { const authStorage = AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }); authStorage.set("mcp:acme-2", { type: "oauth", access: "real-for-acme-2", refresh: "r", expires: Date.now() + 3600_000, endpoint: "https://mcp.acme.test/mcp", }); const { host, panels, statusMessages, errorMessages } = createHost(authStorage); (host as { onMcpAccountLogout?: unknown }).onMcpAccountLogout = vi.fn( async () => "failed" as McpRemoveAccountResult, );
const logoutResult = new ProviderAuthFlows(host).runLogout(); expect(panels).toHaveLength(1); // #2331: the route selector mounts inline, not as an overlay for (const char of "acme-2") { panels[0]?.handleInput?.(char); } panels[0]?.handleInput?.("\r"); await expect(logoutResult).resolves.toBeNull(); expect(errorMessages.join("\n")).toContain("Logout failed"); expect(statusMessages.join("\n")).not.toContain("Logged out of acme-2"); });
// The honest partial-state wording ("could not be saved... try again") is pinned ONCE at the removeAction seam in // mcp-activation-queue.test.ts ("a removeAccount whose record write fails after the logout reports the honest // partial state").
it("the generic /login service option for an MCP account delegates to the guarded host hook", async () => { const authStorage = AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }); const { host } = createHost(authStorage); const delegated = vi.fn(async () => ({ status: "success" as const, providerId: "mcp:acme", providerName: "Acme", authType: "oauth" as const, kind: "service" as const, })); (host as { onMcpAccountLogin?: unknown }).onMcpAccountLogin = delegated;
const result = await new ProviderAuthFlows(host).loginProvider({ id: "mcp:acme", name: "Acme", authType: "oauth", category: "service", });
// The MCP login went through the guarded hook — never a raw dialog writing the final credential directly. expect(delegated).toHaveBeenCalledWith("mcp:acme"); expect(result.status).toBe("success"); });
it("an MCP login without a guarded host hook fails closed without opening a dialog", async () => { const { host, overlays, errorMessages } = createHost( AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }), ); const result = await new ProviderAuthFlows(host).loginProvider({ id: "mcp:my--service", name: "My service", authType: "oauth", category: "service", }); expect(result).toEqual({ status: "failed" }); expect(overlays).toEqual([]); expect(errorMessages.join("\n")).toContain("guarded host"); });
it("an unresolvable MCP login reports an explicit failure — never a raw-dialog fallback", async () => { const authStorage = AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }); const { host } = createHost(authStorage); const delegated = vi.fn(async () => ({ status: "failed" as const })); (host as { onMcpAccountLogin?: unknown }).onMcpAccountLogin = delegated;
const result = await new ProviderAuthFlows(host).loginProvider({ id: "mcp:acme", name: "Acme", authType: "oauth", category: "service", });
// The hook OWNS every MCP login (including unresolvable names): its explicit failed outcome is the route's result — // no raw dialog ever writes the final credential directly. expect(delegated).toHaveBeenCalledWith("mcp:acme"); expect(result).toEqual({ status: "failed" }); });
// A refused stale staged logout staying state-neutral is pinned at the STORE level in mcp-connection-store.test.ts // ("a staged-key logout queued behind a finalize-first move refuses fail-closed" and "a staged-key logout with a // bystander on the real key preserves the account shell") and at the resolution seam in // mcp-activation-queue.test.ts ("logoutMcpAccount resolves exact ids first...").
it("non-MCP logouts stay unchanged: the route removes them directly", async () => { const authStorage = AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }); authStorage.set("anthropic", { type: "api_key", key: "sk-ant-test" }); const { host, panels, statusMessages } = createHost(authStorage); const delegated = vi.fn(async () => "removed" as McpRemoveAccountResult); (host as { onMcpAccountLogout?: unknown }).onMcpAccountLogout = delegated; const logout = vi.spyOn(authStorage, "logout");
const logoutResult = new ProviderAuthFlows(host).runLogout(); expect(panels).toHaveLength(1); // #2331: the route selector mounts inline, not as an overlay for (const char of "anthropic") { panels[0]?.handleInput?.(char); } panels[0]?.handleInput?.("\r"); await expect(logoutResult).resolves.toBe("anthropic"); expect(delegated).not.toHaveBeenCalled(); expect(logout).toHaveBeenCalledWith("anthropic"); expect(statusMessages.join("\n")).toContain("Removed stored API key for anthropic"); });
it("does not offer logout for credentials owned only by the Prime CLI", async () => { writeFileSync(primeConfigPath, JSON.stringify({ api_key: "prime-cli-key" })); const authStorage = AuthStorage.create(authJsonPath, { primeCliConfigPath: primeConfigPath }); const { host, overlays, panels } = createHost(authStorage); await expect(new ProviderAuthFlows(host).runLogout()).resolves.toBeNull(); expect(overlays).toHaveLength(0); expect(panels).toHaveLength(0); expect(JSON.parse(readFileSync(primeConfigPath, "utf-8"))).toEqual({ api_key: "prime-cli-key" }); });
it("includes xAI subscription and API-key entries from the provider registry", () => { const { host } = createHost(AuthStorage.inMemory()); const flows = new ProviderAuthFlows({ ...host, modelRegistry: ModelRegistry.inMemory(host.modelRegistry.authStorage), });
expect(flows.getLoginProviderOptions().filter((provider) => provider.id === "xai")).toEqual([ { id: "xai", name: "xAI (Grok)", authType: "oauth" }, { id: "xai", name: "xAI (Grok)", authType: "api_key" }, ]); });
it("shows the connected account count on multi-account providers", () => { const authStorage = AuthStorage.inMemory({ "openai-codex": [ { type: "oauth", access: "access-token-1", refresh: "refresh-token-1", expires: Date.now() + 60_000, accountId: "account-1", }, { type: "oauth", access: "access-token-2", refresh: "refresh-token-2", expires: Date.now() + 60_000, accountId: "account-2", }, ], }); const selector = new OAuthSelectorComponent( "login", authStorage, [{ id: "openai-codex", name: "ChatGPT Codex", authType: "oauth" }], () => {}, () => {}, );
const output = stripAnsi(selector.render(120).join("\n"));
expect(output).toContain("subscription \u00b7 2 accounts connected"); expect(output).toContain("configured"); });
it("opens login on the requested MCP Connections category", async () => { const authStorage = AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }); const { host, overlays, panels } = createHost(authStorage);
const loginResult = new ProviderAuthFlows(host).runLogin({ initialCategory: "service" });
expect(overlays).toHaveLength(0); expect(panels).toHaveLength(1); expect(panels[0]).toBeInstanceOf(OAuthSelectorComponent); const output = stripAnsi(panels[0]?.render(80).join("\n") ?? ""); expect(output).toContain("Serper (web search)"); expect(output).not.toContain("Anthropic"); panels[0]?.handleInput?.("\x1b"); await expect(loginResult).resolves.toEqual({ status: "cancelled" }); expect(panels).toHaveLength(0); });
it("mounts the API key login dialog inline and unmounts it after submitting", async () => { const authStorage = AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }); const { host, overlays, panels, statusMessages } = createHost(authStorage);
const loginPromise = new ProviderAuthFlows(host).loginProvider({ id: "xai", name: "xAI (Grok)", authType: "api_key", });
await vi.waitFor(() => expect(stripAnsi(panels[0]?.render(80).join("\n") ?? "")).toContain("Enter API key:")); expect(panels[0]).toBeInstanceOf(LoginDialogComponent); expect(overlays).toHaveLength(0); panels[0]?.handleInput?.("test-key"); panels[0]?.handleInput?.("\r");
await expect(loginPromise).resolves.toMatchObject({ status: "success", providerId: "xai" }); expect(panels).toHaveLength(0); expect(authStorage.get("xai")).toEqual({ type: "api_key", key: "test-key" }); expect(statusMessages.join("\n")).toContain("Saved API key for xAI (Grok)"); });
it("restores the login dialog after an in-flow account selection", async () => { const authStorage = AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }); const { host, panels } = createHost(authStorage); const loginSpy = vi.spyOn(authStorage, "login").mockImplementation(async (_providerId, handlers) => { await handlers?.onSelect?.({ message: "Choose an account", options: [ { id: "personal", label: "Personal" }, { id: "business", label: "Business" }, ], }); });
const loginPromise = new ProviderAuthFlows(host).loginProvider({ id: "anthropic", name: "Anthropic", authType: "oauth", });
await vi.waitFor(() => expect(panels).toHaveLength(2)); expect(panels[0]).toBeInstanceOf(LoginDialogComponent); expect(panels[1]).toBeInstanceOf(ExtensionSelectorComponent); const selectorOutput = stripAnsi(panels[1]?.render(80).join("\n") ?? ""); expect(selectorOutput).toContain("Choose an account"); expect(selectorOutput).toContain("Personal"); expect(selectorOutput).toContain("Business");
panels[1]?.handleInput?.("\r");
await expect(loginPromise).resolves.toMatchObject({ status: "success", providerId: "anthropic" }); // Selecting an option unmounts the selector and restores the login dialog, // which is then unmounted when the flow finishes. expect(panels).toHaveLength(0); expect(loginSpy).toHaveBeenCalledOnce(); });});