Password promt

backend/src/app.py

@@ -63,16 +63,30 @@ async def websocket_endpoint(websocket: WebSocket, task_id: str):
     task_manager = get_task_manager()
     queue = await task_manager.subscribe_logs(task_id)
     try:
-        # Send initial logs if any
+        # Stream new logs
+        logger.info(f"Starting log stream for task {task_id}")
+        
+        # Send initial logs first to build context
         initial_logs = task_manager.get_task_logs(task_id)
         for log_entry in initial_logs:
-            # Convert datetime to string for JSON serialization
             log_dict = log_entry.dict()
             log_dict['timestamp'] = log_dict['timestamp'].isoformat()
             await websocket.send_json(log_dict)
 
-        # Stream new logs
-        logger.info(f"Starting log stream for task {task_id}")
+        # Force a check for AWAITING_INPUT status immediately upon connection
+        # This ensures that if the task is already waiting when the user connects, they get the prompt.
+        task = task_manager.get_task(task_id)
+        if task and task.status == "AWAITING_INPUT" and task.input_request:
+             # Construct a synthetic log entry to trigger the frontend handler
+             # This is a bit of a hack but avoids changing the websocket protocol significantly
+             synthetic_log = {
+                 "timestamp": task.logs[-1].timestamp.isoformat() if task.logs else "2024-01-01T00:00:00",
+                 "level": "INFO",
+                 "message": "Task paused for user input (Connection Re-established)",
+                 "context": {"input_request": task.input_request}
+             }
+             await websocket.send_json(synthetic_log)
+
         while True:
             log_entry = await queue.get()
             log_dict = log_entry.dict()
@@ -84,7 +98,9 @@ async def websocket_endpoint(websocket: WebSocket, task_id: str):
             if "Task completed successfully" in log_entry.message or "Task failed" in log_entry.message:
                 # Wait a bit to ensure client receives the last message
                 await asyncio.sleep(2)
-                break
+                # DO NOT BREAK here - allow client to keep connection open if they want to review logs
+                # or until they disconnect. Breaking closes the socket immediately.
+                # break
 
     except WebSocketDisconnect:
         logger.info(f"WebSocket connection disconnected for task {task_id}")