From f47989c776718d18a1a9f50d72363d244909737c Mon Sep 17 00:00:00 2001
From: mohiit1502 <mohiit1502@gmail.com>
Date: Tue, 21 Apr 2026 00:21:55 +0530
Subject: [PATCH] fix(OBS-02): task badge immediately reflects final status
 from WS event
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Root cause: taskStatus was only updated after tasksApi.get() resolved
(100-300ms after WS fires), so badge stayed 'running' during that gap.
Secondary bug: used task.status from the refetch which can be stale.

Fix:
- setPendingTaskIds and taskStatus flip happen synchronously on WS event
  receipt, before the API enrichment call — badge changes instantly
- tasksApi.get() still called to get content + tokenCost, but now
  uses finalStatus from WS type (not potentially-stale task.status)
- .catch() path removed redundant status override (already set above)
---
 webapp/src/components/layout/ChatWorkspace.tsx | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/webapp/src/components/layout/ChatWorkspace.tsx b/webapp/src/components/layout/ChatWorkspace.tsx
index dac1aa87..0a2c0a37 100644
--- a/webapp/src/components/layout/ChatWorkspace.tsx
+++ b/webapp/src/components/layout/ChatWorkspace.tsx
@@ -1916,12 +1916,19 @@ export function ChatWorkspace() {
     const taskId = msg.entity_id
     if (!taskId || !pendingTaskIds.has(taskId)) return
     if (msg.type === 'task.completed' || msg.type === 'task.failed') {
+      const finalStatus = msg.type === 'task.completed' ? 'completed' : 'failed'
+      // Immediately flip status from WS event — stops "Running…" badge without waiting for API
+      setPendingTaskIds(prev => { const n = new Set(prev); n.delete(taskId); return n })
+      setMessages(prev => prev.map(m =>
+        m.taskId === taskId ? { ...m, taskStatus: finalStatus } : m
+      ))
+      // Enrich with full output + token cost once the API call resolves
       tasksApi.get(taskId).then(task => {
-        setPendingTaskIds(prev => { const n = new Set(prev); n.delete(taskId); return n })
         setMessages(prev => prev.map(m =>
           m.taskId === taskId ? {
-            ...m, content: extractOutput(task),
-            taskStatus: task.status,
+            ...m,
+            content: extractOutput(task),
+            taskStatus: finalStatus,          // use WS type, not potentially-stale task.status
             tokenCost: task.token_usage?.cost_usd ?? undefined,
           } : m
         ))
@@ -1929,7 +1936,7 @@ export function ChatWorkspace() {
         qc.invalidateQueries({ queryKey: ['tasks'] })
       }).catch(() => {
         setMessages(prev => prev.map(m =>
-          m.taskId === taskId ? { ...m, taskStatus: msg.type === 'task.failed' ? 'failed' : 'completed', content: 'Could not fetch result.' } : m
+          m.taskId === taskId ? { ...m, content: 'Could not fetch result.' } : m
         ))
       })
     }