AgentAddConnectionDialog.tsx8.3 KBView on GitHub
'use client';

import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
import { Check, Loader2 } from 'lucide-react';
import { useMemo, useState } from 'react';
import { toast } from 'sonner';

import {
  McpConnectionPicker,
  useMcpPickerConnections,
  type McpPickerConnection,
} from '@/modules/integrations/mcp-connection-picker';
import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogHeader,
  DialogTitle,
} from '@/components/ui/dialog';
import { effectiveGrantedServers, nextGrantEntries } from '@/modules/agents/utils/mcp-grants';
import { ScrollArea } from '@/components/ui/scroll-area';
import { useTRPC } from '@/providers/query-provider';
import { Button } from '@/components/ui/button';

/**
 * "Add connection", resolved WHERE IT WAS ASKED.
 *
 * This replaces a deep link to `/settings/connections`. The link was defensible —
 * a connection is owned there and shared by every agent — and it was still the
 * wrong control: you are on this screen because you want THIS agent to reach
 * something, and the link answers by navigating you off the agent, losing the
 * question, and landing you on a page with no idea which agent sent you. You then
 * connect the server and have to find your way back to grant it.
 *
 * What it creates is the SAME `connection` row Settings creates, so a server added
 * here appears at `/settings/connections` immediately. This file adds only the two
 * things that are about an agent: a Grant control per row, and granting whatever
 * you just connected.
 *
 * The three layers a connection actually has, as three visible controls:
 *
 *   THE CONNECTION  — created here, owned by the user, identical to one made in
 *                     Settings.
 *   THE AGENT GRANT — `mcp_servers:` on this agent's frontmatter. Granting here
 *                     grants to THIS agent only.
 *   THE ORG         — `metadata.scope`. 'user', the server is yours; 'org', every
 *                     teammate sees it and may grant it to their own agents. The
 *                     switch writes it through `setMcpConnectionScope`, which
 *                     stamps the same pair (`scope` + `organizationId`) the OAuth
 *                     callback does — one sharing mechanism, two entry points.
 *
 * `grantOnly` on the picker is what makes the middle layer honest: a connection
 * added from an agent is stamped `requiresAgentGrant`, so every OTHER agent gets
 * the OPTION (the row appears in its Connections list, ungranted) and none of them
 * silently gets the TOOL.
 *
 * TODO(mcp-add-step): Settings now adds servers through `systems/add-mcp-step.tsx`,
 * which runs OAuth DISCOVERY and can therefore handle a server needing a manually
 * registered client (Xero) or no OAuth at all. This dialog still runs the older
 * paste-a-token flow. Pointing it at that step would leave one add-a-server flow
 * instead of two, with the Grant control layered over it.
 */
export function AgentAddConnectionDialog({
  agentId,
  targetUserId,
  open,
  onOpenChange,
}: {
  agentId: string;
  /** An org admin/owner viewing a teammate's agent — see `AgentView`'s prop of the same name. */
  targetUserId?: string;
  open: boolean;
  onOpenChange: (next: boolean) => void;
}) {
  const trpc = useTRPC();
  const queryClient = useQueryClient();
  const [pendingGrant, setPendingGrant] = useState<string | null>(null);
  /** Mirrors the picker's own view, for the title. See its `onViewChange`. */
  const [view, setView] = useState<'catalog' | 'custom'>('catalog');

  const { connections, isLoading } = useMcpPickerConnections(open);

  const { data: grants } = useQuery({
    ...trpc.agent.getMcpGrants.queryOptions({ agentId, targetUserId }),
    enabled: open && !!agentId,
  });

  const { mutateAsync: setMcpGrants } = useMutation(trpc.agent.setMcpGrants.mutationOptions());

  const refreshAgent = () => {
    void queryClient.invalidateQueries({ queryKey=[redacted] });
    void queryClient.invalidateQueries({
      queryKey=[redacted] agentId, targetUserId }),
    });
  };

  // ABSENT `mcp_servers:` is not an empty list — see utils/mcp-grants.ts, which owns
  // the three-state read and the "materialise before you change it" rule.
  const grantedNames = useMemo(
    () => effectiveGrantedServers(grants, connections),
    [grants, connections],
  );

  const writeGrants = async (next: string[]) => {
    await setMcpGrants({ agentId, targetUserId, entries: next });
    refreshAgent();
  };

  const toggleGrant = async (name: string, on: boolean) => {
    setPendingGrant(name);
    try {
      await writeGrants(nextGrantEntries(grants, connections, { server: name, granted: on }));
    } catch (err) {
      toast.error(err instanceof Error ? err.message : 'Could not update this agent’s access');
    } finally {
      setPendingGrant(null);
    }
  };

  const grantButton = (c: McpPickerConnection) => {
    const granted = grantedNames.has(c.name);
    return (
      <Button
        type="button"
        size="sm"
        variant={granted ? 'secondary' : 'outline'}
        className="h-7 shrink-0 cursor-pointer gap-1.5 px-2.5 text-xs"
        disabled={pendingGrant === c.name}
        onClick={() => void toggleGrant(c.name, !granted)}
      >
        {pendingGrant === c.name ? (
          <Loader2 className="h-3 w-3 animate-spin" />
        ) : granted ? (
          <Check className="h-3 w-3" />
        ) : null}
        {granted ? 'Granted' : 'Grant'}
      </Button>
    );
  };

  return (
    <Dialog open={open} onOpenChange={onOpenChange}>
      {/* Wider than a list needs, because the FORM inside it needs it: at `max-w-lg`
          every field is a full-width row and the description ends up a lone textarea
          under three inputs. Two columns fit here. */}
      <DialogContent className="max-h-[85vh] gap-0 overflow-hidden p-0 sm:max-w-2xl">
        {/* ONE title, and it names what you are actually looking at. The form used to
            draw its own "Add MCP Connection" heading underneath this header's
            "Connections" — a title inside a title, the same container-in-a-container
            everything else on this surface just lost. No `border-b` either: that rule
            sat directly above the first connector and read as a divider belonging to
            the list rather than to the header. */}
        {/* `px-4`, the same 16px the picker's rows and the form's rows inset their
            own content by — so the title starts on the same vertical line as the
            first field's label. It used to be `px-5` over a form whose rows add
            their own `p-4`, which put every label 36px in against a 20px title. */}
        <DialogHeader className="space-y-1 px-4 py-4">
          <DialogTitle className="text-base">
            {view === 'custom' ? 'Add MCP connection' : 'Connections'}
          </DialogTitle>
          {/* The form says what it wants field by field; a paragraph above it is one
              more thing to read before you can start typing. */}
          {view === 'catalog' && (
            <DialogDescription className="text-xs">
              Grant what this agent may reach. Anything you add is saved to your connections and
              stays available to your other agents — they just will not have it until you grant it
              there too.
            </DialogDescription>
          )}
        </DialogHeader>

        <ScrollArea className="max-h-[65vh]">
          <McpConnectionPicker
            // No horizontal padding here: each view insets its own content, because
            // they inset it DIFFERENTLY — the catalog's rows pull their hover surface
            // outwards past the gutter, the form's rows are the gutter.
            className="pb-4"
            connections={connections}
            isLoading={isLoading}
            grantOnly
            rowAction={grantButton}
            // Grant it to the agent you are standing in. This is the whole reason the
            // modal exists — connecting without granting would land you back where the
            // Settings deep link did.
            onConnected={async ({ name }) =>
              writeGrants(nextGrantEntries(grants, connections, { server: name, granted: true }))
            }
            onChanged={refreshAgent}
            onViewChange={setView}
          />
        </ScrollArea>
      </DialogContent>
    </Dialog>
  );
}