Files
aios/active-runtime/app/aios_v2/task_execution_registry.py
T

608 lines
23 KiB
Python

"""Durable execution truth registry for externally executed AIOS tasks.
A task manifest being present in Git is not evidence that an executor discovered
or ran it. This registry separates declaration, registration, claim, start,
heartbeat and terminal evidence. Release gates must require a terminal success
record with immutable evidence hashes.
"""
from __future__ import annotations
import json
import sqlite3
import threading
import uuid
from dataclasses import asdict, dataclass, field, replace
from datetime import datetime, timedelta, timezone
from enum import Enum
from pathlib import Path
from typing import Any, Mapping, Sequence
def utc_now() -> str:
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
def parse_time(value: str) -> datetime:
normalized = str(value or "").strip()
if normalized.endswith("Z"):
normalized = normalized[:-1] + "+00:00"
parsed = datetime.fromisoformat(normalized)
if parsed.tzinfo is None:
raise ValueError("timestamp requires timezone")
return parsed.astimezone(timezone.utc)
def required(value: str, label: str) -> str:
normalized = str(value or "").strip()
if not normalized:
raise ValueError(f"{label} is required")
return normalized
class ExecutionTruthState(str, Enum):
DECLARED = "DECLARED"
REGISTERED = "REGISTERED"
CLAIMED = "CLAIMED"
RUNNING = "RUNNING"
SUCCEEDED = "SUCCEEDED"
FAILED = "FAILED"
CANCELLED = "CANCELLED"
NOT_STARTED = "NOT_STARTED"
STALLED = "STALLED"
TERMINAL_STATES = frozenset(
{
ExecutionTruthState.SUCCEEDED,
ExecutionTruthState.FAILED,
ExecutionTruthState.CANCELLED,
ExecutionTruthState.NOT_STARTED,
ExecutionTruthState.STALLED,
}
)
@dataclass(frozen=True)
class ExecutionTaskManifest:
task_id: str
executor: str
source_branch: str
source_oid: str
target_version: str
steps: Sequence[str]
start_deadline_seconds: int = 300
heartbeat_timeout_seconds: int = 180
required_evidence: Sequence[str] = ()
notification_events: Sequence[str] = ("STARTED", "SUCCEEDED", "FAILED")
owner_notification_ref: str = "AIOS_OWNER_WECHAT_TARGET_REF"
manifest_version: str = "aios.external-task.v1"
def __post_init__(self) -> None:
for name in (
"task_id",
"executor",
"source_branch",
"source_oid",
"target_version",
"manifest_version",
"owner_notification_ref",
):
object.__setattr__(self, name, required(getattr(self, name), name))
steps = tuple(required(item, "step") for item in self.steps)
if not steps:
raise ValueError("at least one execution step is required")
if self.start_deadline_seconds < 1 or self.heartbeat_timeout_seconds < 1:
raise ValueError("execution deadlines must be positive")
object.__setattr__(self, "steps", steps)
object.__setattr__(
self,
"required_evidence",
tuple(sorted({required(item, "required_evidence") for item in self.required_evidence})),
)
object.__setattr__(
self,
"notification_events",
tuple(sorted({required(item, "notification_event") for item in self.notification_events})),
)
def as_dict(self) -> dict[str, Any]:
value = asdict(self)
value["steps"] = list(self.steps)
value["required_evidence"] = list(self.required_evidence)
value["notification_events"] = list(self.notification_events)
return value
@classmethod
def from_dict(cls, value: Mapping[str, Any]) -> "ExecutionTaskManifest":
return cls(**dict(value))
@dataclass(frozen=True)
class ExecutionTruthRecord:
task_id: str
manifest_json: str
state: ExecutionTruthState
executor_run_id: str | None = None
executor_id: str | None = None
declared_at: str = field(default_factory=utc_now)
registered_at: str | None = None
claimed_at: str | None = None
started_at: str | None = None
heartbeat_at: str | None = None
completed_at: str | None = None
current_step: int = 0
evidence: Mapping[str, str] = field(default_factory=dict)
error: Mapping[str, Any] = field(default_factory=dict)
revision: int = 1
@property
def manifest(self) -> ExecutionTaskManifest:
return ExecutionTaskManifest.from_dict(json.loads(self.manifest_json))
def as_dict(self) -> dict[str, Any]:
return {
**asdict(self),
"state": self.state.value,
"manifest": self.manifest.as_dict(),
"evidence": dict(self.evidence),
"error": dict(self.error),
}
class ExecutionTruthConflict(RuntimeError):
pass
class SQLiteExecutionTruthRegistry:
"""Persistent registration and run-evidence store."""
def __init__(self, path: str | Path) -> None:
self.path = Path(path).resolve()
self.path.parent.mkdir(parents=True, exist_ok=True)
self._lock = threading.RLock()
with self._connect() as connection:
connection.executescript(
"""
CREATE TABLE IF NOT EXISTS execution_truth (
task_id TEXT PRIMARY KEY,
manifest_json TEXT NOT NULL,
state TEXT NOT NULL,
record_json TEXT NOT NULL,
revision INTEGER NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS execution_truth_events (
event_id TEXT PRIMARY KEY,
task_id TEXT NOT NULL,
event_type TEXT NOT NULL,
details_json TEXT NOT NULL,
created_at TEXT NOT NULL,
FOREIGN KEY(task_id) REFERENCES execution_truth(task_id)
);
CREATE INDEX IF NOT EXISTS idx_execution_truth_state
ON execution_truth(state, updated_at);
"""
)
def _connect(self) -> sqlite3.Connection:
connection = sqlite3.connect(str(self.path), timeout=30.0)
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA journal_mode=WAL")
connection.execute("PRAGMA synchronous=FULL")
connection.execute("PRAGMA foreign_keys=ON")
return connection
@staticmethod
def _manifest_json(manifest: ExecutionTaskManifest) -> str:
return json.dumps(manifest.as_dict(), ensure_ascii=False, sort_keys=True, separators=(",", ":"))
@staticmethod
def _record_json(record: ExecutionTruthRecord) -> str:
value = record.as_dict()
value.pop("manifest", None)
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str)
@staticmethod
def _decode(row: sqlite3.Row) -> ExecutionTruthRecord:
value = json.loads(row["record_json"])
value["state"] = ExecutionTruthState(value["state"])
value["evidence"] = dict(value.get("evidence") or {})
value["error"] = dict(value.get("error") or {})
return ExecutionTruthRecord(**value)
def _event(
self,
connection: sqlite3.Connection,
task_id: str,
event_type: str,
details: Mapping[str, Any] | None = None,
) -> None:
connection.execute(
"INSERT INTO execution_truth_events VALUES (?, ?, ?, ?, ?)",
(
f"truth_evt_{uuid.uuid4().hex}",
task_id,
event_type,
json.dumps(dict(details or {}), ensure_ascii=False, sort_keys=True, default=str),
utc_now(),
),
)
def declare(self, manifest: ExecutionTaskManifest) -> ExecutionTruthRecord:
manifest_json = self._manifest_json(manifest)
record = ExecutionTruthRecord(
task_id=manifest.task_id,
manifest_json=manifest_json,
state=ExecutionTruthState.DECLARED,
)
with self._lock, self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing = connection.execute(
"SELECT record_json FROM execution_truth WHERE task_id = ?",
(manifest.task_id,),
).fetchone()
if existing is not None:
current = self._decode(existing)
if current.manifest_json != manifest_json:
previous_manifest = current.manifest.as_dict()
replacement_manifest = manifest.as_dict()
previous_oid = previous_manifest.pop("source_oid")
replacement_oid = replacement_manifest.pop("source_oid")
if (
current.state != ExecutionTruthState.NOT_STARTED
or previous_manifest != replacement_manifest
):
connection.rollback()
raise ExecutionTruthConflict(
"task_id already declared with another manifest"
)
now = utc_now()
updated = replace(
current,
manifest_json=manifest_json,
revision=current.revision + 1,
)
connection.execute(
"UPDATE execution_truth SET manifest_json = ?, record_json = ?, revision = ?, updated_at = ? WHERE task_id = ?",
(
manifest_json,
self._record_json(updated),
updated.revision,
now,
updated.task_id,
),
)
self._event(
connection,
updated.task_id,
"MANIFEST_SUPERSEDED",
{
"previous_source_oid": previous_oid,
"source_oid": replacement_oid,
"state": updated.state.value,
},
)
connection.commit()
return updated
connection.rollback()
return current
connection.execute(
"INSERT INTO execution_truth VALUES (?, ?, ?, ?, ?, ?)",
(
record.task_id,
manifest_json,
record.state.value,
self._record_json(record),
record.revision,
record.declared_at,
),
)
self._event(connection, record.task_id, "DECLARED", {"source_oid": manifest.source_oid})
connection.commit()
return record
def get(self, task_id: str) -> ExecutionTruthRecord:
with self._lock, self._connect() as connection:
row = connection.execute(
"SELECT record_json FROM execution_truth WHERE task_id = ?",
(required(task_id, "task_id"),),
).fetchone()
if row is None:
raise KeyError("execution task not declared")
return self._decode(row)
def _transition(
self,
task_id: str,
*,
allowed: Sequence[ExecutionTruthState],
target: ExecutionTruthState,
executor_run_id: str | None = None,
executor_id: str | None = None,
current_step: int | None = None,
evidence: Mapping[str, str] | None = None,
error: Mapping[str, Any] | None = None,
) -> ExecutionTruthRecord:
with self._lock, self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
row = connection.execute(
"SELECT record_json FROM execution_truth WHERE task_id = ?",
(required(task_id, "task_id"),),
).fetchone()
if row is None:
connection.rollback()
raise KeyError("execution task not declared")
current = self._decode(row)
if current.state not in set(allowed):
connection.rollback()
raise ExecutionTruthConflict(
f"cannot transition {current.state.value} to {target.value}"
)
now = utc_now()
updates: dict[str, Any] = {
"state": target,
"revision": current.revision + 1,
}
if executor_run_id is not None:
if current.executor_run_id and current.executor_run_id != executor_run_id:
connection.rollback()
raise ExecutionTruthConflict("executor_run_id mismatch")
updates["executor_run_id"] = required(executor_run_id, "executor_run_id")
if executor_id is not None:
if current.executor_id and current.executor_id != executor_id:
connection.rollback()
raise ExecutionTruthConflict("executor_id mismatch")
updates["executor_id"] = required(executor_id, "executor_id")
if target == ExecutionTruthState.REGISTERED:
updates["registered_at"] = now
if target == ExecutionTruthState.CLAIMED:
updates["claimed_at"] = now
updates["heartbeat_at"] = now
if target == ExecutionTruthState.RUNNING:
updates["started_at"] = current.started_at or now
updates["heartbeat_at"] = now
if current_step is not None:
if current_step < current.current_step or current_step > len(current.manifest.steps):
connection.rollback()
raise ExecutionTruthConflict("invalid task step progression")
updates["current_step"] = current_step
updates["heartbeat_at"] = now
if evidence is not None:
updates["evidence"] = {
**dict(current.evidence),
**{required(key, "evidence key"): required(value, "evidence hash") for key, value in evidence.items()},
}
if error is not None:
updates["error"] = json.loads(json.dumps(dict(error), default=str))
if target in TERMINAL_STATES:
updates["completed_at"] = now
updated = replace(current, **updates)
connection.execute(
"UPDATE execution_truth SET state = ?, record_json = ?, revision = ?, updated_at = ? WHERE task_id = ?",
(
updated.state.value,
self._record_json(updated),
updated.revision,
now,
updated.task_id,
),
)
self._event(
connection,
updated.task_id,
target.value,
{
"executor_run_id": updated.executor_run_id,
"executor_id": updated.executor_id,
"current_step": updated.current_step,
},
)
connection.commit()
return updated
def register(self, task_id: str, *, executor_run_id: str) -> ExecutionTruthRecord:
return self._transition(
task_id,
allowed=(ExecutionTruthState.DECLARED,),
target=ExecutionTruthState.REGISTERED,
executor_run_id=executor_run_id,
)
def retry_not_started(
self,
task_id: str,
*,
executor_run_id: str,
) -> ExecutionTruthRecord:
"""Start a new audited attempt after the declaration deadline expired."""
with self._lock, self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
row = connection.execute(
"SELECT record_json FROM execution_truth WHERE task_id = ?",
(required(task_id, "task_id"),),
).fetchone()
if row is None:
connection.rollback()
raise KeyError("execution task not declared")
current = self._decode(row)
if current.state != ExecutionTruthState.NOT_STARTED:
connection.rollback()
raise ExecutionTruthConflict(
f"cannot retry task from {current.state.value}"
)
now = utc_now()
updated = replace(
current,
state=ExecutionTruthState.REGISTERED,
executor_run_id=required(executor_run_id, "executor_run_id"),
executor_id=None,
registered_at=now,
claimed_at=None,
started_at=None,
heartbeat_at=None,
completed_at=None,
current_step=0,
evidence={},
error={},
revision=current.revision + 1,
)
connection.execute(
"UPDATE execution_truth SET state = ?, record_json = ?, revision = ?, updated_at = ? WHERE task_id = ?",
(
updated.state.value,
self._record_json(updated),
updated.revision,
now,
updated.task_id,
),
)
self._event(
connection,
updated.task_id,
"RETRIED",
{
"previous_state": current.state.value,
"previous_executor_run_id": current.executor_run_id,
"executor_run_id": updated.executor_run_id,
"previous_error": dict(current.error),
},
)
connection.commit()
return updated
def claim(self, task_id: str, *, executor_run_id: str, executor_id: str) -> ExecutionTruthRecord:
return self._transition(
task_id,
allowed=(ExecutionTruthState.REGISTERED,),
target=ExecutionTruthState.CLAIMED,
executor_run_id=executor_run_id,
executor_id=executor_id,
)
def start(self, task_id: str, *, executor_run_id: str, executor_id: str) -> ExecutionTruthRecord:
return self._transition(
task_id,
allowed=(ExecutionTruthState.CLAIMED,),
target=ExecutionTruthState.RUNNING,
executor_run_id=executor_run_id,
executor_id=executor_id,
current_step=1,
)
def heartbeat(
self,
task_id: str,
*,
executor_run_id: str,
executor_id: str,
current_step: int,
) -> ExecutionTruthRecord:
return self._transition(
task_id,
allowed=(ExecutionTruthState.RUNNING,),
target=ExecutionTruthState.RUNNING,
executor_run_id=executor_run_id,
executor_id=executor_id,
current_step=current_step,
)
def succeed(
self,
task_id: str,
*,
executor_run_id: str,
executor_id: str,
evidence: Mapping[str, str],
) -> ExecutionTruthRecord:
current = self.get(task_id)
missing = sorted(set(current.manifest.required_evidence) - set(evidence))
if missing:
raise ExecutionTruthConflict(f"missing required evidence: {missing}")
if current.current_step != len(current.manifest.steps):
raise ExecutionTruthConflict("all declared steps must complete before success")
return self._transition(
task_id,
allowed=(ExecutionTruthState.RUNNING,),
target=ExecutionTruthState.SUCCEEDED,
executor_run_id=executor_run_id,
executor_id=executor_id,
evidence=evidence,
)
def fail(
self,
task_id: str,
*,
executor_run_id: str,
executor_id: str,
error: Mapping[str, Any],
evidence: Mapping[str, str] | None = None,
) -> ExecutionTruthRecord:
return self._transition(
task_id,
allowed=(ExecutionTruthState.REGISTERED, ExecutionTruthState.CLAIMED, ExecutionTruthState.RUNNING),
target=ExecutionTruthState.FAILED,
executor_run_id=executor_run_id,
executor_id=executor_id,
evidence=evidence,
error=error,
)
def evaluate_deadlines(self, task_id: str, *, now: datetime | None = None) -> ExecutionTruthRecord:
current = self.get(task_id)
current_time = (now or datetime.now(timezone.utc)).astimezone(timezone.utc)
manifest = current.manifest
if current.state in TERMINAL_STATES:
return current
declared = parse_time(current.declared_at)
if current.state in {ExecutionTruthState.DECLARED, ExecutionTruthState.REGISTERED}:
if current_time >= declared + timedelta(seconds=manifest.start_deadline_seconds):
return self._transition(
task_id,
allowed=(current.state,),
target=ExecutionTruthState.NOT_STARTED,
error={"code": "start_deadline_exceeded"},
)
if current.state in {ExecutionTruthState.CLAIMED, ExecutionTruthState.RUNNING}:
heartbeat = parse_time(current.heartbeat_at or current.claimed_at or current.declared_at)
if current_time >= heartbeat + timedelta(seconds=manifest.heartbeat_timeout_seconds):
return self._transition(
task_id,
allowed=(current.state,),
target=ExecutionTruthState.STALLED,
error={"code": "heartbeat_timeout"},
)
return current
def gate(self, task_id: str) -> Mapping[str, Any]:
record = self.evaluate_deadlines(task_id)
passed = record.state == ExecutionTruthState.SUCCEEDED
return {
"task_id": record.task_id,
"state": record.state.value,
"passed": passed,
"release_blocking": not passed,
"executor_run_id": record.executor_run_id,
"executor_id": record.executor_id,
"current_step": record.current_step,
"total_steps": len(record.manifest.steps),
"evidence": dict(record.evidence),
"error": dict(record.error),
}
def events(self, task_id: str) -> list[Mapping[str, Any]]:
with self._lock, self._connect() as connection:
rows = connection.execute(
"SELECT event_id, event_type, details_json, created_at FROM execution_truth_events WHERE task_id = ? ORDER BY created_at, event_id",
(required(task_id, "task_id"),),
).fetchall()
return [
{
"event_id": row["event_id"],
"event_type": row["event_type"],
"details": json.loads(row["details_json"]),
"created_at": row["created_at"],
}
for row in rows
]