diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx index 2419f338036..2159f73f9dc 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -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: { diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 7fddd7c95ad..71f7ff8b877 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, @@ -441,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(); @@ -518,6 +524,30 @@ 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) => { @@ -526,29 +556,29 @@ 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, 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 = ({ 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 = () => {