zcbot/web/routers/software_nodes.py

706 lines
28 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""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 SoftwareContractError, default_capabilities, 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"}