710 lines
28 KiB
Python
710 lines
28 KiB
Python
"""Windows Node MVP 的注册、管理与长连接端点。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import os
|
||
from datetime import datetime
|
||
from hashlib import sha256
|
||
from pathlib import Path
|
||
from uuid import UUID
|
||
|
||
from fastapi import (
|
||
Depends,
|
||
Header,
|
||
HTTPException,
|
||
Request,
|
||
WebSocket,
|
||
WebSocketDisconnect,
|
||
status,
|
||
)
|
||
from fastapi.responses import FileResponse
|
||
|
||
from core.artifact_lifecycle import register_published_artifacts
|
||
from core.paths import from_db_path
|
||
from core.software_contracts import (
|
||
DEFAULT_CAPABILITIES,
|
||
SoftwareContractError,
|
||
get_contract,
|
||
)
|
||
from core.software_jobs import (
|
||
MAX_OUTPUT_ARTIFACT_BYTES,
|
||
MAX_OUTPUT_TOTAL_BYTES,
|
||
SoftwareJobError,
|
||
abandon_offer,
|
||
create_job,
|
||
get_job,
|
||
get_job_input,
|
||
get_job_output_context,
|
||
list_jobs,
|
||
mark_node_jobs_disconnected,
|
||
offer_next_job,
|
||
pending_node_cancellations,
|
||
record_job_terminal,
|
||
replay_succeeded_outputs,
|
||
request_job_cancel,
|
||
respond_to_offer,
|
||
software_job_output_path,
|
||
succeeded_output_upload_matches,
|
||
update_job_state,
|
||
validate_output_manifest,
|
||
)
|
||
from core.software_nodes import (
|
||
SoftwareNodeError,
|
||
authenticate_node,
|
||
create_enrollment,
|
||
delete_node,
|
||
enroll_node,
|
||
list_nodes,
|
||
mark_node_offline,
|
||
set_node_disabled,
|
||
update_node_runtime,
|
||
)
|
||
from web.schemas import (
|
||
SoftwareEnrollmentCreateRequest,
|
||
SoftwareJobCreateRequest,
|
||
SoftwareNodeDisableRequest,
|
||
SoftwareNodeEnrollRequest,
|
||
)
|
||
from web.userfiles import load_user_root, safe_join
|
||
|
||
|
||
class NodeConnectionManager:
|
||
def __init__(self) -> None:
|
||
self._connections: dict[UUID, WebSocket] = {}
|
||
self._send_locks: dict[UUID, asyncio.Lock] = {}
|
||
self._lock = asyncio.Lock()
|
||
|
||
async def activate(self, node_id: UUID, websocket: WebSocket) -> None:
|
||
async with self._lock:
|
||
old = self._connections.get(node_id)
|
||
self._connections[node_id] = websocket
|
||
self._send_locks.setdefault(node_id, asyncio.Lock())
|
||
if old is not None and old is not websocket:
|
||
await old.close(code=4001, reason="replaced by a newer connection")
|
||
|
||
async def remove(self, node_id: UUID, websocket: WebSocket) -> bool:
|
||
async with self._lock:
|
||
if self._connections.get(node_id) is websocket:
|
||
self._connections.pop(node_id, None)
|
||
self._send_locks.pop(node_id, None)
|
||
return True
|
||
return False
|
||
|
||
async def close(self, node_id: UUID) -> None:
|
||
async with self._lock:
|
||
websocket = self._connections.pop(node_id, None)
|
||
self._send_locks.pop(node_id, None)
|
||
if websocket is not None:
|
||
await websocket.close(code=4003, reason="node disabled")
|
||
|
||
async def node_ids(self) -> set[UUID]:
|
||
async with self._lock:
|
||
return set(self._connections)
|
||
|
||
async def send(self, node_id: UUID, message: dict) -> bool:
|
||
async with self._lock:
|
||
websocket = self._connections.get(node_id)
|
||
send_lock = self._send_locks.get(node_id)
|
||
if websocket is None or send_lock is None:
|
||
return False
|
||
async with send_lock:
|
||
await websocket.send_json(message)
|
||
return True
|
||
|
||
async def send_on(self, node_id: UUID, websocket: WebSocket, message: dict) -> bool:
|
||
async with self._lock:
|
||
current = self._connections.get(node_id)
|
||
send_lock = self._send_locks.get(node_id)
|
||
if current is not websocket or send_lock is None:
|
||
return False
|
||
async with send_lock:
|
||
await websocket.send_json(message)
|
||
return True
|
||
|
||
|
||
node_connections = NodeConnectionManager()
|
||
|
||
|
||
def _bearer(authorization: str | None) -> str:
|
||
scheme, _, token = (authorization or "").partition(" ")
|
||
if scheme.lower() != "bearer" or not token:
|
||
raise SoftwareNodeError("missing node bearer token")
|
||
return token
|
||
|
||
|
||
def _authenticate_output_request(
|
||
job_id: UUID,
|
||
authorization: str | None,
|
||
x_node_id: str,
|
||
x_lease_id: str,
|
||
x_request_digest: str,
|
||
) -> tuple[UUID, UUID, dict]:
|
||
try:
|
||
node_id = UUID(x_node_id)
|
||
lease_id = UUID(x_lease_id)
|
||
authenticate_node(node_id, _bearer(authorization))
|
||
except (ValueError, SoftwareNodeError) as exc:
|
||
raise HTTPException(401, "invalid node credentials or job identity") from exc
|
||
context = get_job_output_context(node_id, job_id, lease_id, x_request_digest)
|
||
if context is None:
|
||
raise HTTPException(404, "software job output target not found")
|
||
return node_id, lease_id, context
|
||
|
||
|
||
def _hash_file(path: Path) -> str:
|
||
digest = sha256()
|
||
with path.open("rb") as handle:
|
||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||
digest.update(chunk)
|
||
return digest.hexdigest()
|
||
|
||
|
||
def _reject_symlink_path(root: Path, target: Path) -> None:
|
||
root = root.resolve()
|
||
current = root
|
||
for part in target.relative_to(root).parts:
|
||
current = current / part
|
||
if current.is_symlink():
|
||
raise HTTPException(409, "software job output path contains a symbolic link")
|
||
|
||
|
||
def _task_working_dir(root: Path, stored: str) -> Path:
|
||
"""解析 tasks.working_dir 的 ROOT 相对 DB 形态，并再次守住 user_root 边界。"""
|
||
working_dir = from_db_path(stored).resolve()
|
||
try:
|
||
working_dir.relative_to(root.resolve())
|
||
except ValueError as exc:
|
||
raise SoftwareJobError("software job working directory is outside user root") from exc
|
||
return working_dir
|
||
|
||
|
||
def _staged_output_path(
|
||
staging: Path, capability: str, output_id: str, filename: str
|
||
) -> Path:
|
||
flat = staging / filename
|
||
organized = staging / software_job_output_path(capability, output_id)
|
||
return flat if flat.is_file() else organized
|
||
|
||
|
||
def _organize_staged_metadata(
|
||
staging: Path, capability: str, manifest: list[dict]
|
||
) -> None:
|
||
contract = get_contract(capability)
|
||
metadata = [
|
||
item for item in manifest
|
||
if not contract.output_spec(item["artifact_id"]).publish
|
||
]
|
||
if not metadata:
|
||
return
|
||
(staging / ".meta").mkdir(exist_ok=True)
|
||
for item in metadata:
|
||
source = staging / item["filename"]
|
||
destination = staging / software_job_output_path(capability, item["artifact_id"])
|
||
if not source.is_file():
|
||
if destination.is_file():
|
||
continue
|
||
raise SoftwareJobError(f"uploaded artifact is missing: {item['artifact_id']}")
|
||
if destination.exists():
|
||
if (
|
||
destination.is_file()
|
||
and destination.stat().st_size == item["size_bytes"]
|
||
and _hash_file(destination) == item["sha256"]
|
||
):
|
||
source.unlink()
|
||
continue
|
||
raise SoftwareJobError("software job metadata destination conflicts")
|
||
os.replace(source, destination)
|
||
|
||
|
||
def _publish_software_job_outputs(job_id: UUID, context: dict, manifest: list[dict]) -> list[dict]:
|
||
capability = context.get("capability", DEFAULT_CAPABILITIES[0])
|
||
contract = get_contract(capability)
|
||
root = load_user_root(context["user_id"])
|
||
working_dir = _task_working_dir(root, context["working_dir"])
|
||
staging = safe_join(root, f".zcbot_software_job_staging/{job_id}")
|
||
relative_output = Path(contract.output_namespace) / str(job_id)
|
||
destination = safe_join(working_dir, relative_output.as_posix())
|
||
source = staging if staging.is_dir() else destination
|
||
_reject_symlink_path(root, source)
|
||
_reject_symlink_path(root, destination)
|
||
for item in manifest:
|
||
path = (
|
||
_staged_output_path(
|
||
staging, capability, item["artifact_id"], item["filename"]
|
||
)
|
||
if source == staging
|
||
else destination / software_job_output_path(capability, item["artifact_id"])
|
||
)
|
||
if (
|
||
not path.is_file()
|
||
or path.stat().st_size != item["size_bytes"]
|
||
or _hash_file(path) != item["sha256"]
|
||
):
|
||
raise SoftwareJobError(f"uploaded artifact is missing or invalid: {item['artifact_id']}")
|
||
if source == staging:
|
||
_organize_staged_metadata(staging, capability, manifest)
|
||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||
if destination.exists():
|
||
raise SoftwareJobError("software job output destination already exists unexpectedly")
|
||
os.replace(staging, destination)
|
||
try:
|
||
staging.parent.rmdir()
|
||
except OSError:
|
||
pass
|
||
refs = tuple({
|
||
"path": (
|
||
relative_output / software_job_output_path(capability, item["artifact_id"])
|
||
).as_posix(),
|
||
"label": item["filename"],
|
||
"media_type": item["media_type"],
|
||
} for item in manifest if contract.output_spec(item["artifact_id"]).publish)
|
||
published_refs = register_published_artifacts(
|
||
user_id=context["user_id"],
|
||
task_id=context["task_id"],
|
||
user_root=root,
|
||
working_dir=working_dir,
|
||
refs=refs,
|
||
software_job_id=job_id,
|
||
)
|
||
refs_by_path = {item["path"]: item for item in published_refs}
|
||
return [
|
||
{
|
||
**item,
|
||
"source_artifact_id": item["artifact_id"],
|
||
"artifact_id": (
|
||
refs_by_path.get(
|
||
(
|
||
relative_output
|
||
/ software_job_output_path(capability, item["artifact_id"])
|
||
).as_posix(),
|
||
{},
|
||
).get("artifact_id")
|
||
),
|
||
"path": (
|
||
relative_output / software_job_output_path(capability, item["artifact_id"])
|
||
).as_posix(),
|
||
}
|
||
for item in manifest
|
||
]
|
||
|
||
|
||
def register_software_node_routes(app, *, require_user, require_admin) -> None:
|
||
@app.post(
|
||
"/v1/software-nodes/enroll",
|
||
tags=["software-nodes"],
|
||
status_code=status.HTTP_201_CREATED,
|
||
)
|
||
def node_enroll(body: SoftwareNodeEnrollRequest):
|
||
try:
|
||
return enroll_node(**body.model_dump())
|
||
except SoftwareNodeError as exc:
|
||
raise HTTPException(400, str(exc)) from exc
|
||
|
||
@app.get("/v1/software-jobs/{job_id}/inputs/{input_key}", tags=["software-nodes"])
|
||
def download_software_job_input(
|
||
job_id: UUID,
|
||
input_key: str,
|
||
authorization: str | None = Header(default=None),
|
||
x_node_id: str = Header(default=""),
|
||
):
|
||
try:
|
||
node_id = UUID(x_node_id)
|
||
authenticate_node(node_id, _bearer(authorization))
|
||
except (ValueError, SoftwareNodeError) as exc:
|
||
raise HTTPException(401, "invalid node credentials") from exc
|
||
item = get_job_input(node_id, job_id, input_key)
|
||
if item is None:
|
||
raise HTTPException(404, "software job input not found")
|
||
target = safe_join(load_user_root(item["user_id"]), item["current_path"])
|
||
if not target.is_file():
|
||
raise HTTPException(404, "software job input file not found")
|
||
stat = target.stat()
|
||
if stat.st_size != item["size_bytes"]:
|
||
raise HTTPException(409, "software job input changed after submission")
|
||
digest = sha256()
|
||
with target.open("rb") as handle:
|
||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||
digest.update(chunk)
|
||
if digest.hexdigest() != item["sha256"]:
|
||
raise HTTPException(409, "software job input changed after submission")
|
||
return FileResponse(
|
||
path=str(target),
|
||
filename=item["filename"],
|
||
media_type="application/octet-stream",
|
||
headers={
|
||
"Cache-Control": "no-store",
|
||
"X-Content-SHA256": item["sha256"],
|
||
},
|
||
)
|
||
|
||
@app.put(
|
||
"/v1/software-jobs/{job_id}/outputs/{artifact_id}",
|
||
tags=["software-nodes"],
|
||
status_code=status.HTTP_204_NO_CONTENT,
|
||
)
|
||
async def upload_software_job_output(
|
||
job_id: UUID,
|
||
artifact_id: str,
|
||
request: Request,
|
||
authorization: str | None = Header(default=None),
|
||
x_node_id: str = Header(default=""),
|
||
x_lease_id: str = Header(default=""),
|
||
x_request_digest: str = Header(default=""),
|
||
x_content_sha256: str = Header(default=""),
|
||
x_content_length: int = Header(default=-1),
|
||
):
|
||
_, _, context = await asyncio.to_thread(
|
||
_authenticate_output_request,
|
||
job_id, authorization, x_node_id, x_lease_id, x_request_digest,
|
||
)
|
||
try:
|
||
contract = get_contract(context["capability"])
|
||
output_spec = contract.output_spec(artifact_id)
|
||
requested_ids = set(contract.expected_outputs(context["request"]))
|
||
except SoftwareContractError as exc:
|
||
raise HTTPException(400, str(exc)) from exc
|
||
if artifact_id not in requested_ids:
|
||
raise HTTPException(400, "unsupported output artifact identity")
|
||
filename = output_spec.filename
|
||
if not 1 <= x_content_length <= MAX_OUTPUT_ARTIFACT_BYTES:
|
||
raise HTTPException(400, "output artifact size is invalid")
|
||
if len(x_content_sha256) != 64 or any(c not in "0123456789abcdef" for c in x_content_sha256):
|
||
raise HTTPException(400, "output artifact digest is invalid")
|
||
try:
|
||
if succeeded_output_upload_matches(
|
||
context,
|
||
artifact_id,
|
||
size_bytes=x_content_length,
|
||
digest=x_content_sha256,
|
||
):
|
||
return
|
||
except SoftwareJobError as exc:
|
||
raise HTTPException(409, str(exc)) from exc
|
||
root = load_user_root(context["user_id"])
|
||
working_dir = _task_working_dir(root, context["working_dir"])
|
||
published = safe_join(
|
||
working_dir,
|
||
(
|
||
f"{contract.output_namespace}/{job_id}/"
|
||
f"{software_job_output_path(context['capability'], artifact_id)}"
|
||
),
|
||
)
|
||
if published.is_file():
|
||
if published.stat().st_size == x_content_length and _hash_file(published) == x_content_sha256:
|
||
return
|
||
raise HTTPException(409, "published output conflicts with uploaded artifact")
|
||
staging = safe_join(root, f".zcbot_software_job_staging/{job_id}")
|
||
_reject_symlink_path(root, staging)
|
||
staging.mkdir(parents=True, exist_ok=True)
|
||
organized = staging / software_job_output_path(context["capability"], artifact_id)
|
||
if organized.is_file():
|
||
if organized.stat().st_size == x_content_length and _hash_file(organized) == x_content_sha256:
|
||
return
|
||
raise HTTPException(409, "staged output conflicts with uploaded artifact")
|
||
destination = staging / filename
|
||
if destination.is_file():
|
||
if destination.stat().st_size == x_content_length and _hash_file(destination) == x_content_sha256:
|
||
return
|
||
raise HTTPException(409, "uploaded output conflicts with existing staging file")
|
||
staged_total = sum(item.stat().st_size for item in staging.rglob("*") if item.is_file())
|
||
if staged_total + x_content_length > MAX_OUTPUT_TOTAL_BYTES:
|
||
raise HTTPException(413, "software job outputs exceed the total size limit")
|
||
temporary = destination.with_name(destination.name + ".tmp-" + os.urandom(8).hex())
|
||
digest = sha256()
|
||
total = 0
|
||
try:
|
||
with temporary.open("xb") as handle:
|
||
async for chunk in request.stream():
|
||
total += len(chunk)
|
||
if total > x_content_length or total > MAX_OUTPUT_ARTIFACT_BYTES:
|
||
raise HTTPException(413, "output artifact exceeded declared size")
|
||
digest.update(chunk)
|
||
handle.write(chunk)
|
||
handle.flush()
|
||
os.fsync(handle.fileno())
|
||
if total != x_content_length or digest.hexdigest() != x_content_sha256:
|
||
raise HTTPException(400, "output artifact did not match declared metadata")
|
||
os.replace(temporary, destination)
|
||
finally:
|
||
temporary.unlink(missing_ok=True)
|
||
return
|
||
|
||
@app.post("/v1/software-jobs/{job_id}/outputs/complete", tags=["software-nodes"])
|
||
async def complete_software_job_outputs(
|
||
job_id: UUID,
|
||
request: Request,
|
||
authorization: str | None = Header(default=None),
|
||
x_node_id: str = Header(default=""),
|
||
x_lease_id: str = Header(default=""),
|
||
x_request_digest: str = Header(default=""),
|
||
):
|
||
node_id, lease_id, context = await asyncio.to_thread(
|
||
_authenticate_output_request,
|
||
job_id, authorization, x_node_id, x_lease_id, x_request_digest,
|
||
)
|
||
body = await request.json()
|
||
if not isinstance(body, dict):
|
||
raise HTTPException(400, "output completion body must be an object")
|
||
try:
|
||
manifest = validate_output_manifest(
|
||
context["capability"], context["request"], body.get("artifact_manifest")
|
||
)
|
||
replayed = replay_succeeded_outputs(context, manifest)
|
||
if replayed is not None:
|
||
return {"status": "succeeded", "artifact_manifest": replayed}
|
||
published = await asyncio.to_thread(
|
||
_publish_software_job_outputs, job_id, context, manifest
|
||
)
|
||
terminal = {
|
||
"job_id": str(job_id),
|
||
"lease_id": str(lease_id),
|
||
"request_digest": x_request_digest,
|
||
"status": "succeeded",
|
||
"error": {},
|
||
"artifact_manifest": published,
|
||
}
|
||
await asyncio.to_thread(record_job_terminal, node_id, terminal)
|
||
except (SoftwareJobError, KeyError, TypeError) as exc:
|
||
raise HTTPException(409, str(exc)) from exc
|
||
return {"status": "succeeded", "artifact_manifest": published}
|
||
|
||
@app.websocket("/v1/software-nodes/connect")
|
||
async def node_connect(websocket: WebSocket):
|
||
try:
|
||
node_id = UUID(websocket.headers.get("x-node-id", ""))
|
||
token = _bearer(websocket.headers.get("authorization"))
|
||
identity = await asyncio.to_thread(authenticate_node, node_id, token)
|
||
except (ValueError, SoftwareNodeError):
|
||
# 握手前 close 会被 ASGI 统一表现为 HTTP 403，客户端无法区分
|
||
# “凭据无效”和“代理/路由没有正确转发 WebSocket”。先升级再用
|
||
# 应用关闭码给已持有 Node ID/Token 的节点返回明确诊断。
|
||
await websocket.accept()
|
||
await websocket.close(code=4003, reason="invalid node credentials")
|
||
return
|
||
await websocket.accept()
|
||
await node_connections.activate(node_id, websocket)
|
||
try:
|
||
await node_connections.send_on(
|
||
node_id,
|
||
websocket,
|
||
{"type": "connected", "heartbeat_seconds": 15},
|
||
)
|
||
for cancel in await asyncio.to_thread(pending_node_cancellations, node_id):
|
||
await node_connections.send_on(
|
||
node_id, websocket, {"type": "job_cancel", "payload": cancel}
|
||
)
|
||
while True:
|
||
message = await websocket.receive_json()
|
||
message_type = message.get("type")
|
||
payload = message.get("payload") or {}
|
||
if not isinstance(payload, dict):
|
||
await node_connections.send_on(node_id, websocket,
|
||
{"type": "error", "code": "unsupported_message"}
|
||
)
|
||
continue
|
||
if message_type in {"job_accept", "job_reject"}:
|
||
await asyncio.to_thread(
|
||
respond_to_offer,
|
||
node_id,
|
||
accepted=message_type == "job_accept",
|
||
payload=payload,
|
||
)
|
||
await node_connections.send_on(node_id, websocket,
|
||
{"type": "ack", "message_id": message.get("message_id")}
|
||
)
|
||
continue
|
||
if message_type in {"job_state", "job_terminal"}:
|
||
if message_type == "job_terminal" and payload.get("status") == "succeeded":
|
||
await node_connections.send_on(node_id, websocket,
|
||
{
|
||
"type": "error",
|
||
"code": "outputs_not_published",
|
||
"message_id": message.get("message_id"),
|
||
}
|
||
)
|
||
continue
|
||
handler = (
|
||
update_job_state
|
||
if message_type == "job_state"
|
||
else record_job_terminal
|
||
)
|
||
await asyncio.to_thread(handler, node_id, payload)
|
||
await node_connections.send_on(node_id, websocket,
|
||
{"type": "ack", "message_id": message.get("message_id")}
|
||
)
|
||
continue
|
||
if message_type not in {"hello", "heartbeat"}:
|
||
await node_connections.send_on(node_id, websocket,
|
||
{"type": "error", "code": "unsupported_message"}
|
||
)
|
||
continue
|
||
if payload.get("install_id") and payload["install_id"] != str(
|
||
identity["install_id"]
|
||
):
|
||
await websocket.close(code=1008, reason="install identity mismatch")
|
||
return
|
||
await asyncio.to_thread(
|
||
update_node_runtime,
|
||
node_id,
|
||
status="online",
|
||
runtime=payload,
|
||
)
|
||
await node_connections.send_on(node_id, websocket,
|
||
{"type": "ack", "message_id": message.get("message_id")}
|
||
)
|
||
for cancel in await asyncio.to_thread(pending_node_cancellations, node_id):
|
||
await node_connections.send_on(
|
||
node_id, websocket, {"type": "job_cancel", "payload": cancel}
|
||
)
|
||
offer = await asyncio.to_thread(
|
||
offer_next_job, await node_connections.node_ids()
|
||
)
|
||
if offer is not None:
|
||
delivered = await node_connections.send(
|
||
offer["node_id"],
|
||
{"type": "job_offer", "payload": offer["payload"]},
|
||
)
|
||
if not delivered:
|
||
await asyncio.to_thread(
|
||
abandon_offer, offer["node_id"], offer["payload"]
|
||
)
|
||
except (
|
||
SoftwareJobError,
|
||
SoftwareNodeError,
|
||
WebSocketDisconnect,
|
||
RuntimeError,
|
||
ValueError,
|
||
):
|
||
pass
|
||
finally:
|
||
if await node_connections.remove(node_id, websocket):
|
||
await asyncio.to_thread(mark_node_offline, node_id)
|
||
await asyncio.to_thread(mark_node_jobs_disconnected, node_id)
|
||
|
||
@app.post("/v1/tasks/{task_id}/software-jobs", tags=["software-jobs"])
|
||
async def submit_software_job(
|
||
task_id: UUID,
|
||
body: SoftwareJobCreateRequest,
|
||
user_id: UUID = Depends(require_user), # noqa: B008
|
||
):
|
||
try:
|
||
job, created = await asyncio.to_thread(
|
||
create_job, user_id, task_id, **body.model_dump()
|
||
)
|
||
offer = await asyncio.to_thread(
|
||
offer_next_job, await node_connections.node_ids()
|
||
)
|
||
if offer is not None:
|
||
delivered = await node_connections.send(
|
||
offer["node_id"], {"type": "job_offer", "payload": offer["payload"]}
|
||
)
|
||
if not delivered:
|
||
await asyncio.to_thread(
|
||
abandon_offer, offer["node_id"], offer["payload"]
|
||
)
|
||
return {**job, "created": created}
|
||
except SoftwareJobError as exc:
|
||
detail = str(exc)
|
||
raise HTTPException(404 if detail == "task not found" else 400, detail) from exc
|
||
|
||
@app.get("/v1/software-jobs", tags=["software-jobs"])
|
||
def read_software_jobs(
|
||
task_id: UUID | None = None,
|
||
active_only: bool = False,
|
||
limit: int = 50,
|
||
before_created_at: datetime | None = None,
|
||
before_job_id: UUID | None = None,
|
||
user_id: UUID = Depends(require_user), # noqa: B008
|
||
):
|
||
if (before_created_at is None) != (before_job_id is None):
|
||
raise HTTPException(
|
||
400, "before_created_at and before_job_id must be provided together"
|
||
)
|
||
page_limit = max(1, min(int(limit), 100))
|
||
before = (
|
||
(before_created_at, before_job_id)
|
||
if before_created_at is not None and before_job_id is not None
|
||
else None
|
||
)
|
||
results = list_jobs(
|
||
user_id, task_id=task_id, active_only=active_only,
|
||
limit=page_limit + 1, before=before,
|
||
)
|
||
has_more = len(results) > page_limit
|
||
results = results[:page_limit]
|
||
last = results[-1] if has_more else None
|
||
return {
|
||
"results": results,
|
||
"next_cursor": (
|
||
{"created_at": last["created_at"], "job_id": last["job_id"]}
|
||
if last is not None else None
|
||
),
|
||
}
|
||
|
||
@app.post("/v1/software-jobs/{job_id}/cancel", tags=["software-jobs"])
|
||
async def cancel_software_job(
|
||
job_id: UUID,
|
||
user_id: UUID = Depends(require_user), # noqa: B008
|
||
):
|
||
try:
|
||
job, message = await asyncio.to_thread(request_job_cancel, user_id, job_id)
|
||
except SoftwareJobError as exc:
|
||
raise HTTPException(404, str(exc)) from exc
|
||
if message is not None:
|
||
await node_connections.send(
|
||
message["node_id"],
|
||
{"type": "job_cancel", "payload": message["payload"]},
|
||
)
|
||
return job
|
||
|
||
@app.get("/v1/software-jobs/{job_id}", tags=["software-jobs"])
|
||
def read_software_job(
|
||
job_id: UUID,
|
||
user_id: UUID = Depends(require_user), # noqa: B008
|
||
):
|
||
job = get_job(user_id, job_id)
|
||
if job is None:
|
||
raise HTTPException(404, "software job not found")
|
||
return job
|
||
|
||
@app.post("/v1/admin/software-node-enrollments", tags=["admin"])
|
||
def admin_create_software_enrollment(
|
||
body: SoftwareEnrollmentCreateRequest,
|
||
user_id: UUID = Depends(require_admin), # noqa: B008
|
||
):
|
||
try:
|
||
return create_enrollment(user_id, **body.model_dump())
|
||
except SoftwareNodeError as exc:
|
||
raise HTTPException(400, str(exc)) from exc
|
||
|
||
@app.get("/v1/admin/software-nodes", tags=["admin"])
|
||
def admin_software_nodes(user_id: UUID = Depends(require_admin)): # noqa: B008
|
||
return {"results": list_nodes()}
|
||
|
||
@app.patch("/v1/admin/software-nodes/{node_id}", tags=["admin"])
|
||
async def admin_disable_software_node(
|
||
node_id: UUID,
|
||
body: SoftwareNodeDisableRequest,
|
||
user_id: UUID = Depends(require_admin), # noqa: B008
|
||
):
|
||
if not await asyncio.to_thread(set_node_disabled, node_id, body.disabled):
|
||
raise HTTPException(404, "software node not found")
|
||
if body.disabled:
|
||
await node_connections.close(node_id)
|
||
return {
|
||
"node_id": str(node_id),
|
||
"status": "disabled" if body.disabled else "offline",
|
||
}
|
||
|
||
@app.delete("/v1/admin/software-nodes/{node_id}", tags=["admin"])
|
||
async def admin_delete_software_node(
|
||
node_id: UUID,
|
||
user_id: UUID = Depends(require_admin), # noqa: B008
|
||
):
|
||
# 先撤掉在线连接，避免删除后的旧 socket 继续上报运行态。
|
||
await node_connections.close(node_id)
|
||
if not await asyncio.to_thread(delete_node, node_id):
|
||
raise HTTPException(404, "software node not found")
|
||
return {"node_id": str(node_id), "status": "deleted"}
|