horus-agent 0.1.2__py3-none-any.whl

horus_agent/__init__.py

@@ -0,0 +1,6 @@
+#
+# horus_agent
+# Copyright (c) 2026 Temple Compute
+#
+# MIT License
+#

horus_agent/__main__.py

@@ -0,0 +1,13 @@
+#
+# horus_agent
+# Copyright (c) 2026 Temple Compute
+#
+# MIT License
+#
+"""``python -m horus_agent``."""
+
+import sys
+
+from horus_agent.cli import main
+
+sys.exit(main())

horus_agent/blobs.py

@@ -0,0 +1,89 @@
+#
+# horus_agent
+# Copyright (c) 2026 Temple Compute
+#
+# MIT License
+#
+"""
+Streamed transfers to and from presigned staging URLs.
+
+Used on both sides: the orchestrator and the agent each move large files
+through object storage instead of through the broker.
+"""
+
+import time
+from collections.abc import AsyncIterator
+from pathlib import Path
+
+import httpx
+from horus_runtime.logging import horus_logger
+
+from horus_agent.logs import safe_url
+
+_CHUNK = 1024 * 1024
+MAX_SINGLE_PUT = 5 * 1024**3
+"""S3's limit for one PUT."""
+
+
+async def _read_chunks(path: Path) -> AsyncIterator[bytes]:
+    with path.open("rb") as fh:
+        while chunk := fh.read(_CHUNK):
+            yield chunk
+
+
+async def upload(url: str, source: Path | bytes) -> None:
+    """
+    PUT *source* to the presigned *url*, streaming a file from disk.
+
+    ``Content-Length`` is always sent: S3 refuses a chunked presigned PUT.
+    """
+    if isinstance(source, Path):
+        size = source.stat().st_size
+        if size > MAX_SINGLE_PUT:
+            # shortcut: one presigned PUT, S3's limit; add multipart if
+            # agents need to move larger single files.
+            raise ValueError(
+                f"{source} is {size} bytes; files over 5 GiB cannot be "
+                "moved to or from an agent yet"
+            )
+        content: bytes | AsyncIterator[bytes] = _read_chunks(source)
+    else:
+        size = len(source)
+        content = source
+    log = horus_logger.log
+    log.debug(f"Upload {size} bytes to {safe_url(url)}")
+    started = time.monotonic()
+    async with httpx.AsyncClient(timeout=None) as client:
+        response = await client.put(
+            url, content=content, headers={"Content-Length": str(size)}
+        )
+        response.raise_for_status()
+    log.debug(f"Uploaded {size} bytes in {time.monotonic() - started:.2f}s")
+
+
+async def download_to(url: str, dest: Path) -> None:
+    """GET the presigned *url* into *dest*, streamed, creating parents."""
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    log = horus_logger.log
+    log.debug(f"Download {safe_url(url)} to {dest}")
+    started = time.monotonic()
+    async with (
+        httpx.AsyncClient(timeout=None) as client,
+        client.stream("GET", url) as response,
+    ):
+        response.raise_for_status()
+        with dest.open("wb") as fh:
+            async for chunk in response.aiter_bytes(_CHUNK):
+                fh.write(chunk)
+    log.debug(
+        f"Downloaded {dest.stat().st_size} bytes to {dest} in "
+        f"{time.monotonic() - started:.2f}s"
+    )
+
+
+async def download(url: str) -> bytes:
+    """GET the presigned *url* into memory."""
+    async with httpx.AsyncClient(timeout=None) as client:
+        response = await client.get(url)
+        response.raise_for_status()
+        return response.content

horus_agent/broker.py

@@ -0,0 +1,223 @@
+#
+# horus_agent
+# Copyright (c) 2026 Temple Compute
+#
+# MIT License
+#
+"""
+Orchestrator-side client for the broker that relays ops to agents.
+"""
+
+import asyncio
+import random
+import time
+import weakref
+from typing import Any
+from uuid import uuid4
+
+import httpx
+from horus_runtime.core.task.exceptions import (
+    TargetLostError,
+    TaskExecutionError,
+)
+from horus_runtime.logging import horus_logger
+
+from horus_agent.logs import describe_args
+from horus_agent.protocol import (
+    IDEMPOTENT_METHODS,
+    POLL_WAIT_S,
+    AgentInfo,
+    Op,
+    OpStatus,
+    StagingUrls,
+)
+
+BACKOFF_MAX_S = 30.0
+
+_CLIENTS: weakref.WeakKeyDictionary[
+    asyncio.AbstractEventLoop, dict[tuple[str, str], "BrokerClient"]
+] = weakref.WeakKeyDictionary()
+
+
+_WORK_DIRS: dict[tuple[str, str], str] = {}
+
+
+def agent_work_dir(url: str, token: str, agent: str) -> str:
+    """
+    The working directory *agent* reported, asked once per process.
+
+    Blocking: target paths are computed synchronously, before any op runs.
+    shortcut: a sync request inside the event loop, once per agent per run;
+    move it into an async setup hook if the runtime grows one.
+
+    Raises:
+        TaskExecutionError: The broker does not know the agent, or the agent
+            has not reported its working directory yet.
+    """
+    key = (url, agent)
+    if key in _WORK_DIRS:
+        return _WORK_DIRS[key]
+    try:
+        response = httpx.get(
+            f"{url.rstrip('/')}/agents/info/{agent}",
+            headers={"Authorization": f"Bearer {token}"},
+            timeout=30.0,
+        )
+    except httpx.TransportError as exc:
+        raise TaskExecutionError(f"Broker unreachable: {exc}") from exc
+    if response.status_code >= 400:
+        raise TaskExecutionError(
+            f"Cannot look up agent {agent!r}: HTTP {response.status_code}: "
+            f"{response.text}"
+        )
+    info = AgentInfo.model_validate(response.json())
+    if not info.work_dir:
+        raise TaskExecutionError(
+            f"Agent {agent!r} has not reported its working directory yet; "
+            "start it with horus-agent run"
+        )
+    horus_logger.log.debug(f"Agent {agent!r} runs tasks in {info.work_dir}")
+    _WORK_DIRS[key] = info.work_dir
+    return info.work_dir
+
+
+def broker_client(url: str, token: str) -> "BrokerClient":
+    """
+    The running loop's client for *url* and *token*, so every target in a
+    run shares one connection pool.
+    """
+    clients = _CLIENTS.setdefault(asyncio.get_running_loop(), {})
+    key = (url, token)
+    if key not in clients:
+        clients[key] = BrokerClient(url, token)
+    return clients[key]
+
+
+class BrokerClient:
+    """Submits ops to the broker and waits for their results."""
+
+    def __init__(self, url: str, token: str) -> None:
+        """Bind to the broker at *url*, authenticating with *token*."""
+        self._http = httpx.AsyncClient(
+            base_url=url.rstrip("/") + "/",
+            headers={"Authorization": f"Bearer {token}"},
+            # A long-poll holds the response for up to POLL_WAIT_S.
+            timeout=httpx.Timeout(30.0, read=POLL_WAIT_S + 30.0),
+        )
+
+    async def call(
+        self,
+        agent: str,
+        method: str,
+        args: dict[str, Any],
+        *,
+        offline_timeout: float,
+    ) -> Any:
+        """
+        Run *method* on *agent* and return its result.
+
+        Raises:
+            TargetLostError: The agent has not polled for *offline_timeout*
+                seconds, or restarted while running a non-idempotent op.
+            TaskExecutionError: The op failed on the agent, or the broker
+                refused or stayed unreachable.
+        """
+        log = horus_logger.log
+        started = time.monotonic()
+        while True:
+            op = Op(op_id=uuid4().hex, method=method, args=args)
+            tag = f"Op {op.op_id[:8]} {method} on agent {agent!r}"
+            log.debug(f"{tag}: submit {describe_args(args)}")
+            status = await self._request(
+                "POST",
+                "agents/ops/",
+                offline_timeout,
+                json={"agent": agent, "op": op.model_dump(exclude={"run_id"})},
+            )
+            while not status.done and not status.lost:
+                seen = status.agent_seen_ago
+                if seen is None or seen > offline_timeout:
+                    raise TargetLostError(
+                        f"Agent {agent!r} is offline (last seen: "
+                        f"{'never' if seen is None else f'{seen:.0f}s ago'})"
+                    )
+                log.debug(
+                    f"{tag}: waiting (agent last polled {seen:.0f}s ago)"
+                )
+                status = await self._request(
+                    "GET",
+                    f"agents/ops/{op.op_id}",
+                    offline_timeout,
+                    params={"wait": POLL_WAIT_S},
+                )
+            if status.done:
+                break
+            if method not in IDEMPOTENT_METHODS:
+                raise TargetLostError(
+                    f"Agent {agent!r} restarted while running {method}"
+                )
+            log.debug(f"{tag}: lost on the way; sending it again")
+        log.debug(
+            f"{tag}: {'failed' if status.error else 'done'} after "
+            f"{time.monotonic() - started:.2f}s"
+        )
+        if status.error is not None:
+            raise TaskExecutionError(
+                f"{method} on agent {agent!r} failed: "
+                f"{status.error.type}: {status.error.message}"
+            )
+        return status.result
+
+    async def staging(self, offline_timeout: float) -> StagingUrls:
+        """Presigned URLs for a fresh staging object."""
+        response = await self._send("POST", "agents/staging/", offline_timeout)
+        return StagingUrls.model_validate(response.json())
+
+    async def _request(
+        self,
+        verb: str,
+        path: str,
+        offline_timeout: float,
+        **kwargs: Any,
+    ) -> OpStatus:
+        response = await self._send(verb, path, offline_timeout, **kwargs)
+        return OpStatus.model_validate(response.json())
+
+    async def _send(
+        self,
+        verb: str,
+        path: str,
+        offline_timeout: float,
+        **kwargs: Any,
+    ) -> httpx.Response:
+        """
+        Send one request, retrying with backoff while the broker is down.
+
+        A resent ``POST agents/ops/`` carries the same ``op_id``, which the
+        broker queues only once.
+        """
+        deadline = time.monotonic() + offline_timeout
+        delay = 1.0
+        while True:
+            try:
+                response = await self._http.request(verb, path, **kwargs)
+            except httpx.TransportError as exc:
+                failure = f"{type(exc).__name__}: {exc}"
+            else:
+                if response.status_code < 500:
+                    break
+                failure = f"HTTP {response.status_code}"
+            if time.monotonic() + delay > deadline:
+                raise TaskExecutionError(f"Broker unreachable: {failure}")
+            await asyncio.sleep(delay * random.uniform(0.5, 1.0))
+            delay = min(delay * 2, BACKOFF_MAX_S)
+        if response.status_code >= 400:
+            try:
+                detail = response.json().get("detail", response.text)
+            except ValueError:
+                detail = response.text
+            raise TaskExecutionError(
+                f"Broker refused {verb} {path}: "
+                f"HTTP {response.status_code}: {detail}"
+            )
+        return response

horus_agent/cli.py

@@ -0,0 +1,153 @@
+#
+# horus_agent
+# Copyright (c) 2026 Temple Compute
+#
+# MIT License
+#
+"""
+``horus-agent register`` and ``horus-agent run``.
+"""
+
+import argparse
+import asyncio
+import dataclasses
+import platform
+import signal
+import socket
+import sys
+from pathlib import Path
+
+import httpx
+
+from horus_agent import config, logs
+from horus_agent.daemon import Agent, AgentUnauthorizedError, agent_version
+from horus_agent.protocol import Registration
+
+
+def _work_dir(value: str | None) -> Path:
+    """*value* as an absolute path, created if missing."""
+    path = Path(value).expanduser() if value else config.default_work_dir()
+    path = path.resolve()
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def _register(
+    url: str, token: str, name: str | None, work_dir: str | None
+) -> int:
+    directory = _work_dir(work_dir)
+    body = {
+        "token": token,
+        "hostname": socket.gethostname(),
+        "os": platform.platform(),
+        "version": agent_version(),
+        "work_dir": str(directory),
+    }
+    if name is not None:
+        body["name"] = name
+    response = httpx.post(
+        f"{url.rstrip('/')}/agents/register", json=body, timeout=30.0
+    )
+    if response.status_code >= 400:
+        try:
+            detail = response.json().get("detail", response.text)
+        except ValueError:
+            detail = response.text
+        sys.stderr.write(f"Registration failed: {detail}\n")
+        return 1
+    registration = Registration.model_validate(response.json())
+    path = config.config_path()
+    config.save(
+        config.AgentConfig(
+            url=url,
+            agent_id=registration.agent_id,
+            name=registration.name,
+            token=registration.agent_token,
+            work_dir=str(directory),
+        ),
+        path,
+    )
+    sys.stdout.write(
+        f"Registered as {registration.name!r} (config: {path}).\n"
+        f"Tasks run in {directory}.\n"
+        "Start it with: horus-agent run\n"
+    )
+    return 0
+
+
+async def _run(work_dir: str | None) -> int:
+    path = config.config_path()
+    try:
+        registered = config.load(path)
+    except FileNotFoundError:
+        sys.stderr.write(
+            f"No registration at {path}. Run horus-agent register first.\n"
+        )
+        return 1
+    directory = _work_dir(work_dir or registered.work_dir or None)
+    agent = Agent(dataclasses.replace(registered, work_dir=str(directory)))
+    stop = asyncio.Event()
+    loop = asyncio.get_running_loop()
+
+    def _stop() -> None:
+        # A second signal falls through to the default: exit now.
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            loop.remove_signal_handler(sig)
+        stop.set()
+
+    for sig in (signal.SIGINT, signal.SIGTERM):
+        loop.add_signal_handler(sig, _stop)
+    try:
+        await agent.serve(stop)
+    except AgentUnauthorizedError as exc:
+        sys.stderr.write(f"{exc}\n")
+        return 1
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Entry point of the ``horus-agent`` command."""
+    parser = argparse.ArgumentParser(
+        prog="horus-agent",
+        description="Run Horus tasks on this machine without opening "
+        "inbound connections.",
+    )
+    common = argparse.ArgumentParser(add_help=False)
+    common.add_argument(
+        "--debug",
+        action="store_true",
+        help="Log every connection, request, op and command (also: "
+        f"{logs.DEBUG_ENV}=1). Secret values are never logged.",
+    )
+    commands = parser.add_subparsers(dest="command", required=True)
+    register = commands.add_parser(
+        "register",
+        parents=[common],
+        help="Register this machine with a one-time token.",
+    )
+    register.add_argument("--url", required=True, help="Broker API URL.")
+    register.add_argument(
+        "--token", required=True, help="One-time registration token."
+    )
+    register.add_argument(
+        "--name",
+        help="Agent name, used as the target's 'agent' field, when the "
+        "token does not set one. Default: this machine's hostname.",
+    )
+    work_dir_help = (
+        "Directory where tasks run on this machine, created if missing. "
+        f"Default: {config.default_work_dir()}."
+    )
+    register.add_argument("--work-dir", help=work_dir_help)
+    run = commands.add_parser(
+        "run", parents=[common], help="Poll for work until stopped."
+    )
+    run.add_argument(
+        "--work-dir",
+        help="Use this directory instead of the one given at registration.",
+    )
+    args = parser.parse_args(argv)
+    logs.setup(args.debug)
+    if args.command == "register":
+        return _register(args.url, args.token, args.name, args.work_dir)
+    return asyncio.run(_run(args.work_dir))

horus_agent/config.py

@@ -0,0 +1,62 @@
+#
+# horus_agent
+# Copyright (c) 2026 Temple Compute
+#
+# MIT License
+#
+"""
+The agent's registration, kept in one owner-only JSON file.
+"""
+
+import json
+import os
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+CONFIG_ENV = "HORUS_AGENT_CONFIG"
+
+
+@dataclass(frozen=True)
+class AgentConfig:
+    """What ``horus-agent register`` stores for ``horus-agent run``."""
+
+    url: str
+    """Broker base URL."""
+
+    agent_id: str
+    name: str
+    token: str
+    """The agent's own credential. Never logged."""
+    work_dir: str = ""
+    """
+    Where tasks run on this machine. Empty in configs written before it
+    existed: :func:`default_work_dir` then.
+    """
+
+
+def default_work_dir() -> Path:
+    """``~/horus-agent``: where tasks run unless ``--work-dir`` says else."""
+    return Path.home() / "horus-agent"
+
+
+def config_path() -> Path:
+    """``$HORUS_AGENT_CONFIG``, else ``$XDG_CONFIG_HOME/horus-agent``."""
+    if explicit := os.environ.get(CONFIG_ENV):
+        return Path(explicit)
+    base = os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config"
+    return Path(base) / "horus-agent" / "agent.json"
+
+
+def save(config: AgentConfig, path: Path) -> None:
+    """Write *config* to *path*, readable by its owner only."""
+    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    # O_CREAT's mode does not apply to a file that already existed.
+    os.fchmod(fd, 0o600)
+    with os.fdopen(fd, "w") as fh:
+        json.dump(asdict(config), fh)
+
+
+def load(path: Path) -> AgentConfig:
+    """Read the config written by :func:`save`."""
+    return AgentConfig(**json.loads(path.read_text()))