fix: hide code block when streaming (#2875)

Author: Kitenite

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/assistant-message.tsx

@@ -1,14 +1,14 @@
 import type { ChatMessage } from '@onlook/models';
 import { MessageContent } from './message-content';
 
-export const AssistantMessage = ({ message }: { message: ChatMessage }) => {
+export const AssistantMessage = ({ message, isStreaming }: { message: ChatMessage, isStreaming: boolean }) => {
     return (
         <div className="px-4 py-2 text-small content-start flex flex-col text-wrap gap-2">
             <MessageContent
                 messageId={message.id}
                 parts={message.parts}
                 applied={false}
-                isStream={false}
+                isStream={isStreaming}
             />
         </div>
     );

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/index.tsx

@@ -51,7 +51,7 @@ export const ChatMessages = observer(({
             let messageNode;
             switch (message.role) {
                 case 'assistant':
-                    messageNode = <AssistantMessage key={message.id} message={message} />;
+                    messageNode = <AssistantMessage key={message.id} message={message} isStreaming={isStreaming} />;
                     break;
                 case 'user':
                     messageNode = (
@@ -93,6 +93,10 @@ export const ChatMessages = observer(({
             {messages.map((message, index) => renderMessage(message, index))}
             {streamedMessage && <StreamMessage message={streamedMessage} />}
             {error && <ErrorMessage error={error} />}
+            {isStreaming && <div className="flex w-full h-full flex-row items-center gap-2 px-4 my-2 text-small content-start text-foreground-secondary">
+                <Icons.LoadingSpinner className="animate-spin" />
+                <p>Thinking ...</p>
+            </div>}
         </ChatMessageList>
     );
 },

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/message-content/tool-call-display.tsx

@@ -39,7 +39,7 @@ export const ToolCallDisplay = ({
     const toolName = toolPart.type.split('-')[1];
     const loading = isStream && isLastPart;
 
-    if (isStream && toolPart.state !== 'output-available' && toolPart.state !== 'input-available') {
+    if (isStream || (toolPart.state !== 'output-available' && toolPart.state !== 'input-available')) {
         return (
             <ToolCallSimple
                 toolPart={toolPart}

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/stream-message.tsx

@@ -1,27 +1,15 @@
 import { type ChatMessage } from '@onlook/models';
-import { Icons } from '@onlook/ui/icons';
-import { useMemo } from 'react';
 import { MessageContent } from './message-content';
 
 export const StreamMessage = ({ message }: { message: ChatMessage }) => {
-    const isAssistantStreamMessage = useMemo(() => message?.role === 'assistant', [message?.role]);
-
     return (
-        <>
-            {message && isAssistantStreamMessage && (
-                <div className="px-4 pt-2 text-small content-start flex flex-col text-wrap gap-2">
-                    <MessageContent
-                        messageId={message.id}
-                        parts={message.parts}
-                        applied={false}
-                        isStream={true}
-                    />
-                </div>
-            )}
-            <div className="flex w-full h-full flex-row items-center gap-2 px-4 my-2 text-small content-start text-foreground-secondary">
-                <Icons.LoadingSpinner className="animate-spin" />
-                <p>Thinking ...</p>
-            </div>
-        </>
+        <div className="px-4 pt-2 text-small content-start flex flex-col text-wrap gap-2">
+            <MessageContent
+                messageId={message.id}
+                parts={message.parts}
+                applied={false}
+                isStream={true}
+            />
+        </div>
     );
 };

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/code-display/collapsible-code-block.tsx

@@ -1,5 +1,5 @@
-import { api } from '@/trpc/react';
 import { useEditorEngine } from '@/components/store/editor';
+import { api } from '@/trpc/react';
 import { Button } from '@onlook/ui/button';
 import { Collapsible, CollapsibleContent, CollapsibleTrigger } from '@onlook/ui/collapsible';
 import { Icons } from '@onlook/ui/icons';
@@ -49,7 +49,6 @@ export const CollapsibleCodeBlock = observer(({
 
     return (
         <div className="group relative">
-
             <Collapsible open={isOpen} onOpenChange={setIsOpen}>
                 <div
                     className={cn(
@@ -83,16 +82,13 @@ export const CollapsibleCodeBlock = observer(({
                                 >
                                     <span className="truncate flex-1 min-w-0">{getTruncatedFileName(path)}</span>
                                     {(() => {
-                                        const branch = branchId 
+                                        const branch = branchId
                                             ? editorEngine.branches.allBranches.find(b => b.id === branchId)
                                             : editorEngine.branches.activeBranch;
                                         return branch && (
-                                            <>
-                                                
-                                                 <span className="text-foreground-tertiary group-hover:text-foreground-secondary text-mini ml-0.5 flex-shrink-0 truncate max-w-24">
-                                                     {' • '}{branch.name}
-                                                 </span>
-                                            </>
+                                            <span className="text-foreground-tertiary group-hover:text-foreground-secondary text-mini ml-0.5 flex-shrink-0 truncate max-w-24">
+                                                {' • '}{branch.name}
+                                            </span>
                                         );
                                     })()}
                                 </div>