fix(auth): bind flows to their originating machine

Retain client-owned machine affinity for each interactive auth flow and use it for prompt responses, polling, cancellation, and completion refreshes. This prevents a later machine selection from forwarding secrets to a different remote.\n\nRefs #74
This commit is contained in:
Federico Jaramillo Martinez
2026-07-18 21:09:51 +02:00
parent 67f673b227
commit 62396c3165
4 changed files with 103 additions and 30 deletions
+1 -1
View File
@@ -2,4 +2,4 @@
"@jmfederico/pi-web": patch "@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.
+1 -1
View File
@@ -70,7 +70,7 @@ export type AuthDialogState =
| { step: "method" } | { step: "method" }
| { step: "providers"; mode: "login"; authType?: "oauth" | "api_key"; providers: AuthProviderOption[] } | { step: "providers"; mode: "login"; authType?: "oauth" | "api_key"; providers: AuthProviderOption[] }
| { step: "apiKey"; provider: AuthProviderOption; value: string; saving?: boolean; error?: string } | { 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[] }; | { step: "logout"; providers: AuthProviderOption[] };
export type WorkspaceScopedStateReset = Pick<AppState, export type WorkspaceScopedStateReset = Pick<AppState,
@@ -62,7 +62,7 @@ describe("AuthController", () => {
it("keeps OAuth prompt input and submit state across poll refreshes for the same request", async () => { 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 flow = oauthFlow({ prompt: { requestId: "request-1", message: "Paste callback", kind: "manual" } });
const { controller, getState } = createController( 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"] })) }, { 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 respondCalls: string[] = [];
const { controller } = createController( const { controller } = createController(
{ authDialog: { step: "oauth", flow, inputValue: "" } }, { authDialog: { step: "oauth", flow, machineId: "local", inputValue: "" } },
{ {
respondOAuthFlow: (_flowId, _requestId, value) => { respondOAuthFlow: (_flowId, _requestId, value) => {
respondCalls.push(value); respondCalls.push(value);
@@ -94,7 +94,7 @@ describe("AuthController", () => {
it("resets OAuth prompt input and submit state when the request id changes", async () => { 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 flow = oauthFlow({ prompt: { requestId: "request-1", message: "Paste callback", kind: "manual" } });
const { controller, getState } = createController( 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({ respondOAuthFlow: () => Promise.resolve(oauthFlow({
select: { requestId: "request-2", message: "Choose an account", options: [{ value: "acct-1", label: "Account 1" }] }, 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<typeof defaultApi.status>[0]; machineId: string | undefined }[] = []; const statusCalls: { session: Parameters<typeof defaultApi.status>[0]; machineId: string | undefined }[] = [];
const appliedStatuses: SessionStatus[] = []; const appliedStatuses: SessionStatus[] = [];
const { controller, getState } = createController( 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) => { respondOAuthFlow: (flowId, requestId, value, machineId) => {
respondCalls.push({ 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 () => { 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 flow = oauthFlow({ prompt: { requestId: "request-1", message: "Paste callback", kind: "manual" } });
const { controller, getState } = createController( 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")) }, { respondOAuthFlow: () => Promise.reject(new Error("Invalid callback")) },
); );
@@ -168,7 +168,7 @@ describe("AuthController", () => {
const response = deferred<OAuthFlowState>(); const response = deferred<OAuthFlowState>();
const cancellation = deferred<OAuthFlowState>(); const cancellation = deferred<OAuthFlowState>();
const { controller, getState } = createController( const { controller, getState } = createController(
{ authDialog: { step: "oauth", flow, inputValue: "https://callback" } }, { authDialog: { step: "oauth", flow, machineId: "local", inputValue: "https://callback" } },
{ {
respondOAuthFlow: () => response.promise, respondOAuthFlow: () => response.promise,
cancelOAuthFlow: () => cancellation.promise, cancelOAuthFlow: () => cancellation.promise,
@@ -219,6 +219,78 @@ describe("AuthController", () => {
expect(getState().error).toBe(""); 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 () => { it("does not let a stale OAuth response overwrite a newer flow", async () => {
vi.stubGlobal("window", { setInterval: () => 1, clearInterval: () => undefined }); vi.stubGlobal("window", { setInterval: () => 1, clearInterval: () => undefined });
const oldPrompt = { requestId: "request-1", message: "Paste callback", kind: "manual" } as const; const oldPrompt = { requestId: "request-1", message: "Paste callback", kind: "manual" } as const;
@@ -227,7 +299,7 @@ describe("AuthController", () => {
const response = deferred<OAuthFlowState>(); const response = deferred<OAuthFlowState>();
const providers = [authProvider("anthropic", "oauth")]; const providers = [authProvider("anthropic", "oauth")];
const { controller, getState } = createController( 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, respondOAuthFlow: () => response.promise,
authProviders: () => Promise.resolve({ providers }), authProviders: () => Promise.resolve({ providers }),
@@ -296,7 +368,7 @@ describe("AuthController", () => {
const flow = oauthFlow({ prompt: { requestId: "request-1", message: "Paste callback", kind: "manual" } }); const flow = oauthFlow({ prompt: { requestId: "request-1", message: "Paste callback", kind: "manual" } });
const cancelCalls: { flowId: string; machineId: string | undefined }[] = []; const cancelCalls: { flowId: string; machineId: string | undefined }[] = [];
const { controller, getState } = createController( const { controller, getState } = createController(
{ authDialog: { step: "oauth", flow } }, { authDialog: { step: "oauth", flow, machineId: "local" } },
{ {
cancelOAuthFlow: (flowId, machineId) => { cancelOAuthFlow: (flowId, machineId) => {
cancelCalls.push({ flowId, machineId }); cancelCalls.push({ flowId, machineId });
@@ -390,13 +462,14 @@ function createController(
) { ) {
let state: AppState = { ...initialAppState(), ...statePatch }; let state: AppState = { ...initialAppState(), ...statePatch };
const api = { ...defaultApi, ...apiPatch }; const api = { ...defaultApi, ...apiPatch };
const setState = (patch: Partial<AppState>) => { state = { ...state, ...patch }; };
const controller = new AuthController( const controller = new AuthController(
() => state, () => state,
(patch) => { state = { ...state, ...patch }; }, setState,
applyStatus, applyStatus,
{ api }, { api },
); );
return { controller, getState: () => state }; return { controller, getState: () => state, setState };
} }
async function flushMicrotasks(): Promise<void> { async function flushMicrotasks(): Promise<void> {
+18 -18
View File
@@ -142,10 +142,10 @@ export class AuthController {
delete clean.error; delete clean.error;
this.setState({ authDialog: { ...clean, responding: true } }); this.setState({ authDialog: { ...clean, responding: true } });
try { 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); const current = this.currentOAuthDialog(operationGeneration, flowId);
if (flow.flowId !== flowId || current === undefined || oauthRequestId(current.flow) !== requestId) return; if (flow.flowId !== flowId || current === undefined || oauthRequestId(current.flow) !== requestId) return;
this.updateOAuthFlow(flow); this.updateOAuthFlow(flow, current.machineId);
} catch (error) { } catch (error) {
const current = this.currentOAuthDialog(operationGeneration, flowId); const current = this.currentOAuthDialog(operationGeneration, flowId);
if (current === undefined || oauthRequestId(current.flow) !== requestId) return; if (current === undefined || oauthRequestId(current.flow) !== requestId) return;
@@ -160,7 +160,7 @@ export class AuthController {
return; return;
} }
const flowId = dialog.flow.flowId; const flowId = dialog.flow.flowId;
const machineId = selectedMachineId(this.getState()); const machineId = dialog.machineId;
this.closeDialog(); this.closeDialog();
try { try {
await this.api.cancelOAuthFlow(flowId, machineId); await this.api.cancelOAuthFlow(flowId, machineId);
@@ -217,8 +217,8 @@ export class AuthController {
} }
return; return;
} }
this.updateOAuthFlow(flow); this.updateOAuthFlow(flow, machineId);
if (flow.status === "running") this.startPolling(flow.flowId); if (flow.status === "running") this.startPolling(flow.flowId, machineId);
} catch (error) { } catch (error) {
if (operationGeneration === this.oauthOperationGeneration) this.setState({ error: String(error) }); if (operationGeneration === this.oauthOperationGeneration) this.setState({ error: String(error) });
} }
@@ -232,11 +232,11 @@ export class AuthController {
return true; return true;
} }
private updateOAuthFlow(flow: OAuthFlowState): void { private updateOAuthFlow(flow: OAuthFlowState, machineId: string): void {
if (flow.status === "complete") { if (flow.status === "complete") {
this.stopPolling(); this.stopPolling();
this.closeDialog(); this.closeDialog();
void this.refreshStatus(); void this.refreshStatus(machineId);
return; return;
} }
if (flow.status === "error" || flow.status === "cancelled") { if (flow.status === "error" || flow.status === "cancelled") {
@@ -250,14 +250,14 @@ export class AuthController {
const sameRequest = previousRequestId !== undefined && previousRequestId === newRequestId; const sameRequest = previousRequestId !== undefined && previousRequestId === newRequestId;
const inputValue = sameRequest ? previousInput : ""; const inputValue = sameRequest ? previousInput : "";
const responding = sameRequest && existing?.step === "oauth" ? existing.responding === true : false; 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(); this.stopPolling();
const operationGeneration = this.oauthOperationGeneration; const operationGeneration = this.oauthOperationGeneration;
const pollGeneration = this.pollGeneration; 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 { private stopPolling(): void {
@@ -267,22 +267,22 @@ export class AuthController {
this.pollTimer = undefined; this.pollTimer = undefined;
} }
private async poll(flowId: string, operationGeneration: number, pollGeneration: number): Promise<void> { private async poll(flowId: string, machineId: string, operationGeneration: number, pollGeneration: number): Promise<void> {
if (pollGeneration !== this.pollGeneration) return; if (pollGeneration !== this.pollGeneration) return;
const dialog = this.currentOAuthDialog(operationGeneration, flowId); const dialog = this.currentOAuthDialog(operationGeneration, flowId);
if (dialog === undefined) { if (dialog?.machineId !== machineId) {
this.stopPolling(); this.stopPolling();
return; return;
} }
const requestId = oauthRequestId(dialog.flow); const requestId = oauthRequestId(dialog.flow);
try { 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); const current = this.currentOAuthDialog(operationGeneration, flowId);
if (flow.flowId !== flowId || pollGeneration !== this.pollGeneration || current === undefined || oauthRequestId(current.flow) !== requestId) return; if (flow.flowId !== flowId || pollGeneration !== this.pollGeneration || current?.machineId !== machineId || oauthRequestId(current.flow) !== requestId) return;
this.updateOAuthFlow(flow); this.updateOAuthFlow(flow, machineId);
} catch (error) { } catch (error) {
const current = this.currentOAuthDialog(operationGeneration, flowId); 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.stopPolling();
this.setState({ authDialog: { ...current, error: String(error) } }); this.setState({ authDialog: { ...current, error: String(error) } });
} }
@@ -294,11 +294,11 @@ export class AuthController {
return dialog?.step === "oauth" && dialog.flow.flowId === flowId ? dialog : undefined; return dialog?.step === "oauth" && dialog.flow.flowId === flowId ? dialog : undefined;
} }
private async refreshStatus(): Promise<void> { private async refreshStatus(machineId = selectedMachineId(this.getState())): Promise<void> {
const session = this.session(); const session = this.session();
if (session === undefined) return; if (session === undefined) return;
try { try {
this.applyStatus(await this.api.status(session, selectedMachineId(this.getState()))); this.applyStatus(await this.api.status(session, machineId));
} catch { } catch {
// Status refresh is opportunistic after login completes. // Status refresh is opportunistic after login completes.
} }