428 lines
18 KiB
Python
428 lines
18 KiB
Python
"""Owner-controlled, CAS-bound retry extension for exhausted Online Agent Tasks."""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
import hashlib
|
|
import hmac
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import sqlite3
|
|
from typing import Any, Mapping
|
|
|
|
TASK_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{2,179}$")
|
|
KEY_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{7,159}$")
|
|
MAX_BODY = 64 * 1024
|
|
AUTHORIZE_EXTENSION_ID_ENV = "AIOS_AGENT_TASK_AUTHORIZATION_ID"
|
|
|
|
|
|
class RetryExtensionError(RuntimeError):
|
|
def __init__(self, code: str, status: int = 409) -> None:
|
|
super().__init__(code)
|
|
self.code = code
|
|
self.status = status
|
|
|
|
|
|
def _canonical(value: object) -> bytes:
|
|
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
|
|
|
|
|
def _sha(value: object) -> str:
|
|
return hashlib.sha256(_canonical(value)).hexdigest()
|
|
|
|
|
|
def _required_authorization_id() -> str:
|
|
value = os.environ.get(AUTHORIZE_EXTENSION_ID_ENV, "").strip()
|
|
if not value:
|
|
raise RuntimeError("TASK_RETRY_AUTHORIZATION_ID_MISSING")
|
|
if KEY_RE.fullmatch(value) is None:
|
|
raise RuntimeError("TASK_RETRY_AUTHORIZATION_ID_INVALID")
|
|
return value
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
def _require_string(value: object, code: str, pattern: re.Pattern[str] | None = None) -> str:
|
|
if not isinstance(value, str) or not value.strip():
|
|
raise RetryExtensionError(code, 400)
|
|
value = value.strip()
|
|
if pattern is not None and pattern.fullmatch(value) is None:
|
|
raise RetryExtensionError(code, 400)
|
|
return value
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RetryExtensionRequest:
|
|
expected_status: str
|
|
expected_revision: int
|
|
expected_attempt: int
|
|
expected_max_attempts: int
|
|
next_max_attempts: int
|
|
previous_verification_sha256: str
|
|
authorization_id: str
|
|
idempotency_key: str
|
|
|
|
@classmethod
|
|
def from_mapping(cls, value: Mapping[str, Any]) -> "RetryExtensionRequest":
|
|
allowed = {
|
|
"expected_status", "expected_revision", "expected_attempt", "expected_max_attempts",
|
|
"next_max_attempts", "previous_verification_sha256", "authorization_id", "idempotency_key",
|
|
}
|
|
if set(value) != allowed:
|
|
raise RetryExtensionError("RETRY_EXTENSION_FIELDS_INVALID", 400)
|
|
status = _require_string(value["expected_status"], "EXPECTED_STATUS_INVALID")
|
|
if status != "FAILED":
|
|
raise RetryExtensionError("EXPECTED_STATUS_INVALID", 400)
|
|
ints: dict[str, int] = {}
|
|
for field in ("expected_revision", "expected_attempt", "expected_max_attempts", "next_max_attempts"):
|
|
item = value[field]
|
|
if not isinstance(item, int) or isinstance(item, bool) or item < 0:
|
|
raise RetryExtensionError("RETRY_EXTENSION_INTEGER_INVALID", 400)
|
|
ints[field] = item
|
|
if ints["expected_attempt"] != ints["expected_max_attempts"]:
|
|
raise RetryExtensionError("RETRY_EXTENSION_NOT_EXHAUSTED", 409)
|
|
if ints["next_max_attempts"] != ints["expected_max_attempts"] + 1:
|
|
raise RetryExtensionError("RETRY_EXTENSION_MUST_BE_BOUNDED_ONE_STEP", 400)
|
|
if ints["next_max_attempts"] > 50:
|
|
raise RetryExtensionError("RETRY_EXTENSION_LIMIT_EXCEEDED", 400)
|
|
verification = _require_string(value["previous_verification_sha256"], "VERIFICATION_SHA256_INVALID")
|
|
if re.fullmatch(r"[0-9a-f]{64}", verification) is None:
|
|
raise RetryExtensionError("VERIFICATION_SHA256_INVALID", 400)
|
|
authorization = _require_string(value["authorization_id"], "AUTHORIZATION_ID_INVALID", KEY_RE)
|
|
key = _require_string(value["idempotency_key"], "IDEMPOTENCY_KEY_INVALID", KEY_RE)
|
|
return cls(
|
|
status,
|
|
ints["expected_revision"],
|
|
ints["expected_attempt"],
|
|
ints["expected_max_attempts"],
|
|
ints["next_max_attempts"],
|
|
verification,
|
|
authorization,
|
|
key,
|
|
)
|
|
|
|
def public(self) -> dict[str, Any]:
|
|
return {
|
|
"expected_status": self.expected_status,
|
|
"expected_revision": self.expected_revision,
|
|
"expected_attempt": self.expected_attempt,
|
|
"expected_max_attempts": self.expected_max_attempts,
|
|
"next_max_attempts": self.next_max_attempts,
|
|
"previous_verification_sha256": self.previous_verification_sha256,
|
|
"authorization_id": self.authorization_id,
|
|
"idempotency_key": self.idempotency_key,
|
|
}
|
|
|
|
|
|
class RetryExtensionStore:
|
|
"""Owns the official mutation. Direct/manual DB edits remain forbidden."""
|
|
|
|
def __init__(self, db_path: Path) -> None:
|
|
self.db_path = Path(db_path)
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
connection = sqlite3.connect(self.db_path, timeout=30.0)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA foreign_keys=ON")
|
|
connection.execute("PRAGMA busy_timeout=30000")
|
|
return connection
|
|
|
|
@staticmethod
|
|
def _columns(connection: sqlite3.Connection, table: str) -> dict[str, sqlite3.Row]:
|
|
return {str(row["name"]): row for row in connection.execute(f'PRAGMA table_info("{table}")')}
|
|
|
|
@staticmethod
|
|
def _event_table(connection: sqlite3.Connection) -> tuple[str, dict[str, sqlite3.Row]]:
|
|
tables = {
|
|
row["name"]
|
|
for row in connection.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
|
}
|
|
for name in ("task_events", "events"):
|
|
if name not in tables:
|
|
continue
|
|
columns = RetryExtensionStore._columns(connection, name)
|
|
if "task_id" in columns and ("event_type" in columns or "type" in columns):
|
|
return name, columns
|
|
raise RetryExtensionError("TASK_EVENT_TABLE_UNSUPPORTED", 503)
|
|
|
|
@staticmethod
|
|
def _verification_hash(raw: object) -> str:
|
|
if raw in (None, ""):
|
|
value: object = {}
|
|
elif isinstance(raw, str):
|
|
try:
|
|
value = json.loads(raw)
|
|
except json.JSONDecodeError as exc:
|
|
raise RetryExtensionError("TASK_VERIFICATION_INVALID", 409) from exc
|
|
else:
|
|
value = raw
|
|
return _sha(value)
|
|
|
|
@staticmethod
|
|
def _ensure_extension_table(connection: sqlite3.Connection) -> None:
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS task_retry_extensions (
|
|
task_id TEXT NOT NULL,
|
|
idempotency_key TEXT NOT NULL,
|
|
request_sha256 TEXT NOT NULL,
|
|
authorization_id TEXT NOT NULL,
|
|
previous_verification_sha256 TEXT NOT NULL,
|
|
previous_max_attempts INTEGER NOT NULL,
|
|
next_max_attempts INTEGER NOT NULL,
|
|
expected_revision INTEGER NOT NULL,
|
|
applied_revision INTEGER NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
PRIMARY KEY(task_id, idempotency_key)
|
|
)
|
|
"""
|
|
)
|
|
|
|
@staticmethod
|
|
def _insert_event(
|
|
connection: sqlite3.Connection,
|
|
task_id: str,
|
|
payload: Mapping[str, Any],
|
|
) -> None:
|
|
table, columns = RetryExtensionStore._event_table(connection)
|
|
values: dict[str, Any] = {"task_id": task_id}
|
|
if "event_type" in columns:
|
|
values["event_type"] = "TASK_RETRY_EXTENSION_AUTHORIZED"
|
|
else:
|
|
values["type"] = "TASK_RETRY_EXTENSION_AUTHORIZED"
|
|
payload_text = json.dumps(payload, sort_keys=True, separators=(",", ":"))
|
|
if "payload_json" in columns:
|
|
values["payload_json"] = payload_text
|
|
elif "payload" in columns:
|
|
values["payload"] = payload_text
|
|
else:
|
|
raise RetryExtensionError("TASK_EVENT_PAYLOAD_UNSUPPORTED", 503)
|
|
if "created_at" in columns:
|
|
values["created_at"] = _now()
|
|
if "agent_id" in columns:
|
|
values["agent_id"] = "owner-control"
|
|
if "actor_id" in columns:
|
|
values["actor_id"] = "owner-control"
|
|
names = list(values)
|
|
sql = f'INSERT INTO "{table}" ({",".join(f"""\"{name}\"""" for name in names)}) VALUES ({",".join("?" for _ in names)})'
|
|
connection.execute(sql, tuple(values[name] for name in names))
|
|
|
|
def extend(self, task_id: str, request: RetryExtensionRequest) -> dict[str, Any]:
|
|
task_id = _require_string(task_id, "TASK_ID_INVALID", TASK_RE)
|
|
required_authorization_id = _required_authorization_id()
|
|
if request.authorization_id != required_authorization_id:
|
|
raise RetryExtensionError("RETRY_EXTENSION_AUTHORIZATION_ID_MISMATCH", 403)
|
|
request_payload = request.public()
|
|
request_sha = _sha(request_payload)
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
self._ensure_extension_table(connection)
|
|
replay = connection.execute(
|
|
"SELECT * FROM task_retry_extensions WHERE task_id=? AND idempotency_key=?",
|
|
(task_id, request.idempotency_key),
|
|
).fetchone()
|
|
if replay is not None:
|
|
if replay["request_sha256"] != request_sha:
|
|
raise RetryExtensionError("RETRY_EXTENSION_IDEMPOTENCY_CONFLICT", 409)
|
|
connection.commit()
|
|
return {
|
|
"task_id": task_id,
|
|
"status": "QUEUED",
|
|
"attempt": request.expected_attempt,
|
|
"max_attempts": int(replay["next_max_attempts"]),
|
|
"revision": int(replay["applied_revision"]),
|
|
"idempotent": True,
|
|
"authorization_id": replay["authorization_id"],
|
|
"previous_verification_sha256": replay["previous_verification_sha256"],
|
|
}
|
|
|
|
columns = self._columns(connection, "tasks")
|
|
required = {
|
|
"task_id", "status", "attempt", "max_attempts", "revision", "claimed_by",
|
|
"lease_token_hash", "lease_expires_at", "verification_json",
|
|
}
|
|
if not required.issubset(columns):
|
|
raise RetryExtensionError("TASK_SCHEMA_RETRY_EXTENSION_UNSUPPORTED", 503)
|
|
task = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
if task is None:
|
|
raise RetryExtensionError("TASK_NOT_FOUND", 404)
|
|
actual_verification_sha = self._verification_hash(task["verification_json"])
|
|
if actual_verification_sha != request.previous_verification_sha256:
|
|
raise RetryExtensionError("RETRY_EXTENSION_VERIFICATION_DRIFT", 409)
|
|
if (
|
|
task["status"] != request.expected_status
|
|
or int(task["revision"]) != request.expected_revision
|
|
or int(task["attempt"]) != request.expected_attempt
|
|
or int(task["max_attempts"]) != request.expected_max_attempts
|
|
):
|
|
raise RetryExtensionError("RETRY_EXTENSION_CAS_MISMATCH", 409)
|
|
for lease_field in ("lease_token_hash", "lease_expires_at"):
|
|
if lease_field in columns and task[lease_field] not in (None, ""):
|
|
raise RetryExtensionError("RETRY_EXTENSION_ACTIVE_LEASE_DENY", 409)
|
|
|
|
previous_claimed_by = task["claimed_by"]
|
|
applied_revision = request.expected_revision + 1
|
|
cursor = connection.execute(
|
|
"""
|
|
UPDATE tasks
|
|
SET status='QUEUED', max_attempts=?, revision=?, claimed_by=NULL,
|
|
lease_token_hash=NULL, lease_expires_at=NULL
|
|
WHERE task_id=? AND status=? AND revision=? AND attempt=? AND max_attempts=?
|
|
AND claimed_by IS ?
|
|
AND lease_token_hash IS NULL AND lease_expires_at IS NULL
|
|
""",
|
|
(
|
|
request.next_max_attempts,
|
|
applied_revision,
|
|
task_id,
|
|
request.expected_status,
|
|
request.expected_revision,
|
|
request.expected_attempt,
|
|
request.expected_max_attempts,
|
|
previous_claimed_by,
|
|
),
|
|
)
|
|
if cursor.rowcount != 1:
|
|
raise RetryExtensionError("RETRY_EXTENSION_CAS_MISMATCH", 409)
|
|
|
|
event_payload = {
|
|
"schema": "aios.task-retry-extension.event.v1",
|
|
"task_id": task_id,
|
|
"authorization_id": request.authorization_id,
|
|
"idempotency_key": request.idempotency_key,
|
|
"previous_verification_sha256": request.previous_verification_sha256,
|
|
"previous_max_attempts": request.expected_max_attempts,
|
|
"next_max_attempts": request.next_max_attempts,
|
|
"previous_revision": request.expected_revision,
|
|
"applied_revision": applied_revision,
|
|
"request_sha256": request_sha,
|
|
"previous_claimed_by": previous_claimed_by,
|
|
"stale_terminal_claim_normalized": previous_claimed_by not in (None, ""),
|
|
}
|
|
self._insert_event(connection, task_id, event_payload)
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO task_retry_extensions(
|
|
task_id,idempotency_key,request_sha256,authorization_id,previous_verification_sha256,
|
|
previous_max_attempts,next_max_attempts,expected_revision,applied_revision,created_at
|
|
) VALUES (?,?,?,?,?,?,?,?,?,?)
|
|
""",
|
|
(
|
|
task_id,
|
|
request.idempotency_key,
|
|
request_sha,
|
|
request.authorization_id,
|
|
request.previous_verification_sha256,
|
|
request.expected_max_attempts,
|
|
request.next_max_attempts,
|
|
request.expected_revision,
|
|
applied_revision,
|
|
_now(),
|
|
),
|
|
)
|
|
connection.commit()
|
|
return {
|
|
"task_id": task_id,
|
|
"status": "QUEUED",
|
|
"attempt": request.expected_attempt,
|
|
"max_attempts": request.next_max_attempts,
|
|
"revision": applied_revision,
|
|
"idempotent": False,
|
|
"authorization_id": request.authorization_id,
|
|
"previous_verification_sha256": request.previous_verification_sha256,
|
|
}
|
|
|
|
|
|
class RetryExtensionHandler(BaseHTTPRequestHandler):
|
|
store: RetryExtensionStore | None = None
|
|
control_token: str = ""
|
|
|
|
def log_message(self, *_args: Any) -> None:
|
|
return
|
|
|
|
def _json(self, value: Mapping[str, Any], status: int) -> None:
|
|
data = json.dumps(dict(value), sort_keys=True, separators=(",", ":")).encode("utf-8")
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(data)))
|
|
self.end_headers()
|
|
self.wfile.write(data)
|
|
|
|
def _authorized(self) -> bool:
|
|
header = self.headers.get("Authorization", "")
|
|
return (
|
|
header.startswith("Bearer ")
|
|
and bool(self.control_token)
|
|
and hmac.compare_digest(header[7:], self.control_token)
|
|
)
|
|
|
|
def do_GET(self) -> None:
|
|
if self.path == "/health":
|
|
self._json({"status": "ok", "service": "aios-agent-task-retry-extension"}, 200)
|
|
else:
|
|
self._json({"error": "NOT_FOUND"}, 404)
|
|
|
|
def do_POST(self) -> None:
|
|
match = re.fullmatch(r"/v1/tasks/([^/]+)/retry-extension", self.path)
|
|
if match is None:
|
|
self._json({"error": "NOT_FOUND"}, 404)
|
|
return
|
|
if not self._authorized():
|
|
self._json({"error": "UNAUTHORIZED"}, 401)
|
|
return
|
|
try:
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
except ValueError:
|
|
self._json({"error": "CONTENT_LENGTH_INVALID"}, 400)
|
|
return
|
|
if length < 1 or length > MAX_BODY:
|
|
self._json({"error": "REQUEST_SIZE_INVALID"}, 413)
|
|
return
|
|
try:
|
|
value = json.loads(self.rfile.read(length).decode("utf-8"))
|
|
if not isinstance(value, Mapping):
|
|
raise ValueError
|
|
request = RetryExtensionRequest.from_mapping(value)
|
|
if self.store is None:
|
|
raise RetryExtensionError("RUNTIME_UNAVAILABLE", 503)
|
|
result = self.store.extend(match.group(1), request)
|
|
self._json(result, 200 if result["idempotent"] else 201)
|
|
except RetryExtensionError as exc:
|
|
self._json({"error": exc.code}, exc.status)
|
|
except (UnicodeDecodeError, json.JSONDecodeError, ValueError):
|
|
self._json({"error": "REQUEST_INVALID"}, 400)
|
|
|
|
|
|
def load_control_token(path: Path) -> str:
|
|
if path.is_symlink() or not path.is_file():
|
|
raise RuntimeError("TASK_RETRY_CONTROL_CREDENTIAL_UNAVAILABLE")
|
|
value = path.read_text(encoding="utf-8").strip()
|
|
if len(value) < 32:
|
|
raise RuntimeError("TASK_RETRY_CONTROL_CREDENTIAL_INVALID")
|
|
return value
|
|
|
|
|
|
def serve() -> None:
|
|
db = Path(os.environ.get("AIOS_AGENT_TASK_DB", "/var/lib/aios-agent-task/state.sqlite3"))
|
|
_required_authorization_id()
|
|
token_path = Path(
|
|
os.environ.get(
|
|
"AIOS_AGENT_TASK_CONTROL_TOKEN_FILE",
|
|
str(Path(os.environ.get("CREDENTIALS_DIRECTORY", "/run/credentials")) / "agent_task_control_token"),
|
|
)
|
|
)
|
|
host = os.environ.get("AIOS_TASK_RETRY_HOST", "127.0.0.1")
|
|
port = int(os.environ.get("AIOS_TASK_RETRY_PORT", "18941"))
|
|
RetryExtensionHandler.store = RetryExtensionStore(db)
|
|
RetryExtensionHandler.control_token = load_control_token(token_path)
|
|
ThreadingHTTPServer((host, port), RetryExtensionHandler).serve_forever()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
serve()
|