From 28179038970fee51fa92c9c6f0d3fcc2d2ac2bd7 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Fri, 14 Aug 2026 06:53:30 +0000 Subject: [PATCH 1/7] fix(site): enable MCP server after OAuth --- .../components/AgentChatInput.stories.tsx | 77 +++++++++++++++++++ .../AgentsPage/components/AgentChatInput.tsx | 21 ++++- 2 files changed, 94 insertions(+), 4 deletions(-) diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx index 4c54454efe88e..5089546e55f19 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -756,6 +756,15 @@ const mcpDefaults = { onMCPAuthComplete: fn(), }; +const dispatchMCPOAuthComplete = (serverID: string) => { + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "mcp-oauth2-complete", serverID }, + origin: location.origin, + }), + ); +}; + // ── MCP stories ──────────────────────────────────────────────── /** Input with multiple MCP servers selected — shows icon stack in toolbar. */ @@ -790,6 +799,74 @@ export const WithMCPNeedingAuth: Story = { }, }; +export const MCPAutoEnablesAfterOAuthCompletes: Story = { + args: { + ...mcpDefaults, + mcpServers: [linearMCP, githubMCP], + selectedMCPServerIds: [linearMCP.id], + }, + play: async ({ args }) => { + dispatchMCPOAuthComplete(githubMCP.id); + + await waitFor(() => { + expect(args.onMCPSelectionChange).toHaveBeenCalledWith([ + linearMCP.id, + githubMCP.id, + ]); + expect(args.onMCPAuthComplete).toHaveBeenCalledWith(githubMCP.id); + }); + }, +}; + +export const MCPDoesNotDuplicateSelectionAfterOAuthCompletes: Story = { + args: { + ...mcpDefaults, + mcpServers: [githubMCP], + selectedMCPServerIds: [githubMCP.id], + }, + play: async ({ args }) => { + dispatchMCPOAuthComplete(githubMCP.id); + + await waitFor(() => { + expect(args.onMCPAuthComplete).toHaveBeenCalledWith(githubMCP.id); + }); + expect(args.onMCPSelectionChange).not.toHaveBeenCalled(); + }, +}; + +export const MCPIgnoresDisabledServerAfterOAuthCompletes: Story = { + args: { + ...mcpDefaults, + mcpServers: [{ ...githubMCP, enabled: false }], + selectedMCPServerIds: [], + }, + play: async ({ args }) => { + dispatchMCPOAuthComplete(githubMCP.id); + + await waitFor(() => { + expect(args.onMCPAuthComplete).toHaveBeenCalledWith(githubMCP.id); + }); + expect(args.onMCPSelectionChange).not.toHaveBeenCalled(); + }, +}; + +export const MCPIgnoresUnknownServerAfterOAuthCompletes: Story = { + args: { + ...mcpDefaults, + mcpServers: [linearMCP], + selectedMCPServerIds: [linearMCP.id], + }, + play: async ({ args }) => { + const unknownServerID = "mcp-unknown"; + dispatchMCPOAuthComplete(unknownServerID); + + await waitFor(() => { + expect(args.onMCPAuthComplete).toHaveBeenCalledWith(unknownServerID); + }); + expect(args.onMCPSelectionChange).not.toHaveBeenCalled(); + }, +}; + /** No MCP servers active — shows only "MCP" label with chevron. */ export const WithMCPNoneActive: Story = { args: { diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 0b3d25ba419b9..65269c08c5b34 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -17,6 +17,7 @@ import type React from "react"; import { type FC, useEffect, + useEffectEvent, useImperativeHandle, useRef, useState, @@ -507,6 +508,20 @@ export const AgentChatInput: FC = ({ [], ); + const handleMCPAuthComplete = useEffectEvent((serverID: string) => { + setMcpConnectingId(null); + onMCPAuthComplete?.(serverID); + if ( + onMCPSelectionChange && + selectedMCPServerIds && + mcpServers?.some((server) => server.id === serverID && server.enabled) && + !selectedMCPServerIds.includes(serverID) + ) { + onMCPSelectionChange([...selectedMCPServerIds, serverID]); + } + mcpPopupRef.current = null; + }); + // Listen for OAuth2 completion postMessage from popup. useEffect(() => { const handler = (event: MessageEvent) => { @@ -515,14 +530,12 @@ export const AgentChatInput: FC = ({ event.data?.type === "mcp-oauth2-complete" && typeof event.data.serverID === "string" ) { - setMcpConnectingId(null); - onMCPAuthComplete?.(event.data.serverID); - mcpPopupRef.current = null; + handleMCPAuthComplete(event.data.serverID); } }; window.addEventListener("message", handler); return () => window.removeEventListener("message", handler); - }, [onMCPAuthComplete]); + }, []); // Poll for popup close and clean up on unmount. useEffect(() => { From 0715721a3dd3285da899e7c852f09cc8d51d703b Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Fri, 14 Aug 2026 07:44:12 +0000 Subject: [PATCH 2/7] fix(site): gate MCP auto-select on the initiating OAuth popup --- .../components/AgentChatInput.stories.tsx | 61 ++++++++++++++----- .../AgentsPage/components/AgentChatInput.tsx | 40 +++++++----- 2 files changed, 71 insertions(+), 30 deletions(-) diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx index 5089546e55f19..bfd342cf852a4 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -756,15 +756,28 @@ const mcpDefaults = { onMCPAuthComplete: fn(), }; -const dispatchMCPOAuthComplete = (serverID: string) => { +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. */ @@ -805,8 +818,17 @@ export const MCPAutoEnablesAfterOAuthCompletes: Story = { mcpServers: [linearMCP, githubMCP], selectedMCPServerIds: [linearMCP.id], }, - play: async ({ args }) => { - dispatchMCPOAuthComplete(githubMCP.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([ @@ -824,8 +846,12 @@ export const MCPDoesNotDuplicateSelectionAfterOAuthCompletes: Story = { mcpServers: [githubMCP], selectedMCPServerIds: [githubMCP.id], }, - play: async ({ args }) => { - dispatchMCPOAuthComplete(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); @@ -834,14 +860,14 @@ export const MCPDoesNotDuplicateSelectionAfterOAuthCompletes: Story = { }, }; -export const MCPIgnoresDisabledServerAfterOAuthCompletes: Story = { +export const MCPIgnoresUnsolicitedOAuthComplete: Story = { args: { ...mcpDefaults, - mcpServers: [{ ...githubMCP, enabled: false }], - selectedMCPServerIds: [], + mcpServers: [linearMCP, githubMCP], + selectedMCPServerIds: [linearMCP.id], }, play: async ({ args }) => { - dispatchMCPOAuthComplete(githubMCP.id); + dispatchMCPOAuthComplete(githubMCP.id, window); await waitFor(() => { expect(args.onMCPAuthComplete).toHaveBeenCalledWith(githubMCP.id); @@ -850,18 +876,21 @@ export const MCPIgnoresDisabledServerAfterOAuthCompletes: Story = { }, }; -export const MCPIgnoresUnknownServerAfterOAuthCompletes: Story = { +export const MCPIgnoresMismatchedServerAfterOAuthCompletes: Story = { args: { ...mcpDefaults, - mcpServers: [linearMCP], - selectedMCPServerIds: [linearMCP.id], + mcpServers: [linearMCP, githubMCP], + selectedMCPServerIds: [], }, - play: async ({ args }) => { - const unknownServerID = "mcp-unknown"; - dispatchMCPOAuthComplete(unknownServerID); + beforeEach: () => { + spyOn(window, "open").mockReturnValue(window); + }, + play: async ({ args, canvasElement }) => { + await startMCPOAuthFlow(canvasElement); + dispatchMCPOAuthComplete(linearMCP.id, window); await waitFor(() => { - expect(args.onMCPAuthComplete).toHaveBeenCalledWith(unknownServerID); + expect(args.onMCPAuthComplete).toHaveBeenCalledWith(linearMCP.id); }); expect(args.onMCPSelectionChange).not.toHaveBeenCalled(); }, diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 65269c08c5b34..19c76023af79c 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -508,19 +508,31 @@ export const AgentChatInput: FC = ({ [], ); - const handleMCPAuthComplete = useEffectEvent((serverID: string) => { - setMcpConnectingId(null); - onMCPAuthComplete?.(serverID); - if ( - onMCPSelectionChange && - selectedMCPServerIds && - mcpServers?.some((server) => server.id === serverID && server.enabled) && - !selectedMCPServerIds.includes(serverID) - ) { - onMCPSelectionChange([...selectedMCPServerIds, serverID]); - } - mcpPopupRef.current = null; - }); + const handleMCPAuthComplete = useEffectEvent( + (serverID: string, source: MessageEventSource | null) => { + onMCPAuthComplete?.(serverID); + // Only the popup this input opened expresses intent to use the + // server; a stray same-origin message must not clear an in-flight + // connect or change the selection. + if (source === null || source !== mcpPopupRef.current) { + return; + } + const isInitiatedServer = mcpConnectingId === serverID; + setMcpConnectingId(null); + mcpPopupRef.current = null; + if ( + isInitiatedServer && + onMCPSelectionChange && + selectedMCPServerIds && + mcpServers?.some( + (server) => server.id === serverID && server.enabled, + ) && + !selectedMCPServerIds.includes(serverID) + ) { + onMCPSelectionChange([...selectedMCPServerIds, serverID]); + } + }, + ); // Listen for OAuth2 completion postMessage from popup. useEffect(() => { @@ -530,7 +542,7 @@ export const AgentChatInput: FC = ({ event.data?.type === "mcp-oauth2-complete" && typeof event.data.serverID === "string" ) { - handleMCPAuthComplete(event.data.serverID); + handleMCPAuthComplete(event.data.serverID, event.source); } }; window.addEventListener("message", handler); From e0e0b5e240c479b770ed088e2721bb8cb5e8cd0e Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Fri, 14 Aug 2026 08:00:54 +0000 Subject: [PATCH 3/7] fix(site): keep MCP OAuth flow correlation past popup close --- .../components/AgentChatInput.stories.tsx | 40 ++++++++++++++++++ .../AgentsPage/components/AgentChatInput.tsx | 42 ++++++++++--------- 2 files changed, 62 insertions(+), 20 deletions(-) diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx index bfd342cf852a4..8450aa0a75234 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -840,6 +840,46 @@ export const MCPAutoEnablesAfterOAuthCompletes: Story = { }, }; +// 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, diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 19c76023af79c..c82bebb8b514b 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -431,7 +431,12 @@ export const AgentChatInput: FC = ({ ); const [workspacePickerOpen, setWorkspacePickerOpen] = useState(false); const [mcpConnectingId, setMcpConnectingId] = useState(null); - const mcpPopupRef = useRef(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(null); const queryClient = useQueryClient(); @@ -511,17 +516,15 @@ export const AgentChatInput: FC = ({ const handleMCPAuthComplete = useEffectEvent( (serverID: string, source: MessageEventSource | null) => { onMCPAuthComplete?.(serverID); - // Only the popup this input opened expresses intent to use the - // server; a stray same-origin message must not clear an in-flight - // connect or change the selection. - if (source === null || source !== mcpPopupRef.current) { + // 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; } - const isInitiatedServer = mcpConnectingId === serverID; + mcpAuthFlowRef.current = null; setMcpConnectingId(null); - mcpPopupRef.current = null; if ( - isInitiatedServer && onMCPSelectionChange && selectedMCPServerIds && mcpServers?.some( @@ -549,20 +552,22 @@ export const AgentChatInput: FC = ({ return () => window.removeEventListener("message", handler); }, []); - // 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]); @@ -587,11 +592,8 @@ export const AgentChatInput: FC = ({ chatOrganizationId, server.id, ); - 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 = () => { From 386b8b75da467780e0155aec023a4d64b1ee55d5 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:56:25 +0000 Subject: [PATCH 4/7] refactor(site/src/pages/AgentsPage): extract MCP OAuth popup flow into useMCPOAuthFlow --- .../components/AgentChatInput.stories.tsx | 5 +- .../AgentsPage/components/AgentChatInput.tsx | 102 ++++-------------- .../pages/AgentsPage/hooks/useMCPOAuthFlow.ts | 102 ++++++++++++++++++ 3 files changed, 126 insertions(+), 83 deletions(-) create mode 100644 site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx index 8450aa0a75234..509a035d620dc 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -809,6 +809,9 @@ 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(); }, }; @@ -824,7 +827,7 @@ export const MCPAutoEnablesAfterOAuthCompletes: Story = { play: async ({ args, canvasElement }) => { await startMCPOAuthFlow(canvasElement); expect(window.open).toHaveBeenCalledWith( - `/api/experimental/mcp/servers/${githubMCP.id}/oauth2/connect`, + `/api/experimental/organizations/org-1/mcp-servers/${githubMCP.id}/oauth2/connect`, "_blank", "width=900,height=600", ); diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index c82bebb8b514b..9455a37e9920f 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -17,7 +17,6 @@ import type React from "react"; import { type FC, useEffect, - useEffectEvent, useImperativeHandle, useRef, useState, @@ -25,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"; @@ -63,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 { @@ -430,13 +429,24 @@ export const AgentChatInput: FC = ({ "main", ); const [workspacePickerOpen, setWorkspacePickerOpen] = useState(false); - const [mcpConnectingId, setMcpConnectingId] = useState(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, - ); + // Auto-select the server once its OAuth flow succeeds. + 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(null); const queryClient = useQueryClient(); @@ -513,65 +523,6 @@ export const AgentChatInput: FC = ({ [], ); - 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) => { - if (event.origin !== location.origin) return; - if ( - event.data?.type === "mcp-oauth2-complete" && - typeof event.data.serverID === "string" - ) { - handleMCPAuthComplete(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 (!mcpConnectingId || !mcpAuthFlowRef.current) return; - const interval = setInterval(() => { - if (mcpAuthFlowRef.current?.popup.closed) { - setMcpConnectingId(null); - } - }, 500); - return () => { - clearInterval(interval); - const popup = mcpAuthFlowRef.current?.popup; - if (popup && !popup.closed) { - popup.close(); - mcpAuthFlowRef.current = null; - } - }; - }, [mcpConnectingId]); - const handleMcpToggle = (serverId: string, checked: boolean) => { if (!onMCPSelectionChange || !selectedMCPServerIds) return; if (checked) { @@ -583,19 +534,6 @@ export const AgentChatInput: FC = ({ } }; - const handleMcpConnect = (server: TypesGen.MCPServerConfig) => { - if (!chatOrganizationId) { - return; - } - setMcpConnectingId(server.id); - const connectUrl = mcpServerOAuth2ConnectPath( - chatOrganizationId, - server.id, - ); - const popup = window.open(connectUrl, "_blank", "width=900,height=600"); - mcpAuthFlowRef.current = popup ? { popup, serverID: server.id } : null; - }; - const handleMcpDisconnectConfirm = () => { if (!mcpDisconnectTarget) { return; @@ -1424,7 +1362,7 @@ export const AgentChatInput: FC = ({ 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 } diff --git a/site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts b/site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts new file mode 100644 index 0000000000000..75e3512f3fbc1 --- /dev/null +++ b/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 = { + // Server whose OAuth consent popup is open. + 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( + 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); + }, + ); + + // 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" + ) { + 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 }; +}; From c079e1dbef08b337566bd0b82a39290914dec419 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:08:12 +0000 Subject: [PATCH 5/7] fix(site/src/pages/AgentsPage): drop comment restating the OAuth flow call --- site/src/pages/AgentsPage/components/AgentChatInput.tsx | 1 - 1 file changed, 1 deletion(-) diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 9455a37e9920f..cad62b611b69c 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -429,7 +429,6 @@ export const AgentChatInput: FC = ({ "main", ); const [workspacePickerOpen, setWorkspacePickerOpen] = useState(false); - // Auto-select the server once its OAuth flow succeeds. const { connectingServerId: mcpConnectingId, connect: connectMCPServer } = useMCPOAuthFlow({ organizationId: chatOrganizationId, From 1810f0d742b16a9a4f6e0b607f06a0ff7332355b Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:15:17 +0000 Subject: [PATCH 6/7] fix(site/src/pages/AgentsPage/hooks): drop comments restating hook internals --- site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts b/site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts index 75e3512f3fbc1..8c5014e90d846 100644 --- a/site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts +++ b/site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts @@ -13,7 +13,6 @@ type UseMCPOAuthFlowOptions = { }; type MCPOAuthFlow = { - // Server whose OAuth consent popup is open. connectingServerId: string | null; connect: (serverId: string) => void; }; @@ -51,7 +50,6 @@ export const useMCPOAuthFlow = ({ }, ); - // Listen for OAuth2 completion postMessage from popup. useEffect(() => { const handler = (event: MessageEvent) => { if (event.origin !== location.origin) return; From 6ed0789384d26a30ab8d76f83a15b25cc324c1be Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:22:35 +0000 Subject: [PATCH 7/7] fix(site/src/pages/AgentsPage): use JSDoc for exported hook and correct story helper precondition --- .../AgentsPage/components/AgentChatInput.stories.tsx | 4 ++-- site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts | 10 ++++++---- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx index 509a035d620dc..5a2a2a7f446a7 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -769,8 +769,8 @@ const dispatchMCPOAuthComplete = ( ); }; -// Requires window.open mocked to return `window` so the completion -// message can carry the popup as its 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); diff --git a/site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts b/site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts index 8c5014e90d846..c3d226206491f 100644 --- a/site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts +++ b/site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts @@ -17,10 +17,12 @@ type MCPOAuthFlow = { 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. +/** + * 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,