diff --git a/.changeset/pre/auto-resume-mcp-auth.md b/.changeset/pre/auto-resume-mcp-auth.md new file mode 100644 index 000000000..0785345fb --- /dev/null +++ b/.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. diff --git a/packages/trueforge-ui/src/atoms/adapters/McpAuthPromptAdapter.tsx b/packages/trueforge-ui/src/atoms/adapters/McpAuthPromptAdapter.tsx index 4aa17192e..7f16749ec 100644 --- a/packages/trueforge-ui/src/atoms/adapters/McpAuthPromptAdapter.tsx +++ b/packages/trueforge-ui/src/atoms/adapters/McpAuthPromptAdapter.tsx @@ -9,6 +9,8 @@ export type McpServer = { export type McpAuthPromptProps = { servers: McpServer[]; + connectedServerIds?: ReadonlySet; + 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({ : {server.name} - onConnect(server.id)} className="shrink-0"> - Connect - - + {connectedServerIds?.has(server.id) ? ( + + Connected + + ) : ( + onConnect(server.id)} + className="shrink-0" + > + Connect + + + )} ))} {onContinue && (
- + + {continueLoading ? : null} Continue
diff --git a/packages/trueforge-ui/src/containers/McpAuthContainer.tsx b/packages/trueforge-ui/src/containers/McpAuthContainer.tsx index 0ea6bc7c6..20f5f35d9 100644 --- a/packages/trueforge-ui/src/containers/McpAuthContainer.tsx +++ b/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,7 +11,7 @@ import { useSlot } from '../theme/SlotsProvider.js'; type McpAuthPromptProps = { servers: NonNullable['pending']>['mcpServers']; - onContinue: () => void; + onContinue: () => Promise; readOnly: boolean; }; @@ -18,16 +19,43 @@ function CatalogMcpAuthPrompt({ servers, onContinue, readOnly }: McpAuthPromptPr const McpAuthPrompt = useSlot('McpAuthPrompt'); const { handleAuthorize } = useMCPAuth(); const { refreshConnectors } = useDraftCatalog(); + const [connectedServerIds, setConnectedServerIds] = useState>(() => 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 ; + return ( + + ); } 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 ( { - if (!disabled) void resume(); - }} + onContinue={resume} readOnly={isRunning || disabled} /> ); diff --git a/packages/trueforge-ui/test/containers/McpAuthContainer.test.tsx b/packages/trueforge-ui/test/containers/McpAuthContainer.test.tsx index 5a1cf5ec1..5e612f1e8 100644 --- a/packages/trueforge-ui/test/containers/McpAuthContainer.test.tsx +++ b/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 = ( ); + + return server ? ( + + {content} + + ) : ( + content + ); } describe('McpAuthContainer', () => { @@ -93,4 +106,77 @@ describe('McpAuthContainer', () => { render(); 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(); + + 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(); + + 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( + , + ); + + 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)); + }); });