fix(site): auto-select MCP server after successful OAuth (#28155)

site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx

@@ -756,6 +756,28 @@ const mcpDefaults = {
 	onMCPAuthComplete: fn(),
 };
 
+const dispatchMCPOAuthComplete = (
+	serverID: string,
+	source: MessageEventSource | null = null,
+) => {
+	window.dispatchEvent(
+		new MessageEvent("message", {
+			data: { type: "mcp-oauth2-complete", serverID },
+			origin: location.origin,
+			source,
+		}),
+	);
+};
+
+// Requires window.open mocked to return a Window; the completion
+// message's source must be that same mocked popup to correlate.
+const startMCPOAuthFlow = async (canvasElement: HTMLElement) => {
+	const canvas = within(canvasElement);
+	const body = within(canvasElement.ownerDocument.body);
+	await userEvent.click(canvas.getByRole("button", { name: "More options" }));
+	await userEvent.click(await body.findByRole("button", { name: "Auth" }));
+};
+
 // ── MCP stories ────────────────────────────────────────────────
 
 /** Input with multiple MCP servers selected — shows icon stack in toolbar. */
@@ -787,6 +809,133 @@ export const WithMCPNeedingAuth: Story = {
 			"_blank",
 			"width=900,height=600",
 		);
+		// The popup was blocked (window.open returned null), so the flow
+		// must not enter the connecting state that disables Auth buttons.
+		expect(body.getByRole("button", { name: "Auth" })).toBeEnabled();
+	},
+};
+
+export const MCPAutoEnablesAfterOAuthCompletes: Story = {
+	args: {
+		...mcpDefaults,
+		mcpServers: [linearMCP, githubMCP],
+		selectedMCPServerIds: [linearMCP.id],
+	},
+	beforeEach: () => {
+		spyOn(window, "open").mockReturnValue(window);
+	},
+	play: async ({ args, canvasElement }) => {
+		await startMCPOAuthFlow(canvasElement);
+		expect(window.open).toHaveBeenCalledWith(
+			`/api/experimental/organizations/org-1/mcp-servers/${githubMCP.id}/oauth2/connect`,
+			"_blank",
+			"width=900,height=600",
+		);
+		dispatchMCPOAuthComplete(githubMCP.id, window);
+
+		await waitFor(() => {
+			expect(args.onMCPSelectionChange).toHaveBeenCalledWith([
+				linearMCP.id,
+				githubMCP.id,
+			]);
+			expect(args.onMCPAuthComplete).toHaveBeenCalledWith(githubMCP.id);
+		});
+	},
+};
+
+// The coderd callback page posts the completion message and then closes
+// the popup, so the close poll can observe the closed popup before the
+// queued message is dispatched. An iframe contentWindow stands in for
+// the popup: it is a real Window whose closed becomes true on removal.
+export const MCPAutoEnablesWhenPopupClosesBeforeMessage: Story = {
+	args: {
+		...mcpDefaults,
+		mcpServers: [githubMCP],
+		selectedMCPServerIds: [],
+	},
+	play: async ({ args, canvasElement }) => {
+		const doc = canvasElement.ownerDocument;
+		const iframe = doc.createElement("iframe");
+		doc.body.appendChild(iframe);
+		const popup = iframe.contentWindow;
+		if (!popup) {
+			throw new Error("iframe contentWindow unavailable");
+		}
+		spyOn(window, "open").mockReturnValue(popup);
+
+		await startMCPOAuthFlow(canvasElement);
+		iframe.remove();
+		expect(popup.closed).toBe(true);
+		// Wait for the close poll to clear the connecting state before
+		// delivering the completion message.
+		const body = within(doc.body);
+		await waitFor(
+			() => {
+				expect(body.getByRole("button", { name: "Auth" })).toBeEnabled();
+			},
+			{ timeout: 2_000 },
+		);
+		dispatchMCPOAuthComplete(githubMCP.id, popup);
+
+		await waitFor(() => {
+			expect(args.onMCPSelectionChange).toHaveBeenCalledWith([githubMCP.id]);
+		});
+	},
+};
+
+export const MCPDoesNotDuplicateSelectionAfterOAuthCompletes: Story = {
+	args: {
+		...mcpDefaults,
+		mcpServers: [githubMCP],
+		selectedMCPServerIds: [githubMCP.id],
+	},
+	beforeEach: () => {
+		spyOn(window, "open").mockReturnValue(window);
+	},
+	play: async ({ args, canvasElement }) => {
+		await startMCPOAuthFlow(canvasElement);
+		dispatchMCPOAuthComplete(githubMCP.id, window);
+
+		await waitFor(() => {
+			expect(args.onMCPAuthComplete).toHaveBeenCalledWith(githubMCP.id);
+		});
+		expect(args.onMCPSelectionChange).not.toHaveBeenCalled();
+	},
+};
+
+export const MCPIgnoresUnsolicitedOAuthComplete: Story = {
+	args: {
+		...mcpDefaults,
+		mcpServers: [linearMCP, githubMCP],
+		selectedMCPServerIds: [linearMCP.id],
+	},
+	play: async ({ args }) => {
+		dispatchMCPOAuthComplete(githubMCP.id, window);
+
+		await waitFor(() => {
+			expect(args.onMCPAuthComplete).toHaveBeenCalledWith(githubMCP.id);
+		});
+		expect(args.onMCPSelectionChange).not.toHaveBeenCalled();
+	},
+};
+
+export const MCPIgnoresMismatchedServerAfterOAuthCompletes: Story = {
+	args: {
+		...mcpDefaults,
+		mcpServers: [linearMCP, githubMCP],
+		selectedMCPServerIds: [],
+	},
+	beforeEach: () => {
+		spyOn(window, "open").mockReturnValue(window);
+	},
+	play: async ({ args, canvasElement }) => {
+		await startMCPOAuthFlow(canvasElement);
+		dispatchMCPOAuthComplete(linearMCP.id, window);
+
+		await waitFor(() => {
+			expect(args.onMCPAuthComplete).toHaveBeenCalledWith(linearMCP.id);
+		});
+		expect(args.onMCPSelectionChange).not.toHaveBeenCalled();
 	},
 };
 

site/src/pages/AgentsPage/components/AgentChatInput.tsx

@@ -24,7 +24,6 @@ import {
 import { useMutation, useQueryClient } from "react-query";
 import { Link } from "react-router";
 import { toast } from "sonner";
-import { mcpServerOAuth2ConnectPath } from "#/api/api";
 import { getErrorMessage } from "#/api/errors";
 import { disconnectMCPServerOAuth2 } from "#/api/queries/chats";
 import type * as TypesGen from "#/api/typesGenerated";
@@ -62,6 +61,7 @@ import { cn } from "#/utils/cn";
 import { countInvisibleCharacters } from "#/utils/invisibleUnicode";
 import { isBelowMdViewport, isMobileViewport } from "#/utils/mobile";
 import { chatWidthClass, useChatFullWidth } from "../hooks/useChatFullWidth";
+import { useMCPOAuthFlow } from "../hooks/useMCPOAuthFlow";
 import { useOverflowCount } from "../hooks/useOverflowCount";
 import { useSpeechRecognition } from "../hooks/useSpeechRecognition";
 import {
@@ -429,8 +429,23 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
 		"main",
 	);
 	const [workspacePickerOpen, setWorkspacePickerOpen] = useState(false);
-	const [mcpConnectingId, setMcpConnectingId] = useState<string | null>(null);
-	const mcpPopupRef = useRef<Window | null>(null);
+	const { connectingServerId: mcpConnectingId, connect: connectMCPServer } =
+		useMCPOAuthFlow({
+			organizationId: chatOrganizationId,
+			onAuthComplete: onMCPAuthComplete,
+			onFlowSuccess: (serverID) => {
+				if (
+					onMCPSelectionChange &&
+					selectedMCPServerIds &&
+					mcpServers?.some(
+						(server) => server.id === serverID && server.enabled,
+					) &&
+					!selectedMCPServerIds.includes(serverID)
+				) {
+					onMCPSelectionChange([...selectedMCPServerIds, serverID]);
+				}
+			},
+		});
 	const [mcpDisconnectTarget, setMcpDisconnectTarget] =
 		useState<TypesGen.MCPServerConfig | null>(null);
 	const queryClient = useQueryClient();
@@ -507,41 +522,6 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
 		[],
 	);
 
-	// Listen for OAuth2 completion postMessage from popup.
-	useEffect(() => {
-		const handler = (event: MessageEvent) => {
-			if (event.origin !== location.origin) return;
-			if (
-				event.data?.type === "mcp-oauth2-complete" &&
-				typeof event.data.serverID === "string"
-			) {
-				setMcpConnectingId(null);
-				onMCPAuthComplete?.(event.data.serverID);
-				mcpPopupRef.current = null;
-			}
-		};
-		window.addEventListener("message", handler);
-		return () => window.removeEventListener("message", handler);
-	}, [onMCPAuthComplete]);
-
-	// Poll for popup close and clean up on unmount.
-	useEffect(() => {
-		if (!mcpConnectingId || !mcpPopupRef.current) return;
-		const interval = setInterval(() => {
-			if (mcpPopupRef.current?.closed) {
-				setMcpConnectingId(null);
-				mcpPopupRef.current = null;
-			}
-		}, 500);
-		return () => {
-			clearInterval(interval);
-			if (mcpPopupRef.current && !mcpPopupRef.current.closed) {
-				mcpPopupRef.current.close();
-				mcpPopupRef.current = null;
-			}
-		};
-	}, [mcpConnectingId]);
-
 	const handleMcpToggle = (serverId: string, checked: boolean) => {
 		if (!onMCPSelectionChange || !selectedMCPServerIds) return;
 		if (checked) {
@@ -553,22 +533,6 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
 		}
 	};
 
-	const handleMcpConnect = (server: TypesGen.MCPServerConfig) => {
-		if (!chatOrganizationId) {
-			return;
-		}
-		setMcpConnectingId(server.id);
-		const connectUrl = mcpServerOAuth2ConnectPath(
-			chatOrganizationId,
-			server.id,
-		);
-		mcpPopupRef.current = window.open(
-			connectUrl,
-			"_blank",
-			"width=900,height=600",
-		);
-	};
-
 	const handleMcpDisconnectConfirm = () => {
 		if (!mcpDisconnectTarget) {
 			return;
@@ -1397,7 +1361,7 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
 																	variant="outline"
 																	size="sm"
 																	className="h-6 shrink-0 px-2 text-[10px] leading-none"
-																	onClick={() => handleMcpConnect(server)}
+																	onClick={() => connectMCPServer(server.id)}
 																	disabled={
 																		isDisabled || mcpConnectingId !== null
 																	}

site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts

@@ -0,0 +1,102 @@
+import { useEffect, useEffectEvent, useRef, useState } from "react";
+import { mcpServerOAuth2ConnectPath } from "#/api/api";
+
+type UseMCPOAuthFlowOptions = {
+	organizationId?: string;
+	// Called for every same-origin completion message, regardless of
+	// which flow (if any) produced it, so callers can refresh server
+	// state.
+	onAuthComplete?: (serverId: string) => void;
+	// Called only for a completion posted by the initiating popup for
+	// the initiating server.
+	onFlowSuccess: (serverId: string) => void;
+};
+
+type MCPOAuthFlow = {
+	connectingServerId: string | null;
+	connect: (serverId: string) => void;
+};
+
+/**
+ * Runs the MCP server OAuth2 popup flow: opens the consent popup,
+ * listens for the completion message coderd's callback page posts to
+ * the opener, and reports success only for the initiating popup and
+ * server.
+ */
+export const useMCPOAuthFlow = ({
+	organizationId,
+	onAuthComplete,
+	onFlowSuccess,
+}: UseMCPOAuthFlowOptions): MCPOAuthFlow => {
+	const [connectingServerId, setConnectingServerId] = useState<string | null>(
+		null,
+	);
+	// Correlates a completion message with the initiating flow.
+	// Retained after popup close: the callback page posts before
+	// closing, and the close poll can run before the queued message is
+	// dispatched.
+	const flowRef = useRef<{ popup: Window; serverID: string } | null>(null);
+
+	const handleAuthComplete = useEffectEvent(
+		(serverID: string, source: MessageEventSource | null) => {
+			onAuthComplete?.(serverID);
+			// Only a message from the initiating popup for the initiating
+			// server counts as flow success.
+			const flow = flowRef.current;
+			if (!flow || source !== flow.popup || serverID !== flow.serverID) {
+				return;
+			}
+			flowRef.current = null;
+			setConnectingServerId(null);
+			onFlowSuccess(serverID);
+		},
+	);
+
+	useEffect(() => {
+		const handler = (event: MessageEvent) => {
+			if (event.origin !== location.origin) return;
+			if (
+				event.data?.type === "mcp-oauth2-complete" &&
+				typeof event.data.serverID === "string"
+			) {
+				handleAuthComplete(event.data.serverID, event.source);
+			}
+		};
+		window.addEventListener("message", handler);
+		return () => window.removeEventListener("message", handler);
+	}, []);
+
+	// Clear only the connecting indicator when the popup closes; the
+	// flow ref stays so a completion message posted before close still
+	// correlates.
+	useEffect(() => {
+		if (!connectingServerId || !flowRef.current) return;
+		const interval = setInterval(() => {
+			if (flowRef.current?.popup.closed) {
+				setConnectingServerId(null);
+			}
+		}, 500);
+		return () => {
+			clearInterval(interval);
+			const popup = flowRef.current?.popup;
+			if (popup && !popup.closed) {
+				popup.close();
+				flowRef.current = null;
+			}
+		};
+	}, [connectingServerId]);
+
+	const connect = (serverId: string) => {
+		if (!organizationId) {
+			return;
+		}
+		const connectUrl = mcpServerOAuth2ConnectPath(organizationId, serverId);
+		const popup = window.open(connectUrl, "_blank", "width=900,height=600");
+		// A blocked popup (window.open returns null) must not enter the
+		// connecting state; nothing would ever clear it.
+		flowRef.current = popup ? { popup, serverID: serverId } : null;
+		setConnectingServerId(popup ? serverId : null);
+	};
+
+	return { connectingServerId, connect };
+};