diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx index 4c54454efe8..5a2a2a7f446 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/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(); }, }; diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 0b3d25ba419..cad62b611b6 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/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 = ({ "main", ); const [workspacePickerOpen, setWorkspacePickerOpen] = useState(false); - const [mcpConnectingId, setMcpConnectingId] = useState(null); - const mcpPopupRef = useRef(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(null); const queryClient = useQueryClient(); @@ -507,41 +522,6 @@ export const AgentChatInput: FC = ({ [], ); - // 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 = ({ } }; - 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 = ({ 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 00000000000..c3d22620649 --- /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 = { + 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); + }, + ); + + 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 }; +};