413 lines
14 KiB
Python
413 lines
14 KiB
Python
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, 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:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
class DispatchGenerateRequest(BaseModel):
|
|
device_id: Optional[str] = None
|
|
request: GenerateRequest
|
|
|
|
|
|
class DispatchResponse(BaseModel):
|
|
dispatch_id: str
|
|
device_id: str
|
|
status: str
|
|
created_at: str
|
|
|
|
|
|
class DeviceCommandRequest(BaseModel):
|
|
command: str
|
|
dispatch_id: Optional[str] = None
|
|
branch: Optional[str] = None
|
|
shell_command: Optional[str] = None
|
|
payload: Optional[dict[str, Any]] = None
|
|
|
|
|
|
class DeviceCommandResponse(BaseModel):
|
|
dispatch_id: str
|
|
device_id: str
|
|
command: str
|
|
status: str
|
|
created_at: str
|
|
|
|
|
|
@dataclass
|
|
class EdgeConnection:
|
|
device_id: str
|
|
websocket: WebSocket
|
|
connected_at: str
|
|
last_seen: str
|
|
busy: bool = False
|
|
|
|
|
|
class EdgeDispatchManager:
|
|
def __init__(self) -> None:
|
|
self.connections: dict[str, EdgeConnection] = {}
|
|
self.dispatches: dict[str, dict[str, Any]] = {}
|
|
self.commands: dict[str, dict[str, Any]] = {}
|
|
self.lock = asyncio.Lock()
|
|
|
|
async def register(self, device_id: str, websocket: WebSocket) -> EdgeConnection:
|
|
async with self.lock:
|
|
conn = EdgeConnection(
|
|
device_id=device_id,
|
|
websocket=websocket,
|
|
connected_at=utc_now_iso(),
|
|
last_seen=utc_now_iso(),
|
|
busy=False,
|
|
)
|
|
self.connections[device_id] = conn
|
|
return conn
|
|
|
|
async def unregister(self, device_id: str) -> None:
|
|
async with self.lock:
|
|
self.connections.pop(device_id, None)
|
|
|
|
async def list_devices(self) -> list[dict[str, Any]]:
|
|
async with self.lock:
|
|
return [
|
|
{
|
|
"device_id": conn.device_id,
|
|
"connected_at": conn.connected_at,
|
|
"last_seen": conn.last_seen,
|
|
"busy": conn.busy,
|
|
}
|
|
for conn in self.connections.values()
|
|
]
|
|
|
|
async def select_device(self, preferred: Optional[str]) -> EdgeConnection:
|
|
async with self.lock:
|
|
if preferred:
|
|
conn = self.connections.get(preferred)
|
|
if conn is None:
|
|
raise HTTPException(status_code=404, detail=f"device not found: {preferred}")
|
|
if conn.busy:
|
|
raise HTTPException(status_code=409, detail=f"device is busy: {preferred}")
|
|
return conn
|
|
|
|
for conn in self.connections.values():
|
|
if not conn.busy:
|
|
return conn
|
|
|
|
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 = {
|
|
"dispatch_id": dispatch_id,
|
|
"device_id": conn.device_id,
|
|
"status": "DISPATCHED",
|
|
"request": req.model_dump(),
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
"result": None,
|
|
"error": None,
|
|
"artifact_urls": {},
|
|
}
|
|
|
|
async with self.lock:
|
|
self.dispatches[dispatch_id] = record
|
|
target = self.connections.get(conn.device_id)
|
|
if target is None:
|
|
raise HTTPException(status_code=404, detail=f"device disconnected: {conn.device_id}")
|
|
target.busy = True
|
|
target.last_seen = now
|
|
|
|
payload = {
|
|
"event": "generate",
|
|
"dispatch_id": dispatch_id,
|
|
"request": req.model_dump(),
|
|
}
|
|
try:
|
|
await conn.websocket.send_json(payload)
|
|
except Exception as exc:
|
|
async with self.lock:
|
|
record["status"] = "FAILED"
|
|
record["error"] = f"dispatch send failed: {exc}"
|
|
record["updated_at"] = utc_now_iso()
|
|
target = self.connections.get(conn.device_id)
|
|
if target:
|
|
target.busy = False
|
|
raise
|
|
return record
|
|
|
|
async def create_command(self, conn: EdgeConnection, body: DeviceCommandRequest) -> dict[str, Any]:
|
|
await self._prune_if_needed()
|
|
command = body.command.strip().lower()
|
|
allowed = {"update_code", "restart_service", "ping"}
|
|
if command not in allowed:
|
|
raise HTTPException(status_code=400, detail=f"unsupported command: {command}")
|
|
|
|
dispatch_id = body.dispatch_id or uuid4().hex
|
|
now = utc_now_iso()
|
|
record = {
|
|
"dispatch_id": dispatch_id,
|
|
"device_id": conn.device_id,
|
|
"command": command,
|
|
"status": "DISPATCHED",
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
"request": body.model_dump(),
|
|
"result": None,
|
|
"error": None,
|
|
}
|
|
|
|
payload: dict[str, Any] = {
|
|
"event": command,
|
|
"dispatch_id": dispatch_id,
|
|
}
|
|
if body.branch:
|
|
payload["branch"] = body.branch
|
|
if body.shell_command:
|
|
payload["command"] = body.shell_command
|
|
if body.payload:
|
|
payload.update(body.payload)
|
|
|
|
async with self.lock:
|
|
self.commands[dispatch_id] = record
|
|
target = self.connections.get(conn.device_id)
|
|
if target is None:
|
|
raise HTTPException(status_code=404, detail=f"device disconnected: {conn.device_id}")
|
|
target.last_seen = now
|
|
|
|
try:
|
|
await conn.websocket.send_json(payload)
|
|
except Exception as exc:
|
|
async with self.lock:
|
|
record["status"] = "FAILED"
|
|
record["error"] = f"command send failed: {exc}"
|
|
record["updated_at"] = utc_now_iso()
|
|
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) + len(self.commands)
|
|
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.extend([v for v in self.commands.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)
|
|
self.commands.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:
|
|
conn = self.connections.get(device_id)
|
|
if conn:
|
|
conn.last_seen = now
|
|
|
|
dispatch_id = event.get("dispatch_id")
|
|
if not dispatch_id:
|
|
return
|
|
|
|
record = self.dispatches.get(dispatch_id)
|
|
if record is None:
|
|
record = self.commands.get(dispatch_id)
|
|
if record is None:
|
|
return
|
|
|
|
evt = event.get("event")
|
|
if evt == "accepted":
|
|
record["status"] = "RUNNING"
|
|
elif evt == "status":
|
|
status_val = event.get("status")
|
|
if status_val:
|
|
record["status"] = status_val
|
|
elif evt == "result":
|
|
status_val = event.get("status") or "SUCCEEDED"
|
|
record["status"] = status_val
|
|
record["result"] = event
|
|
if conn:
|
|
conn.busy = False
|
|
elif evt == "error":
|
|
record["status"] = "FAILED"
|
|
record["error"] = event.get("error")
|
|
if conn:
|
|
conn.busy = False
|
|
elif evt == "command_status":
|
|
status_val = event.get("status") or "RUNNING"
|
|
record["status"] = status_val
|
|
elif evt == "command_result":
|
|
status_val = event.get("status") or "SUCCEEDED"
|
|
record["status"] = status_val
|
|
record["result"] = event
|
|
if status_val == "FAILED":
|
|
record["error"] = event.get("error") or event.get("stderr")
|
|
elif evt == "pong":
|
|
record["status"] = "SUCCEEDED"
|
|
record["result"] = event
|
|
|
|
record["updated_at"] = now
|
|
|
|
async def mark_disconnect_failed(self, device_id: str) -> None:
|
|
now = utc_now_iso()
|
|
async with self.lock:
|
|
for record in self.dispatches.values():
|
|
if record["device_id"] == device_id and record["status"] in {"DISPATCHED", "RUNNING", "PENDING"}:
|
|
record["status"] = "FAILED"
|
|
record["error"] = f"device disconnected: {device_id}"
|
|
record["updated_at"] = now
|
|
for record in self.commands.values():
|
|
if record["device_id"] == device_id and record["status"] in {"DISPATCHED", "RUNNING", "PENDING"}:
|
|
record["status"] = "FAILED"
|
|
record["error"] = f"device disconnected: {device_id}"
|
|
record["updated_at"] = now
|
|
|
|
|
|
manager = EdgeDispatchManager()
|
|
app = FastAPI(title="Edge Dispatch Service", version="0.1.0")
|
|
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict[str, Any]:
|
|
devices = await manager.list_devices()
|
|
return {
|
|
"service": "edge_dispatch",
|
|
"status": "ok",
|
|
"connected_devices": len(devices),
|
|
}
|
|
|
|
|
|
@app.get("/devices")
|
|
async def list_devices() -> dict[str, Any]:
|
|
return {"devices": await manager.list_devices()}
|
|
|
|
|
|
@app.post("/dispatch/generate", response_model=DispatchResponse)
|
|
async def dispatch_generate(body: DispatchGenerateRequest) -> DispatchResponse:
|
|
conn = await manager.select_device(body.device_id)
|
|
try:
|
|
record = await manager.create_dispatch(conn, body.request)
|
|
except WebSocketDisconnect as exc:
|
|
await manager.unregister(conn.device_id)
|
|
raise HTTPException(status_code=503, detail=f"device disconnected during dispatch: {conn.device_id}") from exc
|
|
except RuntimeError as exc:
|
|
raise HTTPException(status_code=503, detail=str(exc)) from exc
|
|
|
|
return DispatchResponse(
|
|
dispatch_id=record["dispatch_id"],
|
|
device_id=record["device_id"],
|
|
status=record["status"],
|
|
created_at=record["created_at"],
|
|
)
|
|
|
|
|
|
@app.get("/dispatch/{dispatch_id}")
|
|
async def get_dispatch(dispatch_id: str) -> 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}")
|
|
return record
|
|
|
|
|
|
@app.post("/devices/{device_id}/command", response_model=DeviceCommandResponse)
|
|
async def device_command(device_id: str, body: DeviceCommandRequest) -> DeviceCommandResponse:
|
|
conn = await manager.select_device(device_id)
|
|
try:
|
|
record = await manager.create_command(conn, body)
|
|
except WebSocketDisconnect as exc:
|
|
await manager.unregister(conn.device_id)
|
|
raise HTTPException(status_code=503, detail=f"device disconnected during command: {conn.device_id}") from exc
|
|
except RuntimeError as exc:
|
|
raise HTTPException(status_code=503, detail=str(exc)) from exc
|
|
|
|
return DeviceCommandResponse(
|
|
dispatch_id=record["dispatch_id"],
|
|
device_id=record["device_id"],
|
|
command=record["command"],
|
|
status=record["status"],
|
|
created_at=record["created_at"],
|
|
)
|
|
|
|
|
|
@app.get("/commands/{dispatch_id}")
|
|
async def get_command(dispatch_id: str) -> dict[str, Any]:
|
|
record = manager.commands.get(dispatch_id)
|
|
if record is None:
|
|
raise HTTPException(status_code=404, detail=f"command not found: {dispatch_id}")
|
|
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()
|
|
await manager.register(device_id, websocket)
|
|
|
|
try:
|
|
await websocket.send_json({"event": "registered", "device_id": device_id, "ts": utc_now_iso()})
|
|
while True:
|
|
msg = await websocket.receive_json()
|
|
await manager.mark_event(device_id, msg)
|
|
except WebSocketDisconnect:
|
|
await manager.mark_disconnect_failed(device_id)
|
|
await manager.unregister(device_id)
|