Thanks to visit codestin.com
Credit goes to github.com

Skip to content
149 changes: 149 additions & 0 deletions site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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. */
Expand Down Expand Up @@ -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();
},
};

Expand Down
74 changes: 19 additions & 55 deletions site/src/pages/AgentsPage/components/AgentChatInput.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -429,8 +429,23 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
"main",
);
const [workspacePickerOpen, setWorkspacePickerOpen] = useState(false);
const [mcpConnectingId, setMcpConnectingId] = useState<string | null>(null);
const mcpPopupRef = useRef<Window | null>(null);
const { connectingServerId: mcpConnectingId, connect: connectMCPServer } =
useMCPOAuthFlow({
organizationId: chatOrganizationId,
onAuthComplete: onMCPAuthComplete,
onFlowSuccess: (serverID) => {
if (
onMCPSelectionChange &&
selectedMCPServerIds &&
mcpServers?.some(
(server) => server.id === serverID && server.enabled,
) &&
Comment thread
ibetitsmike marked this conversation as resolved.
!selectedMCPServerIds.includes(serverID)
) {
onMCPSelectionChange([...selectedMCPServerIds, serverID]);
}
},
});
const [mcpDisconnectTarget, setMcpDisconnectTarget] =
useState<TypesGen.MCPServerConfig | null>(null);
const queryClient = useQueryClient();
Expand Down Expand Up @@ -507,41 +522,6 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
[],
);

// 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) {
Expand All @@ -553,22 +533,6 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
}
};

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;
Expand Down Expand Up @@ -1397,7 +1361,7 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
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
}
Expand Down
102 changes: 102 additions & 0 deletions site/src/pages/AgentsPage/hooks/useMCPOAuthFlow.ts
Original file line number Diff line number Diff line change
@@ -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<string | null>(
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 };
};
Loading