From 90c1aba057563a4a7e4006bc9388088aa55a02cc Mon Sep 17 00:00:00 2001
From: Adam Cohen Hillel <adamcohenhillel@gmail.com>
Date: Sat, 19 Apr 2025 12:44:39 +0100
Subject: [PATCH] - [x] When Adding new message / starting new response -> the
 UI doesnt update with the user added message, should show the user added
 message. - [x] After Agent Stop REFETCH Thread Messages - [x] Missing Loading
 State "..." before assistant message gets sent -> just white screen when u
 start new

---
 .../app/dashboard/agents/[threadId]/page.tsx  | 64 +++++++++++++++----
 .../components/thread/file-viewer-modal.tsx   |  2 +-
 2 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/frontend/src/app/dashboard/agents/[threadId]/page.tsx b/frontend/src/app/dashboard/agents/[threadId]/page.tsx
index 68d81aeb..aa2b1710 100644
--- a/frontend/src/app/dashboard/agents/[threadId]/page.tsx
+++ b/frontend/src/app/dashboard/agents/[threadId]/page.tsx
@@ -370,7 +370,7 @@ export default function ThreadPage({ params }: { params: Promise<ThreadParams> }
       thread_id: threadId,
       type: 'user',
       is_llm_message: false,
-      content: JSON.stringify({ role: 'user', content: message }),
+      content: message,
       metadata: '{}',
       created_at: new Date().toISOString(),
       updated_at: new Date().toISOString(),
@@ -397,8 +397,6 @@ export default function ThreadPage({ params }: { params: Promise<ThreadParams> }
       }
 
       const agentResult = results[1].value;
-      setMessages(prev => prev.filter(m => m.message_id !== optimisticUserMessage.message_id));
-
       setAgentRunId(agentResult.agent_run_id);
 
     } catch (err) {
@@ -414,7 +412,36 @@ export default function ThreadPage({ params }: { params: Promise<ThreadParams> }
     console.log(`[PAGE] Requesting agent stop via hook.`);
     setAgentStatus('idle');
     await stopStreaming();
-  }, [stopStreaming]);
+    
+    // Refetch messages after agent stop
+    try {
+      console.log('[PAGE] Refetching messages after agent stop');
+      const messagesData = await getMessages(threadId);
+      if (messagesData) {
+        // Map API message type to UnifiedMessage type
+        const unifiedMessages = (messagesData || [])
+          .filter(msg => msg.type !== 'status') // Filter out status messages
+          .map((msg: ApiMessageType, index: number) => {
+            return {
+              message_id: msg.message_id || null, 
+              thread_id: msg.thread_id || threadId,
+              type: (msg.type || 'system') as UnifiedMessage['type'], 
+              is_llm_message: Boolean(msg.is_llm_message),
+              content: msg.content || '',
+              metadata: msg.metadata || '{}',
+              created_at: msg.created_at || new Date().toISOString(),
+              updated_at: msg.updated_at || new Date().toISOString()
+            };
+          });
+        
+        console.log('[PAGE] Refetched messages after stop:', unifiedMessages.length);
+        setMessages(unifiedMessages);
+        scrollToBottom('smooth');
+      }
+    } catch (err) {
+      console.error('Error refetching messages after agent stop:', err);
+    }
+  }, [stopStreaming, threadId]);
 
   const handleScroll = () => {
     if (!messagesContainerRef.current) return;
@@ -736,11 +763,20 @@ export default function ThreadPage({ params }: { params: Promise<ThreadParams> }
                   return groupedMessages.map((group, groupIndex) => {
                     if (group.type === 'user') {
                       const message = group.messages[0];
-                      const parsedContent = safeJsonParse<ParsedContent>(message.content, {});
+                      // Fix parsing of user message content - handle both formats
+                      const messageContent = (() => {
+                        try {
+                          const parsed = safeJsonParse<ParsedContent>(message.content, { content: message.content });
+                          return parsed.content || message.content;
+                        } catch {
+                          return message.content;
+                        }
+                      })();
+                      
                       return (
                         <div key={group.key} className="flex justify-end">
                           <div className="max-w-[85%] rounded-lg bg-primary/10 px-4 py-3 text-sm">
-                            {parsedContent.content}
+                            {messageContent}
                           </div>
                         </div>
                       );
@@ -992,18 +1028,22 @@ export default function ThreadPage({ params }: { params: Promise<ThreadParams> }
                   });
                 })()}
                  {/* Render thinking indicator ONLY if agent is running AND the last group wasn't assistant OR there are no messages yet */}
-                 {agentStatus === 'running' && !streamingTextContent && !streamingToolCall &&
+                 {(agentStatus === 'running' || agentStatus === 'connecting') && 
                    (messages.length === 0 || messages[messages.length - 1].type === 'user') && (
                       <div ref={latestMessageRef}>
                          <div className="flex items-start gap-3">
                              <div className="flex-shrink-0 w-5 h-5 rounded-full flex items-center justify-center overflow-hidden bg-gray-200">
                                <Image src="/kortix-symbol.svg" alt="Suna" width={14} height={14} className="object-contain"/>
                              </div>
-                        <div className="flex items-center gap-1.5 py-1 mt-3"> {/* Adjusted padding/margin */}
-                          <div className="h-1.5 w-1.5 rounded-full bg-gray-400/50 animate-pulse" />
-                          <div className="h-1.5 w-1.5 rounded-full bg-gray-400/50 animate-pulse delay-150" />
-                          <div className="h-1.5 w-1.5 rounded-full bg-gray-400/50 animate-pulse delay-300" />
-                         </div>
+                        <div className="flex-1 space-y-2">
+                          <div className="max-w-[90%] px-4 py-3 text-sm">
+                            <div className="flex items-center gap-1.5 py-1"> {/* Adjusted padding/margin */}
+                              <div className="h-1.5 w-1.5 rounded-full bg-gray-400/50 animate-pulse" />
+                              <div className="h-1.5 w-1.5 rounded-full bg-gray-400/50 animate-pulse delay-150" />
+                              <div className="h-1.5 w-1.5 rounded-full bg-gray-400/50 animate-pulse delay-300" />
+                            </div>
+                          </div>
+                        </div>
                         </div>
                       </div>
                   )}
diff --git a/frontend/src/components/thread/file-viewer-modal.tsx b/frontend/src/components/thread/file-viewer-modal.tsx
index 0275637f..8ad27b4d 100644
--- a/frontend/src/components/thread/file-viewer-modal.tsx
+++ b/frontend/src/components/thread/file-viewer-modal.tsx
@@ -61,7 +61,7 @@ export function FileViewerModal({
 
   // Handle initial file path when provided
   useEffect(() => {
-    if (open && sandboxId && initialFilePath) {
+    if (open && sandboxId && initialFilePath && typeof initialFilePath === 'string') {
       // Extract the directory path from the file path
       const filePath = initialFilePath.startsWith('/workspace/') 
         ? initialFilePath