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

"""Agent discovery/claim coordinator over diagnostic schedulability evidence."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Mapping
from .schedulability import (
    DiscoveryDiagnostics,
    SchedulabilityProof,
    SchedulabilityService,
)


@dataclass(frozen=True)
class DiscoveryResult:
    status: str
    proof: SchedulabilityProof | None
    blocker_codes: tuple[str, ...] = ()
    candidate_count: int = 0

    def to_dict(self) -> dict[str, Any]:
        value: dict[str, Any] = {
            "status": self.status,
            "blocker_codes": list(self.blocker_codes),
            "candidate_count": self.candidate_count,
            "proof": None,
        }
        if self.proof is not None:
            value["proof"] = {
                "schema": self.proof.schema,
                "proof_id": self.proof.proof_id,
                "task_id": self.proof.task_id,
                "worker_id": self.proof.worker_id,
                "created_at": self.proof.created_at,
                "expires_at": self.proof.expires_at,
                "blockers": [b.to_dict() for b in self.proof.blockers],
                "proof_digest": self.proof.proof_digest,
            }
        return value


def _waiting_status(diag: DiscoveryDiagnostics) -> str:
    codes = set(diag.blocker_codes)
    if diag.candidate_count == 0:
        return "NO_CANDIDATES"
    if not codes:
        return "NO_CLAIMABLE_TASK"
    if codes <= {"CAPACITY_EXHAUSTED"}:
        return "WAITING_CAPACITY"
    if "AUTHORIZATION_BLOCKED" in codes or "EXACT_SCOPE_MISMATCH" in codes:
        return "WAITING_AUTHORIZATION"
    if "DEPENDENCY_BLOCKED" in codes:
        return "WAITING_DEPENDENCY"
    if "FREEZE_BLOCKED" in codes:
        return "WAITING_FREEZE"
    if "CONFLICTING_LEASE" in codes:
        return "WAITING_LEASE"
    if {"STALE_OR_INACTIVE_WORKER", "MISSING_CAPABILITY", "INVALID_AUTHORITY_VALUE"} & codes:
        return "WORKER_NOT_ELIGIBLE"
    if "TASK_STATE_NOT_CLAIMABLE" in codes:
        return "NO_CLAIMABLE_TASK"
    return "BLOCKED"


class AutoClaimCoordinator:
    """No implicit credential derivation; authenticated identity is supplied by host."""

    def __init__(self, schedulability: SchedulabilityService) -> None:
        self.schedulability = schedulability

    def discover(self, *, worker_id: str, actor_identity: str) -> DiscoveryResult:
        proof, diag = self.schedulability.discover_with_diagnostics(
            worker_id=worker_id, actor_identity=actor_identity
        )
        if proof is not None:
            return DiscoveryResult(
                "CLAIMABLE", proof, diag.blocker_codes, diag.candidate_count
            )
        return DiscoveryResult(
            _waiting_status(diag), None, diag.blocker_codes, diag.candidate_count
        )

    def claim(
        self,
        *,
        proof_id: str,
        worker_id: str,
        actor_identity: str,
        lease_seconds: int = 30,
    ) -> Mapping[str, Any]:
        return self.schedulability.claim_with_recheck(
            proof_id=proof_id,
            worker_id=worker_id,
            actor_identity=actor_identity,
            lease_seconds=lease_seconds,
        )