aios/active-runtime/app/aios_v2/execution_tracking.py

"""Git-backed agent execution evidence for AIOS V2.

Evidence is tenant-scoped, append-friendly, and contains no runtime credentials.
It provides a deterministic source for claim, heartbeat, progress, result and
verification state across Codex, OpenClaw and Hermes executors.
"""
from __future__ import annotations

import json
import os
import threading
from dataclasses import asdict, dataclass, field, replace
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Mapping


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _required(value: str, label: str) -> str:
    normalized = str(value or "").strip()
    if not normalized:
        raise ValueError(f"{label} is required")
    return normalized


def _safe(value: Mapping[str, Any] | None) -> dict[str, Any]:
    return json.loads(json.dumps(dict(value or {}), ensure_ascii=False, default=str))


@dataclass(frozen=True)
class ExecutionEvidence:
    execution_id: str
    tenant_id: str
    task_id: str
    agent_id: str
    executor: str
    branch: str
    status: str = "claimed"
    progress_completed: int = 0
    progress_total: int = 1
    current_step: str = "claim"
    last_commit: str | None = None
    test_status: str = "not_run"
    test_summary: Mapping[str, Any] = field(default_factory=dict)
    result_summary: Mapping[str, Any] = field(default_factory=dict)
    claimed_at: str = field(default_factory=utc_now)
    heartbeat_at: str = field(default_factory=utc_now)
    completed_at: str | None = None

    def __post_init__(self) -> None:
        for name in (
            "execution_id",
            "tenant_id",
            "task_id",
            "agent_id",
            "executor",
            "branch",
            "status",
            "current_step",
            "claimed_at",
            "heartbeat_at",
        ):
            object.__setattr__(self, name, _required(getattr(self, name), name))
        if self.progress_completed < 0 or self.progress_total < 1:
            raise ValueError("invalid progress values")
        if self.progress_completed > self.progress_total:
            raise ValueError("progress_completed cannot exceed progress_total")
        object.__setattr__(self, "test_summary", _safe(self.test_summary))
        object.__setattr__(self, "result_summary", _safe(self.result_summary))


class ExecutionTracker:
    """Persist one JSON evidence record per tenant and execution."""

    def __init__(self, root: str | Path) -> None:
        self.root = Path(root).resolve()
        self.root.mkdir(parents=True, exist_ok=True)
        self._lock = threading.RLock()

    @staticmethod
    def _component(value: str) -> str:
        normalized = _required(value, "path component")
        if normalized in {".", ".."} or "/" in normalized or "\\" in normalized:
            raise ValueError("invalid evidence path component")
        return normalized

    def _path(self, tenant_id: str, execution_id: str) -> Path:
        tenant = self._component(tenant_id)
        execution = self._component(execution_id)
        return self.root / tenant / f"{execution}.json"

    def claim(self, evidence: ExecutionEvidence) -> ExecutionEvidence:
        path = self._path(evidence.tenant_id, evidence.execution_id)
        with self._lock:
            if path.exists():
                current = self._read(path)
                if current.task_id != evidence.task_id or current.agent_id != evidence.agent_id:
                    raise RuntimeError("execution_id already belongs to another task")
                return current
            self._write(path, evidence)
            return evidence

    def heartbeat(
        self,
        *,
        tenant_id: str,
        execution_id: str,
        status: str | None = None,
        current_step: str | None = None,
        progress_completed: int | None = None,
        progress_total: int | None = None,
        last_commit: str | None = None,
    ) -> ExecutionEvidence:
        path = self._path(tenant_id, execution_id)
        with self._lock:
            current = self._read(path)
            updated = replace(
                current,
                status=status or current.status,
                current_step=current_step or current.current_step,
                progress_completed=(
                    progress_completed
                    if progress_completed is not None
                    else current.progress_completed
                ),
                progress_total=(
                    progress_total if progress_total is not None else current.progress_total
                ),
                last_commit=last_commit or current.last_commit,
                heartbeat_at=utc_now(),
            )
            self._write(path, updated)
            return updated

    def complete(
        self,
        *,
        tenant_id: str,
        execution_id: str,
        status: str,
        test_status: str,
        test_summary: Mapping[str, Any] | None = None,
        result_summary: Mapping[str, Any] | None = None,
        last_commit: str | None = None,
    ) -> ExecutionEvidence:
        path = self._path(tenant_id, execution_id)
        with self._lock:
            current = self._read(path)
            updated = replace(
                current,
                status=_required(status, "status"),
                current_step="complete",
                progress_completed=current.progress_total,
                last_commit=last_commit or current.last_commit,
                test_status=_required(test_status, "test_status"),
                test_summary=_safe(test_summary),
                result_summary=_safe(result_summary),
                heartbeat_at=utc_now(),
                completed_at=utc_now(),
            )
            self._write(path, updated)
            return updated

    def get(self, *, tenant_id: str, execution_id: str) -> ExecutionEvidence:
        return self._read(self._path(tenant_id, execution_id))

    def list(self, *, tenant_id: str) -> list[ExecutionEvidence]:
        tenant_dir = self.root / self._component(tenant_id)
        if not tenant_dir.exists():
            return []
        records = [self._read(path) for path in tenant_dir.glob("*.json")]
        return sorted(records, key=lambda item: (item.heartbeat_at, item.execution_id), reverse=True)

    @staticmethod
    def _read(path: Path) -> ExecutionEvidence:
        try:
            value = json.loads(path.read_text(encoding="utf-8"))
        except FileNotFoundError as exc:
            raise KeyError("execution evidence not found") from exc
        if not isinstance(value, Mapping):
            raise RuntimeError("invalid execution evidence file")
        return ExecutionEvidence(**dict(value))

    @staticmethod
    def _write(path: Path, evidence: ExecutionEvidence) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        temporary = path.with_suffix(f".tmp-{os.getpid()}")
        temporary.write_text(
            json.dumps(asdict(evidence), ensure_ascii=False, sort_keys=True, indent=2) + "\n",
            encoding="utf-8",
        )
        os.replace(temporary, path)