diff --git a/ui/desktop/src/components/McpApps/useSandboxBridge.ts b/ui/desktop/src/components/McpApps/useSandboxBridge.ts index ba8bf69f049b..f519f50bbe19 100644 --- a/ui/desktop/src/components/McpApps/useSandboxBridge.ts +++ b/ui/desktop/src/components/McpApps/useSandboxBridge.ts @@ -86,6 +86,21 @@ export function useSandboxBridge(options: SandboxBridgeOptions): SandboxBridgeRe case 'ui/notifications/initialized': isGuestInitializedRef.current = true; + // Send any pending tool data that arrived before initialization + if (toolInput) { + sendToSandbox({ + jsonrpc: '2.0', + method: 'ui/notifications/tool-input', + params: { arguments: toolInput.arguments }, + }); + } + if (toolResult) { + sendToSandbox({ + jsonrpc: '2.0', + method: 'ui/notifications/tool-result', + params: toolResult, + }); + } break; case 'ui/notifications/size-changed': { @@ -163,7 +178,16 @@ export function useSandboxBridge(options: SandboxBridgeOptions): SandboxBridgeRe } } }, - [resourceHtml, resourceCsp, resolvedTheme, sendToSandbox, onMcpRequest, onSizeChanged] + [ + resourceHtml, + resourceCsp, + resolvedTheme, + sendToSandbox, + onMcpRequest, + onSizeChanged, + toolInput, + toolResult, + ] ); useEffect(() => { diff --git a/ui/desktop/src/components/ToolCallWithResponse.tsx b/ui/desktop/src/components/ToolCallWithResponse.tsx index 77216979dd70..bbc5b9906e42 100644 --- a/ui/desktop/src/components/ToolCallWithResponse.tsx +++ b/ui/desktop/src/components/ToolCallWithResponse.tsx @@ -1,6 +1,6 @@ import { ToolIconWithStatus, ToolCallStatus } from './ToolCallStatusIndicator'; import { getToolCallIcon } from '../utils/toolIconMapping'; -import React, { useEffect, useRef, useState } from 'react'; +import React, { useEffect, useRef, useState, useMemo } from 'react'; import { Button } from './ui/button'; import { ToolCallArguments, ToolCallArgumentValue } from './ToolCallArguments'; import MarkdownContent from './MarkdownContent'; @@ -71,12 +71,19 @@ function isEmbeddedResource(content: Content): content is EmbeddedResource { return 'resource' in content && typeof (content as Record).resource === 'object'; } -function maybeRenderMCPApp( - toolRequest: ToolRequestMessageContent, - toolResponse: ToolResponseMessageContent | undefined, - sessionId: string, - append?: (value: string) => void -): React.ReactNode { +interface McpAppWrapperProps { + toolRequest: ToolRequestMessageContent; + toolResponse?: ToolResponseMessageContent; + sessionId: string; + append?: (value: string) => void; +} + +function McpAppWrapper({ + toolRequest, + toolResponse, + sessionId, + append, +}: McpAppWrapperProps): React.ReactNode { const requestWithMeta = toolRequest as ToolRequestWithMeta; let resourceUri = requestWithMeta._meta?.['ui/resourceUri']; @@ -87,24 +94,37 @@ function maybeRenderMCPApp( } } - if (!resourceUri) return null; - if (requestWithMeta.toolCall.status !== 'success') return null; + const extensionName = + requestWithMeta.toolCall.status === 'success' + ? requestWithMeta.toolCall.value.name.split('__')[0] + : ''; + + const toolArguments = + requestWithMeta.toolCall.status === 'success' + ? requestWithMeta.toolCall.value.arguments + : undefined; - const extensionName = requestWithMeta.toolCall.value.name.split('__')[0]; + // Memoize toolInput to prevent unnecessary re-renders + const toolInput = useMemo(() => ({ arguments: toolArguments || {} }), [toolArguments]); - let toolResult: CallToolResponse | undefined; - if (toolResponse) { + // Memoize toolResult to prevent unnecessary re-renders + const toolResult = useMemo(() => { + if (!toolResponse) return undefined; const resultWithMeta = toolResponse.toolResult as ToolResultWithMeta; if (resultWithMeta?.status === 'success' && resultWithMeta.value) { - toolResult = resultWithMeta.value; + return resultWithMeta.value; } - } + return undefined; + }, [toolResponse]); + + if (!resourceUri) return null; + if (requestWithMeta.toolCall.status !== 'success') return null; return (
+ )} ); }