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
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.
}