add-mcp-step.tsx26.9 KBView on GitHub
'use client';

import {
  hostOf,
  INSTRUCTIONS_HELP,
  isSameMcpServerUrl,
  isTrustedOAuthMessageOrigin,
  KNOWN_MCP_PROVIDERS,
} from './types';
import { AlertTriangle, Check, Copy, Loader2, Server, ShieldCheck } from 'lucide-react';
import { useMutation, useQueryClient } from '@tanstack/react-query';
import type { KnownMcpProvider, McpProbeResult } from './types';
import { useTRPC } from '@/providers/query-provider';
import { Textarea } from '@/components/ui/textarea';
import { useEffect, useRef, useState } from 'react';
import { Button } from '@/components/ui/button';
import { Label } from '@/components/ui/label';
import { Input } from '@/components/ui/input';
import { ScopeSelect } from './scope-select';
import { cn } from '@/lib/utils';
import { toast } from 'sonner';

/** Turn a bare token into a header value, leaving an already-qualified one alone. */
function toAuthorizationHeader(value: string): string {
  const trimmed = value.trim();
  if (!trimmed) return '';
  return /^[A-Za-z]+\s+/.test(trimmed) ? trimmed : `Bearer ${trimmed}`;
}

/** A vault entry name from a host: the schema allows letters, digits, dots, dashes. */
function vaultNameFromUrl(serverUrl: string): string {
  const host = hostOf(serverUrl).replace(/[^a-zA-Z0-9._-]/g, '-');
  return host ? `${host}-oauth-app` : '';
}

/**
 * Step 2 of the add dialog: connect an MCP server.
 *
 * The server is asked what connecting to it takes (`probeMcpServer`) before the user
 * is asked for anything, because the four answers need four different forms and
 * guessing wrong means an opaque failure at the authorization server. Nothing here
 * is stored until the user submits.
 */
export function AddMcpStep({
  userId,
  isOrgAdmin,
  defaultScope = 'user',
  onOAuthOwnershipChange,
  onDone,
}: {
  userId?: string;
  isOrgAdmin: boolean;
  defaultScope?: 'user' | 'org';
  /**
   * Called with `true` while a sign-in STARTED HERE is in flight, and `false` once it
   * ends, including on unmount.
   *
   * This step and the section behind it both listen for the OAuth callback's
   * postMessage, and exactly one of them may act on it. Ownership follows the flow
   * rather than the dialog: the dialog is also open for a credential form and for an
   * edit, and standing the section down for the whole time it is open swallowed the
   * completion of a Reconnect the section itself had started.
   */
  onOAuthOwnershipChange?: (owned: boolean) => void;
  onDone: () => void;
}) {
  const trpc = useTRPC();
  const queryClient = useQueryClient();

  const [serverUrl, setServerUrl] = useState('');
  /**
   * Carried by the two paths whose write has an org variant: OAuth sign-in
   * (`initiateOAuth` takes a scope) and the manually registered client app, which is
   * stored in the credential vault and so can go through `upsertOrg`. A static-token
   * connection is stored by `addMcpConnection`, which writes a user-scoped row and has
   * no org variant, so its form says so in a line of copy instead of offering a choice
   * the write cannot honour.
   */
  const [scope, setScope] = useState<'user' | 'org'>(isOrgAdmin ? defaultScope : 'user');
  const [probe, setProbe] = useState<McpProbeResult | null>(null);
  const [name, setName] = useState('');
  const [token, setToken] = useState('');
  const [clientId, setClientId] = useState('');
  const [clientSecret, setClientSecret] = useState('');
  const [instructions, setInstructions] = useState('');
  const [copied, setCopied] = useState(false);
  const [isConnecting, setIsConnecting] = useState(false);

  // Read inside the postMessage handler, which is registered once and must not go
  // stale on the instructions the user is still typing when the popup returns. Written
  // in an effect, not during render: a render can be thrown away or replayed under
  // concurrent rendering, and a ref written there would record a value never committed.
  const instructionsRef = useRef(instructions);
  useEffect(() => {
    instructionsRef.current = instructions;
  }, [instructions]);

  /**
   * The sign-in popup, watched so the Connect button comes back when the user simply
   * closes it. Nothing is posted in that case, so waiting for the callback message left
   * the button disabled for as long as the dialog stayed open, which reads as Cedar
   * having hung.
   */
  const popupRef = useRef<Window | null>(null);
  const popupPollRef = useRef<ReturnType<typeof setInterval> | null>(null);

  /**
   * The server URL the in-flight sign-in was started for, as the PROBE reported it, and
   * so also the marker for "this step has a sign-in in flight".
   *
   * Not the `serverUrl` state: the field stays editable while the popup is open, so
   * reading state when the callback returns can look for a server the user has since
   * typed over. This is frozen at the moment the popup opened, which is the URL the
   * authorization actually belongs to.
   */
  const connectingServerUrlRef = useRef<string | null>(null);

  /**
   * The last name this form filled in by itself, so a name the USER typed is not thrown
   * away when the URL changes. Anything else in the name field was typed deliberately.
   */
  const autoFilledNameRef = useRef('');

  // Read through a ref, like the instructions above, so the popup bookkeeping below stays
  // free of reactive dependencies and a caller passing a fresh function each render
  // cannot tear down a live sign-in by changing this prop's identity.
  const ownershipChangeRef = useRef(onOAuthOwnershipChange);
  useEffect(() => {
    ownershipChangeRef.current = onOAuthOwnershipChange;
  }, [onOAuthOwnershipChange]);

  const stopWatchingPopup = () => {
    if (popupPollRef.current !== null) {
      clearInterval(popupPollRef.current);
      popupPollRef.current = null;
    }
    popupRef.current = null;
    connectingServerUrlRef.current = null;
    ownershipChangeRef.current?.(false);
  };

  // Closing the dialog unmounts this step, so the interval and any ownership claim
  // have to go with it.
  useEffect(() => stopWatchingPopup, []);

  const { mutateAsync: probeServer, isPending: isProbing } = useMutation(
    trpc.integrations.probeMcpServer.mutationOptions(),
  );
  const { mutateAsync: addConnection, isPending: isAdding } = useMutation(
    trpc.integrations.addMcpConnection.mutationOptions(),
  );
  const { mutateAsync: updateSettings } = useMutation(
    trpc.integrations.updateMcpConnectionSettings.mutationOptions(),
  );
  const { mutateAsync: initiateOAuth } = useMutation(
    trpc.integrations.initiateOAuth.mutationOptions(),
  );
  const { mutateAsync: upsertCredential, isPending: isSavingUserCredential } = useMutation(
    trpc.credentialVault.upsert.mutationOptions(),
  );
  const { mutateAsync: upsertOrgCredential, isPending: isSavingOrgCredential } = useMutation(
    trpc.credentialVault.upsertOrg.mutationOptions(),
  );
  const isSavingCredential = isSavingUserCredential || isSavingOrgCredential;

  const mcpListOptions = trpc.integrations.listMcpConnections.queryOptions(
    userId ? { userId } : undefined,
  );

  const invalidate = () => {
    void queryClient.invalidateQueries({ queryKey=[redacted] });
    void queryClient.invalidateQueries({ queryKey=[redacted] });
    void queryClient.invalidateQueries({ queryKey=[redacted] });
  };

  /**
   * The OAuth callback writes its own connection row, so the instructions the user
   * typed here are applied to it afterwards. `updateMcpConnectionSettings` merges,
   * so this cannot disturb the token set the callback just stored.
   */
  useEffect(() => {
    const handleMessage = async (event: MessageEvent) => {
      // `data.error` is rendered straight into a toast, so an unchecked listener puts
      // text from any page the user has open on Cedar's screen.
      if (!isTrustedOAuthMessageOrigin(event.origin)) return;
      const data = event.data as { type?: string; success?: boolean; error?: string } | undefined;
      if (data?.type !== 'mcp_oauth_connected') return;
      const target = connectingServerUrlRef.current;
      // Not this step's sign-in: it has none in flight. A Reconnect started from the list
      // behind this dialog is the section's message to finish.
      if (!target) return;
      stopWatchingPopup();
      setIsConnecting(false);

      if (!data.success) {
        toast.error(`Could not connect: ${data.error ?? 'Unknown error'}`);
        return;
      }

      try {
        const fresh = await queryClient.fetchQuery(mcpListOptions);
        // Compared as MCP server URLs, not as strings. The callback stores the URL its
        // driver carries, which for a branded server is not character-for-character what
        // the user pasted, so a string compare finds no row and drops the description.
        const created = fresh.connections.find((c) => isSameMcpServerUrl(c.serverUrl, target));
        if (!created) {
          // The connection itself exists; only the description failed to land. Saying so
          // matters because an entry with no description is one no agent reaches for, and
          // this form made the description mandatory precisely for that reason.
          toast.error(
            "Connected, but the description was not saved. Add it from the entry's Edit.",
          );
        } else if (instructionsRef.current.trim()) {
          await updateSettings({
            connectionId: created.id,
            instructions: instructionsRef.current.trim(),
          });
        }
      } catch (err) {
        toast.error(
          err instanceof Error
            ? `Connected, but the description was not saved: ${err.message}`
            : 'Connected, but the description was not saved.',
        );
      }

      invalidate();
      onDone();
    };

    window.addEventListener('message', handleMessage);
    return () => window.removeEventListener('message', handleMessage);
    // Registered once. Everything the handler reads that can change while the popup is
    // open is read through a ref, so there is nothing here to re-register for.
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, []);

  /**
   * `urlOverride` exists for the known-provider picker, which fills the field and
   * probes in the same click. `serverUrl` would still hold the previous value at that
   * point, so the probe has to be handed the URL rather than read it back out of state.
   */
  const runProbe = async (urlOverride?: string) => {
    const url = (urlOverride ?? serverUrl).trim();
    if (!url) {
      toast.error('Paste the MCP server URL first');
      return;
    }
    try {
      const result = await probeServer({ serverUrl: url });
      setProbe(result);
      setServerUrl(url);
      setName((current) => {
        if (current) return current;
        autoFilledNameRef.current = hostOf(url);
        return autoFilledNameRef.current;
      });
    } catch (err) {
      toast.error(err instanceof Error ? err.message : 'Could not reach that server');
    }
  };

  /**
   * Fills the URL field and runs the probe straight away, landing the user on exactly
   * the verdict they would have reached by pasting that URL. The field is left visible
   * and editable, so what got filled in is on screen and can still be replaced.
   */
  const handlePickProvider = async (provider: KnownMcpProvider) => {
    setServerUrl(provider.serverUrl);
    setProbe(null);
    // Seeded before the probe so the brand name, not the bare host, is the default
    // connection name. `runProbe` only fills a name when one is not already set.
    // A name this form filled in for a previously picked provider is replaced; one the
    // user typed is kept.
    setName((current) => {
      if (current && current !== autoFilledNameRef.current) return current;
      autoFilledNameRef.current = provider.name;
      return autoFilledNameRef.current;
    });
    await runProbe(provider.serverUrl);
  };

  const requireInstructions = () => {
    if (instructions.trim()) return true;
    toast.error('Tell the agent when to use this before saving');
    return false;
  };

  const handleConnectOAuth = async () => {
    if (!probe?.providerId || !requireInstructions()) return;
    try {
      setIsConnecting(true);
      const { oauthUrl } = await initiateOAuth({
        integration: probe.providerId,
        scope,
        ...(userId ? { userId } : {}),
      });
      if (!oauthUrl) {
        setIsConnecting(false);
        toast.error('That server did not return a sign-in URL');
        return;
      }
      const popup = window.open(oauthUrl, '_blank', 'width=600,height=700');
      if (!popup) {
        setIsConnecting(false);
        toast.error('Your browser blocked the sign-in window. Allow popups for Cedar and retry.');
        return;
      }
      stopWatchingPopup();
      popupRef.current = popup;
      // Both of these say "a sign-in this step started is in flight": the URL its result
      // belongs to, and the section's cue to leave the callback message alone.
      connectingServerUrlRef.current = probe.serverUrl;
      ownershipChangeRef.current?.(true);
      popupPollRef.current = setInterval(() => {
        // `closed` is readable cross-origin; nothing else about the popup is.
        if (popupRef.current?.closed !== false) {
          stopWatchingPopup();
          setIsConnecting(false);
        }
      }, 500);
    } catch (err) {
      setIsConnecting(false);
      toast.error(err instanceof Error ? err.message : 'Could not start sign-in');
    }
  };

  const handleSaveStaticToken=[redacted] () => {
    if (!probe || !requireInstructions()) return;
    if (!name.trim()) {
      toast.error('Give this connection a name');
      return;
    }
    try {
      await addConnection({
        name: name.trim(),
        serverUrl: probe.serverUrl,
        instructions: instructions.trim(),
        ...(token.trim() ? { authorizationHeader: toAuthorizationHeader(token) } : {}),
        ...(userId ? { userId } : {}),
      });
      invalidate();
      onDone();
    } catch (err) {
      toast.error(err instanceof Error ? err.message : 'Could not save that connection');
    }
  };

  const handleSaveManualClient = async () => {
    if (!probe || !requireInstructions()) return;
    if (!clientId.trim() || !clientSecret.trim()) {
      toast.error('Paste both the client ID and the client secret');
      return;
    }
    try {
      const payload = {
        name: vaultNameFromUrl(probe.serverUrl),
        credentials: { client_id: clientId.trim(), client_secret=[redacted] },
        instructions: instructions.trim(),
        availableToBackgroundAgents: false,
      };
      // An admin registering the app for the organization gets the ORG entry. The two
      // vault writes differ only in which scope they store under, and sending an org
      // save down the user path left the admin with a personal credential and nothing
      // shared, with no failure to notice.
      if (scope === 'org') await upsertOrgCredential(payload);
      else await upsertCredential(payload);
      invalidate();
      onDone();
    } catch (err) {
      toast.error(err instanceof Error ? err.message : 'Could not save those credentials');
    }
  };

  const copyRedirectUri = async () => {
    if (!probe?.redirectUri) return;
    try {
      await navigator.clipboard.writeText(probe.redirectUri);
      setCopied(true);
      setTimeout(() => setCopied(false), 1500);
    } catch {
      toast.error('Could not copy, select the URL and copy it by hand');
    }
  };

  const oauthScopeField = (
    <ScopeSelect
      id="mcp-scope"
      value={scope}
      onChange={setScope}
      isOrgAdmin={isOrgAdmin}
      orgNote="Everyone in your organization uses this connection through your sign-in, and only an admin can reconnect it when it expires."
    />
  );

  const clientScopeField = (
    <ScopeSelect
      id="mcp-client-scope"
      value={scope}
      onChange={setScope}
      isOrgAdmin={isOrgAdmin}
      orgNote="Everyone in your organization signs in through this app registration, and each of them still authorizes as themselves."
    />
  );

  /**
   * Shown on the token forms when the user came in from the Organization tab.
   *
   * `addMcpConnection` writes a user-scoped row and has no org variant, so there is
   * nothing to send and no selector to offer. Saying so is the whole point: without it
   * an admin picks Organization, saves, sees nothing fail, and finds the entry under
   * Mine.
   */
  const tokenScopeNote = defaultScope === 'org' && (
    <p className="text-muted-foreground text-xs">
      Token connections are saved to your own account. Only sign-in connections can be shared with
      your organization.
    </p>
  );

  const instructionsField = (
    <div className="space-y-2">
      <Label htmlFor="mcp-instructions">When should the agent use this?</Label>
      <p className="text-muted-foreground text-xs">{INSTRUCTIONS_HELP}</p>
      <Textarea
        id="mcp-instructions"
        rows={3}
        placeholder="e.g. Our product docs. Use this to answer customer questions about setup, limits and pricing."
        value={instructions}
        onChange={(e) => setInstructions(e.target.value)}
      />
    </div>
  );

  const staticTokenFields = (
    <>
      <div className="space-y-2">
        <Label htmlFor="mcp-name">Name</Label>
        <Input
          id="mcp-name"
          placeholder="e.g. Product Docs"
          value={name}
          onChange={(e) => setName(e.target.value)}
        />
      </div>
      <div className="space-y-2">
        <Label htmlFor="mcp-token">
          Access token <span className="text-muted-foreground font-normal">(optional)</span>
        </Label>
        <Input
          id="mcp-token"
          type="password"
          placeholder="your_access_token_here"
          value={token}
          onChange={(e) => setToken(e.target.value)}
        />
        <p className="text-muted-foreground text-xs">
          Paste just the token. Cedar sends it as a Bearer header and stores it encrypted. Leave it
          blank if the server needs no auth.
        </p>
      </div>
    </>
  );

  const providerPicker = (
    <div className="space-y-2">
      <p className="text-sm font-medium">Popular systems</p>
      <div className="grid grid-cols-2 gap-2">
        {KNOWN_MCP_PROVIDERS.map((provider) => {
          const isSelected = serverUrl.trim() === provider.serverUrl;
          return (
            <button
              key=[redacted]
              type="button"
              onClick={() => void handlePickProvider(provider)}
              disabled={isProbing}
              className={cn(
                // Sunken, like every other control that takes a choice. These were `bg-card`
                // — RAISED — so inside a dialog they read as chips sitting on top of the sheet
                // rather than as inputs you pick from, and they did not match the URL field
                // directly below them.
                'border-border bg-sunken flex cursor-pointer flex-col gap-2 rounded-lg border p-3 text-left',
                'transition-[background-color,border-color] duration-150',
                'hover:border-foreground/20 disabled:cursor-not-allowed disabled:opacity-60',
                isSelected && 'border-primary ring-primary/40 ring-1',
              )}
            >
              <span className="flex h-5 items-center">
                {provider.logoSrc ? (
                  /*
                   * Flattened to the foreground color rather than shown in brand
                   * color: these files are single-tone wordmarks drawn for a light
                   * page, and this dialog is rendered in both themes. It is the same
                   * treatment logos already get elsewhere in the app's chrome.
                   */
                  <img
                    src={provider.logoSrc}
                    alt={provider.name}
                    className="h-4 w-auto max-w-[110px] object-contain object-left brightness-0 dark:invert"
                  />
                ) : (
                  <Server className="text-muted-foreground h-4 w-4" />
                )}
              </span>
              <span className="min-w-0">
                <span className="block text-sm font-medium">{provider.name}</span>
                <span className="text-muted-foreground block truncate text-xs">
                  {hostOf(provider.serverUrl)}
                </span>
              </span>
            </button>
          );
        })}
      </div>
      <p className="text-muted-foreground text-xs">
        Pick one to fill in its URL and check it, or paste any other MCP server below.
      </p>
    </div>
  );

  return (
    <div className="space-y-4">
      {providerPicker}

      <div className="space-y-2">
        <Label htmlFor="mcp-server-url">MCP server URL</Label>
        <div className="flex gap-2">
          <Input
            id="mcp-server-url"
            placeholder="https://mcp.example.com/mcp"
            value={serverUrl}
            onChange={(e) => {
              setServerUrl(e.target.value);
              // The verdict was about the old URL, and so was any name this form derived
              // from it: picking Notion and then pasting an Acme URL must not leave the
              // new server named "Notion". A name the user typed is theirs and is kept.
              setProbe(null);
              setName((current) => (current === autoFilledNameRef.current ? '' : current));
            }}
          />
          <Button
            type="button"
            variant="outline"
            onClick={() => void runProbe()}
            disabled={isProbing}
          >
            {isProbing && <Loader2 className="mr-2 h-3.5 w-3.5 animate-spin" />}
            Check
          </Button>
        </div>
        <p className="text-muted-foreground text-xs">
          Cedar asks the server what connecting to it takes, then shows only the steps that server
          actually needs.
        </p>
      </div>

      {probe && (probe.verdict === 'dcr' || probe.verdict === 'dcr_confidential') && (
        <div className="space-y-4 rounded-lg border p-4">
          <div className="flex items-start gap-2">
            <ShieldCheck className="mt-0.5 h-4 w-4 flex-shrink-0 text-green-600" />
            <div className="space-y-1 text-sm">
              <p className="font-medium">Ready to connect</p>
              <p className="text-muted-foreground text-xs">
                {probe.issuerHost} registers Cedar automatically, so there is nothing to paste.
                {probe.verdict === 'dcr_confidential'
                  ? ' Cedar registers a confidential client and keeps its secret encrypted.'
                  : ''}
              </p>
              {probe.scopes.length > 0 && (
                <p className="text-muted-foreground text-xs">
                  Cedar will ask for: {probe.scopes.join(', ')}
                </p>
              )}
            </div>
          </div>

          {probe.providerId && oauthScopeField}
          {instructionsField}

          {probe.providerId ? (
            <Button size="sm" onClick={handleConnectOAuth} disabled={isConnecting}>
              {isConnecting && <Loader2 className="mr-2 h-3.5 w-3.5 animate-spin" />}
              Connect
            </Button>
          ) : (
            <div className="space-y-4">
              <p className="rounded-md border border-amber-500/30 bg-amber-500/[0.06] px-2.5 py-2 text-xs text-amber-700 dark:text-amber-500">
                This server supports automatic sign-in, but Cedar cannot start it for a server it
                has not shipped support for yet. Connect it with an access token for now.
              </p>
              {tokenScopeNote}
              {staticTokenFields}
              <Button size="sm" onClick={handleSaveStaticToken} disabled={isAdding}>
                {isAdding && <Loader2 className="mr-2 h-3.5 w-3.5 animate-spin" />}
                Save connection
              </Button>
            </div>
          )}
        </div>
      )}

      {probe?.verdict === 'manual_client' && (
        <div className="space-y-4 rounded-lg border p-4">
          <div className="space-y-1 text-sm">
            <p className="font-medium">This server needs an app you register yourself</p>
            <p className="text-muted-foreground text-xs">
              {probe.issuerHost} does not hand out clients automatically. Create an app in that
              provider&apos;s developer settings, add the redirect URL below to it, then paste the
              client ID and secret it gives you. Cedar stores them encrypted and uses them to sign
              in.
            </p>
          </div>

          {probe.redirectUri && (
            <div className="space-y-1.5">
              <Label>Redirect URL to add to that app</Label>
              <div className="flex items-center gap-1.5">
                <code className="bg-muted min-w-0 flex-1 truncate rounded px-1.5 py-1 text-[11px]">
                  {probe.redirectUri}
                </code>
                <Button
                  type="button"
                  variant="outline"
                  size="icon"
                  className="h-8 w-8 flex-shrink-0"
                  onClick={copyRedirectUri}
                  title="Copy redirect URL"
                >
                  {copied ? (
                    <Check className="h-3 w-3 text-green-600" />
                  ) : (
                    <Copy className="h-3 w-3" />
                  )}
                </Button>
              </div>
            </div>
          )}

          <div className="space-y-2">
            <Label htmlFor="mcp-client-id">Client ID</Label>
            <Input
              id="mcp-client-id"
              value={clientId}
              onChange={(e) => setClientId(e.target.value)}
            />
          </div>
          <div className="space-y-2">
            <Label htmlFor="mcp-client-secret">Client secret</Label>
            <Input
              id="mcp-client-secret"
              type="password"
              value={clientSecret}
              onChange={(e) => setClientSecret(e.target.value)}
            />
          </div>

          {clientScopeField}
          {instructionsField}

          <Button size="sm" onClick={handleSaveManualClient} disabled={isSavingCredential}>
            {isSavingCredential && <Loader2 className="mr-2 h-3.5 w-3.5 animate-spin" />}
            Save credentials
          </Button>
        </div>
      )}

      {probe?.verdict === 'no_oauth' && (
        <div className="space-y-4 rounded-lg border p-4">
          <div className="flex items-start gap-2">
            <AlertTriangle className="mt-0.5 h-4 w-4 flex-shrink-0 text-amber-600 dark:text-amber-500" />
            <div className="space-y-1 text-sm">
              <p className="font-medium">No sign-in offered, use a token</p>
              <p className="text-muted-foreground text-xs">{probe.discoveryError}</p>
            </div>
          </div>

          {tokenScopeNote}
          {staticTokenFields}
          {instructionsField}

          <Button size="sm" onClick={handleSaveStaticToken} disabled={isAdding}>
            {isAdding && <Loader2 className="mr-2 h-3.5 w-3.5 animate-spin" />}
            Save connection
          </Button>
        </div>
      )}
    </div>
  );
}