From 7c6437eef6a582f12f0b13fe66f3bb783df04d1b Mon Sep 17 00:00:00 2001
From: Utkarsh-Patel-13 <utkarshhpatel13@gmail.com>
Date: Tue, 22 Jul 2025 20:36:52 -0700
Subject: [PATCH] Better layout for terminal and chat messages

---
 .../app/services/streaming_service.py         | 12 +----
 .../components/chat_v2/ChatInterface.tsx      | 50 ++++++++++++++-----
 .../components/chat_v2/ChatTerminal.tsx       | 42 +++++++---------
 3 files changed, 57 insertions(+), 47 deletions(-)

diff --git a/surfsense_backend/app/services/streaming_service.py b/surfsense_backend/app/services/streaming_service.py
index 08b31b7..cbd0e6d 100644
--- a/surfsense_backend/app/services/streaming_service.py
+++ b/surfsense_backend/app/services/streaming_service.py
@@ -43,17 +43,7 @@ class StreamingService:
         self.message_annotations[0]["content"].append(message)
 
         # Return only the delta annotation
-        # annotation = {"type": "TERMINAL_INFO", "content": [message]}
-        agent_data = {
-            "agent_name": "Terminal agent",
-            "events": [
-                {
-                    "status": message_type,
-                    "result": message.get("text", ""),
-                }
-            ],
-        }
-        annotation = {"type": "agent_events", "data": agent_data}
+        annotation = {"type": "TERMINAL_INFO", "data": message}
         return f"8:[{json.dumps(annotation)}]\n"
 
     def format_sources_delta(self, sources: List[Dict[str, Any]]) -> str:
diff --git a/surfsense_web/components/chat_v2/ChatInterface.tsx b/surfsense_web/components/chat_v2/ChatInterface.tsx
index be3a85d..6a02ca7 100644
--- a/surfsense_web/components/chat_v2/ChatInterface.tsx
+++ b/surfsense_web/components/chat_v2/ChatInterface.tsx
@@ -7,6 +7,7 @@ import {
     ChatMessages,
     useChatUI,
     ChatMessage,
+    Message,
 } from "@llamaindex/chat-ui";
 import { Document } from "@/hooks/use-documents";
 import { CustomChatInput } from "@/components/chat_v2/ChatInputGroup";
@@ -26,23 +27,48 @@ interface ChatInterfaceProps {
     onResearchModeChange?: (mode: ResearchMode) => void;
 }
 
-function ChatMessageDisplay() {
+function ChatMessageDisplay({
+    message,
+    isLast,
+}: {
+    message: Message;
+    isLast: boolean;
+}) {
+    return (
+        <ChatMessage
+            message={message}
+            isLast={isLast}
+            className="flex flex-col "
+        >
+            {message.role === "assistant" ? (
+                <div className="flex-1 flex flex-col space-y-4">
+                    <TerminalDisplay message={message} />
+                    <ChatMessage.Content className="flex-1">
+                        <ChatMessage.Content.Markdown />
+                    </ChatMessage.Content>
+                    <ChatMessage.Actions className="flex-1 flex flex-row justify-end" />
+                </div>
+            ) : (
+                <ChatMessage.Content className="flex-1">
+                    <ChatMessage.Content.Markdown />
+                </ChatMessage.Content>
+            )}
+        </ChatMessage>
+    );
+}
+
+function ChatMessagesDisplay() {
     const { messages } = useChatUI();
 
     return (
         <ChatMessages className="flex-1">
             <ChatMessages.List className="p-4">
                 {messages.map((message, index) => (
-                    <div key={`Message-${index}`}>
-                        {message.role === "assistant" && (
-                            <TerminalDisplay messages={messages} />
-                        )}
-                        <ChatMessage
-                            key={index}
-                            message={message}
-                            isLast={index === messages.length - 1}
-                        />
-                    </div>
+                    <ChatMessageDisplay
+                        key={`Message-${index}`}
+                        message={message}
+                        isLast={index === messages.length - 1}
+                    />
                 ))}
             </ChatMessages.List>
             <ChatMessages.Loading />
@@ -64,7 +90,7 @@ export default function ChatInterface({
     return (
         <ChatSection handler={handler} className="flex h-full">
             <div className="flex flex-1 flex-col">
-                <ChatMessageDisplay />
+                <ChatMessagesDisplay />
                 <div className="border-t p-4">
                     <CustomChatInput
                         onDocumentSelectionChange={onDocumentSelectionChange}
diff --git a/surfsense_web/components/chat_v2/ChatTerminal.tsx b/surfsense_web/components/chat_v2/ChatTerminal.tsx
index e99f147..61d1a9b 100644
--- a/surfsense_web/components/chat_v2/ChatTerminal.tsx
+++ b/surfsense_web/components/chat_v2/ChatTerminal.tsx
@@ -3,31 +3,24 @@
 import React from "react";
 import { getAnnotationData, Message } from "@llamaindex/chat-ui";
 
-export default function TerminalDisplay({ messages }: { messages: Message[] }) {
+export default function TerminalDisplay({ message }: { message: Message }) {
     const [isCollapsed, setIsCollapsed] = React.useState(true);
 
     // Get the last assistant message that's not being typed
-    const lastAssistantMessage = messages
-        .filter((msg) => msg.role === "assistant")
-        .pop();
-
-    if (!lastAssistantMessage) {
+    if (!message) {
         return <></>;
     }
 
-    const annotations = getAnnotationData(lastAssistantMessage, "agent_events");
-
-    if (annotations.length === 0) {
-        return <></>;
+    interface TerminalInfo {
+        id: number;
+        text: string;
+        type: string;
     }
 
-    // Get all agent events and combine them
-    const events: any[] = [];
-    annotations.forEach((annotation: any) => {
-        if (annotation?.events && Array.isArray(annotation.events)) {
-            events.push(...annotation.events);
-        }
-    });
+    const events = getAnnotationData(
+        message,
+        "TERMINAL_INFO"
+    ) as TerminalInfo[];
 
     if (events.length === 0) {
         return <></>;
@@ -85,16 +78,17 @@ export default function TerminalDisplay({ messages }: { messages: Message[] }) {
             {!isCollapsed && (
                 <div className="h-64 overflow-y-auto p-4 space-y-1 bg-gray-900">
                     {events.map((event, index) => (
-                        <div key={index} className="text-green-400">
+                        <div
+                            key={`${event.id}-${index}`}
+                            className="text-green-400"
+                        >
                             <span className="text-blue-400">$</span>
                             <span className="text-yellow-400 ml-2">
-                                [{event.status || "completed"}]
+                                [{event.type || ""}]
+                            </span>
+                            <span className="text-gray-300 ml-4 mt-1 pl-2 border-l-2 border-gray-600">
+                                {event.text || ""}...
                             </span>
-                            {event.result && (
-                                <span className="text-gray-300 ml-4 mt-1 pl-2 border-l-2 border-gray-600">
-                                    {event.result.slice(0, 100)}...
-                                </span>
-                            )}
                         </div>
                     ))}
                     {events.length === 0 && (