plane-agentd/agentd/transports/__init__.py

"""Transports: how agentd reaches a terminal's machine (files, detached launch, kill).

Interface (duck-typed, see docs/executor-contract.md):
  os, join(*parts), quote(s), default_run_root()
  mkdir(path, clean=False), write_text(path, text), read_text(path) -> str|None,
  grep_first(path, needle) -> str|None, tail(path, n) -> str
  fetch(run_dir, ops) -> [result per op]   (ops: ("read", name) | ("grep", name, needle) | ("tail", name, n))
  runner(run_dir, cwd, command, stdout_name) -> (script_name, script_text)
  start_run(task, run_dir, files, script_name) (remote only), launch(task, script_path, run_dir), kill(task, run_dir),
  forget(task), process_count(name) -> int|None (None = unreachable)

Remote terminals are reached through a channel (design §9): `warpgate-ssh` (agent-access.py run through the
Warpgate gateway) or `direct-ssh` (OpenSSH over Tailscale with a dedicated key, optional warpgate fallback)."""
import base64
import functools
import re
import time


class TransportError(Exception):
    pass


class Unreached(TransportError):
    """The remote command certainly did not run (connect / auth / agent failure, direct path unavailable): the
    only kind of failure a direct-ssh terminal may repeat on its fallback channel."""


class HostKeyError(TransportError):
    """The terminal presented a host key other than the pinned one: a security error, never retried elsewhere."""


class CallLog:
    """One line per remote call -> `path` (remote.log); notable events (fallback, host key) also -> `event`."""

    def __init__(self, path=None, event=None, echo=False):
        self.path, self.event_fn, self.echo = path, event, echo
        self._once = set()

    def call(self, terminal, via, op, ms, rc, note=""):
        line = "term=%s via=%s op=%s ms=%d rc=%s%s" % (terminal, via, op or "-", ms, rc, " " + note if note else "")
        if self.echo:
            print(line, flush=True)
        if self.path:
            try:
                self.path.parent.mkdir(parents=True, exist_ok=True)
                with open(self.path, "a", encoding="utf-8") as fh:
                    fh.write(time.strftime("%Y-%m-%dT%H:%M:%S") + " " + line + "\n")
            except OSError:
                pass

    def event(self, msg, once=None):
        """`once` = a key: the message is logged only the first time per process."""
        if once is not None:
            if once in self._once:
                return
            self._once.add(once)
        if self.event_fn:
            self.event_fn(msg)


NULL_LOG = CallLog()


def operation(name):
    """Method decorator: names the remote operation for the call log (innermost wins, restored on return)."""
    def deco(fn):
        @functools.wraps(fn)
        def wrapper(self, *args, **kwargs):
            prev, self._op = getattr(self, "_op", None), name
            try:
                return fn(self, *args, **kwargs)
            finally:
                self._op = prev
        return wrapper
    return deco


def fetch_each(t, run_dir, ops):
    """`fetch` built from the single-file primitives (one call per op): local transport and fallbacks."""
    out = []
    for o in ops:
        path = t.join(run_dir, o[1])
        if o[0] == "read":
            out.append(t.read_text(path))
        elif o[0] == "grep":
            out.append(t.grep_first(path, o[2]))
        elif o[0] == "tail":
            out.append(t.tail(path, o[2]))
        else:
            raise ValueError("unknown fetch op %r" % (o,))
    return out


def parse_fetch(out, ops, what=""):
    """Decode the `F<i>=<base64>` / `M<i>` (missing) lines of a remote fetch script into results per op:
    read -> text|None, grep -> line|None, tail -> text ("" when missing)."""
    res = []
    for i, o in enumerate(ops):
        m = re.search(r"^F%d=([A-Za-z0-9+/=]*)$" % i, out, re.M)
        if m:
            text = base64.b64decode(m.group(1)).decode("utf-8", "replace") if m.group(1) else ""
        elif re.search(r"^M%d$" % i, out, re.M):
            text = None
        else:
            raise TransportError("fetch %s: unexpected output: %s" % (what, out.strip()[-200:]))
        if o[0] == "grep":
            text = (text or "").rstrip("\n") or None
        elif o[0] == "tail":
            text = (text or "").rstrip("\n")
        res.append(text)
    return res


class RoutedTransport:
    """direct-ssh terminal: every operation runs on the primary (direct) transport; when it fails before the remote
    command ran (`Unreached`) the operation is repeated on the fallback (warpgate-ssh) transport, which then serves
    the rest of this transport's life (one tick). Each direct operation is a single remote call, so a repeat never
    doubles a side effect. Host-key failures and failures after the command may have run propagate unchanged."""
    PURE = ("join", "quote", "default_run_root", "runner")

    def __init__(self, terminal, primary, fallback=None, log=None):
        self.terminal, self.primary, self.fallback = terminal, primary, fallback
        self.log = log or NULL_LOG
        self.on_fallback = False

    def __getattr__(self, name):
        attr = getattr(self.primary, name)
        if name in self.PURE or not callable(attr):
            return attr

        def routed(*args, **kwargs):
            if self.on_fallback:
                return getattr(self.fallback, name)(*args, **kwargs)
            try:
                return attr(*args, **kwargs)
            except Unreached as exc:
                if self.fallback is None:
                    raise
                self.on_fallback = True
                self.log.event("%s: direct-ssh failed before the command ran (%s); using the warpgate-ssh fallback "
                               "for the rest of this tick" % (self.terminal, exc))
                return getattr(self.fallback, name)(*args, **kwargs)
        return routed


def make_transport(cfg, terminal, log=None, force=None):
    """`force` = direct-ssh | warpgate-ssh: that channel only, no fallback (agentd ping)."""
    t = cfg.terminals[terminal]
    kind = force or t["transport"]
    if kind == "local":
        from .local import LocalTransport
        return LocalTransport(cfg, terminal)
    if kind == "warpgate-ssh":
        from .warpgate_ssh import WarpgateSshTransport
        return WarpgateSshTransport(cfg, terminal, log=log)
    if kind == "direct-ssh":
        from .direct_ssh import DirectChannel
        from .warpgate_ssh import WarpgateSshTransport
        primary = WarpgateSshTransport(cfg, terminal, channel=DirectChannel(cfg, terminal, log))
        fallback = None
        if not force and t.get("fallback") == "warpgate-ssh":
            fallback = WarpgateSshTransport(cfg, terminal, log=log, label="fallback")
        return RoutedTransport(terminal, primary, fallback, log)
    raise TransportError("unknown transport %s" % kind)