r2flow-agent 0.3.2__py3-none-any.whl

r2flow_agent/__init__.py

@@ -0,0 +1,5 @@
+"""R2Flow Cloud Agent — runs on each server and communicates with the orchestrator."""
+
+from __future__ import annotations
+
+__version__ = "0.3.2"

r2flow_agent/client.py

@@ -0,0 +1,350 @@
+"""HTTP client for communicating with the R2Flow orchestrator."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+import urllib.parse
+from collections.abc import Callable
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _pkg_version
+from typing import Any, cast
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+POLL_INTERVAL_SECONDS = 5
+HEARTBEAT_INTERVAL_SECONDS = 30
+REAUTH_BACKOFF_S = 60.0
+MAX_PACK_BYTES = 200 * 1024 * 1024
+MAX_ARTIFACT_BYTES = 6 * 1024 * 1024
+
+
+def _redact_for_log(text: str, *, limit: int = 500) -> str:
+    """Truncate server-echo text for logs and mask obvious secrets."""
+    import re as _re
+
+    snippet = text[:limit]
+    snippet = _re.sub(r"(?i)(bearer\s+)[A-Za-z0-9._~-]+", r"\1***", snippet)
+    snippet = _re.sub(r"(?i)(token\s*[:=]\s*)['\"]?[^'\"\s,}]+", r"\1***", snippet)
+    snippet = _re.sub(r"(?i)(password\s*[:=]\s*)['\"]?[^'\"\s,}]+", r"\1***", snippet)
+    return snippet
+
+
+def agent_version() -> str:
+    """Installed r2flow-agent version, "unknown" when not resolvable."""
+    try:
+        return _pkg_version("r2flow-agent")
+    except PackageNotFoundError:
+        return "unknown"
+
+
+def engine_version() -> str | None:
+    """Installed r2flow-engine version (None when not in this venv).
+
+    The engine rides in the *deployed process* venv, not the agent's own;
+    this reports whatever ``r2flow`` is importable from the agent venv —
+    good enough for fleet monitoring until per-run stamping lands.
+    """
+    try:
+        return _pkg_version("r2flow")
+    except PackageNotFoundError:
+        return None
+
+
+class OrchestratorError(Exception):
+    """Raised when the orchestrator returns an unexpected response."""
+
+
+class OrchestratorClient:
+    """Async HTTP client wrapping all orchestrator API calls."""
+
+    def __init__(
+        self,
+        orchestrator_url: str,
+        agent_name: str,
+        agent_url: str,
+        *,
+        join_token: str | None = None,
+        agent_id: str | None = None,
+        agent_secret: str | None = None,
+        on_credentials: Callable[[str, str], None] | None = None,
+    ) -> None:
+        self.orchestrator_url = orchestrator_url.rstrip("/")
+        self.agent_name = agent_name
+        self.agent_url = agent_url
+        self.join_token = join_token
+        self.agent_id: str | None = agent_id
+        self._secret: str | None = agent_secret
+        # Called with (agent_id, secret) whenever the orchestrator issues a
+        # new secret so the caller can persist it and survive restarts
+        # without re-registering (re-registration rotates the secret and a
+        # join token alone can no longer rotate an existing agent's secret).
+        self._on_credentials = on_credentials
+        self._reauth_lock = asyncio.Lock()
+        self._last_reauth_attempt = 0.0
+        self._http = httpx.AsyncClient(
+            base_url=self.orchestrator_url,
+            timeout=httpx.Timeout(30.0),
+        )
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    async def register(self, join_token: str | None = None) -> None:
+        """Register (or re-register) this agent with the orchestrator.
+
+        When the client already holds an agent id + secret they are sent
+        along to prove ownership of the existing agent entry; the server
+        answers with a fresh secret either way.
+        """
+        self.join_token = join_token or self.join_token
+        headers = {"Authorization": f"Bearer {self.join_token}"} if self.join_token else None
+        body: dict[str, Any] = {
+            "name": self.agent_name,
+            "url": self.agent_url,
+            "version": agent_version(),
+        }
+        if self.agent_id and self._secret:
+            body["agent_secret"] = self._secret
+        resp = await self._http.post("/api/agents", json=body, headers=headers)
+        resp.raise_for_status()
+        data = resp.json()
+        self.agent_id = data["id"]
+        self._secret = data.get("secret")
+        if self._on_credentials is not None and self.agent_id and self._secret:
+            try:
+                self._on_credentials(self.agent_id, self._secret)
+            except Exception:  # noqa: BLE001 - persistence must never kill the agent
+                logger.exception("Could not persist agent credentials")
+        logger.info("Registered with orchestrator — agent id=%s", self.agent_id)
+
+    async def _re_authenticate(self) -> None:
+        """Recover from 401s: re-register once (single-flight, with backoff).
+
+        A rotated secret (e.g. after an accidental double registration or a
+        cloud-side reset) used to leave the agent 401-ing forever; the
+        join token in the config lets it recover without a restart. Failed
+        attempts back off so a rejected agent does not hammer the server
+        on every poll cycle.
+        """
+        if time.monotonic() - self._last_reauth_attempt < REAUTH_BACKOFF_S:
+            raise OrchestratorError("Re-registration attempted recently — backing off")
+        async with self._reauth_lock:
+            if time.monotonic() - self._last_reauth_attempt < REAUTH_BACKOFF_S:
+                raise OrchestratorError("Re-registration attempted recently — backing off")
+            if not self.join_token:
+                raise OrchestratorError(
+                    "Unauthorized and no join token available for re-registration"
+                )
+            self._last_reauth_attempt = time.monotonic()
+            logger.warning("401 from orchestrator — re-registering")
+            await self.register()
+
+    async def close(self) -> None:
+        """Shut down the HTTP client."""
+        await self._http.aclose()
+
+    # ------------------------------------------------------------------
+    # Heartbeat
+    # ------------------------------------------------------------------
+
+    async def heartbeat(self) -> None:
+        """Send a single heartbeat to the orchestrator (with our versions)."""
+        body: dict[str, Any] = {"status": "online", "version": agent_version()}
+        engine = engine_version()
+        if engine is not None:
+            body["engine_version"] = engine
+        await self._post(f"/api/agents/{self._agent_id}/heartbeat", json=body)
+
+    # ------------------------------------------------------------------
+    # Assets (credentials)
+    # ------------------------------------------------------------------
+
+    async def get_assets(self) -> list[dict[str, Any]]:
+        """Fetch assets for processes: text values + decrypted credentials."""
+        resp = await self._get(f"/api/agents/{self._agent_id}/assets")
+        if resp.status_code == 401:
+            await self._re_authenticate()
+            resp = await self._get(f"/api/agents/{self._agent_id}/assets")
+        resp.raise_for_status()
+        return cast(list[dict[str, Any]], resp.json())
+
+    # ------------------------------------------------------------------
+    # Packs (bot delivery units)
+    # ------------------------------------------------------------------
+
+    async def fetch_pack(self, name: str, version: str) -> bytes:
+        """Download a pinned pack zip from the orchestrator.
+
+        Cut, byte-identical ``GET /api/packs/{name}/versions/{version}.zip``;
+        the agent secret authenticates the request (agents are pack-readable).
+        Returns the raw archive — extraction and manifest verification happen
+        in :class:`~r2flow_agent.executor.ProcessExecutor`.
+        """
+        quoted_name = urllib.parse.quote(name, safe="")
+        quoted_version = urllib.parse.quote(version, safe="")
+        path = f"/api/packs/{quoted_name}/versions/{quoted_version}.zip"
+        resp = await self._get(path)
+        if resp.status_code == 401:
+            await self._re_authenticate()
+            resp = await self._get(path)
+        resp.raise_for_status()
+        length = resp.headers.get("content-length")
+        if length is not None and length.isdigit() and int(length) > MAX_PACK_BYTES:
+            raise OrchestratorError(f"pack too large ({length} bytes)")
+        data = resp.content
+        if len(data) > MAX_PACK_BYTES:
+            raise OrchestratorError(f"pack too large ({len(data)} bytes)")
+        return data
+
+    # ------------------------------------------------------------------
+    # Polling
+    # ------------------------------------------------------------------
+
+    async def poll(self) -> list[dict[str, Any]]:
+        """Poll the orchestrator for pending commands.
+
+        On a 401 (rotated/stale secret) the agent re-registers once and
+        retries, instead of warning forever and never working again.
+
+        Returns a list of command dicts.  Each command has at least a ``type``
+        key, e.g.::
+
+            {"type": "run", "run_id": "...", "process": {…}}
+        """
+        resp = await self._get(f"/api/agents/{self._agent_id}/poll")
+        if resp.status_code == 401:
+            await self._re_authenticate()
+            resp = await self._get(f"/api/agents/{self.agent_id}/poll")
+        if resp.status_code == 204:
+            return []
+        resp.raise_for_status()
+        return cast(list[dict[str, Any]], resp.json())
+
+    # ------------------------------------------------------------------
+    # Logs & status
+    # ------------------------------------------------------------------
+
+    async def push_logs(self, run_id: str, logs: list[dict[str, Any]]) -> None:
+        """Push a batch of log entries to the orchestrator."""
+        if not logs:
+            return
+        await self._post(
+            f"/api/agents/{self._agent_id}/logs",
+            json={"run_id": run_id, "logs": logs},
+            retry=True,
+        )
+
+    async def report_status(
+        self,
+        run_id: str,
+        status: str,
+        *,
+        error: str | None = None,
+    ) -> None:
+        """Report a run status change to the orchestrator."""
+        payload: dict[str, Any] = {"run_id": run_id, "status": status}
+        if error is not None:
+            payload["error"] = error
+        await self._post(f"/api/agents/{self._agent_id}/status", json=payload, retry=True)
+
+    async def ack_deployment(
+        self,
+        deployment_id: str,
+        status: str,
+        *,
+        error: str | None = None,
+    ) -> None:
+        """Confirm a deployment result (deployed/failed) to the orchestrator."""
+        payload: dict[str, Any] = {"status": status}
+        if error is not None:
+            payload["error"] = error
+        await self._post(
+            f"/api/agents/{self._agent_id}/deployments/{deployment_id}/ack",
+            json=payload,
+        )
+
+    async def push_artifact(
+        self,
+        run_id: str,
+        filename: str,
+        content_type: str,
+        data: bytes,
+    ) -> None:
+        """Upload a binary artifact for a run (e.g. a failure screenshot)."""
+        import base64
+
+        if len(data) > MAX_ARTIFACT_BYTES:
+            raise OrchestratorError(f"artifact too large ({len(data)} bytes)")
+        payload: dict[str, Any] = {
+            "run_id": run_id,
+            "filename": filename,
+            "content_type": content_type,
+            "data_base64": base64.b64encode(data).decode("ascii"),
+        }
+        await self._post(
+            f"/api/agents/{self._agent_id}/runs/{run_id}/artifacts",
+            json=payload,
+            retry=True,
+        )
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    @property
+    def _agent_id(self) -> str:
+        if self.agent_id is None:
+            raise OrchestratorError("Agent not registered — call register() first")
+        return self.agent_id
+
+    @property
+    def agent_secret(self) -> str | None:
+        """Current agent secret (for injecting into deployed processes)."""
+        return self._secret
+
+    @property
+    def _auth_headers(self) -> dict[str, str]:
+        if self._secret is None:
+            return {}
+        return {"Authorization": f"Bearer {self._secret}"}
+
+    async def _get(self, path: str) -> httpx.Response:
+        logger.debug("GET %s", path)
+        return await self._http.get(path, headers=self._auth_headers)
+
+    async def _post(self, path: str, *, json: Any = None, retry: bool = False) -> httpx.Response:
+        import random as _random
+
+        logger.debug("POST %s", path)
+        # Status/log pushes are critical: transient 5xx/network blips are
+        # retried with backoff+jitter so runs don't stick in RUNNING forever
+        # and the fleet doesn't retry in lockstep.
+        attempts = 3 if retry else 1
+        delay = 0.5
+        resp: httpx.Response | None = None
+        for attempt in range(attempts):
+            try:
+                resp = await self._http.post(path, json=json, headers=self._auth_headers)
+            except httpx.HTTPError:
+                logger.warning("POST %s attempt %d failed (network)", path, attempt + 1)
+                resp = None
+            else:
+                if resp.status_code < 500:
+                    break
+                logger.warning("POST %s attempt %d -> %s", path, attempt + 1, resp.status_code)
+            if attempt + 1 < attempts:
+                await asyncio.sleep(delay * (0.5 + _random.random()))
+                delay *= 2.0
+        if resp is None:
+            raise OrchestratorError(f"POST {path} failed after {attempts} attempts (network)")
+        if resp.status_code >= 400:
+            # 4xx is terminal (auth/validation) — warn but don't raise so the
+            # agent loop survives; 5xx after retries also only warns.
+            # Never log raw server echo: it may contain assets/secrets.
+            logger.warning("POST %s -> %s: %s", path, resp.status_code, _redact_for_log(resp.text))
+        return resp

r2flow_agent/config.py

@@ -0,0 +1,133 @@
+"""Agent configuration file: everything the service needs to (re)start.
+
+Stored at ``%LOCALAPPDATA%\\r2flow_agent\\config.json`` (per-user: the agent
+runs inside the user's interactive session, which UI automation requires).
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+CONFIG_DIR = Path(os.environ.get("LOCALAPPDATA", str(Path.home()))) / "r2flow_agent"
+CONFIG_PATH = CONFIG_DIR / "config.json"
+
+DEFAULT_AGENT_PORT = 8001
+
+
+def load_config() -> dict[str, Any]:
+    """Read the config file; empty dict when absent."""
+    try:
+        return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError):
+        return {}
+
+
+def save_config(
+    orchestrator_url: str,
+    agent_name: str,
+    agent_url: str,
+    join_token: str | None = None,
+    log_level: str = "INFO",
+    agent_id: str | None = None,
+    agent_secret: str | None = None,
+) -> Path:
+    """Write (or merge into) the config file.
+
+    Merging matters: the service loop updates agent_id/agent_secret while
+    running and must not lose the rest of the config.
+    """
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+    payload: dict[str, Any] = {}
+    try:
+        existing = json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
+        if isinstance(existing, dict):
+            payload.update(existing)
+    except (OSError, json.JSONDecodeError):
+        pass
+    payload.update(
+        {
+            "orchestrator_url": orchestrator_url,
+            "agent_name": agent_name,
+            "agent_url": agent_url,
+            "log_level": log_level,
+        }
+    )
+    if join_token:
+        payload["join_token"] = join_token
+    if agent_id:
+        payload["agent_id"] = agent_id
+    if agent_secret:
+        payload["agent_secret"] = agent_secret
+    # Atomic write + owner-only permissions: config holds bearer secrets.
+    import os
+    import tempfile
+
+    text = json.dumps(payload, indent=2) + "\n"
+    fd, tmp_name = tempfile.mkstemp(dir=str(CONFIG_DIR), prefix=".config-", suffix=".tmp")
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8") as fh:
+            fh.write(text)
+        with contextlib.suppress(OSError):
+            os.chmod(tmp_name, 0o600)
+        Path(tmp_name).replace(CONFIG_PATH)
+        with contextlib.suppress(OSError):
+            os.chmod(CONFIG_PATH, 0o600)
+    finally:
+        with contextlib.suppress(OSError):
+            Path(tmp_name).unlink()
+    return CONFIG_PATH
+
+
+def delete_config() -> bool:
+    try:
+        CONFIG_PATH.unlink()
+        return True
+    except OSError:
+        return False
+
+
+def detect_local_ip(orchestrator_url: str) -> str:
+    """Local IP the orchestrator would see this machine at (no packets sent)."""
+    import socket
+    from urllib.parse import urlparse
+
+    parsed = urlparse(orchestrator_url)
+    host = parsed.hostname or "localhost"
+    port = parsed.port or (443 if parsed.scheme == "https" else 80)
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        sock.connect((host, port))
+        return str(sock.getsockname()[0])
+    except OSError:
+        return "127.0.0.1"
+    finally:
+        sock.close()
+
+
+def agent_url_for(orchestrator_url: str, port: int = DEFAULT_AGENT_PORT) -> str:
+    """Default public URL: http://<this-machine-ip>:<port>."""
+    return f"http://{detect_local_ip(orchestrator_url)}:{port}"
+
+
+def check_orchestrator_url(url: str) -> None:
+    """Refuse cleartext http orchestrator URLs for non-loopback hosts.
+
+    The join token and agent secret travel as Bearer credentials — plain
+    http would expose them. Loopback stays allowed for local development.
+    """
+    from urllib.parse import urlparse
+
+    parsed = urlparse(url)
+    if parsed.scheme not in ("http", "https"):
+        raise ValueError("orchestrator URL must be http(s)")
+    if parsed.scheme == "http":
+        host = (parsed.hostname or "").lower()
+        if host not in ("localhost", "127.0.0.1", "::1"):
+            raise ValueError(
+                "orchestrator URL must use https:// for non-loopback hosts "
+                "(tokens would travel in cleartext)"
+            )