apsimo 1.3.0__py3-none-any.whl

apsimo/task_queue/routing.py

@@ -0,0 +1,287 @@
+"""Server-owned routing contracts for the shared ``agent_action`` job type."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+
+AGENT_SYNC_ROUTE = "agent_sync:v1"
+ACTION_PLANE_ROUTE = "action_plane:v1"
+WORK_ORDER_ROUTE = "work_order:v1"
+HERMES_RUN_ROUTE = "hermes_run:v1"
+THOUGHT_ROUTE = "thought_engine:v1"
+
+AGENT_ACTION_ROUTE_CAPABILITIES = frozenset({
+    AGENT_SYNC_ROUTE,
+    ACTION_PLANE_ROUTE,
+    WORK_ORDER_ROUTE,
+    HERMES_RUN_ROUTE,
+})
+
+_ROUTE_OWNER_ENV = {
+    AGENT_SYNC_ROUTE: "COLONY_AGENT_SYNC_WORKER_NODE_ID",
+    HERMES_RUN_ROUTE: "COLONY_HERMES_RUN_WORKER_NODE_ID",
+    ACTION_PLANE_ROUTE: "COLONY_ACTION_PLANE_WORKER_NODE_ID",
+    WORK_ORDER_ROUTE: "COLONY_ACTION_PLANE_WORKER_NODE_ID",
+}
+
+
+def generic_agent_job_claims_enabled() -> bool:
+    """Resolve the global generic agent-action containment switch."""
+
+    value = os.environ.get(
+        "COLONY_AGENT_JOB_CLAIMS_ENABLED", "true"
+    ).strip().lower()
+    if value in {"1", "true", "yes", "on"}:
+        return True
+    if value in {"0", "false", "no", "off"}:
+        return False
+    # An invalid containment value never silently enables generic claims.
+    return False
+
+
+def expected_agent_action_routes(job: Any) -> tuple[str, ...]:
+    """Derive exact executor lanes from server-known schema/action policy."""
+
+    payload = getattr(job, "payload", {}) or {}
+    if payload.get("schema") == "WorkOrderV1":
+        from apsimo.work_orders import WorkOrderV1
+
+        order = WorkOrderV1.from_payload(payload)
+        if str(getattr(job, "job_id", "")) != order.work_order_id:
+            raise ValueError(
+                "WorkOrder queue job ID does not match authority digest"
+            )
+        # WorkOrders are receipt-aware Action Plane work. Requiring both
+        # capabilities prevents a generic protocol-aware worker from racing
+        # the deployment's authorized effect executor.
+        return (WORK_ORDER_ROUTE, ACTION_PLANE_ROUTE)
+
+    action_hint = str(payload.get("action_hint") or "").strip()
+    if action_hint:
+        from apsimo.initiatives.action_registry import (
+            RiskTier,
+            get_action,
+        )
+
+        spec = get_action(action_hint)
+        if spec is not None:
+            canonical_risk = spec.risk.value
+            declared_risk = str(payload.get("risk") or "").strip().lower()
+            if declared_risk and declared_risk != canonical_risk:
+                raise ValueError(
+                    "agent_action risk does not match server registry"
+                )
+            if spec.risk is RiskTier.READ_ONLY:
+                from apsimo.task_queue.governor import (
+                    job_declares_effect,
+                )
+                if job_declares_effect(job):
+                    raise ValueError(
+                        "read-only agent action declares effect capabilities"
+                    )
+                return (AGENT_SYNC_ROUTE,)
+            return (ACTION_PLANE_ROUTE,)
+        raise ValueError("agent_action action_hint is not server-registered")
+
+    from apsimo.task_queue.governor import job_declares_effect
+
+    if job_declares_effect(job):
+        raise ValueError(
+            "effectful agent_action has no server-registered executor route"
+        )
+    if payload.get("schema") == "HermesRunV1":
+        return (HERMES_RUN_ROUTE,)
+    raise ValueError(
+        "agent_action requires WorkOrderV1, HermesRunV1, or a registered action"
+    )
+
+
+def bind_agent_action_routes(
+    job: Any,
+    *,
+    allow_legacy_migration: bool = False,
+) -> tuple[str, ...]:
+    """Add derived route requirements and reject caller-selected mismatch."""
+
+    expected = expected_agent_action_routes(job)
+    payload = dict(getattr(job, "payload", None) or {})
+    action_hint = str(payload.get("action_hint") or "").strip()
+    if action_hint and payload.get("schema") != "WorkOrderV1":
+        from apsimo.initiatives.action_registry import get_action
+
+        spec = get_action(action_hint)
+        if spec is not None:
+            payload["risk"] = spec.risk.value
+            job.payload = payload
+    capabilities = list(getattr(job, "capabilities", None) or ())
+    names = {
+        str(getattr(capability, "name", capability))
+        for capability in capabilities
+    }
+    supplied_routes = names & AGENT_ACTION_ROUTE_CAPABILITIES
+    if (
+        supplied_routes
+        and supplied_routes != set(expected)
+        and not (
+            allow_legacy_migration
+            and payload.get("schema") == "WorkOrderV1"
+        )
+    ):
+        raise ValueError(
+            "agent_action route capabilities do not match server policy"
+        )
+
+    from apsimo.task_queue.models import JobCapabilityRequirement
+
+    # Caller/legacy metadata can never weaken a route by marking it preferred,
+    # adding a nonsensical numeric minimum, or duplicating it.  Replace every
+    # route entry with one canonical required requirement.
+    if payload.get("schema") == "WorkOrderV1":
+        from apsimo.work_orders import WorkOrderV1
+
+        order = WorkOrderV1.from_payload(payload)
+        # Compatibility fields consumed by the host are reconstructed from the
+        # verified authority object. Unknown executable-looking extras and
+        # caller-edited description/context never survive queue admission.
+        payload = order.payload()
+        job.payload = payload
+        capability_names = (
+            WORK_ORDER_ROUTE,
+            ACTION_PLANE_ROUTE,
+            *order.capability_allowlist,
+        )
+        # The embedded authority allowlist is the complete executor contract;
+        # arbitrary/missing queue metadata must not broaden or weaken it.
+        capabilities = []
+    else:
+        capability_names = expected
+        capabilities = [
+            capability for capability in capabilities
+            if str(getattr(capability, "name", capability))
+            not in AGENT_ACTION_ROUTE_CAPABILITIES
+        ]
+    capabilities.extend(
+        JobCapabilityRequirement(
+            name=name,
+            minimum=None,
+            preferred=False,
+        )
+        for name in dict.fromkeys(capability_names)
+    )
+    job.capabilities = capabilities
+    tags = dict(getattr(job, "tags", None) or {})
+    tags["agent_action_route"] = "+".join(expected)
+    if payload.get("schema") == "WorkOrderV1":
+        tags.update({
+            "schema": order.schema,
+            "project_id": order.project_id,
+            "step_id": order.step_id,
+            "risk_class": order.risk_class,
+            "idempotency_key": order.idempotency_key,
+            "work_order_digest": order.work_order_digest,
+            "work_order_version": str(order.version),
+            "executor_protocol": WORK_ORDER_ROUTE,
+            "action_result_contract": "ExecutionResultV1",
+        })
+    elif expected == (ACTION_PLANE_ROUTE,):
+        from apsimo.initiatives.approval_authority import (
+            build_action_binding,
+        )
+
+        binding = build_action_binding(
+            job_id=str(getattr(job, "job_id", "")),
+            job_type=str(
+                getattr(getattr(job, "job_type", None), "value", "")
+                or getattr(job, "job_type", "")
+            ),
+            payload=payload,
+        )
+        tags.update({
+            "action_digest": binding.action_digest,
+            "action_result_contract": "ActionReceiptAttestationV1",
+        })
+    owners = {
+        os.environ.get(_ROUTE_OWNER_ENV[route], "").strip()
+        for route in expected
+        if os.environ.get(_ROUTE_OWNER_ENV[route], "").strip()
+    }
+    if len(owners) > 1:
+        raise ValueError("agent_action route owners conflict")
+    if owners:
+        tags["agent_action_route_node"] = next(iter(owners))
+    else:
+        tags.pop("agent_action_route_node", None)
+    job.tags = tags
+    return expected
+
+
+def worker_route_capabilities(*, sync: bool, text: bool) -> set[str]:
+    """Build a generic non-effectful worker's explicit route set."""
+
+    routes: set[str] = set()
+    if sync:
+        routes.add(AGENT_SYNC_ROUTE)
+    if text:
+        routes.add(HERMES_RUN_ROUTE)
+    return routes
+
+
+def bind_thought_route(job: Any) -> str:
+    """Canonicalize the private ThoughtJobV1 worker route."""
+
+    payload = getattr(job, "payload", {}) or {}
+    from apsimo.cognition.goal_spine import ThoughtJobV1
+
+    try:
+        thought = ThoughtJobV1.from_payload(payload)
+    except (KeyError, TypeError, ValueError) as exc:
+        raise ValueError(f"invalid ThoughtJobV1 payload: {exc}") from exc
+    if str(getattr(job, "job_id", "")) != thought.thought_job_id:
+        raise ValueError("thought queue job ID does not match authority digest")
+    from apsimo.task_queue.models import JobCapabilityRequirement
+
+    job.capabilities = [
+        JobCapabilityRequirement(name="cognition_scoped"),
+        JobCapabilityRequirement(name=THOUGHT_ROUTE),
+    ]
+    tags = dict(getattr(job, "tags", None) or {})
+    tags["thought_route"] = THOUGHT_ROUTE
+    tags.update({
+        "schema": thought.schema,
+        "concern_id": thought.concern_id,
+        "viewer_scope": thought.viewer_scope,
+        "shareability": thought.shareability,
+        "thought_job_digest": thought.thought_job_digest,
+        "risk_class": "read_only",
+    })
+    owner = os.environ.get("COLONY_THOUGHT_WORKER_NODE_ID", "").strip()
+    if owner:
+        tags["thought_route_node"] = owner
+    else:
+        tags.pop("thought_route_node", None)
+    job.tags = tags
+    return THOUGHT_ROUTE
+
+
+def configured_route_owners() -> dict[str, str | None]:
+    """Return deployment-pinned owners by logical executor lane."""
+
+    return {
+        "agent_sync": os.environ.get(
+            "COLONY_AGENT_SYNC_WORKER_NODE_ID", ""
+        ).strip() or None,
+        "action_plane": os.environ.get(
+            "COLONY_ACTION_PLANE_WORKER_NODE_ID", ""
+        ).strip() or None,
+        "work_order": os.environ.get(
+            "COLONY_ACTION_PLANE_WORKER_NODE_ID", ""
+        ).strip() or None,
+        "hermes_run": os.environ.get(
+            "COLONY_HERMES_RUN_WORKER_NODE_ID", ""
+        ).strip() or None,
+        "thought": os.environ.get(
+            "COLONY_THOUGHT_WORKER_NODE_ID", ""
+        ).strip() or None,
+    }

apsimo/task_queue/scheduler.py

@@ -0,0 +1,252 @@
+"""Distributed Task Queue — Queen Scheduler.
+
+Runs on the Sovereign node. Drives five scheduling phases per tick:
+1. Expire past deadlines
+2. Abandon silent jobs (heartbeat timeout)
+3. Requeue retryable jobs
+4. Reconcile explicit governance holds
+5. Unblock ready jobs (deps met)
+6. Assign queued jobs to available workers
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from datetime import datetime, timezone
+from typing import Any, Callable, Optional
+
+from apsimo.task_queue.models import WorkerCapabilities
+from apsimo.task_queue.queue_manager import QueueManager
+
+logger = logging.getLogger(__name__)
+
+
+class Scheduler:
+    """Priority-aware, capability-matching job scheduler.
+
+    Runs on the Sovereign node. The Regent runs a passive replica
+    and takes over if the Sovereign is unavailable.
+
+    Usage::
+
+        scheduler = Scheduler(queue=queue_manager)
+        task = asyncio.create_task(scheduler.run())
+        ...
+        await scheduler.stop()
+        await task
+    """
+
+    def __init__(
+        self,
+        queue: QueueManager,
+        tick_interval_secs: float = 2.0,
+        heartbeat_timeout_secs: Optional[float] = None,
+        claim_timeout_secs: float = 30.0,
+        no_worker_warning_secs: float = 300.0,
+        event_bus: Optional[Any] = None,
+        readiness_callback: Optional[Callable[[bool, str], None]] = None,
+        approval_timeout_hours: Optional[float] = None,
+    ) -> None:
+        self._queue = queue
+        self._tick = tick_interval_secs
+        if heartbeat_timeout_secs is not None:
+            # Backward-compatible constructor overrides update the queue-owned
+            # value rather than creating a second liveness interpretation.
+            self._queue.configure_worker_heartbeat_ttl(
+                heartbeat_timeout_secs
+            )
+        self._claim_timeout = claim_timeout_secs
+        self._no_worker_warning = no_worker_warning_secs
+        self._event_bus = event_bus
+        self._readiness_callback = readiness_callback
+        if approval_timeout_hours is None:
+            try:
+                import os
+                approval_timeout_hours = float(os.environ.get(
+                    "COLONY_APPROVAL_TIMEOUT_HOURS", "72"
+                ))
+            except (TypeError, ValueError):
+                approval_timeout_hours = 72.0
+        if approval_timeout_hours <= 0:
+            raise ValueError("approval_timeout_hours must be positive")
+        self._approval_timeout_hours = approval_timeout_hours
+        self._running = False
+        self._last_tick_at: Optional[str] = None
+        self._last_error: Optional[str] = None
+        self._tick_count = 0
+
+    @property
+    def worker_heartbeat_ttl_secs(self) -> float:
+        return self._queue.worker_heartbeat_ttl_secs
+
+    @property
+    def health(self) -> dict:
+        return {
+            "running": self._running,
+            "healthy": self._last_error is None,
+            "last_tick_at": self._last_tick_at,
+            "last_error": self._last_error,
+            "tick_count": self._tick_count,
+        }
+
+    async def run(self) -> None:
+        """Main scheduling loop. Blocks until stop() is called."""
+        self._running = True
+        logger.info("Scheduler started (tick=%.1fs)", self._tick)
+        while self._running:
+            try:
+                await self.tick_once()
+            except Exception as exc:
+                self._last_error = str(exc)[:500]
+                if self._readiness_callback is not None:
+                    self._readiness_callback(False, self._last_error)
+                logger.exception("Scheduler tick failed")
+            await asyncio.sleep(self._tick)
+        logger.info("Scheduler stopped")
+
+    async def stop(self) -> None:
+        """Signal the scheduling loop to exit after the current tick."""
+        self._running = False
+        if self._readiness_callback is not None:
+            self._readiness_callback(False, "scheduler_stopped")
+
+    async def tick_once(self) -> None:
+        """Public single-tick for testing or manual scheduling."""
+        await self._tick_once()
+        self._last_tick_at = datetime.now(timezone.utc).isoformat()
+        self._last_error = None
+        self._tick_count += 1
+        if self._readiness_callback is not None:
+            self._readiness_callback(True, "scheduler_tick_ok")
+
+    async def _tick_once(self) -> None:
+        now = datetime.now(timezone.utc)
+
+        # Expire bounded WorkControl acknowledgement leases first. Otherwise
+        # a dead/uncooperative worker's pending stop could suppress the very
+        # timeout/failure transition needed to close its attempt.
+        expired_controls = await self._queue.reconcile_expired_work_controls(
+            now,
+        )
+        if expired_controls:
+            logger.info(
+                "Scheduler: reconciled %d expired WorkControl commands",
+                expired_controls,
+            )
+
+        # Phase 1: Expire deadlines
+        expired = await self._queue.expire_past_deadlines(now)
+        if expired:
+            logger.info("Scheduler: expired %d deadline-past jobs", expired)
+
+        execution_timeouts = await self._queue.expire_execution_timeouts(now)
+        if execution_timeouts:
+            logger.info(
+                "Scheduler: expired %d execution timeouts",
+                execution_timeouts,
+            )
+
+        approval_repairs = await self._queue.reconcile_blocked_approval_authority()
+        if approval_repairs:
+            logger.info(
+                "Scheduler: reconciled %d owner-approval holds",
+                approval_repairs,
+            )
+
+        approval_timeouts = await self._queue.expire_blocked_approvals(
+            now, timeout_hours=self._approval_timeout_hours,
+        )
+        if approval_timeouts:
+            logger.info(
+                "Scheduler: expired %d owner-approval holds",
+                approval_timeouts,
+            )
+
+        # Phase 2: Detect abandoned jobs (heartbeat timeout)
+        abandoned = await self._queue.abandon_silent_jobs(
+            now,
+            timeout_secs=self.worker_heartbeat_ttl_secs,
+            claim_timeout_secs=self._claim_timeout,
+        )
+        if abandoned:
+            logger.info("Scheduler: abandoned %d silent jobs", abandoned)
+
+        # Phase 3: Redistribute abandoned/retryable jobs
+        requeued = await self._queue.requeue_retryable_jobs(now)
+        if requeued:
+            logger.info("Scheduler: requeued %d retryable jobs", requeued)
+
+        # Phase 4: Re-evaluate only explicit boundary/governor holds. This
+        # safely handles directive lifts and mode rollback without conflating
+        # them with dependency or approval state.
+        governance_released = await self._queue.reconcile_governance_holds()
+        if governance_released:
+            logger.info(
+                "Scheduler: released %d governance-held jobs",
+                governance_released,
+            )
+
+        # Durable competence/governor events survive cancellation and process
+        # restart. Drain them independently from worker execution enablement.
+        # Governor delivery may involve an LLM or outbound notice. Schedule
+        # its bounded reconciler without ever stalling lease/deadline phases.
+        self._queue._schedule_worker_outcome_drain()
+
+        # Phase 5: Unblock jobs whose dependencies are now met
+        unblocked = await self._queue.unblock_ready_jobs()
+        if unblocked:
+            logger.info("Scheduler: unblocked %d dependent jobs", unblocked)
+
+        # Reconcile source-owned runtime holds after dependencies so a
+        # rollback/re-enable cycle converges before worker notification.  The
+        # actual claim/start boundaries independently re-evaluate the fence.
+        source_reconciler = getattr(
+            self._queue, "reconcile_runtime_claim_holds", None,
+        )
+        source_released = (
+            await source_reconciler() if callable(source_reconciler) else 0
+        )
+        if source_released:
+            logger.info(
+                "Scheduler: released %d source-runtime-held jobs",
+                source_released,
+            )
+
+        # Phase 6: Assign QUEUED jobs to available workers
+        await self._assign_queued_jobs(now)
+
+    async def _assign_queued_jobs(self, now: datetime) -> None:
+        queued = await self._queue.get_queued_jobs_sorted(now)
+        if not queued:
+            return
+
+        workers = await self._queue.get_available_workers()
+
+        for job in queued:
+            best_worker: Optional[WorkerCapabilities] = None
+            best_score = 0.0
+
+            for worker in workers:
+                # Headroom bonus: prefer workers with more open slots
+                running_count = round(worker.load * worker.max_concurrent)
+                headroom = worker.max_concurrent - running_count
+                score = worker.affinity_score(job)
+                if headroom >= 2:
+                    score += 0.05 * (headroom - 1)
+                if score > best_score:
+                    best_score = score
+                    best_worker = worker
+
+            if best_worker is None:
+                if job.deadline is not None:
+                    remaining = (job.deadline - now).total_seconds()
+                    if 0 < remaining < self._no_worker_warning:
+                        logger.warning(
+                            "No capable worker for job %s (type=%s, required=%s, "
+                            "deadline in %.0fs)",
+                            job.job_id, job.job_type,
+                            job.required_capabilities(), remaining,
+                        )
+            else:
+                await self._queue.notify_worker(best_worker.node_id, job.job_id)