multi-agent-mux/.agents/skills/lib_py/delegation/remote.py

"""remote.py — deliver a job to an agent on another host over the broker (Issue #18).

The herdr transport injects the prompt into a pane on this machine and reads
the agent's `started`/`progress`/`completed` events back over the broker. A
remote agent shares neither the pane nor the orchestrator's disk, so the
prompt itself has to travel over the wire too:

  orchestrator  --retained `<topic_prefix>/dispatch`-->  remote runner
  remote runner --`<topic_prefix>/events` + `/state`-->  orchestrator's subscriber

`publish_dispatch` (wrapper side) publishes one retained dispatch message
carrying the brief, the job's auth_token / dispatch_token / attempt_id and the
liveness contract (heartbeat N, misses K, ack T). `run_runner` (remote side)
subscribes for its session name, mirrors the job into its own registry,
publishes `started` with the automatic heartbeat, runs the agent command, and
publishes the terminal event. Every event it publishes is signed with the
job's auth_token and carries the orchestrator's dispatch identity, so the
orchestrator's unchanged subscriber accepts it; the subscriber's `--remote`
mode mirrors those events into the orchestrator's registry.

Trust: a dispatch makes a runner execute a command, so it is signed with a
pre-shared secret (`MAM_DISPATCH_SECRET`, from each host's own .mam.env) and a
runner ignores anything unsigned or mis-signed. The auth_token travels inside
that signed message; use this only on a private broker (see MESSAGING.md).
"""
from __future__ import annotations

import argparse
import hashlib
import hmac
import json
import logging
import os
import queue
import shlex
import socket
import subprocess
import sys
import time
from typing import Any, Dict, Optional

from lib_py.delegation import mqtt_common, registry
from lib_py.delegation.mqtt_common import (
    DEFAULT_TOPIC_ROOT,
    SCHEMA_VERSION,
    broker_config_from_env,
    broker_config_from_job,
    default_registry_dir,
    load_job,
    make_client,
)
from lib_py.delegation.publisher import OneShotPublisher

logger = logging.getLogger("delegate_job.remote")

DISPATCH_KIND = "dispatch"
SECRET_ENV = "MAM_DISPATCH_SECRET"
SESSION_PREFIX = "remote:"


def is_remote_session(agent_session: Optional[str]) -> bool:
    return bool(agent_session) and str(agent_session).startswith(SESSION_PREFIX)


def session_name(agent_session: str) -> str:
    return agent_session[len(SESSION_PREFIX):] if is_remote_session(agent_session) else agent_session


def dispatch_topic_for(job: Dict[str, Any]) -> str:
    prefix = job.get("topic_prefix") or mqtt_common.topic_prefix_for(job["job_id"])
    return f"{prefix}/dispatch"


def _canonical(payload: Dict[str, Any]) -> bytes:
    body = {k: v for k, v in payload.items() if k != "sig"}
    return json.dumps(body, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


def sign_dispatch(payload: Dict[str, Any], secret: str) -> Dict[str, Any]:
    payload = dict(payload)
    payload["sig"] = hmac.new(secret.encode("utf-8"), _canonical(payload), hashlib.sha256).hexdigest()
    return payload


def verify_dispatch(payload: Dict[str, Any], secret: Optional[str]) -> bool:
    """True when the dispatch carries a valid signature for ``secret``. With
    no secret configured on the runner side nothing is accepted."""
    if not secret:
        return False
    sig = payload.get("sig")
    if not isinstance(sig, str):
        return False
    expected = hmac.new(secret.encode("utf-8"), _canonical(payload), hashlib.sha256).hexdigest()
    return hmac.compare_digest(sig, expected)


def build_dispatch(job: Dict[str, Any], session: str, instructions: str) -> Dict[str, Any]:
    return {
        "schema_version": SCHEMA_VERSION,
        "kind": DISPATCH_KIND,
        "job_id": job["job_id"],
        "session": session,
        "instructions": instructions,
        "brief": job.get("prompt", ""),
        "role": job.get("role"),
        "agent": job.get("agent"),
        "auth_token": job.get("auth_token"),
        "dispatch_token": job.get("dispatch_token"),
        "attempt_id": job.get("attempt_id", job.get("iteration", 1)),
        "timeout_sec": job.get("timeout_sec", 3600),
        "idle_timeout_sec": job.get("idle_timeout_sec"),
        "heartbeat_sec": job.get("heartbeat_sec", registry.DEFAULT_HEARTBEAT_SEC),
        "heartbeat_misses": job.get("heartbeat_misses", registry.DEFAULT_HEARTBEAT_MISSES),
        "ack_timeout_sec": job.get("ack_timeout_sec", registry.DEFAULT_ACK_TIMEOUT_SEC),
        "topic_prefix": job.get("topic_prefix") or mqtt_common.topic_prefix_for(job["job_id"]),
        "origin": socket.gethostname(),
        "timestamp": mqtt_common._utcnow(),
    }


def publish_dispatch(job_id: str, registry_dir: str, session: str, instructions: str,
                     secret: Optional[str] = None, allow_unsigned: bool = False) -> str:
    """Publish the retained dispatch for ``job_id``; returns the topic."""
    job = load_job(job_id, registry_dir)
    secret = secret or os.environ.get(SECRET_ENV)
    payload = build_dispatch(job, session, instructions)
    if secret:
        payload = sign_dispatch(payload, secret)
    elif not allow_unsigned:
        raise RuntimeError("{} is not set: refusing to publish an unsigned dispatch "
                           "(set it in .mam.env on both hosts, or pass --allow-unsigned)".format(SECRET_ENV))
    topic = dispatch_topic_for(job)
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    mqtt_common.with_retry(OneShotPublisher().publish_raw, attempts=3,
                           exceptions=(OSError, TimeoutError, ConnectionError, ValueError))(
        broker_config_from_job(job), topic, body, True)
    mqtt_common.append_event(job_id, {"event": "dispatched", "topic": topic, "session": session,
                                      "attempt_id": payload["attempt_id"]})
    logger.info("dispatched job %s to remote session %s on %s", job_id, session, topic)
    return topic


def clear_dispatch(config, topic: str) -> None:
    """Retire a retained dispatch (empty retained payload) so a runner that
    restarts does not execute it again."""
    try:
        OneShotPublisher().publish_raw(config, topic, b"", True)
    except Exception as exc:  # pragma: no cover - best effort
        logger.warning("could not clear dispatch %s: %s", topic, exc)


# ---------------------------------------------------------------------------
# remote side
# ---------------------------------------------------------------------------

def mirror_job(d: Dict[str, Any], registry_dir: str, session: str) -> str:
    """Create the local job record a remote agent publishes against.

    The record carries the orchestrator's job_id, auth_token, dispatch_token
    and attempt_id, so `publish_event` on this host produces events the
    orchestrator's subscriber accepts. A previous attempt's record (a
    re-dispatch) is retired first: creation is the only write allowed here.
    """
    jid = d["job_id"]
    path = mqtt_common._job_path(jid, registry_dir)
    if path.exists():
        old = load_job(jid, registry_dir)
        if (old.get("dispatch_token") == d.get("dispatch_token")
                and old.get("status") not in ("completed", "error", "cancelled")):
            return jid  # same attempt, already mirrored (duplicate delivery)
        with mqtt_common.registry_lock(registry_dir):
            path.unlink()
            log = mqtt_common._job_path(jid, registry_dir).with_suffix(".events.log")
            if log.exists():
                log.unlink()
    registry.register_job(
        prompt=d.get("brief", ""),
        agent=d.get("agent") or "remote",
        agent_session=SESSION_PREFIX + session,
        role=d.get("role") or "Worker",
        timeout_sec=int(d.get("timeout_sec") or 3600),
        idle_timeout_sec=d.get("idle_timeout_sec"),
        registry_dir=registry_dir,
        job_id=jid,
        auth_token=d.get("auth_token"),
        heartbeat_sec=int(d.get("heartbeat_sec") or registry.DEFAULT_HEARTBEAT_SEC),
        heartbeat_misses=int(d.get("heartbeat_misses") or registry.DEFAULT_HEARTBEAT_MISSES),
        ack_timeout_sec=int(d.get("ack_timeout_sec") or registry.DEFAULT_ACK_TIMEOUT_SEC),
        dispatch_token=d.get("dispatch_token"),
        attempt_id=int(d.get("attempt_id") or 1),
        transport="remote",
    )
    brief_dir = os.path.join(registry_dir, jid)
    os.makedirs(brief_dir, exist_ok=True)
    with open(os.path.join(brief_dir, "brief.md"), "w", encoding="utf-8") as fh:
        fh.write(d.get("brief", ""))
    return jid


def _publish(jid: str, registry_dir: str, event: str, detail: str, d: Dict[str, Any],
             heartbeat: Optional[int] = None) -> int:
    from lib_py.delegation import publish_event
    argv = ["--registry-dir", registry_dir, "--job", jid, "--event", event, "--detail", detail]
    if d.get("dispatch_token"):
        argv += ["--expect-dispatch-token", d["dispatch_token"]]
    if d.get("attempt_id") is not None:
        argv += ["--attempt-id", str(d["attempt_id"])]
    if heartbeat:
        argv += ["--heartbeat", str(heartbeat)]
    return publish_event.main(argv)


def handle_dispatch(d: Dict[str, Any], registry_dir: str, session: str,
                    exec_cmd: Optional[str], python: Optional[str] = None) -> int:
    """Run one dispatched job to its terminal event. Returns the agent rc."""
    jid = mirror_job(d, registry_dir, session)
    heartbeat = int(d.get("heartbeat_sec") or registry.DEFAULT_HEARTBEAT_SEC)
    rc_started = _publish(jid, registry_dir, "started", "remote runner {} claimed the job".format(session),
                          d, heartbeat=heartbeat)
    if rc_started not in (0, 2):
        logger.error("could not publish started for %s (rc %s)", jid, rc_started)
        return 1
    job = load_job(jid, registry_dir)
    rc = 0
    if exec_cmd:
        py = python or sys.executable
        skills_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        env = dict(os.environ)
        env["PYTHONPATH"] = skills_root + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
        env.update({
            "JOB_ID": jid,
            "MAM_JOB_ID": jid,
            "MAM_JOB_REGISTRY_DIR": registry_dir,
            "MAM_JOB_BRIEF": os.path.join(registry_dir, jid, "brief.md"),
            "MAM_JOB_INSTRUCTIONS": d.get("instructions", ""),
            "MAM_PUBLISH_CMD": " ".join(shlex.quote(x) for x in [
                py, "-m", "lib_py.delegation.publish_event", "--registry-dir", registry_dir, "--job", jid]
                + (["--expect-dispatch-token", d["dispatch_token"]] if d.get("dispatch_token") else [])
                + (["--attempt-id", str(d["attempt_id"])] if d.get("attempt_id") is not None else [])),
        })
        timeout = float(job.get("timeout_sec", 3600))
        try:
            res = subprocess.run(exec_cmd, shell=True, env=env, capture_output=True, text=True,
                                 timeout=timeout, cwd=os.getcwd())
            rc = res.returncode
            tail = (res.stdout or "").strip().splitlines()
            detail = tail[-1][:200] if tail else "exit {}".format(rc)
            if rc != 0:
                err = (res.stderr or "").strip().splitlines()
                detail = "exit {}: {}".format(rc, (err[-1] if err else detail)[:200])
        except subprocess.TimeoutExpired:
            rc, detail = 124, "agent command exceeded timeout_sec {}".format(int(timeout))
        cur = load_job(jid, registry_dir)
        if cur.get("status") not in ("completed", "error", "cancelled"):
            _publish(jid, registry_dir, "completed" if rc == 0 else "error", detail, d)
        return rc
    # No command: an agent on this host publishes the terminal event itself.
    deadline = time.monotonic() + float(job.get("timeout_sec", 3600))
    while time.monotonic() < deadline:
        cur = load_job(jid, registry_dir)
        if cur.get("status") in ("completed", "error", "cancelled"):
            return 0 if cur.get("status") == "completed" else 1
        time.sleep(1.0)
    _publish(jid, registry_dir, "error", "deadline_timeout (remote runner)", d)
    return 2


def run_runner(session: str, registry_dir: str, exec_cmd: Optional[str] = None, once: bool = False,
               secret: Optional[str] = None, allow_unsigned: bool = False,
               timeout: Optional[float] = None, python: Optional[str] = None) -> int:
    secret = secret or os.environ.get(SECRET_ENV)
    if not secret and not allow_unsigned:
        logger.error("%s is not set: a runner must not execute unsigned dispatches "
                     "(pass --allow-unsigned only on a private broker)", SECRET_ENV)
        return 1
    cfg = broker_config_from_env()
    client = make_client("remote-runner", cfg)
    inbox: "queue.Queue[Dict[str, Any]]" = queue.Queue()
    topic = f"{DEFAULT_TOPIC_ROOT}/+/dispatch"

    def on_connect(_c, _u, _flags, reason_code, _props):
        if mqtt_common.reason_code_value(reason_code) != 0:
            logger.error("broker connection failed: rc=%s", reason_code)
            return
        _c.subscribe(topic, qos=1)
        print(f"SUBSCRIBED {topic}", flush=True)

    def on_message(_c, _u, msg):
        if not msg.payload:
            return  # a cleared (retired) dispatch
        try:
            d = json.loads(msg.payload.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError):
            logger.warning("drop unparseable dispatch on %s", msg.topic)
            return
        if not isinstance(d, dict) or d.get("kind") != DISPATCH_KIND or d.get("schema_version") != SCHEMA_VERSION:
            return
        if d.get("session") != session:
            return
        if secret and not verify_dispatch(d, secret):
            logger.warning("drop dispatch for %s: bad or missing signature", d.get("job_id"))
            return
        if not secret and not allow_unsigned:
            return
        d["_topic"] = msg.topic
        inbox.put(d)

    client.on_connect = on_connect
    client.on_message = on_message
    client.reconnect_delay_set(min_delay=1, max_delay=16)
    mqtt_common.with_retry(lambda: client.connect(cfg.host, cfg.port, cfg.keepalive),
                           attempts=3, base_delay=0.2, max_delay=1.0)()
    client.loop_start()
    deadline = time.monotonic() + timeout if timeout else None
    handled = 0
    try:
        while True:
            wait = 1.0 if deadline is None else max(0.0, min(1.0, deadline - time.monotonic()))
            try:
                d = inbox.get(timeout=wait)
            except queue.Empty:
                if deadline is not None and time.monotonic() >= deadline:
                    logger.warning("runner timeout after %d job(s)", handled)
                    return 0 if handled else 2
                continue
            print("DISPATCH job={} attempt={} from={}".format(d.get("job_id"), d.get("attempt_id"), d.get("origin")),
                  flush=True)
            # Retire the retained dispatch before running so a restart cannot replay it.
            clear_dispatch(cfg, d["_topic"])
            rc = handle_dispatch(d, registry_dir, session, exec_cmd, python=python)
            handled += 1
            print("DONE job={} rc={}".format(d.get("job_id"), rc), flush=True)
            if once:
                return 0
    finally:
        client.loop_stop()
        try:
            client.disconnect()
        except Exception:
            pass


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description="remote transport for delegated jobs (Issue #18)")
    sub = parser.add_subparsers(dest="command", required=True)
    p_d = sub.add_parser("dispatch", help="publish the retained dispatch for a registered job")
    p_d.add_argument("--job", required=True)
    p_d.add_argument("--registry-dir", default=default_registry_dir())
    p_d.add_argument("--session", required=True, help="remote session name (without the remote: prefix)")
    p_d.add_argument("--instructions", default="")
    p_d.add_argument("--allow-unsigned", action="store_true")
    p_r = sub.add_parser("run", help="run a remote session: execute dispatches addressed to it")
    p_r.add_argument("--session", required=True)
    p_r.add_argument("--registry-dir", default=default_registry_dir())
    p_r.add_argument("--exec", dest="exec_cmd", default=None,
                     help="shell command run per job (env: JOB_ID, MAM_JOB_BRIEF, MAM_PUBLISH_CMD, ...)")
    p_r.add_argument("--once", action="store_true", help="exit after one job")
    p_r.add_argument("--timeout", type=float, default=None, help="exit after this many seconds")
    p_r.add_argument("--allow-unsigned", action="store_true")
    p_r.add_argument("--python", default=None)
    parser.add_argument("-v", "--verbose", action="store_true")
    args = parser.parse_args(argv)
    mqtt_common.ensure_env_loaded()
    mqtt_common.setup_logging(logging.DEBUG if args.verbose else logging.INFO)
    if args.command == "dispatch":
        try:
            print(publish_dispatch(args.job, args.registry_dir, args.session, args.instructions,
                                   allow_unsigned=args.allow_unsigned))
        except (FileNotFoundError, RuntimeError) as exc:
            logger.error("%s", exc)
            return 1
        except Exception as exc:
            logger.error("dispatch publish failed: %s", exc)
            return 2
        return 0
    return run_runner(args.session, args.registry_dir, args.exec_cmd, args.once,
                      allow_unsigned=args.allow_unsigned, timeout=args.timeout, python=args.python)


if __name__ == "__main__":
    sys.exit(main())