diff --git a/coderd/mcp_test.go b/coderd/mcp_test.go index c3e77212072..5ef5709f700 100644 --- a/coderd/mcp_test.go +++ b/coderd/mcp_test.go @@ -526,9 +526,14 @@ func TestMCPServerConfigsOAuth2Disconnect(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - adminClient := newMCPClient(t) + providerKeys := coderdtest.FakeOpenAICompatProviderAPIKeys(t) + adminClient, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{ + DeploymentValues: mcpDeploymentValues(t), + ChatProviderAPIKeys: &providerKeys, + }) firstUser := coderdtest.CreateFirstUser(t, adminClient) - memberClient, _ := coderdtest.CreateAnotherUser(t, adminClient, firstUser.OrganizationID) + memberClient, member := coderdtest.CreateAnotherUser(t, adminClient, firstUser.OrganizationID) + otherClient, other := coderdtest.CreateAnotherUser(t, adminClient, firstUser.OrganizationID) created, err := adminClient.CreateMCPServerConfig(ctx, codersdk.CreateMCPServerConfigRequest{ DisplayName: "OAuth Disconnect Test", @@ -549,6 +554,36 @@ func TestMCPServerConfigsOAuth2Disconnect(t *testing.T) { // Disconnect should succeed even when no token exists (idempotent). err = memberClient.MCPServerOAuth2Disconnect(ctx, created.ID) require.NoError(t, err) + + for _, userID := range []uuid.UUID{member.ID, other.ID} { + //nolint:gocritic // Seeding test state requires system access. + _, err = db.UpsertMCPServerUserToken(dbauthz.AsSystemRestricted(ctx), database.UpsertMCPServerUserTokenParams{ + MCPServerConfigID: created.ID, + UserID: userID, + AccessToken: "valid-access", + TokenType: "Bearer", + Expiry: sql.NullTime{Time: time.Now().Add(time.Hour), Valid: true}, + }) + require.NoError(t, err) + } + + requireAuthConnected := func(client *codersdk.Client, want bool) { + t.Helper() + configs, err := client.MCPServerConfigs(ctx) + require.NoError(t, err) + require.Len(t, configs, 1) + require.Equal(t, want, configs[0].AuthConnected) + } + requireAuthConnected(memberClient, true) + requireAuthConnected(otherClient, true) + + err = memberClient.MCPServerOAuth2Disconnect(ctx, created.ID) + require.NoError(t, err) + requireAuthConnected(memberClient, false) + requireAuthConnected(otherClient, true) + + err = memberClient.MCPServerOAuth2Disconnect(ctx, created.ID) + require.NoError(t, err) } func TestMCPServerConfigsOAuth2AutoDiscovery(t *testing.T) { diff --git a/site/src/api/api.ts b/site/src/api/api.ts index b4efff76f01..e89a8f1994a 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -3976,6 +3976,12 @@ class ExperimentalApiMethods { ); }; + disconnectMCPServerOAuth2 = async (id: string): Promise => { + await this.axios.delete( + `${mcpServerConfigsPath}/${encodeURIComponent(id)}/oauth2/disconnect`, + ); + }; + getChatCostSummary = async ( user = "me", params?: ChatCostDateParams, diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 8262c8f4918..281f04234b1 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -2070,6 +2070,13 @@ export const deleteMCPServerConfig = (queryClient: QueryClient) => ({ }, }); +export const disconnectMCPServerOAuth2 = (queryClient: QueryClient) => ({ + mutationFn: (id: string) => API.experimental.disconnectMCPServerOAuth2(id), + onSuccess: async () => { + await invalidateMCPServerConfigQueries(queryClient); + }, +}); + type SetChatUserRoleVariables = { chatId: string; userId: string; diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx index 9f010a948eb..ed58e95e641 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -1,7 +1,8 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; import { MonitorDotIcon } from "lucide-react"; import { useEffect, useRef } from "react"; -import { expect, fn, userEvent, waitFor, within } from "storybook/test"; +import { expect, fn, spyOn, userEvent, waitFor, within } from "storybook/test"; +import { API } from "#/api/api"; import type * as TypesGen from "#/api/typesGenerated"; import { MockChatContextClean, @@ -739,6 +740,16 @@ const githubMCP = buildMCPServer({ const githubMCPConnected = { ...githubMCP, auth_connected: true }; +const notionMCPConnected = buildMCPServer({ + id: "mcp-notion", + display_name: "Notion", + slug: "notion", + availability: "default_on", + auth_type: "oauth2", + auth_connected: true, + enabled: true, +}); + const mcpDefaults = { onMCPSelectionChange: fn(), onMCPAuthComplete: fn(), @@ -800,6 +811,109 @@ export const PlusMenuOpen: Story = { }, }; +export const MCPDisconnectControls: Story = { + args: { + ...mcpDefaults, + mcpServers: [linearMCP, githubMCP, notionMCPConnected], + selectedMCPServerIds: [linearMCP.id, notionMCPConnected.id], + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const body = within(canvasElement.ownerDocument.body); + await userEvent.click(canvas.getByRole("button", { name: "More options" })); + expect( + await body.findByRole("button", { name: "Disconnect Notion" }), + ).toBeInTheDocument(); + expect( + body.queryByRole("button", { name: "Disconnect GitHub" }), + ).not.toBeInTheDocument(); + expect(body.getByRole("button", { name: "Auth" })).toBeInTheDocument(); + expect( + body.queryByRole("button", { name: "Disconnect Linear" }), + ).not.toBeInTheDocument(); + }, +}; + +export const MCPDisconnectCancel: Story = { + args: { + ...mcpDefaults, + mcpServers: [githubMCPConnected], + selectedMCPServerIds: [githubMCPConnected.id], + }, + beforeEach: () => { + spyOn(API.experimental, "disconnectMCPServerOAuth2").mockResolvedValue(); + }, + play: async ({ canvasElement }) => { + 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: "Disconnect GitHub" }), + ); + expect(await body.findByText("Disconnect GitHub?")).toBeInTheDocument(); + await userEvent.click(body.getByRole("button", { name: "Cancel" })); + await waitFor(() => + expect(body.queryByText("Disconnect GitHub?")).not.toBeInTheDocument(), + ); + expect(API.experimental.disconnectMCPServerOAuth2).not.toHaveBeenCalled(); + }, +}; + +export const MCPDisconnectConfirm: Story = { + args: { + ...mcpDefaults, + mcpServers: [githubMCPConnected], + selectedMCPServerIds: [githubMCPConnected.id], + }, + beforeEach: () => { + spyOn(API.experimental, "disconnectMCPServerOAuth2").mockResolvedValue(); + }, + play: async ({ canvasElement }) => { + 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: "Disconnect GitHub" }), + ); + await body.findByText("Disconnect GitHub?"); + await userEvent.click(body.getByRole("button", { name: "Disconnect" })); + await waitFor(() => + expect(body.queryByText("Disconnect GitHub?")).not.toBeInTheDocument(), + ); + expect(API.experimental.disconnectMCPServerOAuth2).toHaveBeenCalledTimes(1); + expect(API.experimental.disconnectMCPServerOAuth2).toHaveBeenCalledWith( + githubMCPConnected.id, + ); + }, +}; + +export const MCPDisconnectError: Story = { + args: { + ...mcpDefaults, + mcpServers: [githubMCPConnected], + selectedMCPServerIds: [githubMCPConnected.id], + }, + beforeEach: () => { + spyOn(API.experimental, "disconnectMCPServerOAuth2").mockRejectedValue( + new Error("disconnect failed"), + ); + }, + play: async ({ canvasElement }) => { + 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: "Disconnect GitHub" }), + ); + await body.findByText("Disconnect GitHub?"); + await userEvent.click(body.getByRole("button", { name: "Disconnect" })); + await waitFor(() => + expect(API.experimental.disconnectMCPServerOAuth2).toHaveBeenCalled(), + ); + expect(body.getByText("Disconnect GitHub?")).toBeInTheDocument(); + }, +}; + export const PlanFirstMenuItem: Story = { args: { onPlanModeToggle: fn(), diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 98675297337..335434dacfb 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -10,6 +10,7 @@ import { PlusIcon, ServerIcon, SquareIcon, + UnlinkIcon, XIcon, } from "lucide-react"; import type React from "react"; @@ -20,7 +21,11 @@ import { useRef, useState, } from "react"; +import { useMutation, useQueryClient } from "react-query"; import { Link } from "react-router"; +import { toast } from "sonner"; +import { getErrorMessage } from "#/api/errors"; +import { disconnectMCPServerOAuth2 } from "#/api/queries/chats"; import type * as TypesGen from "#/api/typesGenerated"; import type { AgentChatSendShortcut, @@ -37,6 +42,7 @@ import { CommandItem, CommandList, } from "#/components/Command/Command"; +import { ConfirmDialog } from "#/components/Dialogs/ConfirmDialog/ConfirmDialog"; import { ExternalImage } from "#/components/ExternalImage/ExternalImage"; import { Popover, @@ -428,6 +434,12 @@ export const AgentChatInput: FC = ({ const [workspacePickerOpen, setWorkspacePickerOpen] = useState(false); const [mcpConnectingId, setMcpConnectingId] = useState(null); const mcpPopupRef = useRef(null); + const [mcpDisconnectTarget, setMcpDisconnectTarget] = + useState(null); + const queryClient = useQueryClient(); + const mcpDisconnectMutation = useMutation( + disconnectMCPServerOAuth2(queryClient), + ); const [hasFileReferences, setHasFileReferences] = useState(false); const [cycleIndex, setCycleIndex] = useState(null); @@ -554,6 +566,22 @@ export const AgentChatInput: FC = ({ ); }; + const handleMcpDisconnectConfirm = () => { + if (!mcpDisconnectTarget) { + return; + } + const name = mcpDisconnectTarget.display_name; + mcpDisconnectMutation.mutate(mcpDisconnectTarget.id, { + onSuccess: () => { + setMcpDisconnectTarget(null); + toast.success(`Disconnected ${name}.`); + }, + onError: (error) => { + toast.error(getErrorMessage(error, `Failed to disconnect ${name}.`)); + }, + }); + }; + const selectedWorkspace = workspaceOptions?.find( (ws) => ws.id === selectedWorkspaceId, ); @@ -1400,15 +1428,32 @@ export const AgentChatInput: FC = ({ Auth ) : ( - - handleMcpToggle(server.id, checked) - } - disabled={isDisabled || isForceOn} - aria-label={`${isSelected ? "Disable" : "Enable"} ${server.display_name}`} - /> + <> + {server.auth_type === "oauth2" && ( + + )} + + handleMcpToggle(server.id, checked) + } + disabled={isDisabled || isForceOn} + aria-label={`${isSelected ? "Disable" : "Enable"} ${server.display_name}`} + /> + )} ); @@ -1648,6 +1693,16 @@ export const AgentChatInput: FC = ({ }} /> )} + setMcpDisconnectTarget(null)} + /> ); };