From a81b644a8fc6d1c2904896882372231de3069447 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Wed, 23 Apr 2025 22:07:04 +0700
Subject: [PATCH] enhancement: tool call block should be wrapped in a
 collapsible scroll area

---
 .../ThreadCenterPanel/ChatBody/index.tsx      | 35 +++++++++++++++++--
 .../ThreadCenterPanel/ChatItem/index.tsx      |  2 ++
 .../TextMessage/ToolCallBlock.tsx             | 21 +++++++----
 .../ThreadCenterPanel/TextMessage/index.tsx   |  7 +++-
 4 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx
index 80d33b22f..3f5bf690a 100644
--- a/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx
@@ -4,7 +4,7 @@ import { ThreadMessage } from '@janhq/core'
 import { ScrollArea } from '@janhq/joi'
 import { useVirtualizer } from '@tanstack/react-virtual'
 
-import { useAtomValue } from 'jotai'
+import { useAtomValue, useSetAtom } from 'jotai'
 
 import { loadModelErrorAtom } from '@/hooks/useActiveModel'
 
@@ -12,6 +12,8 @@ import ChatItem from '../ChatItem'
 
 import LoadModelError from '../LoadModelError'
 
+import { toolCallBlockStateAtom } from '../TextMessage/ToolCallBlock'
+
 import EmptyThread from './EmptyThread'
 
 import { getCurrentChatMessagesAtom } from '@/helpers/atoms/ChatMessage.atom'
@@ -65,9 +67,11 @@ const ChatBody = memo(
     const parentRef = useRef<HTMLDivElement>(null)
     const prevScrollTop = useRef(0)
     const isUserManuallyScrollingUp = useRef(false)
+    const isNestedScrollviewExpanding = useRef(false)
     const currentThread = useAtomValue(activeThreadAtom)
     const isBlockingSend = useAtomValue(isBlockingSendAtom)
     const showScrollBar = useAtomValue(showScrollBarAtom)
+    const setToolCallExpanded = useSetAtom(toolCallBlockStateAtom)
 
     const count = useMemo(
       () => (messages?.length ?? 0) + (loadModelError ? 1 : 0),
@@ -97,7 +101,10 @@ const ChatBody = memo(
       _,
       instance
     ) => {
-      if (isUserManuallyScrollingUp.current === true && isBlockingSend)
+      if (
+        isNestedScrollviewExpanding ||
+        (isUserManuallyScrollingUp.current === true && isBlockingSend)
+      )
         return false
       return (
         // item.start < (instance.scrollOffset ?? 0) &&
@@ -130,6 +137,22 @@ const ChatBody = memo(
       [isBlockingSend]
     )
 
+    const preserveScrollOnExpand = (callback: () => void) => {
+      isNestedScrollviewExpanding.current = true
+      const scrollEl = parentRef.current
+      const prevScrollTop = scrollEl?.scrollTop ?? 0
+      const prevScrollHeight = scrollEl?.scrollHeight ?? 0
+
+      callback() // Expand content (e.g. setIsExpanded(true))
+
+      if (scrollEl)
+        requestAnimationFrame(() => {
+          const newScrollHeight = scrollEl?.scrollHeight ?? 0
+          scrollEl.scrollTop =
+            prevScrollTop + (newScrollHeight - prevScrollHeight)
+        })
+    }
+
     return (
       <div className="flex h-full w-full flex-col overflow-x-hidden">
         <ScrollArea
@@ -179,6 +202,14 @@ const ChatBody = memo(
                       isCurrentMessage={
                         virtualRow.index === messages?.length - 1
                       }
+                      onExpand={(props) =>
+                        preserveScrollOnExpand(() => {
+                          setToolCallExpanded((prev) => ({
+                            ...prev,
+                            ...props,
+                          }))
+                        })
+                      }
                     />
                   )}
                 </div>
diff --git a/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
index 75b566ceb..0d55e9ea9 100644
--- a/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
@@ -22,6 +22,7 @@ type Props = {
   loadModelError?: string
   isCurrentMessage?: boolean
   index: number
+  onExpand: (props: { [id: number]: boolean }) => void
 } & ThreadMessage
 
 const ChatItem = forwardRef<Ref, Props>((message, ref) => {
@@ -81,6 +82,7 @@ const ChatItem = forwardRef<Ref, Props>((message, ref) => {
               status={status}
               index={message.index}
               isCurrentMessage={message.isCurrentMessage ?? false}
+              onExpand={message.onExpand}
             />
           </div>
         )}
diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/ToolCallBlock.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/ToolCallBlock.tsx
index 818af5b4d..bd28e01b2 100644
--- a/web/screens/Thread/ThreadCenterPanel/TextMessage/ToolCallBlock.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/ToolCallBlock.tsx
@@ -1,8 +1,11 @@
 import React from 'react'
 
+import { ScrollArea } from '@janhq/joi'
 import { atom, useAtom } from 'jotai'
 import { ChevronDown, ChevronUp, Loader } from 'lucide-react'
 
+import { twMerge } from 'tailwind-merge'
+
 import { MarkdownTextMessage } from './MarkdownTextMessage'
 
 interface Props {
@@ -10,16 +13,18 @@ interface Props {
   name: string
   id: number
   loading: boolean
+  onExpand: (props: { [id: number]: boolean }) => void
 }
 
-const toolCallBlockStateAtom = atom<{ [id: number]: boolean }>({})
+export const toolCallBlockStateAtom = atom<{ [id: number]: boolean }>({})
 
-const ToolCallBlock = ({ id, name, result, loading }: Props) => {
+const ToolCallBlock = ({ id, name, result, loading, onExpand }: Props) => {
   const [collapseState, setCollapseState] = useAtom(toolCallBlockStateAtom)
 
   const isExpanded = collapseState[id] ?? false
   const handleClick = () => {
-    setCollapseState((prev) => ({ ...prev, [id]: !isExpanded }))
+    onExpand({ [id]: !isExpanded })
+    // setCollapseState((prev) => ({ ...prev, [id]: !isExpanded }))
   }
   return (
     <div className="mx-auto w-full">
@@ -38,17 +43,21 @@ const ToolCallBlock = ({ id, name, result, loading }: Props) => {
               <ChevronDown className="h-4 w-4" />
             )}
             <span className="font-medium">
-              {' '}
               View result from <span className="font-bold">{name}</span>
             </span>
           </button>
         </div>
 
-        {isExpanded && (
+        <ScrollArea
+          className={twMerge(
+            'w-full overflow-hidden transition-all duration-300',
+            isExpanded ? 'max-h-96' : 'max-h-0'
+          )}
+        >
           <div className="mt-2 overflow-x-hidden pl-6 text-[hsla(var(--text-secondary))]">
             <span>{result ?? ''} </span>
           </div>
-        )}
+        </ScrollArea>
       </div>
     </div>
   )
diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
index 9d1558ec9..81fe1ce56 100644
--- a/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
@@ -30,7 +30,11 @@ import { selectedModelAtom } from '@/helpers/atoms/Model.atom'
 import { chatWidthAtom } from '@/helpers/atoms/Setting.atom'
 
 const MessageContainer: React.FC<
-  ThreadMessage & { isCurrentMessage: boolean; index: number }
+  ThreadMessage & {
+    isCurrentMessage: boolean
+    index: number
+    onExpand: (props: { [id: number]: boolean }) => void
+  }
 > = (props) => {
   const isUser = props.role === ChatCompletionRole.User
   const isSystem = props.role === ChatCompletionRole.System
@@ -195,6 +199,7 @@ const MessageContainer: React.FC<
                 <>
                   {props.metadata.tool_calls.map((toolCall) => (
                     <ToolCallBlock
+                      onExpand={props.onExpand}
                       id={toolCall.tool?.id}
                       name={toolCall.tool?.function?.name ?? ''}
                       key={toolCall.tool?.id}