chore: streaming tool output (#5237)

* enhancement: tool streaming output

* chore: update memo

* fix: streaming

* chore: update stream tools arguments

* chore: update condition

* fix: style

* fix: style

* chore: fix stop button

* chore: update color accent and hide arrow button

---------

Co-authored-by: Louis <louis@jan.ai>

web-app/src/containers/RenderMarkdown.tsx

@@ -22,6 +22,7 @@ interface MarkdownProps {
   components?: Components
   enableRawHtml?: boolean
   isUser?: boolean
+  isWrapping?: boolean
 }
 
 function RenderMarkdownComponent({
@@ -30,6 +31,7 @@ function RenderMarkdownComponent({
   className,
   isUser,
   components,
+  isWrapping,
 }: MarkdownProps) {
   const { codeBlockStyle, showLineNumbers } = useCodeblock()
 
@@ -117,9 +119,13 @@ function RenderMarkdownComponent({
               showLineNumbers={showLineNumbers}
               wrapLines={true}
               // Temporary comment we try calculate main area width on __root
-              // lineProps={{
+              lineProps={
-              //   style: { wordBreak: 'break-all', whiteSpace: 'pre-wrap' },
+                isWrapping
-              // }}
+                  ? {
+                      style: { wordBreak: 'break-all', whiteSpace: 'pre-wrap' },
+                    }
+                  : {}
+              }
               customStyle={{
                 margin: 0,
                 padding: '8px',

web-app/src/containers/StreamingContent.tsx

@@ -32,6 +32,12 @@ export const StreamingContent = memo(({ threadId }: Props) => {
     return extractReasoningSegment(text)
   }, [streamingContent])
 
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const streamingTools: any = useMemo(() => {
+    const calls = streamingContent?.metadata?.tool_calls
+    return calls
+  }, [streamingContent])
+
   const lastAssistant = useMemo(() => {
     return [...messages].reverse().find((m) => m.role === 'assistant')
   }, [messages])
@@ -52,6 +58,14 @@ export const StreamingContent = memo(({ threadId }: Props) => {
   // The streaming content is always the last message
   return (
     <ThreadContent
+      streamTools={{
+        tool_calls: {
+          function: {
+            name: streamingTools?.[0]?.function?.name as string,
+            arguments: streamingTools?.[0]?.function?.arguments as string,
+          },
+        },
+      }}
       {...streamingContent}
       isLastMessage={true}
       showAssistant={

web-app/src/containers/ThreadContent.tsx

@@ -77,6 +77,8 @@ export const ThreadContent = memo(
       isLastMessage?: boolean
       index?: number
       showAssistant?: boolean
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      streamTools?: any
     }
   ) => {
     const [message, setMessage] = useState(item.content?.[0]?.text?.value || '')
@@ -316,8 +318,17 @@ export const ThreadContent = memo(
                 {(item.metadata.tool_calls as ToolCall[]).map((toolCall) => (
                   <ToolCallBlock
                     id={toolCall.tool?.id ?? 0}
-                    name={toolCall.tool?.function?.name ?? ''}
                     key={toolCall.tool?.id}
+                    name={
+                      (item.streamTools?.tool_calls?.function?.name ||
+                        toolCall.tool?.function?.name) ??
+                      ''
+                    }
+                    args={
+                      item.streamTools?.tool_calls?.function?.arguments ||
+                      toolCall.tool?.function?.arguments ||
+                      undefined
+                    }
                     result={JSON.stringify(toolCall.response)}
                     loading={toolCall.state === 'pending'}
                   />
@@ -332,7 +343,7 @@ export const ThreadContent = memo(
                     'flex items-center gap-2',
                     item.isLastMessage &&
                       streamingContent &&
-                      'opacity-0 visinility-hidden pointer-events-none'
+                      'opacity-0 visibility-hidden pointer-events-none'
                   )}
                 >
                   <CopyButton text={item.content?.[0]?.text.value || ''} />

web-app/src/containers/ToolCallBlock.tsx

@@ -9,10 +9,12 @@ import {
   DialogHeader,
   DialogTitle,
 } from '@/components/ui/dialog'
+import { twMerge } from 'tailwind-merge'
 
 interface Props {
   result: string
   name: string
+  args: object
   id: number
   loading: boolean
 }
@@ -90,7 +92,7 @@ const ContentItemRenderer = ({
   if (item.type === 'image' && item.data && item.mimeType) {
     const imageUrl = createDataUrl(item.data, item.mimeType)
     return (
-      <div key={index} className="mt-3">
+      <div key={index} className="my-3">
         <img
           src={imageUrl}
           alt={`Result image ${index + 1}`}
@@ -123,9 +125,9 @@ const ContentItemRenderer = ({
   )
 }
 
-const ToolCallBlock = ({ id, name, result, loading }: Props) => {
+const ToolCallBlock = ({ id, name, result, loading, args }: Props) => {
   const { collapseState, setCollapseState } = useToolCallBlockStore()
-  const isExpanded = collapseState[id] ?? false
+  const isExpanded = collapseState[id] ?? (loading ? true : false)
   const [modalImage, setModalImage] = useState<{
     url: string
     alt: string
@@ -145,10 +147,9 @@ const ToolCallBlock = ({ id, name, result, loading }: Props) => {
   }
 
   // Parse the MCP response and extract content items
-  const { parsedResult, contentItems, hasStructuredContent, parseError } =
+  const { parsedResult, contentItems, hasStructuredContent } = useMemo(() => {
-    useMemo(() => {
+    return parseMCPResponse(result)
-      return parseMCPResponse(result)
+  }, [result])
-    }, [result])
 
   return (
     <div
@@ -158,17 +159,36 @@ const ToolCallBlock = ({ id, name, result, loading }: Props) => {
       <div className="rounded-lg bg-main-view-fg/4 border border-dashed border-main-view-fg/10">
         <div className="flex items-center gap-3 p-2" onClick={handleClick}>
           {loading && (
-            <Loader className="size-4 animate-spin text-main-view-fg/60" />
+            <div className="w-4 h-4">
+              <Loader className="size-4 animate-spin text-main-view-fg/60" />
+            </div>
           )}
           <button className="flex items-center gap-2 focus:outline-none">
-            {isExpanded ? (
+            {!loading && (
-              <ChevronUp className="h-4 w-4" />
+              <>
-            ) : (
+                {isExpanded ? (
-              <ChevronDown className="h-4 w-4" />
+                  <>
+                    <div className="ml-1 w-4 h-4">
+                      <ChevronUp className="h-4 w-4" />
+                    </div>
+                  </>
+                ) : (
+                  <div className="ml-1 w-4 h-4">
+                    <ChevronDown className="h-4 w-4" />
+                  </div>
+                )}
+              </>
             )}
             <span className="font-medium text-main-view-fg/80">
-              View result from{' '}
+              <span className="font-medium text-main-view-fg mr-2">{name}</span>
-              <span className="font-medium text-main-view-fg">{name}</span>
+              <span
+                className={twMerge(
+                  'text-xs bg-main-view-fg/4 rounded-sm p-1',
+                  loading ? 'text-main-view-fg/40' : 'text-accent'
+                )}
+              >
+                {loading ? 'Calling tool' : 'Completed'}{' '}
+              </span>
             </span>
           </button>
         </div>
@@ -179,36 +199,43 @@ const ToolCallBlock = ({ id, name, result, loading }: Props) => {
             isExpanded ? '' : 'max-h-0 overflow-hidden'
           )}
         >
-          <div className="mt-2 text-main-view-fg/60">
+          <div className="mt-2 text-main-view-fg/60 max-w-[89%] overflow-hidden">
-            {hasStructuredContent ? (
+            {args && Object.keys(args).length > 3 && (
-              /* Render each content item individually based on its type */
+              <>
-              <div className="space-y-2">
+                <p className="mb-3">Arguments:</p>
-                {contentItems.map((item, index) => (
+                <RenderMarkdown
-                  <ContentItemRenderer
+                  isWrapping={true}
-                    key={index}
+                  content={'```json\n' + args + '\n```'}
-                    item={item}
+                />
-                    index={index}
+              </>
-                    onImageClick={handleImageClick}
+            )}
+
+            {result && (
+              <>
+                <p>Output:</p>
+                {hasStructuredContent ? (
+                  /* Render each content item individually based on its type */
+                  <div className="space-y-2">
+                    {contentItems.map((item, index) => (
+                      <ContentItemRenderer
+                        key={index}
+                        item={item}
+                        index={index}
+                        onImageClick={handleImageClick}
+                      />
+                    ))}
+                  </div>
+                ) : (
+                  /* Fallback: render as JSON for valid JSON but unstructured responses */
+                  <RenderMarkdown
+                    content={
+                      '```json\n' +
+                      JSON.stringify(parsedResult, null, 2) +
+                      '\n```'
+                    }
                   />
-                ))}
+                )}
-              </div>
+              </>
-            ) : parseError ? (
-              /* Handle JSON parse error - render as plain text */
-              <div className="mt-3 p-3 bg-main-view-fg/5 rounded-md border border-main-view-fg/10">
-                <div className="text-sm font-medium text-main-view-fg/80 mb-2">
-                  Raw Response:
-                </div>
-                <div className="whitespace-pre-wrap font-mono text-sm">
-                  {parsedResult as string}
-                </div>
-              </div>
-            ) : (
-              /* Fallback: render as JSON for valid JSON but unstructured responses */
-              <RenderMarkdown
-                content={
-                  '```json\n' + JSON.stringify(parsedResult, null, 2) + '\n```'
-                }
-              />
             )}
           </div>
         </div>

web-app/src/hooks/useAppState.ts

@@ -2,6 +2,7 @@ import { create } from 'zustand'
 import { ThreadMessage } from '@janhq/core'
 import { MCPTool } from '@/types/completion'
 import { useAssistant } from './useAssistant'
+import { ChatCompletionMessageToolCall } from 'openai/resources'
 
 type AppState = {
   streamingContent?: ThreadMessage
@@ -10,8 +11,12 @@ type AppState = {
   serverStatus: 'running' | 'stopped' | 'pending'
   abortControllers: Record<string, AbortController>
   tokenSpeed?: TokenSpeed
+  currentToolCall?: ChatCompletionMessageToolCall
   setServerStatus: (value: 'running' | 'stopped' | 'pending') => void
   updateStreamingContent: (content: ThreadMessage | undefined) => void
+  updateCurrentToolCall: (
+    toolCall: ChatCompletionMessageToolCall | undefined
+  ) => void
   updateLoadingModel: (loading: boolean) => void
   updateTools: (tools: MCPTool[]) => void
   setAbortController: (threadId: string, controller: AbortController) => void
@@ -26,6 +31,7 @@ export const useAppState = create<AppState>()((set) => ({
   serverStatus: 'stopped',
   abortControllers: {},
   tokenSpeed: undefined,
+  currentToolCall: undefined,
   updateStreamingContent: (content: ThreadMessage | undefined) => {
     set(() => ({
       streamingContent: content
@@ -40,6 +46,11 @@ export const useAppState = create<AppState>()((set) => ({
         : undefined,
     }))
   },
+  updateCurrentToolCall: (toolCall) => {
+    set(() => ({
+      currentToolCall: toolCall,
+    }))
+  },
   updateLoadingModel: (loading) => {
     set({ loadingModel: loading })
   },

web-app/src/hooks/useChat.ts

@@ -181,8 +181,21 @@ export const useChat = () => {
           } else {
             for await (const part of completion) {
               const delta = part.choices[0]?.delta?.content || ''
+
               if (part.choices[0]?.delta?.tool_calls) {
-                extractToolCall(part, currentCall, toolCalls)
+                const calls = extractToolCall(part, currentCall, toolCalls)
+                const currentContent = newAssistantThreadContent(
+                  activeThread.id,
+                  accumulatedText,
+                  {
+                    tool_calls: calls.map((e) => ({
+                      ...e,
+                      state: 'pending',
+                    })),
+                  }
+                )
+                updateStreamingContent(currentContent)
+                await new Promise((resolve) => setTimeout(resolve, 0))
               }
               if (delta) {
                 accumulatedText += delta
@@ -190,7 +203,13 @@ export const useChat = () => {
                 // Use a timeout to prevent React from batching updates too quickly
                 const currentContent = newAssistantThreadContent(
                   activeThread.id,
-                  accumulatedText
+                  accumulatedText,
+                  {
+                    tool_calls: toolCalls.map((e) => ({
+                      ...e,
+                      state: 'pending',
+                    })),
+                  }
                 )
                 updateStreamingContent(currentContent)
                 updateTokenSpeed(currentContent)
@@ -225,6 +244,7 @@ export const useChat = () => {
             allowAllMCPPermissions
           )
           addMessage(updatedMessage ?? finalContent)
+          updateStreamingContent(emptyThreadContent)
           updateThreadTimestamp(activeThread.id)
 
           isCompleted = !toolCalls.length

web-app/src/lib/completion.ts

@@ -63,7 +63,8 @@ export const newUserThreadContent = (
  */
 export const newAssistantThreadContent = (
   threadId: string,
-  content: string
+  content: string,
+  metadata: Record<string, unknown> = {}
 ): ThreadMessage => ({
   type: 'text',
   role: ChatCompletionRole.Assistant,
@@ -82,6 +83,7 @@ export const newAssistantThreadContent = (
   status: MessageStatus.Ready,
   created_at: 0,
   completed_at: 0,
+  metadata,
 })
 
 /**

web-app/src/routes/__root.tsx

@@ -43,7 +43,7 @@ const AppLayout = () => {
           <div
             className={cn(
               'h-full flex w-full p-1',
-              isLeftPanelOpen && 'w-[calc(100%-192px)]'
+              isLeftPanelOpen && 'w-[calc(100%-198px)]'
             )}
           >
             <div className="bg-main-view text-main-view-fg border border-main-view-fg/5 w-full rounded-lg overflow-hidden">

web-app/src/types/message.d.ts

@@ -3,6 +3,7 @@ type ToolCall = {
     id?: number
     function?: {
       name?: string
+      arguments?: object
     }
   }
   response?: unknown