"""Core data models for the task plane.""" from __future__ import annotations import hashlib import json import re import time import uuid import urllib.parse from enum import StrEnum from dataclasses import dataclass, field from datetime import datetime, timezone from typing import Any SENSITIVE_KEYS = { "secret", "password", "token", "api_key", "access_token", "refresh_token", "authorization", "private_key", "client_secret", "credential", } IDENTIFIER_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,127}") TYPE_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}") RISK_LEVELS = {"low", "medium", "high"} DATA_CLASSIFICATIONS = { "public", "internal", "confidential", "credential-secret", } SECRET_TEXT_PATTERNS = ( re.compile(r"(?i)\bAuthorization\s*:\s*Bearer\s+\S+"), re.compile(r"(?i)\b(?:password|passwd|api[_-]?key|access[_-]?token|secret)\s*[:=]\s*\S+"), re.compile(r"\bsk-[A-Za-z0-9_-]{12,}\b"), re.compile(r"\bAKIA[0-9A-Z]{16}\b"), re.compile(r"-----BEGIN (?:RSA |EC |OPENSSH )?PRIVATE KEY-----"), ) TASK_PLAN_KEYS = { "schema_version", "task_id", "task_type", "issuer", "intent", "required_capabilities", "preferred_workers", "risk", "approval_policy", "dependencies", "max_delegation_depth", "budgets", "evidence_requirements", "parent_task_id", "correlation_id", "metadata", "data_classification", "execution_zone", "network_policy", "logging_policy", "allowed_input_paths", "allowed_output_paths", "output_policy", } def _now_ts() -> float: return time.time() def canonical_json(value: Any) -> str: return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) def compute_sha256_hex(value: Any) -> str: return hashlib.sha256(canonical_json(value).encode("utf-8")).hexdigest() def contains_secret_text(value: Any) -> bool: if isinstance(value, str): return any(pattern.search(value) for pattern in SECRET_TEXT_PATTERNS) if isinstance(value, dict): return any(contains_secret_text(item) for item in value.values()) if isinstance(value, (list, tuple)): return any(contains_secret_text(item) for item in value) return False def validate_relative_paths(values: Any, name: str) -> list[str]: paths = _normalize_list(values) for value in paths: decoded = urllib.parse.unquote(value) parts = value.replace("\\", "/").split("/") if ( decoded != value or "\\" in value or value.startswith("/") or re.match(r"^[A-Za-z]:", value) or any(part in {"", ".", ".."} for part in parts) ): raise ValueError(f"invalid {name}") return paths def validate_identifier(value: Any, name: str) -> str: text = str(value or "") if ( not IDENTIFIER_RE.fullmatch(text) or urllib.parse.unquote(text) != text or any(ord(ch) < 32 or ord(ch) == 127 for ch in text) ): raise ValueError(f"invalid {name}") return text def validate_type(value: Any, name: str) -> str: text = str(value or "") if not TYPE_RE.fullmatch(text): raise ValueError(f"invalid {name}") return text def redact_value(value: Any) -> Any: if isinstance(value, dict): redacted = {} for key, inner in value.items(): lowered = str(key).lower() if lowered in SENSITIVE_KEYS or lowered.endswith( ("_password", "_secret", "_access_token", "_refresh_token", "_private_key") ): redacted[key] = "[REDACTED]" else: redacted[key] = redact_value(inner) return redacted if isinstance(value, list): return [redact_value(v) for v in value] if isinstance(value, tuple): return tuple(redact_value(v) for v in value) return value class TaskStateTransitionError(ValueError): pass class TaskStatus(StrEnum): CREATED = "CREATED" QUEUED = "QUEUED" CLAIMED = "CLAIMED" RUNNING = "RUNNING" WAITING_CHILD = "WAITING_CHILD" VERIFYING = "VERIFYING" COMPLETED = "COMPLETED" FAILED_RECOVERABLE = "FAILED_RECOVERABLE" FAILED_FINAL = "FAILED_FINAL" CANCELLED = "CANCELLED" MANUAL_REVIEW_REQUIRED = "MANUAL_REVIEW_REQUIRED" TASK_STATES = { status.value for status in TaskStatus } TASK_NEXT_STATES = { "CREATED": {"QUEUED", "CLAIMED", "CANCELLED", "MANUAL_REVIEW_REQUIRED"}, "QUEUED": {"CLAIMED", "CANCELLED", "FAILED_RECOVERABLE", "FAILED_FINAL"}, "CLAIMED": {"RUNNING", "FAILED_RECOVERABLE", "FAILED_FINAL", "CANCELLED"}, "RUNNING": {"VERIFYING", "WAITING_CHILD", "COMPLETED", "FAILED_RECOVERABLE", "FAILED_FINAL", "CANCELLED"}, "WAITING_CHILD": {"RUNNING", "VERIFYING", "FAILED_RECOVERABLE", "FAILED_FINAL", "CANCELLED"}, "VERIFYING": {"COMPLETED", "FAILED_RECOVERABLE", "FAILED_FINAL", "CANCELLED"}, "COMPLETED": set(), "FAILED_RECOVERABLE": {"CLAIMED", "RUNNING", "CANCELLED"}, "FAILED_FINAL": set(), "CANCELLED": set(), "MANUAL_REVIEW_REQUIRED": {"RUNNING", "CANCELLED"}, } @dataclass class Budget: time_seconds: int | None = None attempts: int | None = None tokens: int | None = None cost: float | None = None @classmethod def from_dict(cls, data: dict[str, Any] | None) -> "Budget": if not data: return cls() if not isinstance(data, dict) or set(data) - { "time_seconds", "attempts", "tokens", "cost", }: raise ValueError("invalid budget") for name in ("time_seconds", "attempts", "tokens", "cost"): value = data.get(name) if value is not None: try: if float(value) < 0: raise ValueError("budget values must be non-negative") except (TypeError, ValueError) as exc: raise ValueError("invalid budget value") from exc return cls( time_seconds=_to_positive_int(data.get("time_seconds")), attempts=_to_positive_int(data.get("attempts")), tokens=_to_positive_int(data.get("tokens")), cost=_to_float(data.get("cost")), ) def to_dict(self) -> dict[str, Any]: return { "time_seconds": self.time_seconds, "attempts": self.attempts, "tokens": self.tokens, "cost": self.cost, } def merge(self, other: "Budget") -> "Budget": return Budget( time_seconds=_first_non_none(self.time_seconds, other.time_seconds), attempts=_first_non_none(self.attempts, other.attempts), tokens=_first_non_none(self.tokens, other.tokens), cost=_first_non_none(self.cost, other.cost), ) @dataclass class TaskPlan: schema_version: int = 1 task_id: str | None = None task_type: str = "code.change" issuer: dict[str, Any] = field(default_factory=dict) intent: str | None = None required_capabilities: list[str] = field(default_factory=list) preferred_workers: list[str] = field(default_factory=list) risk: str = "medium" approval_policy: str = "patch-only" dependencies: list[str] = field(default_factory=list) max_delegation_depth: int = 5 budgets: Budget = field(default_factory=Budget) evidence_requirements: list[str] = field(default_factory=list) parent_task_id: str | None = None correlation_id: str | None = None metadata: dict[str, Any] = field(default_factory=dict) data_classification: str = "internal" execution_zone: str = "isolated-development" network_policy: str = "controlled" logging_policy: str = "redacted" allowed_input_paths: list[str] = field(default_factory=list) allowed_output_paths: list[str] = field(default_factory=list) output_policy: str = "standard" @classmethod def from_dict(cls, data: dict[str, Any] | None) -> "TaskPlan": if not isinstance(data, dict): raise ValueError("task plan payload must be an object") unknown = set(data) - TASK_PLAN_KEYS if unknown: raise ValueError("unknown task plan fields") if "task_id" not in data or data["task_id"] is None: data = dict(data) data["task_id"] = f"task-{uuid.uuid4()}" risk = str(data.get("risk", "medium")) if risk not in RISK_LEVELS: raise ValueError("invalid risk") issuer = data.get("issuer", {}) if not isinstance(issuer, dict): raise ValueError("issuer must be an object") schema_version = int(data.get("schema_version", 1)) if schema_version != 1: raise ValueError("unsupported schema_version") max_depth = _to_positive_int(data.get("max_delegation_depth"), default=5) if max_depth is None or max_depth > 32: raise ValueError("invalid max_delegation_depth") classification = str(data.get("data_classification", "internal")) if classification not in DATA_CLASSIFICATIONS: raise ValueError("invalid data_classification") execution_zone = validate_type( data.get("execution_zone", "isolated-development"), "execution_zone" ) network_policy = validate_type( data.get("network_policy", "controlled"), "network_policy" ) logging_policy = validate_type( data.get("logging_policy", "redacted"), "logging_policy" ) output_policy = validate_type( data.get("output_policy", "standard"), "output_policy" ) if contains_secret_text( { "intent": data.get("intent"), "metadata": data.get("metadata", {}), } ): raise ValueError("task content contains secret material") if classification == "credential-secret": if execution_zone != "secure-file-transformer": raise ValueError("credential-secret requires secure-file-transformer") if network_policy != "denied": raise ValueError("credential-secret requires network denied") if logging_policy != "metadata-only": raise ValueError("credential-secret requires metadata-only logging") if output_policy != "hash-count-status-only": raise ValueError("credential-secret requires restricted output") return cls( schema_version=schema_version, task_id=validate_identifier(data["task_id"], "task_id"), task_type=validate_type(data.get("task_type", "code.change"), "task_type"), issuer=issuer, intent=data.get("intent"), required_capabilities=_normalize_list(data.get("required_capabilities", [])), preferred_workers=_normalize_list(data.get("preferred_workers", [])), risk=risk, approval_policy=validate_type( data.get("approval_policy", "patch-only"), "approval_policy" ), dependencies=_normalize_list(data.get("dependencies", [])), max_delegation_depth=max_depth, budgets=Budget.from_dict(data.get("budgets")), evidence_requirements=_normalize_list(data.get("evidence_requirements", [])), parent_task_id=data.get("parent_task_id"), correlation_id=data.get("correlation_id"), metadata=data.get("metadata", {}) if isinstance(data.get("metadata", {}), dict) else {}, data_classification=classification, execution_zone=execution_zone, network_policy=network_policy, logging_policy=logging_policy, allowed_input_paths=validate_relative_paths( data.get("allowed_input_paths", []), "allowed_input_paths" ), allowed_output_paths=validate_relative_paths( data.get("allowed_output_paths", []), "allowed_output_paths" ), output_policy=output_policy, ) def to_dict(self, *, redact: bool = False) -> dict[str, Any]: payload = { "schema_version": self.schema_version, "task_id": self.task_id, "task_type": self.task_type, "issuer": self.issuer, "intent": self.intent, "required_capabilities": list(self.required_capabilities), "preferred_workers": list(self.preferred_workers), "risk": self.risk, "approval_policy": self.approval_policy, "dependencies": list(self.dependencies), "max_delegation_depth": self.max_delegation_depth, "budgets": self.budgets.to_dict(), "evidence_requirements": list(self.evidence_requirements), "parent_task_id": self.parent_task_id, "correlation_id": self.correlation_id, "metadata": dict(self.metadata), "data_classification": self.data_classification, "execution_zone": self.execution_zone, "network_policy": self.network_policy, "logging_policy": self.logging_policy, "allowed_input_paths": list(self.allowed_input_paths), "allowed_output_paths": list(self.allowed_output_paths), "output_policy": self.output_policy, } if redact: payload = redact_value(payload) return payload def operation_digest(self) -> str: return compute_sha256_hex(self.to_dict()) @dataclass class TaskRecord: task_id: str status: str revision: int task_plan: TaskPlan created_at: float updated_at: float root_task_id: str idempotency_key: str | None = None failure_reason: str | None = None def with_status(self, status: str) -> "TaskRecord": allowed = TASK_NEXT_STATES.get(self.status, set()) if status != self.status and (not allowed and self.status not in {"COMPLETED", "FAILED_FINAL", "CANCELLED"}): raise TaskStateTransitionError(f"Cannot transition {self.status} -> {status}") if status not in TASK_STATES: raise TaskStateTransitionError(f"Invalid status {status}") if status != self.status and self.status in TASK_NEXT_STATES and status not in TASK_NEXT_STATES[self.status]: raise TaskStateTransitionError(f"Cannot transition {self.status} -> {status}") return TaskRecord( task_id=self.task_id, status=status, revision=self.revision + 1, task_plan=self.task_plan, created_at=self.created_at, updated_at=_now_ts(), root_task_id=self.root_task_id, idempotency_key=self.idempotency_key, failure_reason=self.failure_reason, ) def to_dict(self, *, include_task=True, redact_payload: bool = True) -> dict[str, Any]: payload = { "task_id": self.task_id, "status": self.status, "revision": self.revision, "root_task_id": self.root_task_id, "parent_task_id": self.task_plan.parent_task_id, "created_at": _fmt_ts(self.created_at), "updated_at": _fmt_ts(self.updated_at), "failure_reason": self.failure_reason, } if include_task: payload["plan"] = self.task_plan.to_dict(redact=redact_payload) if self.idempotency_key: payload["idempotency_key"] = self.idempotency_key return payload @dataclass class WorkerRegistration: worker_id: str worker_type: str identity: str capabilities: list[str] trust_zone: str max_concurrency: int heartbeat_interval_seconds: int supported_protocols: list[str] registered_at: float last_heartbeat: float status: str = "active" disabled: bool = False owner_identity: str = "" registration_revision: int = 1 registration_policy: str = "self" created_at: float = 0.0 updated_at: float = 0.0 disabled_reason: str | None = None @classmethod def from_dict(cls, data: dict[str, Any]) -> "WorkerRegistration": if not isinstance(data, dict): raise ValueError("worker registration payload must be an object") now = _now_ts() worker_id = validate_identifier(data.get("worker_id"), "worker_id") worker_type = validate_type(data.get("worker_type"), "worker_type") identity = str(data.get("identity") or "").strip() if not identity or any(ord(ch) < 32 for ch in identity): raise ValueError("invalid identity") max_concurrency = _to_positive_int(data.get("max_concurrency"), default=1) heartbeat_interval = _to_positive_int( data.get("heartbeat_interval_seconds"), default=15 ) if not max_concurrency: raise ValueError("max_concurrency must be positive") if not heartbeat_interval: raise ValueError("heartbeat_interval_seconds must be positive") return cls( worker_id=worker_id, worker_type=worker_type, identity=identity, capabilities=_normalize_list(data.get("capabilities", [])), trust_zone=str(data.get("trust_zone") or "default"), max_concurrency=max_concurrency, heartbeat_interval_seconds=heartbeat_interval, supported_protocols=_normalize_list( data.get("supported_protocols", ["aios-task-v1"]) ), registered_at=float(data.get("registered_at") or now), last_heartbeat=float(data.get("last_heartbeat") or now), status=str(data.get("status") or "active"), disabled=bool(data.get("disabled", False)), owner_identity=str(data.get("owner_identity") or identity), registration_revision=int(data.get("registration_revision") or 1), registration_policy=str(data.get("registration_policy") or "self"), created_at=float(data.get("created_at") or data.get("registered_at") or now), updated_at=float(data.get("updated_at") or now), disabled_reason=data.get("disabled_reason"), ) def to_dict(self, redact_payload: bool = True) -> dict[str, Any]: payload = { "worker_id": self.worker_id, "worker_type": self.worker_type, "identity": self.identity, "capabilities": list(self.capabilities), "trust_zone": self.trust_zone, "max_concurrency": self.max_concurrency, "heartbeat_interval_seconds": self.heartbeat_interval_seconds, "supported_protocols": list(self.supported_protocols), "registered_at": _fmt_ts(self.registered_at), "last_heartbeat": _fmt_ts(self.last_heartbeat), "status": self.status, "disabled": self.disabled, "owner_identity": self.owner_identity, "registration_revision": self.registration_revision, "registration_policy": self.registration_policy, "created_at": _fmt_ts(self.created_at or self.registered_at), "updated_at": _fmt_ts(self.updated_at or self.last_heartbeat), "disabled_reason": self.disabled_reason, } if redact_payload: payload = redact_value(payload) return payload @dataclass(frozen=True) class EvidenceEnvelope: schema_version: int evidence_id: str task_id: str claim_id: str kind: str worker_identity: str payload: dict[str, Any] created_at: float previous_evidence_hash: str evidence_digest: str signature: str | None = None attestation_ref: str | None = None @classmethod def create( cls, *, task_id: str, claim_id: str, kind: str, worker_identity: str, payload: dict[str, Any], previous_evidence_hash: str, signature: str | None = None, attestation_ref: str | None = None, evidence_id: str | None = None, created_at: float | None = None, ) -> "EvidenceEnvelope": envelope = cls( schema_version=1, evidence_id=evidence_id or f"evd-{uuid.uuid4()}", task_id=validate_identifier(task_id, "task_id"), claim_id=validate_identifier(claim_id, "claim_id"), kind=validate_type(kind, "evidence kind"), worker_identity=str(worker_identity), payload=redact_value(payload), created_at=float(created_at or _now_ts()), previous_evidence_hash=str(previous_evidence_hash or ""), evidence_digest="", signature=signature, attestation_ref=attestation_ref, ) return dataclass_replace( envelope, evidence_digest=compute_sha256_hex(envelope.digest_payload()) ) def digest_payload(self) -> dict[str, Any]: return { "schema_version": self.schema_version, "evidence_id": self.evidence_id, "task_id": self.task_id, "claim_id": self.claim_id, "kind": self.kind, "worker_identity": self.worker_identity, "payload": self.payload, "created_at": self.created_at, "previous_evidence_hash": self.previous_evidence_hash, "attestation_ref": self.attestation_ref, } def verify( self, *, expected_task_id: str, expected_claim_id: str, expected_previous_hash: str, attestation_verifier: Any = None, ) -> bool: if ( self.schema_version != 1 or self.task_id != expected_task_id or self.claim_id != expected_claim_id or self.previous_evidence_hash != expected_previous_hash or self.evidence_digest != compute_sha256_hex(self.digest_payload()) ): return False if self.signature or self.attestation_ref: return True if attestation_verifier is None else bool(attestation_verifier(self)) return True def to_dict(self) -> dict[str, Any]: return { **self.digest_payload(), "evidence_digest": self.evidence_digest, "signature": self.signature, } class ProbeOutcome(StrEnum): NOT_APPLIED = "NOT_APPLIED" APPLIED_VERIFIED = "APPLIED_VERIFIED" INDETERMINATE = "INDETERMINATE" CONFLICT = "CONFLICT" @dataclass(frozen=True) class ProbeEvidence: schema_version: int probe_id: str task_id: str claim_id: str worker_id: str worker_identity: str operation_digest: str outcome: ProbeOutcome observed_at: float expires_at: float probe_executor_identity: str evidence_digest: str signature: str | None = None attestation_ref: str | None = None @classmethod def create(cls, **data: Any) -> "ProbeEvidence": outcome = data["outcome"] if not isinstance(outcome, ProbeOutcome): outcome = ProbeOutcome(str(outcome)) probe = cls( schema_version=1, probe_id=data.get("probe_id") or f"probe-{uuid.uuid4()}", task_id=validate_identifier(data["task_id"], "task_id"), claim_id=validate_identifier(data["claim_id"], "claim_id"), worker_id=validate_identifier(data["worker_id"], "worker_id"), worker_identity=str(data["worker_identity"]), operation_digest=str(data["operation_digest"]), outcome=outcome, observed_at=float(data["observed_at"]), expires_at=float(data["expires_at"]), probe_executor_identity=str(data["probe_executor_identity"]), evidence_digest="", signature=data.get("signature"), attestation_ref=data.get("attestation_ref"), ) return dataclass_replace( probe, evidence_digest=compute_sha256_hex(probe.digest_payload()) ) def digest_payload(self) -> dict[str, Any]: return { "schema_version": self.schema_version, "probe_id": self.probe_id, "task_id": self.task_id, "claim_id": self.claim_id, "worker_id": self.worker_id, "worker_identity": self.worker_identity, "operation_digest": self.operation_digest, "outcome": self.outcome.value, "observed_at": self.observed_at, "expires_at": self.expires_at, "probe_executor_identity": self.probe_executor_identity, "attestation_ref": self.attestation_ref, } def verify(self, *, now: float, attestation_verifier: Any) -> bool: if ( self.schema_version != 1 or self.expires_at <= now or self.observed_at > now or self.evidence_digest != compute_sha256_hex(self.digest_payload()) or not (self.signature or self.attestation_ref) ): return False return bool(attestation_verifier and attestation_verifier(self)) def to_dict(self) -> dict[str, Any]: return { **self.digest_payload(), "evidence_digest": self.evidence_digest, "signature": self.signature, } def dataclass_replace(value: Any, **changes: Any) -> Any: from dataclasses import replace return replace(value, **changes) @dataclass class WorkerClaim: claim_id: str task_id: str worker_id: str worker_identity: str claimed_at: float lease_until: float attempt: int task_revision: int operation_digest: str status: str = "ACTIVE" def to_dict(self, redact_payload: bool = True) -> dict[str, Any]: payload = { "claim_id": self.claim_id, "task_id": self.task_id, "worker_id": self.worker_id, "worker_identity": self.worker_identity, "claimed_at": _fmt_ts(self.claimed_at), "lease_until": _fmt_ts(self.lease_until), "attempt": self.attempt, "task_revision": self.task_revision, "operation_digest": self.operation_digest, "status": self.status, } if redact_payload: payload = redact_value(payload) return payload @dataclass class ResultEnvelope: task_id: str claim_id: str worker_identity: str status: str outputs: list[dict[str, Any]] evidence: list[dict[str, Any]] started_at: str completed_at: str result_digest: str | None = None signature: str | None = None attestation_ref: str | None = None operation_digest: str | None = None trust_status: str = "untrusted_shadow_result" def digest_payload(self) -> dict[str, Any]: return { "task_id": self.task_id, "claim_id": self.claim_id, "worker_identity": self.worker_identity, "status": self.status, "outputs": self.outputs, "evidence": self.evidence, "started_at": self.started_at, "completed_at": self.completed_at, "operation_digest": self.operation_digest, "attestation_ref": self.attestation_ref, } def computed_digest(self) -> str: return compute_sha256_hex(self.digest_payload()) def to_dict(self, redact_payload: bool = True) -> dict[str, Any]: payload = { "task_id": self.task_id, "claim_id": self.claim_id, "worker_identity": self.worker_identity, "status": self.status, "outputs": self.outputs, "evidence": self.evidence, "started_at": self.started_at, "completed_at": self.completed_at, "result_digest": self.result_digest, "signature": self.signature, "attestation_ref": self.attestation_ref, "operation_digest": self.operation_digest, "trust_status": self.trust_status, } if redact_payload: payload = redact_value(payload) return payload def _normalize_list(value: Any) -> list[str]: if not isinstance(value, list): raise ValueError("expected list") if len(value) > 128 or any( not isinstance(v, str) or not v or len(v) > 256 or any(ord(ch) < 32 for ch in v) for v in value ): raise ValueError("invalid list item") return list(value) def _to_positive_int(value: Any, default: int | None = None) -> int | None: if value is None: return default try: number = int(value) except (TypeError, ValueError): return default return number if number >= 0 else default def _to_float(value: Any) -> float | None: if value is None: return None try: return float(value) except (TypeError, ValueError): return None def _first_non_none(primary: Any, fallback: Any) -> Any: if primary is not None: return primary return fallback def iso8601_now() -> str: return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z") def _fmt_ts(value: float | None) -> str | None: if value is None: return None return datetime.fromtimestamp(value, tz=timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")