diff --git a/src/server/sessions/oauthLoginFlowService.test.ts b/src/server/sessions/oauthLoginFlowService.test.ts index 24e3ab8..6bc75bb 100644 --- a/src/server/sessions/oauthLoginFlowService.test.ts +++ b/src/server/sessions/oauthLoginFlowService.test.ts @@ -1,9 +1,9 @@ -import type { OAuthLoginCallbacks } from "@earendil-works/pi-ai"; -import type { AuthStorage } from "@earendil-works/pi-coding-agent"; +import type { AuthInteraction } 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; +type LoginHandler = (providerId: string, interaction: AuthInteraction) => Promise; afterEach(() => { vi.useRealTimers(); @@ -17,18 +17,18 @@ 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(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" }); const afterRespond = service.respond(state.flowId, prompt.requestId, "abc123"); @@ -41,14 +41,30 @@ describe("OAuthLoginFlowService", () => { 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" }); + await interaction.prompt({ type: "text", message: "Waiting" }); + }), + }); + + expect(service.get(state.flowId)).toMatchObject({ auth: { url: "https://example.test/device", instructions: "Enter code: WXYZ-1234" } }); + 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" }], }); @@ -73,10 +89,8 @@ describe("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) => { + manualValue = await interaction.prompt({ type: "manual_code", message: "Paste the callback URL or authorization code" }); }), }); @@ -92,15 +106,44 @@ describe("OAuthLoginFlowService", () => { service.dispose(); }); + it("rejects a pending prompt when its own signal aborts without ending the flow", async () => { + const promptRejected = deferred(); + const service = new OAuthLoginFlowService(); + const controller = new AbortController(); + 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" }); + + 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(); 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 +165,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; @@ -145,8 +188,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" }); }), }); @@ -165,9 +208,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; @@ -187,8 +230,10 @@ describe("OAuthLoginFlowService", () => { }); }); -function fakeAuthStorage(login: LoginHandler): Pick { - return { login }; +function fakeRuntime(login: LoginHandler): Pick { + return { + login: (providerId, _type, interaction) => login(providerId, interaction).then(() => ({ type: "oauth", refresh: "r", access: "a", expires: 0 })), + }; } async function flushAsyncLogin(): Promise { diff --git a/src/server/sessions/oauthLoginFlowService.ts b/src/server/sessions/oauthLoginFlowService.ts index 02035b5..cf150bf 100644 --- a/src/server/sessions/oauthLoginFlowService.ts +++ b/src/server/sessions/oauthLoginFlowService.ts @@ -1,15 +1,16 @@ 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 } 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; +/** The single runtime capability this service drives — narrowed for testable DI. */ +type OAuthLoginRuntime = Pick; type TimerHandle = ReturnType; interface PendingOAuthRequest { requestId: string; allowEmpty: boolean; - resolve: (value: string | undefined) => void; + resolve: (value: string) => void; reject: (error: Error) => void; } @@ -46,7 +47,7 @@ export class OAuthLoginFlowService { start(options: { providerId: string; providerName: string; - authStorage: OAuthLoginStorage; + runtime: OAuthLoginRuntime; onComplete?: () => void; }): OAuthFlowState { const flowId = crypto.randomUUID(); @@ -66,28 +67,16 @@ 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) + void options.runtime.login(options.providerId, "oauth", interaction) .then(() => { if (!this.isCurrentRunning(record)) return; record.pending = undefined; @@ -147,14 +136,51 @@ export class OAuthLoginFlowService { this.flows.clear(); } - private waitForPrompt(record: OAuthFlowRecord, prompt: OAuthPrompt, kind: "prompt" | "manual"): Promise { + private handlePrompt(record: OAuthFlowRecord, prompt: AuthPrompt): Promise { + if (prompt.type === "select") { + return this.waitForSelect(record, prompt.message, prompt.options, prompt.signal); + } + // `manual_code` is the paste-back path for callback-server flows; text/secret + // are ordinary interactive entry. Both map to the single web-UI prompt shape. + const kind = prompt.type === "manual_code" ? "manual" : "prompt"; + return this.waitForPrompt(record, { + message: prompt.message, + ...(prompt.placeholder === undefined ? {} : { placeholder: prompt.placeholder }), + ...(prompt.signal === undefined ? {} : { signal: prompt.signal }), + }, kind); + } + + 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; + // 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. + case "device_code": + this.updateState(record, { ...record.state, auth: { url: event.verificationUri, instructions: `Enter code: ${event.userCode}` } }); + return; + case "info": + case "progress": + this.updateState(record, { ...record.state, progress: [...record.state.progress, event.message] }); + return; + } + } + + private waitForPrompt(record: OAuthFlowRecord, prompt: { message: string; placeholder?: string; signal?: AbortSignal }, kind: "prompt" | "manual"): Promise { return new Promise((resolve, reject) => { if (!this.isCurrentRunning(record)) { reject(new Error("Login cancelled")); return; } + if (prompt.signal?.aborted === true) { + reject(new Error("Prompt cancelled")); + return; + } const requestId = crypto.randomUUID(); - record.pending = { requestId, allowEmpty: prompt.allowEmpty === true, resolve: (value) => { resolve(value ?? ""); }, reject }; + record.pending = { requestId, allowEmpty: false, resolve, reject }; + this.bindPromptSignal(record, requestId, prompt.signal); const base = withoutInteraction(record.state); this.updateState(record, { ...base, @@ -163,26 +189,44 @@ export class OAuthLoginFlowService { message: prompt.message, kind, ...(prompt.placeholder === undefined ? {} : { placeholder: prompt.placeholder }), - ...(prompt.allowEmpty === true ? { allowEmpty: true } : {}), }, }); }); } - private waitForSelect(record: OAuthFlowRecord, prompt: OAuthSelectPrompt): Promise { + private waitForSelect(record: OAuthFlowRecord, message: string, promptOptions: readonly { id: string; label: string; description?: string }[], signal?: AbortSignal): Promise { return new Promise((resolve, reject) => { if (!this.isCurrentRunning(record)) { reject(new Error("Login cancelled")); return; } + if (signal?.aborted === true) { + reject(new Error("Prompt cancelled")); + return; + } const requestId = crypto.randomUUID(); - const options: CommandOption[] = prompt.options.map((option) => ({ value: option.id, label: option.label })); + const options: CommandOption[] = promptOptions.map((option) => ({ value: option.id, label: option.label })); record.pending = { requestId, allowEmpty: true, resolve, reject }; + this.bindPromptSignal(record, requestId, signal); const base = withoutInteraction(record.state); - this.updateState(record, { ...base, select: { requestId, message: prompt.message, options } }); + this.updateState(record, { ...base, select: { requestId, 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, requestId: string, signal?: AbortSignal): void { + if (signal === undefined) return; + signal.addEventListener("abort", () => { + const pending = record.pending; + if (pending?.requestId !== requestId) return; + record.pending = undefined; + if (this.isCurrentRunning(record)) this.updateState(record, withoutInteraction(record.state)); + pending.reject(new Error("Prompt cancelled")); + }, { once: true }); + } + private isCurrentRunning(record: OAuthFlowRecord): boolean { return this.flows.get(record.flowId) === record && record.state.status === "running"; }