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

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 37 additions & 2 deletions coderd/mcp_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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) {
Expand Down
6 changes: 6 additions & 0 deletions site/src/api/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3976,6 +3976,12 @@ class ExperimentalApiMethods {
);
};

disconnectMCPServerOAuth2 = async (id: string): Promise<void> => {
await this.axios.delete(
`${mcpServerConfigsPath}/${encodeURIComponent(id)}/oauth2/disconnect`,
);
};

getChatCostSummary = async (
user = "me",
params?: ChatCostDateParams,
Expand Down
7 changes: 7 additions & 0 deletions site/src/api/queries/chats.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
116 changes: 115 additions & 1 deletion site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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(),
Expand Down
73 changes: 64 additions & 9 deletions site/src/pages/AgentsPage/components/AgentChatInput.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import {
PlusIcon,
ServerIcon,
SquareIcon,
UnlinkIcon,
XIcon,
} from "lucide-react";
import type React from "react";
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -428,6 +434,12 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
const [workspacePickerOpen, setWorkspacePickerOpen] = useState(false);
const [mcpConnectingId, setMcpConnectingId] = useState<string | null>(null);
const mcpPopupRef = useRef<Window | null>(null);
const [mcpDisconnectTarget, setMcpDisconnectTarget] =
useState<TypesGen.MCPServerConfig | null>(null);
const queryClient = useQueryClient();
const mcpDisconnectMutation = useMutation(
disconnectMCPServerOAuth2(queryClient),
);

const [hasFileReferences, setHasFileReferences] = useState(false);
const [cycleIndex, setCycleIndex] = useState<number | null>(null);
Expand Down Expand Up @@ -554,6 +566,22 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
);
};

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,
);
Expand Down Expand Up @@ -1400,15 +1428,32 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
Auth
</Button>
) : (
<Switch
size="sm"
checked={isSelected}
onCheckedChange={(checked) =>
handleMcpToggle(server.id, checked)
}
disabled={isDisabled || isForceOn}
aria-label={`${isSelected ? "Disable" : "Enable"} ${server.display_name}`}
/>
<>
{server.auth_type === "oauth2" && (
<Button
variant="subtle"
size="icon"
className="size-6 shrink-0 text-content-secondary [&>svg]:size-3"
onClick={() => {
setPlusMenuOpen(false);
setMcpDisconnectTarget(server);
}}
disabled={isDisabled}
aria-label={`Disconnect ${server.display_name}`}
>
<UnlinkIcon />
</Button>
)}
<Switch
size="sm"
checked={isSelected}
onCheckedChange={(checked) =>
handleMcpToggle(server.id, checked)
}
disabled={isDisabled || isForceOn}
aria-label={`${isSelected ? "Disable" : "Enable"} ${server.display_name}`}
/>
</>
)}
</div>
);
Expand Down Expand Up @@ -1648,6 +1693,16 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
}}
/>
)}
<ConfirmDialog
open={mcpDisconnectTarget !== null}
title={`Disconnect ${mcpDisconnectTarget?.display_name ?? "MCP server"}?`}
description="This removes your credentials for this MCP server from Coder. You can authenticate again later."
type="delete"
confirmText="Disconnect"
confirmLoading={mcpDisconnectMutation.isPending}
onConfirm={handleMcpDisconnectConfirm}
onClose={() => setMcpDisconnectTarget(null)}
/>
</>
);
};
Expand Down
Loading