diff --git a/.changeset/cancel-stale-auth-flows.md b/.changeset/cancel-stale-auth-flows.md index 02aad37..0b4cfc5 100644 --- a/.changeset/cancel-stale-auth-flows.md +++ b/.changeset/cancel-stale-auth-flows.md @@ -2,4 +2,4 @@ "@jmfederico/pi-web": patch --- -Cancel auth flows created after their browser start operation becomes stale, preventing abandoned provider polling or callback listeners after the dialog closes. +Keep auth interactions bound to their originating machine and cancel flows created after their browser start operation becomes stale, preventing secrets from reaching the wrong remote or abandoned provider resources from surviving a closed dialog. diff --git a/src/client/src/appState.ts b/src/client/src/appState.ts index 037ba74..9b1eb68 100644 --- a/src/client/src/appState.ts +++ b/src/client/src/appState.ts @@ -70,7 +70,7 @@ export type AuthDialogState = | { step: "method" } | { step: "providers"; mode: "login"; authType?: "oauth" | "api_key"; providers: AuthProviderOption[] } | { step: "apiKey"; provider: AuthProviderOption; value: string; saving?: boolean; error?: string } - | { step: "oauth"; flow: OAuthFlowState; responding?: boolean; inputValue?: string; error?: string } + | { step: "oauth"; flow: OAuthFlowState; machineId: string; responding?: boolean; inputValue?: string; error?: string } | { step: "logout"; providers: AuthProviderOption[] }; export type WorkspaceScopedStateReset = Pick { it("keeps OAuth prompt input and submit state across poll refreshes for the same request", async () => { const flow = oauthFlow({ prompt: { requestId: "request-1", message: "Paste callback", kind: "manual" } }); const { controller, getState } = createController( - { authDialog: { step: "oauth", flow, inputValue: "https://callback", responding: true } }, + { authDialog: { step: "oauth", flow, machineId: "local", inputValue: "https://callback", responding: true } }, { respondOAuthFlow: () => Promise.resolve(oauthFlow({ prompt: { requestId: "request-1", message: "Paste callback", kind: "manual" }, progress: ["Still waiting"] })) }, ); @@ -77,7 +77,7 @@ describe("AuthController", () => { }); const respondCalls: string[] = []; const { controller } = createController( - { authDialog: { step: "oauth", flow, inputValue: "" } }, + { authDialog: { step: "oauth", flow, machineId: "local", inputValue: "" } }, { respondOAuthFlow: (_flowId, _requestId, value) => { respondCalls.push(value); @@ -94,7 +94,7 @@ describe("AuthController", () => { it("resets OAuth prompt input and submit state when the request id changes", async () => { const flow = oauthFlow({ prompt: { requestId: "request-1", message: "Paste callback", kind: "manual" } }); const { controller, getState } = createController( - { authDialog: { step: "oauth", flow, inputValue: "https://callback", responding: true } }, + { authDialog: { step: "oauth", flow, machineId: "local", inputValue: "https://callback", responding: true } }, { respondOAuthFlow: () => Promise.resolve(oauthFlow({ select: { requestId: "request-2", message: "Choose an account", options: [{ value: "acct-1", label: "Account 1" }] }, @@ -121,7 +121,7 @@ describe("AuthController", () => { const statusCalls: { session: Parameters[0]; machineId: string | undefined }[] = []; const appliedStatuses: SessionStatus[] = []; const { controller, getState } = createController( - { selectedSession: session, authDialog: { step: "oauth", flow, inputValue: "https://callback" } }, + { selectedSession: session, authDialog: { step: "oauth", flow, machineId: "local", inputValue: "https://callback" } }, { respondOAuthFlow: (flowId, requestId, value, machineId) => { respondCalls.push({ flowId, requestId, value, machineId }); @@ -147,7 +147,7 @@ describe("AuthController", () => { it("leaves the OAuth dialog ready to retry if responding fails", async () => { const flow = oauthFlow({ prompt: { requestId: "request-1", message: "Paste callback", kind: "manual" } }); const { controller, getState } = createController( - { authDialog: { step: "oauth", flow, inputValue: "https://callback", responding: true } }, + { authDialog: { step: "oauth", flow, machineId: "local", inputValue: "https://callback", responding: true } }, { respondOAuthFlow: () => Promise.reject(new Error("Invalid callback")) }, ); @@ -168,7 +168,7 @@ describe("AuthController", () => { const response = deferred(); const cancellation = deferred(); const { controller, getState } = createController( - { authDialog: { step: "oauth", flow, inputValue: "https://callback" } }, + { authDialog: { step: "oauth", flow, machineId: "local", inputValue: "https://callback" } }, { respondOAuthFlow: () => response.promise, cancelOAuthFlow: () => cancellation.promise, @@ -219,6 +219,78 @@ describe("AuthController", () => { expect(getState().error).toBe(""); }); + it("keeps prompt responses and cancellation bound to the flow's originating machine", async () => { + const prompt = { requestId: "request-1", message: "Enter secret", kind: "prompt", promptType: "secret" } as const; + const flow = oauthFlow({ providerId: "amazon-bedrock", prompt }); + const respondCalls: { value: string; machineId: string | undefined }[] = []; + const cancelCalls: { flowId: string; machineId: string | undefined }[] = []; + const { controller } = createController( + { + selectedMachine: remoteMachine("remote-2"), + authDialog: { step: "oauth", flow, machineId: "remote-1", inputValue: "secret-value" }, + }, + { + respondOAuthFlow: (_flowId, _requestId, value, machineId) => { + respondCalls.push({ value, machineId }); + return Promise.resolve(flow); + }, + cancelOAuthFlow: (flowId, machineId) => { + cancelCalls.push({ flowId, machineId }); + return Promise.resolve(oauthFlow({ status: "cancelled" })); + }, + }, + ); + + await controller.respondOAuth(); + await controller.cancelOAuth(); + + expect(respondCalls).toEqual([{ value: "secret-value", machineId: "remote-1" }]); + expect(cancelCalls).toEqual([{ flowId: "flow-1", machineId: "remote-1" }]); + }); + + it("keeps polling bound to the machine where an interactive flow started", async () => { + let pollCallback: (() => void) | undefined; + vi.stubGlobal("window", { + setInterval: (callback: () => void) => { + pollCallback = callback; + return 1; + }, + clearInterval: () => undefined, + }); + const prompt = { requestId: "request-1", message: "Enter secret", kind: "prompt", promptType: "secret" } as const; + const flow = oauthFlow({ providerId: "amazon-bedrock", prompt }); + const provider: AuthProviderOption = { ...authProvider("amazon-bedrock", "api_key"), loginFlow: "interactive" }; + const pollMachines: (string | undefined)[] = []; + const { controller, getState, setState } = createController( + { + selectedMachine: remoteMachine("remote-1"), + authDialog: { step: "providers", mode: "login", authType: "api_key", providers: [provider] }, + }, + { + startInteractiveApiKeyLogin: () => Promise.resolve(flow), + oauthFlow: (_flowId, machineId) => { + pollMachines.push(machineId); + return Promise.resolve(flow); + }, + }, + ); + + try { + await controller.selectLoginProvider(provider.id, provider.authType); + expect(getState().authDialog).toMatchObject({ step: "oauth", machineId: "remote-1" }); + + setState({ selectedMachine: remoteMachine("remote-2") }); + if (pollCallback === undefined) throw new Error("Expected auth polling to start"); + pollCallback(); + await flushMicrotasks(); + + expect(pollMachines).toEqual(["remote-1"]); + } finally { + controller.dispose(); + vi.unstubAllGlobals(); + } + }); + it("does not let a stale OAuth response overwrite a newer flow", async () => { vi.stubGlobal("window", { setInterval: () => 1, clearInterval: () => undefined }); const oldPrompt = { requestId: "request-1", message: "Paste callback", kind: "manual" } as const; @@ -227,7 +299,7 @@ describe("AuthController", () => { const response = deferred(); const providers = [authProvider("anthropic", "oauth")]; const { controller, getState } = createController( - { authDialog: { step: "oauth", flow: oldFlow, inputValue: "https://old-callback" } }, + { authDialog: { step: "oauth", flow: oldFlow, machineId: "local", inputValue: "https://old-callback" } }, { respondOAuthFlow: () => response.promise, authProviders: () => Promise.resolve({ providers }), @@ -296,7 +368,7 @@ describe("AuthController", () => { const flow = oauthFlow({ prompt: { requestId: "request-1", message: "Paste callback", kind: "manual" } }); const cancelCalls: { flowId: string; machineId: string | undefined }[] = []; const { controller, getState } = createController( - { authDialog: { step: "oauth", flow } }, + { authDialog: { step: "oauth", flow, machineId: "local" } }, { cancelOAuthFlow: (flowId, machineId) => { cancelCalls.push({ flowId, machineId }); @@ -390,13 +462,14 @@ function createController( ) { let state: AppState = { ...initialAppState(), ...statePatch }; const api = { ...defaultApi, ...apiPatch }; + const setState = (patch: Partial) => { state = { ...state, ...patch }; }; const controller = new AuthController( () => state, - (patch) => { state = { ...state, ...patch }; }, + setState, applyStatus, { api }, ); - return { controller, getState: () => state }; + return { controller, getState: () => state, setState }; } async function flushMicrotasks(): Promise { diff --git a/src/client/src/controllers/authController.ts b/src/client/src/controllers/authController.ts index 105f202..b96839d 100644 --- a/src/client/src/controllers/authController.ts +++ b/src/client/src/controllers/authController.ts @@ -142,10 +142,10 @@ export class AuthController { delete clean.error; this.setState({ authDialog: { ...clean, responding: true } }); try { - const flow = await this.api.respondOAuthFlow(flowId, requestId, responseValue, selectedMachineId(this.getState())); + const flow = await this.api.respondOAuthFlow(flowId, requestId, responseValue, dialog.machineId); const current = this.currentOAuthDialog(operationGeneration, flowId); if (flow.flowId !== flowId || current === undefined || oauthRequestId(current.flow) !== requestId) return; - this.updateOAuthFlow(flow); + this.updateOAuthFlow(flow, current.machineId); } catch (error) { const current = this.currentOAuthDialog(operationGeneration, flowId); if (current === undefined || oauthRequestId(current.flow) !== requestId) return; @@ -160,7 +160,7 @@ export class AuthController { return; } const flowId = dialog.flow.flowId; - const machineId = selectedMachineId(this.getState()); + const machineId = dialog.machineId; this.closeDialog(); try { await this.api.cancelOAuthFlow(flowId, machineId); @@ -217,8 +217,8 @@ export class AuthController { } return; } - this.updateOAuthFlow(flow); - if (flow.status === "running") this.startPolling(flow.flowId); + this.updateOAuthFlow(flow, machineId); + if (flow.status === "running") this.startPolling(flow.flowId, machineId); } catch (error) { if (operationGeneration === this.oauthOperationGeneration) this.setState({ error: String(error) }); } @@ -232,11 +232,11 @@ export class AuthController { return true; } - private updateOAuthFlow(flow: OAuthFlowState): void { + private updateOAuthFlow(flow: OAuthFlowState, machineId: string): void { if (flow.status === "complete") { this.stopPolling(); this.closeDialog(); - void this.refreshStatus(); + void this.refreshStatus(machineId); return; } if (flow.status === "error" || flow.status === "cancelled") { @@ -250,14 +250,14 @@ export class AuthController { const sameRequest = previousRequestId !== undefined && previousRequestId === newRequestId; const inputValue = sameRequest ? previousInput : ""; const responding = sameRequest && existing?.step === "oauth" ? existing.responding === true : false; - this.setState({ authDialog: { step: "oauth", flow, inputValue, responding } }); + this.setState({ authDialog: { step: "oauth", flow, machineId, inputValue, responding } }); } - private startPolling(flowId: string): void { + private startPolling(flowId: string, machineId: string): void { this.stopPolling(); const operationGeneration = this.oauthOperationGeneration; const pollGeneration = this.pollGeneration; - this.pollTimer = window.setInterval(() => { void this.poll(flowId, operationGeneration, pollGeneration); }, this.pollIntervalMs); + this.pollTimer = window.setInterval(() => { void this.poll(flowId, machineId, operationGeneration, pollGeneration); }, this.pollIntervalMs); } private stopPolling(): void { @@ -267,22 +267,22 @@ export class AuthController { this.pollTimer = undefined; } - private async poll(flowId: string, operationGeneration: number, pollGeneration: number): Promise { + private async poll(flowId: string, machineId: string, operationGeneration: number, pollGeneration: number): Promise { if (pollGeneration !== this.pollGeneration) return; const dialog = this.currentOAuthDialog(operationGeneration, flowId); - if (dialog === undefined) { + if (dialog?.machineId !== machineId) { this.stopPolling(); return; } const requestId = oauthRequestId(dialog.flow); try { - const flow = await this.api.oauthFlow(flowId, selectedMachineId(this.getState())); + const flow = await this.api.oauthFlow(flowId, machineId); const current = this.currentOAuthDialog(operationGeneration, flowId); - if (flow.flowId !== flowId || pollGeneration !== this.pollGeneration || current === undefined || oauthRequestId(current.flow) !== requestId) return; - this.updateOAuthFlow(flow); + if (flow.flowId !== flowId || pollGeneration !== this.pollGeneration || current?.machineId !== machineId || oauthRequestId(current.flow) !== requestId) return; + this.updateOAuthFlow(flow, machineId); } catch (error) { const current = this.currentOAuthDialog(operationGeneration, flowId); - if (pollGeneration !== this.pollGeneration || current === undefined || oauthRequestId(current.flow) !== requestId) return; + if (pollGeneration !== this.pollGeneration || current?.machineId !== machineId || oauthRequestId(current.flow) !== requestId) return; this.stopPolling(); this.setState({ authDialog: { ...current, error: String(error) } }); } @@ -294,11 +294,11 @@ export class AuthController { return dialog?.step === "oauth" && dialog.flow.flowId === flowId ? dialog : undefined; } - private async refreshStatus(): Promise { + private async refreshStatus(machineId = selectedMachineId(this.getState())): Promise { const session = this.session(); if (session === undefined) return; try { - this.applyStatus(await this.api.status(session, selectedMachineId(this.getState()))); + this.applyStatus(await this.api.status(session, machineId)); } catch { // Status refresh is opportunistic after login completes. }