Merge pull request #73 from jmfederico/agent/pr-64-audit

fix(auth): preserve browser auth-flow ownership

.changeset/cancel-stale-auth-flows.md

@@ -0,0 +1,5 @@
+---
+"@jmfederico/pi-web": patch
+---
+
+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

@@ -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

@@ -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 });
@@ -144,10 +144,73 @@ describe("AuthController", () => {
     expect(appliedStatuses).toEqual([refreshedStatus]);
   });
 
+  it("does not refresh a session from another selected machine when a flow completes", async () => {
+    const flow = oauthFlow({ prompt: { requestId: "request-1", message: "Enter secret", kind: "prompt", promptType: "secret" } });
+    const respondMachines: (string | undefined)[] = [];
+    const statusMachines: (string | undefined)[] = [];
+    const appliedStatuses: SessionStatus[] = [];
+    const { controller } = createController(
+      {
+        selectedMachine: remoteMachine("remote-2"),
+        selectedSession: sessionInfo("session-2"),
+        authDialog: { step: "oauth", flow, machineId: "remote-1", inputValue: "secret-value" },
+      },
+      {
+        respondOAuthFlow: (_flowId, _requestId, _value, machineId) => {
+          respondMachines.push(machineId);
+          return Promise.resolve(oauthFlow({ status: "complete" }));
+        },
+        status: (_session, machineId) => {
+          statusMachines.push(machineId);
+          return Promise.resolve(sessionStatus("session-2"));
+        },
+      },
+      (status) => { appliedStatuses.push(status); },
+    );
+
+    await controller.respondOAuth();
+    await flushMicrotasks();
+
+    expect(respondMachines).toEqual(["remote-1"]);
+    expect(statusMachines).toEqual([]);
+    expect(appliedStatuses).toEqual([]);
+  });
+
+  it("does not apply an auth status refresh after the selected session changes", async () => {
+    const flow = oauthFlow({ prompt: { requestId: "request-1", message: "Enter secret", kind: "prompt", promptType: "secret" } });
+    const originalSession = sessionInfo("session-1");
+    const statusResponse = deferred<SessionStatus>();
+    const statusCalls: { session: Parameters<typeof defaultApi.status>[0]; machineId: string | undefined }[] = [];
+    const appliedStatuses: SessionStatus[] = [];
+    const { controller, setState } = createController(
+      {
+        selectedMachine: remoteMachine("remote-1"),
+        selectedSession: originalSession,
+        authDialog: { step: "oauth", flow, machineId: "remote-1", inputValue: "secret-value" },
+      },
+      {
+        respondOAuthFlow: () => Promise.resolve(oauthFlow({ status: "complete" })),
+        status: (session, machineId) => {
+          statusCalls.push({ session, machineId });
+          return statusResponse.promise;
+        },
+      },
+      (status) => { appliedStatuses.push(status); },
+    );
+
+    await controller.respondOAuth();
+    setState({ selectedMachine: remoteMachine("remote-2"), selectedSession: sessionInfo("session-2") });
+    statusResponse.resolve(sessionStatus(originalSession.id));
+    await flushMicrotasks();
+
+    expect(statusCalls).toEqual([{ session: originalSession, machineId: "remote-1" }]);
+    expect(appliedStatuses).toEqual([]);
+  });
+
   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 +231,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,
@@ -191,6 +254,106 @@ describe("AuthController", () => {
     expect(getState().authDialog).toBeUndefined();
   });
 
+  it("best-effort cancels a running flow whose start response arrives after the dialog closes", async () => {
+    const start = deferred<OAuthFlowState>();
+    const provider: AuthProviderOption = { ...authProvider("amazon-bedrock", "api_key"), loginFlow: "interactive" };
+    const cancelCalls: { flowId: string; machineId: string | undefined }[] = [];
+    const { controller, getState } = createController(
+      {
+        selectedMachine: remoteMachine("remote-1"),
+        authDialog: { step: "providers", mode: "login", authType: "api_key", providers: [provider] },
+      },
+      {
+        startInteractiveApiKeyLogin: () => start.promise,
+        cancelOAuthFlow: (flowId, machineId) => {
+          cancelCalls.push({ flowId, machineId });
+          return Promise.reject(new Error("Cancel unavailable"));
+        },
+      },
+    );
+
+    const pendingStart = controller.selectLoginProvider(provider.id, provider.authType);
+    controller.closeDialog();
+    start.resolve(oauthFlow({ flowId: "stale-flow", providerId: provider.id, providerName: provider.name }));
+    await pendingStart;
+
+    expect(cancelCalls).toEqual([{ flowId: "stale-flow", machineId: "remote-1" }]);
+    expect(getState().authDialog).toBeUndefined();
+    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;
@@ -199,7 +362,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 }),
@@ -268,7 +431,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 });
@@ -362,13 +525,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

@@ -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);
@@ -205,9 +205,20 @@ export class AuthController {
       const flow = provider.authType === "oauth"
         ? await this.api.startOAuthLogin(provider.id, machineId)
         : await this.api.startInteractiveApiKeyLogin(provider.id, machineId);
-      if (operationGeneration !== this.oauthOperationGeneration) return;
+      if (operationGeneration !== this.oauthOperationGeneration) {
-      this.updateOAuthFlow(flow);
+        // Sessiond has already allocated this flow. Do not orphan its timer,
-      if (flow.status === "running") this.startPolling(flow.flowId);
+        // provider polling, or callback listener when the UI operation is stale.
+        if (flow.status === "running") {
+          try {
+            await this.api.cancelOAuthFlow(flow.flowId, machineId);
+          } catch {
+            // Best-effort cleanup; the obsolete flow must not restore UI state.
+          }
+        }
+        return;
+      }
+      this.updateOAuthFlow(flow, machineId);
+      if (flow.status === "running") this.startPolling(flow.flowId, machineId);
     } catch (error) {
       if (operationGeneration === this.oauthOperationGeneration) this.setState({ error: String(error) });
     }
@@ -221,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") {
@@ -239,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 {
@@ -256,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;
+      if (flow.flowId !== flowId || pollGeneration !== this.pollGeneration || current?.machineId !== machineId || oauthRequestId(current.flow) !== requestId) return;
-      this.updateOAuthFlow(flow);
+      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) } });
     }
@@ -283,18 +294,23 @@ 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();
+    const session = this.selectedSessionForMachine(machineId);
     if (session === undefined) return;
     try {
-      this.applyStatus(await this.api.status(session, selectedMachineId(this.getState())));
+      const status = await this.api.status(session, machineId);
+      const current = this.selectedSessionForMachine(machineId);
+      if (current?.id !== session.id || current.cwd !== session.cwd) return;
+      this.applyStatus(status);
     } catch {
       // Status refresh is opportunistic after login completes.
     }
   }
 
-  private session() {
+  private selectedSessionForMachine(machineId: string) {
-    const session = this.getState().selectedSession;
+    const state = this.getState();
+    if (selectedMachineId(state) !== machineId) return undefined;
+    const session = state.selectedSession;
     if (session === undefined || session.archived === true) return undefined;
     return session;
   }