@eeacms/volto-eea-chatbot 2.0.1 → 2.0.3

package/src/ChatBlock/packets/MultiToolRenderer.tsx

@@ -1,13 +1,13 @@
-import type { Packet } from '../types/streamingModels';
-import type { Message } from '../types/interfaces';
+import type { Packet } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import type { Message } from '@eeacms/volto-eea-chatbot/ChatBlock/types/interfaces';
 import React, { useState, useEffect, useMemo } from 'react';
 import cx from 'classnames';
-import { PacketType } from '../types/streamingModels';
+import { PacketType } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
 import { RendererComponent } from './RendererComponent';
-import { useToolDisplayTiming } from '../hooks/useToolDisplayTiming';
-import SVGIcon from '../components/Icon';
-import DoneIcon from '../../icons/done.svg';
-import ChevronIcon from '../../icons/chevron.svg';
+import { useToolDisplayTiming } from '@eeacms/volto-eea-chatbot/ChatBlock/hooks/useToolDisplayTiming';
+import SVGIcon from '@eeacms/volto-eea-chatbot/ChatBlock/components/Icon';
+import DoneIcon from '@eeacms/volto-eea-chatbot/icons/done.svg';
+import ChevronIcon from '@eeacms/volto-eea-chatbot/icons/chevron.svg';
 
 interface MultiToolRendererProps {
   toolGroups: { ind: number; packets: Packet[] }[];
@@ -19,7 +19,7 @@ interface MultiToolRendererProps {
 
 export function MultiToolRenderer({
   toolGroups,
-  showTools = [PacketType.SEARCH_TOOL_START],
+  showTools = [PacketType.SEARCH_TOOL_START, PacketType.REASONING_START],
   message,
   libs,
   onAllToolsDisplayed,
@@ -28,23 +28,33 @@ export function MultiToolRenderer({
   const { isFinalMessageComing = false, isComplete = false } = message;
 
   // Filter tool groups based on allowed tool types
-  const filteredToolGroups = useMemo(
-    () =>
-      toolGroups.filter(
-        (group) =>
-          group.packets?.some(
-            (packet) => showTools?.includes(packet.obj.type as PacketType),
-          ),
+  const filteredToolGroups = useMemo(() => {
+    const expandedShowTools = [...(showTools || [])];
+    if (showTools?.includes(PacketType.SEARCH_TOOL_START)) {
+      expandedShowTools.push(
+        PacketType.SEARCH_TOOL_START_V3,
+        PacketType.SEARCH_TOOL_QUERIES_DELTA,
+        PacketType.SEARCH_TOOL_DOCUMENTS_DELTA,
+        PacketType.SEARCH_TOOL_DELTA,
+      );
+    }
+    if (showTools?.includes(PacketType.REASONING_START)) {
+      expandedShowTools.push(
+        PacketType.REASONING_DELTA,
+        PacketType.REASONING_DONE,
+        PacketType.REASONING_END as any,
+      );
+    }
+    return toolGroups.filter((group) =>
+      group.packets?.some((packet) =>
+        expandedShowTools.includes(packet.obj.type as PacketType),
       ),
-    [toolGroups, showTools],
-  );
+    );
+  }, [toolGroups, showTools]);
 
   // Manage tool display timing
-  const { allToolsDisplayed, handleToolComplete } = useToolDisplayTiming(
-    filteredToolGroups,
-    isFinalMessageComing,
-    isComplete,
-  );
+  const { allToolsDisplayed, handleToolComplete, toolStates } =
+    useToolDisplayTiming(filteredToolGroups, isFinalMessageComing, isComplete);
 
   // Notify parent when all tools are displayed
   useEffect(() => {
@@ -65,27 +75,29 @@ export function MultiToolRenderer({
 
   if (filteredToolGroups.length === 0) return null;
 
-  const isStreaming = !allToolsDisplayed;
+  const isOverallStreaming = !allToolsDisplayed;
 
   const count = filteredToolGroups.length;
 
-  const ariaLabel = `${count} ${isStreaming ? 'processing' : 'completed'} ${
-    count === 1 ? 'step' : 'steps'
-  }, ${isExpanded ? 'expanded' : 'collapsed'}`;
+  const ariaLabel = `${count} ${
+    isOverallStreaming ? 'processing' : 'completed'
+  } ${count === 1 ? 'step' : 'steps'}, ${isExpanded ? 'expanded' : 'collapsed'}`;
 
   return (
     <div
       className={cx('multi-tool-renderer', {
-        streaming: isStreaming,
-        complete: !isStreaming,
+        streaming: isOverallStreaming,
+        complete: !isOverallStreaming,
       })}
     >
       {/* Header */}
-      <div className={cx({ 'tools-container collapsed-view': isStreaming })}>
+      <div
+        className={cx({ 'tools-container collapsed-view': isOverallStreaming })}
+      >
         <div
           className={cx({
-            'tools-collapsed-header': isStreaming,
-            'tools-summary-header': !isStreaming,
+            'tools-collapsed-header': isOverallStreaming,
+            'tools-summary-header': !isOverallStreaming,
           })}
           onClick={toggleExpanded}
           role="button"
@@ -110,21 +122,25 @@ export function MultiToolRenderer({
         {/* Tools List */}
         <div
           className={cx({
-            'tools-collapsed-list': isStreaming,
-            'tools-expanded-content': !isStreaming,
-            expanded: isExpanded && isStreaming,
-            visible: isExpanded && !isStreaming,
+            'tools-collapsed-list': isOverallStreaming,
+            'tools-expanded-content': !isOverallStreaming,
+            expanded: isExpanded && isOverallStreaming,
+            visible: isExpanded && !isOverallStreaming,
           })}
         >
-          <div className={cx({ 'tools-list': isStreaming })}>
+          <div className={cx({ 'tools-list': isOverallStreaming })}>
             <div>
               {filteredToolGroups.map((toolGroup, index) => {
                 const isLastItem = index === filteredToolGroups.length - 1;
+                const toolState = toolStates.get(toolGroup.ind);
+                const isToolCompleted = toolState?.isCompleted;
 
                 return (
                   <div
                     key={toolGroup.ind}
-                    className={cx({ 'tool-collaps ed-wrapper': isStreaming })}
+                    className={cx({
+                      'tool-collapsed-wrapper': isOverallStreaming,
+                    })}
                   >
                     <RendererComponent
                       packets={toolGroup.packets}
@@ -144,32 +160,46 @@ export function MultiToolRenderer({
                         ) : (
                           <span
                             className={cx({
-                              'tool-icon-dot': isStreaming,
-                              'tool-icon-default': !isStreaming,
+                              'tool-icon-dot': isOverallStreaming,
+                              'tool-icon-default': !isOverallStreaming,
                             })}
                           />
                         );
 
-                        // Streaming: collapsed view (status only)
-                        if (isStreaming) {
-                          return (
-                            <div
-                              className={cx('tool-item-collapsed', {
-                                active: isLastItem,
-                                completed: !isLastItem,
-                              })}
-                            >
-                              <div className="tool-collapsed-icon">
-                                {finalIcon}
-                              </div>
-                              <span className="tool-collapsed-status">
-                                {status}
-                              </span>
-                            </div>
+                        // If tool is not completed and we are overall streaming, show collapsed view
+                        // EXCEPT for reasoning and search which we want to see while they stream/progress
+                        if (isOverallStreaming && !isToolCompleted) {
+                          const isDetailedTool = toolGroup.packets.some(
+                            (p) =>
+                              p.obj.type === PacketType.REASONING_START ||
+                              p.obj.type === PacketType.REASONING_DELTA ||
+                              p.obj.type === PacketType.SEARCH_TOOL_START ||
+                              p.obj.type === PacketType.SEARCH_TOOL_START_V3 ||
+                              p.obj.type === PacketType.SEARCH_TOOL_DELTA ||
+                              p.obj.type === PacketType.SEARCH_TOOL_QUERIES_DELTA ||
+                              p.obj.type === PacketType.SEARCH_TOOL_DOCUMENTS_DELTA,
                           );
+
+                          if (!isDetailedTool || !content) {
+                            return (
+                              <div
+                                className={cx('tool-item-collapsed', {
+                                  active: isLastItem,
+                                  completed: isToolCompleted,
+                                })}
+                              >
+                                <div className="tool-collapsed-icon">
+                                  {finalIcon}
+                                </div>
+                                <span className="tool-collapsed-status">
+                                  {status}
+                                </span>
+                              </div>
+                            );
+                          }
                         }
 
-                        // Complete: expanded view (full content)
+                        // Expanded view (full content) - used for completed tools or when overall complete
                         return (
                           <div className="tool-item-expanded">
                             <div className="tool-connector-line" />

package/src/ChatBlock/packets/RendererComponent.tsx

@@ -1,6 +1,9 @@
-import type { Packet } from '../types/streamingModels';
-import type { RendererResult, Message } from '../types/interfaces';
-import { PacketType } from '../types/streamingModels';
+import type { Packet } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import type {
+  RendererResult,
+  Message,
+} from '@eeacms/volto-eea-chatbot/ChatBlock/types/interfaces';
+import { PacketType } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
 import {
   MessageTextRenderer,
   SearchToolRenderer,
@@ -23,7 +26,10 @@ function isChatPacket(packet: Packet): boolean {
 }
 
 function isSearchToolPacket(packet: Packet): boolean {
-  return packet.obj.type === PacketType.SEARCH_TOOL_START;
+  return (
+    packet.obj.type === PacketType.SEARCH_TOOL_START ||
+    packet.obj.type === PacketType.SEARCH_TOOL_START_V3
+  );
 }
 
 function isImageToolPacket(packet: Packet): boolean {
@@ -41,7 +47,9 @@ function isFetchToolPacket(packet: Packet): boolean {
 function isReasoningPacket(packet: Packet): boolean {
   return (
     packet.obj.type === PacketType.REASONING_START ||
-    packet.obj.type === PacketType.REASONING_DELTA
+    packet.obj.type === PacketType.REASONING_DELTA ||
+    packet.obj.type === PacketType.REASONING_DONE ||
+    packet.obj.type === PacketType.REASONING_END
   );
 }
 

package/src/ChatBlock/packets/renderers/CustomToolRenderer.tsx

@@ -2,10 +2,10 @@ import type {
   CustomToolPacket,
   CustomToolStart,
   CustomToolDelta,
-} from '../../types/streamingModels';
-import type { MessageRenderer } from '../../types/interfaces';
+} from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import type { MessageRenderer } from '@eeacms/volto-eea-chatbot/ChatBlock/types/interfaces';
 import { useEffect } from 'react';
-import { PacketType } from '../../types/streamingModels';
+import { PacketType } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
 
 export const CustomToolRenderer: MessageRenderer<CustomToolPacket> = ({
   packets,

package/src/ChatBlock/packets/renderers/FetchToolRenderer.tsx

@@ -1,10 +1,10 @@
 import type {
   FetchToolPacket,
   FetchToolStart,
-} from '../../types/streamingModels';
-import type { MessageRenderer } from '../../types/interfaces';
+} from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import type { MessageRenderer } from '@eeacms/volto-eea-chatbot/ChatBlock/types/interfaces';
 import { useEffect } from 'react';
-import { PacketType } from '../../types/streamingModels';
+import { PacketType } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
 
 export const FetchToolRenderer: MessageRenderer<FetchToolPacket> = ({
   packets,

package/src/ChatBlock/packets/renderers/ImageToolRenderer.tsx

@@ -3,9 +3,9 @@ import type {
   ImageGenerationToolPacket,
   ImageGenerationToolDelta,
   GeneratedImage,
-} from '../../types/streamingModels';
-import type { MessageRenderer } from '../../types/interfaces';
-import { PacketType } from '../../types/streamingModels';
+} from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import type { MessageRenderer } from '@eeacms/volto-eea-chatbot/ChatBlock/types/interfaces';
+import { PacketType } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
 
 export const ImageToolRenderer: MessageRenderer<ImageGenerationToolPacket> = ({
   packets,

package/src/ChatBlock/packets/renderers/MessageTextRenderer.tsx

@@ -1,11 +1,11 @@
-import type { ChatPacket } from '../../types/streamingModels';
-import type { MessageRenderer } from '../../types/interfaces';
+import type { ChatPacket } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import type { MessageRenderer } from '@eeacms/volto-eea-chatbot/ChatBlock/types/interfaces';
 import { useEffect, useMemo, useState } from 'react';
 import loadable from '@loadable/component';
-import { components } from '../../components/markdown';
-import { isFinalAnswerComplete } from '../../services/packetUtils';
-import { PacketType } from '../../types/streamingModels';
-import { BlinkingDot } from '../../components/BlinkingDot';
+import { components } from '@eeacms/volto-eea-chatbot/ChatBlock/components/markdown';
+import { isFinalAnswerComplete } from '@eeacms/volto-eea-chatbot/ChatBlock/services/packetUtils';
+import { PacketType } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import { BlinkingDot } from '@eeacms/volto-eea-chatbot/ChatBlock/components/BlinkingDot';
 
 const Markdown: any = loadable(() => import('react-markdown'));
 
@@ -90,7 +90,11 @@ export const MessageTextRenderer: MessageRenderer<ChatPacket> = ({
           // If we're far behind, catch up faster
           const increment =
             remaining > CATCH_UP_THRESHOLD ? PACKETS_PER_TICK : 1;
-          return Math.min(prev + increment, packets.length);
+          const next = Math.min(prev + increment, packets.length);
+          if (isStreamFinished && next === packets.length) {
+            console.log(`[MessageTextRenderer] Animation finished: ${next}/${packets.length}`);
+          }
+          return next;
         });
       }, PACKET_DELAY_MS);
 
@@ -104,8 +108,8 @@ export const MessageTextRenderer: MessageRenderer<ChatPacket> = ({
       const resetCount = isStreamFinished
         ? packets.length // Show all if stream is finished
         : packets.length > 0
-        ? 1
-        : 0;
+          ? 1
+          : 0;
       setDisplayedPacketCount(resetCount);
     }
   }, [animate, packets.length, displayedPacketCount, isStreamFinished]);
@@ -121,6 +125,7 @@ export const MessageTextRenderer: MessageRenderer<ChatPacket> = ({
       ) {
         return;
       }
+      console.log(`[MessageTextRenderer] Calling onComplete: packets=${packets.length}, finished=${isStreamFinished}`);
       onComplete();
     }
   }, [

package/src/ChatBlock/packets/renderers/ReasoningRenderer.tsx

@@ -1,13 +1,13 @@
 import type {
   ReasoningPacket,
   ReasoningDelta,
-} from '../../types/streamingModels';
-import type { MessageRenderer } from '../../types/interfaces';
+} from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import type { MessageRenderer } from '@eeacms/volto-eea-chatbot/ChatBlock/types/interfaces';
 import { useEffect, useState, useRef, useMemo } from 'react';
 import loadable from '@loadable/component';
-import { PacketType } from '../../types/streamingModels';
-import { components } from '../../components/markdown';
-import { addCitations } from '../../utils/citations';
+import { PacketType } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import { components } from '@eeacms/volto-eea-chatbot/ChatBlock/components/markdown';
+import { addCitations } from '@eeacms/volto-eea-chatbot/ChatBlock/utils/citations';
 
 const Markdown: any = loadable(() => import('react-markdown'));
 
@@ -21,6 +21,7 @@ function constructCurrentReasoningState(packets: ReasoningPacket[]) {
   const hasEnd = packets.some(
     (p) =>
       p.obj.type === PacketType.SECTION_END ||
+      p.obj.type === PacketType.REASONING_DONE ||
       // Support either convention for reasoning completion
       (p.obj as any).type === PacketType.REASONING_END,
   );

package/src/ChatBlock/packets/renderers/SearchToolRenderer.tsx

@@ -4,17 +4,17 @@ import type {
   SearchToolDelta,
   SectionEnd,
   OnyxDocument,
-} from '../../types/streamingModels';
-import type { MessageRenderer } from '../../types/interfaces';
+} from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import type { MessageRenderer } from '@eeacms/volto-eea-chatbot/ChatBlock/types/interfaces';
 import { useEffect, useState, useRef, useMemo } from 'react';
-import { PacketType } from '../../types/streamingModels';
-import { SourceChip } from '../../components/SourceChip';
-import { BlinkingDot } from '../../components/BlinkingDot';
-import SVGIcon from '../../components/Icon';
-import { WebResultIcon } from '../../components/WebResultIcon';
-import SearchIcon from '../../../icons/search.svg';
-import GlobeIcon from '../../../icons/globe.svg';
-import FileIcon from '../../../icons/file.svg';
+import { PacketType } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import { SourceChip } from '@eeacms/volto-eea-chatbot/ChatBlock/components/SourceChip';
+import { BlinkingDot } from '@eeacms/volto-eea-chatbot/ChatBlock/components/BlinkingDot';
+import SVGIcon from '@eeacms/volto-eea-chatbot/ChatBlock/components/Icon';
+import { WebResultIcon } from '@eeacms/volto-eea-chatbot/ChatBlock/components/WebResultIcon';
+import SearchIcon from '@eeacms/volto-eea-chatbot/icons/search.svg';
+import GlobeIcon from '@eeacms/volto-eea-chatbot/icons/globe.svg';
+import FileIcon from '@eeacms/volto-eea-chatbot/icons/file.svg';
 
 const INITIAL_RESULTS_TO_SHOW = 3;
 const RESULTS_PER_EXPANSION = 10;
@@ -55,26 +55,33 @@ const constructCurrentSearchState = (
   isInternetSearch: boolean;
 } => {
   const searchStart = packets.find(
-    (packet) => packet.obj.type === PacketType.SEARCH_TOOL_START,
+    (packet) =>
+      packet.obj.type === PacketType.SEARCH_TOOL_START ||
+      packet.obj.type === PacketType.SEARCH_TOOL_START_V3,
   )?.obj as SearchToolStart | null;
 
   const searchDeltas = packets
-    .filter((packet) => packet.obj.type === PacketType.SEARCH_TOOL_DELTA)
-    .map((packet) => packet.obj as SearchToolDelta);
+    .filter(
+      (packet) =>
+        packet.obj.type === PacketType.SEARCH_TOOL_DELTA ||
+        packet.obj.type === PacketType.SEARCH_TOOL_QUERIES_DELTA ||
+        packet.obj.type === PacketType.SEARCH_TOOL_DOCUMENTS_DELTA,
+    )
+    .map((packet) => packet.obj);
 
   const searchEnd = packets.find(
     (packet) => packet.obj.type === PacketType.SECTION_END,
   )?.obj as SectionEnd | null;
 
-  // Extract queries from ToolDelta packets
+  // Extract queries from various delta packets
   const queries = searchDeltas
-    .flatMap((delta) => delta?.queries || [])
+    .flatMap((delta: any) => delta?.queries || [])
     .filter((query, index, arr) => arr.indexOf(query) === index); // Remove duplicates
 
   const seenDocIds = new Set<string>();
   const results = searchDeltas
-    .flatMap((delta) => delta?.documents || [])
-    .filter((doc) => {
+    .flatMap((delta: any) => delta?.documents || [])
+    .filter((doc: OnyxDocument) => {
       if (!doc || !doc.document_id) return false;
       if (seenDocIds.has(doc.document_id)) return false;
       seenDocIds.add(doc.document_id);
@@ -202,8 +209,8 @@ export const SearchToolRenderer: MessageRenderer<SearchToolPacket> = ({
     <SVGIcon name={isInternetSearch ? GlobeIcon : SearchIcon} size={size} />
   );
 
-  // Don't render anything if search hasn't started
-  if (queries.length === 0) {
+  // Don't render anything if search hasn't started or has no data yet
+  if (queries.length === 0 && results.length === 0) {
     return children({
       icon: IconComponent,
       status: status,
@@ -218,7 +225,9 @@ export const SearchToolRenderer: MessageRenderer<SearchToolPacket> = ({
       <div className="search-tool-renderer">
         <div className="queries-section">
           <div className="queries-header">
-            <strong>Queries</strong>
+            <strong>
+              {isInternetSearch ? 'Web Queries' : 'Internal Search Queries'}
+            </strong>
           </div>
           <div className="queries-list">
             {queries.slice(0, queriesToShow).map((query, index) => (
@@ -262,7 +271,7 @@ export const SearchToolRenderer: MessageRenderer<SearchToolPacket> = ({
 
         <div className="results-section">
           <div className="results-header">
-            <strong>{isInternetSearch ? 'Results' : 'Documents'}</strong>
+            <strong>{isInternetSearch ? 'Web Results' : 'Documents'}</strong>
           </div>
 
           <div className="results-list">

package/src/ChatBlock/{schema.js → schema.jsx}

@@ -104,6 +104,7 @@ export function ChatBlockSchema({ assistants, data }) {
           'scrollToInput',
           'showAssistantTitle',
           'showAssistantDescription',
+          'onyxVersion',
         ],
       },
     ],
@@ -368,6 +369,18 @@ range is from 0 to 100`,
         type: 'boolean',
         default: true,
       },
+      onyxVersion: {
+        title: 'Onyx API version',
+        choices: [
+          ['2', 'Onyx 2.x'],
+          ['3', 'Onyx 3.x'],
+        ],
+        default: '2',
+        description:
+          'Select which Onyx API version the backend is running. ' +
+          'Onyx 2.x uses the legacy send-message payload; ' +
+          'Onyx 3.x uses the new send-message payload with placement-based streaming.',
+      },
       showAssistantPrompts: {
         title: 'Show predefined prompts',
         type: 'boolean',

package/src/ChatBlock/services/messageProcessor.ts

@@ -5,14 +5,17 @@ import type {
   OnyxDocument,
   Packet,
   StreamingCitation,
-} from '../types/streamingModels';
-import type { Message, ToolCallMetadata } from '../types/interfaces';
+} from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
+import type {
+  Message,
+  ToolCallMetadata,
+} from '@eeacms/volto-eea-chatbot/ChatBlock/types/interfaces';
 import {
   getSynteticPacket,
   isToolPacket,
   isDisplayPacket,
 } from './packetUtils';
-import { PacketType } from '../types/streamingModels';
+import { PacketType } from '@eeacms/volto-eea-chatbot/ChatBlock/types/streamingModels';
 
 /**
  * Process streaming packets into a message object
@@ -108,17 +111,26 @@ export class MessageProcessor {
       this.indicesStarted.push(packet.ind);
     }
 
-    // Send synthetic SECTION_END when needed
-    if (
-      packet.obj.type === PacketType.SECTION_END &&
-      this.indicesStarted.length > 0
-    ) {
-      processedPacket = getSynteticPacket(
-        this.indicesStarted.shift()!,
-        PacketType.SECTION_END,
-      );
-    } else if (packet.obj.type === PacketType.SECTION_END) {
-      return;
+    // Handle SECTION_END (v2/v3 compatible)
+    if (packet.obj.type === PacketType.SECTION_END) {
+      let targetInd = packet.ind;
+
+      // If ind is -1 (common in v2), use synthetic logic
+      if (targetInd === -1 && this.indicesStarted.length > 0) {
+        targetInd = this.indicesStarted.shift()!;
+      } else if (targetInd !== -1) {
+        // If ind is provided (v3), remove it from indicesStarted if present
+        const startIdx = this.indicesStarted.indexOf(targetInd);
+        if (startIdx > -1) {
+          this.indicesStarted.splice(startIdx, 1);
+        }
+      }
+
+      if (targetInd === -1) {
+        return; // Skip if we can't map it to a turn
+      }
+
+      processedPacket = getSynteticPacket(targetInd, PacketType.SECTION_END);
     }
 
     const { ind } = processedPacket;
@@ -126,6 +138,10 @@ export class MessageProcessor {
     // Store processed packet for later aggregation
     this.packets.push(processedPacket);
 
+    if (processedPacket.obj.type === PacketType.MESSAGE_START || processedPacket.obj.type === PacketType.SECTION_END) {
+        console.log(`[MessageProcessor] Processed ${processedPacket.obj.type} for ind=${processedPacket.ind}`);
+    }
+
     // Group packets by index for later processing
     if (!this.groupedPackets.has(ind)) {
       this.groupedPackets.set(ind, []);
@@ -171,6 +187,7 @@ export class MessageProcessor {
         PacketType.MESSAGE_START,
         PacketType.SEARCH_TOOL_DELTA,
         PacketType.FETCH_TOOL_START,
+        PacketType.SEARCH_TOOL_DOCUMENTS_DELTA,
       ].includes(packet.obj.type)
     ) {
       return;
@@ -178,7 +195,8 @@ export class MessageProcessor {
     let newDocuments = false;
     const data = packet.obj as any;
     const documents = data.final_documents || data.documents;
-    if (documents) {
+
+    if (documents && Array.isArray(documents)) {
       documents.forEach((doc: OnyxDocument) => {
         const docId = doc.document_id;
         if (docId && !this.documentMap.has(docId)) {
@@ -187,8 +205,23 @@ export class MessageProcessor {
         }
       });
     }
+
     if (newDocuments) {
       this._documents = Array.from(this.documentMap.values());
+
+      // If we have final_documents and no citations yet, create a fallback mapping
+      // This ensures the Sources tab shows up in v3 when citation_info is missing
+      if (
+        packet.obj.type === PacketType.MESSAGE_START &&
+        data.final_documents &&
+        this._citations.size === 0
+      ) {
+        data.final_documents.forEach((doc: OnyxDocument, index: number) => {
+          if (doc.document_id) {
+            this._citations.set(index + 1, doc.document_id);
+          }
+        });
+      }
     }
   }
 
@@ -197,6 +230,16 @@ export class MessageProcessor {
    * Updates the internal citation collection and notifies when new citations are added
    */
   private processCitations(packet: Packet) {
+    if (packet.obj.type === PacketType.CITATION_INFO) {
+      const citationInfo = packet.obj as any;
+      if (citationInfo.citation_number && citationInfo.document_id) {
+        this._citations.set(
+          citationInfo.citation_number,
+          citationInfo.document_id,
+        );
+      }
+      return;
+    }
     if (packet.obj.type !== PacketType.CITATION_DELTA) {
       return;
     }
@@ -234,6 +277,13 @@ export class MessageProcessor {
    */
   private processStreamEnd(packet: Packet) {
     if ([PacketType.STOP, PacketType.ERROR].includes(packet.obj.type)) {
+      // Close any remaining open sections (especially the last one)
+      while (this.indicesStarted.length > 0) {
+        const ind = this.indicesStarted.shift()!;
+        console.log(`[MessageProcessor] Stream ended. Synthesizing section_end for ind=${ind}`);
+        const synthetic = getSynteticPacket(ind, PacketType.SECTION_END);
+        this.processPacket(synthetic);
+      }
       this._isComplete = true;
     }
   }
@@ -245,7 +295,9 @@ export class MessageProcessor {
   private extractToolCall(packets: Packet[]): ToolCallMetadata | null {
     // Look for search tool packets
     const searchToolStart = packets.find(
-      (p) => p.obj.type === PacketType.SEARCH_TOOL_START,
+      (p) =>
+        p.obj.type === PacketType.SEARCH_TOOL_START ||
+        p.obj.type === PacketType.SEARCH_TOOL_START_V3,
     );
 
     if (!searchToolStart) {
@@ -257,7 +309,10 @@ export class MessageProcessor {
     const processedDocs: Record<string, any> = {};
 
     for (const packet of packets) {
-      if (packet.obj.type === PacketType.SEARCH_TOOL_DELTA) {
+      if (
+        packet.obj.type === PacketType.SEARCH_TOOL_DELTA ||
+        packet.obj.type === PacketType.SEARCH_TOOL_DOCUMENTS_DELTA
+      ) {
         const delta = packet.obj as any;
         if (delta.documents && Array.isArray(delta.documents)) {
           delta.documents.forEach((doc: any) => {