SciMesh/scimesh/worker/config.py

"""Configuration parsing for the worker command."""

from __future__ import annotations

from dataclasses import dataclass
from math import isfinite
from pathlib import Path
import os
import socket
from typing import Mapping
from urllib.parse import urlsplit


def _clean_url(value: object | None) -> str | None:
    """Normalise an optional URL: drop a blank one, strip a trailing slash."""
    if value is None:
        return None
    text = str(value).strip()
    return text.rstrip("/") or None


def _positive_number(value: object, name: str, *, allow_zero: bool = False) -> None:
    if (
        isinstance(value, bool)
        or not isinstance(value, (int, float))
        or not isfinite(value)
        or value < 0
        or (not allow_zero and value == 0)
    ):
        qualifier = "non-negative" if allow_zero else "positive"
        raise ValueError(f"{name} must be {qualifier}")


@dataclass(frozen=True)
class WorkerConfig:
    coordinator_url: str
    worker_id: str | None
    work_dir: Path
    worker_name: str = "scimesh-worker"
    cpu_count: int = 1
    memory_mb: int | None = None
    poll_interval: float = 2.0
    request_timeout: float = 30.0
    heartbeat_interval: float = 15.0
    bearer_token: str | None = None
    # A long-lived per-user credential. When set (with userservice_url), the
    # worker exchanges it for short-lived JWTs instead of using bearer_token,
    # binding the worker to that user's account.
    worker_key: str | None = None
    userservice_url: str | None = None
    cleanup_after_seconds: float | None = None
    max_tasks: int | None = None
    exit_when_idle: bool = False
    # Distributed similarity-graph requires triangular block-pair planning and
    # is deliberately not advertised until CTX-10. A normal worker must never
    # make a multi-shard graph job appear scientifically complete.
    # The local CLI uses hyphens; the first coordinator contract used
    # underscores, so retain the search alias during migration.
    capabilities: tuple[str, ...] = (
        "similarity-search",
        "similarity_search",
    )

    def __post_init__(self) -> None:
        parsed = urlsplit(self.coordinator_url)
        if parsed.scheme not in {"http", "https"} or not parsed.hostname:
            raise ValueError("coordinator_url must be an absolute HTTP(S) URL")
        if not isinstance(self.worker_name, str) or not self.worker_name.strip():
            raise ValueError("worker_name must be non-empty")
        if self.userservice_url is not None:
            us = urlsplit(self.userservice_url)
            if us.scheme not in {"http", "https"} or not us.hostname:
                raise ValueError("userservice_url must be an absolute HTTP(S) URL")
        if self.worker_key is not None and not self.userservice_url:
            raise ValueError("worker_key requires userservice_url (SCIMESH_USERSERVICE_URL)")
        if isinstance(self.cpu_count, bool) or not isinstance(self.cpu_count, int) or self.cpu_count < 1:
            raise ValueError("cpu_count must be positive")
        if self.worker_id is not None and not isinstance(self.worker_id, str):
            raise ValueError("worker_id must be a string when set")
        if self.memory_mb is not None and (
            isinstance(self.memory_mb, bool)
            or not isinstance(self.memory_mb, int)
            or self.memory_mb < 1
        ):
            raise ValueError("memory_mb must be positive when set")
        _positive_number(self.poll_interval, "poll_interval")
        _positive_number(self.request_timeout, "request_timeout")
        _positive_number(self.heartbeat_interval, "heartbeat_interval")
        if self.cleanup_after_seconds is not None:
            _positive_number(self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True)
        if self.max_tasks is not None:
            if (
                isinstance(self.max_tasks, bool)
                or not isinstance(self.max_tasks, int)
                or self.max_tasks < 1
            ):
                raise ValueError("max_tasks must be positive when set")
        if not isinstance(self.exit_when_idle, bool):
            raise ValueError("exit_when_idle must be a boolean")
        if not self.capabilities:
            raise ValueError("capabilities cannot be empty")
        # Runner subprocesses use a task directory as their cwd. Keep the
        # configured root absolute so input/output paths remain valid there
        # even when the CLI received a convenient relative --work-dir value.
        object.__setattr__(self, "work_dir", self.work_dir.expanduser().resolve())

    @classmethod
    def from_environment(
        cls, overrides: Mapping[str, object] | None = None
    ) -> "WorkerConfig":
        """Build config from environment, allowing typed CLI values to override it."""
        values = overrides or {}

        def value(name: str, environment: str, default: object | None = None) -> object | None:
            override = values.get(name)
            return override if override is not None else os.getenv(environment, default)

        url = value("coordinator_url", "SCIMESH_COORDINATOR_URL")
        if not isinstance(url, str) or not url:
            raise ValueError("SCIMESH_COORDINATOR_URL or --coordinator-url is required")
        cleanup = value("cleanup_after_seconds", "SCIMESH_CLEANUP_AFTER_SECONDS")
        cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1)
        memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB")
        max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS")
        return cls(
            coordinator_url=url.rstrip("/"),
            worker_id=value("worker_id", "SCIMESH_WORKER_ID"),
            work_dir=Path(value("work_dir", "SCIMESH_WORK_DIR", "./scimesh-worker-data")),
            worker_name=str(value("worker_name", "SCIMESH_WORKER_NAME", socket.gethostname())),
            cpu_count=int(cpu_count),
            memory_mb=int(memory_mb) if memory_mb is not None else None,
            poll_interval=float(value("poll_interval", "SCIMESH_POLL_INTERVAL", "2")),
            request_timeout=float(value("request_timeout", "SCIMESH_REQUEST_TIMEOUT", "30")),
            heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")),
            bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"),
            worker_key=value("worker_key", "SCIMESH_WORKER_KEY"),
            userservice_url=_clean_url(value("userservice_url", "SCIMESH_USERSERVICE_URL")),
            cleanup_after_seconds=float(cleanup) if cleanup else None,
            max_tasks=int(max_tasks) if max_tasks is not None else None,
            exit_when_idle=bool(values.get("exit_when_idle", False)),
        )