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

fix(auth): preserve browser auth-flow ownership

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;
-      this.updateOAuthFlow(flow);
-      if (flow.status === "running") this.startPolling(flow.flowId);
+      if (operationGeneration !== this.oauthOperationGeneration) {
+        // Sessiond has already allocated this flow. Do not orphan its timer,
+        // 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;
-      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) } });
     }
@@ -283,18 +294,23 @@ export class AuthController {
     return dialog?.step === "oauth" && dialog.flow.flowId === flowId ? dialog : undefined;
   }
 
-  private async refreshStatus(): Promise<void> {
-    const session = this.session();
+  private async refreshStatus(machineId = selectedMachineId(this.getState())): Promise<void> {
+    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() {
-    const session = this.getState().selectedSession;
+  private selectedSessionForMachine(machineId: string) {
+    const state = this.getState();
+    if (selectedMachineId(state) !== machineId) return undefined;
+    const session = state.selectedSession;
     if (session === undefined || session.archived === true) return undefined;
     return session;
   }