VoiceStudio/backend/worker/agent.py

"""Worker mode — the other half of the feature.

A worker is the ordinary backend with this agent running alongside it. That is
the whole point of not writing a slim agent: the engines, sidecar venvs, model
downloads, and VRAM budgeting the executor needs are already there.

The interesting problem here is bootstrapping trust. The control plane has a
self-signed certificate, so the worker has nothing to validate it against —
except the fingerprint baked into the enrollment token. So on first contact the
worker fetches the certificate the server presents, checks it against that
fingerprint, and only then uses it as the *sole* trusted root for every later
connection. Trust on first use, with the token as the anchor that makes the
"first use" safe.

If the fingerprint does not match, the agent stops. It does not warn and
continue: a mismatch is precisely the attack pinning exists to catch.
"""
from __future__ import annotations

import asyncio
import base64
import binascii
import errno
import hashlib
import json
import logging
import os
import ssl
import tempfile
from typing import Optional
from urllib.parse import urlsplit

from worker.async_utils import to_thread_and_drain_on_cancel

logger = logging.getLogger("omnivoice.worker")

# How often the idle-engine sweep runs. Well under the ten-minute idle
# threshold it enforces, so a model is released promptly after it goes cold
# rather than up to a full interval later. Override with
# OMNIVOICE_IDLE_SWEEP_SECONDS when shortening the threshold for testing —
# leaving the interval at 60s while the threshold is 30s means waiting a full
# minute to observe a thirty-second rule, which reads as a broken sweep.
def _sweep_seconds_from_env(default: float = 60.0, floor: float = 1.0) -> float:
    raw = (os.environ.get("OMNIVOICE_IDLE_SWEEP_SECONDS") or "").strip()
    if not raw:
        return default
    try:
        value = float(raw)
    except ValueError:
        logger.warning("Ignoring OMNIVOICE_IDLE_SWEEP_SECONDS=%r: not a number.", raw)
        return default
    if value < floor:
        logger.warning(
            "Ignoring OMNIVOICE_IDLE_SWEEP_SECONDS=%s: below the %ss floor.", value, floor
        )
        return default
    return value


IDLE_SWEEP_INTERVAL_SECONDS = _sweep_seconds_from_env()


_WORKER_MODE_SETTING = "worker_mode_enabled"
_LOWER_HEX = frozenset("0123456789abcdef")
_ENROLLMENT_MANIFEST_SCHEMA = 1


class EnrollmentStateError(RuntimeError):
    """A committed enrollment exists but cannot be trusted or used."""


class EnrollmentRollbackError(EnrollmentStateError):
    """A failed join could not restore the enrollment it replaced."""


def worker_mode_enabled() -> bool:
    """Is this machine lending its GPU to someone else's control plane?

    Environment first so a headless box can be a worker with no UI at all, then
    settings for the desktop case — the same precedence as
    ``service.remote_workers_enabled``. The settings half is what lets a user
    join from the app and still be a worker after a restart; before it, joining
    meant setting OMNIVOICE_WORKER_MODE and OMNIVOICE_WORKER_TOKEN by hand and
    relaunching, which is the whole reason the feature was unreachable in the
    UI.
    """
    env = (os.environ.get("OMNIVOICE_WORKER_MODE") or "").strip().lower()
    if env in ("1", "true", "yes", "on"):
        return True
    if env in ("0", "false", "no", "off"):
        return False
    try:
        from services import settings_store  # noqa: PLC0415

        stored = (settings_store.get_text(_WORKER_MODE_SETTING, "") or "").strip().lower()
        return stored in ("1", "true", "yes", "on")
    except Exception:
        return False


def set_worker_mode_enabled(enabled: bool) -> None:
    from services import settings_store  # noqa: PLC0415

    settings_store.set_text(_WORKER_MODE_SETTING, "true" if enabled else "false")


_ENDPOINT_SETTING = "worker_endpoint"


def _stored_endpoint() -> str:
    try:
        from services import settings_store  # noqa: PLC0415

        return (settings_store.get_text(_ENDPOINT_SETTING, "") or "").strip()
    except Exception:
        return ""


def _remember_endpoint(endpoint: str) -> None:
    if not endpoint:
        return
    try:
        from services import settings_store  # noqa: PLC0415

        settings_store.set_text(_ENDPOINT_SETTING, endpoint)
    except Exception:
        logger.debug("Could not remember the control-plane endpoint.", exc_info=True)


def _snapshot_file_generation(path: str, description: str) -> dict:
    """Capture exact bytes/presence for one file a join may overwrite."""
    try:
        with open(path, "rb") as fh:
            content = fh.read()
        return {"path": path, "present": True, "bytes": content}
    except FileNotFoundError:
        return {"path": path, "present": False, "bytes": None}
    except OSError as exc:
        raise EnrollmentStateError(
            f"The current worker {description} cannot be backed up safely. "
            "Fix its file permissions, then try joining again."
        ) from exc


def snapshot_enrollment() -> dict:
    """Everything a join overwrites, kept so a failed one can be undone."""
    locations = _paths()
    manifest_path = _enrollment_manifest_path(locations)
    manifest_snapshot = _snapshot_file_generation(manifest_path, "enrollment")
    root = os.path.dirname(os.path.abspath(manifest_path))
    file_snapshots = {
        "pinned_cert": _snapshot_file_generation(
            locations["pinned_cert"], "pinned certificate"
        ),
        "worker_id": _snapshot_file_generation(
            locations.get("worker_id") or os.path.join(root, "worker-id"),
            "identity",
        ),
        "token_hash": _snapshot_file_generation(
            locations.get("enrollment_token_hash")
            or os.path.join(root, "enrollment-token.sha256"),
            "enrollment-token marker",
        ),
    }
    preserve_enrollment_files = False
    try:
        manifest = _load_enrollment_manifest(manifest_path)
    except EnrollmentStateError:
        # A fresh explicit join can replace corrupt committed state safely.
        # Until acceptance nothing is written, so a failed repair should leave
        # those files untouched rather than "restoring" stale legacy mirrors.
        manifest = None
        preserve_enrollment_files = True
    certificate: Optional[bytes] = manifest["certificate"] if manifest else None
    if certificate is None and not preserve_enrollment_files:
        certificate = file_snapshots["pinned_cert"]["bytes"]
    try:
        from services import settings_store  # noqa: PLC0415

        worker_mode_setting_present, stored_mode = settings_store.get_text_state(
            _WORKER_MODE_SETTING
        )
        endpoint_setting_present, stored_endpoint_setting = (
            settings_store.get_text_state(_ENDPOINT_SETTING)
        )
    except Exception as exc:
        raise EnrollmentStateError(
            "The current worker settings cannot be backed up safely. "
            "Fix the settings database, then try joining again."
        ) from exc
    return {
        "certificate": certificate,
        "endpoint": (
            manifest["endpoint"]
            if manifest
            else ("" if preserve_enrollment_files else stored_endpoint_setting.strip())
        ),
        "preserve_enrollment_files": preserve_enrollment_files,
        "manifest_path": manifest_path,
        "manifest_present": manifest_snapshot["present"],
        "manifest_bytes": manifest_snapshot["bytes"],
        "manifest_snapshot_captured": True,
        "file_snapshots": file_snapshots,
        "endpoint_setting": stored_endpoint_setting,
        "endpoint_setting_present": endpoint_setting_present,
        "agent_endpoint": getattr(agent, "endpoint", ""),
        # The RAW setting, not worker_mode_enabled(): restoring "" as "false"
        # would persist a decision the machine never made.
        "worker_mode": stored_mode,
        "worker_mode_setting_present": worker_mode_setting_present,
    }


def _restore_file_generation(snapshot: dict) -> None:
    path = snapshot["path"]
    if snapshot.get("present"):
        content = snapshot.get("bytes")
        if not isinstance(content, bytes):
            raise OSError("the enrollment snapshot is incomplete")
        _replace_enrollment_manifest(path, content)
        return
    try:
        os.remove(path)
    except FileNotFoundError:
        return
    _fsync_parent_directory(os.path.dirname(os.path.abspath(path)))


async def _abort_enrollment_rollback(exc: BaseException, storage: str) -> None:
    message = (
        "The previous worker enrollment could not be restored safely. "
        f"Worker mode remains stopped; fix the {storage}, then try again."
    )
    try:
        await agent.stop()
    except Exception:
        logger.warning(
            "Worker agent also failed to stop after rollback failed.",
            exc_info=True,
        )
    agent.endpoint = ""
    agent.last_error = message
    raise EnrollmentRollbackError(message) from exc


async def restore_enrollment(previous: dict) -> None:
    """Put a machine back the way a failed rejoin found it.

    Compatibility mirrors are best effort, but the authoritative manifest is
    not: reconnecting after that restoration fails could silently run the
    replacement enrollment that this request just reported as rejected.
    """
    locations = _paths()
    manifest_path = previous.get("manifest_path") or _enrollment_manifest_path(locations)
    if previous.get("manifest_snapshot_captured"):
        try:
            _restore_file_generation(
                {
                    "path": manifest_path,
                    "present": previous.get("manifest_present"),
                    "bytes": previous.get("manifest_bytes"),
                }
            )
        except OSError as exc:
            await _abort_enrollment_rollback(exc, "enrollment storage")

    file_snapshots = previous.get("file_snapshots")
    if isinstance(file_snapshots, dict):
        for description, snapshot in file_snapshots.items():
            try:
                _restore_file_generation(snapshot)
            except OSError as exc:
                if not previous.get("manifest_present"):
                    await _abort_enrollment_rollback(
                        exc, f"worker {description.replace('_', ' ')} storage"
                    )
                logger.warning(
                    "Could not restore the previous worker %s.",
                    description.replace("_", " "),
                    exc_info=True,
                )
    else:
        certificate = previous.get("certificate")
        path = locations["pinned_cert"]
        if not previous.get("preserve_enrollment_files"):
            try:
                if certificate is None:
                    if os.path.exists(path):
                        os.remove(path)
                else:
                    _durable_makedirs(os.path.dirname(path))
                    with open(path, "wb") as fh:
                        fh.write(certificate)
            except OSError as exc:
                if not previous.get("manifest_present"):
                    await _abort_enrollment_rollback(
                        exc, "worker pinned-certificate storage"
                    )
                logger.warning(
                    "Could not restore the previous pinned certificate.", exc_info=True
                )

    has_exact_settings_snapshot = "worker_mode_setting_present" in previous
    if "endpoint_setting" in previous:
        try:
            from services import settings_store  # noqa: PLC0415

            if previous.get("endpoint_setting_present"):
                settings_store.set_text(
                    _ENDPOINT_SETTING, previous.get("endpoint_setting") or ""
                )
            else:
                settings_store.clear_text(_ENDPOINT_SETTING)
        except Exception as exc:
            if previous.get("manifest_snapshot_captured") and not previous.get(
                "manifest_present"
            ):
                await _abort_enrollment_rollback(exc, "settings storage")
            logger.warning(
                "Could not restore the previous worker endpoint.", exc_info=True
            )
    else:
        _remember_endpoint(previous.get("endpoint", ""))
    agent.endpoint = previous.get("agent_endpoint") or previous.get("endpoint", "")
    if has_exact_settings_snapshot:
        try:
            from services import settings_store  # noqa: PLC0415

            if previous.get("worker_mode_setting_present"):
                settings_store.set_text(
                    _WORKER_MODE_SETTING, previous.get("worker_mode") or ""
                )
            else:
                settings_store.clear_text(_WORKER_MODE_SETTING)
        except Exception as exc:
            if previous.get("manifest_snapshot_captured") and not previous.get(
                "manifest_present"
            ):
                await _abort_enrollment_rollback(exc, "settings storage")
            logger.warning(
                "Could not restore the previous worker-mode setting.", exc_info=True
            )
    stored_mode = (previous.get("worker_mode") or "").strip().lower()
    if stored_mode not in ("1", "true", "yes", "on"):
        if stored_mode and not has_exact_settings_snapshot:
            set_worker_mode_enabled(False)
        return
    if not has_exact_settings_snapshot:
        set_worker_mode_enabled(True)
    try:
        await agent.start()
    except Exception:
        # It was working a moment ago; if it will not come back the panel
        # already shows the join error and the user can retry.
        logger.warning("Could not resume the previous control plane.", exc_info=True)


def enrolled() -> bool:
    """Has this machine ever completed a join?

    The atomic manifest is committed only after acceptance. Its presence still
    counts when corrupt so the UI reports broken saved state instead of
    pretending this is a never-enrolled machine; the certificate is the legacy
    pre-manifest signal.
    """
    try:
        locations = _paths()
        try:
            manifest = _load_enrollment_manifest(
                _enrollment_manifest_path(locations)
            )
        except EnrollmentStateError:
            # Do not present corrupt committed state as a never-enrolled
            # machine and then fall through to stale compatibility mirrors.
            return True
        return bool(manifest or os.path.exists(locations["pinned_cert"]))
    except Exception:
        return False


def _paths() -> dict[str, str]:
    from worker.service import paths  # noqa: PLC0415

    locations = paths()
    locations["pinned_cert"] = os.path.join(locations["root"], "control-plane.pinned.crt")
    # The server-assigned id, remembered so a restarted worker can prove who it
    # is. The challenge signature binds to this id, so a worker that forgets it
    # cannot authenticate with the key it already enrolled.
    locations["worker_id"] = os.path.join(locations["root"], "worker-id")
    locations["enrollment_token_hash"] = os.path.join(
        locations["root"], "enrollment-token.sha256"
    )
    locations["enrollment_manifest"] = os.path.join(
        locations["root"], "enrollment.json"
    )
    return locations


def _enrollment_manifest_path(locations: dict[str, str]) -> str:
    configured = locations.get("enrollment_manifest")
    if configured:
        return configured
    root = locations.get("root")
    if not root:
        root = os.path.dirname(os.path.abspath(locations["pinned_cert"]))
    return os.path.join(root, "enrollment.json")


def _load_enrollment_manifest(path: str) -> Optional[dict]:
    """Read the atomic enrollment generation, or fall back to legacy files."""
    try:
        with open(path, encoding="utf-8") as fh:
            record = json.load(fh)
        if (
            not isinstance(record, dict)
            or record.get("schema") != _ENROLLMENT_MANIFEST_SCHEMA
        ):
            raise ValueError("unsupported enrollment schema")
        endpoint = record.get("endpoint")
        worker_id = record.get("worker_id")
        token_hash = record.get("token_hash")
        encoded_certificate = record.get("certificate_pem")
        if not isinstance(endpoint, str) or not endpoint.strip():
            raise ValueError("missing enrollment endpoint")
        if not isinstance(worker_id, str) or not worker_id.strip():
            raise ValueError("missing worker id")
        if not isinstance(token_hash, str) or (
            token_hash
            and (
                len(token_hash) != 64
                or any(character not in _LOWER_HEX for character in token_hash)
            )
        ):
            raise ValueError("invalid enrollment-token hash")
        if not isinstance(encoded_certificate, str):
            raise ValueError("missing pinned certificate")
        certificate = base64.b64decode(encoded_certificate, validate=True)
        if not certificate:
            raise ValueError("empty pinned certificate")
        return {
            "endpoint": endpoint.strip(),
            "worker_id": worker_id.strip(),
            "token_hash": token_hash,
            "certificate": certificate,
        }
    except FileNotFoundError:
        return None
    except (
        OSError,
        UnicodeError,
        json.JSONDecodeError,
        binascii.Error,
        ValueError,
    ) as exc:
        raise EnrollmentStateError(
            "Saved worker enrollment is unreadable. Restore or remove its enrollment "
            "file, then join the control plane again."
        ) from exc


def _save_enrollment_manifest(
    path: str,
    *,
    endpoint: str,
    certificate: bytes,
    worker_id: str,
    token_hash: str,
) -> None:
    """Atomically commit all state needed to reconnect after acceptance."""
    endpoint = endpoint.strip()
    worker_id = worker_id.strip()
    if not endpoint or not worker_id or not certificate:
        raise ValueError("Accepted enrollment state is incomplete")
    if token_hash and (
        len(token_hash) != 64 or any(character not in _LOWER_HEX for character in token_hash)
    ):
        raise ValueError("Accepted enrollment token hash is invalid")
    encoded = json.dumps(
        {
            "schema": _ENROLLMENT_MANIFEST_SCHEMA,
            "endpoint": endpoint,
            "worker_id": worker_id,
            "token_hash": token_hash,
            "certificate_pem": base64.b64encode(certificate).decode("ascii"),
        },
        separators=(",", ":"),
        sort_keys=True,
    ).encode("utf-8")
    _replace_enrollment_manifest(path, encoded)


def _replace_enrollment_manifest(path: str, encoded: bytes) -> None:
    """Durably replace one complete enrollment generation."""
    directory = os.path.dirname(os.path.abspath(path))
    _durable_makedirs(directory)
    file_descriptor, temporary = tempfile.mkstemp(
        prefix=".enrollment.", suffix=".tmp", dir=directory
    )
    try:
        with os.fdopen(file_descriptor, "wb") as fh:
            fh.write(encoded)
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(temporary, path)
        _fsync_parent_directory(directory)
    except Exception:
        try:
            os.unlink(temporary)
        except OSError:  # best effort: preserve the original write/replace error
            pass
        raise


def _fsync_parent_directory(directory: str) -> None:
    """Make a preceding directory-entry replacement durable where supported."""
    directory_flag = getattr(os, "O_DIRECTORY", None)
    if directory_flag is None:
        return
    unsupported = {
        errno.EINVAL,
        getattr(errno, "ENOTSUP", errno.EINVAL),
        getattr(errno, "EOPNOTSUPP", errno.EINVAL),
    }
    try:
        descriptor = os.open(directory, os.O_RDONLY | directory_flag)
    except OSError as exc:
        if exc.errno in unsupported:
            return
        raise
    try:
        os.fsync(descriptor)
    except OSError as exc:
        if exc.errno not in unsupported:
            raise
    finally:
        os.close(descriptor)


def _durable_makedirs(path: str) -> None:
    """Create a directory chain and persist each new parent entry."""
    target = os.path.abspath(path)
    if os.path.isdir(target):
        return
    if os.path.exists(target):
        raise FileExistsError(f"Not a directory: {target}")

    missing: list[str] = []
    cursor = target
    while not os.path.exists(cursor):
        missing.append(cursor)
        parent = os.path.dirname(cursor)
        if parent == cursor:
            break
        cursor = parent
    if not os.path.isdir(cursor):
        raise NotADirectoryError(cursor)

    for directory in reversed(missing):
        parent = os.path.dirname(directory)
        try:
            os.mkdir(directory)
        except FileExistsError:
            if not os.path.isdir(directory):
                raise
        # A file fsync cannot make the directory entry that names its parent
        # durable. Persist each level so a fresh workers/ tree survives power
        # loss before the enrollment manifest is acknowledged.
        _fsync_parent_directory(parent)


def _persisted_endpoint() -> str:
    try:
        locations = _paths()
        manifest = _load_enrollment_manifest(_enrollment_manifest_path(locations))
        if manifest:
            return manifest["endpoint"]
    except EnrollmentStateError:
        return ""
    except Exception:
        pass
    return _stored_endpoint()


def load_worker_id(path: str) -> str:
    try:
        with open(path, encoding="utf-8") as fh:
            return fh.read().strip()
    except (FileNotFoundError, PermissionError):
        return ""


def save_worker_id(path: str, worker_id: str) -> None:
    if not worker_id:
        return
    _durable_makedirs(os.path.dirname(os.path.abspath(path)))
    with open(path, "w", encoding="utf-8") as fh:
        fh.write(worker_id)


def _token_hash(token_text: str) -> str:
    return hashlib.sha256(token_text.encode("utf-8")).hexdigest() if token_text else ""


def _load_consumed_token_hash(path: str) -> str:
    try:
        with open(path, encoding="ascii") as fh:
            digest = fh.read()
    except (OSError, UnicodeError):
        return ""
    if len(digest) != 64 or any(character not in _LOWER_HEX for character in digest):
        return ""
    return digest


def _save_consumed_token_hash(path: str, token_text: str) -> None:
    digest = _token_hash(token_text)
    if not digest:
        return
    directory = os.path.dirname(os.path.abspath(path))
    _durable_makedirs(directory)
    file_descriptor, temporary = tempfile.mkstemp(
        prefix=".enrollment-token.", suffix=".tmp", dir=directory
    )
    try:
        with os.fdopen(file_descriptor, "w", encoding="ascii") as fh:
            fh.write(digest)
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(temporary, path)
    except Exception:
        try:
            os.unlink(temporary)
        except OSError:  # best effort: preserve the original write/replace error
            pass
        raise


def _save_pinned_certificate(path: str, certificate: bytes) -> None:
    """Atomically replace the certificate only after enrollment is accepted."""
    directory = os.path.dirname(os.path.abspath(path))
    _durable_makedirs(directory)
    file_descriptor, temporary = tempfile.mkstemp(
        prefix=".control-plane.", suffix=".tmp", dir=directory
    )
    try:
        with os.fdopen(file_descriptor, "wb") as fh:
            fh.write(certificate)
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(temporary, path)
    except Exception:
        try:
            os.unlink(temporary)
        except OSError:  # best effort: preserve the original write/replace error
            pass
        raise


def fetch_server_certificate(endpoint: str, *, timeout: float = 10.0) -> bytes:
    """Retrieve the certificate the control plane presents, unvalidated.

    Unvalidated on purpose and safe only because the caller immediately checks
    it against the token's fingerprint — this is the fetch half of pin-on-first-
    use, not a trust decision.
    """
    try:
        parsed = urlsplit(f"//{endpoint}")
        host = parsed.hostname
        port = parsed.port
    except ValueError as exc:
        raise ValueError(f"Endpoint must be host:port — got {endpoint!r}") from exc
    if (
        not host
        or port is None
        or parsed.username is not None
        or parsed.password is not None
        or parsed.path
        or parsed.query
        or parsed.fragment
    ):
        raise ValueError(f"Endpoint must be host:port — got {endpoint!r}")
    from worker import tls  # noqa: PLC0415

    context = tls.unverified_client_context()
    with ssl.create_connection((host, port), timeout=timeout) as raw:
        with context.wrap_socket(raw, server_hostname=host) as tls:
            der = tls.getpeercert(binary_form=True)
    if not der:
        raise ConnectionError("The control plane presented no certificate.")
    return ssl.DER_cert_to_PEM_cert(der).encode("ascii")


def _verify_enrollment_token(token_text: str) -> tuple[str, bytes]:
    """Resolve a token into a verified endpoint and certificate without persisting it.

    Raises ``ValueError`` when the presented certificate does not match the
    token. There is deliberately no override.
    """
    from worker.identity import EnrollmentToken  # noqa: PLC0415
    from worker.transport.client import verify_pin  # noqa: PLC0415

    token = EnrollmentToken.decode(token_text)
    if token.expired():
        raise ValueError("This enrollment token has expired. Generate a new one.")

    certificate = fetch_server_certificate(token.endpoint)
    if not verify_pin(certificate, token.cert_fingerprint):
        raise ValueError(
            "The control plane's certificate does not match this enrollment token. "
            "Stop — this is what the token's fingerprint exists to catch. Generate a "
            "fresh token on the control plane and try again."
        )

    return token.endpoint, certificate


def pin_certificate(token_text: str, *, cert_path: Optional[str] = None) -> tuple[str, bytes]:
    """Verify and atomically pin the certificate from an enrollment token."""
    endpoint, certificate = _verify_enrollment_token(token_text)
    _save_pinned_certificate(cert_path or _paths()["pinned_cert"], certificate)
    return endpoint, certificate


def _classify_legacy_environment_token(
    token_text: str,
    *,
    endpoint: str,
    cert_path: str,
    certificate: Optional[bytes] = None,
) -> tuple[bool, str]:
    """Return ``(is_replacement, decoded_endpoint)`` for a legacy token.

    Releases before the token-hash marker persisted the original Compose token
    beside a worker id. Re-spending that token breaks every restart, while
    ignoring every unmarked token prevents a non-revoked worker from moving to
    another control plane. The token's endpoint and certificate fingerprint
    identify which enrollment it belongs to without contacting either server.

    A malformed token is treated as the already-spent legacy value. That keeps
    a working key enrollment usable. A valid same-origin token remains
    ambiguous and available only if a later key reconnect is rejected.
    """
    from worker.identity import EnrollmentToken  # noqa: PLC0415
    from worker.transport.client import verify_pin  # noqa: PLC0415

    try:
        token = EnrollmentToken.decode(token_text)
    except (TypeError, ValueError):
        return False, ""

    token_endpoint = token.endpoint.strip()
    if endpoint.strip() and token_endpoint != endpoint.strip():
        return True, token_endpoint
    try:
        if certificate is None:
            with open(cert_path, "rb") as fh:
                certificate = fh.read()
        return not verify_pin(certificate, token.cert_fingerprint), token_endpoint
    except (OSError, TypeError, ValueError):
        # A valid token can recover an enrollment whose local trust material is
        # absent or corrupt; it cannot be the token for that unusable state.
        return True, token_endpoint


class WorkerAgent:
    """Keeps this machine connected to a control plane and running its work."""

    def __init__(self) -> None:
        self._task: Optional[asyncio.Task] = None
        self._idle_sweep: Optional[asyncio.Task] = None
        self._client = None
        # Why the last start failed, kept so the panel can say "that token has
        # expired" instead of leaving a toggle that silently springs back.
        self.last_error: str = ""
        self.endpoint: str = ""
        # Set only after the control plane activates the Control stream. A
        # Register response saves identity but is still provisional server-side.
        self._registered = asyncio.Event()
        # One lock for every lifecycle change. Two concurrent requests — a join
        # and a toggle, or two joins — used to interleave their stop/start
        # pairs, and `start()` returns early when something is already running,
        # so a join could report success for a control plane it never dialled.
        self.lifecycle = asyncio.Lock()

    @property
    def running(self) -> bool:
        return self._task is not None and not self._task.done()

    def status(self) -> dict:
        """What the "lend this machine" panel renders, in one call."""
        return {
            "worker_mode": worker_mode_enabled(),
            "running": self.running,
            "enrolled": enrolled(),
            "endpoint": self.endpoint
            or (os.environ.get("OMNIVOICE_WORKER_ENDPOINT") or "").strip()
            or _persisted_endpoint(),
            "last_error": self.last_error,
            "env_pinned": bool((os.environ.get("OMNIVOICE_WORKER_MODE") or "").strip()),
        }

    def readiness(self) -> dict[str, object]:
        """Whether a worker-only deployment activated its first Control stream."""
        if not worker_mode_enabled():
            return {"ready": False, "status": "disabled"}
        if self.running and self._registered.is_set():
            return {"ready": True, "status": "ready"}
        if self.last_error:
            return {"ready": False, "status": "failed"}
        if self.running:
            return {"ready": False, "status": "registering"}
        return {"ready": False, "status": "stopped"}

    async def start(self, *, token_text: str = "", endpoint: str = "") -> None:
        from worker import capabilities  # noqa: PLC0415
        from worker.executor import TaskExecutor  # noqa: PLC0415
        from worker.identity import load_or_create_worker_key  # noqa: PLC0415
        from worker.transport.client import (  # noqa: PLC0415
            TerminalRegistrationError,
            WorkerClient,
            WorkerConfig,
            describe_host,
        )

        if self.running:
            return

        # A terminal connection failure ends `_task` but not the independent
        # idle-unload loop. Retire that whole generation before replacing any
        # references, or each retry leaks one infinite sweep that stop() can no
        # longer find (and leaves the failed client's work undrained).
        await self._retire_generation()

        self._registered.clear()
        self.last_error = ""
        locations = _paths()
        _durable_makedirs(locations["root"])
        # Generated once and never transmitted; this is the worker's identity
        # for the life of the machine.
        keypair = load_or_create_worker_key(locations["worker_key"])

        # A container's environment persists across process restarts, while an
        # enrollment token is single-use. Once registration has persisted a
        # worker id, reconnect with the identity key instead of trying to
        # spend that environment token again. An explicit token still means
        # "re-enrol", which keeps the in-app join flow able to replace an
        # existing control plane.
        explicit_token = token_text.strip()
        environment_token = (os.environ.get("OMNIVOICE_WORKER_TOKEN") or "").strip()
        manifest_path = _enrollment_manifest_path(locations)
        replace_invalid_manifest = False
        try:
            manifest = _load_enrollment_manifest(manifest_path)
        except EnrollmentStateError:
            if not explicit_token:
                raise
            # An explicit join is user-authorised repair. Treat none of the
            # separate compatibility files as authoritative; only the freshly
            # verified token may replace the corrupt generation on acceptance.
            manifest = None
            replace_invalid_manifest = True
        worker_id = manifest["worker_id"] if manifest else load_worker_id(
            locations["worker_id"]
        )
        token_hash_path = locations.get("enrollment_token_hash") or os.path.join(
            locations["root"], "enrollment-token.sha256"
        )
        consumed_token_hash = (
            manifest["token_hash"]
            if manifest
            else _load_consumed_token_hash(token_hash_path)
        )
        stored_endpoint = manifest["endpoint"] if manifest else _stored_endpoint()
        current_endpoint = (
            stored_endpoint
            or endpoint
            or (os.environ.get("OMNIVOICE_WORKER_ENDPOINT") or "").strip()
        )
        legacy_token_is_new = False
        legacy_token_endpoint = ""
        if environment_token and worker_id and not consumed_token_hash:
            legacy_token_is_new, legacy_token_endpoint = (
                _classify_legacy_environment_token(
                    environment_token,
                    endpoint=current_endpoint,
                    cert_path=locations["pinned_cert"],
                    certificate=(manifest["certificate"] if manifest else None),
                )
            )
        environment_token_is_new = bool(
            environment_token
            and (
                (
                    consumed_token_hash
                    and _token_hash(environment_token) != consumed_token_hash
                )
                or legacy_token_is_new
            )
        )
        token_text = explicit_token or environment_token
        should_enroll = bool(
            token_text and (explicit_token or not worker_id or environment_token_is_new)
        )
        legacy_token_fallback = bool(
            not should_enroll
            and environment_token
            and worker_id
            and not consumed_token_hash
            and not legacy_token_is_new
            and legacy_token_endpoint
        )
        if should_enroll:
            # Verify in memory, but do not replace a working enrollment until
            # the new control plane accepts this worker. A rejected token must
            # leave the certificate, endpoint, id, and token marker untouched.
            endpoint, certificate = await asyncio.to_thread(
                _verify_enrollment_token, token_text
            )
        else:
            token_text = ""
            # Already enrolled: reuse the certificate pinned at join time.
            if manifest:
                certificate = manifest["certificate"]
            else:
                try:
                    with open(locations["pinned_cert"], "rb") as fh:
                        certificate = fh.read()
                except (FileNotFoundError, PermissionError) as exc:
                    raise RuntimeError(
                        "This machine has not been enrolled yet. Ask for a join code on the "
                        "control plane (Settings → System → Remote workers) and paste it into "
                        "Lend this machine's GPU."
                    ) from exc
            endpoint = (
                endpoint
                or (os.environ.get("OMNIVOICE_WORKER_ENDPOINT") or "").strip()
                or stored_endpoint
                or legacy_token_endpoint
            )
            if not endpoint:
                raise RuntimeError(
                    "This machine does not know which control plane to dial. Paste a fresh "
                    "join code, or set OMNIVOICE_WORKER_ENDPOINT to its host:port."
                )
        if not should_enroll:
            self.endpoint = endpoint

        # Unavailable engines are reported too, so the control plane can tell
        # "this worker has no such engine" from "it has it but the weights
        # aren't downloaded" — a row that never arrives can only look like the
        # former, and the download-first flow has nothing to offer.
        discovered = capabilities.discover(include_unavailable=True)
        host = describe_host()
        host["gpus"] = capabilities.describe_gpus()

        config = WorkerConfig(
            endpoint=endpoint,
            cert_fingerprint="",
            certificate_pem=certificate,
            keypair=keypair,
            # Worker ids are panel-local.  A fresh enrollment may be moving
            # this key to a different control plane, whose accepted id is not
            # known until Register succeeds.  Sending the previous panel's id
            # also makes a lost first Register response unrecoverable: the
            # spent-token retry is signed over that stale id instead of the
            # newly enrolled identity.  Keep the durable old generation for
            # rollback, but enroll on the wire without it.
            worker_id="" if should_enroll else worker_id,
            enrollment_token=token_text,
            max_concurrent_tasks=capabilities.max_concurrent_tasks(discovered),
            capabilities=discovered,
            host=host,
        )
        # No reporters here on purpose: the client injects a pair bound to each
        # assignment's ref (``execute(assignment, on_progress=…,
        # on_model_loading=…)``), which is the only way a multi-slot worker can
        # say which task a progress fraction belongs to. Those reports are what
        # renews the server's progress lease.
        executor = TaskExecutor()
        accepted_enrollment = {
            # An unmarked same-origin legacy token is ambiguous: it may be the
            # old spent Compose token or a fresh token minted after a control-
            # plane reset. A successful key reconnect proves neither case, so
            # only record the hash after the token itself is accepted.
            "token": token_text if should_enroll else "",
            "certificate": certificate,
            "endpoint": endpoint,
        }

        def handle_registered(registered_worker_id: str) -> None:
            self._on_registered(
                locations["worker_id"],
                manifest_path=manifest_path,
                token_hash_path=token_hash_path,
                enrollment_token=accepted_enrollment["token"],
                consumed_token_hash=consumed_token_hash,
                cert_path=locations["pinned_cert"],
                certificate=accepted_enrollment["certificate"],
                endpoint=accepted_enrollment["endpoint"],
                replace_invalid_manifest=replace_invalid_manifest,
            )(registered_worker_id)

        def handle_activated(worker_id: str) -> None:
            # The durability callback runs on a worker thread. Publish only
            # this loop-owned status field once Config proves activation.
            self.endpoint = accepted_enrollment["endpoint"]
            self._on_activated(worker_id)

        client = WorkerClient(
            config,
            execute=executor.execute,
            # Re-probed on every reconnect so a model loaded (or evicted) since
            # the last connection is reported honestly rather than from a
            # snapshot taken at startup.
            capability_probe=lambda: capabilities.discover(include_unavailable=True),
            on_registered=handle_registered,
            on_activated=handle_activated,
            drain_active_work=executor.drain_active_work,
        )
        self._client = client

        async def run_with_legacy_token_recovery() -> None:
            try:
                await client.run_forever()
                return
            except TerminalRegistrationError as exc:
                error_code = str(exc).partition(":")[0].strip()
                if (
                    not legacy_token_fallback
                    or error_code != "AUTH_FAILED"
                ):
                    raise

            # The token looked identical to a legacy spent value, so only a
            # rejected key proves it is worth trying. Verify trust and expiry
            # now, then redeem it once; a second terminal response propagates.
            retry_endpoint, retry_certificate = await asyncio.to_thread(
                _verify_enrollment_token, environment_token
            )
            client.config.endpoint = retry_endpoint
            client.config.certificate_pem = retry_certificate
            # This is a fresh enrollment after the original panel identity was
            # rejected.  Its old id is panel-local and would make recovery of
            # a lost first token response sign the spent-token retry over the
            # wrong identity.
            client.config.worker_id = ""
            client.config.enrollment_token = environment_token
            accepted_enrollment.update(
                token=environment_token,
                certificate=retry_certificate,
                endpoint=retry_endpoint,
            )
            await client.run_forever()

        self._task = asyncio.create_task(
            run_with_legacy_token_recovery(), name="worker-agent"
        )
        self._task.add_done_callback(self._record_task_failure)
        self._idle_sweep = asyncio.create_task(
            self._unload_idle_engines(), name="worker-idle-unload"
        )
        logger.info(
            "Worker agent connecting to %s with %d engine(s)", endpoint, len(discovered)
        )

    def _on_registered(
        self,
        worker_id_path: str,
        *,
        manifest_path: str,
        token_hash_path: str = "",
        enrollment_token: str = "",
        consumed_token_hash: str = "",
        cert_path: str = "",
        certificate: bytes = b"",
        endpoint: str = "",
        replace_invalid_manifest: bool = False,
    ):
        def handle(worker_id) -> None:
            token_hash = _token_hash(enrollment_token) or consumed_token_hash
            enrollment = {
                "endpoint": endpoint.strip(),
                "worker_id": worker_id.strip(),
                "token_hash": token_hash,
                "certificate": certificate,
            }
            # One atomic replace is the source of truth. The compatibility
            # mirrors below may fail independently without creating mixed
            # enrollment state on the next restart.
            try:
                current_enrollment = _load_enrollment_manifest(manifest_path)
            except EnrollmentStateError:
                if not replace_invalid_manifest:
                    raise
                current_enrollment = None
            if current_enrollment != enrollment:
                _save_enrollment_manifest(
                    manifest_path,
                    endpoint=endpoint,
                    certificate=certificate,
                    worker_id=worker_id,
                    token_hash=token_hash,
                )
            if cert_path:
                try:
                    with open(cert_path, "rb") as fh:
                        saved_certificate = fh.read()
                except OSError:
                    saved_certificate = None
                if saved_certificate != certificate:
                    try:
                        _save_pinned_certificate(cert_path, certificate)
                    except OSError:
                        logger.warning(
                            "Could not update the legacy pinned-certificate mirror."
                        )
            if endpoint and _stored_endpoint() != endpoint.strip():
                try:
                    _remember_endpoint(endpoint)
                except Exception:
                    logger.warning("Could not update the legacy endpoint mirror.")
            if load_worker_id(worker_id_path) != worker_id.strip():
                try:
                    save_worker_id(worker_id_path, worker_id)
                except OSError:
                    logger.warning("Could not update the legacy worker-id mirror.")
            if (
                token_hash_path
                and enrollment_token
                and _load_consumed_token_hash(token_hash_path) != token_hash
            ):
                try:
                    _save_consumed_token_hash(token_hash_path, enrollment_token)
                except OSError:
                    logger.warning(
                        "Could not update the legacy enrollment-token mirror."
                    )

        return handle

    def _on_activated(self, _worker_id: str) -> None:
        """The first ConfigUpdate proves Control published this generation."""
        self._registered.set()

    def _record_task_failure(self, task: asyncio.Task) -> None:
        if task.cancelled():
            return
        try:
            error = task.exception()
        except asyncio.CancelledError:
            return
        if error is not None:
            self.last_error = str(error)

    async def wait_until_registered(self, timeout: float = 20.0) -> None:
        """Block until the control plane has activated this worker.

        `start()` returns as soon as the connection is scheduled, so a join
        that the control plane rejects — expired token, wrong address, a
        server that never answers — looked exactly like a successful one:
        the setting was persisted, the panel said "connected", and the machine
        retried in the background forever. Raises with the connection's own
        reason so the caller can undo the join and show it.
        """
        if self._task is None:
            raise RuntimeError("The worker agent is not running.")
        registered = asyncio.ensure_future(self._registered.wait())
        try:
            done, _ = await asyncio.wait(
                {registered, self._task},
                timeout=timeout,
                return_when=asyncio.FIRST_COMPLETED,
            )
        finally:
            if not registered.done():
                registered.cancel()
        if registered in done:
            return
        if self._task in done:
            # The connection loop gave up. Surface ITS error, not a timeout.
            error = self._task.exception() if not self._task.cancelled() else None
            raise RuntimeError(
                str(error) if error else "The connection to the control plane stopped."
            )
        raise RuntimeError(
            "The control plane did not answer in time. Check that it is running and "
            "reachable at this address, then try the code again."
        )

    # ── Idle unloading ────────────────────────────────────────────────────

    async def _unload_idle_engines(self) -> None:
        await idle_unload_loop(self._refresh_capabilities)

    async def _refresh_capabilities(self) -> None:
        if self._client is not None:
            await self._client.refresh_capabilities()

    async def _retire_generation(self) -> None:
        client = self._client
        tasks = [task for task in (self._task, self._idle_sweep) if task is not None]
        stop_error: Optional[BaseException] = None
        if client is not None:
            try:
                await client.stop()
            except BaseException as exc:
                stop_error = exc
        for task in tasks:
            if not task.done():
                task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        self._task = None
        self._idle_sweep = None
        self._client = None
        if stop_error is not None:
            raise stop_error

    async def stop(self) -> None:
        await self._retire_generation()


async def idle_unload_loop(on_released=None) -> None:
    """Hand back engines this machine has not used for ten minutes.

    Module-level because BOTH transports need it and only one of them used to
    have it: the sweep lived inside the dial-out agent, which an inbound-only
    node never starts, so a machine lending its GPU to panels that dial IN held
    several GB of weights forever against a task that might never come. That is
    the exact cost this exists to avoid, and it was silently absent in the mode
    most likely to be a shared box.

    Local behaviour is unchanged: nothing sweeps unless a worker role is
    running. `on_released` re-advertises capabilities when something was
    actually freed, so a control plane's view of what is resident does not go
    stale the moment it becomes useful.
    """
    from services import tts_backend  # noqa: PLC0415

    while True:
        await asyncio.sleep(IDLE_SWEEP_INTERVAL_SECONDS)
        try:
            # This process serves the desktop user as well as remote
            # assignments. Local work runs through the shared GPU pool but
            # does not enter the worker executor's per-engine guard.
            from services import model_manager  # noqa: PLC0415

            local = model_manager.gpu_pool_stats()
            if local.get("running", 0) or local.get("queued", 0):
                continue
            # unload() frees device caches and reaps sidecars — blocking, so it
            # must not run on the loop that answers heartbeats.
            released = await to_thread_and_drain_on_cancel(
                tts_backend.release_idle_engines
            )
            if released and on_released is not None:
                await on_released()
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception("Idle engine sweep failed (the worker continues)")


agent = WorkerAgent()


async def start_if_worker_mode() -> None:
    """Called from the app lifespan on the worker machine.

    Never fatal: a machine that cannot reach its control plane is still a
    perfectly good OmniVoice install for the person sitting at it.
    """
    if not worker_mode_enabled():
        return
    try:
        await agent.start()
    except Exception as exc:
        agent.last_error = str(exc)
        logger.exception("Worker agent failed to start (the app continues normally)")


async def stop() -> None:
    try:
        await agent.stop()
    except Exception:
        logger.exception("Worker agent failed to stop cleanly")


__all__ = [
    "WorkerAgent",
    "agent",
    "fetch_server_certificate",
    "pin_certificate",
    "start_if_worker_mode",
    "stop",
    "worker_mode_enabled",
]