feat(agent): Stream agent messages, sources, tool calls, etc.

.github/copilot-instructions.md

@@ -149,3 +149,4 @@ When working on this codebase, you might need to:
 - You can use the context7 tool to get help using the following identifiers for libraries used in this project
   - `/langchain-ai/langchainjs` for LangChain
   - `/langchain-ai/langgraph` for LangGraph
+  - `/quantizor/markdown-to-jsx` for Markdown to JSX conversion

src/app/api/chat/route.ts

@@ -108,7 +108,7 @@ const handleEmitterEvents = async (
       writer.write(
         encoder.encode(
           JSON.stringify({
-            type: 'message',
+            type: 'response',
             data: parsedData.data,
             messageId: aiMessageId,
           }) + '\n',
@@ -138,20 +138,23 @@ const handleEmitterEvents = async (
       );
 
       sources = parsedData.data;
-    }
+    } else if (parsedData.type === 'tool_call') {
-  });
+      // Handle tool call events - stream them directly to the client AND accumulate for database
-
-  stream.on('agent_action', (data) => {
       writer.write(
         encoder.encode(
           JSON.stringify({
-          type: 'agent_action',
+            type: 'tool_call',
-          data: data.data,
+            data: parsedData.data,
-          messageId: userMessageId,
+            messageId: aiMessageId,
           }) + '\n',
         ),
       );
+
+      // Add tool call content to the received message for database storage
+      recievedMessage += parsedData.data.content;
+    }
   });
+
   let modelStats: ModelStats = {
     modelName: '',
   };

src/components/AgentActionDisplay.tsx

@@ -1,263 +0,0 @@
-'use client';
-
-import { useState } from 'react';
-import { cn } from '@/lib/utils';
-import {
-  ChevronDown,
-  ChevronUp,
-  Bot,
-  Search,
-  Zap,
-  Microscope,
-  Ban,
-  CircleCheck,
-  ListPlus,
-} from 'lucide-react';
-import { AgentActionEvent } from './ChatWindow';
-
-interface AgentActionDisplayProps {
-  events: AgentActionEvent[];
-  messageId: string;
-  isLoading: boolean;
-}
-
-const AgentActionDisplay = ({
-  events,
-  messageId,
-  isLoading,
-}: AgentActionDisplayProps) => {
-  const [isExpanded, setIsExpanded] = useState(false);
-
-  // Get the most recent event for collapsed view
-  const latestEvent = events[events.length - 1];
-
-  // Common function to format action names
-  const formatActionName = (action: string) => {
-    return action.replace(/_/g, ' ').toLocaleLowerCase();
-  };
-
-  // Function to get appropriate icon based on action type
-  const getActionIcon = (action: string, size: number = 20) => {
-    switch (action) {
-      case 'ANALYZING_PREVIEW_CONTENT':
-        return <Search size={size} className="text-[#9C27B0]" />;
-      case 'PROCESSING_PREVIEW_CONTENT':
-        return <Zap size={size} className="text-[#9C27B0]" />;
-      case 'PROCEEDING_WITH_FULL_ANALYSIS':
-        return <Microscope size={size} className="text-[#9C27B0]" />;
-      case 'SKIPPING_IRRELEVANT_SOURCE':
-        return <Ban size={size} className="text-red-600 dark:text-red-500" />;
-      case 'CONTEXT_UPDATED':
-        return (
-          <ListPlus
-            size={size}
-            className="text-green-600 dark:text-green-500"
-          />
-        );
-      case 'INFORMATION_GATHERING_COMPLETE':
-        return (
-          <CircleCheck
-            size={size}
-            className="text-green-600 dark:text-green-500"
-          />
-        );
-      default:
-        return <Bot size={size} className="text-[#9C27B0]" />;
-    }
-  };
-
-  if (!latestEvent) {
-    return null;
-  }
-
-  return (
-    <div className="my-4 bg-light-secondary/50 dark:bg-dark-secondary/50 rounded-xl border border-light-200 dark:border-dark-200 overflow-hidden">
-      <button
-        onClick={() => setIsExpanded(!isExpanded)}
-        className="w-full flex items-center justify-between px-4 py-3 text-black/90 dark:text-white/90 hover:bg-light-200 dark:hover:bg-dark-200 transition duration-200"
-      >
-        <div className="flex items-center space-x-2">
-          {getActionIcon(latestEvent.action)}
-          <span className="font-medium text-base text-black/70 dark:text-white/70 tracking-wide capitalize flex items-center">
-            {!isLoading ||
-            latestEvent.action === 'INFORMATION_GATHERING_COMPLETE'
-              ? 'Agent Log'
-              : formatActionName(latestEvent.action)}
-            {/* {isLoading &&
-              latestEvent.action !== 'INFORMATION_GATHERING_COMPLETE' && (
-                <span className="ml-2 inline-block align-middle">
-                  <span className="animate-spin inline-block w-4 h-4 border-2 border-t-transparent border-[#9C27B0] rounded-full align-middle"></span>
-                </span>
-              )} */}
-          </span>
-        </div>
-        {isExpanded ? (
-          <ChevronUp size={18} className="text-black/70 dark:text-white/70" />
-        ) : (
-          <ChevronDown size={18} className="text-black/70 dark:text-white/70" />
-        )}
-      </button>
-
-      {isExpanded && (
-        <div className="px-4 py-3 text-black/80 dark:text-white/80 text-base border-t border-light-200 dark:border-dark-200 bg-light-100/50 dark:bg-dark-100/50">
-          <div className="space-y-3">
-            {events.map((event, index) => (
-              <div
-                key={`${messageId}-${index}-${event.action}`}
-                className="flex flex-col space-y-1 p-3 bg-white/50 dark:bg-black/20 rounded-lg border border-light-200/50 dark:border-dark-200/50"
-              >
-                <div className="flex items-center space-x-2">
-                  {getActionIcon(event.action, 16)}
-                  <span className="font-medium text-sm text-black/70 dark:text-white/70 capitalize tracking-wide">
-                    {formatActionName(event.action)}
-                  </span>
-                </div>
-
-                {event.message && event.message.length > 0 && (
-                  <p className="text-base">{event.message}</p>
-                )}
-
-                {/* Display relevant details based on event type */}
-                {event.details && Object.keys(event.details).length > 0 && (
-                  <div className="mt-2 text-sm text-black/60 dark:text-white/60">
-                    {event.details.sourceUrl && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Source:
-                        </span>
-                        <span className="truncate">
-                          <a href={event.details.sourceUrl} target="_blank">
-                            {event.details.sourceUrl}
-                          </a>
-                        </span>
-                      </div>
-                    )}
-                    {event.details.skipReason && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Reason:
-                        </span>
-                        <span>{event.details.skipReason}</span>
-                      </div>
-                    )}
-                    {event.details.searchQuery &&
-                      event.details.searchQuery !== event.details.query && (
-                        <div className="flex space-x-1">
-                          <span className="font-bold whitespace-nowrap">
-                            Search Query:
-                          </span>
-                          <span className="italic">
-                            &quot;{event.details.searchQuery}&quot;
-                          </span>
-                        </div>
-                      )}
-                    {event.details.sourcesFound !== undefined && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Sources Found:
-                        </span>
-                        <span>{event.details.sourcesFound}</span>
-                      </div>
-                    )}
-                    {/* {(event.details.documentCount !== undefined && event.details.documentCount > 0) && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">Documents:</span>
-                        <span>{event.details.documentCount}</span>
-                      </div>
-                    )} */}
-                    {event.details.contentLength !== undefined && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Content Length:
-                        </span>
-                        <span>{event.details.contentLength} characters</span>
-                      </div>
-                    )}
-                    {event.details.searchInstructions !== undefined && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Search Instructions:
-                        </span>
-                        <span>{event.details.searchInstructions}</span>
-                      </div>
-                    )}
-                    {/* {event.details.previewCount !== undefined && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">Preview Sources:</span>
-                        <span>{event.details.previewCount}</span>
-                      </div>
-                    )} */}
-                    {event.details.processingType && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Processing Type:
-                        </span>
-                        <span className="capitalize">
-                          {event.details.processingType.replace('-', ' ')}
-                        </span>
-                      </div>
-                    )}
-                    {event.details.insufficiencyReason && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Reason:
-                        </span>
-                        <span>{event.details.insufficiencyReason}</span>
-                      </div>
-                    )}
-                    {event.details.reason && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Reason:
-                        </span>
-                        <span>{event.details.reason}</span>
-                      </div>
-                    )}
-                    {/* {event.details.taskCount !== undefined && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">Tasks:</span>
-                        <span>{event.details.taskCount}</span>
-                      </div>
-                    )} */}
-                    {event.details.currentTask && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Current Task:
-                        </span>
-                        <span className="italic">
-                          &quot;{event.details.currentTask}&quot;
-                        </span>
-                      </div>
-                    )}
-                    {event.details.nextTask && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Next:
-                        </span>
-                        <span className="italic">
-                          &quot;{event.details.nextTask}&quot;
-                        </span>
-                      </div>
-                    )}
-                    {event.details.currentSearchFocus && (
-                      <div className="flex space-x-1">
-                        <span className="font-bold whitespace-nowrap">
-                          Search Focus:
-                        </span>
-                        <span className="italic">
-                          &quot;{event.details.currentSearchFocus}&quot;
-                        </span>
-                      </div>
-                    )}
-                  </div>
-                )}
-              </div>
-            ))}
-          </div>
-        </div>
-      )}
-    </div>
-  );
-};
-
-export default AgentActionDisplay;

src/components/Chat.tsx

@@ -5,7 +5,6 @@ import { File, Message } from './ChatWindow';
 import MessageBox from './MessageBox';
 import MessageBoxLoading from './MessageBoxLoading';
 import MessageInput from './MessageInput';
-import AgentActionDisplay from './AgentActionDisplay';
 
 const Chat = ({
   loading,
@@ -25,6 +24,7 @@ const Chat = ({
   analysisProgress,
   systemPromptIds,
   setSystemPromptIds,
+  onThinkBoxToggle,
 }: {
   messages: Message[];
   sendMessage: (
@@ -54,6 +54,11 @@ const Chat = ({
   } | null;
   systemPromptIds: string[];
   setSystemPromptIds: (ids: string[]) => void;
+  onThinkBoxToggle: (
+    messageId: string,
+    thinkBoxId: string,
+    expanded: boolean,
+  ) => void;
 }) => {
   const [isAtBottom, setIsAtBottom] = useState(true);
   const [manuallyScrolledUp, setManuallyScrolledUp] = useState(false);
@@ -224,30 +229,8 @@ const Chat = ({
               rewrite={rewrite}
               sendMessage={sendMessage}
               handleEditMessage={handleEditMessage}
+              onThinkBoxToggle={onThinkBoxToggle}
             />
-            {/* Show agent actions after user messages - either completed or in progress */}
-            {msg.role === 'user' && (
-              <>
-                {/* Show agent actions if they exist */}
-                {msg.agentActions && msg.agentActions.length > 0 && (
-                  <AgentActionDisplay
-                    messageId={msg.messageId}
-                    events={msg.agentActions}
-                    isLoading={loading}
-                  />
-                )}
-                {/* Show empty agent action display if this is the last user message and we're loading */}
-                {loading &&
-                  isLast &&
-                  (!msg.agentActions || msg.agentActions.length === 0) && (
-                    <AgentActionDisplay
-                      messageId={msg.messageId}
-                      events={[]}
-                      isLoading={loading}
-                    />
-                  )}
-              </>
-            )}
             {!isLast && msg.role === 'assistant' && (
               <div className="h-px w-full bg-light-secondary dark:bg-dark-secondary" />
             )}

src/components/ChatWindow.tsx

@@ -36,13 +36,13 @@ export type Message = {
   modelStats?: ModelStats;
   searchQuery?: string;
   searchUrl?: string;
-  agentActions?: AgentActionEvent[];
   progress?: {
     message: string;
     current: number;
     total: number;
     subMessage?: string;
   };
+  expandedThinkBoxes?: Set<string>;
 };
 
 export interface File {
@@ -467,33 +467,6 @@ const ChatWindow = ({ id }: { id?: string }) => {
         return;
       }
 
-      if (data.type === 'agent_action') {
-        const agentActionEvent: AgentActionEvent = {
-          action: data.data.action,
-          message: data.data.message,
-          details: data.data.details || {},
-          timestamp: new Date(),
-        };
-
-        // Update the user message with agent actions
-        setMessages((prev) =>
-          prev.map((message) => {
-            if (
-              message.messageId === data.messageId &&
-              message.role === 'user'
-            ) {
-              const updatedActions = [
-                ...(message.agentActions || []),
-                agentActionEvent,
-              ];
-              return { ...message, agentActions: updatedActions };
-            }
-            return message;
-          }),
-        );
-        return;
-      }
-
       if (data.type === 'sources') {
         sources = data.data;
         if (!added) {
@@ -512,23 +485,68 @@ const ChatWindow = ({ id }: { id?: string }) => {
           ]);
           added = true;
           setScrollTrigger((prev) => prev + 1);
+        } else {
+          // set the sources
+          setMessages((prev) =>
+            prev.map((message) => {
+              if (message.messageId === data.messageId) {
+                return { ...message, sources: sources };
+              }
+              return message;
+            }),
+          );
         }
       }
 
-      if (data.type === 'message') {
+      if (data.type === 'tool_call') {
+        // Add the tool content to the current assistant message (already formatted with newlines)
+        const toolContent = data.data.content;
+
+        if (!added) {
+          // Create initial message with tool content
+          setMessages((prevMessages) => [
+            ...prevMessages,
+            {
+              content: toolContent,
+              messageId: data.messageId, // Use the AI message ID from the backend
+              chatId: chatId!,
+              role: 'assistant',
+              sources: sources,
+              createdAt: new Date(),
+            },
+          ]);
+          added = true;
+        } else {
+          // Append tool content to existing message
+          setMessages((prev) =>
+            prev.map((message) => {
+              if (message.messageId === data.messageId) {
+                return {
+                  ...message,
+                  content: message.content + toolContent,
+                };
+              }
+              return message;
+            }),
+          );
+        }
+
+        recievedMessage += toolContent;
+        setScrollTrigger((prev) => prev + 1);
+        return;
+      }
+
+      if (data.type === 'response') {
         if (!added) {
           setMessages((prevMessages) => [
             ...prevMessages,
             {
               content: data.data,
-              messageId: data.messageId,
+              messageId: data.messageId, // Use the AI message ID from the backend
               chatId: chatId!,
               role: 'assistant',
               sources: sources,
               createdAt: new Date(),
-              modelStats: {
-                modelName: data.modelName,
-              },
             },
           ]);
           added = true;
@@ -703,6 +721,27 @@ const ChatWindow = ({ id }: { id?: string }) => {
     }
   };
 
+  const handleThinkBoxToggle = (
+    messageId: string,
+    thinkBoxId: string,
+    expanded: boolean,
+  ) => {
+    setMessages((prev) =>
+      prev.map((message) => {
+        if (message.messageId === messageId) {
+          const expandedThinkBoxes = new Set(message.expandedThinkBoxes || []);
+          if (expanded) {
+            expandedThinkBoxes.add(thinkBoxId);
+          } else {
+            expandedThinkBoxes.delete(thinkBoxId);
+          }
+          return { ...message, expandedThinkBoxes };
+        }
+        return message;
+      }),
+    );
+  };
+
   useEffect(() => {
     if (isReady && initialMessage && isConfigReady) {
       // Check if we have an initial query and apply saved search settings
@@ -788,6 +827,7 @@ const ChatWindow = ({ id }: { id?: string }) => {
               analysisProgress={analysisProgress}
               systemPromptIds={systemPromptIds}
               setSystemPromptIds={setSystemPromptIds}
+              onThinkBoxToggle={handleThinkBoxToggle}
             />
           </>
         ) : (

src/components/MarkdownRenderer.tsx

@@ -2,7 +2,7 @@
 'use client';
 
 import { cn } from '@/lib/utils';
-import { CheckCheck, Copy as CopyIcon, Brain } from 'lucide-react';
+import { CheckCheck, Copy as CopyIcon, Search, FileText, Globe, Settings } from 'lucide-react';
 import Markdown, { MarkdownToJSX } from 'markdown-to-jsx';
 import { useState } from 'react';
 import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
@@ -15,13 +15,13 @@ import ThinkBox from './ThinkBox';
 
 // Helper functions for think overlay
 const extractThinkContent = (content: string): string | null => {
-  const thinkRegex = /<think>([\s\S]*?)<\/think>/g;
+  const thinkRegex = /<think[^>]*>([\s\S]*?)<\/think>/g;
   const matches = content.match(thinkRegex);
   if (!matches) return null;
 
   // Extract content between think tags and join if multiple
   const extractedContent = matches
-    .map((match) => match.replace(/<\/?think>/g, ''))
+    .map((match) => match.replace(/<\/?think[^>]*>/g, ''))
     .join('\n\n');
 
   // Return null if content is empty or only whitespace
@@ -29,24 +29,150 @@ const extractThinkContent = (content: string): string | null => {
 };
 
 const removeThinkTags = (content: string): string => {
-  return content.replace(/<think>[\s\S]*?<\/think>/g, '').trim();
+  return content.replace(/<think[^>]*>[\s\S]*?<\/think>/g, '').trim();
+};
+
+// Add stable IDs to think tags if they don't already have them
+const addThinkBoxIds = (content: string): string => {
+  let thinkCounter = 0;
+  return content.replace(/<think(?![^>]*\sid=)/g, () => {
+    return `<think id="think-${thinkCounter++}"`;
+  });
+};
+
+interface MarkdownRendererProps {
+  content: string;
+  className?: string;
+  showThinking?: boolean;
+  messageId?: string;
+  expandedThinkBoxes?: Set<string>;
+  onThinkBoxToggle?: (
+    messageId: string,
+    thinkBoxId: string,
+    expanded: boolean,
+  ) => void;
+}
+
+// Custom ToolCall component for markdown
+const ToolCall = ({
+  type,
+  query,
+  urls,
+  count,
+  children,
+}: {
+  type?: string;
+  query?: string;
+  urls?: string;
+  count?: string;
+  children?: React.ReactNode;
+}) => {
+  const getIcon = (toolType: string) => {
+    switch (toolType) {
+      case 'search':
+      case 'web_search':
+        return (
+          <Search size={16} className="text-blue-600 dark:text-blue-400" />
+        );
+      case 'file':
+      case 'file_search':
+        return (
+          <FileText size={16} className="text-green-600 dark:text-green-400" />
+        );
+      case 'url':
+      case 'url_summarization':
+        return (
+          <Globe size={16} className="text-purple-600 dark:text-purple-400" />
+        );
+      default:
+        return (
+          <Settings size={16} className="text-gray-600 dark:text-gray-400" />
+        );
+    }
+  };
+
+  const formatToolMessage = () => {
+    if (type === 'search' || type === 'web_search') {
+      return (
+        <>
+          <span className="mr-2">{getIcon(type)}</span>
+          <span className="text-black/60 dark:text-white/60">Web search:</span>
+          <span className="ml-2 px-2 py-0.5 bg-black/5 dark:bg-white/5 rounded font-mono text-sm">
+            {query || children}
+          </span>
+        </>
+      );
+    }
+
+    if (type === 'file' || type === 'file_search') {
+      return (
+        <>
+          <span className="mr-2">{getIcon(type)}</span>
+          <span className="text-black/60 dark:text-white/60">File search:</span>
+          <span className="ml-2 px-2 py-0.5 bg-black/5 dark:bg-white/5 rounded font-mono text-sm">
+            {query || children}
+          </span>
+        </>
+      );
+    }
+
+    if (type === 'url' || type === 'url_summarization') {
+      const urlCount = count ? parseInt(count) : 1;
+      return (
+        <>
+          <span className="mr-2">{getIcon(type)}</span>
+          <span className="text-black/60 dark:text-white/60">
+            Analyzing {urlCount} web page{urlCount === 1 ? '' : 's'} for
+            additional details
+          </span>
+        </>
+      );
+    }
+
+    // Fallback for unknown tool types
+    return (
+      <>
+        <span className="mr-2">{getIcon(type || 'default')}</span>
+        <span className="text-black/60 dark:text-white/60">Using tool:</span>
+        <span className="ml-2 px-2 py-0.5 bg-black/5 dark:bg-white/5 rounded font-mono text-sm border">
+          {type || 'unknown'}
+        </span>
+      </>
+    );
+  };
+
+  return (
+    <div className="my-3 px-4 py-3 bg-gradient-to-r from-blue-50/50 to-purple-50/50 dark:from-blue-900/20 dark:to-purple-900/20 border border-blue-200/30 dark:border-blue-700/30 rounded-lg">
+      <div className="flex items-center text-sm font-medium">
+        {formatToolMessage()}
+      </div>
+    </div>
+  );
 };
 
 const ThinkTagProcessor = ({
   children,
-  isOverlayMode = false,
+  id,
+  isExpanded,
+  onToggle,
 }: {
   children: React.ReactNode;
-  isOverlayMode?: boolean;
+  id?: string;
+  isExpanded?: boolean;
+  onToggle?: (thinkBoxId: string, expanded: boolean) => void;
 }) => {
-  // In overlay mode, don't render anything (content will be handled by overlay)
+  return (
-  if (isOverlayMode) {
+    <ThinkBox 
-    return null;
+      content={children} 
+      expanded={isExpanded}
+      onToggle={() => {
+        if (id && onToggle) {
+          onToggle(id, !isExpanded);
         }
-  return <ThinkBox content={children} />;
+      }}
-};
+    />
-
+  );
-const CodeBlock = ({
+};const CodeBlock = ({
   className,
   children,
 }: {
@@ -115,31 +241,55 @@ const CodeBlock = ({
   );
 };
 
-interface MarkdownRendererProps {
-  content: string;
-  className?: string;
-  thinkOverlay?: boolean;
-}
-
 const MarkdownRenderer = ({
   content,
   className,
-  thinkOverlay = false,
+  showThinking = true,
+  messageId,
+  expandedThinkBoxes,
+  onThinkBoxToggle,
 }: MarkdownRendererProps) => {
-  const [showThinkBox, setShowThinkBox] = useState(false);
+  // Preprocess content to add stable IDs to think tags
+  const processedContent = addThinkBoxIds(content);
 
-  // Extract think content from the markdown
+  // Check if a think box is expanded
-  const thinkContent = thinkOverlay ? extractThinkContent(content) : null;
+  const isThinkBoxExpanded = (thinkBoxId: string) => {
-  const contentWithoutThink = thinkOverlay ? removeThinkTags(content) : content;
+    return expandedThinkBoxes?.has(thinkBoxId) || false;
+  };
+
+  // Handle think box toggle
+  const handleThinkBoxToggle = (thinkBoxId: string, expanded: boolean) => {
+    if (messageId && onThinkBoxToggle) {
+      onThinkBoxToggle(messageId, thinkBoxId, expanded);
+    }
+  };
+
+  // Determine what content to render based on showThinking parameter
+  const contentToRender = showThinking 
+    ? processedContent
+    : removeThinkTags(processedContent);
   // Markdown formatting options
   const markdownOverrides: MarkdownToJSX.Options = {
     overrides: {
+      ToolCall: {
+        component: ToolCall,
+      },
       think: {
-        component: ({ children }) => (
+        component: ({ children, id, ...props }) => {
-          <ThinkTagProcessor isOverlayMode={thinkOverlay}>
+          // Use the id from the HTML attribute
+          const thinkBoxId = id || 'think-unknown';
+          const isExpanded = isThinkBoxExpanded(thinkBoxId);
+
+          return (
+            <ThinkTagProcessor
+              id={thinkBoxId}
+              isExpanded={isExpanded}
+              onToggle={handleThinkBoxToggle}
+            >
               {children}
             </ThinkTagProcessor>
-        ),
+          );
+        },
       },
       code: {
         component: ({ className, children }) => {
@@ -181,17 +331,6 @@ const MarkdownRenderer = ({
 
   return (
     <div className="relative">
-      {/* Think box when expanded - shows above markdown */}
-      {thinkOverlay && thinkContent && showThinkBox && (
-        <div className="mb-4">
-          <ThinkBox
-            content={thinkContent}
-            expanded={true}
-            onToggle={() => setShowThinkBox(false)}
-          />
-        </div>
-      )}
-
       <Markdown
         className={cn(
           'prose prose-h1:mb-3 prose-h2:mb-2 prose-h2:mt-6 prose-h2:font-[800] prose-h3:mt-4 prose-h3:mb-1.5 prose-h3:font-[600] dark:prose-invert prose-p:leading-relaxed prose-pre:p-0 font-[400]',
@@ -203,22 +342,8 @@ const MarkdownRenderer = ({
         )}
         options={markdownOverrides}
       >
-        {thinkOverlay ? contentWithoutThink : content}
+        {contentToRender}
       </Markdown>
-
-      {/* Overlay icon when think box is collapsed */}
-      {thinkOverlay && thinkContent && !showThinkBox && (
-        <button
-          onClick={() => setShowThinkBox(true)}
-          className="absolute top-2 right-2 p-2 rounded-lg bg-black/20 dark:bg-white/20 backdrop-blur-sm opacity-30 hover:opacity-100 transition-opacity duration-200 group"
-          title="Show thinking process"
-        >
-          <Brain
-            size={16}
-            className="text-gray-700 dark:text-gray-300 group-hover:text-blue-600 dark:group-hover:text-blue-400 transition-colors"
-          />
-        </button>
-      )}
     </div>
   );
 };

src/components/MessageBox.tsx

@@ -13,6 +13,7 @@ const MessageBox = ({
   rewrite,
   sendMessage,
   handleEditMessage,
+  onThinkBoxToggle,
 }: {
   message: Message;
   messageIndex: number;
@@ -29,6 +30,11 @@ const MessageBox = ({
     },
   ) => void;
   handleEditMessage: (messageId: string, content: string) => void;
+  onThinkBoxToggle: (
+    messageId: string,
+    thinkBoxId: string,
+    expanded: boolean,
+  ) => void;
 }) => {
   // Local state for editing functionality
   const [isEditing, setIsEditing] = useState(false);
@@ -123,6 +129,7 @@ const MessageBox = ({
           loading={loading}
           rewrite={rewrite}
           sendMessage={sendMessage}
+          onThinkBoxToggle={onThinkBoxToggle}
         />
       )}
     </div>

src/components/MessageSources.tsx

@@ -4,74 +4,90 @@ import { File, Zap, Microscope, FileText, Sparkles } from 'lucide-react';
 
 const MessageSources = ({ sources }: { sources: Document[] }) => {
   return (
-    <div className="grid grid-cols-2 lg:grid-cols-4 gap-2">
+    <div className="flex flex-col space-y-3">
       {sources.map((source, i) => (
         <a
-          className="bg-light-100 hover:bg-light-200 dark:bg-dark-100 dark:hover:bg-dark-200 transition duration-200 rounded-lg p-3 flex flex-col space-y-2 font-medium"
+          className="bg-light-100 hover:bg-light-200 dark:bg-dark-100 dark:hover:bg-dark-200 transition duration-200 rounded-lg p-4 flex flex-row space-x-3 font-medium"
           key={i}
           href={source.metadata.url}
           target="_blank"
         >
-          <p className="dark:text-white text-xs overflow-hidden whitespace-nowrap text-ellipsis">
+          {/* Left side: Favicon/Icon and source number */}
-            {source.metadata.title}
+          <div className="flex flex-col items-center space-y-2 flex-shrink-0">
-          </p>
-          <div className="flex flex-row items-center justify-between">
-            <div className="flex flex-row items-center space-x-1">
             {source.metadata.url === 'File' ? (
-                <div className="bg-dark-200 hover:bg-dark-100 transition duration-200 flex items-center justify-center w-6 h-6 rounded-full">
+              <div className="bg-dark-200 hover:bg-dark-100 transition duration-200 flex items-center justify-center w-8 h-8 rounded-full">
-                  <File size={12} className="text-white/70" />
+                <File size={16} className="text-white/70" />
               </div>
             ) : (
               <img
                 src={`https://s2.googleusercontent.com/s2/favicons?domain_url=${source.metadata.url}`}
-                  width={16}
+                width={28}
-                  height={16}
+                height={28}
                 alt="favicon"
-                  className="rounded-lg h-4 w-4"
+                className="rounded-lg h-7 w-7"
               />
             )}
-              <p className="text-xs text-black/50 dark:text-white/50 overflow-hidden whitespace-nowrap text-ellipsis">
-                {source.metadata.url.replace(/.+\/\/|www.|\..+/g, '')}
-              </p>
-            </div>
             <div className="flex flex-row items-center space-x-1 text-black/50 dark:text-white/50 text-xs">
-              <div className="bg-black/50 dark:bg-white/50 h-[4px] w-[4px] rounded-full" />
+              <span className="font-semibold">{i + 1}</span>
-              <span>{i + 1}</span>
               {/* Processing type indicator */}
               {source.metadata.processingType === 'preview-only' && (
                 <span title="Partial content analyzed" className="inline-flex">
                   <Zap
-                    size={14}
+                    size={12}
-                    className="text-black/40 dark:text-white/40 ml-1"
+                    className="text-black/40 dark:text-white/40"
                   />
                 </span>
               )}
               {source.metadata.processingType === 'full-content' && (
                 <span title="Full content analyzed" className="inline-flex">
                   <Microscope
-                    size={14}
+                    size={12}
-                    className="text-black/40 dark:text-white/40 ml-1"
+                    className="text-black/40 dark:text-white/40"
                   />
                 </span>
               )}
               {source.metadata.processingType === 'url-direct-content' && (
                 <span title="Direct URL content" className="inline-flex">
                   <FileText
-                    size={14}
+                    size={12}
-                    className="text-black/40 dark:text-white/40 ml-1"
+                    className="text-black/40 dark:text-white/40"
                   />
                 </span>
               )}
               {source.metadata.processingType === 'url-content-extraction' && (
                 <span title="Summarized URL content" className="inline-flex">
                   <Sparkles
-                    size={14}
+                    size={12}
-                    className="text-black/40 dark:text-white/40 ml-1"
+                    className="text-black/40 dark:text-white/40"
                   />
                 </span>
               )}
             </div>
           </div>
+
+          {/* Right side: Content */}
+          <div className="flex-1 flex flex-col space-y-2">
+            {/* Title */}
+            <h3 className="dark:text-white text-sm font-semibold leading-tight">
+              {source.metadata.title}
+            </h3>
+            
+            {/* URL */}
+            <p className="text-xs text-black/50 dark:text-white/50">
+              {source.metadata.url.replace(/.+\/\/|www.|\..+/g, '')}
+            </p>
+
+            {/* Preview content */}
+            <p className="text-xs text-black/70 dark:text-white/70 leading-relaxed overflow-hidden" style={{ display: '-webkit-box', WebkitLineClamp: 3, WebkitBoxOrient: 'vertical' }}>
+              {/* Use snippet for preview-only content, otherwise use pageContent */}
+              {source.metadata.processingType === 'preview-only' && source.metadata.snippet
+                ? source.metadata.snippet
+                : source.pageContent?.length > 250 
+                  ? source.pageContent.slice(0, 250) + '...'
+                  : source.pageContent || 'No preview available'
+              }
+            </p>
+          </div>
         </a>
       ))}
     </div>

src/components/MessageTabs.tsx

@@ -43,6 +43,11 @@ interface SearchTabsProps {
       suggestions?: string[];
     },
   ) => void;
+  onThinkBoxToggle: (
+    messageId: string,
+    thinkBoxId: string,
+    expanded: boolean,
+  ) => void;
 }
 
 const MessageTabs = ({
@@ -54,6 +59,7 @@ const MessageTabs = ({
   loading,
   rewrite,
   sendMessage,
+  onThinkBoxToggle,
 }: SearchTabsProps) => {
   const [activeTab, setActiveTab] = useState<TabType>('text');
   const [imageCount, setImageCount] = useState(0);
@@ -273,9 +279,14 @@ const MessageTabs = ({
         {/* Answer Tab */}
         {activeTab === 'text' && (
           <div className="flex flex-col space-y-4 animate-fadeIn">
-            <MarkdownRenderer content={parsedMessage} className="px-4" />
+            <MarkdownRenderer 
-
+              content={parsedMessage} 
-            {loading && isLast ? null : (
+              className="px-4" 
+              messageId={message.messageId}
+              expandedThinkBoxes={message.expandedThinkBoxes}
+              onThinkBoxToggle={onThinkBoxToggle}
+              showThinking={true}
+            />            {loading && isLast ? null : (
               <div className="flex flex-row items-center justify-between w-full text-black dark:text-white px-4 py-4">
                 <div className="flex flex-row items-center space-x-1">
                   <Rewrite rewrite={rewrite} messageId={message.messageId} />

src/components/ThinkBox.tsx

@@ -27,14 +27,14 @@ const ThinkBox = ({ content, expanded, onToggle }: ThinkBoxProps) => {
     <div className="my-4 bg-light-secondary/50 dark:bg-dark-secondary/50 rounded-xl border border-light-200 dark:border-dark-200 overflow-hidden">
       <button
         onClick={handleToggle}
-        className="w-full flex items-center justify-between px-4 py-1 text-black/90 dark:text-white/90 hover:bg-light-200 dark:hover:bg-dark-200 transition duration-200"
+        className="w-full flex items-center justify-between px-4 py-4 text-black/90 dark:text-white/90 hover:bg-light-200 dark:hover:bg-dark-200 transition duration-200"
       >
         <div className="flex items-center space-x-2">
           <BrainCircuit
             size={20}
             className="text-[#9C27B0] dark:text-[#CE93D8]"
           />
-          <p className="font-medium text-sm">Thinking Process</p>
+          <span className="font-medium text-sm">Thinking Process</span>
         </div>
         {isExpanded ? (
           <ChevronUp size={18} className="text-black/70 dark:text-white/70" />

src/components/dashboard/WidgetConfigModal.tsx

@@ -6,8 +6,9 @@ import {
   DialogTitle,
   Transition,
   TransitionChild,
+  Switch,
 } from '@headlessui/react';
-import { X, Plus, Trash2, Play, Save } from 'lucide-react';
+import { X, Plus, Trash2, Play, Save, Brain } from 'lucide-react';
 import { Fragment, useState, useEffect } from 'react';
 import MarkdownRenderer from '@/components/MarkdownRenderer';
 import ModelSelector from '@/components/MessageInputActions/ModelSelector';
@@ -72,6 +73,7 @@ const WidgetConfigModal = ({
     model: string;
   } | null>(null);
   const [selectedTools, setSelectedTools] = useState<string[]>([]);
+  const [showThinking, setShowThinking] = useState(false);
 
   // Update config when editingWidget changes
   useEffect(() => {
@@ -415,6 +417,25 @@ const WidgetConfigModal = ({
                       <h4 className="text-sm font-medium text-black dark:text-white">
                         Preview
                       </h4>
+                      <div className="flex items-center gap-2">
+                        <div className="flex items-center gap-2">
+                          <Brain size={16} className="text-gray-600 dark:text-gray-400" />
+                          <span className="text-sm text-gray-700 dark:text-gray-300">Thinking</span>
+                          <Switch
+                            checked={showThinking}
+                            onChange={setShowThinking}
+                            className="bg-light-secondary dark:bg-dark-secondary border border-light-200/70 dark:border-dark-200 relative inline-flex h-5 w-10 sm:h-6 sm:w-11 items-center rounded-full"
+                          >
+                            <span className="sr-only">Show thinking tags</span>
+                            <span
+                              className={`${
+                                showThinking
+                                  ? 'translate-x-6 bg-purple-600'
+                                  : 'translate-x-1 bg-black/50 dark:bg-white/50'
+                              } inline-block h-3 w-3 sm:h-4 sm:w-4 transform rounded-full transition-all duration-200`}
+                            />
+                          </Switch>
+                        </div>
                         <button
                           onClick={handlePreview}
                           disabled={isPreviewLoading}
@@ -424,12 +445,13 @@ const WidgetConfigModal = ({
                           {isPreviewLoading ? 'Loading...' : 'Run Preview'}
                         </button>
                       </div>
+                    </div>
 
                     <div className="h-80 p-4 border border-light-200 dark:border-dark-200 rounded-md bg-light-secondary dark:bg-dark-secondary overflow-y-auto max-w-full">
                       {previewContent ? (
                         <div className="prose prose-sm dark:prose-invert max-w-full">
                           <MarkdownRenderer
-                            thinkOverlay={true}
+                            showThinking={showThinking}
                             content={previewContent}
                           />
                         </div>

src/components/dashboard/WidgetDisplay.tsx

@@ -119,7 +119,7 @@ const WidgetDisplay = ({
             </div>
           ) : widget.content ? (
             <div className="prose prose-sm dark:prose-invert max-w-none">
-              <MarkdownRenderer content={widget.content} thinkOverlay={true} />
+              <MarkdownRenderer content={widget.content} showThinking={false} />
             </div>
           ) : (
             <div className="flex items-center justify-center py-8 text-gray-500 dark:text-gray-400">

src/lib/search/agentSearch.ts

@@ -47,19 +47,6 @@ export class AgentSearch {
   ): Promise<void> {
     console.log('AgentSearch: Using simplified agent implementation');
 
-    // Emit agent action to indicate simplified agent usage
-    this.emitter.emit(
-      'data',
-      JSON.stringify({
-        type: 'agent_action',
-        data: {
-          action: 'agent_implementation_selection',
-          message: 'Using simplified agent implementation (experimental)',
-          details: `Focus mode: ${this.focusMode}, Files: ${fileIds.length}`,
-        },
-      }),
-    );
-
     // Delegate to simplified agent with focus mode
     await this.simplifiedAgent.searchAndAnswer(
       query,

src/lib/search/simplifiedAgent.ts

@@ -4,6 +4,7 @@ import {
   BaseMessage,
   HumanMessage,
   SystemMessage,
+  AIMessage,
 } from '@langchain/core/messages';
 import { Embeddings } from '@langchain/core/embeddings';
 import { EventEmitter } from 'events';
@@ -17,6 +18,7 @@ import {
 } from '@/lib/tools/agents';
 import { formatDateForLLM } from '../utils';
 import { getModelName } from '../utils/modelUtils';
+import { removeThinkingBlocks } from '../utils/contentUtils';
 
 /**
  * Simplified Agent using createReactAgent
@@ -451,19 +453,6 @@ Use all available tools strategically to provide comprehensive, well-researched,
       // Initialize agent with the provided focus mode and file context
       const agent = this.initializeAgent(focusMode, fileIds);
 
-      // Emit initial agent action
-      this.emitter.emit(
-        'data',
-        JSON.stringify({
-          type: 'agent_action',
-          data: {
-            action: 'simplified_agent_start',
-            message: `Starting simplified agent search in ${focusMode} mode`,
-            details: `Processing query with ${fileIds.length} files available`,
-          },
-        }),
-      );
-
       // Prepare initial state
       const initialState = {
         messages: [...history, new HumanMessage(query)],
@@ -489,25 +478,165 @@ Use all available tools strategically to provide comprehensive, well-researched,
         signal: this.signal,
       };
 
-      // Execute the agent
+      // Use streamEvents to capture both tool calls and token-level streaming
-      const result = await agent.invoke(initialState, config);
+      const eventStream = agent.streamEvents(initialState, {
+        ...config,
+        version: 'v2',
+      });
 
-      // Collect relevant documents from tool execution history
+      let finalResult: any = null;
       let collectedDocuments: any[] = [];
+      let currentResponseBuffer = '';
 
-      // Get the relevant docs from the current agent state
+      // Process the event stream
-      if (result && result.relevantDocuments) {
+      for await (const event of eventStream) {
-        collectedDocuments.push(...result.relevantDocuments);
+        // Handle different event types
+        if (
+          event.event === 'on_chain_end' &&
+          event.name === 'RunnableSequence'
+        ) {
+          finalResult = event.data.output;
+          // Collect relevant documents from the final result
+          if (finalResult && finalResult.relevantDocuments) {
+            collectedDocuments.push(...finalResult.relevantDocuments);
+          }
         }
 
-      // Add collected documents to result for source tracking
+        // Collect sources from tool results
-      const finalResult = {
-        ...result,
-        relevantDocuments: collectedDocuments,
-      };
-
-      // Extract final message and emit as response
         if (
+          event.event === 'on_chain_end' &&
+          (event.name.includes('search') ||
+            event.name.includes('Search') ||
+            event.name.includes('tool') ||
+            event.name.includes('Tool'))
+        ) {
+          // Handle LangGraph state updates with relevantDocuments
+          if (event.data?.output && Array.isArray(event.data.output)) {
+            for (const item of event.data.output) {
+              if (
+                item.update &&
+                item.update.relevantDocuments &&
+                Array.isArray(item.update.relevantDocuments)
+              ) {
+                collectedDocuments.push(...item.update.relevantDocuments);
+              }
+            }
+          }
+        }
+
+        // Emit sources as we collect them
+        if (collectedDocuments.length > 0) {
+          this.emitter.emit(
+            'data',
+            JSON.stringify({
+              type: 'sources',
+              data: collectedDocuments,
+              searchQuery: '',
+              searchUrl: '',
+            }),
+          );
+        }
+
+        // Handle streaming tool calls (for thought messages)
+        if (event.event === 'on_chat_model_end' && event.data.output) {
+          const output = event.data.output;
+          if (
+            output._getType() === 'ai' &&
+            output.tool_calls &&
+            output.tool_calls.length > 0
+          ) {
+            const aiMessage = output as AIMessage;
+
+            // Process each tool call and emit thought messages
+            for (const toolCall of aiMessage.tool_calls || []) {
+              if (toolCall && toolCall.name) {
+                const toolName = toolCall.name;
+                const toolArgs = toolCall.args || {};
+
+                // Create user-friendly messages for different tools using markdown components
+                let toolMarkdown = '';
+                switch (toolName) {
+                  case 'web_search':
+                    toolMarkdown = `<ToolCall type="search" query="${(toolArgs.query || 'relevant information').replace(/"/g, '&quot;')}"></ToolCall>`;
+                    break;
+                  case 'file_search':
+                    toolMarkdown = `<ToolCall type="file" query="${(toolArgs.query || 'relevant information').replace(/"/g, '&quot;')}"></ToolCall>`;
+                    break;
+                  case 'url_summarization':
+                    if (Array.isArray(toolArgs.urls)) {
+                      toolMarkdown = `<ToolCall type="url" count="${toolArgs.urls.length}"></ToolCall>`;
+                    } else {
+                      toolMarkdown = `<ToolCall type="url" count="1"></ToolCall>`;
+                    }
+                    break;
+                  default:
+                    toolMarkdown = `<ToolCall type="${toolName}"></ToolCall>`;
+                }
+
+                // Emit the thought message
+                this.emitter.emit(
+                  'data',
+                  JSON.stringify({
+                    type: 'tool_call',
+                    data: {
+                      // messageId: crypto.randomBytes(7).toString('hex'),
+                      content: toolMarkdown,
+                    },
+                  }),
+                );
+              }
+            }
+          }
+        }
+
+        // Handle token-level streaming for the final response
+        if (event.event === 'on_chat_model_stream' && event.data.chunk) {
+          const chunk = event.data.chunk;
+          if (chunk.content && typeof chunk.content === 'string') {
+            // If this is the first token, emit sources if we have them
+            if (currentResponseBuffer === '' && collectedDocuments.length > 0) {
+              this.emitter.emit(
+                'data',
+                JSON.stringify({
+                  type: 'sources',
+                  data: collectedDocuments,
+                  searchQuery: '',
+                  searchUrl: '',
+                }),
+              );
+            }
+
+            // Add the token to our buffer
+            currentResponseBuffer += chunk.content;
+
+            // Emit the individual token
+            this.emitter.emit(
+              'data',
+              JSON.stringify({
+                type: 'response',
+                data: chunk.content,
+              }),
+            );
+          }
+        }
+      }
+
+      // Emit the final sources used for the response
+      if (collectedDocuments.length > 0) {
+        this.emitter.emit(
+          'data',
+          JSON.stringify({
+            type: 'sources',
+            data: collectedDocuments,
+            searchQuery: '',
+            searchUrl: '',
+          }),
+        );
+      }
+
+      // If we didn't get any streamed tokens but have a final result, emit it
+      if (
+        currentResponseBuffer === '' &&
         finalResult &&
         finalResult.messages &&
         finalResult.messages.length > 0
@@ -516,23 +645,7 @@ Use all available tools strategically to provide comprehensive, well-researched,
           finalResult.messages[finalResult.messages.length - 1];
 
         if (finalMessage && finalMessage.content) {
-          console.log('SimplifiedAgent: Emitting final response');
+          console.log('SimplifiedAgent: Emitting complete response (fallback)');
-
-          // Emit the sources used for the response
-          if (
-            finalResult.relevantDocuments &&
-            finalResult.relevantDocuments.length > 0
-          ) {
-            this.emitter.emit(
-              'data',
-              JSON.stringify({
-                type: 'sources',
-                data: finalResult.relevantDocuments,
-                searchQuery: '',
-                searchUrl: '',
-              }),
-            );
-          }
 
           this.emitter.emit(
             'data',
@@ -541,8 +654,17 @@ Use all available tools strategically to provide comprehensive, well-researched,
               data: finalMessage.content,
             }),
           );
-        } else {
+        }
-          console.warn('SimplifiedAgent: No valid final message found');
+      }
+
+      // If we still have no response, emit a fallback message
+      if (
+        currentResponseBuffer === '' &&
+        (!finalResult ||
+          !finalResult.messages ||
+          finalResult.messages.length === 0)
+      ) {
+        console.warn('SimplifiedAgent: No valid response found');
         this.emitter.emit(
           'data',
           JSON.stringify({
@@ -551,16 +673,6 @@ Use all available tools strategically to provide comprehensive, well-researched,
           }),
         );
       }
-      } else {
-        console.warn('SimplifiedAgent: No result messages found');
-        this.emitter.emit(
-          'data',
-          JSON.stringify({
-            type: 'response',
-            data: 'I encountered an issue while processing your request. Please try again with a different query.',
-          }),
-        );
-      }
 
       // Emit model stats and end signal after streaming is complete
       const modelName = getModelName(this.llm);
@@ -577,7 +689,7 @@ Use all available tools strategically to provide comprehensive, well-researched,
       console.error('SimplifiedAgent: Error during search and answer:', error);
 
       // Handle specific error types
-      if (error.name === 'AbortError') {
+      if (error.name === 'AbortError' || this.signal.aborted) {
         console.warn('SimplifiedAgent: Operation was aborted');
         this.emitter.emit(
           'data',