set the run status to running as soon as event streaming starts (#149)

---------
Co-authored-by: thucpn <thucsh2@gmail.com>

.changeset/empty-mails-laugh.md

@@ -0,0 +1,5 @@
+---
+'@llamaindex/chat-ui': patch
+---
+
+set the run status to running as soon as event streaming starts

.prettierignore

@@ -0,0 +1 @@
+.venv/

packages/chat-ui/src/hook/use-workflow/hook.ts

@@ -42,9 +42,6 @@ export function useWorkflow<E extends WorkflowEvent = WorkflowEvent>(
       await fetchTaskEvents<E>(
         { client, deploymentName, task: inputTask },
         {
-          onStart: () => {
-            setStatus('running')
-          },
           onData: event => {
             setEvents(prev => [...prev, event])
             onData?.(event)
@@ -78,6 +75,7 @@ export function useWorkflow<E extends WorkflowEvent = WorkflowEvent>(
 
       setTask(existingTask)
       setIsInitialized(true)
+      setStatus('running')
       await streamTaskEvents(existingTask)
     }
 
@@ -104,6 +102,7 @@ export function useWorkflow<E extends WorkflowEvent = WorkflowEvent>(
         workflow,
       })
       setTask(newTask) // update new task with new session when trigger start event
+      setStatus('running') // set status running as soon as start streaming events
       await streamTaskEvents(newTask)
     },
     [client, deploymentName, streamTaskEvents, workflow]

packages/chat-ui/src/hook/use-workflow/types.ts

@@ -20,7 +20,7 @@ export interface WorkflowHookParams<E extends WorkflowEvent = WorkflowEvent> {
 
 export interface WorkflowHookHandler<E extends WorkflowEvent = WorkflowEvent> {
   runId?: string // Task ID used internally, will be the same for the whole session
-  start: (eventData: E['data']) => Promise<void> // Create new task with start event data, updates sessionId
+  start: (eventData: E['data']) => Promise<void> // Create new task with start event data, updates sessionId. Returned promise resolves on task stop (completion)
   stop: (data?: E['data']) => Promise<void> // Send stop event to stop current task
   sendEvent: (event: E) => Promise<void> // Function to send a new event to the current session, throws error if session is not created yet
   events: E[] // events of the current run

packages/chat-ui/src/test/__tests__/use-workflow.test.tsx

@@ -51,6 +51,63 @@ describe('useWorkflow', () => {
   })
 
   describe('Starting workflow', () => {
+    it('should set status to running after task created', async () => {
+      const originalHandlers = [...server.listHandlers()]
+      let streamController: ReadableStreamDefaultController<Uint8Array>
+
+      try {
+        server.use(
+          http.get(
+            'http://127.0.0.1:4501/deployments/test-workflow/tasks/test-run-id/events',
+            () => {
+              // Create a proper SSE stream with headers
+              const stream = new ReadableStream({
+                start(controller) {
+                  streamController = controller
+                  // Don't send any data initially - simulating long running task
+                },
+              })
+
+              return new Response(stream, {
+                status: 200,
+                headers: {
+                  'Content-Type': 'text/event-stream',
+                  'Cache-Control': 'no-cache',
+                  Connection: 'keep-alive',
+                },
+              })
+            }
+          )
+        )
+
+        const { result } = renderHook(() =>
+          useWorkflow<TestEvent>(mockWorkflowParams)
+        )
+        expect(result.current.status).toBeUndefined()
+
+        act(() => {
+          result.current.start({ message: 'test start' })
+        })
+
+        await waitFor(() => {
+          expect(result.current.runId).toBe('test-run-id')
+          // should set to running despite no events received
+          expect(result.current.status).toBe('running')
+        })
+
+        // Close the stream to complete the workflow
+        act(() => {
+          streamController.close()
+        })
+
+        await waitFor(() => {
+          expect(result.current.status).toBe('complete')
+        })
+      } finally {
+        server.resetHandlers(...originalHandlers)
+      }
+    })
+
     it('should create a new task and start streaming events', async () => {
       const { result } = renderHook(() =>
         useWorkflow<TestEvent>(mockWorkflowParams)
@@ -119,6 +176,64 @@ describe('useWorkflow', () => {
         expect(result.current.status).toBe('complete')
       })
     })
+
+    it('should set status to running immediately when reconnecting to existing task', async () => {
+      const originalHandlers = [...server.listHandlers()]
+      let streamController: ReadableStreamDefaultController<Uint8Array>
+
+      try {
+        server.use(
+          http.get(
+            'http://127.0.0.1:4501/deployments/test-workflow/tasks/test-existing-run-id/events',
+            () => {
+              // Create a proper SSE stream with headers
+              const stream = new ReadableStream({
+                start(controller) {
+                  streamController = controller
+                  // Don't send any data initially - simulating long running task
+                },
+              })
+
+              return new Response(stream, {
+                status: 200,
+                headers: {
+                  'Content-Type': 'text/event-stream',
+                  'Cache-Control': 'no-cache',
+                  Connection: 'keep-alive',
+                },
+              })
+            }
+          )
+        )
+
+        const paramsWithRunId = {
+          ...mockWorkflowParams,
+          runId: 'test-existing-run-id',
+        }
+
+        const { result } = renderHook(() =>
+          useWorkflow<TestEvent>(paramsWithRunId)
+        )
+
+        // Should set status to running immediately when reconnecting to existing task
+        await waitFor(() => {
+          expect(result.current.runId).toBe('test-existing-run-id')
+          // should set to running after headers received but before events
+          expect(result.current.status).toBe('running')
+        })
+
+        // Close the stream to complete the workflow
+        act(() => {
+          streamController.close()
+        })
+
+        await waitFor(() => {
+          expect(result.current.status).toBe('complete')
+        })
+      } finally {
+        server.resetHandlers(...originalHandlers)
+      }
+    })
   })
 
   describe('Error handling', () => {