coder · ibetitsmike · Aug 14, 2026 · Aug 14, 2026 · Aug 14, 2026
@@ -755,6 +755,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 `window` so the completion
+// message can carry the popup as its source.
+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. */
@@ -775,6 +797,130 @@ export const WithMCPNeedingAuth: Story = {
 	},
 };
 
+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/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();
+	},
+};
+
 /** No MCP servers active — shows only "MCP" label with chevron. */
 export const WithMCPNoneActive: Story = {
 	args: {

@@ -17,6 +17,7 @@ import type React from "react";
 import {
 	type FC,
 	useEffect,
+	useEffectEvent,
 	useImperativeHandle,
 	useRef,
 	useState,
@@ -441,7 +442,12 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
 	);
 	const [workspacePickerOpen, setWorkspacePickerOpen] = useState(false);
 	const [mcpConnectingId, setMcpConnectingId] = useState<string | null>(null);
-	const mcpPopupRef = useRef<Window | null>(null);
+	// Correlates a completion message with the initiating OAuth flow.
+	// Retained after popup close: the callback page posts before closing,
+	// and the close poll can run before the queued message is dispatched.
+	const mcpAuthFlowRef = useRef<{ popup: Window; serverID: string } | null>(
+		null,
+	);
 	const [mcpDisconnectTarget, setMcpDisconnectTarget] =
 		useState<TypesGen.MCPServerConfig | null>(null);
 	const queryClient = useQueryClient();
@@ -518,6 +524,30 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
 		[],
 	);
 
+	const handleMCPAuthComplete = useEffectEvent(
+		(serverID: string, source: MessageEventSource | null) => {
+			onMCPAuthComplete?.(serverID);
+			// Only a message from the initiating popup for the initiating
+			// server may change the selection.
+			const flow = mcpAuthFlowRef.current;
+			if (!flow || source !== flow.popup || serverID !== flow.serverID) {
+				return;
+			}
+			mcpAuthFlowRef.current = null;
+			setMcpConnectingId(null);
+			if (
+				onMCPSelectionChange &&
+				selectedMCPServerIds &&
+				mcpServers?.some(
+					(server) => server.id === serverID && server.enabled,
+				) &&
+				!selectedMCPServerIds.includes(serverID)
+			) {
+				onMCPSelectionChange([...selectedMCPServerIds, serverID]);
+			}
+		},
+	);
+
 	// Listen for OAuth2 completion postMessage from popup.
 	useEffect(() => {
 		const handler = (event: MessageEvent) => {
@@ -526,29 +556,29 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
 				event.data?.type === "mcp-oauth2-complete" &&
 				typeof event.data.serverID === "string"
 			) {
-				setMcpConnectingId(null);
-				onMCPAuthComplete?.(event.data.serverID);
-				mcpPopupRef.current = null;
+				handleMCPAuthComplete(event.data.serverID, event.source);
 			}
 		};
 		window.addEventListener("message", handler);
 		return () => window.removeEventListener("message", handler);
-	}, [onMCPAuthComplete]);
+	}, []);
 
-	// Poll for popup close and clean up on unmount.
+	// Clear only the connecting indicator when the popup closes; the flow
+	// ref stays so a completion message posted before close still
+	// correlates.
 	useEffect(() => {
-		if (!mcpConnectingId || !mcpPopupRef.current) return;
+		if (!mcpConnectingId || !mcpAuthFlowRef.current) return;
 		const interval = setInterval(() => {
-			if (mcpPopupRef.current?.closed) {
+			if (mcpAuthFlowRef.current?.popup.closed) {
 				setMcpConnectingId(null);
-				mcpPopupRef.current = null;
 			}
 		}, 500);
 		return () => {
 			clearInterval(interval);
-			if (mcpPopupRef.current && !mcpPopupRef.current.closed) {
-				mcpPopupRef.current.close();
-				mcpPopupRef.current = null;
+			const popup = mcpAuthFlowRef.current?.popup;
+			if (popup && !popup.closed) {
+				popup.close();
+				mcpAuthFlowRef.current = null;
 			}
 		};
 	}, [mcpConnectingId]);
@@ -567,11 +597,8 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
 	const handleMcpConnect = (server: TypesGen.MCPServerConfig) => {
 		setMcpConnectingId(server.id);
 		const connectUrl = `/api/experimental/mcp/servers/${encodeURIComponent(server.id)}/oauth2/connect`;
-		mcpPopupRef.current = window.open(
-			connectUrl,
-			"_blank",
-			"width=900,height=600",
-		);
+		const popup = window.open(connectUrl, "_blank", "width=900,height=600");
+		mcpAuthFlowRef.current = popup ? { popup, serverID: server.id } : null;
 	};
 
 	const handleMcpDisconnectConfirm = () => {