feat: add unit tests for chat-ui (#137)

---------
Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.changeset/slimy-mails-smash.md

@@ -0,0 +1,6 @@
+---
+'@llamaindex/chat-ui': patch
+'@llamaindex/chat-ui-docs': patch
+---
+
+feat: specify workflow to run of the given deployment

.github/workflows/ci.yml

@@ -40,3 +40,7 @@ jobs:
 
       - name: Type Check
         run: pnpm type-check
+
+      - name: Run Unit Tests
+        run: pnpm test
+        working-directory: packages/chat-ui

docs/chat-ui/hooks.mdx

@@ -412,8 +412,7 @@ import { useWorkflow } from '@llamaindex/chat-ui'
 
 function WorkflowComponent() {
   const {
-    sessionId,
-    taskId,
+    runId,
     start,
     stop,
     sendEvent,
@@ -421,6 +420,7 @@ function WorkflowComponent() {
     status,
   } = useWorkflow({
     baseUrl: 'http://localhost:8000',
+    deployment: 'my-deployment',
     workflow: 'my-workflow',
     onStopEvent: (event) => {
       console.log('Workflow stopped:', event)
@@ -447,15 +447,14 @@ function WorkflowComponent() {
 
   return (
     <div>
-      <p>Session ID: {sessionId}</p>
-      <p>Task ID: {taskId}</p>
+      <p>Run ID: {runId}</p>
       <p>Status: {status}</p>
       <p>Events received: {events.length}</p>
 
       <button onClick={handleStartWorkflow} disabled={status === 'running'}>
         Start Workflow
       </button>
-      <button onClick={handleSendCustomEvent} disabled={!sessionId}>
+      <button onClick={handleSendCustomEvent} disabled={!runId}>
         Send Event
       </button>
       <button onClick={handleStopWorkflow} disabled={status !== 'running'}>
@@ -478,20 +477,20 @@ function WorkflowComponent() {
 **Options:**
 
 - `baseUrl` - Base URL for the workflow API (optional)
-- `workflow` - Name of the registered deployment
-- `taskId` - Optional task ID for resuming an existing workflow task
+- `deployment` - Name of the registered deployment
+- `workflow` - Set the service to run - if not set uses the default workflow (optional)
+- `runId` - Optional run ID for resuming a previous workflow run
 - `onStopEvent` - Callback when workflow receives a stop event
 - `onError` - Error callback for workflow errors
 
 **Returned Properties:**
 
-- `sessionId` - Current workflow session ID
-- `taskId` - Current workflow task ID
+- `runId` - Current workflow run ID
 - `start` - Function to start a new workflow with initial event data
 - `stop` - Function to send stop event to current workflow
 - `sendEvent` - Function to send custom events to the workflow
 - `events` - Array of all received workflow events
-- `status` - Current workflow status ('idle' | 'running' | 'complete' | 'error')
+- `status` - Current workflow run status ('idle' | 'running' | 'complete' | 'error')
 
 **Workflow Event Types:**
 
@@ -515,8 +514,8 @@ Resume an existing workflow task:
 ```tsx
 function ResumeWorkflowComponent() {
   const workflow = useWorkflow({
-    workflow: 'my-workflow',
-    taskId: 'existing-task-id', // Resume existing task
+    deployment: 'my-deployment',
+    runId: 'existing-run-id', // Resume existing task
     onStopEvent: (event) => {
       console.log('Workflow completed:', event)
     },
@@ -529,7 +528,7 @@ function ResumeWorkflowComponent() {
 
   return (
     <div>
-      <p>Resumed task: {workflow.taskId}</p>
+      <p>Resumed run: {workflow.runId}</p>
       <p>Events: {workflow.events.length}</p>
     </div>
   )
@@ -550,6 +549,7 @@ interface CustomWorkflowEvent extends WorkflowEvent {
 
 function CustomWorkflowComponent() {
   const workflow = useWorkflow<CustomWorkflowEvent>({
+    deployment: 'progress-deployment',
     workflow: 'progress-workflow',
     onStopEvent: (event) => {
       if (event.type === 'workflow.ResultEvent') {

examples/llama-deploy/llama_deploy.yml

@@ -11,7 +11,13 @@ services:
     source:
       type: local
       name: src
-    path: src/workflow:echo_workflow
+    path: src/echo_workflow:workflow
+  adhoc_workflow:
+    name: Adhoc Workflow
+    source:
+      type: local
+      name: src
+    path: src/adhoc_workflow:workflow
 
 ui:
   name: My Nextjs App

examples/llama-deploy/src/adhoc_workflow.py

@@ -18,7 +18,7 @@ class AdhocEvent(Event):
     pass
 
 
-class EchoWorkflow(Workflow):
+class AdhocWorkflow(Workflow):
     """A dummy workflow with only one step sending back the input given."""
 
     def __init__(self, **kwargs):
@@ -62,4 +62,4 @@ class EchoWorkflow(Workflow):
         )
 
 
-echo_workflow = EchoWorkflow()
+workflow = AdhocWorkflow()

examples/llama-deploy/src/echo_workflow.py

@@ -0,0 +1,31 @@
+import asyncio
+
+from llama_index.core.workflow import (
+    Workflow,
+    StartEvent,
+    StopEvent,
+    step,
+    Context,
+    Event,
+)
+
+
+class UIEvent(Event):
+    data: dict
+
+
+class EchoWorkflow(Workflow):
+    @step()
+    async def run_step(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        message = str(ev.get("message", ""))
+
+        # Loop, increase the counter to 10 and show it in UI
+        for i in range(10):
+            ctx.write_event_to_stream(UIEvent(data={"counter": i}))
+            await ctx.set("counter", i + 1)
+            await asyncio.sleep(1)
+
+        return StopEvent(result=f"Message received: {message}")
+
+
+workflow = EchoWorkflow()

examples/llama-deploy/ui/app/page.tsx

@@ -4,22 +4,25 @@ import { useState } from 'react'
 import { useWorkflow } from '@llamaindex/chat-ui'
 
 const baseUrl = process.env.NEXT_PUBLIC_BASE_PATH ?? ''
-const workflow = process.env.NEXT_PUBLIC_DEPLOYMENT_NAME || 'QuickStart'
+const deployment = process.env.NEXT_PUBLIC_DEPLOYMENT_NAME || 'QuickStart'
+const defaultWorkflow =
+  process.env.NEXT_PUBLIC_WORKFLOW_NAME || 'adhoc_workflow'
 
 export default function Home() {
   const [userInput, setUserInput] = useState('Please run task')
+  const [workflow, setWorkflow] = useState(defaultWorkflow)
 
-  const { sessionId, taskId, start, stop, sendEvent, events, status } =
-    useWorkflow({
-      baseUrl,
-      workflow,
-      onStopEvent: event => {
-        console.log('Stop event:', event)
-      },
-      onError: error => {
-        console.error('Error:', error)
-      },
-    })
+  const { runId, start, stop, sendEvent, events, status } = useWorkflow({
+    baseUrl,
+    deployment,
+    workflow,
+    onStopEvent: event => {
+      console.log('Stop event:', event)
+    },
+    onError: error => {
+      console.error('Error:', error)
+    },
+  })
 
   const handleStart = async () => {
     await start({ message: userInput })
@@ -27,33 +30,52 @@ export default function Home() {
 
   const handleRetrieve = async () => {
     // AdhocEvent is defined in workflow definition
-    await sendEvent({ type: 'workflow.AdhocEvent' })
+    await sendEvent({ type: 'adhoc_workflow.AdhocEvent' })
   }
 
   const handleStop = async () => {
     await stop()
   }
 
+  const handleWorkflowSwitch = () => {
+    const newWorkflow =
+      workflow === 'adhoc_workflow' ? 'echo_workflow' : 'adhoc_workflow'
+    setWorkflow(newWorkflow)
+  }
+
   return (
     <div className="mx-auto h-screen w-full max-w-4xl px-4 py-4">
       <h1 className="mb-6 text-2xl font-bold">Llama Deploy with Chat UI</h1>
 
+      {/* Workflow Switcher */}
+      <div className="mb-4 flex items-center gap-4">
+        <button
+          type="button"
+          onClick={handleWorkflowSwitch}
+          disabled={status === 'running'}
+          className="ml-auto rounded-full bg-blue-500 px-3 py-1.5 text-white shadow-lg hover:bg-blue-600 disabled:opacity-50"
+        >
+          Switch to{' '}
+          {workflow === 'adhoc_workflow' ? 'echo_workflow' : 'adhoc_workflow'}
+        </button>
+      </div>
+
       {/* Status Panel */}
       <div className="mb-6 rounded-lg bg-gray-100 p-4">
-        <h2 className="mb-2 text-lg font-semibold">Workflow Status</h2>
         <div className="grid grid-cols-2 gap-4 text-sm">
           <div>
-            <strong>Session ID:</strong> {sessionId || 'Not created'}
-          </div>
-          <div>
-            <strong>Task ID:</strong> {taskId || 'Not created'}
+            <strong>Workflow:</strong> {workflow}
           </div>
 
           <div>
-            <strong>Task Events:</strong> {events.length}
+            <strong>Run ID:</strong> {runId || 'Not created'}
+          </div>
+
+          <div>
+            <strong>Event Count:</strong> {events.length}
           </div>
           <div>
-            <strong>Task Status:</strong>{' '}
+            <strong>Run Status:</strong>{' '}
             <span
               className={`text-sm ${
                 status === 'idle' || !status
@@ -88,7 +110,9 @@ export default function Home() {
         <input
           type="text"
           value={userInput}
-          onChange={e => setUserInput(e.target.value)}
+          onChange={e => {
+            setUserInput(e.target.value)
+          }}
           className="flex-1 rounded border border-gray-300 p-2"
           placeholder="Type your message..."
           disabled={status === 'running'}
@@ -101,14 +125,16 @@ export default function Home() {
         >
           Start
         </button>
-        <button
-          type="button"
-          onClick={handleRetrieve}
-          disabled={status !== 'running'}
-          className="rounded-full bg-yellow-500 px-6 py-2 text-white shadow-2xl hover:bg-yellow-600 disabled:opacity-50"
-        >
-          Retrieve
-        </button>
+        {workflow === 'adhoc_workflow' && (
+          <button
+            type="button"
+            onClick={handleRetrieve}
+            disabled={status !== 'running'}
+            className="rounded-full bg-yellow-500 px-6 py-2 text-white shadow-2xl hover:bg-yellow-600 disabled:opacity-50"
+          >
+            Retrieve
+          </button>
+        )}
         <button
           type="button"
           onClick={handleStop}

packages/chat-ui/package.json

@@ -36,7 +36,10 @@
     "type-check": "tsc --noEmit",
     "clean": "rm -rf dist",
     "copy": "cp -r ../../README.md ../../LICENSE .",
-    "postbuild": "pnpm run copy"
+    "postbuild": "pnpm run copy",
+    "test": "vitest",
+    "test:watch": "vitest --watch",
+    "test:ui": "vitest --ui"
   },
   "devDependencies": {
     "@llamaindex/eslint-config": "workspace:*",
@@ -47,7 +50,15 @@
     "postcss-cli": "^11.0.0",
     "tailwindcss": "^4.0.7",
     "bunchee": "^6.5.2",
-    "typescript": "^5.3.3"
+    "typescript": "^5.3.3",
+    "@testing-library/react": "^16.3.0",
+    "@testing-library/react-hooks": "^8.0.1",
+    "@testing-library/jest-dom": "^6.6.3",
+    "@vitejs/plugin-react": "^4.5.2",
+    "jsdom": "^26.1.0",
+    "vitest": "^3.2.4",
+    "@vitest/ui": "^3.2.4",
+    "msw": "^2.10.2"
   },
   "dependencies": {
     "@llamaindex/llama-deploy": "workspace:*",

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

@@ -42,12 +42,16 @@ export async function createTask<E extends WorkflowEvent>(params: {
   client: Client
   deploymentName: string
   eventData: E['data']
+  workflow?: string // create task in default service if not provided
 }): Promise<WorkflowTask> {
   const data =
     await createDeploymentTaskNowaitDeploymentsDeploymentNameTasksCreatePost({
       client: params.client,
       path: { deployment_name: params.deploymentName },
-      body: { input: JSON.stringify(params.eventData ?? {}) },
+      body: {
+        input: JSON.stringify(params.eventData ?? {}),
+        service_id: params.workflow,
+      },
     })
 
   const { task_id, session_id, service_id, input } = data.data ?? {}

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

@@ -11,7 +11,7 @@ import {
   WorkflowEventType,
   WorkflowHookHandler,
   WorkflowHookParams,
-  WorkflowStatus,
+  RunStatus,
   WorkflowTask,
 } from './types'
 
@@ -20,8 +20,9 @@ export function useWorkflow<E extends WorkflowEvent = WorkflowEvent>(
 ): WorkflowHookHandler<E> {
   const {
     baseUrl,
-    workflow: deploymentName,
-    taskId: initialTaskId,
+    deployment: deploymentName,
+    workflow,
+    runId: initialTaskId,
     onStopEvent,
     onError,
   } = params
@@ -29,7 +30,7 @@ export function useWorkflow<E extends WorkflowEvent = WorkflowEvent>(
   const [isInitialized, setIsInitialized] = useState(false)
   const [task, setTask] = useState<WorkflowTask>()
   const [events, setEvents] = useState<E[]>([])
-  const [status, setStatus] = useState<WorkflowStatus>()
+  const [status, setStatus] = useState<RunStatus>()
 
   const client = useMemo(() => {
     return createClient(createConfig({ baseUrl }))
@@ -94,11 +95,16 @@ export function useWorkflow<E extends WorkflowEvent = WorkflowEvent>(
   const start = useCallback(
     async (eventData: E['data']) => {
       setEvents([]) // reset events when start a new task
-      const newTask = await createTask({ client, deploymentName, eventData })
+      const newTask = await createTask({
+        client,
+        deploymentName,
+        eventData,
+        workflow,
+      })
       setTask(newTask) // update new task with new session when trigger start event
       await streamTaskEvents(newTask)
     },
-    [client, deploymentName, streamTaskEvents]
+    [client, deploymentName, streamTaskEvents, workflow]
   )
 
   const stop = useCallback(
@@ -110,8 +116,7 @@ export function useWorkflow<E extends WorkflowEvent = WorkflowEvent>(
   )
 
   return {
-    sessionId: task?.session_id,
-    taskId: task?.task_id,
+    runId: task?.task_id,
     sendEvent,
     start,
     stop,

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

@@ -6,24 +6,24 @@ export interface WorkflowEvent {
   data?: JSONValue | undefined
 }
 
-export type WorkflowStatus = 'idle' | 'running' | 'complete' | 'error'
+export type RunStatus = 'idle' | 'running' | 'complete' | 'error'
 
 export interface WorkflowHookParams<E extends WorkflowEvent = WorkflowEvent> {
   baseUrl?: string // Optional base URL for the workflow API
-  workflow: string // Name of the registered deployment
-  taskId?: string // Optional task ID for resuming a workflow task
+  deployment: string // Name of the registered deployment
+  runId?: string // Optional task ID for resuming a workflow task
+  workflow?: string // Set the default service to run
   onStopEvent?: (event: E) => void
   onError?: (error: any) => void
 }
 
 export interface WorkflowHookHandler<E extends WorkflowEvent = WorkflowEvent> {
-  sessionId?: string // Session ID once the workflow session starts
-  taskId?: string // Task ID used internally, will be the same for the whole session
+  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
   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[]
-  status?: WorkflowStatus
+  events: E[] // events of the current run
+  status?: RunStatus // status of the current run
 }
 
 // extend TaskDefinition with sessionId, taskId, serviceId are required

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

@@ -0,0 +1,243 @@
+import { act, renderHook, waitFor } from '@testing-library/react'
+import { vi, describe, it, expect, beforeEach } from 'vitest'
+import { useWorkflow } from '../../hook/use-workflow'
+import { WorkflowEvent } from '../../hook/use-workflow/types'
+import { server } from '../setup'
+import { http, HttpResponse } from 'msw'
+
+interface TestEvent extends WorkflowEvent {
+  type: string
+  data?: any
+}
+
+const mockWorkflowParams = {
+  baseUrl: 'http://127.0.0.1:4501',
+  deployment: 'test-workflow',
+}
+
+describe('useWorkflow', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  describe('Basic functionality', () => {
+    it('should initialize with correct default state', () => {
+      const { result } = renderHook(() =>
+        useWorkflow<TestEvent>(mockWorkflowParams)
+      )
+
+      expect(result.current.runId).toBeUndefined()
+      expect(result.current.events).toEqual([])
+      expect(result.current.status).toBeUndefined()
+      expect(typeof result.current.start).toBe('function')
+      expect(typeof result.current.stop).toBe('function')
+      expect(typeof result.current.sendEvent).toBe('function')
+    })
+
+    it('should provide a stable API interface', () => {
+      const { result, rerender } = renderHook(() =>
+        useWorkflow<TestEvent>(mockWorkflowParams)
+      )
+
+      const firstRender = result.current
+      rerender()
+      const secondRender = result.current
+
+      // Functions should be referentially stable
+      expect(firstRender.start).toBe(secondRender.start)
+      expect(firstRender.stop).toBe(secondRender.stop)
+      expect(firstRender.sendEvent).toBe(secondRender.sendEvent)
+    })
+  })
+
+  describe('Starting workflow', () => {
+    it('should create a new task and start streaming events', async () => {
+      const { result } = renderHook(() =>
+        useWorkflow<TestEvent>(mockWorkflowParams)
+      )
+
+      await act(async () => {
+        await result.current.start({ message: 'test start' })
+      })
+
+      // Should have received all mock events
+      await waitFor(() => {
+        expect(result.current.runId).toBe('test-run-id')
+        expect(result.current.events).toHaveLength(12)
+        expect(result.current.status).toBe('complete')
+      })
+
+      // First event should be a UIEvent
+      const firstEvent = result.current.events[0]
+      expect(firstEvent.type).toBe('workflow.UIEvent')
+
+      // Last event should be a StopEvent
+      const lastEvent = result.current.events[result.current.events.length - 1]
+      expect(lastEvent.type).toBe('llama_index.core.workflow.events.StopEvent')
+    })
+
+    it('should reset events when starting a new task', async () => {
+      const { result } = renderHook(() =>
+        useWorkflow<TestEvent>(mockWorkflowParams)
+      )
+
+      // Start first task
+      await act(async () => {
+        await result.current.start({ message: 'first task' })
+      })
+      await waitFor(() => {
+        expect(result.current.events).toHaveLength(12)
+        expect(result.current.status).toBe('complete')
+      })
+
+      // Start second task - should reset events
+      await act(async () => {
+        await result.current.start({ message: 'second task' })
+      })
+      await waitFor(() => {
+        expect(result.current.events).toHaveLength(12) // New set of events
+        expect(result.current.status).toBe('complete')
+      })
+    })
+  })
+
+  describe('Existing task initialization', () => {
+    it('should initialize with existing task ID', async () => {
+      const paramsWithRunId = {
+        ...mockWorkflowParams,
+        runId: 'test-existing-run-id',
+      }
+
+      const { result } = renderHook(() =>
+        useWorkflow<TestEvent>(paramsWithRunId)
+      )
+
+      // Then check for the full expected state
+      await waitFor(() => {
+        expect(result.current.events).toHaveLength(12)
+        expect(result.current.runId).toBe('test-existing-run-id')
+        expect(result.current.status).toBe('complete')
+      })
+    })
+  })
+
+  describe('Error handling', () => {
+    it('should handle streaming errors', async () => {
+      const onError = vi.fn()
+      const paramsWithErrorHandler = {
+        ...mockWorkflowParams,
+        onError,
+      }
+
+      // Mock streaming error
+      server.use(
+        http.get(
+          'http://127.0.0.1:4501/deployments/test-workflow/tasks/test-run-id/events',
+          () => {
+            return HttpResponse.json(
+              { error: 'Streaming failed' },
+              { status: 500 }
+            )
+          }
+        )
+      )
+
+      const { result } = renderHook(() =>
+        useWorkflow<TestEvent>(paramsWithErrorHandler)
+      )
+
+      await act(async () => {
+        await result.current.start({ message: 'test start' }).catch(error => {
+          expect(error).toBeDefined()
+        })
+      })
+
+      await waitFor(() => {
+        expect(onError).toHaveBeenCalled()
+        expect(result.current.status).toBe('error')
+      })
+    })
+  })
+
+  describe('Event handling', () => {
+    it('should throw error when sending event without task', async () => {
+      const { result } = renderHook(() =>
+        useWorkflow<TestEvent>(mockWorkflowParams)
+      )
+
+      const customEvent: TestEvent = { type: 'custom.AdHocEvent' }
+
+      await expect(
+        act(async () => {
+          await result.current.sendEvent(customEvent)
+        })
+      ).rejects.toThrow('Task is not initialized')
+    })
+
+    it('should send events to existing task', async () => {
+      const { result } = renderHook(() =>
+        useWorkflow<TestEvent>(mockWorkflowParams)
+      )
+
+      // Start a task first
+      // eslint-disable-next-line @typescript-eslint/require-await -- no wait to test long running task
+      await act(async () => {
+        result.current.start({
+          message: 'test long running task',
+          delay: true, // flag to make task long running
+        })
+      })
+
+      // Send a custom event
+      await act(async () => {
+        await result.current.sendEvent({ type: 'custom.AdHocEvent1' })
+      })
+
+      // wait for 3 seconds to make sure stream is complete
+      await act(async () => {
+        await new Promise(resolve => {
+          setTimeout(resolve, 3000)
+        })
+      })
+
+      await waitFor(() => expect(result.current.events).toHaveLength(13))
+    })
+
+    it('should send stop event', async () => {
+      const { result } = renderHook(() =>
+        useWorkflow<TestEvent>(mockWorkflowParams)
+      )
+
+      // Start a task first
+      // eslint-disable-next-line @typescript-eslint/require-await -- no wait to test long running task
+      await act(async () => {
+        result.current.start({
+          message: 'test long running task',
+          delay: true, // flag to make task long running
+        })
+      })
+
+      // wait for 100ms to have some events in the stream
+      await act(async () => {
+        await new Promise(resolve => {
+          setTimeout(resolve, 100)
+        })
+      })
+
+      // Send a stop event
+      await act(async () => {
+        await result.current.stop()
+      })
+
+      // wait for 3 seconds to make sure stream is complete
+      await act(async () => {
+        await new Promise(resolve => {
+          setTimeout(resolve, 2000)
+        })
+      })
+
+      // events list should not be completed
+      await waitFor(() => expect(result.current.events.length).toBeLessThan(13))
+    })
+  })
+})

packages/chat-ui/src/test/mocks/llama-deploy-server.ts

@@ -0,0 +1,151 @@
+import { http, HttpResponse } from 'msw'
+import { RawEvent } from '../../hook/use-workflow/types'
+
+const BASE_URL = 'http://127.0.0.1:4501'
+
+// Mock task data
+const mockTask = {
+  task_id: 'test-run-id',
+  session_id: 'test-session-id',
+  service_id: 'echo_workflow',
+  input: '{"message": "Please run task"}',
+}
+
+const mockExistingTask = {
+  task_id: 'test-existing-run-id',
+  session_id: 'test-existing-session-id',
+  service_id: 'echo_workflow',
+  input: '{"message": "Please run existing task"}',
+}
+
+const mockTasks = [mockExistingTask, mockTask]
+
+const welcomeEvent = {
+  __is_pydantic: true,
+  value: { data: { request: 'Welcome to the workflow' } },
+  qualified_name: 'workflow.UIEvent',
+}
+
+// Mock event data
+const mockEvents = [
+  // Mock an UIEvent after task is created
+  {
+    __is_pydantic: true,
+    value: { data: { request: 'Please run task' } },
+    qualified_name: 'workflow.UIEvent',
+  },
+
+  // Mock some events for counter
+  ...Array.from({ length: 9 }, (_, index) => ({
+    __is_pydantic: true,
+    value: { data: { counter: index } },
+    qualified_name: 'workflow.UIEvent',
+  })),
+
+  // Mock StopEvent
+  {
+    __is_pydantic: true,
+    value: {},
+    qualified_name: 'llama_index.core.workflow.events.StopEvent',
+  },
+]
+
+// Global state to simulate events queue
+const eventsQueue: RawEvent[] = []
+let delayAfterStart = 10 // delay in ms
+
+export const mockLlamaDeployServer = [
+  // GET /deployments/{deployment_name}/tasks - Get existing tasks
+  http.get(`${BASE_URL}/deployments/:deploymentName/tasks`, () => {
+    return HttpResponse.json(mockTasks)
+  }),
+
+  // POST /deployments/{deployment_name}/tasks/create - Create new task
+  http.post(
+    `${BASE_URL}/deployments/:deploymentName/tasks/create`,
+    async ({ request }) => {
+      const { input } = (await request.json()) as { input: string }
+      const delay = JSON.parse(input).delay
+      if (delay) {
+        delayAfterStart = 100
+      }
+      return HttpResponse.json(mockTask)
+    }
+  ),
+
+  // GET /deployments/{deployment_name}/tasks/{task_id}/events - Stream events
+  http.get(
+    `${BASE_URL}/deployments/:deploymentName/tasks/:taskId/events`,
+    () => {
+      // Return a readable stream of events
+      const encoder = new TextEncoder()
+      const stream = new ReadableStream({
+        start(controller) {
+          controller.enqueue(
+            encoder.encode(`${JSON.stringify(welcomeEvent)}\n`)
+          )
+
+          let eventIndex = 0
+
+          const sendNextEvent = () => {
+            if (eventIndex >= mockEvents.length) {
+              controller.close()
+              return
+            }
+
+            const event = mockEvents[eventIndex]
+            controller.enqueue(encoder.encode(`${JSON.stringify(event)}\n`))
+            eventIndex++
+
+            if (eventsQueue.length > 0) {
+              // continuously consume events from the queue
+              const nextEvent = eventsQueue.shift()
+              controller.enqueue(
+                encoder.encode(`${JSON.stringify(nextEvent)}\n`)
+              )
+              if (
+                nextEvent?.qualified_name ===
+                'llama_index.core.workflow.events.StopEvent'
+              ) {
+                controller.close()
+                return
+              }
+            }
+
+            if (
+              event.qualified_name ===
+              'llama_index.core.workflow.events.StopEvent'
+            ) {
+              controller.close()
+              return
+            }
+
+            setTimeout(sendNextEvent, delayAfterStart)
+          }
+
+          // Start sending events
+          sendNextEvent()
+        },
+      })
+
+      return new HttpResponse(stream, {
+        headers: {
+          'Content-Type': 'application/json',
+        },
+      })
+    }
+  ),
+
+  // POST /deployments/{deployment_name}/tasks/{task_id}/events - Send event
+  http.post(
+    `${BASE_URL}/deployments/:deploymentName/tasks/:taskId/events`,
+    async ({ request }) => {
+      const event = (await request.json()) as { event_obj_str: string }
+      if (event) {
+        // simulate an event queue
+        eventsQueue.push(JSON.parse(event.event_obj_str) as RawEvent)
+      }
+      return HttpResponse.json({ data: { success: true } })
+    }
+  ),
+]

packages/chat-ui/src/test/setup.ts

@@ -0,0 +1,14 @@
+import '@testing-library/jest-dom'
+import { beforeAll, afterEach, afterAll } from 'vitest'
+import { setupServer } from 'msw/node'
+import { mockLlamaDeployServer } from './mocks/llama-deploy-server'
+
+// Setup MSW server
+export const server = setupServer(...mockLlamaDeployServer)
+
+beforeAll(() => server.listen())
+afterEach(() => server.resetHandlers())
+afterAll(() => server.close())
+
+// Mock fetch globally
+global.fetch = global.fetch || (() => Promise.resolve({} as Response))

packages/chat-ui/vitest.config.ts

@@ -0,0 +1,16 @@
+import { defineConfig } from 'vitest/config'
+import react from '@vitejs/plugin-react'
+
+export default defineConfig({
+  plugins: [react()],
+  test: {
+    environment: 'jsdom',
+    setupFiles: ['./src/test/setup.ts'],
+    globals: true,
+  },
+  resolve: {
+    alias: {
+      '@': './src',
+    },
+  },
+})