271 lines
14 KiB
Python
271 lines
14 KiB
Python
"""Durable governed task binding for CREATE_NEW and CONTINUE_EXISTING."""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, asdict
|
|
from contextlib import contextmanager
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import re
|
|
from typing import Any, Mapping
|
|
|
|
from aios_gitops_job_controller import GitOpsJobController, JobStore, sha256_json
|
|
from aios_online_agent_execution import (
|
|
GovernedExecutionBinding,
|
|
project_execution_envelope,
|
|
project_online_task,
|
|
)
|
|
|
|
OID_RE = re.compile(r"^[0-9a-f]{40}$")
|
|
BRANCH_RE = re.compile(r"^agent/[A-Za-z0-9][A-Za-z0-9._/-]{2,179}$")
|
|
|
|
|
|
class BootstrapError(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class GovernedTaskBinding:
|
|
task_id: str
|
|
repository: str
|
|
base_oid: str
|
|
target_branch: str
|
|
risk: str
|
|
description: str
|
|
idempotency_key: str
|
|
expected_branch_oid: str
|
|
branch_mode: str
|
|
generation: int
|
|
job_id: str | None = None
|
|
|
|
def response(self, *, idempotent: bool = False) -> dict[str, Any]:
|
|
return {"task_id": self.task_id, "repository": self.repository, "base_oid": self.base_oid, "target_branch": self.target_branch, "expected_branch_oid": self.expected_branch_oid, "generation": self.generation, "branch_mode": self.branch_mode, "create_target_branch": self.branch_mode == "CREATE_NEW", "worktree_mode": "detached_at_base_oid" if self.branch_mode == "CREATE_NEW" else "detached_at_expected_head", "cas_ref_update_required": True, "idempotent": idempotent, "state": "BOUND", "job_id": self.job_id}
|
|
|
|
|
|
def _legacy_safe(value: str) -> str:
|
|
return re.sub(r"[^A-Za-z0-9_.-]+", "-", value).strip("-")[:120] or "task"
|
|
|
|
|
|
def _task_key(value: str) -> str:
|
|
readable = _legacy_safe(value)[:72]
|
|
return f"{readable}-{hashlib.sha256(value.encode('utf-8')).hexdigest()[:24]}"
|
|
|
|
|
|
class GovernedTaskBootstrap:
|
|
def __init__(self, authority: Any, root: Path, *, remote_name: str, token_path: Path, ssh_key_path: Path, api_base: str = "https://g.cnnetmkt.top/git/api/v1") -> None:
|
|
self.authority = authority; self.root = Path(root); self.remote_name = remote_name; self.token_path = Path(token_path); self.ssh_key_path = Path(ssh_key_path); self.api_base = api_base; self.task_meta_root = self.root / "task-bindings"; self.job_root = self.root / "jobs"; self.lock_root = self.root / "task-locks"; self.task_meta_root.mkdir(parents=True, exist_ok=True); self.job_root.mkdir(parents=True, exist_ok=True); self.lock_root.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _store(self, task_id: str) -> tuple[JobStore, GitOpsJobController]:
|
|
legacy = self.job_root / _legacy_safe(task_id)
|
|
keyed = self.job_root / _task_key(task_id)
|
|
root = legacy if legacy.is_dir() and not keyed.exists() else keyed
|
|
store = JobStore(root, payload_root=root / "payloads"); return store, GitOpsJobController(store)
|
|
|
|
def _metadata(self, task_id: str) -> Path:
|
|
return self.task_meta_root / (_task_key(task_id) + ".json")
|
|
|
|
def _legacy_metadata(self, task_id: str) -> Path:
|
|
return self.task_meta_root / (_legacy_safe(task_id) + ".json")
|
|
|
|
@contextmanager
|
|
def _task_lock(self, task_id: str):
|
|
path = self.lock_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)
|
|
|
|
@contextmanager
|
|
def _branch_lock(self, repository: str, target_branch: str):
|
|
key = hashlib.sha256((repository + "\0" + target_branch).encode("utf-8")).hexdigest()
|
|
path = self.lock_root / ("branch-" + key + ".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)
|
|
|
|
def _read_existing(self, task_id: str) -> GovernedTaskBinding | None:
|
|
path = self._metadata(task_id)
|
|
if not path.exists():
|
|
path = self._legacy_metadata(task_id)
|
|
if not path.is_file() or path.is_symlink(): return None
|
|
binding = GovernedTaskBinding(**json.loads(path.read_text(encoding="utf-8")))
|
|
if binding.task_id != task_id:
|
|
raise BootstrapError("TASK_BINDING_KEY_COLLISION")
|
|
return binding
|
|
|
|
def get_binding(self, task_id: str) -> GovernedTaskBinding:
|
|
"""Return the exact durable binding; never synthesize missing lineage."""
|
|
existing = self._read_existing(str(task_id))
|
|
if existing is None or not existing.job_id:
|
|
raise BootstrapError("TASK_EXECUTION_PROJECTION_MISSING")
|
|
return existing
|
|
|
|
def execution_binding(self, binding: GovernedTaskBinding) -> GovernedExecutionBinding:
|
|
expected = binding.expected_branch_oid
|
|
mode = binding.branch_mode
|
|
if mode == "CREATE_NEW":
|
|
current = self.authority.read_ref("refs/heads/" + binding.target_branch)
|
|
if current is not None:
|
|
store, _ = self._store(binding.task_id)
|
|
with store._connect() as connection:
|
|
rows = connection.execute(
|
|
"SELECT current_head_oid,generation,state,lineage_anchor_oid "
|
|
"FROM branch_leases WHERE job_id=? AND branch=? AND state='ACTIVE' "
|
|
"ORDER BY revision DESC",
|
|
(binding.job_id, binding.target_branch),
|
|
).fetchall()
|
|
if len(rows) != 1:
|
|
raise BootstrapError("SAME_JOB_CONTINUATION_LEASE_INVALID")
|
|
row = rows[0]
|
|
if (
|
|
str(row["current_head_oid"] or "") != current
|
|
or int(row["generation"] or 0) != binding.generation
|
|
or str(row["lineage_anchor_oid"] or "") != binding.base_oid
|
|
):
|
|
raise BootstrapError("SAME_JOB_CONTINUATION_BINDING_INVALID")
|
|
expected = current
|
|
mode = "CONTINUE_EXISTING"
|
|
return GovernedExecutionBinding.from_mapping(
|
|
{
|
|
"task_id": binding.task_id,
|
|
"job_id": binding.job_id,
|
|
"repository": binding.repository,
|
|
"base_oid": binding.base_oid,
|
|
"expected_branch_oid": expected,
|
|
"target_branch": binding.target_branch,
|
|
"branch_mode": mode,
|
|
"generation": binding.generation,
|
|
"idempotency_key": binding.idempotency_key,
|
|
}
|
|
)
|
|
|
|
def execution_envelope(
|
|
self,
|
|
task_id: str,
|
|
execution_request: Mapping[str, Any] | None = None,
|
|
*,
|
|
required_capabilities: tuple[str, ...] = (),
|
|
expected_job_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
binding = self.get_binding(task_id)
|
|
if expected_job_id and binding.job_id != expected_job_id:
|
|
raise BootstrapError("JOB_BINDING_CONFLICT")
|
|
return project_execution_envelope(
|
|
self.execution_binding(binding),
|
|
execution_request,
|
|
required_capabilities=required_capabilities,
|
|
)
|
|
|
|
def online_task_projection(
|
|
self,
|
|
task_id: str,
|
|
task_request: Mapping[str, Any],
|
|
execution_request: Mapping[str, Any] | None = None,
|
|
*,
|
|
expected_job_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
binding = self.get_binding(task_id)
|
|
if expected_job_id and binding.job_id != expected_job_id:
|
|
raise BootstrapError("JOB_BINDING_CONFLICT")
|
|
return project_online_task(self.execution_binding(binding), task_request, execution_request)
|
|
|
|
def _write(self, binding: GovernedTaskBinding) -> None:
|
|
path = self._metadata(binding.task_id); temporary = path.with_suffix(".tmp"); temporary.write_text(json.dumps(asdict(binding), sort_keys=True), encoding="utf-8"); temporary.replace(path)
|
|
|
|
def create(self, request: Mapping[str, Any]) -> dict[str, Any]:
|
|
task_id = str(request.get("task_id") or "")
|
|
if not task_id:
|
|
raise BootstrapError("TASK_BINDING_INVALID")
|
|
repository = str(request.get("repository") or "")
|
|
target_branch = str(request.get("target_branch") or "")
|
|
with self._task_lock(task_id):
|
|
with self._branch_lock(repository, target_branch):
|
|
return self._create_locked(request)
|
|
|
|
def _create_locked(self, request: Mapping[str, Any]) -> dict[str, Any]:
|
|
values = {key: str(request[key]) for key in ("task_id", "repository", "base_oid", "target_branch", "risk", "description", "idempotency_key")}
|
|
if values["repository"] != "aios/aios" or not OID_RE.fullmatch(values["base_oid"]): raise BootstrapError("TASK_BINDING_INVALID")
|
|
if not BRANCH_RE.fullmatch(values["target_branch"]): raise BootstrapError("TARGET_BRANCH_INVALID")
|
|
existing = self._read_existing(values["task_id"])
|
|
if existing is not None:
|
|
if any(getattr(existing, key) != values[key] for key in ("repository", "base_oid", "target_branch", "risk", "description", "idempotency_key")): raise BootstrapError("TASK_BINDING_CONFLICT")
|
|
optional = {"expected_branch_oid": existing.expected_branch_oid, "branch_mode": existing.branch_mode, "generation": existing.generation}
|
|
for key, durable in optional.items():
|
|
if key in request and str(request[key]) != str(durable): raise BootstrapError("TASK_BINDING_CONFLICT")
|
|
return existing.response(idempotent=True)
|
|
if not self.authority.commit_exists(values["base_oid"]): raise BootstrapError("BASE_COMMIT_NOT_FOUND")
|
|
current = self.authority.read_ref("refs/heads/" + values["target_branch"])
|
|
expected = str(request.get("expected_branch_oid") or values["base_oid"])
|
|
if current is None:
|
|
mode, generation = "CREATE_NEW", 1
|
|
if expected != values["base_oid"]: raise BootstrapError("EXPECTED_BASE_CONFLICT")
|
|
else:
|
|
if expected != current: raise BootstrapError("EXISTING_BRANCH_OID_CONFLICT")
|
|
parent_task_id = str(request.get("parent_task_id") or "")
|
|
parent_job_id = str(request.get("parent_job_id") or "")
|
|
if not parent_task_id or not parent_job_id:
|
|
raise BootstrapError("EXISTING_BRANCH_OWNERSHIP_REQUIRED")
|
|
parent = self._read_existing(parent_task_id)
|
|
if parent is None or parent.job_id != parent_job_id or parent.repository != values["repository"] or parent.target_branch != values["target_branch"]:
|
|
raise BootstrapError("EXISTING_BRANCH_OWNERSHIP_CONFLICT")
|
|
store, _ = self._store(parent_task_id)
|
|
with store._connect() as connection:
|
|
rows = connection.execute(
|
|
"SELECT current_head_oid,generation,state FROM branch_leases WHERE job_id=? AND branch=? AND state='ACTIVE' ORDER BY revision DESC",
|
|
(parent_job_id, values["target_branch"]),
|
|
).fetchall()
|
|
if len(rows) != 1 or str(rows[0]["current_head_oid"] or "") != current:
|
|
raise BootstrapError("EXISTING_BRANCH_OWNERSHIP_CONFLICT")
|
|
mode, generation = "CONTINUE_EXISTING", int(rows[0]["generation"]) + 1
|
|
if "generation" in request and int(request["generation"]) != generation:
|
|
raise BootstrapError("GENERATION_BINDING_CONFLICT")
|
|
online_request = request.get("online_task")
|
|
execution_request = request.get("execution")
|
|
projection_required = (
|
|
online_request is not None
|
|
or execution_request is not None
|
|
or request.get("require_online_task_projection") is True
|
|
)
|
|
if projection_required:
|
|
if not isinstance(online_request, Mapping):
|
|
raise BootstrapError("ONLINE_TASK_REQUEST_INVALID")
|
|
if execution_request is not None and not isinstance(execution_request, Mapping):
|
|
raise BootstrapError("EXECUTION_REQUEST_INVALID")
|
|
provisional_job_id = "job-" + sha256_json(
|
|
{"client_idempotency_key": values["idempotency_key"]}
|
|
)[:32]
|
|
try:
|
|
provisional_binding = GovernedExecutionBinding.from_mapping(
|
|
{
|
|
"task_id": values["task_id"],
|
|
"job_id": provisional_job_id,
|
|
"repository": values["repository"],
|
|
"base_oid": values["base_oid"],
|
|
"expected_branch_oid": expected,
|
|
"target_branch": values["target_branch"],
|
|
"branch_mode": mode,
|
|
"generation": generation,
|
|
"idempotency_key": values["idempotency_key"],
|
|
}
|
|
)
|
|
project_online_task(
|
|
provisional_binding,
|
|
dict(online_request),
|
|
dict(execution_request or {}),
|
|
)
|
|
except Exception as exc:
|
|
raise BootstrapError("EXECUTION_BINDING_INVALID") from exc
|
|
|
|
payload = json.dumps(values, sort_keys=True, separators=(",", ":")).encode()
|
|
_, controller = self._store(values["task_id"])
|
|
receipt = controller.submit({"idempotency_key": values["idempotency_key"], "campaign_id": values["task_id"], "task_id": values["task_id"], "repository": values["repository"], "operation": "GOVERNED_TASK_CREATE", "branch": values["target_branch"], "expected_branch_oid": expected, "payload_digest": hashlib.sha256(payload).hexdigest(), "lineage_anchor_oid": values["base_oid"]}, payload=payload)
|
|
binding = GovernedTaskBinding(**values, expected_branch_oid=expected, branch_mode=mode, generation=generation, job_id=receipt.job_id)
|
|
self._write(binding)
|
|
return binding.response()
|