From 9fad42ad762687e9879c7ff5d1ba25cf4e752934 Mon Sep 17 00:00:00 2001
From: EAniwa <30927054+EAniwa@users.noreply.github.com>
Date: Wed, 19 Nov 2025 20:16:11 -0500
Subject: [PATCH] fix: Resolve real-time agent progress update issues

- Modify StreamHandler to send current agent statuses immediately upon WebSocket connection
- Switch AnalysisService to use asynchronous graph streaming (astream) to prevent event loop blocking
- Clean up debug logging in frontend and backend
---
 backend/api/services/analysis_service.py |  2 +-
 backend/api/websocket/stream_handler.py  | 10 ++++++++++
 frontend/app/analysis/[id]/page.tsx      | 12 +++++-------
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/backend/api/services/analysis_service.py b/backend/api/services/analysis_service.py
index 1d65610c..94bfeb40 100644
--- a/backend/api/services/analysis_service.py
+++ b/backend/api/services/analysis_service.py
@@ -143,7 +143,7 @@ class AnalysisService:
             
             # Stream the analysis
             trace = []
-            for chunk in graph.graph.stream(init_agent_state, **args):
+            async for chunk in graph.graph.astream(init_agent_state, **args):
                 if update_callback:
                     await self._process_chunk(chunk, update_callback, analysis_id, request.analysts)
                 trace.append(chunk)
diff --git a/backend/api/websocket/stream_handler.py b/backend/api/websocket/stream_handler.py
index e0346db1..9ff503e9 100644
--- a/backend/api/websocket/stream_handler.py
+++ b/backend/api/websocket/stream_handler.py
@@ -1,4 +1,5 @@
 import json
+from datetime import datetime
 from typing import Dict, Set
 from fastapi import WebSocket, WebSocketDisconnect
 from ..models.schemas import StreamUpdate
@@ -64,6 +65,15 @@ class StreamHandler:
                 "timestamp": ""
             })
             
+            # Send current agent statuses
+            current_statuses = self.analysis_service.agent_statuses.get(analysis_id, {})
+            for agent, status in current_statuses.items():
+                await websocket.send_json({
+                    "type": "agent_status",
+                    "data": {"agent": agent, "status": status},
+                    "timestamp": datetime.now().isoformat()
+                })
+            
             # Keep connection alive and forward updates
             while True:
                 # Wait for any incoming messages (ping/pong or close)
diff --git a/frontend/app/analysis/[id]/page.tsx b/frontend/app/analysis/[id]/page.tsx
index 209c182f..34085d3a 100644
--- a/frontend/app/analysis/[id]/page.tsx
+++ b/frontend/app/analysis/[id]/page.tsx
@@ -94,16 +94,14 @@ export default function AnalysisPage() {
             Analysis: {statusData?.ticker || "Loading..."} - {statusData?.analysis_date || ""}
           </h1>
           <div className="flex items-center space-x-4">
-            <span className={`px-3 py-1 rounded-full text-sm ${
-              status === "completed" ? "bg-green-100 text-green-800" :
+            <span className={`px-3 py-1 rounded-full text-sm ${status === "completed" ? "bg-green-100 text-green-800" :
               status === "running" ? "bg-blue-100 text-blue-800" :
-              "bg-gray-100 text-gray-800"
-            }`}>
+                "bg-gray-100 text-gray-800"
+              }`}>
               {status}
             </span>
-            <span className={`px-3 py-1 rounded-full text-sm ${
-              isConnected ? "bg-green-100 text-green-800" : "bg-red-100 text-red-800"
-            }`}>
+            <span className={`px-3 py-1 rounded-full text-sm ${isConnected ? "bg-green-100 text-green-800" : "bg-red-100 text-red-800"
+              }`}>
               {isConnected ? "Connected" : "Disconnected"}
             </span>
           </div>