From 8cee56c1c2c7037816cd9441b7e0d037d6d44941 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/3] 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 2419f338036..cc8ec81fc08 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -755,6 +755,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. */ @@ -775,6 +784,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 7fddd7c95ad..06179ededaf 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, @@ -518,6 +519,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) => { @@ -526,14 +541,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 83894d4116dc8d3357a9006707c1932d45af1d94 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/3] 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 cc8ec81fc08..1a2834e5647 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -755,15 +755,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. */ @@ -790,8 +803,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([ @@ -809,8 +831,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); @@ -819,14 +845,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); @@ -835,18 +861,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 06179ededaf..8918145a540 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -519,19 +519,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(() => { @@ -541,7 +553,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 62b641092a355f5af5d207aaaa6bce4989417f84 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/3] 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 1a2834e5647..2159f73f9dc 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -825,6 +825,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 8918145a540..71f7ff8b877 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -442,7 +442,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(); @@ -522,17 +527,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( @@ -560,20 +563,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]); @@ -592,11 +597,8 @@ export const AgentChatInput: FC = ({ 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 = () => {