levain-runtime 0.1.2__py3-none-any.whl

levain_runtime/__init__.py

@@ -0,0 +1,5 @@
+"""Self-hosted runtime for Levain agents.
+
+Keep this module light: the CLI reads the package version from
+metadata and the server imports its own modules directly.
+"""

levain_runtime/__main__.py

@@ -0,0 +1,3 @@
+from .cli import app
+
+app()

levain_runtime/chisel.py

@@ -0,0 +1,153 @@
+"""The tunnel client: a pinned chisel binary run under supervision."""
+
+from __future__ import annotations
+
+import asyncio
+import gzip
+import hashlib
+import logging
+import os
+import platform
+import stat
+from pathlib import Path
+
+import httpx
+
+from .config import Credentials
+from .process import host_env
+from .protocol import CHISEL_VERSION, RUNNER_HTTP_PORT
+
+log = logging.getLogger("levain_runtime.chisel")
+
+_RELEASES = "https://github.com/jpillora/chisel/releases/download"
+_BACKOFF_MAX_S = 30.0
+
+
+def _asset_name() -> str:
+    system = platform.system().lower()
+    machine = platform.machine().lower()
+    arch = {
+        "x86_64": "amd64",
+        "amd64": "amd64",
+        "aarch64": "arm64",
+        "arm64": "arm64",
+    }.get(machine)
+    if system not in ("linux", "darwin") or arch is None:
+        raise RuntimeError(
+            f"no chisel build for {system}/{machine}; set CHISEL_BIN"
+        )
+    return f"chisel_{CHISEL_VERSION}_{system}_{arch}"
+
+
+def download(root: Path) -> Path:
+    """Fetch the pinned release into ``<root>/bin`` after verifying its checksum."""
+    name = _asset_name()
+    base = f"{_RELEASES}/v{CHISEL_VERSION}"
+    with httpx.Client(follow_redirects=True, timeout=120) as client:
+        sums = client.get(f"{base}/chisel_{CHISEL_VERSION}_checksums.txt")
+        sums.raise_for_status()
+        expected = None
+        for line in sums.text.splitlines():
+            parts = line.split()
+            if len(parts) == 2 and parts[1] == f"{name}.gz":
+                expected = parts[0]
+        if expected is None:
+            raise RuntimeError(f"{name}.gz missing from the release checksums")
+        blob = client.get(f"{base}/{name}.gz")
+        blob.raise_for_status()
+    actual = hashlib.sha256(blob.content).hexdigest()
+    if actual != expected:
+        raise RuntimeError(f"chisel checksum mismatch: {actual} != {expected}")
+    target = root / "bin" / "chisel"
+    target.parent.mkdir(parents=True, exist_ok=True)
+    tmp = target.with_suffix(".tmp")
+    tmp.write_bytes(gzip.decompress(blob.content))
+    tmp.chmod(tmp.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
+    os.replace(tmp, target)
+    return target
+
+
+def resolve_binary(root: Path) -> Path:
+    override = os.environ.get("CHISEL_BIN")
+    if override:
+        return Path(override)
+    local = root / "bin" / "chisel"
+    if local.is_file():
+        return local
+    log.info("downloading chisel %s", CHISEL_VERSION)
+    return download(root)
+
+
+def server_url(gateway_url: str) -> str:
+    """The gateway as chisel wants it: an http(s) URL it upgrades itself.
+
+    Chisel prepends ``http://`` to any other scheme, so a ``wss://`` origin
+    would be dialled as host ``wss:`` on port 80.
+    """
+    for ws, http in (("wss://", "https://"), ("ws://", "http://")):
+        if gateway_url.startswith(ws):
+            return http + gateway_url[len(ws) :]
+    return gateway_url
+
+
+def client_argv(binary: Path, creds: Credentials) -> list[str]:
+    remotes = [f"R:{creds.reverse_port}:127.0.0.1:{RUNNER_HTTP_PORT}"]
+    remotes += [f"{f.local_port}:{f.remote}" for f in creds.forwards]
+    return [
+        str(binary),
+        "client",
+        "--keepalive",
+        "25s",
+        # The gateway pod is replaced on every platform deploy; chisel's
+        # default lets the reconnect wait grow to five minutes.
+        "--max-retry-interval",
+        "30s",
+        server_url(creds.gateway_url),
+        *remotes,
+    ]
+
+
+async def supervise(
+    binary: Path, creds: Credentials, stop: asyncio.Event
+) -> None:
+    """Keep the tunnel client alive; the secret travels via the AUTH env var."""
+    env = {**host_env(), "AUTH": f"{creds.user}:{creds.secret}"}
+    backoff = 1.0
+    while not stop.is_set():
+        started = asyncio.get_running_loop().time()
+        proc = await asyncio.create_subprocess_exec(
+            *client_argv(binary, creds),
+            env=env,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+        waiter = asyncio.create_task(stop.wait())
+        assert proc.stdout is not None
+        try:
+            while True:
+                reader = asyncio.create_task(proc.stdout.readline())
+                done, _ = await asyncio.wait(
+                    {reader, waiter}, return_when=asyncio.FIRST_COMPLETED
+                )
+                if waiter in done:
+                    reader.cancel()
+                    proc.terminate()
+                    break
+                line = reader.result()
+                if not line:
+                    break
+                log.info("chisel: %s", line.decode(errors="replace").rstrip())
+        finally:
+            waiter.cancel()
+        rc = await proc.wait()
+        if stop.is_set():
+            return
+        # A long-lived connection that died resets the backoff.
+        if asyncio.get_running_loop().time() - started > 60:
+            backoff = 1.0
+        log.warning("chisel exited with %s; restarting in %.0fs", rc, backoff)
+        try:
+            await asyncio.wait_for(stop.wait(), timeout=backoff)
+            return
+        except TimeoutError:
+            backoff = min(backoff * 2, _BACKOFF_MAX_S)

levain_runtime/cli.py

@@ -0,0 +1,365 @@
+"""``levain-runtime`` — set up, connect and run a self-hosted runtime."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import shlex
+import subprocess
+import sys
+from importlib import resources
+from pathlib import Path
+from typing import Annotated
+
+import httpx
+import typer
+import uvicorn
+
+from . import chisel, platform_api
+from .config import Credentials, Mode, RuntimeConfig, runtime_home
+from .heartbeat import package_version, run_heartbeat
+from .protocol import RUNNER_HTTP_PORT
+from .server import create_app
+from .sessions import SessionManager
+
+app = typer.Typer(
+    name="levain-runtime",
+    help="Run Levain agents on your own machine or container.",
+    add_completion=False,
+)
+
+_HOME_DIRS = ("recipes", "sdk", "sandbox", "sessions", "agents", "bin")
+_SYSTEM_LINKS = ("recipes", "sdk", "sandbox", "workspace")
+_SYNTHETIC_CONF = Path("/etc/synthetic.conf")
+_VENDORED = ("node_runner.py", "run_logging.py")
+
+
+def _version_callback(value: bool) -> None:
+    if value:
+        print(f"levain-runtime {package_version()}")
+        raise typer.Exit()
+
+
+@app.callback()
+def _main(
+    version: Annotated[
+        bool,
+        typer.Option("--version", callback=_version_callback, is_eager=True),
+    ] = False,
+) -> None:
+    """Self-hosted runtime for Levain agents."""
+
+
+def _load_config() -> RuntimeConfig:
+    home = runtime_home()
+    if not (home / "config.json").is_file():
+        raise typer.BadParameter(
+            f"no runtime at {home}; run `levain-runtime init` first"
+        )
+    return RuntimeConfig.load(home)
+
+
+def write_sandbox_files(root: Path) -> None:
+    """Materialise the node entrypoint the platform invokes per node."""
+    pkg = resources.files("levain_runtime").joinpath("sandbox")
+    (root / "sandbox").mkdir(parents=True, exist_ok=True)
+    # Shipped with a non-module suffix so linters treat them as data.
+    for name in _VENDORED:
+        data = pkg.joinpath(f"{name}.vendored").read_bytes()
+        (root / "sandbox" / name).write_bytes(data)
+
+
+def prepare_home(root: Path) -> None:
+    for name in _HOME_DIRS:
+        (root / name).mkdir(parents=True, exist_ok=True)
+    write_sandbox_files(root)
+    idle = root / "sessions" / ".idle" / "workspace"
+    idle.mkdir(parents=True, exist_ok=True)
+    link = root / "workspace"
+    if not link.is_symlink():
+        link.symlink_to(idle)
+
+
+def system_link_commands(root: Path) -> list[list[str]]:
+    """``ln`` commands still needed for ``/recipes`` & co.; empty when done."""
+    cmds: list[list[str]] = []
+    for name in _SYSTEM_LINKS:
+        path = Path("/") / name
+        target = root / name
+        if path.is_symlink():
+            if os.readlink(path) == str(target):
+                continue
+        elif path.exists():
+            raise typer.BadParameter(
+                f"{path} exists and is not a symlink; move it aside first"
+            )
+        cmds.append(["ln", "-sfn", str(target), str(path)])
+    return cmds
+
+
+def synthetic_conf_lines(root: Path, existing: str) -> list[str]:
+    """Entries ``/etc/synthetic.conf`` still needs for the system links."""
+    present: dict[str, str] = {}
+    for raw in existing.splitlines():
+        parts = raw.split("\t")
+        if parts[0].strip():
+            present[parts[0].strip()] = (
+                parts[1].strip() if len(parts) > 1 else ""
+            )
+    lines: list[str] = []
+    for name in _SYSTEM_LINKS:
+        target = str(root / name)
+        if name in present:
+            if present[name] != target:
+                raise typer.BadParameter(
+                    f"{_SYNTHETIC_CONF} maps /{name} to "
+                    f"{present[name] or 'a directory'}; change it to {target}"
+                )
+            continue
+        lines.append(f"{name}\t{target}")
+    return lines
+
+
+def ensure_system_links(root: Path) -> None:
+    if sys.platform == "darwin":
+        _ensure_synthetic_links(root)
+        return
+    sudo = [] if os.access("/", os.W_OK) else ["sudo"]
+    for cmd in system_link_commands(root):
+        full = [*sudo, *cmd]
+        print(shlex.join(full))
+        subprocess.run(full, check=True)
+
+
+def _ensure_synthetic_links(root: Path) -> None:
+    # The macOS system volume is read-only: root-level links come from
+    # /etc/synthetic.conf and appear after a restart.
+    if not system_link_commands(root):
+        return
+    existing = _SYNTHETIC_CONF.read_text() if _SYNTHETIC_CONF.is_file() else ""
+    lines = synthetic_conf_lines(root, existing)
+    if lines:
+        payload = "".join(line + "\n" for line in lines)
+        subprocess.run(
+            [
+                "sudo",
+                "sh",
+                "-c",
+                f"cat >> {shlex.quote(str(_SYNTHETIC_CONF))}",
+            ],
+            input=payload.encode(),
+            check=True,
+        )
+    print(
+        "Restart the machine to activate /recipes, /sdk, /sandbox and "
+        "/workspace, then run `levain-runtime connect`."
+    )
+
+
+@app.command()
+def init(
+    mode: Annotated[Mode, typer.Option(help="direct or docker")] = "direct",
+    image: Annotated[
+        str | None, typer.Option(help="Container image (docker mode)")
+    ] = None,
+    share_claude_login: Annotated[
+        bool,
+        typer.Option(
+            "--share-claude-login/--no-share-claude-login",
+            help="Copy this host's Claude login into docker sessions",
+        ),
+    ] = True,
+) -> None:
+    """Prepare the runtime home and, in direct mode, the system links."""
+    if mode == "docker" and not image:
+        raise typer.BadParameter("--image is required in docker mode")
+    home = runtime_home()
+    prepare_home(home)
+    existing = (
+        RuntimeConfig.load(home) if (home / "config.json").is_file() else None
+    )
+    config = RuntimeConfig(
+        root=home,
+        mode=mode,
+        image=image,
+        share_claude_login=share_claude_login,
+        credentials=existing.credentials if existing else None,
+    )
+    config.save()
+    if mode == "direct":
+        ensure_system_links(home)
+    print(f"runtime ready at {home} ({mode})")
+
+
+@app.command()
+def connect(
+    credentials: Annotated[
+        str, typer.Option(help="Credentials JSON file, or - for stdin")
+    ],
+    no_start: Annotated[
+        bool, typer.Option("--no-start", help="Store credentials only")
+    ] = False,
+) -> None:
+    """Store the platform credentials and start the runtime."""
+    config = _load_config()
+    raw = (
+        sys.stdin.read()
+        if credentials == "-"
+        else Path(credentials).read_text()
+    )
+    config.credentials = Credentials.from_dict(json.loads(raw))
+    config.save()
+    print(f"connected runtime {config.credentials.runtime_id}")
+    if not no_start:
+        _serve(config)
+
+
+@app.command()
+def register(
+    name: Annotated[
+        str, typer.Option(help="Runtime name, unique per workspace")
+    ],
+    mode: Annotated[
+        Mode | None,
+        typer.Option(
+            help="direct or docker; defaults to the initialised mode"
+        ),
+    ] = None,
+    image: Annotated[
+        str | None, typer.Option(help="Container image (docker mode)")
+    ] = None,
+    model_auth: Annotated[
+        str,
+        typer.Option(
+            help="platform: model calls on your Levain plan; "
+            "local: the provider login on this machine"
+        ),
+    ] = "platform",
+    api_url: Annotated[
+        str | None,
+        typer.Option(help="Levain API base URL", envvar="LEVAIN_API"),
+    ] = None,
+) -> None:
+    """Register this machine with Levain and store its credentials.
+
+    Needs a workspace API key in LEVAIN_API_KEY (prompted when unset).
+    """
+    config = _load_config()
+    if mode is not None and mode != config.mode:
+        raise typer.BadParameter(
+            f"this runtime was initialised in {config.mode} mode; "
+            f"run `levain-runtime init --mode {mode}` first"
+        )
+    if model_auth not in ("platform", "local"):
+        raise typer.BadParameter("--model-auth must be platform or local")
+    if config.credentials is not None:
+        raise typer.BadParameter(
+            f"already registered as runtime {config.credentials.runtime_id}; "
+            "delete it on Levain and remove credentials from config.json "
+            "to register again"
+        )
+    api_key = os.environ.get("LEVAIN_API_KEY") or typer.prompt(
+        "Workspace API key", hide_input=True
+    )
+    body: dict[str, object] = {
+        "name": name,
+        "mode": config.mode,
+        "model_auth": model_auth,
+    }
+    if config.mode == "docker":
+        body["image"] = image or config.image
+    base = api_url or platform_api.DEFAULT_API_URL
+    try:
+        created = platform_api.register_runtime(base, api_key, body)
+    except (httpx.HTTPError, ValueError) as exc:
+        raise typer.BadParameter(f"registration failed: {exc}") from None
+    creds = dict(created["credentials"])
+    creds.setdefault("api_url", base)
+    config.credentials = Credentials.from_dict(creds)
+    config.save()
+    print(f"registered runtime {name} ({config.credentials.runtime_id})")
+    print("next: levain-runtime up")
+
+
+@app.command()
+def up() -> None:
+    """Start the node server, the tunnel and the heartbeat."""
+    config = _load_config()
+    if config.credentials is None:
+        raise typer.BadParameter(
+            "no credentials; run `levain-runtime connect`"
+        )
+    _serve(config)
+
+
+@app.command()
+def status() -> None:
+    """Show the configuration and the sessions on this runtime."""
+    config = _load_config()
+    manager = SessionManager(config)
+    print(f"home: {config.root}")
+    print(
+        f"mode: {config.mode}" + (f" ({config.image})" if config.image else "")
+    )
+    if config.credentials:
+        print(f"runtime: {config.credentials.runtime_id}")
+        print(f"gateway: {config.credentials.gateway_url}")
+    else:
+        print("runtime: not connected")
+    if config.mode == "direct":
+        pending = system_link_commands(config.root)
+        print("system links: " + ("ok" if not pending else "missing"))
+    print(
+        f"chisel: {os.environ.get('CHISEL_BIN') or config.root / 'bin' / 'chisel'}"
+    )
+    sessions = manager.store.all()
+    print(
+        f"sessions: {len(sessions)} ({len(manager.active_sessions())} active)"
+    )
+    for meta in sessions:
+        print(f"  {meta.sid} {meta.state} {meta.workspace_dir}")
+
+
+def _serve(config: RuntimeConfig) -> None:
+    logging.basicConfig(
+        level=logging.INFO, format="%(levelname)s %(name)s %(message)s"
+    )
+    creds = config.credentials
+    assert creds is not None
+    # The forward list lives on the platform; pick up changes on start.
+    if platform_api.refresh_connection(creds):
+        config.save()
+    binary = chisel.resolve_binary(config.root)
+    asyncio.run(_run_all(config, creds, binary))
+
+
+async def _run_all(
+    config: RuntimeConfig, creds: Credentials, binary: Path
+) -> None:
+    application = create_app(config)
+    manager: SessionManager = application.state.manager
+    server = uvicorn.Server(
+        uvicorn.Config(
+            application,
+            host="127.0.0.1",
+            port=RUNNER_HTTP_PORT,
+            log_level="warning",
+            lifespan="off",
+        )
+    )
+    stop = asyncio.Event()
+
+    async def follow_server() -> None:
+        # uvicorn owns the signal handlers; mirror its exit flag.
+        while not server.should_exit:
+            await asyncio.sleep(0.5)
+        stop.set()
+
+    await asyncio.gather(
+        server.serve(),
+        follow_server(),
+        chisel.supervise(binary, creds, stop),
+        run_heartbeat(creds, manager, stop),
+    )

levain_runtime/config.py

@@ -0,0 +1,117 @@
+"""Runtime configuration: home directory, mode, and platform credentials."""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, Literal
+
+Mode = Literal["direct", "docker"]
+
+CONFIG_FILE = "config.json"
+
+
+def runtime_home() -> Path:
+    raw = os.environ.get("LEVAIN_RUNTIME_HOME")
+    return (
+        Path(raw).expanduser()
+        if raw
+        else Path("~/.levain/runtime").expanduser()
+    )
+
+
+@dataclass(frozen=True)
+class Forward:
+    local_port: int
+    remote: str
+
+
+@dataclass
+class Credentials:
+    runtime_id: str
+    api_url: str
+    gateway_url: str
+    user: str
+    secret: str
+    reverse_port: int
+    forwards: list[Forward] = field(default_factory=list)
+    heartbeat_seconds: int = 30
+
+    @classmethod
+    def from_dict(cls, raw: dict[str, Any]) -> Credentials:
+        forwards = [
+            Forward(local_port=int(f["local_port"]), remote=str(f["remote"]))
+            for f in raw.get("forwards") or []
+        ]
+        return cls(
+            runtime_id=str(raw["runtime_id"]),
+            api_url=str(raw["api_url"]).rstrip("/"),
+            gateway_url=str(raw["gateway_url"]),
+            user=str(raw["user"]),
+            secret=str(raw["secret"]),
+            reverse_port=int(raw["reverse_port"]),
+            forwards=forwards,
+            heartbeat_seconds=int(raw.get("heartbeat_seconds") or 30),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    def apply_connection(self, raw: dict[str, Any]) -> bool:
+        """Take the platform's current connection settings; True if changed."""
+        fresh = Credentials.from_dict(
+            {
+                **self.to_dict(),
+                **{k: v for k, v in raw.items() if k != "secret"},
+            }
+        )
+        changed = fresh != self
+        self.gateway_url = fresh.gateway_url
+        self.user = fresh.user
+        self.reverse_port = fresh.reverse_port
+        self.forwards = fresh.forwards
+        self.heartbeat_seconds = fresh.heartbeat_seconds
+        return changed
+
+
+@dataclass
+class RuntimeConfig:
+    root: Path
+    mode: Mode = "direct"
+    image: str | None = None
+    share_claude_login: bool = True
+    credentials: Credentials | None = None
+
+    @property
+    def path(self) -> Path:
+        return self.root / CONFIG_FILE
+
+    @classmethod
+    def load(cls, root: Path) -> RuntimeConfig:
+        raw = json.loads((root / CONFIG_FILE).read_text())
+        creds = raw.get("credentials")
+        return cls(
+            root=root,
+            mode=raw.get("mode", "direct"),
+            image=raw.get("image"),
+            share_claude_login=bool(raw.get("share_claude_login", True)),
+            credentials=Credentials.from_dict(creds) if creds else None,
+        )
+
+    def save(self) -> None:
+        self.root.mkdir(parents=True, exist_ok=True)
+        payload = {
+            "mode": self.mode,
+            "image": self.image,
+            "share_claude_login": self.share_claude_login,
+            "credentials": self.credentials.to_dict()
+            if self.credentials
+            else None,
+        }
+        tmp = self.path.with_suffix(".json.tmp")
+        tmp.write_text(json.dumps(payload, indent=2))
+        # Credentials live here; keep the file private.
+        os.chmod(tmp, 0o600)
+        os.replace(tmp, self.path)

levain_runtime/heartbeat.py

@@ -0,0 +1,47 @@
+"""Periodic liveness report to the platform API."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from importlib.metadata import PackageNotFoundError, version
+
+import httpx
+
+from .config import Credentials
+from .sessions import SessionManager
+
+log = logging.getLogger("levain_runtime.heartbeat")
+
+
+def package_version() -> str:
+    try:
+        return version("levain-runtime")
+    except PackageNotFoundError:
+        return "dev"
+
+
+async def run_heartbeat(
+    creds: Credentials, manager: SessionManager, stop: asyncio.Event
+) -> None:
+    url = f"{creds.api_url}/api/v1/runtimes/{creds.runtime_id}/heartbeat"
+    headers = {"Authorization": f"Bearer {creds.secret}"}
+    async with httpx.AsyncClient(timeout=15) as client:
+        while not stop.is_set():
+            payload = {
+                "version": package_version(),
+                "mode": manager.mode,
+                "active_sessions": len(manager.active_sessions()),
+            }
+            try:
+                resp = await client.post(url, json=payload, headers=headers)
+                if resp.status_code >= 400:
+                    log.warning("heartbeat rejected: %s", resp.status_code)
+            except httpx.HTTPError as exc:
+                log.warning("heartbeat failed: %s", exc)
+            try:
+                await asyncio.wait_for(
+                    stop.wait(), timeout=creds.heartbeat_seconds
+                )
+            except TimeoutError:
+                continue