CedarCopilot/cedar-mail

/**
 * The description a user types into the add-MCP form is applied to the row the OAuth
 * callback writes, and the two ends of that round trip do not agree on the server URL
 * string.
 *
 * `probeMcpServer` echoes back exactly what was typed; the callback stores the URL its
 * driver carries, which for Pylon has a trailing slash the user never typed. Matching
 * those as raw strings finds nothing, and the form used to close as though it had
 * worked, leaving a connection with the callback's generic auto-text. An entry with no
 * real description is one no agent reaches for, so this is silent data loss, not a
 * cosmetic miss.
 *
 * These tests pin both halves: the row IS found across that difference, and a genuine
 * miss is reported rather than swallowed.
 */
import { act, fireEvent, render, screen, waitFor } from '@testing-library/react';

// --- Mocks -----------------------------------------------------------------

const PROBED_URL = 'https://mcp.usepylon.com';
/** What Cedar's callback actually persists: the driver's URL, trailing slash and all. */
const STORED_URL = 'https://mcp.usepylon.com/';

const mockProbeMcpServer = jest.fn();
const mockInitiateOAuth = jest.fn();
const mockUpdateMcpConnectionSettings = jest.fn();
const mockAddMcpConnection = jest.fn();
const mockUpsertCredential = jest.fn();
const mockUpsertOrgCredential = jest.fn();
const mockFetchQuery = jest.fn();
const mockToastError = jest.fn();

jest.mock('sonner', () => ({
  toast: {
    error: (...args: unknown[]) => mockToastError(...args),
    success: jest.fn(),
    info: jest.fn(),
  },
}));

jest.mock('@tanstack/react-query', () => ({
  // The component only ever reads `mutateAsync` and `isPending`, and each call site
  // gets its options straight from the tRPC mock below, so handing back the mocked
  // mutationFn wires each button to its own spy.
  useMutation: (options: { mutationFn: (...args: unknown[]) => unknown }) => ({
    mutateAsync: options.mutationFn,
    isPending: false,
  }),
  useQueryClient: () => ({ fetchQuery: mockFetchQuery, invalidateQueries: jest.fn() }),
}));

jest.mock('@/providers/query-provider', () => ({
  useTRPC: () => ({
    integrations: {
      probeMcpServer: { mutationOptions: () => ({ mutationFn: mockProbeMcpServer }) },
      addMcpConnection: { mutationOptions: () => ({ mutationFn: mockAddMcpConnection }) },
      updateMcpConnectionSettings: {
        mutationOptions: () => ({ mutationFn: mockUpdateMcpConnectionSettings }),
      },
      initiateOAuth: { mutationOptions: () => ({ mutationFn: mockInitiateOAuth }) },
      listMcpConnections: {
        queryOptions: () => ({ queryKey: ['listMcpConnections'] }),
      },
      list: { queryKey: () => ['integrations.list'] },
    },
    credentialVault: {
      upsert: { mutationOptions: () => ({ mutationFn: mockUpsertCredential }) },
      upsertOrg: { mutationOptions: () => ({ mutationFn: mockUpsertOrgCredential }) },
      list: { queryKey: () => ['credentialVault.list'] },
    },
  }),
}));

// Imported after the mocks so the component resolves them.
import { AddMcpStep } from '../add-mcp-step';

/** A popup handle the poll loop reads as still open. */
const openPopup = { closed: false } as unknown as Window;

function renderStep() {
  const onDone = jest.fn();
  const onOAuthOwnershipChange = jest.fn();
  render(
    <AddMcpStep
      isOrgAdmin={false}
      onOAuthOwnershipChange={onOAuthOwnershipChange}
      onDone={onDone}
    />,
  );
  return { onDone, onOAuthOwnershipChange };
}

/** The message the OAuth callback page posts back into this window. */
function postCallbackSuccess() {
  act(() => {
    window.dispatchEvent(
      new MessageEvent('message', {
        origin: window.location.origin,
        data: { type: 'mcp_oauth_connected', success: true },
      }),
    );
  });
}

/** Paste the URL, probe it, type a description and press Connect. */
async function walkToConnected(instructions: string) {
  fireEvent.change(screen.getByLabelText('MCP server URL'), { target: { value: PROBED_URL } });
  fireEvent.click(screen.getByRole('button', { name: 'Check' }));
  await screen.findByText('Ready to connect');

  fireEvent.change(screen.getByLabelText('When should the agent use this?'), {
    target: { value: instructions },
  });
  fireEvent.click(screen.getByRole('button', { name: 'Connect' }));
  await waitFor(() => expect(mockInitiateOAuth).toHaveBeenCalled());

  postCallbackSuccess();
}

describe('AddMcpStep, applying the typed description after OAuth', () => {
  beforeEach(() => {
    jest.clearAllMocks();
    window.open = jest.fn(() => openPopup);
    mockProbeMcpServer.mockResolvedValue({
      verdict: 'dcr',
      providerId: 'pylon',
      // The server echoes the input back verbatim.
      serverUrl: PROBED_URL,
      issuerHost: 'app.usepylon.com',
      scopes: [],
      redirectUri: null,
      discoveryError: null,
    });
    mockInitiateOAuth.mockResolvedValue({ oauthUrl: 'https://app.usepylon.com/oauth/authorize' });
    mockUpdateMcpConnectionSettings.mockResolvedValue({ success: true });
  });

  it('finds the row the callback wrote even though its URL carries a trailing slash', async () => {
    mockFetchQuery.mockResolvedValue({
      connections: [{ id: 'conn-1', name: 'Pylon', serverUrl: STORED_URL }],
    });

    const { onDone } = renderStep();
    await walkToConnected('Our support desk. Use it for ticket history.');

    await waitFor(() =>
      expect(mockUpdateMcpConnectionSettings).toHaveBeenCalledWith({
        connectionId: 'conn-1',
        instructions: 'Our support desk. Use it for ticket history.',
      }),
    );
    await waitFor(() => expect(onDone).toHaveBeenCalled());
  });

  it('ignores a callback for a sign-in it did not start, and reports when it owns one', async () => {
    mockFetchQuery.mockResolvedValue({ connections: [] });
    const { onDone, onOAuthOwnershipChange } = renderStep();

    // A Reconnect running in the list behind this dialog finishes. That message belongs
    // to the section, which has stood down only while THIS step has a flow of its own.
    postCallbackSuccess();
    expect(mockFetchQuery).not.toHaveBeenCalled();
    expect(onDone).not.toHaveBeenCalled();
    expect(mockToastError).not.toHaveBeenCalled();
    expect(onOAuthOwnershipChange).not.toHaveBeenCalledWith(true);

    await walkToConnected('Our support desk.');
    expect(onOAuthOwnershipChange).toHaveBeenCalledWith(true);
    await waitFor(() => expect(onOAuthOwnershipChange).toHaveBeenLastCalledWith(false));
  });

  it('says so when no row matches, rather than closing as though it saved', async () => {
    mockFetchQuery.mockResolvedValue({
      connections: [{ id: 'other', name: 'Notion', serverUrl: 'https://mcp.notion.com/sse' }],
    });

    renderStep();
    await walkToConnected('Our support desk.');

    await waitFor(() =>
      expect(mockToastError).toHaveBeenCalledWith(
        "Connected, but the description was not saved. Add it from the entry's Edit.",
      ),
    );
    expect(mockUpdateMcpConnectionSettings).not.toHaveBeenCalled();
  });
});