From dcb3f794d329fcbfaa34a17f47e4795be0ddd3d0 Mon Sep 17 00:00:00 2001
From: Sam Hoang Van <samhv.ict@gmail.com>
Date: Tue, 10 Jun 2025 15:06:42 +0700
Subject: [PATCH] Feat: render mcp content (#5229)

---
 web-app/src/containers/ToolCallBlock.tsx | 176 +++++++++++++++++++++--
 1 file changed, 167 insertions(+), 9 deletions(-)

diff --git a/web-app/src/containers/ToolCallBlock.tsx b/web-app/src/containers/ToolCallBlock.tsx
index f0f6ca934..f1e50c939 100644
--- a/web-app/src/containers/ToolCallBlock.tsx
+++ b/web-app/src/containers/ToolCallBlock.tsx
@@ -2,6 +2,13 @@ import { ChevronDown, ChevronUp, Loader } from 'lucide-react'
 import { cn } from '@/lib/utils'
 import { create } from 'zustand'
 import { RenderMarkdown } from './RenderMarkdown'
+import { useMemo, useState } from 'react'
+import {
+  Dialog,
+  DialogContent,
+  DialogHeader,
+  DialogTitle,
+} from '@/components/ui/dialog'
 
 interface Props {
   result: string
@@ -26,23 +33,130 @@ const useToolCallBlockStore = create<ToolCallBlockState>((set) => ({
     })),
 }))
 
+// Types for MCP response content
+interface MCPContentItem {
+  type: string
+  data?: string
+  text?: string
+  mimeType?: string
+}
+
+interface MCPResponse {
+  content?: MCPContentItem[]
+}
+
+// Utility function to create data URL from base64 and mimeType
+const createDataUrl = (base64Data: string, mimeType: string): string => {
+  // Handle case where base64 data might already include data URL prefix
+  if (base64Data.startsWith('data:')) {
+    return base64Data
+  }
+  return `data:${mimeType};base64,${base64Data}`
+}
+
+// Parse MCP response and extract content items
+const parseMCPResponse = (result: string) => {
+  try {
+    const parsed: MCPResponse = JSON.parse(result)
+    const content = parsed.content || []
+
+    return {
+      parsedResult: parsed,
+      contentItems: content,
+      hasStructuredContent: content.length > 0,
+      parseError: false,
+    }
+  } catch {
+    // Fallback: JSON parsing failed, treat as plain text
+    return {
+      parsedResult: result,
+      contentItems: [],
+      hasStructuredContent: false,
+      parseError: true,
+    }
+  }
+}
+
+// Component to render individual content items based on type
+const ContentItemRenderer = ({
+  item,
+  index,
+  onImageClick,
+}: {
+  item: MCPContentItem
+  index: number
+  onImageClick?: (imageUrl: string, alt: string) => void
+}) => {
+  if (item.type === 'image' && item.data && item.mimeType) {
+    const imageUrl = createDataUrl(item.data, item.mimeType)
+    return (
+      <div key={index} className="mt-3">
+        <img
+          src={imageUrl}
+          alt={`Result image ${index + 1}`}
+          className="max-w-full max-h-64 object-contain rounded-md border border-main-view-fg/10 cursor-pointer hover:opacity-80 transition-opacity"
+          onError={(e) => {
+            // Hide broken images
+            e.currentTarget.style.display = 'none'
+          }}
+          onClick={() => onImageClick?.(imageUrl, `Result image ${index + 1}`)}
+        />
+      </div>
+    )
+  }
+
+  if (item.type === 'text' && item.text) {
+    return (
+      <div key={index} className="mt-3">
+        <RenderMarkdown content={item.text} />
+      </div>
+    )
+  }
+
+  // For any other types, render as JSON
+  return (
+    <div key={index} className="mt-3">
+      <RenderMarkdown
+        content={'```json\n' + JSON.stringify(item, null, 2) + '\n```'}
+      />
+    </div>
+  )
+}
+
 const ToolCallBlock = ({ id, name, result, loading }: Props) => {
   const { collapseState, setCollapseState } = useToolCallBlockStore()
   const isExpanded = collapseState[id] ?? false
+  const [modalImage, setModalImage] = useState<{
+    url: string
+    alt: string
+  } | null>(null)
 
   const handleClick = () => {
     const newExpandedState = !isExpanded
     setCollapseState(id, newExpandedState)
   }
 
+  const handleImageClick = (imageUrl: string, alt: string) => {
+    setModalImage({ url: imageUrl, alt })
+  }
+
+  const closeModal = () => {
+    setModalImage(null)
+  }
+
+  // Parse the MCP response and extract content items
+  const { parsedResult, contentItems, hasStructuredContent, parseError } =
+    useMemo(() => {
+      return parseMCPResponse(result)
+    }, [result])
+
   return (
     <div
       className="mx-auto w-full cursor-pointer break-words"
-      onClick={handleClick}
       data-tool-call-block={id}
     >
       <div className="rounded-lg bg-main-view-fg/4 border border-dashed border-main-view-fg/10">
-        <div className="flex items-center gap-3 p-2">
+        <div className="flex items-center gap-3 p-2" onClick={handleClick}>
           {loading && (
             <Loader className="size-4 animate-spin text-main-view-fg/60" />
           )}
@@ -66,16 +180,60 @@ const ToolCallBlock = ({ id, name, result, loading }: Props) => {
           )}
         >
           <div className="mt-2 text-main-view-fg/60">
-            <RenderMarkdown
-              content={
-                '```json\n' +
-                JSON.stringify(result ? JSON.parse(result) : null, null, 2) +
-                '\n```'
-              }
-            />
+            {hasStructuredContent ? (
+              /* Render each content item individually based on its type */
+              <div className="space-y-2">
+                {contentItems.map((item, index) => (
+                  <ContentItemRenderer
+                    key={index}
+                    item={item}
+                    index={index}
+                    onImageClick={handleImageClick}
+                  />
+                ))}
+              </div>
+            ) : parseError ? (
+              /* Handle JSON parse error - render as plain text */
+              <div className="mt-3 p-3 bg-main-view-fg/5 rounded-md border border-main-view-fg/10">
+                <div className="text-sm font-medium text-main-view-fg/80 mb-2">
+                  Raw Response:
+                </div>
+                <div className="whitespace-pre-wrap font-mono text-sm">
+                  {parsedResult as string}
+                </div>
+              </div>
+            ) : (
+              /* Fallback: render as JSON for valid JSON but unstructured responses */
+              <RenderMarkdown
+                content={
+                  '```json\n' + JSON.stringify(parsedResult, null, 2) + '\n```'
+                }
+              />
+            )}
           </div>
         </div>
       </div>
+
+      {/* Image Modal */}
+      <Dialog open={!!modalImage} onOpenChange={(open) => !open && closeModal()}>
+        <DialogContent className="max-w-4xl max-h-[90vh] p-0">
+          <DialogHeader className="p-6 pb-2">
+            <DialogTitle>{modalImage?.alt || 'Image'}</DialogTitle>
+          </DialogHeader>
+          <div className="flex justify-center items-center p-6 pt-2">
+            {modalImage && (
+              <img
+                src={modalImage.url}
+                alt={modalImage.alt}
+                className="max-w-full max-h-[70vh] object-contain rounded-md"
+                onError={(e) => {
+                  e.currentTarget.style.display = 'none'
+                }}
+              />
+            )}
+          </div>
+        </DialogContent>
+      </Dialog>
     </div>
   )
 }