871 lines
35 KiB
Python
871 lines
35 KiB
Python
#!/usr/bin/env python3
|
|
"""Durable online task protocol for external AIOS agents.
|
|
|
|
The store is deliberately independent from Git/Gitea mutation. It provides a
|
|
small durable state machine for publishing work, atomic claim/lease ownership,
|
|
heartbeat, idempotent result submission, explicit server-side verification,
|
|
lease recovery, and audit events.
|
|
|
|
No plaintext agent credential or task lease token is persisted.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta, timezone
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
from pathlib import Path
|
|
import secrets
|
|
import sqlite3
|
|
from typing import Any, Callable, Iterable, Mapping
|
|
import uuid
|
|
|
|
from aios_capability_contract import has_capabilities
|
|
from aios_online_agent_execution import (
|
|
ExecutionProjectionError,
|
|
validate_execution_envelope,
|
|
)
|
|
|
|
|
|
TERMINAL_STATES = frozenset({"SUCCEEDED", "FAILED", "EXPIRED", "CANCELLED"})
|
|
ACTIVE_LEASE_STATES = frozenset({"CLAIMED", "RUNNING"})
|
|
TASK_STATES = frozenset({
|
|
"QUEUED", "CLAIMED", "RUNNING", "VERIFYING",
|
|
"SUCCEEDED", "FAILED", "EXPIRED", "CANCELLED",
|
|
})
|
|
VERIFY_DECISIONS = frozenset({"SUCCEEDED", "NEEDS_REWORK", "FAILED"})
|
|
|
|
|
|
GOVERNED_WRITE_CAPABILITIES = frozenset({
|
|
"governed.write.branch",
|
|
"governed.write.file",
|
|
})
|
|
EXECUTION_PROJECTION_SCHEMA = "aios.online-agent.task-projection.v1"
|
|
|
|
|
|
def _validate_write_execution_binding(
|
|
task_id: str,
|
|
required_capabilities: list[str],
|
|
metadata: Mapping[str, Any],
|
|
) -> None:
|
|
required = set(required_capabilities)
|
|
if not required.intersection(GOVERNED_WRITE_CAPABILITIES):
|
|
return
|
|
execution = metadata.get("execution")
|
|
projection = metadata.get("execution_projection")
|
|
if not isinstance(execution, Mapping) or not isinstance(projection, Mapping):
|
|
raise TaskProtocolError("GOVERNED_EXECUTION_BINDING_REQUIRED", 409)
|
|
try:
|
|
canonical = validate_execution_envelope(dict(execution))
|
|
except (ExecutionProjectionError, TypeError, ValueError, KeyError) as exc:
|
|
raise TaskProtocolError("EXECUTION_ENVELOPE_INVALID", 409) from exc
|
|
if canonical.get("task_id") != task_id:
|
|
raise TaskProtocolError("TASK_BINDING_CONFLICT", 409)
|
|
if set(canonical.get("capability_scope") or ()) != required:
|
|
raise TaskProtocolError("EXECUTION_CAPABILITY_SCOPE_CONFLICT", 409)
|
|
if (
|
|
projection.get("schema") != EXECUTION_PROJECTION_SCHEMA
|
|
or projection.get("task_id") != task_id
|
|
or projection.get("job_id") != canonical.get("job_id")
|
|
or projection.get("execution_digest") != canonical.get("execution_digest")
|
|
):
|
|
raise TaskProtocolError("EXECUTION_PROJECTION_INVALID", 409)
|
|
|
|
|
|
class TaskProtocolError(RuntimeError):
|
|
def __init__(self, code: str, http_status: int = 409) -> None:
|
|
super().__init__(code)
|
|
self.code = code
|
|
self.http_status = http_status
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Principal:
|
|
kind: str
|
|
identifier: str
|
|
|
|
|
|
def utcnow() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def iso(value: datetime) -> str:
|
|
return value.astimezone(timezone.utc).isoformat()
|
|
|
|
|
|
def parse_time(value: str | None) -> datetime | None:
|
|
if not value:
|
|
return None
|
|
parsed = datetime.fromisoformat(value)
|
|
if parsed.tzinfo is None:
|
|
parsed = parsed.replace(tzinfo=timezone.utc)
|
|
return parsed.astimezone(timezone.utc)
|
|
|
|
|
|
def canonical_json(value: Any) -> str:
|
|
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
|
|
|
|
|
def digest_json(value: Any) -> str:
|
|
return hashlib.sha256(canonical_json(value).encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _string_list(value: Iterable[object] | None, field: str, *, max_items: int = 128) -> list[str]:
|
|
if value is None:
|
|
return []
|
|
if isinstance(value, (str, bytes)) or not isinstance(value, Iterable):
|
|
raise TaskProtocolError(field.upper() + "_INVALID", 400)
|
|
result: list[str] = []
|
|
for item in value:
|
|
text = str(item).strip()
|
|
if not text or len(text) > 160:
|
|
raise TaskProtocolError(field.upper() + "_INVALID", 400)
|
|
result.append(text)
|
|
if len(result) > max_items:
|
|
raise TaskProtocolError(field.upper() + "_TOO_LARGE", 400)
|
|
return sorted(set(result))
|
|
|
|
|
|
def _json_object(value: object | None, field: str, *, max_bytes: int = 128 * 1024) -> dict[str, Any]:
|
|
if value is None:
|
|
return {}
|
|
if not isinstance(value, Mapping):
|
|
raise TaskProtocolError(field.upper() + "_INVALID", 400)
|
|
result = dict(value)
|
|
if len(canonical_json(result).encode("utf-8")) > max_bytes:
|
|
raise TaskProtocolError(field.upper() + "_TOO_LARGE", 413)
|
|
return result
|
|
|
|
|
|
def _criteria(value: object | None) -> list[Any]:
|
|
if value is None:
|
|
return []
|
|
if isinstance(value, (str, bytes)) or not isinstance(value, list):
|
|
raise TaskProtocolError("SUCCESS_CRITERIA_INVALID", 400)
|
|
if len(value) > 128:
|
|
raise TaskProtocolError("SUCCESS_CRITERIA_TOO_LARGE", 413)
|
|
if len(canonical_json(value).encode("utf-8")) > 128 * 1024:
|
|
raise TaskProtocolError("SUCCESS_CRITERIA_TOO_LARGE", 413)
|
|
return list(value)
|
|
|
|
|
|
class TaskStore:
|
|
"""SQLite-backed durable task store.
|
|
|
|
SQLite BEGIN IMMEDIATE serializes claims so two agents cannot successfully
|
|
claim the same queued task. WAL mode allows concurrent readers.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
path: str | Path,
|
|
*,
|
|
now: Callable[[], datetime] = utcnow,
|
|
token_pepper: str = "",
|
|
) -> None:
|
|
self.path = Path(path)
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.now = now
|
|
self.token_pepper = token_pepper
|
|
self._initialize()
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
connection = sqlite3.connect(self.path, timeout=10, isolation_level=None)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA foreign_keys=ON")
|
|
connection.execute("PRAGMA journal_mode=WAL")
|
|
connection.execute("PRAGMA synchronous=FULL")
|
|
connection.execute("PRAGMA busy_timeout=10000")
|
|
return connection
|
|
|
|
def _initialize(self) -> None:
|
|
with self._connect() as connection:
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS agents (
|
|
agent_id TEXT PRIMARY KEY,
|
|
agent_type TEXT NOT NULL,
|
|
agent_version TEXT NOT NULL,
|
|
capabilities_json TEXT NOT NULL,
|
|
metadata_json TEXT NOT NULL,
|
|
token_hash TEXT NOT NULL UNIQUE,
|
|
status TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
last_seen_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS tasks (
|
|
task_id TEXT PRIMARY KEY,
|
|
goal TEXT NOT NULL,
|
|
instructions TEXT NOT NULL,
|
|
success_criteria_json TEXT NOT NULL,
|
|
required_capabilities_json TEXT NOT NULL,
|
|
constraints_json TEXT NOT NULL,
|
|
metadata_json TEXT NOT NULL,
|
|
deadline TEXT,
|
|
priority INTEGER NOT NULL,
|
|
max_attempts INTEGER NOT NULL,
|
|
attempt INTEGER NOT NULL DEFAULT 0,
|
|
status TEXT NOT NULL,
|
|
claimed_by TEXT,
|
|
lease_token_hash TEXT,
|
|
lease_expires_at TEXT,
|
|
result_json TEXT,
|
|
result_digest TEXT,
|
|
result_idempotency_key TEXT,
|
|
verification_json TEXT,
|
|
revision INTEGER NOT NULL DEFAULT 1,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL,
|
|
FOREIGN KEY(claimed_by) REFERENCES agents(agent_id)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS idempotency (
|
|
scope TEXT NOT NULL,
|
|
key TEXT NOT NULL,
|
|
payload_digest TEXT NOT NULL,
|
|
response_json TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
PRIMARY KEY(scope, key)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS events (
|
|
event_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
task_id TEXT,
|
|
agent_id TEXT,
|
|
event_type TEXT NOT NULL,
|
|
payload_json TEXT NOT NULL,
|
|
created_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_tasks_queue
|
|
ON tasks(status, priority DESC, created_at);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_tasks_lease
|
|
ON tasks(status, lease_expires_at);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_events_task
|
|
ON events(task_id, event_id);
|
|
"""
|
|
)
|
|
|
|
def _token_hash(self, token: str) -> str:
|
|
return hashlib.sha256((self.token_pepper + "\0" + token).encode("utf-8")).hexdigest()
|
|
|
|
@staticmethod
|
|
def _row_task(row: sqlite3.Row, *, include_result: bool = True) -> dict[str, Any]:
|
|
value: dict[str, Any] = {
|
|
"task_id": row["task_id"],
|
|
"goal": row["goal"],
|
|
"instructions": row["instructions"],
|
|
"success_criteria": json.loads(row["success_criteria_json"]),
|
|
"required_capabilities": json.loads(row["required_capabilities_json"]),
|
|
"constraints": json.loads(row["constraints_json"]),
|
|
"metadata": json.loads(row["metadata_json"]),
|
|
"deadline": row["deadline"],
|
|
"priority": row["priority"],
|
|
"max_attempts": row["max_attempts"],
|
|
"attempt": row["attempt"],
|
|
"status": row["status"],
|
|
"claimed_by": row["claimed_by"],
|
|
"lease_expires_at": row["lease_expires_at"],
|
|
"revision": row["revision"],
|
|
"created_at": row["created_at"],
|
|
"updated_at": row["updated_at"],
|
|
}
|
|
if include_result:
|
|
value["result"] = json.loads(row["result_json"]) if row["result_json"] else None
|
|
value["verification"] = (
|
|
json.loads(row["verification_json"]) if row["verification_json"] else None
|
|
)
|
|
return value
|
|
|
|
def _event(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
event_type: str,
|
|
*,
|
|
task_id: str | None = None,
|
|
agent_id: str | None = None,
|
|
payload: Mapping[str, Any] | None = None,
|
|
) -> None:
|
|
connection.execute(
|
|
"INSERT INTO events(task_id,agent_id,event_type,payload_json,created_at)"
|
|
" VALUES(?,?,?,?,?)",
|
|
(task_id, agent_id, event_type, canonical_json(dict(payload or {})), iso(self.now())),
|
|
)
|
|
|
|
def register_agent(
|
|
self,
|
|
*,
|
|
agent_id: str,
|
|
agent_type: str,
|
|
agent_version: str,
|
|
capabilities: Iterable[object],
|
|
metadata: Mapping[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
agent_id = agent_id.strip()
|
|
agent_type = agent_type.strip()
|
|
agent_version = agent_version.strip()
|
|
if not agent_id or len(agent_id) > 160 or not agent_type or not agent_version:
|
|
raise TaskProtocolError("AGENT_IDENTITY_INVALID", 400)
|
|
caps = _string_list(capabilities, "capabilities")
|
|
meta = _json_object(metadata, "metadata")
|
|
token = "agt_" + secrets.token_urlsafe(32)
|
|
now = iso(self.now())
|
|
try:
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
connection.execute(
|
|
"INSERT INTO agents(agent_id,agent_type,agent_version,capabilities_json,"
|
|
"metadata_json,token_hash,status,created_at,last_seen_at)"
|
|
" VALUES(?,?,?,?,?,?,?,?,?)",
|
|
(
|
|
agent_id,
|
|
agent_type,
|
|
agent_version,
|
|
canonical_json(caps),
|
|
canonical_json(meta),
|
|
self._token_hash(token),
|
|
"ACTIVE",
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
self._event(
|
|
connection,
|
|
"AGENT_REGISTERED",
|
|
agent_id=agent_id,
|
|
payload={"agent_type": agent_type, "agent_version": agent_version},
|
|
)
|
|
connection.commit()
|
|
except sqlite3.IntegrityError as exc:
|
|
raise TaskProtocolError("AGENT_ALREADY_EXISTS", 409) from exc
|
|
return {
|
|
"agent_id": agent_id,
|
|
"agent_type": agent_type,
|
|
"agent_version": agent_version,
|
|
"capabilities": caps,
|
|
"agent_token": token,
|
|
"credential_returned_once": True,
|
|
}
|
|
|
|
def authenticate_agent(self, token: str) -> dict[str, Any]:
|
|
if not token:
|
|
raise TaskProtocolError("UNAUTHORIZED", 401)
|
|
token_hash = self._token_hash(token)
|
|
with self._connect() as connection:
|
|
row = connection.execute(
|
|
"SELECT * FROM agents WHERE token_hash=? AND status='ACTIVE'",
|
|
(token_hash,),
|
|
).fetchone()
|
|
if row is None:
|
|
raise TaskProtocolError("UNAUTHORIZED", 401)
|
|
connection.execute(
|
|
"UPDATE agents SET last_seen_at=? WHERE agent_id=?",
|
|
(iso(self.now()), row["agent_id"]),
|
|
)
|
|
return {
|
|
"agent_id": row["agent_id"],
|
|
"agent_type": row["agent_type"],
|
|
"agent_version": row["agent_version"],
|
|
"capabilities": json.loads(row["capabilities_json"]),
|
|
"metadata": json.loads(row["metadata_json"]),
|
|
}
|
|
|
|
def _idempotent_lookup(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
*,
|
|
scope: str,
|
|
key: str,
|
|
payload_digest: str,
|
|
) -> dict[str, Any] | None:
|
|
row = connection.execute(
|
|
"SELECT payload_digest,response_json FROM idempotency WHERE scope=? AND key=?",
|
|
(scope, key),
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
if not hmac.compare_digest(row["payload_digest"], payload_digest):
|
|
raise TaskProtocolError("IDEMPOTENCY_CONFLICT", 409)
|
|
return json.loads(row["response_json"])
|
|
|
|
def _idempotent_store(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
*,
|
|
scope: str,
|
|
key: str,
|
|
payload_digest: str,
|
|
response: Mapping[str, Any],
|
|
) -> None:
|
|
connection.execute(
|
|
"INSERT INTO idempotency(scope,key,payload_digest,response_json,created_at)"
|
|
" VALUES(?,?,?,?,?)",
|
|
(scope, key, payload_digest, canonical_json(dict(response)), iso(self.now())),
|
|
)
|
|
|
|
def create_task(
|
|
self,
|
|
*,
|
|
task_id: str | None,
|
|
goal: str,
|
|
instructions: str,
|
|
success_criteria: object | None,
|
|
required_capabilities: Iterable[object] | None = None,
|
|
constraints: Mapping[str, Any] | None = None,
|
|
metadata: Mapping[str, Any] | None = None,
|
|
deadline: str | None = None,
|
|
priority: int = 0,
|
|
max_attempts: int = 3,
|
|
idempotency_key: str,
|
|
) -> dict[str, Any]:
|
|
task_id = (task_id or ("task-" + uuid.uuid4().hex)).strip()
|
|
goal = goal.strip()
|
|
instructions = instructions.strip()
|
|
idempotency_key = idempotency_key.strip()
|
|
if not task_id or len(task_id) > 180 or not goal or not instructions:
|
|
raise TaskProtocolError("TASK_BINDING_INVALID", 400)
|
|
if not idempotency_key or len(idempotency_key) > 160:
|
|
raise TaskProtocolError("IDEMPOTENCY_KEY_INVALID", 400)
|
|
if not isinstance(priority, int) or isinstance(priority, bool) or not -100 <= priority <= 100:
|
|
raise TaskProtocolError("PRIORITY_INVALID", 400)
|
|
if not isinstance(max_attempts, int) or isinstance(max_attempts, bool) or not 1 <= max_attempts <= 20:
|
|
raise TaskProtocolError("MAX_ATTEMPTS_INVALID", 400)
|
|
parsed_deadline = parse_time(deadline)
|
|
if parsed_deadline and parsed_deadline <= self.now():
|
|
raise TaskProtocolError("DEADLINE_INVALID", 400)
|
|
criteria = _criteria(success_criteria)
|
|
required = _string_list(required_capabilities, "required_capabilities")
|
|
constraints_value = _json_object(constraints, "constraints")
|
|
metadata_value = _json_object(metadata, "metadata")
|
|
_validate_write_execution_binding(task_id, required, metadata_value)
|
|
payload = {
|
|
"task_id": task_id,
|
|
"goal": goal,
|
|
"instructions": instructions,
|
|
"success_criteria": criteria,
|
|
"required_capabilities": required,
|
|
"constraints": constraints_value,
|
|
"metadata": metadata_value,
|
|
"deadline": iso(parsed_deadline) if parsed_deadline else None,
|
|
"priority": priority,
|
|
"max_attempts": max_attempts,
|
|
}
|
|
payload_digest = digest_json(payload)
|
|
now = iso(self.now())
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
replay = self._idempotent_lookup(
|
|
connection,
|
|
scope="create_task",
|
|
key=idempotency_key,
|
|
payload_digest=payload_digest,
|
|
)
|
|
if replay is not None:
|
|
connection.commit()
|
|
return {**replay, "idempotent": True}
|
|
try:
|
|
connection.execute(
|
|
"INSERT INTO tasks(task_id,goal,instructions,success_criteria_json,"
|
|
"required_capabilities_json,constraints_json,metadata_json,deadline,priority,"
|
|
"max_attempts,attempt,status,revision,created_at,updated_at)"
|
|
" VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
|
|
(
|
|
task_id,
|
|
goal,
|
|
instructions,
|
|
canonical_json(criteria),
|
|
canonical_json(required),
|
|
canonical_json(constraints_value),
|
|
canonical_json(metadata_value),
|
|
payload["deadline"],
|
|
priority,
|
|
max_attempts,
|
|
0,
|
|
"QUEUED",
|
|
1,
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
except sqlite3.IntegrityError as exc:
|
|
raise TaskProtocolError("TASK_ID_CONFLICT", 409) from exc
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
result = self._row_task(row)
|
|
self._idempotent_store(
|
|
connection,
|
|
scope="create_task",
|
|
key=idempotency_key,
|
|
payload_digest=payload_digest,
|
|
response=result,
|
|
)
|
|
self._event(connection, "TASK_CREATED", task_id=task_id, payload={"priority": priority})
|
|
connection.commit()
|
|
return result
|
|
|
|
def _expire_locked(self, connection: sqlite3.Connection) -> None:
|
|
now_dt = self.now()
|
|
now = iso(now_dt)
|
|
rows = connection.execute(
|
|
"SELECT * FROM tasks WHERE status IN ('QUEUED','CLAIMED','RUNNING')"
|
|
).fetchall()
|
|
for row in rows:
|
|
deadline = parse_time(row["deadline"])
|
|
if deadline is not None and deadline <= now_dt:
|
|
connection.execute(
|
|
"UPDATE tasks SET status='EXPIRED',lease_token_hash=NULL,"
|
|
"lease_expires_at=NULL,revision=revision+1,updated_at=? WHERE task_id=?",
|
|
(now, row["task_id"]),
|
|
)
|
|
self._event(connection, "TASK_EXPIRED", task_id=row["task_id"])
|
|
continue
|
|
if row["status"] in ACTIVE_LEASE_STATES:
|
|
lease_expires = parse_time(row["lease_expires_at"])
|
|
if lease_expires is not None and lease_expires <= now_dt:
|
|
if row["attempt"] >= row["max_attempts"]:
|
|
next_state = "FAILED"
|
|
else:
|
|
next_state = "QUEUED"
|
|
if next_state == "FAILED":
|
|
connection.execute(
|
|
"UPDATE tasks SET status=?,lease_token_hash=NULL,"
|
|
"lease_expires_at=NULL,revision=revision+1,updated_at=? WHERE task_id=?",
|
|
(next_state, now, row["task_id"]),
|
|
)
|
|
else:
|
|
connection.execute(
|
|
"UPDATE tasks SET status=?,claimed_by=NULL,lease_token_hash=NULL,"
|
|
"lease_expires_at=NULL,revision=revision+1,updated_at=? WHERE task_id=?",
|
|
(next_state, now, row["task_id"]),
|
|
)
|
|
self._event(
|
|
connection,
|
|
"LEASE_EXPIRED",
|
|
task_id=row["task_id"],
|
|
agent_id=row["claimed_by"],
|
|
payload={"next_state": next_state, "attempt": row["attempt"]},
|
|
)
|
|
|
|
def expire_leases(self) -> int:
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
before = connection.total_changes
|
|
self._expire_locked(connection)
|
|
changed = connection.total_changes - before
|
|
connection.commit()
|
|
return changed
|
|
|
|
def poll_tasks(self, *, agent_id: str, limit: int = 20) -> list[dict[str, Any]]:
|
|
if not 1 <= int(limit) <= 100:
|
|
raise TaskProtocolError("LIMIT_INVALID", 400)
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
self._expire_locked(connection)
|
|
agent = connection.execute(
|
|
"SELECT capabilities_json FROM agents WHERE agent_id=? AND status='ACTIVE'",
|
|
(agent_id,),
|
|
).fetchone()
|
|
if agent is None:
|
|
raise TaskProtocolError("AGENT_NOT_FOUND", 404)
|
|
capabilities = json.loads(agent["capabilities_json"])
|
|
rows = connection.execute(
|
|
"SELECT * FROM tasks WHERE status='QUEUED' ORDER BY priority DESC,created_at ASC"
|
|
).fetchall()
|
|
result = []
|
|
for row in rows:
|
|
required = json.loads(row["required_capabilities_json"])
|
|
if has_capabilities(capabilities, required):
|
|
result.append(self._row_task(row, include_result=False))
|
|
if len(result) >= int(limit):
|
|
break
|
|
connection.commit()
|
|
return result
|
|
|
|
def _agent_capabilities(self, connection: sqlite3.Connection, agent_id: str) -> set[str]:
|
|
row = connection.execute(
|
|
"SELECT capabilities_json FROM agents WHERE agent_id=? AND status='ACTIVE'",
|
|
(agent_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
raise TaskProtocolError("AGENT_NOT_FOUND", 404)
|
|
return set(json.loads(row["capabilities_json"]))
|
|
|
|
def claim_task(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
agent_id: str,
|
|
lease_seconds: int = 300,
|
|
) -> dict[str, Any]:
|
|
if not isinstance(lease_seconds, int) or isinstance(lease_seconds, bool) or not 30 <= lease_seconds <= 3600:
|
|
raise TaskProtocolError("LEASE_SECONDS_INVALID", 400)
|
|
lease_token = "lease_" + secrets.token_urlsafe(32)
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
self._expire_locked(connection)
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
if row is None:
|
|
raise TaskProtocolError("TASK_NOT_FOUND", 404)
|
|
if row["status"] != "QUEUED":
|
|
raise TaskProtocolError("TASK_NOT_CLAIMABLE", 409)
|
|
capabilities = self._agent_capabilities(connection, agent_id)
|
|
required = json.loads(row["required_capabilities_json"])
|
|
if not has_capabilities(capabilities, required):
|
|
raise TaskProtocolError("AGENT_CAPABILITY_MISMATCH", 403)
|
|
attempt = int(row["attempt"]) + 1
|
|
if attempt > int(row["max_attempts"]):
|
|
raise TaskProtocolError("TASK_ATTEMPTS_EXHAUSTED", 409)
|
|
expires = self.now() + timedelta(seconds=lease_seconds)
|
|
now = iso(self.now())
|
|
updated = connection.execute(
|
|
"UPDATE tasks SET status='CLAIMED',claimed_by=?,lease_token_hash=?,"
|
|
"lease_expires_at=?,attempt=?,revision=revision+1,updated_at=? "
|
|
"WHERE task_id=? AND status='QUEUED'",
|
|
(
|
|
agent_id,
|
|
self._token_hash(lease_token),
|
|
iso(expires),
|
|
attempt,
|
|
now,
|
|
task_id,
|
|
),
|
|
)
|
|
if updated.rowcount != 1:
|
|
raise TaskProtocolError("TASK_CLAIM_CONFLICT", 409)
|
|
self._event(
|
|
connection,
|
|
"TASK_CLAIMED",
|
|
task_id=task_id,
|
|
agent_id=agent_id,
|
|
payload={"attempt": attempt, "lease_expires_at": iso(expires)},
|
|
)
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
connection.commit()
|
|
return {
|
|
**self._row_task(row, include_result=False),
|
|
"lease_token": lease_token,
|
|
"lease_token_returned_once": True,
|
|
}
|
|
|
|
def _validate_lease(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
*,
|
|
task_id: str,
|
|
agent_id: str,
|
|
lease_token: str,
|
|
allowed_states: frozenset[str] = ACTIVE_LEASE_STATES,
|
|
) -> sqlite3.Row:
|
|
self._expire_locked(connection)
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
if row is None:
|
|
raise TaskProtocolError("TASK_NOT_FOUND", 404)
|
|
if row["status"] not in allowed_states:
|
|
raise TaskProtocolError("TASK_LEASE_NOT_ACTIVE", 409)
|
|
if row["claimed_by"] != agent_id:
|
|
raise TaskProtocolError("TASK_LEASE_OWNER_MISMATCH", 403)
|
|
if not row["lease_token_hash"] or not hmac.compare_digest(
|
|
row["lease_token_hash"], self._token_hash(lease_token)
|
|
):
|
|
raise TaskProtocolError("TASK_LEASE_TOKEN_INVALID", 403)
|
|
expires = parse_time(row["lease_expires_at"])
|
|
if expires is None or expires <= self.now():
|
|
raise TaskProtocolError("TASK_LEASE_EXPIRED", 409)
|
|
return row
|
|
|
|
def heartbeat(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
agent_id: str,
|
|
lease_token: str,
|
|
progress: Mapping[str, Any] | None = None,
|
|
lease_seconds: int = 300,
|
|
) -> dict[str, Any]:
|
|
if not 30 <= int(lease_seconds) <= 3600:
|
|
raise TaskProtocolError("LEASE_SECONDS_INVALID", 400)
|
|
progress_value = _json_object(progress, "progress", max_bytes=64 * 1024)
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
self._validate_lease(
|
|
connection,
|
|
task_id=task_id,
|
|
agent_id=agent_id,
|
|
lease_token=lease_token,
|
|
)
|
|
expires = self.now() + timedelta(seconds=int(lease_seconds))
|
|
now = iso(self.now())
|
|
connection.execute(
|
|
"UPDATE tasks SET status='RUNNING',lease_expires_at=?,revision=revision+1,"
|
|
"updated_at=? WHERE task_id=?",
|
|
(iso(expires), now, task_id),
|
|
)
|
|
self._event(
|
|
connection,
|
|
"TASK_HEARTBEAT",
|
|
task_id=task_id,
|
|
agent_id=agent_id,
|
|
payload={"progress": progress_value, "lease_expires_at": iso(expires)},
|
|
)
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
connection.commit()
|
|
return self._row_task(row, include_result=False)
|
|
|
|
def submit_result(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
agent_id: str,
|
|
lease_token: str,
|
|
idempotency_key: str,
|
|
result: Mapping[str, Any],
|
|
) -> dict[str, Any]:
|
|
idempotency_key = idempotency_key.strip()
|
|
if not idempotency_key or len(idempotency_key) > 160:
|
|
raise TaskProtocolError("IDEMPOTENCY_KEY_INVALID", 400)
|
|
result_value = _json_object(result, "result", max_bytes=1024 * 1024)
|
|
result_digest = digest_json(result_value)
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
if row is None:
|
|
raise TaskProtocolError("TASK_NOT_FOUND", 404)
|
|
if row["result_idempotency_key"] == idempotency_key:
|
|
if not row["result_digest"] or not hmac.compare_digest(row["result_digest"], result_digest):
|
|
raise TaskProtocolError("IDEMPOTENCY_CONFLICT", 409)
|
|
connection.commit()
|
|
return {**self._row_task(row), "idempotent": True}
|
|
self._validate_lease(
|
|
connection,
|
|
task_id=task_id,
|
|
agent_id=agent_id,
|
|
lease_token=lease_token,
|
|
)
|
|
now = iso(self.now())
|
|
connection.execute(
|
|
"UPDATE tasks SET status='VERIFYING',result_json=?,result_digest=?,"
|
|
"result_idempotency_key=?,lease_token_hash=NULL,lease_expires_at=NULL,"
|
|
"revision=revision+1,updated_at=? WHERE task_id=?",
|
|
(
|
|
canonical_json(result_value),
|
|
result_digest,
|
|
idempotency_key,
|
|
now,
|
|
task_id,
|
|
),
|
|
)
|
|
self._event(
|
|
connection,
|
|
"RESULT_SUBMITTED",
|
|
task_id=task_id,
|
|
agent_id=agent_id,
|
|
payload={"result_digest": result_digest, "idempotency_key": idempotency_key},
|
|
)
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
connection.commit()
|
|
return self._row_task(row)
|
|
|
|
def verify_task(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
decision: str,
|
|
verification: Mapping[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
decision = decision.strip().upper()
|
|
if decision not in VERIFY_DECISIONS:
|
|
raise TaskProtocolError("VERIFICATION_DECISION_INVALID", 400)
|
|
verification_value = _json_object(verification, "verification")
|
|
verification_value = {"decision": decision, **verification_value}
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
if row is None:
|
|
raise TaskProtocolError("TASK_NOT_FOUND", 404)
|
|
if row["status"] != "VERIFYING":
|
|
raise TaskProtocolError("TASK_NOT_VERIFYING", 409)
|
|
if decision == "SUCCEEDED":
|
|
next_state = "SUCCEEDED"
|
|
elif decision == "FAILED":
|
|
next_state = "FAILED"
|
|
elif int(row["attempt"]) >= int(row["max_attempts"]):
|
|
next_state = "FAILED"
|
|
verification_value["reason"] = verification_value.get(
|
|
"reason", "MAX_ATTEMPTS_EXHAUSTED"
|
|
)
|
|
else:
|
|
next_state = "QUEUED"
|
|
now = iso(self.now())
|
|
if next_state in {"SUCCEEDED", "FAILED"}:
|
|
connection.execute(
|
|
"UPDATE tasks SET status=?,lease_token_hash=NULL,"
|
|
"lease_expires_at=NULL,verification_json=?,revision=revision+1,updated_at=? "
|
|
"WHERE task_id=?",
|
|
(next_state, canonical_json(verification_value), now, task_id),
|
|
)
|
|
else:
|
|
connection.execute(
|
|
"UPDATE tasks SET status=?,claimed_by=NULL,lease_token_hash=NULL,"
|
|
"lease_expires_at=NULL,verification_json=?,revision=revision+1,updated_at=? "
|
|
"WHERE task_id=?",
|
|
(next_state, canonical_json(verification_value), now, task_id),
|
|
)
|
|
self._event(
|
|
connection,
|
|
"TASK_VERIFIED",
|
|
task_id=task_id,
|
|
agent_id=row["claimed_by"],
|
|
payload={"decision": decision, "next_state": next_state},
|
|
)
|
|
updated = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
connection.commit()
|
|
return self._row_task(updated)
|
|
|
|
def cancel_task(self, *, task_id: str, reason: str = "") -> dict[str, Any]:
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
if row is None:
|
|
raise TaskProtocolError("TASK_NOT_FOUND", 404)
|
|
if row["status"] in TERMINAL_STATES:
|
|
connection.commit()
|
|
return self._row_task(row)
|
|
connection.execute(
|
|
"UPDATE tasks SET status='CANCELLED',lease_token_hash=NULL,"
|
|
"lease_expires_at=NULL,revision=revision+1,updated_at=? WHERE task_id=?",
|
|
(iso(self.now()), task_id),
|
|
)
|
|
self._event(connection, "TASK_CANCELLED", task_id=task_id, payload={"reason": reason})
|
|
updated = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
connection.commit()
|
|
return self._row_task(updated)
|
|
|
|
def get_task(self, task_id: str) -> dict[str, Any]:
|
|
self.expire_leases()
|
|
with self._connect() as connection:
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
if row is None:
|
|
raise TaskProtocolError("TASK_NOT_FOUND", 404)
|
|
return self._row_task(row)
|
|
|
|
def get_events(self, task_id: str, *, limit: int = 200) -> list[dict[str, Any]]:
|
|
with self._connect() as connection:
|
|
rows = connection.execute(
|
|
"SELECT * FROM events WHERE task_id=? ORDER BY event_id DESC LIMIT ?",
|
|
(task_id, max(1, min(int(limit), 1000))),
|
|
).fetchall()
|
|
return [
|
|
{
|
|
"event_id": row["event_id"],
|
|
"task_id": row["task_id"],
|
|
"agent_id": row["agent_id"],
|
|
"event_type": row["event_type"],
|
|
"payload": json.loads(row["payload_json"]),
|
|
"created_at": row["created_at"],
|
|
}
|
|
for row in reversed(rows)
|
|
]
|