316 lines
13 KiB
Python
316 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""HTTP API for AIOS external-agent online task distribution.
|
|
|
|
The service is intended to bind to loopback and be published through the
|
|
existing TLS reverse proxy. Control operations use a server-managed control
|
|
Bearer token. Agents receive individual Bearer tokens at registration; only
|
|
token hashes are persisted.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import hmac
|
|
import json
|
|
import os
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
import re
|
|
from typing import Any, Mapping
|
|
from urllib.parse import urlparse
|
|
|
|
from aios_agent_task_protocol import TaskProtocolError, TaskStore
|
|
|
|
|
|
MAX_BODY = 1024 * 1024
|
|
TASK_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{2,179}$")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class APIConfig:
|
|
db_path: Path
|
|
control_token_path: Path
|
|
host: str = "127.0.0.1"
|
|
port: int = 18940
|
|
token_pepper: str = ""
|
|
|
|
@classmethod
|
|
def from_env(cls) -> "APIConfig":
|
|
return cls(
|
|
db_path=Path(os.getenv(
|
|
"AIOS_AGENT_TASK_DB",
|
|
"/var/lib/aios-agent-task/state.sqlite3",
|
|
)),
|
|
control_token_path=Path(os.getenv(
|
|
"AIOS_AGENT_TASK_CONTROL_TOKEN_FILE",
|
|
"/run/credentials/aios-agent-task-api.service/control_token",
|
|
)),
|
|
host=os.getenv("AIOS_AGENT_TASK_HOST", "127.0.0.1"),
|
|
port=int(os.getenv("AIOS_AGENT_TASK_PORT", "18940")),
|
|
token_pepper=os.getenv("AIOS_AGENT_TASK_TOKEN_PEPPER", ""),
|
|
)
|
|
|
|
|
|
class AgentTaskApplication:
|
|
def __init__(self, config: APIConfig) -> None:
|
|
self.config = config
|
|
self.store = TaskStore(
|
|
config.db_path,
|
|
token_pepper=config.token_pepper,
|
|
)
|
|
|
|
def _control_token(self) -> str:
|
|
path = self.config.control_token_path
|
|
if path.is_symlink() or not path.is_file():
|
|
raise TaskProtocolError("CONTROL_CREDENTIAL_UNAVAILABLE", 503)
|
|
value = path.read_text(encoding="utf-8").strip()
|
|
if len(value) < 32:
|
|
raise TaskProtocolError("CONTROL_CREDENTIAL_INVALID", 503)
|
|
return value
|
|
|
|
@staticmethod
|
|
def _bearer(headers: Mapping[str, str]) -> str:
|
|
value = str(headers.get("Authorization") or "")
|
|
if not value.startswith("Bearer "):
|
|
return ""
|
|
return value[7:].strip()
|
|
|
|
def authenticate_control(self, headers: Mapping[str, str]) -> None:
|
|
supplied = self._bearer(headers)
|
|
expected = self._control_token()
|
|
if not supplied or not hmac.compare_digest(supplied, expected):
|
|
raise TaskProtocolError("UNAUTHORIZED", 401)
|
|
|
|
def authenticate_agent(self, headers: Mapping[str, str]) -> dict[str, Any]:
|
|
return self.store.authenticate_agent(self._bearer(headers))
|
|
|
|
def principal(self, headers: Mapping[str, str]) -> tuple[str, dict[str, Any] | None]:
|
|
supplied = self._bearer(headers)
|
|
if not supplied:
|
|
raise TaskProtocolError("UNAUTHORIZED", 401)
|
|
try:
|
|
expected = self._control_token()
|
|
except TaskProtocolError:
|
|
expected = ""
|
|
if expected and hmac.compare_digest(supplied, expected):
|
|
return "control", None
|
|
return "agent", self.store.authenticate_agent(supplied)
|
|
|
|
def health(self) -> dict[str, Any]:
|
|
return {
|
|
"status": "ok",
|
|
"schema": "aios.agent-task.v1",
|
|
"task_states": [
|
|
"QUEUED", "CLAIMED", "RUNNING", "VERIFYING",
|
|
"SUCCEEDED", "FAILED", "EXPIRED", "CANCELLED",
|
|
],
|
|
"lease_recovery": True,
|
|
"atomic_claim": True,
|
|
"idempotent_result_submission": True,
|
|
"server_side_verification_required": True,
|
|
"plaintext_credentials_persisted": False,
|
|
}
|
|
|
|
def register_agent(self, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
|
|
self.authenticate_control(headers)
|
|
return self.store.register_agent(
|
|
agent_id=str(body.get("agent_id") or ""),
|
|
agent_type=str(body.get("agent_type") or ""),
|
|
agent_version=str(body.get("agent_version") or ""),
|
|
capabilities=body.get("capabilities") or [],
|
|
metadata=body.get("metadata") if isinstance(body.get("metadata"), Mapping) else {},
|
|
)
|
|
|
|
def create_task(self, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
|
|
self.authenticate_control(headers)
|
|
return self.store.create_task(
|
|
task_id=str(body.get("task_id") or "") or None,
|
|
goal=str(body.get("goal") or ""),
|
|
instructions=str(body.get("instructions") or ""),
|
|
success_criteria=body.get("success_criteria"),
|
|
required_capabilities=body.get("required_capabilities") or [],
|
|
constraints=body.get("constraints") if isinstance(body.get("constraints"), Mapping) else {},
|
|
metadata=body.get("metadata") if isinstance(body.get("metadata"), Mapping) else {},
|
|
deadline=str(body.get("deadline") or "") or None,
|
|
priority=body.get("priority", 0),
|
|
max_attempts=body.get("max_attempts", 3),
|
|
idempotency_key=str(body.get("idempotency_key") or ""),
|
|
)
|
|
|
|
def poll(self, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
|
|
agent = self.authenticate_agent(headers)
|
|
tasks = self.store.poll_tasks(
|
|
agent_id=agent["agent_id"],
|
|
limit=int(body.get("limit", 20)),
|
|
)
|
|
return {"agent_id": agent["agent_id"], "tasks": tasks}
|
|
|
|
def claim(self, task_id: str, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
|
|
agent = self.authenticate_agent(headers)
|
|
return self.store.claim_task(
|
|
task_id=task_id,
|
|
agent_id=agent["agent_id"],
|
|
lease_seconds=int(body.get("lease_seconds", 300)),
|
|
)
|
|
|
|
def heartbeat(self, task_id: str, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
|
|
agent = self.authenticate_agent(headers)
|
|
return self.store.heartbeat(
|
|
task_id=task_id,
|
|
agent_id=agent["agent_id"],
|
|
lease_token=str(body.get("lease_token") or ""),
|
|
progress=body.get("progress") if isinstance(body.get("progress"), Mapping) else {},
|
|
lease_seconds=int(body.get("lease_seconds", 300)),
|
|
)
|
|
|
|
def result(self, task_id: str, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
|
|
agent = self.authenticate_agent(headers)
|
|
result = body.get("result")
|
|
if not isinstance(result, Mapping):
|
|
raise TaskProtocolError("RESULT_INVALID", 400)
|
|
return self.store.submit_result(
|
|
task_id=task_id,
|
|
agent_id=agent["agent_id"],
|
|
lease_token=str(body.get("lease_token") or ""),
|
|
idempotency_key=str(body.get("idempotency_key") or ""),
|
|
result=result,
|
|
)
|
|
|
|
def verify(self, task_id: str, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
|
|
self.authenticate_control(headers)
|
|
return self.store.verify_task(
|
|
task_id=task_id,
|
|
decision=str(body.get("decision") or ""),
|
|
verification=body.get("verification") if isinstance(body.get("verification"), Mapping) else {},
|
|
)
|
|
|
|
def cancel(self, task_id: str, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
|
|
self.authenticate_control(headers)
|
|
return self.store.cancel_task(task_id=task_id, reason=str(body.get("reason") or ""))
|
|
|
|
def get_task(self, task_id: str, headers: Mapping[str, str]) -> dict[str, Any]:
|
|
kind, agent = self.principal(headers)
|
|
task = self.store.get_task(task_id)
|
|
if kind == "agent" and task.get("claimed_by") != agent["agent_id"]:
|
|
raise TaskProtocolError("FORBIDDEN", 403)
|
|
return task
|
|
|
|
def events(self, task_id: str, headers: Mapping[str, str]) -> dict[str, Any]:
|
|
self.authenticate_control(headers)
|
|
return {"task_id": task_id, "events": self.store.get_events(task_id)}
|
|
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
app: AgentTaskApplication | None = None
|
|
server_version = "AIOSAgentTaskAPI/1"
|
|
|
|
def log_message(self, _format: str, *_args: Any) -> None:
|
|
return
|
|
|
|
def _headers(self) -> dict[str, str]:
|
|
return {key: value for key, value in self.headers.items()}
|
|
|
|
def _json(self, value: Any, status: int = 200) -> None:
|
|
payload = json.dumps(value, ensure_ascii=False, sort_keys=True).encode("utf-8")
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json; charset=utf-8")
|
|
self.send_header("Cache-Control", "no-store")
|
|
self.send_header("Content-Length", str(len(payload)))
|
|
self.end_headers()
|
|
self.wfile.write(payload)
|
|
|
|
def _body(self) -> dict[str, Any]:
|
|
try:
|
|
size = int(self.headers.get("Content-Length", "0"))
|
|
except ValueError as exc:
|
|
raise TaskProtocolError("CONTENT_LENGTH_INVALID", 400) from exc
|
|
if size < 0 or size > MAX_BODY:
|
|
raise TaskProtocolError("REQUEST_TOO_LARGE", 413)
|
|
raw = self.rfile.read(size) if size else b"{}"
|
|
try:
|
|
value = json.loads(raw.decode("utf-8"))
|
|
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
raise TaskProtocolError("JSON_INVALID", 400) from exc
|
|
if not isinstance(value, dict):
|
|
raise TaskProtocolError("JSON_OBJECT_REQUIRED", 400)
|
|
return value
|
|
|
|
@staticmethod
|
|
def _task_route(path: str) -> tuple[str, str | None]:
|
|
match = re.fullmatch(
|
|
r"/v1/tasks/([A-Za-z0-9][A-Za-z0-9._:-]{2,179})(?:/(claim|heartbeat|result|verify|cancel|events))?",
|
|
path,
|
|
)
|
|
if match:
|
|
return match.group(1), match.group(2)
|
|
return "", None
|
|
|
|
def _run(self, callback) -> None:
|
|
try:
|
|
callback()
|
|
except TaskProtocolError as exc:
|
|
self._json({"error": exc.code}, exc.http_status)
|
|
except Exception as exc:
|
|
self._json({"error": "INTERNAL_ERROR", "detail": type(exc).__name__}, 500)
|
|
|
|
def do_GET(self) -> None:
|
|
def action() -> None:
|
|
if self.app is None:
|
|
raise TaskProtocolError("SERVICE_UNAVAILABLE", 503)
|
|
path = urlparse(self.path).path.rstrip("/") or "/"
|
|
if path == "/health":
|
|
self._json(self.app.health())
|
|
return
|
|
task_id, suffix = self._task_route(path)
|
|
if task_id and suffix is None:
|
|
self._json(self.app.get_task(task_id, self._headers()))
|
|
return
|
|
if task_id and suffix == "events":
|
|
self._json(self.app.events(task_id, self._headers()))
|
|
return
|
|
self._json({"error": "NOT_FOUND"}, 404)
|
|
self._run(action)
|
|
|
|
def do_POST(self) -> None:
|
|
def action() -> None:
|
|
if self.app is None:
|
|
raise TaskProtocolError("SERVICE_UNAVAILABLE", 503)
|
|
path = urlparse(self.path).path.rstrip("/") or "/"
|
|
body = self._body()
|
|
if path == "/v1/agents/register":
|
|
self._json(self.app.register_agent(self._headers(), body), 201)
|
|
return
|
|
if path == "/v1/tasks":
|
|
self._json(self.app.create_task(self._headers(), body), 201)
|
|
return
|
|
if path == "/v1/tasks/poll":
|
|
self._json(self.app.poll(self._headers(), body))
|
|
return
|
|
task_id, suffix = self._task_route(path)
|
|
if not task_id:
|
|
self._json({"error": "NOT_FOUND"}, 404)
|
|
return
|
|
if suffix == "claim":
|
|
self._json(self.app.claim(task_id, self._headers(), body))
|
|
elif suffix == "heartbeat":
|
|
self._json(self.app.heartbeat(task_id, self._headers(), body))
|
|
elif suffix == "result":
|
|
self._json(self.app.result(task_id, self._headers(), body))
|
|
elif suffix == "verify":
|
|
self._json(self.app.verify(task_id, self._headers(), body))
|
|
elif suffix == "cancel":
|
|
self._json(self.app.cancel(task_id, self._headers(), body))
|
|
else:
|
|
self._json({"error": "NOT_FOUND"}, 404)
|
|
self._run(action)
|
|
|
|
|
|
def serve(config: APIConfig | None = None) -> None:
|
|
config = config or APIConfig.from_env()
|
|
application = AgentTaskApplication(config)
|
|
Handler.app = application
|
|
ThreadingHTTPServer((config.host, config.port), Handler).serve_forever()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
serve()
|