fix: 新增数据模块

video_worker/app/edge_dispatch_service.py

@@ -1,13 +1,16 @@
 import asyncio
 from dataclasses import dataclass
 from datetime import datetime, timezone
+from pathlib import Path
 from typing import Any, Optional
 from uuid import uuid4
 
-from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, File, Form, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
 from pydantic import BaseModel
 
+from app.oss_client import oss_uploader
 from app.schemas import GenerateRequest
+from app.settings import settings
 
 
 def utc_now_iso() -> str:
@@ -86,6 +89,7 @@ class EdgeDispatchManager:
         raise HTTPException(status_code=409, detail="no idle edge device available")
 
     async def create_dispatch(self, conn: EdgeConnection, req: GenerateRequest) -> dict[str, Any]:
+        await self._prune_if_needed()
         dispatch_id = uuid4().hex
         now = utc_now_iso()
         record = {
@@ -97,6 +101,7 @@ class EdgeDispatchManager:
             "updated_at": now,
             "result": None,
             "error": None,
+            "artifact_urls": {},
         }
 
         async with self.lock:
@@ -125,6 +130,18 @@ class EdgeDispatchManager:
             raise
         return record
 
+    async def _prune_if_needed(self) -> None:
+        max_records = max(100, int(settings.edge_max_dispatch_records))
+        async with self.lock:
+            total = len(self.dispatches)
+            if total < max_records:
+                return
+            over = total - max_records + 1
+            done = [v for v in self.dispatches.values() if v.get("status") in {"SUCCEEDED", "FAILED"}]
+            done.sort(key=lambda x: x.get("updated_at", ""))
+            for rec in done[:over]:
+                self.dispatches.pop(rec["dispatch_id"], None)
+
     async def mark_event(self, device_id: str, event: dict[str, Any]) -> None:
         now = utc_now_iso()
         async with self.lock:
@@ -217,6 +234,54 @@ async def get_dispatch(dispatch_id: str) -> dict[str, Any]:
     return record
 
 
+@app.post("/dispatch/{dispatch_id}/artifacts")
+async def upload_artifacts(
+    dispatch_id: str,
+    task_id: str = Form(default=""),
+    status: str = Form(default="SUCCEEDED"),
+    files: list[UploadFile] = File(default_factory=list),
+) -> dict[str, Any]:
+    record = manager.dispatches.get(dispatch_id)
+    if record is None:
+        raise HTTPException(status_code=404, detail=f"dispatch not found: {dispatch_id}")
+    if not oss_uploader.enabled:
+        raise HTTPException(status_code=400, detail="OSS upload is disabled, set OSS_ENABLED=true")
+    if not files:
+        raise HTTPException(status_code=400, detail="no files uploaded")
+
+    uploaded: dict[str, dict[str, str]] = {}
+    for file in files:
+        name = file.filename or "artifact.bin"
+        try:
+            result = await asyncio.to_thread(oss_uploader.upload_fileobj, dispatch_id, name, file.file)
+            uploaded[Path(name).name] = result
+        finally:
+            await file.close()
+
+    now = utc_now_iso()
+    async with manager.lock:
+        record["artifact_urls"] = uploaded
+        record["result"] = {
+            "event": "result",
+            "dispatch_id": dispatch_id,
+            "task_id": task_id or None,
+            "status": status,
+            "artifact_urls": uploaded,
+        }
+        record["status"] = status
+        record["updated_at"] = now
+        conn = manager.connections.get(record["device_id"])
+        if conn:
+            conn.busy = False
+
+    return {
+        "dispatch_id": dispatch_id,
+        "status": status,
+        "artifact_urls": uploaded,
+        "updated_at": now,
+    }
+
+
 @app.websocket("/ws/edge/{device_id}")
 async def edge_socket(websocket: WebSocket, device_id: str) -> None:
     await websocket.accept()