"""SQLite-backed durable store with CAS revisions, WAL mode, and append-only events.""" from __future__ import annotations import hashlib import json import os import sqlite3 import threading import time import uuid from contextlib import contextmanager from typing import Any from .capabilities import has_capabilities from .errors import ( ConflictError, LockedError, NotFoundError, ValidationError, ) from .models import ( Budget, EvidenceEnvelope, ProbeEvidence, ProbeOutcome, TaskPlan, TaskRecord, WorkerClaim, WorkerRegistration, canonical_json, redact_value, validate_identifier, ) EventEnvelope = tuple[str, str, str, str, str, float] OutboxRow = tuple[str, str, str, str, str, str, str, float, float | None, int, str | None] def _canonical(v: Any) -> str: return canonical_json(v) def _hash_hex(value: str) -> str: return hashlib.sha256(value.encode("utf-8")).hexdigest() def _now() -> float: return time.time() def _validated_id(value: str, name: str) -> str: try: return validate_identifier(value, name) except ValueError as exc: raise ValidationError(f"invalid {name}") from exc class TaskPlaneStore: """Durable SQLite store with atomic operations and audit-ready events.""" def __init__(self, db_path: str = ":memory:", *, high_assurance: bool = False) -> None: self.db_path = db_path self.high_assurance = high_assurance self.fault_injector = None self.probe_attestation_verifier = None self._lock = threading.RLock() self._tx_state = threading.local() self._closed = False self._init_db(db_path) def _init_db(self, path: str) -> None: if path != ":memory:": parent = os.path.dirname(os.path.abspath(path)) os.makedirs(parent, mode=0o700, exist_ok=True) if os.path.islink(parent): raise PermissionError("task plane database directory must not be symlink") os.chmod(parent, 0o700) self.conn = sqlite3.connect( path, check_same_thread=False, isolation_level=None ) if path != ":memory:": os.chmod(path, 0o600) self.conn.row_factory = sqlite3.Row self.conn.execute("PRAGMA busy_timeout=5000;") if path != ":memory:": self.conn.execute("PRAGMA journal_mode=WAL;") self.conn.execute( f"PRAGMA synchronous={'FULL' if self.high_assurance else 'NORMAL'};" ) self.conn.execute("PRAGMA foreign_keys=ON;") self._create_schema() def close(self) -> None: with self._lock: if not self._closed: self.conn.close() self._closed = True def __del__(self) -> None: try: self.close() except Exception: pass @contextmanager def _transaction(self): """Serialize one connection and keep nested writes in one transaction.""" with self._lock: depth = getattr(self._tx_state, "depth", 0) if depth == 0: self.conn.execute("BEGIN IMMEDIATE") self._tx_state.depth = depth + 1 try: yield except Exception: self._tx_state.depth -= 1 if depth == 0: self.conn.rollback() raise else: self._tx_state.depth -= 1 if depth == 0: self.conn.commit() def _create_schema(self) -> None: with self._transaction(): self.conn.execute( """ CREATE TABLE IF NOT EXISTS tasks ( task_id TEXT PRIMARY KEY, revision INTEGER NOT NULL, status TEXT NOT NULL, task_json TEXT NOT NULL, task_type TEXT NOT NULL, root_task_id TEXT NOT NULL, parent_task_id TEXT, correlation_id TEXT, idempotency_key TEXT UNIQUE, created_at REAL NOT NULL, updated_at REAL NOT NULL, failure_reason TEXT ) """ ) self.conn.execute( """ CREATE TABLE IF NOT EXISTS workers ( worker_id TEXT PRIMARY KEY, worker_type TEXT NOT NULL, identity TEXT NOT NULL, capabilities TEXT NOT NULL, trust_zone TEXT NOT NULL, max_concurrency INTEGER NOT NULL, heartbeat_interval_seconds INTEGER NOT NULL, supported_protocols TEXT NOT NULL, registered_at REAL NOT NULL, last_heartbeat REAL NOT NULL, status TEXT NOT NULL DEFAULT 'active', disabled INTEGER NOT NULL DEFAULT 0 ) """ ) self.conn.execute( """ CREATE TABLE IF NOT EXISTS claims ( claim_id TEXT PRIMARY KEY, task_id TEXT NOT NULL, worker_id TEXT NOT NULL, worker_identity TEXT NOT NULL, claimed_at REAL NOT NULL, lease_until REAL NOT NULL, attempt INTEGER NOT NULL, task_revision INTEGER NOT NULL, operation_digest TEXT NOT NULL, status TEXT NOT NULL, completed_at REAL, result_payload TEXT, UNIQUE(task_id, worker_id, claim_id) ) """ ) self.conn.execute( """ CREATE TABLE IF NOT EXISTS events ( event_id TEXT PRIMARY KEY, task_id TEXT NOT NULL, seq INTEGER NOT NULL, type TEXT NOT NULL, event_data TEXT NOT NULL, subject TEXT, created_at REAL NOT NULL, prev_hash TEXT NOT NULL, hash TEXT NOT NULL ) """ ) self.conn.execute( """ CREATE TABLE IF NOT EXISTS outbox ( message_id TEXT PRIMARY KEY, task_id TEXT NOT NULL, task_status TEXT, event_type TEXT NOT NULL, stage TEXT, correlation_id TEXT, payload TEXT NOT NULL, delivery_status TEXT NOT NULL DEFAULT 'NOT_DELIVERED', created_at REAL NOT NULL, delivered_at REAL, attempts INTEGER NOT NULL DEFAULT 0, error TEXT, receipt TEXT ) """ ) self.conn.execute( """ CREATE TABLE IF NOT EXISTS evidence ( evidence_id TEXT PRIMARY KEY, task_id TEXT NOT NULL, kind TEXT NOT NULL, payload TEXT NOT NULL, digest TEXT NOT NULL, worker_identity TEXT NOT NULL, created_at REAL NOT NULL ) """ ) outbox_columns = { row["name"] for row in self.conn.execute("PRAGMA table_info(outbox)") } if "receipt" not in outbox_columns: self.conn.execute("ALTER TABLE outbox ADD COLUMN receipt TEXT") self.conn.execute( """ CREATE TABLE IF NOT EXISTS delegations ( parent_task_id TEXT NOT NULL, child_task_id TEXT PRIMARY KEY, created_at REAL NOT NULL ) """ ) self.conn.execute( """ CREATE TABLE IF NOT EXISTS task_budget ( task_id TEXT PRIMARY KEY, time_seconds INTEGER, attempts INTEGER, tokens INTEGER, cost REAL ) """ ) self.conn.execute( """ CREATE TABLE IF NOT EXISTS schema_meta ( key TEXT PRIMARY KEY, value TEXT NOT NULL ) """ ) self.conn.execute( """ CREATE TABLE IF NOT EXISTS probe_evidence ( probe_id TEXT PRIMARY KEY, task_id TEXT NOT NULL, claim_id TEXT NOT NULL, envelope_json TEXT NOT NULL, created_at REAL NOT NULL ) """ ) self._ensure_columns( "workers", { "owner_identity": "TEXT NOT NULL DEFAULT ''", "registration_revision": "INTEGER NOT NULL DEFAULT 1", "registration_policy": "TEXT NOT NULL DEFAULT 'self'", "created_at": "REAL NOT NULL DEFAULT 0", "updated_at": "REAL NOT NULL DEFAULT 0", "disabled_reason": "TEXT", }, ) self._ensure_columns( "evidence", { "claim_id": "TEXT NOT NULL DEFAULT ''", "previous_evidence_hash": "TEXT NOT NULL DEFAULT ''", "envelope_json": "TEXT", }, ) self.conn.execute( """ INSERT INTO schema_meta(key,value) VALUES('schema_version','2') ON CONFLICT(key) DO UPDATE SET value=excluded.value """ ) self.conn.execute("CREATE INDEX IF NOT EXISTS ix_tasks_status ON tasks(status)") self.conn.execute("CREATE INDEX IF NOT EXISTS ix_tasks_parent ON tasks(parent_task_id)") self.conn.execute("CREATE INDEX IF NOT EXISTS ix_workers_status ON workers(status)") self.conn.execute("CREATE INDEX IF NOT EXISTS ix_claims_task ON claims(task_id)") self.conn.execute( "CREATE UNIQUE INDEX IF NOT EXISTS ux_events_task_seq ON events(task_id, seq)" ) self.conn.execute("CREATE INDEX IF NOT EXISTS ix_outbox_status ON outbox(delivery_status)") def _ensure_columns(self, table: str, columns: dict[str, str]) -> None: existing = { row["name"] for row in self.conn.execute(f"PRAGMA table_info({table})") } for name, declaration in columns.items(): if name not in existing: self.conn.execute( f"ALTER TABLE {table} ADD COLUMN {name} {declaration}" ) def _fault(self, point: str) -> None: if self.fault_injector is not None: self.fault_injector(point) def schema_version(self) -> int: row = self.conn.execute( "SELECT value FROM schema_meta WHERE key='schema_version'" ).fetchone() return int(row["value"]) if row else 0 def integrity_check(self) -> str: return str(self.conn.execute("PRAGMA integrity_check").fetchone()[0]) def checkpoint_and_backup(self, destination: str) -> None: target = os.path.abspath(destination) os.makedirs(os.path.dirname(target), mode=0o700, exist_ok=True) with self._lock: if self.db_path != ":memory:": self.conn.execute("PRAGMA wal_checkpoint(FULL)") backup = sqlite3.connect(target) try: self.conn.backup(backup) finally: backup.close() os.chmod(target, 0o600) def _row_to_task(self, row: sqlite3.Row | None) -> TaskRecord: if row is None: raise NotFoundError("task") task_json = json.loads(row["task_json"]) task_plan = TaskPlan.from_dict(task_json) task_plan.parent_task_id = row["parent_task_id"] task_plan.correlation_id = row["correlation_id"] return TaskRecord( task_id=row["task_id"], status=row["status"], revision=row["revision"], task_plan=task_plan, created_at=row["created_at"], updated_at=row["updated_at"], root_task_id=row["root_task_id"], idempotency_key=row["idempotency_key"], failure_reason=row["failure_reason"], ) def _row_to_worker(self, row: sqlite3.Row | None) -> WorkerRegistration: if row is None: raise NotFoundError("worker") return WorkerRegistration( worker_id=row["worker_id"], worker_type=row["worker_type"], identity=row["identity"], capabilities=json.loads(row["capabilities"]), trust_zone=row["trust_zone"], max_concurrency=row["max_concurrency"], heartbeat_interval_seconds=row["heartbeat_interval_seconds"], supported_protocols=json.loads(row["supported_protocols"]), registered_at=row["registered_at"], last_heartbeat=row["last_heartbeat"], status=row["status"], disabled=bool(row["disabled"]), owner_identity=row["owner_identity"] or row["identity"], registration_revision=row["registration_revision"], registration_policy=row["registration_policy"], created_at=row["created_at"] or row["registered_at"], updated_at=row["updated_at"] or row["last_heartbeat"], disabled_reason=row["disabled_reason"], ) def create_task(self, task_plan: TaskPlan, idempotency_key: str | None = None) -> TaskRecord: task_id = task_plan.task_id or f"task-{uuid.uuid4()}" task_plan.task_id = task_id task_json = redact_value(task_plan.to_dict(redact=False)) now = _now() with self._transaction(): if idempotency_key: existing = self.conn.execute( "SELECT task_id,task_json FROM tasks WHERE idempotency_key=?", (idempotency_key,), ).fetchone() if existing: if _canonical(json.loads(existing["task_json"])) != _canonical(task_json): raise ConflictError("idempotency key payload mismatch") return self.get_task(existing["task_id"]) root_task_id = task_id for dependency in task_plan.dependencies: dependency_id = _validated_id(dependency, "dependency task_id") if dependency_id == task_id: raise ConflictError("task dependency cycle") self.get_task(dependency_id) if task_plan.parent_task_id: parent = self.get_task(task_plan.parent_task_id) root_task_id = parent.root_task_id self.conn.execute( """ INSERT INTO tasks (task_id, revision, status, task_json, task_type, root_task_id, parent_task_id, correlation_id, idempotency_key, created_at, updated_at, failure_reason) VALUES (?, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL) """, ( task_id, "CREATED", json.dumps(task_json, ensure_ascii=False), task_plan.task_type, root_task_id, task_plan.parent_task_id, task_plan.correlation_id, idempotency_key, now, now, ), ) self._set_budget(task_id, task_plan.budgets) if task_plan.parent_task_id: self.conn.execute( "INSERT INTO delegations (parent_task_id, child_task_id, created_at) VALUES (?, ?, ?)", (task_plan.parent_task_id, task_id, now), ) return self.get_task(task_id) def create_task_with_event( self, task_plan: TaskPlan, idempotency_key: str | None = None ) -> TaskRecord: with self._transaction(): task = self.create_task(task_plan, idempotency_key=idempotency_key) if not self.list_events(task.task_id): self._fault("create.after_task_before_event") self.append_event( task.task_id, "TASK_CREATED", {"task_id": task.task_id}, subject=f"task/{task.task_id}", ) return self.get_task(task.task_id) def get_task(self, task_id: str) -> TaskRecord: task_id = _validated_id(task_id, "task_id") with self._lock: row = self.conn.execute( "SELECT * FROM tasks WHERE task_id=?", (task_id,) ).fetchone() if row is None: raise NotFoundError(f"task {task_id} not found") return self._row_to_task(row) def list_tasks(self) -> list[TaskRecord]: with self._lock: rows = self.conn.execute( "SELECT * FROM tasks ORDER BY created_at DESC" ).fetchall() return [self._row_to_task(row) for row in rows] def get_task_budget(self, task_id: str) -> Budget: with self._lock: row = self.conn.execute( "SELECT * FROM task_budget WHERE task_id=?", (task_id,), ).fetchone() if row is None: return Budget() return Budget( time_seconds=row["time_seconds"], attempts=row["attempts"], tokens=row["tokens"], cost=row["cost"], ) def _set_budget(self, task_id: str, budget: Budget) -> None: self.conn.execute( """ INSERT INTO task_budget(task_id, time_seconds, attempts, tokens, cost) VALUES (?, ?, ?, ?, ?) ON CONFLICT(task_id) DO UPDATE SET time_seconds=excluded.time_seconds, attempts=excluded.attempts, tokens=excluded.tokens, cost=excluded.cost """, (task_id, budget.time_seconds, budget.attempts, budget.tokens, budget.cost), ) def consume_budget(self, task_id: str, needed: Budget) -> Budget: with self._transaction(): row = self.conn.execute( "SELECT time_seconds, attempts, tokens, cost FROM task_budget WHERE task_id=?", (task_id,), ).fetchone() if row is None: raise NotFoundError("task budget") current = Budget(row["time_seconds"], row["attempts"], row["tokens"], row["cost"]) for key in ("time_seconds", "attempts", "tokens"): needed_value = getattr(needed, key) current_value = getattr(current, key) if needed_value is not None and current_value is not None: if needed_value is not None and current_value < needed_value: raise ConflictError(f"budget exceeded for {key}") updated = Budget( time_seconds=current.time_seconds - needed.time_seconds if needed.time_seconds is not None and current.time_seconds is not None else current.time_seconds, attempts=current.attempts - needed.attempts if needed.attempts is not None and current.attempts is not None else current.attempts, tokens=current.tokens - needed.tokens if needed.tokens is not None and current.tokens is not None else current.tokens, cost=current.cost - needed.cost if needed.cost is not None and current.cost is not None else current.cost, ) self.conn.execute( """ UPDATE task_budget SET time_seconds=?, attempts=?, tokens=?, cost=? WHERE task_id=? """, (updated.time_seconds, updated.attempts, updated.tokens, updated.cost, task_id), ) return updated def update_task_status( self, task_id: str, expected_revision: int, status: str, failure_reason: str | None = None, ) -> TaskRecord: with self._transaction(): current = self.get_task(task_id) if current.revision != expected_revision: raise ConflictError("task revision mismatch") current.with_status(status) now = _now() self.conn.execute( """ UPDATE tasks SET status=?, revision=revision+1, updated_at=?, failure_reason=? WHERE task_id=? """, (status, now, failure_reason, task_id), ) return self.get_task(task_id) def register_worker( self, registration: WorkerRegistration | dict[str, Any], *, actor_identity: str | None = None, is_admin: bool = False, expected_revision: int | None = None, ) -> WorkerRegistration: if isinstance(registration, dict): registration = WorkerRegistration.from_dict(registration) actor_identity = actor_identity or registration.identity with self._transaction(): existing_row = self.conn.execute( "SELECT * FROM workers WHERE worker_id=?", (registration.worker_id,) ).fetchone() if existing_row is not None: existing = self._row_to_worker(existing_row) if existing.disabled and not is_admin: raise LockedError("worker disabled") if existing.identity != registration.identity: raise ConflictError("worker identity takeover rejected") if actor_identity != existing.owner_identity and not is_admin: raise ConflictError("worker owner identity mismatch") if expected_revision is None or expected_revision != existing.registration_revision: raise ConflictError("worker registration revision mismatch") escalated = ( not set(registration.capabilities).issubset(existing.capabilities) or registration.max_concurrency > existing.max_concurrency or registration.trust_zone != existing.trust_zone ) if escalated and not is_admin: from .errors import ForbiddenError raise ForbiddenError("worker privilege escalation requires admin") revision = existing.registration_revision + 1 created_at = existing.created_at else: if actor_identity != registration.identity and not is_admin: from .errors import ForbiddenError raise ForbiddenError("self registration identity mismatch") revision = 1 created_at = registration.created_at or registration.registered_at now = _now() self.conn.execute( """ INSERT INTO workers (worker_id, worker_type, identity, capabilities, trust_zone, max_concurrency, heartbeat_interval_seconds, supported_protocols, registered_at, last_heartbeat, status, disabled, owner_identity, registration_revision, registration_policy, created_at, updated_at, disabled_reason) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(worker_id) DO UPDATE SET worker_type=excluded.worker_type, capabilities=excluded.capabilities, trust_zone=excluded.trust_zone, max_concurrency=excluded.max_concurrency, heartbeat_interval_seconds=excluded.heartbeat_interval_seconds, supported_protocols=excluded.supported_protocols, last_heartbeat=excluded.last_heartbeat, status=excluded.status, disabled=excluded.disabled, registration_revision=excluded.registration_revision, registration_policy=excluded.registration_policy, updated_at=excluded.updated_at, disabled_reason=excluded.disabled_reason """, ( registration.worker_id, registration.worker_type, registration.identity, json.dumps(registration.capabilities, ensure_ascii=False), registration.trust_zone, registration.max_concurrency, registration.heartbeat_interval_seconds, json.dumps(registration.supported_protocols, ensure_ascii=False), registration.registered_at, registration.last_heartbeat, registration.status, int(registration.disabled), registration.owner_identity or registration.identity, revision, registration.registration_policy, created_at, now, registration.disabled_reason, ), ) self._append_worker_event( registration.worker_id, "WORKER_REGISTERED" if existing_row is None else "WORKER_UPDATED", actor_identity, revision, ) return self.get_worker(registration.worker_id) def _append_worker_event( self, worker_id: str, event_type: str, actor_identity: str, revision: int ) -> None: # Worker audit records live in an append-only synthetic task stream. audit_task_id = "task-worker-registry" row = self.conn.execute( "SELECT task_id FROM tasks WHERE task_id=?", (audit_task_id,) ).fetchone() if row is None: plan = TaskPlan.from_dict( { "task_id": audit_task_id, "task_type": "worker.registry.audit", "issuer": {"type": "service", "id": "task-plane"}, } ) self.create_task(plan) self.append_event( audit_task_id, event_type, { "worker_id": worker_id, "actor_identity": actor_identity, "registration_revision": revision, }, subject=f"worker/{worker_id}", ) def list_workers(self) -> list[WorkerRegistration]: with self._lock: rows = self.conn.execute( "SELECT * FROM workers ORDER BY registered_at DESC" ).fetchall() return [self._row_to_worker(r) for r in rows] def get_worker(self, worker_id: str) -> WorkerRegistration: worker_id = _validated_id(worker_id, "worker_id") with self._lock: row = self.conn.execute( "SELECT * FROM workers WHERE worker_id=?", (worker_id,) ).fetchone() if row is None: raise NotFoundError(f"worker {worker_id} not found") return self._row_to_worker(row) def heartbeat_worker( self, worker_id: str, heartbeat_interval_seconds: int | None = None, *, actor_identity: str | None = None, expected_revision: int | None = None, is_admin: bool = False, ) -> WorkerRegistration: with self._transaction(): worker = self.get_worker(worker_id) if worker.disabled: raise LockedError("worker disabled") if not is_admin and actor_identity != worker.identity: from .errors import ForbiddenError raise ForbiddenError("worker heartbeat identity mismatch") if expected_revision is None or expected_revision != worker.registration_revision: raise ConflictError("worker registration revision mismatch") now = _now() heartbeat_interval = heartbeat_interval_seconds or worker.heartbeat_interval_seconds self.conn.execute( """ UPDATE workers SET last_heartbeat=?, heartbeat_interval_seconds=?, registration_revision=registration_revision+1, updated_at=? WHERE worker_id=? """, (now, heartbeat_interval, now, worker_id), ) updated = self.get_worker(worker_id) self._append_worker_event( worker_id, "WORKER_HEARTBEAT", actor_identity or "", updated.registration_revision ) return updated def disable_worker( self, worker_id: str, *, actor_identity: str, expected_revision: int, reason: str, ) -> WorkerRegistration: with self._transaction(): worker = self.get_worker(worker_id) if worker.registration_revision != expected_revision: raise ConflictError("worker registration revision mismatch") now = _now() self.conn.execute( """ UPDATE workers SET disabled=1,status='disabled',disabled_reason=?, registration_revision=registration_revision+1,updated_at=? WHERE worker_id=? """, (reason, now, worker_id), ) updated = self.get_worker(worker_id) self._append_worker_event( worker_id, "WORKER_DISABLED", actor_identity, updated.registration_revision ) return updated def restore_worker( self, worker_id: str, *, actor_identity: str, expected_revision: int, ) -> WorkerRegistration: with self._transaction(): worker = self.get_worker(worker_id) if worker.registration_revision != expected_revision: raise ConflictError("worker registration revision mismatch") now = _now() self.conn.execute( """ UPDATE workers SET disabled=0,status='active',disabled_reason=NULL, registration_revision=registration_revision+1,updated_at=? WHERE worker_id=? """, (now, worker_id), ) updated = self.get_worker(worker_id) self._append_worker_event( worker_id, "WORKER_RESTORED", actor_identity, updated.registration_revision ) return updated def active_claim(self, task_id: str) -> WorkerClaim | None: with self._lock: row = self.conn.execute( """ SELECT * FROM claims WHERE task_id=? AND status='ACTIVE' AND lease_until>? ORDER BY claimed_at DESC LIMIT 1 """, (task_id, _now()), ).fetchone() if row is None: return None return self._row_to_claim(row) def get_claim(self, claim_id: str) -> WorkerClaim: claim_id = _validated_id(claim_id, "claim_id") with self._lock: row = self.conn.execute( "SELECT * FROM claims WHERE claim_id=?", (claim_id,) ).fetchone() if row is None: raise NotFoundError(f"claim {claim_id} not found") return self._row_to_claim(row) def _row_to_claim(self, row: sqlite3.Row) -> WorkerClaim: return WorkerClaim( claim_id=row["claim_id"], task_id=row["task_id"], worker_id=row["worker_id"], worker_identity=row["worker_identity"], claimed_at=row["claimed_at"], lease_until=row["lease_until"], attempt=row["attempt"], task_revision=row["task_revision"], operation_digest=row["operation_digest"], status=row["status"], ) def claim_task( self, task_id: str, worker_id: str, worker_identity: str, lease_seconds: int, expected_revision: int, operation_digest: str, ) -> WorkerClaim: now = _now() lease_until = now + max(1, int(lease_seconds)) with self._transaction(): worker = self.get_worker(worker_id) if worker.disabled: raise LockedError("worker disabled") if worker.identity != worker_identity: raise ConflictError("worker identity mismatch") task = self.get_task(task_id) if task.revision != expected_revision: raise ConflictError("task revision mismatch") if task.status not in {"CREATED", "QUEUED", "FAILED_RECOVERABLE"}: raise ConflictError(f"task in invalid state for claim: {task.status}") if task.status == "FAILED_RECOVERABLE" and not self.task_has_retryable_probe( task_id ): raise LockedError("recoverable task lacks NOT_APPLIED probe evidence") if self.active_claim(task_id): raise ConflictError("task already claimed") if task.task_plan.data_classification == "credential-secret": if ( worker.worker_type != "secure_file_transformer" and "secure_file_transformer" not in worker.capabilities ): from .errors import ForbiddenError raise ForbiddenError( "credential-secret task requires secure file transformer" ) for dependency_id in task.task_plan.dependencies: if self.get_task(dependency_id).status != "COMPLETED": raise LockedError("task dependency incomplete") required = task.task_plan.required_capabilities if required and not has_capabilities(worker.capabilities, required): raise ConflictError("worker lacks required capabilities") attempts = self.conn.execute( "SELECT COUNT(1) FROM claims WHERE task_id=? AND status='ACTIVE' AND lease_until>?", (task_id, now), ).fetchone()[0] if attempts > 0: raise ConflictError("existing active claim") claim_id = f"claim-{uuid.uuid4()}" active_for_worker = self.conn.execute( "SELECT COUNT(1) FROM claims WHERE worker_id=? AND status='ACTIVE' AND lease_until>?", (worker_id, now), ).fetchone()[0] if active_for_worker >= worker.max_concurrency: raise LockedError("worker concurrency exceeded") attempt = self.conn.execute( "SELECT COUNT(1) FROM claims WHERE task_id=?", (task_id,), ).fetchone()[0] + 1 self.conn.execute( "INSERT INTO claims(claim_id, task_id, worker_id, worker_identity, claimed_at, lease_until, attempt, task_revision, operation_digest, status) " "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'ACTIVE')", ( claim_id, task_id, worker_id, worker_identity, now, lease_until, attempt, expected_revision, operation_digest, ), ) schedule_payload = { "worker_id": worker_id, "worker_identity": worker_identity, "operation_digest": operation_digest, "attempt": attempt, } existing_decision = self.conn.execute( """ SELECT 1 FROM events WHERE task_id=? AND type='SCHEDULE_DECISION' AND event_data=? """, ( task_id, json.dumps(redact_value(schedule_payload), ensure_ascii=False), ), ).fetchone() if existing_decision is None: self.append_event( task_id, "SCHEDULE_DECISION", schedule_payload, subject=f"task/{task_id}", ) self.conn.execute( "UPDATE tasks SET status='CLAIMED', revision=revision+1, updated_at=? WHERE task_id=?", (now, task_id), ) self.append_event(task_id, "TASK_CLAIMED", {"claim_id": claim_id, "worker_id": worker_id}, subject=f"task/{task_id}") return self._row_to_claim(self.conn.execute( "SELECT * FROM claims WHERE claim_id=?", (claim_id,), ).fetchone()) def heartbeat_claim(self, task_id: str, claim_id: str, lease_seconds: int) -> WorkerClaim: now = _now() lease_until = now + max(1, lease_seconds) with self._transaction(): row = self.conn.execute( "SELECT * FROM claims WHERE claim_id=? AND task_id=?", (claim_id, task_id), ).fetchone() if row is None: raise NotFoundError("claim") claim = self._row_to_claim(row) if claim.status != "ACTIVE": raise ConflictError("claim not active") if claim.lease_until <= now: raise LockedError("lease expired") if self.get_task(task_id).status in {"CANCELLED", "COMPLETED", "FAILED_FINAL"}: raise ConflictError("task is terminal") updated = self.conn.execute( """ UPDATE claims SET lease_until=? WHERE claim_id=? AND task_id=? AND status='ACTIVE' AND lease_until>? """, (lease_until, claim_id, task_id, now), ) if updated.rowcount != 1: raise LockedError("claim changed during heartbeat") self.append_event(task_id, "TASK_PROGRESS", {"claim_id": claim_id, "stage": "heartbeat"}, subject=f"task/{task_id}") return self.get_claim(claim_id) def record_progress( self, task_id: str, claim_id: str, worker_identity: str, *, progress: int | float, stage: str, message: str | None, lease_seconds: int, ) -> WorkerClaim: if not isinstance(progress, (int, float)) or not 0 <= progress <= 100: raise ValidationError("progress must be between 0 and 100") now = _now() with self._transaction(): claim = self.get_claim(claim_id) if claim.task_id != task_id: raise ConflictError("claim/task mismatch") if claim.worker_identity != worker_identity: raise ConflictError("claim identity mismatch") if claim.status != "ACTIVE": raise ConflictError("claim not active") if claim.lease_until < now: raise LockedError("lease expired") task = self.get_task(task_id) if task.status == "CLAIMED": self.conn.execute( "UPDATE tasks SET status='RUNNING', revision=revision+1, updated_at=? WHERE task_id=?", (now, task_id), ) self.append_event( task_id, "TASK_STARTED", {"claim_id": claim_id}, subject=f"task/{task_id}", ) elif task.status != "RUNNING": raise ConflictError("task is not running") self.conn.execute( "UPDATE claims SET lease_until=? WHERE claim_id=?", (now + max(1, int(lease_seconds)), claim_id), ) self.append_event( task_id, "TASK_PROGRESS", { "claim_id": claim_id, "progress": progress, "stage": stage, "message": message, }, subject=f"task/{task_id}", ) return self.get_claim(claim_id) def recover_expired_claim( self, task_id: str, claim_id: str, probe: ProbeEvidence ) -> TaskRecord: """Compatibility name retaining the typed, attested recovery contract.""" if not isinstance(probe, ProbeEvidence): raise ValidationError("typed ProbeEvidence required") return self.recover_claim_with_probe_and_event(task_id, claim_id, probe) def recover_claim_with_probe_and_event( self, task_id: str, claim_id: str, probe: ProbeEvidence ) -> TaskRecord: now = _now() with self._transaction(): claim = self.get_claim(claim_id) if claim.task_id != task_id: raise ConflictError("claim/task mismatch") if claim.status != "ACTIVE": raise ConflictError("claim not active") if claim.lease_until >= now: raise LockedError("lease not expired") if ( probe.task_id != task_id or probe.claim_id != claim_id or probe.worker_id != claim.worker_id or probe.worker_identity != claim.worker_identity or probe.operation_digest != claim.operation_digest ): raise ConflictError("probe binding mismatch") if not probe.verify( now=now, attestation_verifier=self.probe_attestation_verifier ): from .errors import ForbiddenError raise ForbiddenError("probe evidence invalid") self.conn.execute( """ INSERT INTO probe_evidence (probe_id,task_id,claim_id,envelope_json,created_at) VALUES(?,?,?,?,?) """, ( probe.probe_id, task_id, claim_id, json.dumps(probe.to_dict(), ensure_ascii=False), now, ), ) self._fault("recover.after_probe_before_state") current_task = self.get_task(task_id) if current_task.status in {"CANCELLED", "COMPLETED", "FAILED_FINAL"}: # Recovery records what happened; it cannot undo a terminal Owner # decision or reopen a finished task. This check shares the state # transaction with the claim close and the appended probe event. task_status = current_task.status claim_status = "CANCELLED" if current_task.status == "CANCELLED" else "CLOSED_BY_RECOVERY" elif probe.outcome == ProbeOutcome.APPLIED_VERIFIED: task_status = "COMPLETED" claim_status = "COMPLETED" elif probe.outcome == ProbeOutcome.CONFLICT: task_status = "MANUAL_REVIEW_REQUIRED" claim_status = "CONFLICT" elif probe.outcome == ProbeOutcome.INDETERMINATE: task_status = "FAILED_RECOVERABLE" claim_status = "RECONCILIATION_REQUIRED" else: task_status = "FAILED_RECOVERABLE" claim_status = "EXPIRED_NOT_APPLIED" self.conn.execute( "UPDATE claims SET status=?,completed_at=? WHERE claim_id=?", (claim_status, now, claim_id), ) self.conn.execute( """ UPDATE tasks SET status=?,revision=revision+1,updated_at=? WHERE task_id=? """, (task_status, now, task_id), ) self.append_event( task_id, "CLAIM_RECOVERY_PROBED", { "claim_id": claim_id, "probe_id": probe.probe_id, "outcome": probe.outcome.value, }, subject=f"task/{task_id}", ) return self.get_task(task_id) def task_has_retryable_probe(self, task_id: str) -> bool: row = self.conn.execute( """ SELECT envelope_json FROM probe_evidence WHERE task_id=? ORDER BY created_at DESC LIMIT 1 """, (task_id,), ).fetchone() if row is None: return False try: payload = json.loads(row["envelope_json"]) return ( payload["outcome"] == ProbeOutcome.NOT_APPLIED.value and float(payload["expires_at"]) > _now() ) except (KeyError, TypeError, ValueError, json.JSONDecodeError): return False def complete_claim( self, task_id: str, claim_id: str, task_status: str, result_payload: dict[str, Any], evidence_entries: list[tuple[str, dict[str, Any], str]] | None = None, ) -> WorkerClaim: now = _now() with self._transaction(): claim = self.get_claim(claim_id) if claim.status != "ACTIVE": raise ConflictError("claim already closed") if claim.task_id != task_id: raise ConflictError("claim/task mismatch") if claim.lease_until < now: raise LockedError("lease expired") task = self.get_task(task_id) if claim.operation_digest != task.task_plan.operation_digest(): raise ConflictError("operation digest mismatch") task.with_status(task_status) for kind, payload, worker_identity in evidence_entries or []: self.add_evidence(task_id, kind, payload, worker_identity) self.conn.execute( "UPDATE claims SET status='COMPLETED', completed_at=?, result_payload=? WHERE claim_id=?", (now, json.dumps(result_payload, ensure_ascii=False), claim_id), ) self.conn.execute( "UPDATE tasks SET status=?, revision=revision+1, updated_at=? WHERE task_id=?", (task_status, now, task_id), ) event_type = "TASK_COMPLETED" if task_status == "COMPLETED" else "TASK_FAILED" self.append_event(task_id, event_type, {"claim_id": claim_id}, subject=f"task/{task_id}") return self.get_claim(claim_id) def complete_claim_with_evidence_and_event( self, task_id: str, claim_id: str, task_status: str, result_payload: dict[str, Any], evidence_envelopes: list[EvidenceEnvelope], ) -> WorkerClaim: now = _now() with self._transaction(): claim = self.get_claim(claim_id) if claim.task_id != task_id: raise ConflictError("claim/task mismatch") if claim.status != "ACTIVE": raise ConflictError("claim already closed") if claim.lease_until < now: raise LockedError("lease expired") task = self.get_task(task_id) if claim.operation_digest != task.task_plan.operation_digest(): raise ConflictError("operation digest mismatch") previous = self._last_evidence_hash(task_id) for envelope in evidence_envelopes: if not envelope.verify( expected_task_id=task_id, expected_claim_id=claim_id, expected_previous_hash=previous, ): raise ConflictError("evidence envelope invalid") self.add_evidence_envelope(envelope) previous = envelope.evidence_digest self._fault("complete.after_evidence_before_state") self.conn.execute( """ UPDATE claims SET status='COMPLETED',completed_at=?,result_payload=? WHERE claim_id=? """, (now, json.dumps(redact_value(result_payload), ensure_ascii=False), claim_id), ) self.conn.execute( """ UPDATE tasks SET status=?,revision=revision+1,updated_at=? WHERE task_id=? """, (task_status, now, task_id), ) self.append_event( task_id, "TASK_COMPLETED" if task_status == "COMPLETED" else "TASK_FAILED", {"claim_id": claim_id}, subject=f"task/{task_id}", ) return self.get_claim(claim_id) def fail_claim(self, task_id: str, claim_id: str, task_status: str) -> WorkerClaim: return self.complete_claim(task_id, claim_id, task_status, {"status": task_status}) def fail_claim_with_event( self, task_id: str, claim_id: str, task_status: str, reason: str ) -> WorkerClaim: with self._transaction(): claim = self.complete_claim( task_id, claim_id, task_status, {"status": task_status} ) # complete_claim emitted the single TASK_FAILED event. self._fault("fail.after_state_before_commit") return claim def cancel_task(self, task_id: str, expected_revision: int, reason: str) -> TaskRecord: return self.update_task_status(task_id, expected_revision, "CANCELLED", failure_reason=reason) def cancel_task_with_event( self, task_id: str, expected_revision: int, reason: str ) -> TaskRecord: with self._transaction(): updated = self.cancel_task(task_id, expected_revision, reason) self._fault("cancel.after_state_before_event") self.append_event( task_id, "TASK_CANCELLED", {"reason": reason}, subject=f"task/{task_id}", ) return updated def delegate_task_with_budget_and_events( self, parent_task_id: str, child_plan: TaskPlan, *, caller: str ) -> TaskRecord: with self._transaction(): parent_budget = self.get_task_budget(parent_task_id) needed = Budget() for name in ("time_seconds", "attempts", "tokens", "cost"): available = getattr(parent_budget, name) requested = getattr(child_plan.budgets, name) if name == "attempts" and requested is None: requested = 1 if available is not None: requested = 0 if requested is None else requested if requested > available: raise ConflictError(f"budget exceeded for {name}") setattr(child_plan.budgets, name, requested) setattr(needed, name, requested) self.consume_budget(parent_task_id, needed) self._fault("delegate.after_budget_before_child") child_plan.parent_task_id = parent_task_id child = self.create_task(child_plan) self.append_event( parent_task_id, "TASK_DELEGATED", { "from_task_id": parent_task_id, "to_task_id": child.task_id, "delegated_by": caller, }, subject=f"task/{parent_task_id}", ) self.append_event( child.task_id, "TASK_CREATED", {"task_id": child.task_id, "parent_task_id": parent_task_id}, subject=f"task/{child.task_id}", ) return child def append_event( self, task_id: str, event_type: str, event_data: dict[str, Any], *, subject: str | None = None, ) -> None: if len(json.dumps(event_data)) > 256_000: raise ValidationError("event data too large", code="event_payload_too_large") with self._transaction(): task = self.get_task(task_id) stored_data = redact_value(event_data) task_data = self.conn.execute( "SELECT seq, hash FROM events WHERE task_id=? ORDER BY seq DESC LIMIT 1", (task_id,), ).fetchone() prev_hash = task_data["hash"] if task_data else "" seq = (task_data["seq"] + 1) if task_data else 1 event_id = f"evt-{uuid.uuid4()}" created_at = _now() resolved_subject = subject or f"task/{task_id}" payload = { "event_id": event_id, "task_id": task_id, "seq": seq, "type": event_type, "subject": resolved_subject, "created_at": created_at, "data": stored_data, "prev_hash": prev_hash, } ev_hash = _hash_hex(_canonical(payload)) self.conn.execute( """ INSERT INTO events(event_id, task_id, seq, type, event_data, subject, created_at, prev_hash, hash) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( event_id, task_id, seq, event_type, json.dumps(stored_data, ensure_ascii=False), resolved_subject, created_at, prev_hash, ev_hash, ), ) self.conn.execute( """ INSERT INTO outbox(message_id, task_id, task_status, event_type, stage, correlation_id, payload, delivery_status, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, 'NOT_DELIVERED', ?) """, ( f"msg-{uuid.uuid4()}", task_id, task.status, event_type, stored_data.get("stage"), stored_data.get("correlation_id") or task.task_plan.correlation_id or task_id, json.dumps( { "specversion": "1.0", "id": event_id, "type": event_type, "source": stored_data.get("source", "aios.task-plane"), "subject": subject or f"task/{task_id}", "time": time.strftime( "%Y-%m-%dT%H:%M:%SZ", time.gmtime(created_at) ), "datacontenttype": "application/json", "data": {"task_id": task_id, "event": stored_data}, }, ensure_ascii=False, ), _now(), ), ) def list_events(self, task_id: str) -> list[dict[str, Any]]: task_id = _validated_id(task_id, "task_id") with self._lock: rows = self.conn.execute( "SELECT event_id, seq, type, event_data, subject, created_at, hash, prev_hash FROM events WHERE task_id=? ORDER BY seq ASC", (task_id,), ).fetchall() result: list[dict[str, Any]] = [] for r in rows: result.append( { "event_id": r["event_id"], "seq": r["seq"], "type": r["type"], "subject": r["subject"], "created_at": r["created_at"], "hash": r["hash"], "prev_hash": r["prev_hash"], "data": json.loads(r["event_data"]), } ) return result def list_outbox(self, status: str = "NOT_DELIVERED") -> list[dict[str, Any]]: with self._lock: rows = self.conn.execute( "SELECT * FROM outbox WHERE delivery_status=? ORDER BY created_at ASC", (status,), ).fetchall() return [dict(row) for row in rows] def get_outbox_message(self, message_id: str) -> dict[str, Any]: message_id = _validated_id(message_id, "message_id") with self._lock: row = self.conn.execute( "SELECT * FROM outbox WHERE message_id=?", (message_id,) ).fetchone() if row is None: raise NotFoundError("outbox message not found") return dict(row) def record_outbox_delivery( self, message_id: str, *, receipt: str | None, error: str | None = None, ) -> dict[str, Any]: """Only a non-empty gateway receipt can close delivery as DELIVERED.""" message_id = _validated_id(message_id, "message_id") with self._transaction(): row = self.conn.execute( "SELECT * FROM outbox WHERE message_id=?", (message_id,) ).fetchone() if row is None: raise NotFoundError("outbox message not found") if row["delivery_status"] == "DELIVERED": if receipt and row["receipt"] == receipt: return dict(row) raise ConflictError("delivery receipt conflict") status = "DELIVERED" if receipt else "NOT_DELIVERED" delivered_at = _now() if receipt else None self.conn.execute( """ UPDATE outbox SET delivery_status=?, delivered_at=?, attempts=attempts+1, error=?, receipt=? WHERE message_id=? """, (status, delivered_at, error, receipt, message_id), ) updated = self.conn.execute( "SELECT * FROM outbox WHERE message_id=?", (message_id,) ).fetchone() return dict(updated) def verify_event_chain(self, task_id: str) -> bool: events = self.list_events(task_id) prev = "" for event in events: current = _hash_hex( _canonical( { "event_id": event["event_id"], "task_id": task_id, "seq": event["seq"], "type": event["type"], "subject": event["subject"], "created_at": event["created_at"], "data": event["data"], "prev_hash": event["prev_hash"], } ) ) if event["prev_hash"] != prev or event["hash"] != current: return False prev = event["hash"] return True def add_evidence(self, task_id: str, kind: str, payload: dict[str, Any], worker_identity: str) -> str: claim = self.active_claim(task_id) claim_id = claim.claim_id if claim else "claim-unbound-shadow" envelope = EvidenceEnvelope.create( task_id=task_id, claim_id=claim_id, kind=kind, worker_identity=worker_identity, payload=payload, previous_evidence_hash=self._last_evidence_hash(task_id), ) self.add_evidence_envelope(envelope) return envelope.evidence_id def _last_evidence_hash(self, task_id: str) -> str: row = self.conn.execute( """ SELECT digest FROM evidence WHERE task_id=? ORDER BY created_at DESC,evidence_id DESC LIMIT 1 """, (task_id,), ).fetchone() return str(row["digest"]) if row else "" def add_evidence_envelope(self, envelope: EvidenceEnvelope) -> str: with self._transaction(): self.get_task(envelope.task_id) expected_previous = self._last_evidence_hash(envelope.task_id) if not envelope.verify( expected_task_id=envelope.task_id, expected_claim_id=envelope.claim_id, expected_previous_hash=expected_previous, ): raise ConflictError("evidence envelope invalid") self.conn.execute( """ INSERT INTO evidence (evidence_id,task_id,kind,payload,digest,worker_identity,created_at, claim_id,previous_evidence_hash,envelope_json) VALUES(?,?,?,?,?,?,?,?,?,?) """, ( envelope.evidence_id, envelope.task_id, envelope.kind, json.dumps(envelope.payload, ensure_ascii=False), envelope.evidence_digest, envelope.worker_identity, envelope.created_at, envelope.claim_id, envelope.previous_evidence_hash, json.dumps(envelope.to_dict(), ensure_ascii=False), ), ) return envelope.evidence_id def list_evidence(self, task_id: str) -> list[dict[str, Any]]: task_id = _validated_id(task_id, "task_id") with self._lock: rows = self.conn.execute( "SELECT * FROM evidence WHERE task_id=? ORDER BY created_at ASC", (task_id,), ).fetchall() return [dict(r) for r in rows] def verify_evidence_integrity(self, task_id: str) -> bool: previous = "" for evidence in self.list_evidence(task_id): try: raw = evidence.get("envelope_json") if not raw: return False payload = json.loads(raw) envelope = EvidenceEnvelope( schema_version=int(payload["schema_version"]), evidence_id=payload["evidence_id"], task_id=payload["task_id"], claim_id=payload["claim_id"], kind=payload["kind"], worker_identity=payload["worker_identity"], payload=payload["payload"], created_at=float(payload["created_at"]), previous_evidence_hash=payload["previous_evidence_hash"], evidence_digest=payload["evidence_digest"], signature=payload.get("signature"), attestation_ref=payload.get("attestation_ref"), ) except (KeyError, TypeError, ValueError, json.JSONDecodeError): return False if ( evidence["task_id"] != task_id or evidence["digest"] != envelope.evidence_digest or evidence["kind"] != envelope.kind or evidence["worker_identity"] != envelope.worker_identity or json.loads(evidence["payload"]) != envelope.payload or not envelope.verify( expected_task_id=task_id, expected_claim_id=envelope.claim_id, expected_previous_hash=previous, ) ): return False previous = envelope.evidence_digest return True def list_children(self, task_id: str) -> list[str]: task_id = _validated_id(task_id, "task_id") with self._lock: rows = self.conn.execute( "SELECT child_task_id FROM delegations WHERE parent_task_id=? ORDER BY created_at ASC", (task_id,), ).fetchall() return [row["child_task_id"] for row in rows] def parent_of(self, task_id: str) -> str | None: task_id = _validated_id(task_id, "task_id") with self._lock: row = self.conn.execute( "SELECT parent_task_id FROM tasks WHERE task_id=?", (task_id,) ).fetchone() return row["parent_task_id"] if row else None def _compute_depth(self, task_id: str) -> int: depth = 0 visited = set() current = task_id while current: if current in visited: raise ValidationError("parent cycle detected in task chain") visited.add(current) parent = self.parent_of(current) if parent is None: return depth depth += 1 current = parent return depth