Fix persistence

Author: Sg312

apps/sim/hooks/use-task-events.test.ts

@@ -0,0 +1,62 @@
+/**
+ * @vitest-environment node
+ */
+
+import type { QueryClient } from '@tanstack/react-query'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { taskKeys } from '@/hooks/queries/tasks'
+import { handleTaskStatusEvent } from '@/hooks/use-task-events'
+
+describe('handleTaskStatusEvent', () => {
+  const queryClient = {
+    invalidateQueries: vi.fn().mockResolvedValue(undefined),
+  } satisfies Pick<QueryClient, 'invalidateQueries'>
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('invalidates the task list and completed chat detail', () => {
+    handleTaskStatusEvent(
+      queryClient,
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'completed',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(2)
+    expect(queryClient.invalidateQueries).toHaveBeenNthCalledWith(1, {
+      queryKey: taskKeys.lists(),
+    })
+    expect(queryClient.invalidateQueries).toHaveBeenNthCalledWith(2, {
+      queryKey: taskKeys.detail('chat-1'),
+    })
+  })
+
+  it('keeps list invalidation only for non-completed task events', () => {
+    handleTaskStatusEvent(
+      queryClient,
+      JSON.stringify({
+        chatId: 'chat-1',
+        type: 'started',
+        timestamp: Date.now(),
+      })
+    )
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(1)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.lists(),
+    })
+  })
+
+  it('preserves list invalidation when task event payload is invalid', () => {
+    handleTaskStatusEvent(queryClient, '{')
+
+    expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(1)
+    expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
+      queryKey: taskKeys.lists(),
+    })
+  })
+})

apps/sim/hooks/use-task-events.ts

@@ -1,12 +1,60 @@
 import { useEffect } from 'react'
 import { createLogger } from '@sim/logger'
+import type { QueryClient } from '@tanstack/react-query'
 import { useQueryClient } from '@tanstack/react-query'
 import { taskKeys } from '@/hooks/queries/tasks'
 
 const logger = createLogger('TaskEvents')
 
+interface TaskStatusEventPayload {
+  chatId?: string
+  type?: 'started' | 'completed' | 'created' | 'deleted' | 'renamed'
+}
+
+function parseTaskStatusEventPayload(data: unknown): TaskStatusEventPayload | null {
+  let parsed = data
+
+  if (typeof parsed === 'string') {
+    try {
+      parsed = JSON.parse(parsed)
+    } catch {
+      return null
+    }
+  }
+
+  if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) {
+    return null
+  }
+
+  const record = parsed as Record<string, unknown>
+
+  return {
+    ...(typeof record.chatId === 'string' ? { chatId: record.chatId } : {}),
+    ...(typeof record.type === 'string'
+      ? { type: record.type as TaskStatusEventPayload['type'] }
+      : {}),
+  }
+}
+
+export function handleTaskStatusEvent(
+  queryClient: Pick<QueryClient, 'invalidateQueries'>,
+  data: unknown
+): void {
+  queryClient.invalidateQueries({ queryKey: taskKeys.lists() })
+
+  const payload = parseTaskStatusEventPayload(data)
+  if (!payload) {
+    logger.warn('Received invalid task_status payload')
+    return
+  }
+
+  if (payload.type === 'completed' && payload.chatId) {
+    queryClient.invalidateQueries({ queryKey: taskKeys.detail(payload.chatId) })
+  }
+}
+
 /**
- * Subscribes to task status SSE events and invalidates the task list on changes.
+ * Subscribes to task status SSE events and invalidates task caches on changes.
  */
 export function useTaskEvents(workspaceId: string | undefined) {
   const queryClient = useQueryClient()
@@ -18,8 +66,8 @@ export function useTaskEvents(workspaceId: string | undefined) {
       `/api/mothership/events?workspaceId=${encodeURIComponent(workspaceId)}`
     )
 
-    eventSource.addEventListener('task_status', () => {
-      queryClient.invalidateQueries({ queryKey: taskKeys.lists() })
+    eventSource.addEventListener('task_status', (event) => {
+      handleTaskStatusEvent(queryClient, event instanceof MessageEvent ? event.data : undefined)
     })
 
     eventSource.onerror = () => {