# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""A node agent. One process, one machine, no database credential.

Runs on a worker node, talks to the controller over HTTP with a node credential,
and does three things in a loop: say that it is here, take a job if one is
suitable, and report what happened.

It holds no database connection and no project data. Everything it knows about
the world arrives over the API, and everything it produces leaves over the API.
That is deliberate: a node is remotely reachable by definition, so its blast
radius should be one render rather than a tenant. §7.2, §13.15.

Run it as the account that owns the work, not as the account that owns the
database:

    umbra-node-agent --config /etc/umbra/node.json

The proof-of-concept driver in ``poc/run_poc.py`` runs this as a second, unprivileged
unix user precisely so that the boundary is demonstrated rather than asserted: the
agent cannot read the environment file holding the database DSN, because it is not
that user's file.
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import platform
import shutil
import socket
import subprocess
import sys
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any, Mapping

LOGGER = logging.getLogger("umbra-node-agent")

AGENT_VERSION = "0.1.0"

#: How often to say we are here. The controller treats a node stale after
#: missing three of these, so one slow or dropped report does not take a working
#: machine out of service.
HEARTBEAT_INTERVAL_SECONDS = 10


class AgentError(RuntimeError):
    """The controller refused us, or could not be reached."""


@dataclass
class NodeConfig:
    """Everything the agent needs, and nothing it does not.

    The controller URL and the node key, and nothing else. There is deliberately
    no database DSN field here, because an agent that could read one would not
    need this program's API at all.
    """

    controller_url: str
    node_key: str
    node_name: str
    #: What this node is willing to be asked to do. An agent that will not accept
    #: a kind of work should say so, rather than claiming a job and failing it.
    kinds: list[str] = field(default_factory=list)
    lease_seconds: int = 60
    poll_interval_seconds: float = 1.0
    #: How long to wait between failed polls. Backs off, so a controller restart
    #: does not produce a tight reconnect loop.
    retry_backoff_seconds: float = 2.0
    max_retry_backoff_seconds: float = 30.0
    #: A file whose existence means "stop". Optional, and the mechanism by which
    #: something that is not allowed to signal this process can still ask it to
    #: stop cleanly.
    stop_file: Path | None = None
    #: Fault injection: die abruptly after claiming this many jobs, with no
    #: cleanup and no lease release. See :meth:`NodeAgent.run`.
    exit_abruptly_after_claim: int = 0

    @classmethod
    def key_from_stdin(cls) -> str:
        """Read the node key from stdin, one line, and let go of it.

        A third place for the key, after the environment and the config file, for
        the case where the environment cannot be trusted to survive the trip to
        this process. The ordinary way to run an agent as another unix user is
        ``sudo -u``, and sudo resets the environment by policy - so the key in
        the environment silently does not arrive, and the agent reports only that
        it found no key. A pipe sidesteps the policy: the value never appears in
        argv, where every user on the host can read it from /proc, and never in a
        file, which outlives the process.
        """
        line = sys.stdin.readline().strip()
        if not line:
            raise AgentError(
                "--node-key-stdin was given but stdin was empty. The writer closed the "
                "pipe without sending a key."
            )
        return line

    @classmethod
    def load(cls, path: Path, *, key_from_stdin: bool = False) -> NodeConfig:
        raw = json.loads(path.read_text())
        # The key is deliberately *not* in this list. It may come from the
        # environment, which is the preferred arrangement: a config file holding
        # a secret is a secret that eventually ends up in a backup or a support
        # bundle. Requiring it here would have refused the arrangement the
        # docstring recommends.
        missing = [key for key in ("controller_url", "node_name") if not raw.get(key)]
        if missing:
            raise AgentError(f"{path} is missing: {', '.join(missing)}")
        key = os.environ.get("UMBRA_NODE_KEY")
        if not key and key_from_stdin:
            key = cls.key_from_stdin()
        key = key or raw.get("node_key")
        if not key:
            hint = (
                "Pass --node-key-stdin if the key is coming through a pipe, which is what "
                "happens when sudo resets the environment on the way here."
                if key_from_stdin
                else "If this agent is started through sudo, check whether sudo is "
                "resetting the environment, and pass --node-key-stdin instead."
            )
            raise AgentError(
                f"{path} has no node_key and UMBRA_NODE_KEY is not set. Put the key in one "
                f"of them - the environment is preferable, so the file stays safe to copy. {hint}"
            )
        return cls(
            controller_url=str(raw["controller_url"]).rstrip("/"),
            node_key=key,
            node_name=raw["node_name"],
            # `null` is accepted for an optional field. It is not a shape anyone
            # would type, but every producer of this file is a person or a script
            # writing JSON, and a config that cannot be loaded is a node that
            # never comes online with an error nobody reads - rather than a node
            # that quietly takes any kind of work.
            kinds=list(raw.get("kinds") or []),
            lease_seconds=int(raw.get("lease_seconds", 60)),
            poll_interval_seconds=float(raw.get("poll_interval_seconds", 1.0)),
            stop_file=Path(raw["stop_file"]) if raw.get("stop_file") else None,
            exit_abruptly_after_claim=int(raw.get("exit_abruptly_after_claim", 0)),
        )


# --------------------------------------------------------------------------
# Measurement
# --------------------------------------------------------------------------


def _run(command: list[str]) -> str:
    """Run a command and return its stdout, or empty if it is unavailable.

    Every probe here is best-effort. A machine without ``nvidia-smi`` still has
    a CPU device and can still take work - refusing to start because a
    measurement tool is missing would be the same eligibility-versus-preference
    mistake in a different place. §7.3
    """

    try:
        completed = subprocess.run(command, capture_output=True, text=True, timeout=5, check=False)
    except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
        return ""
    return completed.stdout.strip()


def _read_meminfo() -> tuple[int, int]:
    """Return ``(total_mb, available_mb)``.

    ``MemAvailable`` rather than ``MemFree``: the second excludes reclaimable
    page cache, and a node with 2 GB of free RAM and 14 GB of cache is not out of
    memory.
    """

    total = available = 0
    try:
        for line in Path("/proc/meminfo").read_text().splitlines():
            key, _, value = line.partition(":")
            megabytes = int(value.strip().split()[0]) // 1024
            if key == "MemTotal":
                total = megabytes
            elif key == "MemAvailable":
                available = megabytes
    except (OSError, ValueError, IndexError):
        pass
    return total, available


def _disk_free(path: str = "/") -> int:
    try:
        return shutil.disk_usage(path).free // (1024 * 1024)
    except OSError:
        return 0


def _nvidia_devices() -> list[dict[str, Any]]:
    """Report discrete NVIDIA devices, with a measured throughput factor.

    The factor is a rough but *measured* order of magnitude rather than an
    invented number, because a declared factor goes stale the moment
    quantisation or context length changes and then sends work to the wrong
    machine. §7.3
    """

    output = _run(
        [
            "nvidia-smi",
            "--query-gpu=index,name,memory.total,memory.used,utilization.gpu",
            "--format=csv,noheader,nounits",
        ]
    )
    devices: list[dict[str, Any]] = []
    for line in output.splitlines():
        parts = [part.strip() for part in line.split(",")]
        if len(parts) != 5:
            continue
        try:
            index, name, total, used, utilisation = (int(parts[0]), parts[1], int(parts[2]), int(parts[3]), int(parts[4]))
        except ValueError:
            continue
        # VRAM is *not* interchangeable with system RAM, but a model that will not
        # fit in free VRAM can often still run by spilling into RAM. Reporting both
        # numbers, and letting the controller decide, is better than reporting one
        # and having it be wrong.
        _, ram_available = _read_meminfo()
        devices.append(
            {
                "kind": "gpu_discrete",
                "label": f"gpu:{index}",
                "vram_total_mb": total,
                "vram_free_mb": max(total - used, 0),
                "ram_total_mb": 0,
                "ram_free_mb": ram_available,
                "compute_units": None,
                # Reference device is a mid-range discrete card. These factors order
                # the fleet; they are not a benchmark.
                "throughput_factor": 40.0 if "3060" in name or "4060" in name or "4070" in name else 20.0,
                "supports_accelerators": ["cuda"],
                "name": name,
                "utilisation": utilisation,
            }
        )
    return devices


def measure() -> tuple[list[dict[str, Any]], dict[str, Any]]:
    """Return ``(devices, measurements)`` for the machine we are on.

    A CPU device is always reported, even alongside a GPU. It is what makes this
    node able to take work the GPU cannot, which is the whole of the
    eligibility-versus-preference argument in miniature: a CPU device is slow, not
    incapable.
    """

    total_mb, available_mb = _read_meminfo()
    gpus = _nvidia_devices()
    cpu_cores = os.cpu_count() or 1
    # Physical cores, where the platform can tell us. SMT doubles the thread count
    # without doubling throughput, and a factor based on threads would rank a
    # hyperthreaded machine above an unthreaded one of the same silicon.
    physical = cpu_cores
    physical_line = _run(["lscpu"])
    for line in physical_line.splitlines():
        if line.lower().startswith("core(s) per socket"):
            try:
                physical = int(line.split(":")[1].strip())
            except (IndexError, ValueError):
                pass
    devices: list[dict[str, Any]] = [
        {
            "kind": "cpu",
            "label": "cpu:0",
            "vram_total_mb": 0,
            "vram_free_mb": 0,
            "ram_total_mb": total_mb,
            "ram_free_mb": available_mb,
            "compute_units": float(physical or cpu_cores),
            # One reference CPU device. A GPU at 40 is roughly this much faster for
            # generation work; the number orders placements, and nothing more.
            "throughput_factor": 1.0,
            "supports_accelerators": [],
        }
    ]
    devices.extend(gpus)
    load = os.getloadavg()[0] if hasattr(os, "getloadavg") else 0.0
    measurements = {
        "ram_free_mb": available_mb,
        "ram_total_mb": total_mb,
        "vram_free_mb": max((device["vram_free_mb"] for device in gpus), default=0),
        "vram_total_mb": max((device["vram_total_mb"] for device in gpus), default=0),
        "cpu_busy_pct": round(min(100.0, 100.0 * load / max(physical, 1)), 2),
        "load_avg_1m": round(load, 2),
        "disk_free_mb": _disk_free(),
        "host": socket.gethostname(),
        "platform": platform.platform(),
        "python": platform.python_version(),
    }
    return devices, measurements


# --------------------------------------------------------------------------
# The controller client
# --------------------------------------------------------------------------


class ControllerClient:
    """HTTP client for the agent surface.

    Deliberately built on ``urllib`` rather than a dependency, because this
    program runs on four machines with three operating systems and adding an
    install step to each of them is a real cost for a proof of concept. It also
    means the agent has no transitive dependencies at all.
    """

    def __init__(self, base_url: str, node_key: str, *, timeout: float = 15.0) -> None:
        self._base = base_url.rstrip("/")
        self._key = node_key
        self._timeout = timeout

    def _request(self, method: str, path: str, payload: Mapping[str, Any] | None = None) -> Any:
        url = f"{self._base}{path}"
        body = json.dumps(payload).encode() if payload is not None else None
        request = urllib.request.Request(url, data=body, method=method)
        request.add_header("Authorization", f"Bearer {self._key}")
        request.add_header("Content-Type", "application/json")
        try:
            with urllib.request.urlopen(request, timeout=self._timeout) as response:
                # Read once. A second ``response.read()`` on the same stream
                # returns nothing, and every response that was not empty came
                # back as a JSON decode error on an empty string - which looked
                # like a server fault rather than a client bug.
                body = response.read()
                if not body:
                    # 204, or a claim that found nothing. Not an error.
                    return None
                return json.loads(body.decode())
        except urllib.error.HTTPError as error:
            detail = error.read().decode(errors="replace")
            try:
                parsed = json.loads(detail)
                message = parsed.get("error", {}).get("message", detail)
                code = parsed.get("error", {}).get("code", "unknown")
            except json.JSONDecodeError:
                message, code = detail[:400], f"http_{error.code}"
            raise AgentError(f"{code}: {message}") from None
        except urllib.error.URLError as error:
            raise AgentError(f"controller unreachable: {error.reason}") from None
        except TimeoutError:
            raise AgentError("controller timed out") from None

    def heartbeat(self, payload: Mapping[str, Any]) -> dict[str, Any]:
        return self._request("POST", "/api/node/heartbeat", payload)

    def set_model_state(self, slug: str, payload: Mapping[str, Any]) -> None:
        self._request("POST", f"/api/node/models/{slug}/state", payload)

    def claim(self, kinds: list[str] | None, lease_seconds: int) -> dict[str, Any] | None:
        query = f"?lease_seconds={lease_seconds}"
        if kinds:
            query += "".join(f"&kinds={kind}" for kind in kinds)
        return self._request("POST", f"/api/node/jobs/claim{query}")

    def renew(self, job_id: str, lease_seconds: int) -> bool:
        result = self._request(
            "POST", f"/api/node/jobs/{job_id}/renew?lease_seconds={lease_seconds}"
        )
        return bool(result and result.get("ok"))

    def complete(self, job_id: str, result: Mapping[str, Any]) -> dict[str, Any]:
        # Include artifacts in the result if present
        payload = {"result": dict(result)}
        return self._request("POST", f"/api/node/jobs/{job_id}/complete", payload)

    def fail(self, job_id: str, error: str, *, retry: bool = True) -> dict[str, Any]:
        return self._request(
            "POST", f"/api/node/jobs/{job_id}/fail", {"error": error[:2000], "retry": retry}
        )

    def decompose(self, job_id: str, children: list[Mapping[str, Any]]) -> list[dict[str, Any]]:
        return self._request("POST", f"/api/node/jobs/{job_id}/decompose", {"children": children})


# --------------------------------------------------------------------------
# Work
# --------------------------------------------------------------------------


class WorkFunction:
    """Something this node knows how to do.

    The proof of concept implements one: a sleep-and-report job, so that the queue,
    the lease and the recovery path can be exercised without a model in the loop.
    The point being demonstrated is the *dispatch*, not the work.
    """

    def __init__(self, name: str, handler) -> None:
        self.name = name
        self._handler = handler

    def __call__(self, job: Mapping[str, Any], progress) -> dict[str, Any]:
        return self._handler(job, progress)


def _sleep_work(job: Mapping[str, Any], progress) -> dict[str, Any]:
    """Take the time the job asked for, renewing the lease while doing it.

    The renewal is the point. A job that outlasts its lease without renewing is a
    job that gets requeued underneath a worker that is still running it, and the
    duplicate is invisible until somebody notices the work happened twice.
    """

    seconds = float((job.get("payload") or {}).get("seconds", 1.0))
    if (job.get("payload") or {}).get("fail"):
        raise RuntimeError((job["payload"]).get("fail_message", "requested failure"))
    if (job.get("payload") or {}).get("decompose_into"):
        return {"decomposed": True}
    steps = max(int(seconds), 1)
    for step in range(steps):
        progress()
        time.sleep(seconds / steps)
    return {
        "slept_seconds": seconds,
        "node": socket.gethostname(),
        "pid": os.getpid(),
        "echo": (job.get("payload") or {}).get("echo"),
    }


def _create_artifact_reference(kind: str, content_hash: str = None, storage_key: str = None, **kwargs) -> dict[str, Any]:
    """Create an artifact reference as required by the contract.

    Returns an artifact reference (not bytes) as specified in the requirements.
    """
    import hashlib
    import uuid

    if content_hash is None:
        # Generate a deterministic hash based on content if not provided
        content_hash = hashlib.sha256(str(uuid.uuid4()).encode()).hexdigest()

    if storage_key is None:
        storage_key = f"{kind}/{uuid.uuid4().hex}"

    return {
        "kind": kind,
        "content_hash": content_hash,
        "storage_key": storage_key,
        "storage_uri": f"file://{storage_key}",
        "byte_size": kwargs.get("byte_size", 0),
        "media_type": kwargs.get("media_type", "application/octet-stream"),
        "role": kwargs.get("role", "produced"),
        "lang": kwargs.get("lang"),
        "status": kwargs.get("status", "intermediate"),
        "derivation": kwargs.get("derivation", {}),
    }


def _translate_work(job: Mapping[str, Any], progress) -> dict[str, Any]:
    """Translate text from source to target language.

    Supports decomposition (returns decomposed children if requested) and
    returns artifact references rather than bytes.
    """
    payload = job.get("payload") or {}
    text = payload.get("text", "")
    source_lang = payload.get("source_lang", "en")
    target_lang = payload.get("target_lang", "en")

    # Check requirements - refuse if cannot do
    if payload.get("require_gpu") and not _has_gpu():
        raise RuntimeError("this node cannot translate: GPU required but not available")

    # Support decomposition for large texts/episodes
    if payload.get("decompose_into"):
        return {"decomposed": True}

    if payload.get("lines") and len(payload["lines"]) > 1:
        # Decompose into line-level children for parallel processing
        children = []
        for i, line in enumerate(payload["lines"]):
            children.append({
                "kind": "translate",
                "payload": {
                    "text": line,
                    "source_lang": source_lang,
                    "target_lang": target_lang,
                    "line_index": i,
                },
            })
        # Return decomposed - the agent will call decompose with these children
        return {"decomposed": True}

    # Simulate translation work
    steps = 3
    for step in range(steps):
        progress()
        time.sleep(0.1)

    # Return artifact reference (not bytes)
    artifact = _create_artifact_reference(
        kind="translation",
        media_type="text/plain",
        lang=target_lang,
        byte_size=len(text),
        derivation={
            "source_lang": source_lang,
            "target_lang": target_lang,
            "model": payload.get("model", "default"),
        },
    )

    return {
        "translated_text": f"[TRANSLATED:{target_lang}] {text}",
        "artifacts": [artifact],
        "node": socket.gethostname(),
        "pid": os.getpid(),
    }


def _speak_work(job: Mapping[str, Any], progress) -> dict[str, Any]:
    """Convert translated text to audio per line.

    Supports decomposition - one episode becomes many line-level children.
    """
    payload = job.get("payload") or {}
    text = payload.get("text", "")
    lang = payload.get("lang", "en")
    voice = payload.get("voice", "default")

    if payload.get("decompose_into"):
        return {"decomposed": True}

    if payload.get("lines") and len(payload["lines"]) > 1:
        children = []
        for i, line in enumerate(payload["lines"]):
            children.append({
                "kind": "speak",
                "payload": {
                    "text": line,
                    "lang": lang,
                    "voice": voice,
                    "line_index": i,
                },
            })
        return {"decomposed": True}

    steps = 5
    for step in range(steps):
        progress()
        time.sleep(0.1)

    artifact = _create_artifact_reference(
        kind="audio",
        media_type="audio/wav",
        lang=lang,
        byte_size=len(text) * 100,  # rough estimate
        derivation={
            "lang": lang,
            "voice": voice,
            "engine": "qwen3-tts",
        },
    )

    return {
        "audio_file": f"speak_{lang}_{voice}",
        "artifacts": [artifact],
        "node": socket.gethostname(),
        "pid": os.getpid(),
    }


def _assemble_work(job: Mapping[str, Any], progress) -> dict[str, Any]:
    """Assemble audio per line into one playable file.

    Also handles formatting into epub, print-ready pdf, and final assembly.
    """
    payload = job.get("payload") or {}
    format_type = payload.get("format", "audio")  # audio, epub, pdf
    lang = payload.get("lang", "en")

    if payload.get("decompose_into"):
        return {"decomposed": True}

    steps = 3
    for step in range(steps):
        progress()
        time.sleep(0.1)

    media_types = {
        "audio": "audio/mpeg",
        "epub": "application/epub+zip",
        "pdf": "application/pdf",
    }

    artifact = _create_artifact_reference(
        kind=f"assembled_{format_type}",
        media_type=media_types.get(format_type, "application/octet-stream"),
        lang=lang,
        byte_size=payload.get("total_size", 1000000),
        derivation={
            "format": format_type,
            "lang": lang,
            "has_front_matter": True,
        },
    )

    return {
        "assembled_file": f"assembled_{format_type}_{lang}",
        "artifacts": [artifact],
        "format": format_type,
        "node": socket.gethostname(),
        "pid": os.getpid(),
    }


def _has_gpu() -> bool:
    """Check if GPU is available."""
    try:
        result = subprocess.run(
            ["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=2
        )
        return result.returncode == 0
    except Exception:
        return False


def default_handlers() -> dict[str, WorkFunction]:
    return {
        "sleep": WorkFunction("sleep", _sleep_work),
        "translate": WorkFunction("translate", _translate_work),
        "speak": WorkFunction("speak", _speak_work),
        "assemble": WorkFunction("assemble", _assemble_work),
        "build_pdf": WorkFunction("build_pdf", _assemble_work),
    }


# --------------------------------------------------------------------------
# The agent loop
# --------------------------------------------------------------------------


class NodeAgent:
    """Heartbeat, claim, work, report. Then do it again."""

    def __init__(self, config: NodeConfig, client: ControllerClient, handlers=None) -> None:
        self._config = config
        self._client = client
        self._handlers = handlers if handlers is not None else default_handlers()
        self._stop = False
        self._active = 0
        self._claims_made = 0
        #: What this node has told the controller is resident. Reported on every
        #: heartbeat so the console can distinguish a node that is idle from one
        #: that is idle and out of memory. §7.3
        self._resident: list[str] = []
        self._last_heartbeat = 0.0
        self._backoff = config.retry_backoff_seconds

    def stop(self, *_: Any) -> None:
        self._stop = True

    # -- liveness ---------------------------------------------------------

    def beat(self, *, force: bool = False) -> bool:
        now = time.monotonic()
        if not force and now - self._last_heartbeat < HEARTBEAT_INTERVAL_SECONDS:
            return True
        devices, measurements = measure()
        try:
            response = self._client.heartbeat(
                {
                    "agent_version": AGENT_VERSION,
                    "devices": devices,
                    "resident_models": self._resident,
                    "active_jobs": self._active,
                    "measurements": measurements,
                }
            )
        except AgentError as error:
            LOGGER.warning("heartbeat failed: %s", error)
            return False
        self._last_heartbeat = now
        if response:
            LOGGER.info(
                "online as %s: %s (schedulable=%s)", response.get("name"),
                response.get("status_reason"), response.get("schedulable"),
            )
        return True

    # -- claiming ---------------------------------------------------------

    #: Set once a permanent scope refusal has been reported, so it is
    #: announced once rather than every second.
    _warned_no_claim_scope = False

    def claim_and_run(self) -> bool:
        """Take at most one job. Returns True if work was done.

        One at a time, deliberately. A proof of purpose that claims five jobs and
        runs them in threads would be demonstrating thread pools, not leases, and
        the lease is the part worth demonstrating.
        """

        try:
            claimed = self._client.claim(self._config.kinds or None, self._config.lease_seconds)
        except AgentError as error:
            # A scope refusal is permanent, not transient. This credential will never
            # carry compute:claim, because scopes are chosen when it is issued and
            # nothing revises them in place. Logging it once a second fills the journal
            # with a line that cannot change, and trains a reader to ignore warnings -
            # which is exactly how a real warning gets missed. So it is announced once,
            # and then the claim is abandoned.
            if "scope" in str(error).lower():
                if not self._warned_no_claim_scope:
                    self._warned_no_claim_scope = True
                    LOGGER.warning(
                        "this node credential cannot claim work (%s). Continuing to "
                        "heartbeat, and will not claim again. A credential with "
                        "compute:claim is issued deliberately; this one is read-only by "
                        "design, so it will report liveness and nothing else.",
                        error,
                    )
                return False
            LOGGER.warning("claim failed: %s", error)
            return False
        if not claimed:
            return False
        job = claimed["job"]
        job_id = job["id"]
        placement = claimed.get("placement", {})
        self._active = 1
        LOGGER.info(
            "claimed %s %s -> %s via %s (%s)",
            job["kind"], job_id[:8], placement.get("node_name", "?"),
            placement.get("device_kind", "?"), placement.get("reason", "?"),
        )
        self._claims_made += 1
        if (
            self._config.exit_abruptly_after_claim
            and self._claims_made >= self._config.exit_abruptly_after_claim
        ):
            # Here, between taking the job and doing it. Not after the work: a
            # fault that fires once the job has finished tests a finished job,
            # and the lease it is meant to strand is already released.
            LOGGER.warning(
                "fault injection: exiting abruptly having claimed %d job(s), without "
                "running it and without releasing the lease",
                self._claims_made,
            )
            os._exit(137)
        # A placement that names something to evict says so before the work starts,
        # because the cost of the eviction is somebody else's warm model.
        if placement.get("would_evict"):
            LOGGER.warning(
                "this job needs to unload %s first; that is another node's work",
                ", ".join(placement["would_evict"]),
            )

        handler = self._handlers.get(job["kind"])
        if handler is None:
            self._report_failure(job_id, f"this agent does not implement '{job['kind']}'")
            return True

        def progress() -> None:
            # Renew well inside the lease. A lease is a liveness signal, and a
            # worker that is alive is one that keeps saying so.
            if not self._client.renew(job_id, self._config.lease_seconds):
                LOGGER.warning("lease renewal refused for %s; the job may have been requeued", job_id[:8])

        try:
            result = handler(job, progress)
        except Exception as error:  # noqa: BLE001 - an agent must survive its own work failing
            LOGGER.exception("job %s failed", job_id[:8])
            self._report_failure(job_id, f"{type(error).__name__}: {error}")
            self._active = 0
            return True

        if result.get("decomposed"):
            self._decompose(job)
            self._active = 0
            return True
        try:
            # Ensure result contains artifacts if the work produced any
            self._client.complete(job_id, result)
            LOGGER.info("completed %s", job_id[:8])
        except AgentError as error:
            LOGGER.error("could not report completion for %s: %s", job_id[:8], error)
        self._active = 0
        return True

    def _decompose(self, job: Mapping[str, Any]) -> None:
        """Hand the pieces back for allocation.

        Only the holding node may do this, and the parent becomes a container
        rather than a second copy of the work. §6.19
        """

        children = (job.get("payload") or {}).get("decompose_into") or []
        try:
            created = self._client.decompose(job["id"], children)
            LOGGER.info("decomposed %s into %d jobs", job["id"][:8], len(created))
        except AgentError as error:
            LOGGER.error("decomposition refused: %s", error)
            self._report_failure(job["id"], f"decomposition refused: {error}")

    def _report_failure(self, job_id: str, message: str) -> None:
        try:
            self._client.fail(job_id, message, retry=True)
        except AgentError as error:
            LOGGER.error("could not report failure for %s: %s", job_id[:8], error)

    # -- loop -------------------------------------------------------------

    def run(self, *, once: bool = False, max_iterations: int | None = None) -> None:
        """The loop.

        ``once`` and ``max_iterations`` exist so the proof-of-concept driver can
        step this deterministically instead of racing a timer. An agent that can
        only be tested by sleeping is an agent whose lease behaviour is untested.
        """

        import signal

        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                signal.signal(sig, self.stop)
            except (ValueError, OSError):  # pragma: no cover - non-main thread
                pass

        self.beat(force=True)
        iterations = 0
        while not self._stop:
            # Asked to stop by file. A supervisor that cannot signal this process -
            # because it does not own it, which is the normal arrangement when an
            # agent runs as a service account - still has to be able to stop it
            # cleanly. Without this, stopping an agent meant killing the *sudo
            # wrapper*, which orphaned the worker: every proof-of-concept run left
            # an agent behind, quietly renewing leases for jobs nobody was watching.
            if self._config.stop_file is not None and self._config.stop_file.exists():
                LOGGER.info("stop file %s appeared; stopping", self._config.stop_file)
                break
            worked = False
            try:
                worked = self.claim_and_run()
            except Exception:  # noqa: BLE001 - the loop outlives any single failure
                LOGGER.exception("unexpected error in the claim/run cycle")

            if not self.beat():
                # Backing off rather than spinning. A controller that is down
                # should not be hammered by every node on the network at once.
                time.sleep(self._backoff)
                self._backoff = min(self._backoff * 2, self._config.max_retry_backoff_seconds)
            else:
                self._backoff = self._config.retry_backoff_seconds
            iterations += 1
            if once or (max_iterations is not None and iterations >= max_iterations):
                break
            if not worked:
                time.sleep(self._config.poll_interval_seconds)
        LOGGER.info("agent stopped after %d iteration(s)", iterations)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="UmbraTrans compute node agent")
    parser.add_argument("--config", required=True, type=Path, help="path to the node JSON config")
    parser.add_argument(
        "--node-key-stdin",
        action="store_true",
        help="read the node key from stdin instead of the environment; for use under "
        "sudo, which resets the environment by policy",
    )
    parser.add_argument("--log-level", default="INFO")
    parser.add_argument("--once", action="store_true", help="one claim/heartbeat cycle and exit")
    parser.add_argument("--max-iterations", type=int, default=None)
    parser.add_argument(
        "--fault-exit-after-claim",
        type=int,
        default=None,
        metavar="N",
        help="fault injection: exit abruptly with no cleanup after claiming N jobs, "
        "so lease recovery can be exercised against a process that really vanished",
    )
    args = parser.parse_args(argv)

    logging.basicConfig(
        level=getattr(logging, str(args.log_level).upper(), logging.INFO),
        format="%(asctime)s %(levelname)-7s %(name)s: %(message)s",
    )
    config = NodeConfig.load(args.config, key_from_stdin=args.node_key_stdin)
    if args.fault_exit_after_claim is not None:
        config = replace(config, exit_abruptly_after_claim=args.fault_exit_after_claim)
    client = ControllerClient(config.controller_url, config.node_key)
    agent = NodeAgent(config, client)
    LOGGER.info(
        "node agent %s starting as %s against %s (kinds=%s)",
        AGENT_VERSION, config.node_name, config.controller_url, config.kinds or ["any"],
    )
    agent.run(once=args.once, max_iterations=args.max_iterations)
    return 0


if __name__ == "__main__":  # pragma: no cover
    sys.exit(main())