From 9a4158c649d922deb4aced4c65ac63546e5743e8 Mon Sep 17 00:00:00 2001
From: Vivek R <123vivekr@gmail.com>
Date: Sun, 22 Jun 2025 22:39:53 +0530
Subject: [PATCH] feat: implement tool call/result mapping with collapsible UI

---
 src/components/AgentExecution.tsx        |  24 ++-
 src/components/ClaudeCodeSession.tsx     |  22 ++-
 src/components/CollapsibleToolResult.tsx | 191 +++++++++++++++++++++++
 src/components/SessionOutputViewer.tsx   |  24 ++-
 src/components/StreamMessage.tsx         | 142 ++++++++++-------
 src/types/enhanced-messages.ts           | 143 +++++++++++++++++
 6 files changed, 469 insertions(+), 77 deletions(-)
 create mode 100644 src/components/CollapsibleToolResult.tsx
 create mode 100644 src/types/enhanced-messages.ts

diff --git a/src/components/AgentExecution.tsx b/src/components/AgentExecution.tsx
index 1655486..e057640 100644
--- a/src/components/AgentExecution.tsx
+++ b/src/components/AgentExecution.tsx
@@ -24,6 +24,7 @@ import { listen, type UnlistenFn } from "@tauri-apps/api/event";
 import { StreamMessage } from "./StreamMessage";
 import { ExecutionControlBar } from "./ExecutionControlBar";
 import { ErrorBoundary } from "./ErrorBoundary";
+import { enhanceMessages, type EnhancedMessage } from "@/types/enhanced-messages";
 
 interface AgentExecutionProps {
   /**
@@ -73,6 +74,7 @@ export const AgentExecution: React.FC<AgentExecutionProps> = ({
   const [model, setModel] = useState(agent.model || "sonnet");
   const [isRunning, setIsRunning] = useState(false);
   const [messages, setMessages] = useState<ClaudeStreamMessage[]>([]);
+  const [enhancedMessages, setEnhancedMessages] = useState<EnhancedMessage[]>([]);
   const [rawJsonlOutput, setRawJsonlOutput] = useState<string[]>([]);
   const [error, setError] = useState<string | null>(null);
   const [copyPopoverOpen, setCopyPopoverOpen] = useState(false);
@@ -159,6 +161,12 @@ export const AgentExecution: React.FC<AgentExecutionProps> = ({
     setTotalTokens(tokens);
   }, [messages]);
 
+  // Enhance messages whenever they change
+  useEffect(() => {
+    const enhanced = enhanceMessages(messages);
+    setEnhancedMessages(enhanced);
+  }, [messages]);
+
   const handleSelectPath = async () => {
     try {
       const selected = await open({
@@ -594,7 +602,7 @@ export const AgentExecution: React.FC<AgentExecutionProps> = ({
             }}
           >
             <div ref={messagesContainerRef}>
-            {messages.length === 0 && !isRunning && (
+            {enhancedMessages.length === 0 && !isRunning && (
               <div className="flex flex-col items-center justify-center h-full text-center">
                 <Terminal className="h-16 w-16 text-muted-foreground mb-4" />
                 <h3 className="text-lg font-medium mb-2">Ready to Execute</h3>
@@ -604,7 +612,7 @@ export const AgentExecution: React.FC<AgentExecutionProps> = ({
               </div>
             )}
 
-            {isRunning && messages.length === 0 && (
+            {isRunning && enhancedMessages.length === 0 && (
               <div className="flex items-center justify-center h-full">
                 <div className="flex items-center gap-3">
                   <Loader2 className="h-6 w-6 animate-spin" />
@@ -614,7 +622,7 @@ export const AgentExecution: React.FC<AgentExecutionProps> = ({
             )}
 
             <AnimatePresence>
-              {messages.map((message, index) => (
+              {enhancedMessages.map((message, index) => (
                 <motion.div
                   key={index}
                   initial={{ opacity: 0, y: 10 }}
@@ -623,7 +631,7 @@ export const AgentExecution: React.FC<AgentExecutionProps> = ({
                   className="mb-4"
                 >
                   <ErrorBoundary>
-                    <StreamMessage message={message} streamMessages={messages} />
+                    <StreamMessage message={message} streamMessages={enhancedMessages} />
                   </ErrorBoundary>
                 </motion.div>
               ))}
@@ -724,7 +732,7 @@ export const AgentExecution: React.FC<AgentExecutionProps> = ({
                 }
               }}
             >
-              {messages.length === 0 && !isRunning && (
+              {enhancedMessages.length === 0 && !isRunning && (
                 <div className="flex flex-col items-center justify-center h-full text-center">
                   <Terminal className="h-16 w-16 text-muted-foreground mb-4" />
                   <h3 className="text-lg font-medium mb-2">Ready to Execute</h3>
@@ -734,7 +742,7 @@ export const AgentExecution: React.FC<AgentExecutionProps> = ({
                 </div>
               )}
 
-              {isRunning && messages.length === 0 && (
+              {isRunning && enhancedMessages.length === 0 && (
                 <div className="flex items-center justify-center h-full">
                   <div className="flex items-center gap-3">
                     <Loader2 className="h-6 w-6 animate-spin" />
@@ -744,7 +752,7 @@ export const AgentExecution: React.FC<AgentExecutionProps> = ({
               )}
 
               <AnimatePresence>
-                {messages.map((message, index) => (
+                {enhancedMessages.map((message, index) => (
                   <motion.div
                     key={index}
                     initial={{ opacity: 0, y: 10 }}
@@ -753,7 +761,7 @@ export const AgentExecution: React.FC<AgentExecutionProps> = ({
                     className="mb-4"
                   >
                     <ErrorBoundary>
-                      <StreamMessage message={message} streamMessages={messages} />
+                      <StreamMessage message={message} streamMessages={enhancedMessages} />
                     </ErrorBoundary>
                   </motion.div>
                 ))}
diff --git a/src/components/ClaudeCodeSession.tsx b/src/components/ClaudeCodeSession.tsx
index ce16785..a81cb52 100644
--- a/src/components/ClaudeCodeSession.tsx
+++ b/src/components/ClaudeCodeSession.tsx
@@ -26,6 +26,7 @@ import { TimelineNavigator } from "./TimelineNavigator";
 import { CheckpointSettings } from "./CheckpointSettings";
 import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogDescription, DialogFooter } from "@/components/ui/dialog";
 import type { ClaudeStreamMessage } from "./AgentExecution";
+import { enhanceMessages, type EnhancedMessage } from "@/types/enhanced-messages";
 
 interface ClaudeCodeSessionProps {
   /**
@@ -60,6 +61,7 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
 }) => {
   const [projectPath, setProjectPath] = useState(initialProjectPath || session?.project_path || "");
   const [messages, setMessages] = useState<ClaudeStreamMessage[]>([]);
+  const [enhancedMessages, setEnhancedMessages] = useState<EnhancedMessage[]>([]);
   const [isLoading, setIsLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
   const [rawJsonlOutput, setRawJsonlOutput] = useState<string[]>([]);
@@ -115,10 +117,16 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
     }
   }, [session]);
 
+  // Enhance messages whenever they change
+  useEffect(() => {
+    const enhanced = enhanceMessages(messages);
+    setEnhancedMessages(enhanced);
+  }, [messages]);
+
   // Auto-scroll to bottom when new messages arrive
   useEffect(() => {
     messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
-  }, [messages]);
+  }, [enhancedMessages]);
 
   // Calculate total tokens from messages
   useEffect(() => {
@@ -513,7 +521,7 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
               </>
             )}
             
-            {messages.length > 0 && (
+            {enhancedMessages.length > 0 && (
               <Popover
                 trigger={
                   <Button
@@ -611,7 +619,7 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
 
         {/* Messages Display */}
         <div className="flex-1 overflow-y-auto p-4 space-y-2 pb-40">
-          {messages.length === 0 && !isLoading && (
+          {enhancedMessages.length === 0 && !isLoading && (
             <div className="flex flex-col items-center justify-center h-full text-center">
               <Terminal className="h-16 w-16 text-muted-foreground mb-4" />
               <h3 className="text-lg font-medium mb-2">Ready to Start</h3>
@@ -624,7 +632,7 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
             </div>
           )}
 
-          {isLoading && messages.length === 0 && (
+          {isLoading && enhancedMessages.length === 0 && (
             <div className="flex items-center justify-center h-full">
               <div className="flex items-center gap-3">
                 <Loader2 className="h-6 w-6 animate-spin" />
@@ -636,7 +644,7 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
           )}
 
           <AnimatePresence>
-            {messages.map((message, index) => (
+            {enhancedMessages.map((message, index) => (
               <motion.div
                 key={index}
                 initial={{ opacity: 0, y: 10 }}
@@ -644,14 +652,14 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
                 transition={{ duration: 0.2 }}
               >
                 <ErrorBoundary>
-                  <StreamMessage message={message} streamMessages={messages} />
+                  <StreamMessage message={message} streamMessages={enhancedMessages} />
                 </ErrorBoundary>
               </motion.div>
             ))}
           </AnimatePresence>
           
           {/* Show loading indicator when processing, even if there are messages */}
-          {isLoading && messages.length > 0 && (
+          {isLoading && enhancedMessages.length > 0 && (
             <div className="flex items-center gap-2 p-4">
               <Loader2 className="h-4 w-4 animate-spin" />
               <span className="text-sm text-muted-foreground">Processing...</span>
diff --git a/src/components/CollapsibleToolResult.tsx b/src/components/CollapsibleToolResult.tsx
new file mode 100644
index 0000000..c904c96
--- /dev/null
+++ b/src/components/CollapsibleToolResult.tsx
@@ -0,0 +1,191 @@
+import React, { useState } from "react";
+import { motion, AnimatePresence } from "framer-motion";
+import { 
+  ChevronDown, 
+  ChevronRight, 
+  Loader2, 
+  CheckCircle2, 
+  AlertCircle,
+  Terminal,
+  FileText,
+  Search,
+  Edit,
+  FolderOpen,
+  Code
+} from "lucide-react";
+import { cn } from "@/lib/utils";
+import type { ToolCall, ToolResult } from "@/types/enhanced-messages";
+
+interface CollapsibleToolResultProps {
+  toolCall: ToolCall;
+  toolResult?: ToolResult;
+  className?: string;
+  children?: React.ReactNode;
+}
+
+// Map tool names to icons
+const toolIcons: Record<string, React.ReactNode> = {
+  read: <FileText className="h-4 w-4" />,
+  write: <Edit className="h-4 w-4" />,
+  edit: <Edit className="h-4 w-4" />,
+  multiedit: <Edit className="h-4 w-4" />,
+  bash: <Terminal className="h-4 w-4" />,
+  ls: <FolderOpen className="h-4 w-4" />,
+  glob: <Search className="h-4 w-4" />,
+  grep: <Search className="h-4 w-4" />,
+  task: <Code className="h-4 w-4" />,
+  default: <Terminal className="h-4 w-4" />
+};
+
+// Get tool icon based on tool name
+function getToolIcon(toolName: string): React.ReactNode {
+  const lowerName = toolName.toLowerCase();
+  return toolIcons[lowerName] || toolIcons.default;
+}
+
+// Get display name for tools
+function getToolDisplayName(toolName: string): string {
+  const displayNames: Record<string, string> = {
+    ls: "List directory",
+    read: "Read file",
+    write: "Write file",
+    edit: "Edit file",
+    multiedit: "Multi-edit file",
+    bash: "Run command",
+    glob: "Find files",
+    grep: "Search files",
+    task: "Run task",
+    todowrite: "Update todos",
+    todoread: "Read todos",
+    websearch: "Search web",
+    webfetch: "Fetch webpage"
+  };
+  
+  const lowerName = toolName.toLowerCase();
+  return displayNames[lowerName] || toolName;
+}
+
+// Get a brief description of the tool call
+function getToolDescription(toolCall: ToolCall): string {
+  const name = toolCall.name.toLowerCase();
+  const input = toolCall.input;
+  
+  switch (name) {
+    case "read":
+      return input?.file_path ? `${input.file_path}` : "Reading file";
+    case "write":
+      return input?.file_path ? `${input.file_path}` : "Writing file";
+    case "edit":
+    case "multiedit":
+      return input?.file_path ? `${input.file_path}` : "Editing file";
+    case "bash":
+      return input?.command ? `${input.command}` : "Running command";
+    case "ls":
+      return input?.path ? `${input.path}` : "Listing directory";
+    case "glob":
+      return input?.pattern ? `${input.pattern}` : "Finding files";
+    case "grep":
+      return input?.pattern ? `${input.pattern}` : "Searching files";
+    case "task":
+      return input?.description || "Running task";
+    default:
+      return toolCall.name;
+  }
+}
+
+export const CollapsibleToolResult: React.FC<CollapsibleToolResultProps> = ({
+  toolCall,
+  toolResult,
+  className,
+  children
+}) => {
+  const [isExpanded, setIsExpanded] = useState(false);
+  const isPending = !toolResult;
+  const isError = toolResult?.isError;
+  
+  return (
+    <div className={cn("space-y-2", className)}>
+      {/* Tool Call Header */}
+      <div
+        className={cn(
+          "flex items-center gap-2 p-2 rounded-md border cursor-pointer transition-colors",
+          "hover:bg-muted/50",
+          isPending && "border-muted-foreground/20",
+          !isPending && !isError && "border-green-500/20",
+          isError && "border-destructive/20"
+        )}
+        onClick={() => setIsExpanded(!isExpanded)}
+      >
+        {/* Expand/Collapse Icon */}
+        <motion.div
+          animate={{ rotate: isExpanded ? 90 : 0 }}
+          transition={{ duration: 0.2 }}
+        >
+          <ChevronRight className="h-3 w-3 text-muted-foreground" />
+        </motion.div>
+        
+        {/* Tool Icon */}
+        <div className="text-muted-foreground">
+          {getToolIcon(toolCall.name)}
+        </div>
+        
+        {/* Tool Name */}
+        <span className="text-sm font-medium">
+          {getToolDisplayName(toolCall.name)}
+        </span>
+        
+        {/* Tool Description */}
+        <span className="text-xs text-muted-foreground flex-1 truncate">
+          {getToolDescription(toolCall)}
+        </span>
+        
+        {/* Status Icon */}
+        <div className="ml-auto">
+          {isPending ? (
+            <Loader2 className="h-4 w-4 animate-spin text-muted-foreground" />
+          ) : isError ? (
+            <AlertCircle className="h-4 w-4 text-destructive" />
+          ) : (
+            <CheckCircle2 className="h-4 w-4 text-green-500" />
+          )}
+        </div>
+      </div>
+      
+      {/* Tool Result (collapsible) */}
+      <AnimatePresence>
+        {isExpanded && toolResult && (
+          <motion.div
+            initial={{ height: 0, opacity: 0 }}
+            animate={{ height: "auto", opacity: 1 }}
+            exit={{ height: 0, opacity: 0 }}
+            transition={{ duration: 0.2 }}
+            className="overflow-hidden"
+          >
+            <div className={cn(
+              "ml-6 p-2 rounded-md border",
+              isError ? "border-destructive/20 bg-destructive/5" : "border-green-500/20 bg-green-500/5"
+            )}>
+              <div className="flex items-center gap-2 mb-2">
+                {isError ? (
+                  <AlertCircle className="h-4 w-4 text-destructive" />
+                ) : (
+                  <CheckCircle2 className="h-4 w-4 text-green-500" />
+                )}
+                <span className="text-sm font-medium">
+                  {isError ? "Tool Error" : "Tool Result"}
+                </span>
+              </div>
+              
+              {/* Result Content */}
+              <div className="text-xs font-mono overflow-x-auto whitespace-pre-wrap">
+                {typeof toolResult.content === 'string' 
+                  ? toolResult.content 
+                  : JSON.stringify(toolResult.content, null, 2)}
+              </div>
+            </div>
+          </motion.div>
+        )}
+      </AnimatePresence>
+    </div>
+  );
+};
\ No newline at end of file
diff --git a/src/components/SessionOutputViewer.tsx b/src/components/SessionOutputViewer.tsx
index d36d0e8..ff601fd 100644
--- a/src/components/SessionOutputViewer.tsx
+++ b/src/components/SessionOutputViewer.tsx
@@ -12,6 +12,7 @@ import type { AgentRun } from '@/lib/api';
 import { listen, type UnlistenFn } from '@tauri-apps/api/event';
 import { StreamMessage } from './StreamMessage';
 import { ErrorBoundary } from './ErrorBoundary';
+import { enhanceMessages, type EnhancedMessage } from '@/types/enhanced-messages';
 
 interface SessionOutputViewerProps {
   session: AgentRun;
@@ -39,6 +40,7 @@ export interface ClaudeStreamMessage {
 
 export function SessionOutputViewer({ session, onClose, className }: SessionOutputViewerProps) {
   const [messages, setMessages] = useState<ClaudeStreamMessage[]>([]);
+  const [enhancedMessages, setEnhancedMessages] = useState<EnhancedMessage[]>([]);
   const [rawJsonlOutput, setRawJsonlOutput] = useState<string[]>([]);
   const [loading, setLoading] = useState(false);
   const [isFullscreen, setIsFullscreen] = useState(false);
@@ -89,6 +91,12 @@ export function SessionOutputViewer({ session, onClose, className }: SessionOutp
     }
   }, [messages, hasUserScrolled, isFullscreen]);
 
+  // Enhance messages whenever they change
+  useEffect(() => {
+    const enhanced = enhanceMessages(messages);
+    setEnhancedMessages(enhanced);
+  }, [messages]);
+
   const loadOutput = async (skipCache = false) => {
     if (!session.id) return;
 
@@ -309,13 +317,13 @@ export function SessionOutputViewer({ session, onClose, className }: SessionOutp
                       </Badge>
                     )}
                     <span className="text-xs text-muted-foreground">
-                      {messages.length} messages
+                      {enhancedMessages.length} messages
                     </span>
                   </div>
                 </div>
               </div>
               <div className="flex items-center space-x-2">
-                {messages.length > 0 && (
+                {enhancedMessages.length > 0 && (
                   <>
                     <Button
                       variant="outline"
@@ -402,7 +410,7 @@ export function SessionOutputViewer({ session, onClose, className }: SessionOutp
                   }
                 }}
               >
-                {messages.length === 0 ? (
+                {enhancedMessages.length === 0 ? (
                   <div className="flex flex-col items-center justify-center h-full text-center">
                     {session.status === 'running' ? (
                       <>
@@ -431,7 +439,7 @@ export function SessionOutputViewer({ session, onClose, className }: SessionOutp
                 ) : (
                   <>
                     <AnimatePresence>
-                      {messages.map((message, index) => (
+                      {enhancedMessages.map((message, index) => (
                         <motion.div
                           key={index}
                           initial={{ opacity: 0, y: 10 }}
@@ -439,7 +447,7 @@ export function SessionOutputViewer({ session, onClose, className }: SessionOutp
                           transition={{ duration: 0.2 }}
                         >
                           <ErrorBoundary>
-                            <StreamMessage message={message} streamMessages={messages} />
+                            <StreamMessage message={message} streamMessages={enhancedMessages} />
                           </ErrorBoundary>
                         </motion.div>
                       ))}
@@ -536,7 +544,7 @@ export function SessionOutputViewer({ session, onClose, className }: SessionOutp
                 }
               }}
             >
-              {messages.length === 0 ? (
+              {enhancedMessages.length === 0 ? (
                 <div className="flex flex-col items-center justify-center h-full text-center">
                   {session.status === 'running' ? (
                     <>
@@ -555,7 +563,7 @@ export function SessionOutputViewer({ session, onClose, className }: SessionOutp
               ) : (
                 <>
                   <AnimatePresence>
-                    {messages.map((message, index) => (
+                    {enhancedMessages.map((message, index) => (
                       <motion.div
                         key={index}
                         initial={{ opacity: 0, y: 10 }}
@@ -563,7 +571,7 @@ export function SessionOutputViewer({ session, onClose, className }: SessionOutp
                         transition={{ duration: 0.2 }}
                       >
                         <ErrorBoundary>
-                          <StreamMessage message={message} streamMessages={messages} />
+                          <StreamMessage message={message} streamMessages={enhancedMessages} />
                         </ErrorBoundary>
                       </motion.div>
                     ))}
diff --git a/src/components/StreamMessage.tsx b/src/components/StreamMessage.tsx
index 096dd30..7eb048a 100644
--- a/src/components/StreamMessage.tsx
+++ b/src/components/StreamMessage.tsx
@@ -13,6 +13,8 @@ import remarkGfm from "remark-gfm";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { claudeSyntaxTheme } from "@/lib/claudeSyntaxTheme";
 import type { ClaudeStreamMessage } from "./AgentExecution";
+import { CollapsibleToolResult } from "./CollapsibleToolResult";
+import type { EnhancedMessage } from "@/types/enhanced-messages";
 import {
   TodoWidget,
   LSWidget,
@@ -37,9 +39,9 @@ import {
 } from "./ToolWidgets";
 
 interface StreamMessageProps {
-  message: ClaudeStreamMessage;
+  message: ClaudeStreamMessage | EnhancedMessage;
   className?: string;
-  streamMessages: ClaudeStreamMessage[];
+  streamMessages: (ClaudeStreamMessage | EnhancedMessage)[];
 }
 
 /**
@@ -72,6 +74,9 @@ export const StreamMessage: React.FC<StreamMessageProps> = ({ message, className
     // Assistant message
     if (message.type === "assistant" && message.message) {
       const msg = message.message;
+      const enhancedMsg = message as EnhancedMessage;
+      const hasToolCalls = enhancedMsg.toolCalls && enhancedMsg.toolCalls.length > 0;
+      
       return (
         <Card className={cn("border-primary/20 bg-primary/5", className)}>
           <CardContent className="p-4">
@@ -121,62 +126,91 @@ export const StreamMessage: React.FC<StreamMessageProps> = ({ message, className
                     const toolName = content.name?.toLowerCase();
                     const input = content.input;
                     
-                    // Task tool - for sub-agent tasks
-                    if (toolName === "task" && input) {
-                      return <TaskWidget key={idx} description={input.description} prompt={input.prompt} />;
+                    // Function to render the appropriate tool widget
+                    const renderToolWidget = () => {
+                      // Task tool - for sub-agent tasks
+                      if (toolName === "task" && input) {
+                        return <TaskWidget description={input.description} prompt={input.prompt} />;
+                      }
+                      
+                      // Edit tool
+                      if (toolName === "edit" && input?.file_path) {
+                        return <EditWidget {...input} />;
+                      }
+                      
+                      // MultiEdit tool
+                      if (toolName === "multiedit" && input?.file_path && input?.edits) {
+                        return <MultiEditWidget {...input} />;
+                      }
+                      
+                      // MCP tools (starting with mcp__)
+                      if (content.name?.startsWith("mcp__")) {
+                        return <MCPWidget toolName={content.name} input={input} />;
+                      }
+                      
+                      // TodoWrite tool
+                      if (toolName === "todowrite" && input?.todos) {
+                        return <TodoWidget todos={input.todos} />;
+                      }
+                      
+                      // LS tool
+                      if (toolName === "ls" && input?.path) {
+                        return <LSWidget path={input.path} />;
+                      }
+                      
+                      // Read tool
+                      if (toolName === "read" && input?.file_path) {
+                        return <ReadWidget filePath={input.file_path} />;
+                      }
+                      
+                      // Glob tool
+                      if (toolName === "glob" && input?.pattern) {
+                        return <GlobWidget pattern={input.pattern} />;
+                      }
+                      
+                      // Bash tool
+                      if (toolName === "bash" && input?.command) {
+                        return <BashWidget command={input.command} description={input.description} />;
+                      }
+                      
+                      // Write tool
+                      if (toolName === "write" && input?.file_path && input?.content) {
+                        return <WriteWidget filePath={input.file_path} content={input.content} />;
+                      }
+                      
+                      // Grep tool
+                      if (toolName === "grep" && input?.pattern) {
+                        return <GrepWidget pattern={input.pattern} include={input.include} path={input.path} exclude={input.exclude} />;
+                      }
+                      
+                      // Default - return null, will be handled by CollapsibleToolResult
+                      return null;
+                    };
+                    
+                    // Check if we have enhanced message with tool results
+                    if (hasToolCalls && enhancedMsg.toolResults && content.id) {
+                      const toolCall = enhancedMsg.toolCalls?.find(tc => tc.id === content.id);
+                      const toolResult = enhancedMsg.toolResults.get(content.id);
+                      
+                      if (toolCall && toolResult) {
+                        // Only use collapsible widget when we have both tool call AND result
+                        return (
+                          <CollapsibleToolResult
+                            key={idx}
+                            toolCall={toolCall}
+                            toolResult={toolResult}
+                          />
+                        );
+                      }
                     }
                     
-                    // Edit tool
-                    if (toolName === "edit" && input?.file_path) {
-                      return <EditWidget key={idx} {...input} />;
+                    // Render the normal tool widget (for pending tool calls or non-enhanced messages)
+                    const widget = renderToolWidget();
+                    if (widget) {
+                      return <div key={idx}>{widget}</div>;
                     }
                     
-                    // MultiEdit tool
-                    if (toolName === "multiedit" && input?.file_path && input?.edits) {
-                      return <MultiEditWidget key={idx} {...input} />;
-                    }
-                    
-                    // MCP tools (starting with mcp__)
-                    if (content.name?.startsWith("mcp__")) {
-                      return <MCPWidget key={idx} toolName={content.name} input={input} />;
-                    }
-                    
-                    // TodoWrite tool
-                    if (toolName === "todowrite" && input?.todos) {
-                      return <TodoWidget key={idx} todos={input.todos} />;
-                    }
-                    
-                    // LS tool
-                    if (toolName === "ls" && input?.path) {
-                      return <LSWidget key={idx} path={input.path} />;
-                    }
-                    
-                    // Read tool
-                    if (toolName === "read" && input?.file_path) {
-                      return <ReadWidget key={idx} filePath={input.file_path} />;
-                    }
-                    
-                    // Glob tool
-                    if (toolName === "glob" && input?.pattern) {
-                      return <GlobWidget key={idx} pattern={input.pattern} />;
-                    }
-                    
-                    // Bash tool
-                    if (toolName === "bash" && input?.command) {
-                      return <BashWidget key={idx} command={input.command} description={input.description} />;
-                    }
-                    
-                    // Write tool
-                    if (toolName === "write" && input?.file_path && input?.content) {
-                      return <WriteWidget key={idx} filePath={input.file_path} content={input.content} />;
-                    }
-                    
-                    // Grep tool
-                    if (toolName === "grep" && input?.pattern) {
-                      return <GrepWidget key={idx} pattern={input.pattern} include={input.include} path={input.path} exclude={input.exclude} />;
-                    }
-                    
-                    // Default tool display
+                    // Fallback to basic tool display
                     return (
                       <div key={idx} className="space-y-2">
                         <div className="flex items-center gap-2">
diff --git a/src/types/enhanced-messages.ts b/src/types/enhanced-messages.ts
new file mode 100644
index 0000000..02079db
--- /dev/null
+++ b/src/types/enhanced-messages.ts
@@ -0,0 +1,143 @@
+// Enhanced message types that map tool calls with their results
+export interface ToolCall {
+  id: string;
+  name: string;
+  input: any;
+  timestamp?: number;
+}
+
+export interface ToolResult {
+  toolUseId: string;
+  content: any;
+  isError?: boolean;
+  timestamp?: number;
+}
+
+export interface EnhancedMessage {
+  type: "system" | "assistant" | "user" | "result";
+  subtype?: string;
+  message?: {
+    content?: any[];
+    usage?: {
+      input_tokens: number;
+      output_tokens: number;
+    };
+  };
+  usage?: {
+    input_tokens: number;
+    output_tokens: number;
+  };
+  // Enhanced fields for tool call mapping
+  toolCalls?: ToolCall[];
+  toolResults?: Map<string, ToolResult>;
+  [key: string]: any;
+}
+
+// Helper function to extract tool calls from assistant messages
+export function extractToolCalls(message: any): ToolCall[] {
+  const toolCalls: ToolCall[] = [];
+  
+  if (message.type === "assistant" && message.message?.content && Array.isArray(message.message.content)) {
+    for (const content of message.message.content) {
+      if (content.type === "tool_use" && content.id) {
+        toolCalls.push({
+          id: content.id,
+          name: content.name || "unknown",
+          input: content.input,
+          timestamp: Date.now()
+        });
+      }
+    }
+  }
+  
+  return toolCalls;
+}
+
+// Helper function to extract tool results from user messages
+export function extractToolResult(message: any): ToolResult | null {
+  if (message.type === "user" && message.message?.content && Array.isArray(message.message.content)) {
+    for (const content of message.message.content) {
+      if (content.type === "tool_result" && content.tool_use_id) {
+        return {
+          toolUseId: content.tool_use_id,
+          content: content.content,
+          isError: content.is_error || false,
+          timestamp: Date.now()
+        };
+      }
+    }
+  }
+  
+  return null;
+}
+
+// Function to enhance messages with tool call/result mapping
+export function enhanceMessages(rawMessages: any[]): EnhancedMessage[] {
+  const enhanced: EnhancedMessage[] = [];
+  
+  // First pass: create enhanced messages and collect all tool calls
+  const toolCallMap = new Map<string, { message: EnhancedMessage, toolCall: ToolCall }>();
+  
+  for (let i = 0; i < rawMessages.length; i++) {
+    const message = rawMessages[i];
+    const enhancedMessage: EnhancedMessage = { ...message };
+    
+    // Extract tool calls from assistant messages
+    const toolCalls = extractToolCalls(message);
+    if (toolCalls.length > 0) {
+      enhancedMessage.toolCalls = toolCalls;
+      enhancedMessage.toolResults = new Map();
+      
+      // Store reference to tool calls for later mapping
+      for (const toolCall of toolCalls) {
+        toolCallMap.set(toolCall.id, { message: enhancedMessage, toolCall });
+      }
+    }
+    
+    enhanced.push(enhancedMessage);
+  }
+  
+  // Second pass: extract tool results and attach them to corresponding tool calls
+  for (let i = 0; i < rawMessages.length; i++) {
+    const message = rawMessages[i];
+    
+    // Extract tool results from user messages
+    if (message.type === "user" && message.message?.content && Array.isArray(message.message.content)) {
+      let hasOnlyMappedToolResults = true;
+      let hasAnyContent = false;
+      
+      for (const content of message.message.content) {
+        if (content.type === "tool_result" && content.tool_use_id) {
+          hasAnyContent = true;
+          const toolCallInfo = toolCallMap.get(content.tool_use_id);
+          if (toolCallInfo) {
+            // Create tool result
+            const toolResult: ToolResult = {
+              toolUseId: content.tool_use_id,
+              content: content.content,
+              isError: content.is_error || false,
+              timestamp: Date.now()
+            };
+            
+            // Attach result to the assistant message that contains the tool call
+            toolCallInfo.message.toolResults?.set(content.tool_use_id, toolResult);
+          } else {
+            // This tool result doesn't have a matching tool call
+            hasOnlyMappedToolResults = false;
+          }
+        } else if (content.type !== "tool_result") {
+          // This message has non-tool-result content
+          hasOnlyMappedToolResults = false;
+          hasAnyContent = true;
+        }
+      }
+      
+      // Only mark as meta if the message contains ONLY tool results that have been mapped
+      if (hasAnyContent && hasOnlyMappedToolResults) {
+        enhanced[i].isMeta = true;
+      }
+    }
+  }
+  
+  return enhanced;
+}
\ No newline at end of file