Chore: enhance MCP authentication flow with automatic resume and connection status updates

.changeset/pre/auto-resume-mcp-auth.md

@@ -0,0 +1,6 @@
+---
+"@truefoundry/trueforge-ui": patch
+---
+
+Show successful MCP authentication in chat, automatically continue after every required server connects, and indicate
+while the turn is starting.

packages/trueforge-ui/src/atoms/adapters/McpAuthPromptAdapter.tsx

@@ -9,6 +9,8 @@ export type McpServer = {
 
 export type McpAuthPromptProps = {
   servers: McpServer[];
+  connectedServerIds?: ReadonlySet<string>;
+  continueLoading?: boolean;
   onConnect: (serverId: string) => void;
   onContinue?: () => void;
   readOnly?: boolean;
@@ -21,6 +23,8 @@ const DEFAULT_TITLE = 'MCP Authentication Required';
 
 export function McpAuthPrompt({
   servers,
+  connectedServerIds,
+  continueLoading = false,
   onConnect,
   onContinue,
   readOnly = false,
@@ -46,15 +50,27 @@ export function McpAuthPrompt({
               <span className="shrink-0 text-xs font-semibold text-text-secondary">:</span>
               <span className="truncate font-sans font-medium text-text-primary">{server.name}</span>
             </div>
-            <Button.Primary size="small" disabled={readOnly} onClick={() => onConnect(server.id)} className="shrink-0">
-              Connect
-              <Icon name="external-link" size="0.75em" className="ml-1" />
-            </Button.Primary>
+            {connectedServerIds?.has(server.id) ? (
+              <Button.Primary size="small" disabled className="shrink-0">
+                Connected
+              </Button.Primary>
+            ) : (
+              <Button.Primary
+                size="small"
+                disabled={readOnly}
+                onClick={() => onConnect(server.id)}
+                className="shrink-0"
+              >
+                Connect
+                <Icon name="external-link" size="0.75em" className="ml-1" />
+              </Button.Primary>
+            )}
           </div>
         ))}
         {onContinue && (
           <div className="flex justify-end border-t border-border pt-2">
-            <Button.Primary size="small" disabled={readOnly} onClick={onContinue}>
+            <Button.Primary size="small" disabled={readOnly || continueLoading} onClick={onContinue}>
+              {continueLoading ? <Icon name="loader" className="animate-spin" /> : null}
               Continue
             </Button.Primary>
           </div>

packages/trueforge-ui/src/containers/McpAuthContainer.tsx

@@ -2,6 +2,7 @@
 
 import { useThreadIsRunning } from '@assistant-ui/core/react';
 import { useTrueFoundryMcpAuth } from '@truefoundry/assistant-ui-runtime';
+import { useRef, useState } from 'react';
 
 import { useDraftCatalog } from '@/atoms/draft/DraftCatalogProvider.js';
 import { useMCPAuth } from '@/hooks/useMcpAuth.js';
@@ -10,24 +11,51 @@ import { useSlot } from '../theme/SlotsProvider.js';
 
 type McpAuthPromptProps = {
   servers: NonNullable<ReturnType<typeof useTrueFoundryMcpAuth>['pending']>['mcpServers'];
-  onContinue: () => void;
+  onContinue: () => Promise<void>;
   readOnly: boolean;
 };
 
 function CatalogMcpAuthPrompt({ servers, onContinue, readOnly }: McpAuthPromptProps) {
   const McpAuthPrompt = useSlot('McpAuthPrompt');
   const { handleAuthorize } = useMCPAuth();
   const { refreshConnectors } = useDraftCatalog();
+  const [connectedServerIds, setConnectedServerIds] = useState<ReadonlySet<string>>(() => new Set());
+  const connectedServerIdsRef = useRef(connectedServerIds);
+  const [isResuming, setIsResuming] = useState(false);
+  const resumedRef = useRef(false);
+
+  const startResume = () => {
+    if (readOnly || resumedRef.current) return;
+    resumedRef.current = true;
+    setIsResuming(true);
+    void onContinue().catch(() => {
+      resumedRef.current = false;
+      setIsResuming(false);
+    });
+  };
 
   const handleConnect = (serverId: string) => {
     void handleAuthorize(serverId, isSuccess => {
       if (isSuccess) {
+        const nextConnectedServerIds = new Set([...connectedServerIdsRef.current, serverId]);
+        connectedServerIdsRef.current = nextConnectedServerIds;
+        setConnectedServerIds(nextConnectedServerIds);
         void refreshConnectors();
+        if (servers.every(server => nextConnectedServerIds.has(server.id))) startResume();
       }
     });
   };
 
-  return <McpAuthPrompt servers={servers} onConnect={handleConnect} onContinue={onContinue} readOnly={readOnly} />;
+  return (
+    <McpAuthPrompt
+      servers={servers}
+      connectedServerIds={connectedServerIds}
+      continueLoading={isResuming}
+      onConnect={handleConnect}
+      onContinue={startResume}
+      readOnly={readOnly}
+    />
+  );
 }
 
 export function McpAuthContainer({ disabled = false }: { disabled?: boolean }) {
@@ -39,12 +67,12 @@ export function McpAuthContainer({ disabled = false }: { disabled?: boolean }) {
   if (!pending) return null;
 
   if (catalog) {
+    const pendingServerKey = JSON.stringify(pending.mcpServers.map(server => server.id));
     return (
       <CatalogMcpAuthPrompt
+        key={pendingServerKey}
         servers={pending.mcpServers}
-        onContinue={() => {
-          if (!disabled) void resume();
-        }}
+        onContinue={resume}
         readOnly={isRunning || disabled}
       />
     );

packages/trueforge-ui/test/containers/McpAuthContainer.test.tsx

@@ -1,25 +1,30 @@
 // @vitest-environment jsdom
 import { AssistantRuntimeProvider, useExternalStoreRuntime, type ThreadMessageLike } from '@assistant-ui/react';
-import { fireEvent, render, screen } from '@testing-library/react';
+import { fireEvent, render, screen, waitFor } from '@testing-library/react';
 import { trueFoundryExtras, type TrueFoundryRuntimeExtras } from '@truefoundry/assistant-ui-runtime';
 import { describe, expect, it, vi } from 'vitest';
 
+import { DraftCatalogProvider } from '@/atoms/draft/DraftCatalogProvider.js';
 import { McpAuthContainer } from '@/containers/McpAuthContainer.js';
+import { ServerProvider } from '@/server/ServerContext.js';
+import type { AgentUIServer } from '@/server/types.js';
+import { createMockAgentUIServer, createMockCatalog } from '../server/mockServer.js';
 
-const SERVERS = [
-  { id: 'srv-1', name: 'github', authUrl: 'https://example.com/auth/github' },
-  { id: 'srv-2', name: 'slack', authUrl: 'https://example.com/auth/slack' },
-];
+const GITHUB_SERVER = { id: 'srv-1', name: 'github', authUrl: 'https://example.com/auth/github' };
+const SLACK_SERVER = { id: 'srv-2', name: 'slack', authUrl: 'https://example.com/auth/slack' };
+const SERVERS = [GITHUB_SERVER, SLACK_SERVER];
 const PENDING = { mcpServers: SERVERS };
 
 function McpAuthHarness({
   pendingMcpAuth,
   resumeMcpAuth,
   isRunning = false,
+  server,
 }: {
   pendingMcpAuth: TrueFoundryRuntimeExtras['pendingMcpAuth'];
   resumeMcpAuth: TrueFoundryRuntimeExtras['resumeMcpAuth'];
   isRunning?: boolean;
+  server?: AgentUIServer;
 }) {
   const messages: ThreadMessageLike[] = [];
   const runtime = useExternalStoreRuntime({
@@ -47,11 +52,19 @@ function McpAuthHarness({
     }),
   });
 
-  return (
+  const content = (
     <AssistantRuntimeProvider runtime={runtime}>
       <McpAuthContainer />
     </AssistantRuntimeProvider>
   );
+
+  return server ? (
+    <ServerProvider server={server}>
+      <DraftCatalogProvider>{content}</DraftCatalogProvider>
+    </ServerProvider>
+  ) : (
+    content
+  );
 }
 
 describe('McpAuthContainer', () => {
@@ -93,4 +106,77 @@ describe('McpAuthContainer', () => {
     render(<McpAuthHarness pendingMcpAuth={PENDING} resumeMcpAuth={vi.fn()} isRunning={true} />);
     expect(screen.getByRole('button', { name: /continue/i })).toBeDisabled();
   });
+
+  it('shows each successful catalog connection and resumes once all servers are connected', async () => {
+    const resumeMcpAuth = vi.fn().mockResolvedValue(undefined);
+    const authenticateConnector = vi.fn().mockResolvedValue({ status: 'AUTHENTICATED' });
+    const catalog = createMockCatalog({
+      connectorCatalog: {
+        ...createMockCatalog().connectorCatalog,
+        authenticateConnector,
+      },
+    });
+    const server = createMockAgentUIServer({ catalog });
+
+    render(<McpAuthHarness pendingMcpAuth={PENDING} resumeMcpAuth={resumeMcpAuth} server={server} />);
+
+    const firstConnect = screen.getAllByRole('button', { name: 'Connect' })[0];
+    if (!firstConnect) throw new Error('Expected the first MCP Connect button');
+    fireEvent.click(firstConnect);
+
+    await waitFor(() => expect(screen.getByRole('button', { name: 'Connected' })).toBeDisabled());
+    expect(resumeMcpAuth).not.toHaveBeenCalled();
+
+    fireEvent.click(screen.getByRole('button', { name: 'Connect' }));
+
+    await waitFor(() => expect(screen.getAllByRole('button', { name: 'Connected' })).toHaveLength(2));
+    await waitFor(() => expect(screen.getByRole('button', { name: 'Continue' })).toBeDisabled());
+    expect(resumeMcpAuth).toHaveBeenCalledTimes(1);
+  });
+
+  it('keeps a failed catalog connection available without resuming', async () => {
+    const resumeMcpAuth = vi.fn().mockResolvedValue(undefined);
+    const authenticateConnector = vi.fn().mockRejectedValue(new Error('Authorization failed'));
+    const catalog = createMockCatalog({
+      connectorCatalog: {
+        ...createMockCatalog().connectorCatalog,
+        authenticateConnector,
+      },
+    });
+    const server = createMockAgentUIServer({ catalog });
+
+    render(<McpAuthHarness pendingMcpAuth={PENDING} resumeMcpAuth={resumeMcpAuth} server={server} />);
+
+    const firstConnect = screen.getAllByRole('button', { name: 'Connect' })[0];
+    if (!firstConnect) throw new Error('Expected the first MCP Connect button');
+    fireEvent.click(firstConnect);
+
+    await waitFor(() => expect(authenticateConnector).toHaveBeenCalledTimes(1));
+    expect(screen.getAllByRole('button', { name: 'Connect' })).toHaveLength(2);
+    expect(resumeMcpAuth).not.toHaveBeenCalled();
+  });
+
+  it('allows retrying Continue when resume fails', async () => {
+    const resumeMcpAuth = vi.fn().mockRejectedValue(new Error('Resume failed'));
+    const authenticateConnector = vi.fn().mockResolvedValue({ status: 'AUTHENTICATED' });
+    const catalog = createMockCatalog({
+      connectorCatalog: {
+        ...createMockCatalog().connectorCatalog,
+        authenticateConnector,
+      },
+    });
+    const server = createMockAgentUIServer({ catalog });
+
+    render(
+      <McpAuthHarness pendingMcpAuth={{ mcpServers: [GITHUB_SERVER] }} resumeMcpAuth={resumeMcpAuth} server={server} />,
+    );
+
+    fireEvent.click(screen.getByRole('button', { name: 'Connect' }));
+
+    await waitFor(() => expect(resumeMcpAuth).toHaveBeenCalledTimes(1));
+    await waitFor(() => expect(screen.getByRole('button', { name: 'Continue' })).toBeEnabled());
+
+    fireEvent.click(screen.getByRole('button', { name: 'Continue' }));
+    await waitFor(() => expect(resumeMcpAuth).toHaveBeenCalledTimes(2));
+  });
 });