mcp-connection-picker.tsx19.7 KBView on GitHub 'use client';
import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
import { useEffect, useMemo, useState, type ReactNode } from 'react';
import { Loader2, Plus, Server } from 'lucide-react';
import { toast } from 'sonner';
import {
buildPolicyFromRows,
patchToolRow,
toolRowsFromProbe,
type McpProbeStatus,
type McpToolRow,
} from './mcp-tool-policy';
import { ConnectionInitial, connectionLogo } from '@/modules/agents/utils/connection-logo';
import { McpConnectionForm, type McpConnectionFormValues } from './mcp-connection-form';
import { KNOWN_MCP_PROVIDERS, type KnownMcpProvider } from './systems/types';
import { useTRPC } from '@/providers/query-provider';
import { Switch } from '@/components/ui/switch';
import { Button } from '@/components/ui/button';
import { Field } from '@/components/ui/field';
import { cn } from '@/lib/utils';
/**
* "Connect an MCP server", from inside an agent.
*
* Settings adds servers through `systems/add-mcp-step.tsx` now — the discovery
* flow, which asks the server how to authenticate before offering anything. This
* picker is the agent-side entry point: the same connection rows, plus the two
* things that are about an agent.
*
* The differences are PROPS, not forks, and stay that way so the discovery step
* can take over the connect half without disturbing the grant half:
* `rowAction` — an extra control per connected row (the agent renders Grant).
* `onConnected` — run after a successful connect (the agent grants it there).
* `grantOnly` — stamp new connections `requiresAgentGrant`.
*
* `grantOnly` is the load-bearing one: you added this server for ONE agent, so the
* others get the row in their list, ungranted, rather than the capability. A server
* added on a workspace-wide screen passes false and stays available to any agent
* that inherits.
*/
/** A connected server, keyed by the SAME name the grant layer uses server-side. */
export interface McpPickerConnection {
id: string;
/** `connection.name`, falling back to the row id exactly as the server does. */
name: string;
serverUrl: string;
providerId: string;
ownedByMe: boolean;
orgShared: boolean;
requiresAgentGrant: boolean;
}
/**
* The user's MCP connections, normalised.
*
* `connection.name` is nullable and every server-side grant check falls back to
* `row.name ?? row.id`. Anything else here — the URL, an empty string — would make
* a picker's checkbox state and the grant it writes disagree about which server
* they mean.
*/
export function useMcpPickerConnections(enabled = true) {
const trpc = useTRPC();
const { data, isLoading } = useQuery({
...trpc.integrations.listMcpConnections.queryOptions(undefined),
enabled,
});
const connections = useMemo<McpPickerConnection[]>(
() =>
(data?.connections ?? []).map((c) => ({
id: c.id,
name: c.name || c.id,
serverUrl: c.serverUrl,
providerId: c.providerId,
ownedByMe: c.isOwnedByCaller,
orgShared: c.scope === 'org',
requiresAgentGrant: c.requiresAgentGrant,
})),
[data],
);
return { connections, isLoading };
}
/** The brand logo for a server, falling back to a coloured initial. */
function ConnectionMark({ name, className }: { name: string; className?: string }) {
const Logo = connectionLogo(name.toLowerCase(), name);
if (Logo) return <Logo className={cn('h-5 w-5 shrink-0', className)} />;
return (
<ConnectionInitial
name={name}
className={cn(
'inline-flex h-5 w-5 shrink-0 items-center justify-center rounded text-[10px] font-bold leading-none text-white',
className,
)}
/>
);
}
const EMPTY_FORM: McpConnectionFormValues = {
name: '',
serverUrl: '',
authorizationHeader: '',
instructions: '',
};
/**
* Compare two MCP server URLs as the same server. Trailing slashes and case in the host
* are not identity: `https://mcp.usepylon.com/` and `https://mcp.usepylon.com` are one
* server, and treating them as two is how the catalog offers to connect something the
* user has already connected.
*/
function normalizeServerUrl(value: string): string {
return value.trim().toLowerCase().replace(/\/+$/, '');
}
/** Bare token → `Bearer <token>`; an already-scheme'd value passes through. */
function normalizeAuthorizationHeader(value: string): string {
const trimmed = value.trim();
if (!trimmed) return '';
return /^[A-Za-z]+\s+/.test(trimmed) ? trimmed : `Bearer ${trimmed}`;
}
export interface McpConnectionPickerProps {
connections: McpPickerConnection[];
isLoading?: boolean;
/** See the module note — true from an agent, false from Settings. */
grantOnly?: boolean;
/** Extra control on the right of a connected row. The agent renders Grant here. */
rowAction?: (connection: McpPickerConnection) => ReactNode;
/** Runs after a successful connect, before the toast. The agent grants here. */
onConnected?: (connection: { id: string; name: string }) => Promise<void>;
/** Called after any write, so the caller can invalidate what it owns. */
onChanged?: () => void;
/**
* Which half of the picker is showing.
*
* Reported UP rather than titled here, because the container is a dialog and a
* dialog already has a header: a heading inside this component put a second title
* under the first one, which is a container inside a container. The custom view
* switches from a LIST of connections to ONE form, so the dialog's own title has
* to change with it — and that is the only thing the caller needs to know.
*/
onViewChange?: (view: 'catalog' | 'custom') => void;
className?: string;
}
export function McpConnectionPicker({
connections,
isLoading,
grantOnly = false,
rowAction,
onConnected,
onChanged,
onViewChange,
className,
}: McpConnectionPickerProps) {
const trpc = useTRPC();
const queryClient = useQueryClient();
const [view, setView] = useState<'catalog' | 'custom'>('catalog');
const [form, setForm] = useState<McpConnectionFormValues>(EMPTY_FORM);
const [probeStatus, setProbeStatus] = useState<McpProbeStatus>('idle');
const [probeError, setProbeError] = useState<string | null>(null);
const [toolRows, setToolRows] = useState<McpToolRow[]>([]);
const [shareWithOrg, setShareWithOrg] = useState(false);
const [connecting, setConnecting] = useState<string | null>(null);
const { mutateAsync: addConnection, isPending: isAdding } = useMutation(
trpc.integrations.addMcpConnection.mutationOptions(),
);
const { mutateAsync: setToolPolicy } = useMutation(
trpc.integrations.setMcpToolPolicy.mutationOptions(),
);
const { mutateAsync: setScope } = useMutation(
trpc.integrations.setMcpConnectionScope.mutationOptions(),
);
// Two probes, two questions. `probeServerTools` enumerates what a server EXPOSES
// and needs a credential; `probeServerAuth` asks the server how to AUTHENTICATE to
// it and is what resolves a catalog entry to a shipped OAuth driver.
const { mutateAsync: probeServerTools } = useMutation(
trpc.integrations.probeMcpServerTools.mutationOptions(),
);
const { mutateAsync: probeServerAuth } = useMutation(
trpc.integrations.probeMcpServer.mutationOptions(),
);
const { mutateAsync: initiateOAuth } = useMutation(
trpc.integrations.initiateOAuth.mutationOptions(),
);
useEffect(() => {
onViewChange?.(view);
}, [view, onViewChange]);
const refresh = () => {
void queryClient.invalidateQueries({
queryKey=[redacted],
});
void queryClient.invalidateQueries({ queryKey=[redacted] });
onChanged?.();
};
/**
* Connected servers, then the known providers not yet connected.
*
* An already-connected provider appears ONCE, as its connected row — listing the
* catalog and the connections separately is exactly how someone connects Notion
* a second time. Matched on SERVER URL rather than on a provider id: the catalog
* (`KNOWN_MCP_PROVIDERS`) is a list of URLs, and a URL with no shipped driver is
* stored under the generic `mcp` provider id, which every such row would share.
*/
const connectable = useMemo(() => {
const have = new Set(connections.map((c) => normalizeServerUrl(c.serverUrl)));
return KNOWN_MCP_PROVIDERS.filter((p) => !have.has(normalizeServerUrl(p.serverUrl)));
}, [connections]);
/**
* Connect a catalog entry.
*
* The catalog stores a URL, not a provider id, because whether Cedar ships an OAuth
* driver for a server is a property of the SERVER and is answered by discovery. So
* the id is resolved at click time: a driver-backed server goes straight to the OAuth
* popup, and anything else drops into the custom form pre-filled, which is where the
* access-token and manual-client paths already live.
*/
const connectKnown = async (provider: KnownMcpProvider) => {
setConnecting(provider.serverUrl);
try {
const probe = await probeServerAuth({ serverUrl: provider.serverUrl });
if (probe.providerId) {
setConnecting(null);
await connectOAuth(probe.providerId, provider.name);
return;
}
setForm({ ...EMPTY_FORM, name: provider.name, serverUrl: provider.serverUrl });
setView('custom');
} catch (err) {
toast.error(err instanceof Error ? err.message : `Could not reach ${provider.name}`);
} finally {
setConnecting(null);
}
};
const resetCustom = () => {
setForm(EMPTY_FORM);
setProbeStatus('idle');
setProbeError(null);
setToolRows([]);
setView('catalog');
};
const handleProbe = async () => {
const serverUrl = form.serverUrl.trim();
if (!serverUrl) return;
setProbeStatus('probing');
setProbeError(null);
try {
const result = await probeServerTools({
serverUrl,
...(form.authorizationHeader.trim()
? { authorizationHeader: normalizeAuthorizationHeader(form.authorizationHeader) }
: {}),
});
if (result.success) {
setToolRows(toolRowsFromProbe(result.tools));
setProbeStatus('ready');
} else {
setToolRows([]);
setProbeError(result.error ?? 'The server did not respond to tools/list.');
setProbeStatus('failed');
}
} catch (err) {
setToolRows([]);
setProbeError(err instanceof Error ? err.message : 'Failed to reach the server');
setProbeStatus('failed');
}
};
const submitCustom = async () => {
const name = form.name.trim();
if (!name || !form.serverUrl.trim()) {
toast.error('Name and Server URL are required');
return;
}
try {
const created = await addConnection({
name,
serverUrl: form.serverUrl.trim(),
...(form.authorizationHeader.trim()
? { authorizationHeader: normalizeAuthorizationHeader(form.authorizationHeader) }
: {}),
...(form.instructions.trim() ? { instructions: form.instructions.trim() } : {}),
...(grantOnly ? { requiresAgentGrant: true } : {}),
...(shareWithOrg ? { orgShared: true } : {}),
});
if (created?.connection?.id && toolRows.length > 0) {
try {
await setToolPolicy({
connectionId: created.connection.id,
policy: buildPolicyFromRows(toolRows, {
discoveredAt: created.toolPolicy?.discoveredAt,
}),
});
} catch (policyError) {
// The connection exists and is inert (deny-by-default). Say which half
// failed rather than leaving the user unsure what saved.
toast.error(
policyError instanceof Error
? `Connected, but permissions did not save: ${policyError.message}`
: 'Connected, but permissions did not save',
);
}
}
if (created?.connection?.id) await onConnected?.({ id: created.connection.id, name });
resetCustom();
refresh();
toast.success(`${name} connected`);
} catch (err) {
toast.error(err instanceof Error ? err.message : 'Failed to add connection');
}
};
/**
* OAuth providers are connected by a popup whose CALLBACK writes the row, so the
* scope stamp and any grant can only land once it reports back.
*/
const connectOAuth = async (providerId: string, providerName: string) => {
setConnecting(providerId);
const onMessage = async (event: MessageEvent) => {
if (event.data?.type !== 'mcp_oauth_connected' || event.data?.provider !== providerId) return;
window.removeEventListener('message', onMessage);
setConnecting(null);
if (!event.data.success) {
toast.error(`Failed to connect ${providerName}: ${event.data.error ?? 'Unknown error'}`);
return;
}
try {
// `staleTime: 0` is load-bearing. The provider sets a 1-minute global
// staleTime and this list was fetched seconds ago, so a plain fetchQuery
// would hand back the cache from BEFORE the popup wrote the row, find no
// connection, and silently skip the scope stamp and the grant.
const fresh = await queryClient.fetchQuery({
...trpc.integrations.listMcpConnections.queryOptions(undefined),
staleTime: 0,
});
const row = fresh.connections.find((c) => c.providerId === providerId);
if (row) {
if (grantOnly || shareWithOrg) {
await setScope({
connectionId: row.id,
...(grantOnly ? { requiresAgentGrant: true } : {}),
...(shareWithOrg ? { orgShared: true } : {}),
});
}
// `row.name || row.id` — the same fallback the grant layer uses server-side.
// The provider's display name is not the connection's identity.
await onConnected?.({ id: row.id, name: row.name || row.id });
toast.success(`${providerName} connected`);
}
} catch (err) {
toast.error(
err instanceof Error
? `Connected, but the settings did not save: ${err.message}`
: 'Connected, but the settings did not save',
);
}
refresh();
};
window.addEventListener('message', onMessage);
try {
const { oauthUrl } = await initiateOAuth({ integration: providerId });
if (!oauthUrl) {
window.removeEventListener('message', onMessage);
setConnecting(null);
toast.error(`Failed to get ${providerName} OAuth URL`);
return;
}
window.open(oauthUrl, '_blank', 'width=600,height=700');
} catch (err) {
window.removeEventListener('message', onMessage);
setConnecting(null);
toast.error(err instanceof Error ? err.message : `Failed to connect ${providerName}`);
}
};
const toggleOrgShare = async (connectionId: string, on: boolean) => {
try {
await setScope({ connectionId, orgShared: on });
refresh();
} catch (err) {
toast.error(err instanceof Error ? err.message : 'Could not change sharing');
}
};
if (view === 'custom') {
return (
// No back link. Cancel already IS back — both call `resetCustom` — and two
// controls for one act is one of them you have to think about.
<div className={className}>
<McpConnectionForm
mode="add"
orgShare={
<Field
label="Share this server with my organization"
hint="Teammates can grant it to their own agents."
>
<Switch
checked={shareWithOrg}
onCheckedChange={setShareWithOrg}
aria-label="Share this server with my organization"
/>
</Field>
}
values={form}
onChange={(patch) => setForm((f) => ({ ...f, ...patch }))}
probeStatus={probeStatus}
probeError={probeError}
tools={toolRows}
onProbe={handleProbe}
onToggleTool={(toolName, allowed) =>
setToolRows((rows) => patchToolRow(rows, toolName, { allowed }))
}
onToolInstructionChange={(toolName, instruction) =>
setToolRows((rows) => patchToolRow(rows, toolName, { instruction }))
}
onSubmit={submitCustom}
onCancel={resetCustom}
isSaving={isAdding}
/>
</div>
);
}
return (
<div className={cn('flex flex-col px-4', className)}>
{/* No dividers. Rows in one list are separated by space and their own hover
surface; a rule between each pair turns a picker into a table. */}
<div className="flex flex-col">
{isLoading && (
<div className="text-muted-foreground flex items-center gap-2 py-6 text-sm">
<Loader2 className="h-3.5 w-3.5 animate-spin" /> Loading connections…
</div>
)}
{connections.map((c) => (
<div key=[redacted] className="-mx-2 flex items-start gap-3 rounded-md px-2 py-3">
<ConnectionMark name={c.name} className="mt-0.5" />
<div className="min-w-0 flex-1">
<div className="flex items-center gap-2">
<p className="truncate text-sm font-medium">{c.name}</p>
{!c.ownedByMe && (
<span className="text-muted-foreground shrink-0 text-[11px]">
shared with your org
</span>
)}
</div>
<p className="text-muted-foreground truncate text-xs">{c.serverUrl}</p>
{/* Only the OWNER may unshare: a teammate revoking a server for the
whole org from a screen that never said it was theirs is not a
control, it is an accident waiting to happen. */}
{c.ownedByMe && (
<label className="mt-1.5 flex w-fit cursor-pointer items-center gap-2">
<Switch
checked={c.orgShared}
onCheckedChange={(on) => void toggleOrgShare(c.id, on)}
className="origin-left scale-[0.7] cursor-pointer"
aria-label={`Share ${c.name} with the organization`}
/>
<span className="text-muted-foreground text-[11px]">Share with organization</span>
</label>
)}
</div>
{rowAction?.(c)}
</div>
))}
{connectable.map((p) => (
<div key=[redacted] className="-mx-2 flex items-start gap-3 rounded-md px-2 py-3">
<ConnectionMark name={p.name} className="mt-0.5" />
<div className="min-w-0 flex-1">
<p className="text-sm font-medium">{p.name}</p>
<p className="text-muted-foreground truncate text-xs">{p.serverUrl}</p>
</div>
<Button
type="button"
size="sm"
variant="outline"
className="h-7 shrink-0 cursor-pointer px-2.5 text-xs"
disabled={connecting === p.serverUrl}
onClick={() => void connectKnown(p)}
>
{connecting === p.serverUrl && <Loader2 className="mr-1.5 h-3 w-3 animate-spin" />}
Connect
</Button>
</div>
))}
<button
type="button"
onClick={() => setView('custom')}
className="hover:bg-sunken -mx-2 flex w-[calc(100%+1rem)] cursor-pointer items-center gap-3 rounded-md px-2 py-3 text-left transition-colors"
>
<span className="bg-muted flex h-5 w-5 shrink-0 items-center justify-center rounded">
<Server className="h-3 w-3" />
</span>
<span className="min-w-0 flex-1">
<span className="block text-sm font-medium">Custom MCP server</span>
<span className="text-muted-foreground block text-xs">
Any MCP-speaking server, by URL and token.
</span>
</span>
<Plus className="text-muted-foreground h-4 w-4 shrink-0" />
</button>
</div>
</div>
);
}