snowspeeder/pi-web
Archived
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(auth): bind flows to their originating machine

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/cancel-stale-auth-flows.md
+1 -1
View File
@@ -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.
src/client/src/appState.ts
+1 -1
View File
@@ -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<AppState,
src/client/src/controllers/authController.test.ts
+83 -10
View File
@@ -62,7 +62,7 @@ describe("AuthController", () => {
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<typeof defaultApi.status>[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<OAuthFlowState>();
const cancellation = deferred<OAuthFlowState>();
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<OAuthFlowState>();
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<AppState>) => { 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<void> {
src/client/src/controllers/authController.ts
+18 -18
View File
@@ -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<void> {
private async poll(flowId: string, machineId: string, operationGeneration: number, pollGeneration: number): Promise<void> {
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<void> {
private async refreshStatus(machineId = selectedMachineId(this.getState())): Promise<void> {
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.
}