324 lines
12 KiB
Python
324 lines
12 KiB
Python
"""AIOS V2 Agent 执行证据。
|
|
|
|
Evidence 按租户隔离并使用原子 JSON 写入。生产模式下，claim、heartbeat
|
|
和 complete 在落盘前必须通过最小权限 Evidence Grant 授权。
|
|
"""
|
|
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, Protocol
|
|
|
|
from .provenance import ProvenanceDenied, digest_payload
|
|
|
|
|
|
class EvidenceAuthorizationWriter(Protocol):
|
|
def authorize_evidence_write(self, **values: Any) -> Any: ...
|
|
|
|
|
|
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
|
|
provenance_id: str | None = None
|
|
run_id: str | None = None
|
|
last_authorization_id: 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))
|
|
for optional in ("provenance_id", "run_id", "last_authorization_id"):
|
|
value = getattr(self, optional)
|
|
object.__setattr__(self, optional, str(value).strip() if value else None)
|
|
|
|
|
|
class ExecutionTracker:
|
|
"""每个租户、每次执行持久保存一份证据记录。"""
|
|
|
|
def __init__(
|
|
self,
|
|
root: str | Path,
|
|
*,
|
|
provenance_writer: EvidenceAuthorizationWriter | None = None,
|
|
require_provenance: bool = False,
|
|
) -> None:
|
|
self.root = Path(root).resolve()
|
|
self.root.mkdir(parents=True, exist_ok=True)
|
|
self.provenance_writer = provenance_writer
|
|
self.require_provenance = bool(require_provenance)
|
|
if self.require_provenance and self.provenance_writer is None:
|
|
raise ValueError("production evidence tracking requires a provenance writer")
|
|
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 _authorize(
|
|
self,
|
|
evidence: ExecutionEvidence,
|
|
*,
|
|
operation: str,
|
|
principal_id: str | None,
|
|
provenance_id: str | None,
|
|
run_id: str | None,
|
|
changes: Mapping[str, Any],
|
|
) -> str | None:
|
|
if not self.require_provenance:
|
|
return None
|
|
if self.provenance_writer is None:
|
|
raise RuntimeError("Run Provenance writer is not configured")
|
|
effective_provenance = provenance_id or evidence.provenance_id
|
|
effective_run = run_id or evidence.run_id
|
|
if not effective_provenance or not effective_run or not principal_id:
|
|
raise RuntimeError(
|
|
"production evidence write requires provenance_id, run_id and principal_id"
|
|
)
|
|
request_digest = digest_payload(
|
|
{
|
|
"operation": operation,
|
|
"execution_id": evidence.execution_id,
|
|
"tenant_id": evidence.tenant_id,
|
|
"task_id": evidence.task_id,
|
|
"agent_id": evidence.agent_id,
|
|
"executor": evidence.executor,
|
|
"branch": evidence.branch,
|
|
"changes": dict(changes),
|
|
}
|
|
)
|
|
try:
|
|
authorization = self.provenance_writer.authorize_evidence_write(
|
|
provenance_id=effective_provenance,
|
|
run_id=effective_run,
|
|
tenant_id=evidence.tenant_id,
|
|
task_id=evidence.task_id,
|
|
agent_id=evidence.agent_id,
|
|
principal_id=principal_id,
|
|
request_digest=request_digest,
|
|
)
|
|
except ProvenanceDenied as exc:
|
|
raise RuntimeError(f"Run Provenance denied evidence write: {exc}") from exc
|
|
return authorization.authorization_id
|
|
|
|
def claim(
|
|
self,
|
|
evidence: ExecutionEvidence,
|
|
*,
|
|
principal_id: str | None = None,
|
|
provenance_id: str | None = None,
|
|
run_id: str | None = None,
|
|
) -> 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")
|
|
if self.require_provenance and (
|
|
current.provenance_id != (provenance_id or evidence.provenance_id)
|
|
or current.run_id != (run_id or evidence.run_id)
|
|
):
|
|
raise RuntimeError(
|
|
"execution_id already belongs to another provenance context"
|
|
)
|
|
return current
|
|
effective_provenance = provenance_id or evidence.provenance_id
|
|
effective_run = run_id or evidence.run_id
|
|
authorization_id = self._authorize(
|
|
evidence,
|
|
operation="claim",
|
|
principal_id=principal_id,
|
|
provenance_id=effective_provenance,
|
|
run_id=effective_run,
|
|
changes={"status": evidence.status, "current_step": evidence.current_step},
|
|
)
|
|
persisted = replace(
|
|
evidence,
|
|
provenance_id=effective_provenance,
|
|
run_id=effective_run,
|
|
last_authorization_id=authorization_id,
|
|
)
|
|
self._write(path, persisted)
|
|
return persisted
|
|
|
|
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,
|
|
principal_id: str | None = None,
|
|
provenance_id: str | None = None,
|
|
run_id: str | None = None,
|
|
) -> ExecutionEvidence:
|
|
path = self._path(tenant_id, execution_id)
|
|
with self._lock:
|
|
current = self._read(path)
|
|
changes = {
|
|
"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,
|
|
}
|
|
authorization_id = self._authorize(
|
|
current,
|
|
operation="heartbeat",
|
|
principal_id=principal_id,
|
|
provenance_id=provenance_id,
|
|
run_id=run_id,
|
|
changes=changes,
|
|
)
|
|
updated = replace(
|
|
current,
|
|
**changes,
|
|
last_authorization_id=authorization_id or current.last_authorization_id,
|
|
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,
|
|
principal_id: str | None = None,
|
|
provenance_id: str | None = None,
|
|
run_id: str | None = None,
|
|
) -> ExecutionEvidence:
|
|
path = self._path(tenant_id, execution_id)
|
|
with self._lock:
|
|
current = self._read(path)
|
|
changes = {
|
|
"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),
|
|
}
|
|
authorization_id = self._authorize(
|
|
current,
|
|
operation="complete",
|
|
principal_id=principal_id,
|
|
provenance_id=provenance_id,
|
|
run_id=run_id,
|
|
changes=changes,
|
|
)
|
|
updated = replace(
|
|
current,
|
|
**changes,
|
|
last_authorization_id=authorization_id or current.last_authorization_id,
|
|
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.chmod(temporary, 0o640)
|
|
os.replace(temporary, path)
|