Update app.py

app.py

@@ -12,6 +12,7 @@ import asyncio
 import os
 import uuid
 import toml
+import time
 from datetime import datetime
 from json import dumps
 
@@ -121,24 +122,29 @@ async def task_events(task_id: str):
         queue = task_manager.queues[task_id]
         task = task_manager.tasks.get(task_id)
         if task:
-            logger.info(f"Sending initial task status for: {task_id}")
             yield f"event: status\ndata: {dumps({'type': 'status', 'status': task.status, 'steps': task.steps})}\n\n"
 
+        last_event_time = time.time()
+        
         while True:
             try:
-                event = await queue.get()
-                formatted_event = dumps(event)
-                yield ": heartbeat\n\n"
-                yield f"event: {event['type']}\ndata: {formatted_event}\n\n"
-
-                if event["type"] in ["complete", "error"]:
+                # wait up to 5 seconds for new events
+                try:
+                    event = await asyncio.wait_for(queue.get(), timeout=5.0)
+                    formatted_event = dumps(event)
+                    yield f"event: {event['type']}\ndata: {formatted_event}\n\n"
+                    last_event_time = time.time()
+
+                    if event["type"] in ["complete", "error"]:
                     break
+                except asyncio.TimeoutError:
+                    # Send heartbeat to keep connection alive
+                    yield ": heartbeat\n\n"
+                
 
             except asyncio.CancelledError:
-                logger.warning("Client disconnected")
                 break
             except Exception as e:
-                logger.error(f"Event stream error: {e}")
                 yield f"event: error\ndata: {dumps({'message': str(e)})}\n\n"
                 break