188 lines
7.4 KiB
Python
188 lines
7.4 KiB
Python
"""Durable, authority-bound deterministic execution plans for Online Agent Tasks."""
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, Mapping
|
|
|
|
from aios_online_agent_execution import (
|
|
GovernedExecutionBinding,
|
|
project_execution_envelope,
|
|
)
|
|
|
|
|
|
class ExecutionPlanError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _canonical(value: object) -> bytes:
|
|
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
|
|
|
|
|
def _digest(value: object) -> str:
|
|
return hashlib.sha256(_canonical(value)).hexdigest()
|
|
|
|
|
|
def _task_key(task_id: str) -> str:
|
|
readable = "".join(ch if ch.isalnum() or ch in "._-" else "-" for ch in task_id).strip("-")[:72] or "task"
|
|
return readable + "-" + hashlib.sha256(task_id.encode("utf-8")).hexdigest()[:24]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExecutionPlanRecord:
|
|
task_id: str
|
|
job_id: str
|
|
source_head_oid: str
|
|
generation: int
|
|
idempotency_key: str
|
|
plan_digest: str
|
|
execution: dict[str, Any]
|
|
|
|
def public(self) -> dict[str, Any]:
|
|
return {
|
|
"schema": "aios.online-agent.plan-record.v1",
|
|
"task_id": self.task_id,
|
|
"job_id": self.job_id,
|
|
"source_head_oid": self.source_head_oid,
|
|
"generation": self.generation,
|
|
"idempotency_key": self.idempotency_key,
|
|
"plan_digest": self.plan_digest,
|
|
"execution": dict(self.execution),
|
|
}
|
|
|
|
|
|
class ExecutionPlanStore:
|
|
def __init__(self, root: Path) -> None:
|
|
self.root = Path(root)
|
|
self.root.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _path(self, task_id: str) -> Path:
|
|
return self.root / (_task_key(task_id) + ".json")
|
|
|
|
@contextmanager
|
|
def _lock(self, task_id: str):
|
|
path = self.root / (_task_key(task_id) + ".lock")
|
|
with path.open("a+b") as handle:
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
|
|
try:
|
|
yield
|
|
finally:
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
|
|
|
@staticmethod
|
|
def _source_head(binding: GovernedExecutionBinding) -> str:
|
|
return binding.expected_branch_oid if binding.branch_mode == "CONTINUE_EXISTING" else binding.base_oid
|
|
|
|
def _read(self, task_id: str) -> ExecutionPlanRecord | None:
|
|
path = self._path(task_id)
|
|
if not path.is_file() or path.is_symlink():
|
|
return None
|
|
value = json.loads(path.read_text(encoding="utf-8"))
|
|
if not isinstance(value, Mapping) or value.get("schema") != "aios.online-agent.plan-record.v1":
|
|
raise ExecutionPlanError("EXECUTION_PLAN_RECORD_INVALID")
|
|
execution = value.get("execution")
|
|
if not isinstance(execution, Mapping):
|
|
raise ExecutionPlanError("EXECUTION_PLAN_RECORD_INVALID")
|
|
return ExecutionPlanRecord(
|
|
task_id=str(value.get("task_id") or ""),
|
|
job_id=str(value.get("job_id") or ""),
|
|
source_head_oid=str(value.get("source_head_oid") or ""),
|
|
generation=int(value.get("generation") or 0),
|
|
idempotency_key=str(value.get("idempotency_key") or ""),
|
|
plan_digest=str(value.get("plan_digest") or ""),
|
|
execution=dict(execution),
|
|
)
|
|
|
|
def register(
|
|
self,
|
|
binding: GovernedExecutionBinding,
|
|
request: Mapping[str, Any],
|
|
) -> dict[str, Any]:
|
|
with self._lock(binding.task_id):
|
|
return self._register_locked(binding, request)
|
|
|
|
def _register_locked(
|
|
self,
|
|
binding: GovernedExecutionBinding,
|
|
request: Mapping[str, Any],
|
|
) -> dict[str, Any]:
|
|
if request.get("schema") != "aios.online-agent.plan.v1":
|
|
raise ExecutionPlanError("EXECUTION_PLAN_SCHEMA_INVALID")
|
|
if request.get("task_id") != binding.task_id or request.get("job_id") != binding.job_id:
|
|
raise ExecutionPlanError("EXECUTION_PLAN_BINDING_CONFLICT")
|
|
if int(request.get("generation") or 0) != binding.generation:
|
|
raise ExecutionPlanError("EXECUTION_PLAN_GENERATION_CONFLICT")
|
|
source_head = str(request.get("source_head_oid") or "")
|
|
if source_head != self._source_head(binding):
|
|
raise ExecutionPlanError("EXECUTION_PLAN_SOURCE_HEAD_CONFLICT")
|
|
key = str(request.get("idempotency_key") or "")
|
|
if len(key) < 8:
|
|
raise ExecutionPlanError("EXECUTION_PLAN_IDEMPOTENCY_INVALID")
|
|
execution_request = request.get("execution")
|
|
if not isinstance(execution_request, Mapping) or execution_request.get("action") != "deterministic":
|
|
raise ExecutionPlanError("EXECUTION_PLAN_DETERMINISTIC_REQUIRED")
|
|
envelope = project_execution_envelope(binding, execution_request)
|
|
material = {
|
|
"task_id": binding.task_id,
|
|
"job_id": binding.job_id,
|
|
"source_head_oid": source_head,
|
|
"generation": binding.generation,
|
|
"idempotency_key": key,
|
|
"execution": dict(execution_request),
|
|
}
|
|
plan_digest = _digest(material)
|
|
existing = self._read(binding.task_id)
|
|
expected_previous = request.get("expected_previous_plan_digest")
|
|
if existing is not None:
|
|
if existing.idempotency_key == key:
|
|
if existing.plan_digest != plan_digest:
|
|
raise ExecutionPlanError("EXECUTION_PLAN_IDEMPOTENCY_CONFLICT")
|
|
return {**existing.public(), "idempotent": True, "execution_digest": envelope["execution_digest"]}
|
|
if expected_previous != existing.plan_digest:
|
|
raise ExecutionPlanError("EXECUTION_PLAN_REPLACEMENT_CAS_MISMATCH")
|
|
if existing.source_head_oid != source_head:
|
|
raise ExecutionPlanError("EXECUTION_PLAN_STALE_REPLACEMENT_DENY")
|
|
elif expected_previous not in (None, ""):
|
|
raise ExecutionPlanError("EXECUTION_PLAN_REPLACEMENT_CAS_MISMATCH")
|
|
|
|
record = ExecutionPlanRecord(
|
|
task_id=binding.task_id,
|
|
job_id=binding.job_id,
|
|
source_head_oid=source_head,
|
|
generation=binding.generation,
|
|
idempotency_key=key,
|
|
plan_digest=plan_digest,
|
|
execution=dict(execution_request),
|
|
)
|
|
path = self._path(binding.task_id)
|
|
temporary = path.with_suffix(".tmp")
|
|
temporary.write_text(
|
|
json.dumps(record.public(), sort_keys=True, separators=(",", ":")),
|
|
encoding="utf-8",
|
|
)
|
|
temporary.replace(path)
|
|
return {**record.public(), "idempotent": False, "execution_digest": envelope["execution_digest"]}
|
|
|
|
def resolve(self, binding: GovernedExecutionBinding) -> dict[str, Any] | None:
|
|
record = self._read(binding.task_id)
|
|
if record is None:
|
|
return None
|
|
if (
|
|
record.task_id != binding.task_id
|
|
or record.job_id != binding.job_id
|
|
or record.generation != binding.generation
|
|
):
|
|
raise ExecutionPlanError("EXECUTION_PLAN_BINDING_CONFLICT")
|
|
if record.source_head_oid != self._source_head(binding):
|
|
raise ExecutionPlanError("EXECUTION_PLAN_SOURCE_HEAD_DRIFT")
|
|
envelope = project_execution_envelope(binding, record.execution)
|
|
return {
|
|
"execution": dict(record.execution),
|
|
"plan_digest": record.plan_digest,
|
|
"execution_digest": envelope["execution_digest"],
|
|
}
|