apsimo 1.3.0__py3-none-any.whl

apsimo/task_queue/handlers/monitoring.py

@@ -0,0 +1,116 @@
+"""MonitoringHandler — health-check probe jobs."""
+
+from __future__ import annotations
+
+import ipaddress
+import logging
+import os
+import time
+from typing import Any, Dict
+from urllib.parse import urlsplit
+
+from apsimo.task_queue.handlers.base import JobHandler, Job
+
+logger = logging.getLogger(__name__)
+
+
+def _host_allowlisted(host: str) -> bool:
+    entries = {
+        item.strip().lower()
+        for item in os.environ.get(
+            "COLONY_MONITORING_HOST_ALLOWLIST", ""
+        ).split(",")
+        if item.strip()
+    }
+    host = host.lower().rstrip(".")
+    return host in entries or any(
+        entry.startswith("*.") and host.endswith(entry[1:])
+        for entry in entries
+    )
+
+
+def _unsafe_address(value: str) -> bool:
+    address = ipaddress.ip_address(value)
+    return not address.is_global
+
+
+async def _validate_endpoint(url: str) -> str:
+    parsed = urlsplit(str(url or ""))
+    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
+        raise ValueError("monitoring endpoint must be an http(s) URL")
+    if parsed.username is not None or parsed.password is not None:
+        raise ValueError("monitoring endpoint cannot contain credentials")
+    host = parsed.hostname.rstrip(".")
+    try:
+        address = str(ipaddress.ip_address(host))
+    except ValueError:
+        # DNS resolution followed by an ordinary HTTP client resolution is a
+        # rebinding TOCTOU. Hostnames therefore require an exact/operator
+        # wildcard allowlist entry; literal public IPs need no DNS at all.
+        if not _host_allowlisted(host):
+            raise ValueError(
+                "monitoring hostname must be explicitly configured in "
+                "COLONY_MONITORING_HOST_ALLOWLIST"
+            )
+        return url
+    if _unsafe_address(address) and not _host_allowlisted(host):
+        raise ValueError(
+            "monitoring endpoint resolves to a non-public address; configure "
+            "COLONY_MONITORING_HOST_ALLOWLIST for an intentional internal probe"
+        )
+    return url
+
+
+class MonitoringHandler(JobHandler):
+    """Run a health-check probe.
+
+    Job payload:
+        endpoint (str): HTTP URL to GET.
+        expected_status (int, default 200): Expected HTTP status.
+        timeout_secs (float, default 10): Per-request timeout.
+
+    Returns:
+        {"metrics": {"status": int, "latency_ms": float}, "alerts": list[str]}
+    """
+
+    async def execute(self, job: Job) -> Dict[str, Any]:
+        try:
+            import aiohttp
+        except ImportError as exc:
+            raise RuntimeError(
+                "MonitoringHandler requires 'aiohttp'. Install with: pip install aiohttp"
+            ) from exc
+
+        url = await _validate_endpoint(job.payload["endpoint"])
+        expected = job.payload.get("expected_status", 200)
+        timeout = job.payload.get("timeout_secs", 10.0)
+
+        t0 = time.monotonic()
+        async with aiohttp.ClientSession() as session:
+            async with session.get(
+                url,
+                timeout=aiohttp.ClientTimeout(total=timeout),
+                allow_redirects=False,
+            ) as resp:
+                latency_ms = (time.monotonic() - t0) * 1000
+                alerts = []
+                if resp.status != expected:
+                    alerts.append(
+                        f"Expected status {expected}, got {resp.status}"
+                    )
+                succeeded = not alerts
+                return {
+                    "status": "completed" if succeeded else "failed",
+                    "summary": (
+                        f"probe returned expected HTTP {resp.status}"
+                        if succeeded else "; ".join(alerts)
+                    ),
+                    "action_plane": {
+                        "state": "completed" if succeeded else "failed",
+                    },
+                    "metrics": {
+                        "status": resp.status,
+                        "latency_ms": round(latency_ms, 2),
+                    },
+                    "alerts": alerts,
+                }

apsimo/task_queue/handlers/registry.py

@@ -0,0 +1,75 @@
+"""Default handler registry for the Colony task queue."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from apsimo.task_queue.models import JobType
+from apsimo.task_queue.worker import JobHandler
+from apsimo.task_queue.handlers.monitoring import MonitoringHandler
+from apsimo.task_queue.handlers.subtask_handler import SubtaskHandler
+from apsimo.task_queue.handlers.system_maintenance import SystemMaintenanceHandler
+
+if TYPE_CHECKING:
+    from apsimo.router.router import LLMRouter
+    from apsimo.world_model.store import WorldModelStore
+    from apsimo.contacts.store import ContactStore
+    from apsimo.task_queue.handlers.inference import _InferenceGateSessionStore
+
+
+def build_default_handlers(
+    router: Optional["LLMRouter"] = None,
+    world_model_store: Optional["WorldModelStore"] = None,
+    contact_store: Optional["ContactStore"] = None,
+    response_gate: Optional[Any] = None,
+    gate_session_store: Optional["_InferenceGateSessionStore"] = None,
+    desktop_config: Optional[Any] = None,
+    browser_config: Optional[Any] = None,
+    node_id: str = "",
+) -> Dict[JobType, JobHandler]:
+    """Assemble the default set of job handlers.
+
+    Args:
+        router: LLMRouter instance. Required for INFERENCE jobs.
+                If None, InferenceHandler is omitted from the registry.
+        world_model_store: Optional WorldModelStore for context enrichment.
+        contact_store: Optional ContactStore for contact resolution.
+        response_gate: Optional ResponseGate for outbound response filtering.
+        gate_session_store: Session store adapter used by the ResponseGate.
+        desktop_config: Retired unsupported parameter; must be None.
+        browser_config: Retired unsupported parameter; must be None.
+        node_id: Legacy compatibility parameter, unused by supported handlers.
+
+    Returns:
+        Dict mapping JobType → JobHandler instance.
+    """
+    if desktop_config is not None or browser_config is not None:
+        raise ValueError(
+            "Desktop/browser queue workers are unsupported. Remove desktop_config "
+            "and browser_config; use the native runtime's registered tools."
+        )
+    handlers: Dict[JobType, JobHandler] = {
+        JobType.MONITORING: MonitoringHandler(),
+        JobType.SYSTEM_MAINTENANCE: SystemMaintenanceHandler(),
+        JobType.CUSTOM: SubtaskHandler(),
+    }
+
+    if router is not None:
+        from apsimo.task_queue.handlers.inference import (
+            InferenceHandler,
+            ThoughtOnlyInferenceHandler,
+        )
+        inference_handler = InferenceHandler(
+            router,
+            world_model_store=world_model_store,
+            contact_store=contact_store,
+            response_gate=response_gate,
+            gate_session_store=gate_session_store,
+        )
+        handlers[JobType.INFERENCE] = inference_handler
+        # ThoughtJobV1 shares only the router transport. A distinct strict
+        # handler makes it impossible for a registry mix-up to run a generic
+        # inference payload on the owner-private cognition lane.
+        handlers[JobType.THOUGHT] = ThoughtOnlyInferenceHandler(router)
+
+    return handlers

apsimo/task_queue/handlers/subtask_handler.py

@@ -0,0 +1,173 @@
+"""Custom subtask handler for the Colony task queue.
+
+Handles job_type=CUSTOM jobs dispatched by GoalQueueBridge.
+For simple subtasks, spawns a lightweight Sonnet agent to execute the task
+and returns the result as a job output dict.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from datetime import datetime, timezone
+from typing import Any, Dict
+
+from apsimo.task_queue.models import Job
+from apsimo.task_queue.worker import JobHandler
+
+logger = logging.getLogger(__name__)
+
+_SONNET_MODEL = "claude-sonnet-4-6"
+_MAX_TOKENS = 1024
+
+_SYSTEM_PROMPT = (
+    "You are an Apsimo subtask worker with NO tools: you cannot browse, run "
+    "commands, send messages, or touch external systems. Produce the work "
+    "product itself: the analysis, draft, plan, or answer the task asks for, "
+    "concise and self-contained. Never claim to have performed an external "
+    "action. If the task needs information or access you do not have, state "
+    "exactly what is missing in one line and deliver the best partial result "
+    "from what is given. Do not ask clarifying questions; make the most "
+    "reasonable assumption and note it."
+)
+
+
+class SubtaskHandler(JobHandler):
+    """Executes custom subtasks dispatched by GoalQueueBridge.
+
+    The job payload is expected to contain:
+        goal_id     — parent goal identifier
+        subtask_id  — subtask identifier
+        description — natural-language task description (preferred)
+        title / task / action / objective — fallback description keys
+
+    For each subtask the handler calls the Anthropic API with the Sonnet
+    model (restricted to text-only, no tools) and returns the model's
+    completion report as the job output.
+    """
+
+    def __init__(self, api_key: str | None = None) -> None:
+        self._api_key = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
+
+    async def execute(self, job: Job) -> Dict[str, Any]:
+        goal_id = job.payload.get("goal_id", "")
+        subtask_id = job.payload.get("subtask_id", "")
+        task_description = job.payload.get("description", "") or _build_description(job.payload)
+
+        logger.info(
+            "SubtaskHandler: executing subtask %s for goal %s",
+            subtask_id,
+            goal_id,
+        )
+
+        result = await self._run_agent(task_description, goal_id, subtask_id)
+
+        logger.info(
+            "SubtaskHandler: subtask %s finished (goal=%s, success=%s)",
+            subtask_id,
+            goal_id,
+            result.get("success"),
+        )
+        return result
+
+    async def _run_agent(
+        self,
+        task_description: str,
+        goal_id: str,
+        subtask_id: str,
+    ) -> Dict[str, Any]:
+        """Call a lightweight Sonnet agent to execute the task."""
+        if not self._api_key:
+            logger.warning(
+                "SubtaskHandler: ANTHROPIC_API_KEY not set — subtask %s skipped",
+                subtask_id,
+            )
+            return {
+                "success": False,
+                "status": "failed",
+                "action_plane": {"state": "failed"},
+                "error": "ANTHROPIC_API_KEY not configured",
+                "goal_id": goal_id,
+                "subtask_id": subtask_id,
+            }
+
+        try:
+            import anthropic
+        except ImportError:
+            logger.error("SubtaskHandler: anthropic package not installed")
+            return {
+                "success": False,
+                "status": "failed",
+                "action_plane": {"state": "failed"},
+                "error": "anthropic package not available",
+                "goal_id": goal_id,
+                "subtask_id": subtask_id,
+            }
+
+        # Try ModelRegistry utility role for model/credentials; fall back to hardcoded
+        _model = _SONNET_MODEL
+        _api_key = self._api_key
+        try:
+            from colony_cli.model_config import load_model_registry
+            _registry = load_model_registry()
+            _rt = _registry.resolve_runtime(role="utility")
+            if _rt.get("model"):
+                _model = _rt["model"]
+            if _rt.get("api_key"):
+                _api_key = _rt["api_key"]
+        except Exception:
+            pass
+
+        client = anthropic.AsyncAnthropic(api_key=_api_key)
+        try:
+            message = await client.messages.create(
+                model=_model,
+                max_tokens=_MAX_TOKENS,
+                system=_SYSTEM_PROMPT,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": (
+                            f"Goal ID: {goal_id}\n"
+                            f"Subtask ID: {subtask_id}\n"
+                            f"Task: {task_description}"
+                        ),
+                    }
+                ],
+            )
+            output_text = message.content[0].text if message.content else ""
+            return {
+                "success": True,
+                "status": "completed",
+                "summary": output_text,
+                "action_plane": {"state": "completed"},
+                "output": output_text,
+                "goal_id": goal_id,
+                "subtask_id": subtask_id,
+                "model": _model,
+                "completed_at": datetime.now(timezone.utc).isoformat(),
+            }
+        except Exception as exc:
+            logger.error(
+                "SubtaskHandler: agent call failed for subtask %s: %s",
+                subtask_id,
+                exc,
+            )
+            return {
+                "success": False,
+                "status": "failed",
+                "action_plane": {"state": "failed"},
+                "error": str(exc),
+                "goal_id": goal_id,
+                "subtask_id": subtask_id,
+            }
+
+
+def _build_description(payload: Dict[str, Any]) -> str:
+    """Construct a task description from payload when 'description' is absent."""
+    for key in ("title", "task", "action", "objective"):
+        if payload.get(key):
+            return str(payload[key])
+    skip = {"goal_id", "subtask_id"}
+    parts = [f"{k}: {v}" for k, v in payload.items() if k not in skip and v]
+    return " — ".join(parts) or "Execute subtask"

apsimo/task_queue/handlers/system_maintenance.py

@@ -0,0 +1,147 @@
+"""SystemMaintenanceHandler — safe allow-listed maintenance jobs."""
+
+from __future__ import annotations
+
+import logging
+import os
+import shutil
+from pathlib import Path
+from typing import Any, Dict, List
+
+from apsimo.task_queue.handlers.base import JobHandler, Job, run_subprocess
+
+logger = logging.getLogger(__name__)
+
+
+class SystemMaintenanceHandler(JobHandler):
+    """Run safe system maintenance commands (allow-listed only).
+
+    Job payload:
+        action (str): One of: "disk_cleanup", "log_rotate", "db_vacuum".
+        target_path (str, optional): Path scoped to the colony data directory.
+
+    Returns:
+        {"actions_taken": list[str], "errors": list[str]}
+    """
+
+    ALLOWED_ACTIONS = {"disk_cleanup", "log_rotate", "db_vacuum"}
+
+    def __init__(self, allowed_roots: List[Path] | None = None) -> None:
+        if allowed_roots is None:
+            configured = os.environ.get("COLONY_MAINTENANCE_ROOTS", "")
+            roots = [
+                Path(item.strip()).expanduser()
+                for item in configured.split(",") if item.strip()
+            ]
+            if not roots:
+                roots = [Path(
+                    os.environ.get("COLONY_STATE_DIR")
+                    or os.environ.get("COLONY_HOME", "~/.colony")
+                ).expanduser()]
+        else:
+            roots = list(allowed_roots)
+        self._allowed_roots = tuple(root.resolve() for root in roots)
+
+    def _target(self, value: str | None, *, default_name: str) -> Path:
+        root = self._allowed_roots[0]
+        raw = Path(value).expanduser() if value else root / default_name
+        if not raw.is_absolute():
+            raw = root / raw
+        resolved = raw.resolve(strict=False)
+        if not any(resolved.is_relative_to(base) for base in self._allowed_roots):
+            raise ValueError(
+                "maintenance target must stay inside COLONY_MAINTENANCE_ROOTS"
+            )
+        if resolved in self._allowed_roots:
+            raise ValueError("maintenance target cannot be a configured root")
+        return resolved
+
+    async def execute(self, job: Job) -> Dict[str, Any]:
+        action = job.payload.get("action")
+        if action not in self.ALLOWED_ACTIONS:
+            raise ValueError(
+                f"Disallowed maintenance action: {action!r}. "
+                f"Allowed: {sorted(self.ALLOWED_ACTIONS)}"
+            )
+
+        target_path: str | None = job.payload.get("target_path")
+        actions_taken: List[str] = []
+        errors: List[str] = []
+
+        if action == "db_vacuum":
+            if not target_path:
+                errors.append("db_vacuum requires 'target_path' in payload")
+            else:
+                try:
+                    path = self._target(target_path, default_name="state.db")
+                    import aiosqlite
+                    async with aiosqlite.connect(path) as db:
+                        await db.execute("VACUUM")
+                    actions_taken.append(f"vacuumed {path}")
+                    logger.info("SystemMaintenance: vacuumed %s", path)
+                except ImportError:
+                    errors.append(
+                        "db_vacuum requires 'aiosqlite'. Install with: pip install aiosqlite"
+                    )
+                except Exception as exc:
+                    errors.append(f"db_vacuum failed: {exc}")
+                    logger.warning(
+                        "SystemMaintenance: db_vacuum error on %s: %s",
+                        target_path, exc,
+                    )
+
+        elif action == "disk_cleanup":
+            path = self._target(target_path, default_name="tmp")
+            try:
+                cleaned = 0
+                if path.is_dir():
+                    for item in path.iterdir():
+                        try:
+                            if item.is_symlink() or item.is_file():
+                                item.unlink()
+                                cleaned += 1
+                            elif item.is_dir():
+                                shutil.rmtree(item)
+                                cleaned += 1
+                        except Exception as exc:
+                            errors.append(f"Could not remove {item}: {exc}")
+                actions_taken.append(f"disk_cleanup: removed {cleaned} items from {path}")
+                logger.info("SystemMaintenance: disk_cleanup removed %d items from %s", cleaned, path)
+            except Exception as exc:
+                errors.append(f"disk_cleanup failed: {exc}")
+                logger.warning("SystemMaintenance: disk_cleanup error: %s", exc)
+
+        elif action == "log_rotate":
+            path = self._target(target_path, default_name="logs")
+            try:
+                rotated = 0
+                if path.is_dir():
+                    for log_file in path.glob("*.log"):
+                        rotated_name = log_file.with_suffix(".log.bak")
+                        try:
+                            log_file.rename(rotated_name)
+                            rotated += 1
+                        except Exception as exc:
+                            errors.append(f"Could not rotate {log_file}: {exc}")
+                actions_taken.append(f"log_rotate: rotated {rotated} log files in {path}")
+                logger.info("SystemMaintenance: log_rotate rotated %d files in %s", rotated, path)
+            except Exception as exc:
+                errors.append(f"log_rotate failed: {exc}")
+                logger.warning("SystemMaintenance: log_rotate error: %s", exc)
+
+        succeeded = not errors
+        return {
+            "status": "completed" if succeeded else "failed",
+            "summary": (
+                "; ".join(actions_taken)
+                if succeeded else "; ".join(errors)
+            ),
+            "operations": [
+                "delete" if action == "disk_cleanup" else "write"
+            ],
+            "actions_taken": actions_taken,
+            "errors": errors,
+            "action_plane": {
+                "state": "completed" if succeeded else "failed",
+            },
+        }

apsimo/task_queue/mesh_integration.py

@@ -0,0 +1,111 @@
+"""Mesh layer integration for the Colony distributed task queue.
+
+Bridges SWIM node-dead events and mesh role-change events into
+queue scheduler actions. Wires the task queue into the existing
+colony/mesh/ layer without modifying it.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Optional
+
+from apsimo.task_queue.queue_manager import QueueManager
+from apsimo.task_queue.scheduler import Scheduler
+
+logger = logging.getLogger(__name__)
+
+
+_BACKGROUND_TASKS: set = set()
+
+
+def _retain(task) -> None:
+    """Keep a strong reference to a fire-and-forget task until it finishes.
+
+    The event loop holds only a weak reference; an unreferenced task can be
+    garbage-collected (cancelled) mid-flight — for the Sovereign scheduler
+    loop that means job scheduling silently halts with no error.
+    """
+    _BACKGROUND_TASKS.add(task)
+    task.add_done_callback(_BACKGROUND_TASKS.discard)
+
+
+class QueueMeshEventHandler:
+    """Handles mesh-layer events that affect the task queue.
+
+    Subscribe this handler to the Colony event bus for MeshEvent types.
+
+    Responsibilities:
+    - On NODE_DEAD: immediately abandon all jobs claimed by the dead node.
+    - On ROLE_CHANGE to SOVEREIGN (self): assume the scheduler role.
+    """
+
+    def __init__(
+        self,
+        queue: QueueManager,
+        scheduler: Scheduler,
+        own_node_id: str,
+        event_bus: Optional[Any] = None,
+    ) -> None:
+        self._queue = queue
+        self._scheduler = scheduler
+        self._own_node_id = own_node_id
+        self._event_bus = event_bus
+
+    async def on_node_dead(self, node_id: str) -> None:
+        """Immediately abandon all jobs claimed by the dead node.
+
+        Called when SWIM emits a NODE_DEAD event. Does not wait for the
+        heartbeat timeout — fast-path abandonment.
+        """
+        abandoned = await self._queue.abandon_jobs_for_node(node_id)
+        if abandoned:
+            logger.info(
+                "QueueMeshEventHandler: abandoned %d jobs from dead node %s: %s",
+                len(abandoned), node_id, abandoned,
+            )
+        # Trigger immediate redistribution
+        await self._queue.requeue_retryable_jobs(
+            __import__("datetime").datetime.now(__import__("datetime").timezone.utc)
+        )
+
+    async def on_role_change(self, new_role: str, node_id: str) -> None:
+        """Handle mesh role-change events.
+
+        If this node becomes SOVEREIGN (Queen), start the scheduler loop.
+        """
+        if new_role == "sovereign" and node_id == self._own_node_id:
+            logger.info(
+                "QueueMeshEventHandler: this node (%s) became Sovereign — "
+                "starting scheduler", self._own_node_id,
+            )
+            import asyncio
+            _retain(asyncio.create_task(self._scheduler.run()))
+
+    def handle_mesh_event(self, event: Any) -> None:
+        """Synchronous event bus callback. Dispatches to async handlers."""
+        import asyncio
+        node_id = getattr(event, "node_id", "")
+        event_type = getattr(event, "event_type", "")
+
+        if event_type == "node_dead" or event_type == "NODE_DEAD":
+            _retain(asyncio.create_task(self.on_node_dead(node_id)))
+        elif event_type in {"role_changed", "ROLE_CHANGE"}:
+            new_role = ""
+            new_role_attr = getattr(event, "new_role", None)
+            if new_role_attr is not None:
+                new_role = str(new_role_attr.value if hasattr(new_role_attr, "value") else new_role_attr)
+            _retain(asyncio.create_task(self.on_role_change(new_role, node_id)))
+
+    def subscribe_to_event_bus(self) -> None:
+        """Register this handler on the Colony event bus if one is configured."""
+        if self._event_bus is None:
+            return
+        from apsimo.events.types import MeshEvent
+        self._event_bus.subscribe(
+            handler=self.handle_mesh_event,
+            event_types=[MeshEvent],
+        )
+        logger.info(
+            "QueueMeshEventHandler: subscribed to MeshEvent on event bus"
+        )