aios/active-runtime/app/task_plane/schedulability.py

"""SchedulabilityProof model and ProofBuilder for governed claim.

A SchedulabilityProof is a durable, server-side diagnostic snapshot that records
the state evaluated before an Agent may claim a task.  It is NOT a bearer
credential; it is persisted for auditability and idempotent read-back.  The five
sections captured are binding / worker / runtime / observations / status,
following ANNEX-2.3 R2 ARCH-AUTO-CLAIM-R2-20260914 §4.1.
"""

from __future__ import annotations

import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Callable

from .models import compute_sha256_hex


class ProofStatus:
    PENDING = "PENDING"
    ACCEPTED = "ACCEPTED"
    REJECTED = "REJECTED"
    CLAIMED = "CLAIMED"
    FAILED = "FAILED"


@dataclass
class Generation:
    task_revision: int
    attempt: int | None = None
    lease_generation: int | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "task_revision": self.task_revision,
            "attempt": self.attempt,
            "lease_generation": self.lease_generation,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any] | None) -> "Generation":
        data = data or {}
        return cls(
            task_revision=int(data.get("task_revision") or 0),
            attempt=data.get("attempt"),
            lease_generation=data.get("lease_generation"),
        )


@dataclass
class SchedulabilityProof:
    proof_id: str
    task_id: str
    worker_id: str
    generated_at: float
    binding: dict[str, Any]
    worker: dict[str, Any]
    runtime: dict[str, Any]
    observations: list[dict[str, Any]]
    status: str = ProofStatus.PENDING
    schema_version: int = 1
    claim_id: str | None = None
    digest: str = ""

    def payload(self) -> dict[str, Any]:
        """The digest-stable payload (excludes the digest field itself)."""
        return {
            "proof_id": self.proof_id,
            "schema_version": self.schema_version,
            "task_id": self.task_id,
            "worker_id": self.worker_id,
            "claim_id": self.claim_id,
            "generated_at": self.generated_at,
            "binding": self.binding,
            "worker": self.worker,
            "runtime": self.runtime,
            "observations": self.observations,
            "status": self.status,
        }

    def to_dict(self) -> dict[str, Any]:
        return {**self.payload(), "digest": self.digest}


def compute_proof_digest(proof_or_dict: Any) -> str:
    """sha256 over canonical_json of the digest-stable proof payload."""
    if isinstance(proof_or_dict, SchedulabilityProof):
        return compute_sha256_hex(proof_or_dict.payload())
    payload = dict(proof_or_dict)
    payload.pop("digest", None)
    return compute_sha256_hex(payload)


def observation(source: str, field_: str, value: Any, ok: bool, at: float) -> dict[str, Any]:
    return {"source": source, "field": field_, "value": value, "ok": ok, "at": at}


def default_tenant_of_task(task: Any) -> str:
    return task.task_plan.metadata.get("tenant_id", "default")


def default_tenant_of_worker(worker: Any) -> str:
    return worker.trust_zone or "default"


class ProofBuilder:
    """Builds a SchedulabilityProof from live store snapshots."""

    def __init__(
        self,
        store: Any,
        *,
        tenant_of_task: Callable[[Any], str] | None = None,
        tenant_of_worker: Callable[[Any], str] | None = None,
    ) -> None:
        self.store = store
        self.tenant_of_task = tenant_of_task or default_tenant_of_task
        self.tenant_of_worker = tenant_of_worker or default_tenant_of_worker

    def build(
        self,
        task: Any,
        worker: Any,
        *,
        now: float | None = None,
        claim_id: str | None = None,
    ) -> SchedulabilityProof:
        now = now if now is not None else time.time()
        active = self.store.get_active_claim_count(worker.worker_id)
        capacity_ok = active < worker.max_concurrency
        online = (not worker.disabled) and (
            now <= worker.last_heartbeat + (worker.heartbeat_interval_seconds * 3)
        )
        tenant_worker = self.tenant_of_worker(worker)
        tenant_task = self.tenant_of_task(task)

        binding = {
            "worker_identity": worker.identity,
            "capabilities": list(worker.capabilities),
            "trust_zone": worker.trust_zone,
            "tenant_id": tenant_worker,
        }
        binding["binding_digest"] = compute_sha256_hex(binding)

        worker_snapshot = {
            "status": worker.status,
            "disabled": worker.disabled,
            "last_heartbeat": worker.last_heartbeat,
            "heartbeat_interval_seconds": worker.heartbeat_interval_seconds,
            "online": online,
            "worker_type": worker.worker_type,
            "max_concurrency": worker.max_concurrency,
            "current_active": active,
        }

        runtime = {
            "max_concurrency": worker.max_concurrency,
            "current_active": active,
            "capacity_ok": capacity_ok,
            "generation": {
                "task_revision": task.revision,
                "attempt": None,
                "lease_generation": None,
            },
        }

        observations = [
            observation("tenant", "tenant_match", {"task": tenant_task, "worker": tenant_worker}, tenant_task == tenant_worker, now),
            observation("principal", "worker_identity", worker.identity, True, now),
            observation("authorization", "authorization_ok", None, True, now),
            observation("task", "task_id", task.task_id, task.status in {"CREATED", "QUEUED", "FAILED_RECOVERABLE"}, now),
            observation("worker", "worker_status", worker.status, not worker.disabled, now),
            observation("generation", "task_revision", task.revision, True, now),
            observation("capacity", "capacity_ok", capacity_ok, capacity_ok, now),
        ]

        proof = SchedulabilityProof(
            proof_id=f"proof-{uuid.uuid4()}",
            task_id=task.task_id,
            worker_id=worker.worker_id,
            generated_at=now,
            binding=binding,
            worker=worker_snapshot,
            runtime=runtime,
            observations=observations,
            status=ProofStatus.PENDING,
            claim_id=claim_id,
        )
        proof.digest = compute_sha256_hex(proof.payload())
        return proof