Merge pull request #2237 from modelcontextprotocol/v2/fix/2223-oauth-recovery-reset-during-render

fix(web): reset pending OAuth slots during render, not in an effect

clients/web/src/hooks/useOAuthRecovery.test.tsx

@@ -9,7 +9,7 @@ import type {
 import type { AuthChallenge } from "@inspector/core/auth/challenge.js";
 import { AuthRecoveryRequiredError } from "@inspector/core/auth/challenge.js";
 import { EmaClientNotConfiguredError } from "@inspector/core/auth/ema/clientConfigError.js";
-import { useLayoutEffect, useRef } from "react";
+import { useEffect, useLayoutEffect, useRef } from "react";
 import { renderWithMantine, act, waitFor } from "../test/renderWithMantine";
 import {
   OAUTH_RESUME_KEY,
@@ -169,9 +169,22 @@ interface HarnessProps {
   noFetchLog?: boolean;
 }
 
+/** One committed render's view of the two server-scoped pending slots. */
+interface PendingCommit {
+  stepUpServerId: string | undefined;
+  reauthServerId: string | undefined;
+}
+
 interface Harness {
   api: () => OAuthRecovery;
   rerender: (next: HarnessProps) => void;
+  /**
+   * Every committed render's pending slots, in order. Recorded from an effect
+   * with no dependency array, so one entry lands per commit — which is what
+   * lets a test assert that a reset was visible in the *first* frame after a
+   * server switch rather than merely eventually (#2223).
+   */
+  commits: PendingCommit[];
   spies: {
     setActiveServerId: ReturnType<typeof vi.fn>;
     setFailedServerId: ReturnType<typeof vi.fn>;
@@ -183,6 +196,7 @@ interface Harness {
 
 function harness(initial: HarnessProps = {}): Harness {
   let latest: OAuthRecovery | undefined;
+  const commits: PendingCommit[] = [];
   const spies = {
     setActiveServerId: vi.fn(),
     setFailedServerId: vi.fn(),
@@ -236,6 +250,15 @@ function harness(initial: HarnessProps = {}): Harness {
       clearResultPanels: spies.clearResultPanels,
       setSourceScopedError: spies.setSourceScopedError,
     });
+    // No dependency array: one entry per commit. `pendingReauth` is not part
+    // of the hook's public surface, so it is read back through the session ref
+    // the hook mirrors it into.
+    useEffect(() => {
+      commits.push({
+        stepUpServerId: latest?.pendingStepUp?.serverId,
+        reauthServerId: sessionRef.current.pendingReauth?.serverId,
+      });
+    });
     return null;
   }
 
@@ -246,6 +269,7 @@ function harness(initial: HarnessProps = {}): Harness {
       return latest;
     },
     rerender: (next) => rerender(<Probe p={next} />),
+    commits,
     spies,
   };
 }
@@ -322,6 +346,94 @@ describe("useOAuthRecovery", () => {
       expect(client.handleAuthChallenge).not.toHaveBeenCalled();
     });
 
+    it("clears both slots in the first committed frame after a server switch", async () => {
+      visibility.visible = false;
+      const client = fakeClient();
+      const servers = [entry("a"), entry("b")];
+      const h = harness({ servers, activeServerId: "a", client });
+
+      await act(async () => {
+        await h
+          .api()
+          .handleCommandScopedAuthRecovery(
+            new AuthRecoveryRequiredError(
+              AUTH_URL,
+              challenge("insufficient_scope"),
+            ),
+            { serverId: "a", source: "tool" },
+          );
+      });
+      await act(async () => {
+        client.emit("authChallengeInteractive", {
+          challenge: challenge(),
+          authorizationUrl: AUTH_URL,
+        });
+      });
+      await waitFor(() => {
+        const last = h.commits[h.commits.length - 1];
+        expect(last?.stepUpServerId).toBe("a");
+        expect(last?.reauthServerId).toBe("a");
+      });
+
+      const before = h.commits.length;
+      await act(async () => {
+        h.rerender({ servers, activeServerId: "b", client });
+      });
+
+      // The point of the test: the *first* frame after the switch already
+      // shows both slots empty. Resetting them in an effect instead would
+      // commit one frame still carrying server "a"'s step-up prompt and
+      // deferred re-auth, and would pass a `waitFor` assertion just the same.
+      expect(h.commits.length).toBeGreaterThan(before);
+      expect(h.commits[before]).toEqual({
+        stepUpServerId: undefined,
+        reauthServerId: undefined,
+      });
+    });
+
+    it("does not carry the left server's retry into a later ambient step-up", async () => {
+      const client = fakeClient({
+        handleAuthChallenge: vi.fn().mockResolvedValue({ kind: "satisfied" }),
+      });
+      const servers = [entry("a", {}, true), entry("b", {}, true)];
+      const h = harness({ servers, activeServerId: "a", client });
+      const retryA = vi.fn().mockResolvedValue(undefined);
+
+      await act(async () => {
+        await h
+          .api()
+          .handleCommandScopedAuthRecovery(
+            new AuthRecoveryRequiredError(
+              AUTH_URL,
+              challenge("insufficient_scope"),
+            ),
+            { serverId: "a", source: "tool", retryOperation: retryA },
+          );
+      });
+      expect(h.api().pendingStepUp?.serverId).toBe("a");
+
+      await act(async () => {
+        h.rerender({ servers, activeServerId: "b", client });
+      });
+      expect(h.api().pendingStepUp).toBeNull();
+
+      // Server B raises its own step-up, which carries no retry of its own.
+      await act(async () => {
+        client.emit("authChallengeInteractive", {
+          challenge: challenge("insufficient_scope"),
+          authorizationUrl: AUTH_URL,
+        });
+      });
+      await waitFor(() => expect(h.api().pendingStepUp?.serverId).toBe("b"));
+
+      await act(async () => {
+        await h.api().handleStepUpAuthorize();
+      });
+      // Authorizing B must not run the command A was left mid-way through.
+      expect(retryA).not.toHaveBeenCalled();
+      expect(toastTitles()).toContain("Permissions updated");
+    });
+
     it("refuses a second step-up while one is open", async () => {
       const client = fakeClient();
       const h = harness({ servers: [entry("a")], activeServerId: "a", client });

clients/web/src/hooks/useOAuthRecovery.ts

@@ -54,6 +54,7 @@ import {
 import { OAUTH_CALLBACK_PATH } from "../utils/oauthFlow";
 import { INSPECTOR_SERVERS_TAB } from "../utils/inspectorTabs";
 import type { PendingReauth } from "../utils/pendingReauth";
+import { useValueChange } from "./useValueChange";
 import {
   authRecoveryRestoredMessage,
   authRecoveryAbandonedMessage,
@@ -318,19 +319,52 @@ export function useOAuthRecovery({
     sessionRef.current.pendingReauth = pendingReauth;
   });
 
-  useEffect(() => {
+  // Both slots are server-scoped, so switching servers has to clear whichever
-    const pending = sessionRef.current.pendingReauth;
+  // one belongs to the server we just left. Adjusted **during render** rather
-    if (pending && pending.serverId !== activeServerId) {
+  // than in an effect: an effect only runs after the commit, so the frame that
-      setPendingReauth(null);
+  // shows the new server would still paint the previous server's re-auth
-    }
+  // banner or step-up prompt — a modal-grade authorization affordance attached
-    const stepUp = sessionRef.current.pendingStepUp;
+  // to the wrong server (#2223). `activeServerId` is a primitive, so it is
-    if (stepUp && stepUp.serverId !== activeServerId) {
+  // referentially stable across renders that mean "no change", which is what
-      setPendingStepUp(null);
+  // `useValueChange`'s `Object.is` comparison requires. The previous value is
-    }
+  // read from the updater argument rather than from `sessionRef.current`,
-  }, [sessionRef, activeServerId]);
+  // keeping the callback pure — a render can be replayed or abandoned, so it
+  // must not depend on external mutable state.
+  useValueChange(activeServerId, () => {
+    setPendingReauth((current) =>
+      current && current.serverId !== activeServerId ? null : current,
+    );
+    setPendingStepUp((current) =>
+      current && current.serverId !== activeServerId ? null : current,
+    );
+  });
 
+  /**
+   * Opens the step-up prompt, or refuses when one is already up.
+   *
+   * The retry is installed here rather than by the caller because it belongs
+   * to the prompt and must be written in the same synchronous step that opens
+   * it (#2237). Two things follow, and both are load-bearing:
+   *
+   * - A **refused** prompt installs nothing, so a command that was just told
+   *   it cannot start does not replace the open prompt's operation (#2165).
+   * - An **ambient** prompt passes `null` and so *overwrites* whatever the
+   *   previous prompt left in the ref. Nothing else has to clear it: the paths
+   *   that drop the slot as a consequence of something else — the
+   *   server-switch reset, `resetOAuthRecoveryState` on disconnect — can leave
+   *   a stale closure behind, because the only thing that reads it is
+   *   `handleStepUpAuthorize`, which needs a prompt, and every prompt comes
+   *   through here. Clearing it from an effect instead would race: an async
+   *   continuation can open a prompt and install its retry before a previous
+   *   commit's passive effect flushes, and that effect's stale
+   *   `pendingStepUp === null` snapshot would then delete the new prompt's
+   *   retry (Copilot on #2237).
+   */
   const trySetPendingStepUp = useCallback(
-    (next: PendingStepUp): boolean => {
+    (
+      next: PendingStepUp,
+      retryOperation: (() => Promise<unknown>) | null,
+    ): boolean => {
       if (sessionRef.current.pendingStepUp !== null) {
         notifications.show({
           title: "Step-up authorization in progress",
@@ -342,6 +376,7 @@ export function useOAuthRecovery({
         return false;
       }
       setPendingStepUp(next);
+      pendingStepUpRetryRef.current = retryOperation;
       return true;
     },
     [sessionRef],
@@ -619,13 +654,18 @@ export function useOAuthRecovery({
     }) => {
       const server = sessionRef.current.servers.find((s) => s.id === serverId);
       if (isStepUpConfirmation(challenge, server)) {
-        trySetPendingStepUp({
+        // Ambient: nothing to retry, and passing `null` is what evicts a
-          challenge,
+        // previous prompt's closure.
-          authorizationUrl,
+        trySetPendingStepUp(
-          serverId,
+          {
-          source,
+            challenge,
-          enterpriseManaged: isEmaStepUp(challenge, server),
+            authorizationUrl,
-        });
+            serverId,
+            source,
+            enterpriseManaged: isEmaStepUp(challenge, server),
+          },
+          null,
+        );
         return;
       }
       prepareOAuthRedirect({
@@ -690,23 +730,16 @@ export function useOAuthRecovery({
         ) {
           return true;
         }
-        // The retry belongs to the prompt, so it is installed only once the
+        trySetPendingStepUp(
-        // prompt is actually open (#2165). `trySetPendingStepUp` REFUSES a
+          {
-        // second prompt while one is already up — writing the ref first meant
-        // the refused command's operation replaced the open prompt's, and
-        // authorizing that prompt then ran the command the user was just told
-        // could not start.
-        if (
-          trySetPendingStepUp({
             challenge: error.authChallenge,
             authorizationUrl: error.authorizationUrl,
             serverId: options.serverId,
             source: options.source,
             enterpriseManaged: isEmaStepUp(error.authChallenge, server),
-          })
+          },
-        ) {
+          options.retryOperation ?? null,
-          pendingStepUpRetryRef.current = options.retryOperation ?? null;
+        );
-        }
         return false;
       }