fix:优化服务

video_worker/app/api.py

@@ -1,4 +1,7 @@
-from fastapi import APIRouter, HTTPException
+import asyncio
+from typing import Any
+
+from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect
 
 from app.schemas import GenerateRequest, HealthResponse, TaskResultResponse, TaskStatusResponse
 
@@ -46,3 +49,88 @@ def health_check():
         ltx_loaded=ltx_backend.is_loaded(),
         hunyuan_loaded=hunyuan_backend.is_loaded(),
     )
+
+
+@router.websocket("/ws/generate")
+async def ws_generate(websocket: WebSocket):
+    await websocket.accept()
+    task_manager = router.task_manager
+
+    try:
+        message: dict[str, Any] = await websocket.receive_json()
+        action = message.get("action")
+
+        if action == "watch":
+            task_id = message.get("task_id")
+            if not task_id:
+                await websocket.send_json({"event": "error", "error": "task_id is required for watch"})
+                await websocket.close(code=1008)
+                return
+        else:
+            payload = message.get("payload", message)
+            try:
+                req = GenerateRequest.model_validate(payload)
+            except Exception as exc:
+                await websocket.send_json({"event": "error", "error": f"invalid request: {exc}"})
+                await websocket.close(code=1008)
+                return
+
+            created = await task_manager.create_task(req)
+            task_id = created.task_id
+            await websocket.send_json(
+                {
+                    "event": "accepted",
+                    "task_id": task_id,
+                    "status": created.status,
+                    "backend": created.backend,
+                    "model_name": created.model_name,
+                    "progress": created.progress,
+                    "created_at": created.created_at.isoformat(),
+                    "updated_at": created.updated_at.isoformat(),
+                }
+            )
+
+        last_status = None
+        while True:
+            status = task_manager.get_status(task_id)
+            if status.status != last_status:
+                await websocket.send_json(
+                    {
+                        "event": "status",
+                        "task_id": status.task_id,
+                        "status": status.status,
+                        "backend": status.backend,
+                        "model_name": status.model_name,
+                        "progress": status.progress,
+                        "created_at": status.created_at.isoformat(),
+                        "updated_at": status.updated_at.isoformat(),
+                    }
+                )
+                last_status = status.status
+
+            if status.status in {"SUCCEEDED", "FAILED"}:
+                result = task_manager.get_result(task_id)
+                await websocket.send_json(
+                    {
+                        "event": "result",
+                        "task_id": result.task_id,
+                        "status": result.status,
+                        "video_path": result.video_path,
+                        "first_frame_path": result.first_frame_path,
+                        "metadata_path": result.metadata_path,
+                        "log_path": result.log_path,
+                        "error": result.error,
+                    }
+                )
+                await websocket.close(code=1000)
+                break
+
+            await asyncio.sleep(1)
+    except WebSocketDisconnect:
+        return
+    except KeyError as exc:
+        await websocket.send_json({"event": "error", "error": str(exc)})
+        await websocket.close(code=1008)
+    except Exception as exc:
+        await websocket.send_json({"event": "error", "error": f"unexpected error: {exc}"})
+        await websocket.close(code=1011)