levain-runtime 0.1.2__py3-none-any.whl

levain_runtime/platform_api.py

@@ -0,0 +1,83 @@
+"""Calls to the platform API other than the heartbeat."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import httpx
+
+from .config import Credentials
+
+log = logging.getLogger("levain_runtime.platform")
+
+DEFAULT_API_URL = "https://api.levainlabs.com"
+# Tests swap in an ``httpx.MockTransport``.
+TRANSPORT: httpx.BaseTransport | None = None
+
+
+def _client(**kwargs: Any) -> httpx.Client:
+    return httpx.Client(timeout=30, transport=TRANSPORT, **kwargs)
+
+
+def fetch_connection(creds: Credentials) -> dict[str, Any]:
+    """The platform's current gateway, port and forwards for this runtime."""
+    with _client() as client:
+        resp = client.get(
+            f"{creds.api_url}/api/v1/runtimes/{creds.runtime_id}/connection",
+            headers={"Authorization": f"Bearer {creds.secret}"},
+        )
+    resp.raise_for_status()
+    data = resp.json()
+    if not isinstance(data, dict):
+        raise ValueError("connection response is not an object")
+    return data
+
+
+def refresh_connection(creds: Credentials) -> bool:
+    """Apply the platform's connection settings in place; False if unreachable.
+
+    A runner that cannot reach the API keeps its stored settings so a
+    platform outage never blocks a tunnel that would otherwise come up.
+    """
+    try:
+        data = fetch_connection(creds)
+    except (httpx.HTTPError, ValueError) as exc:
+        log.warning("could not refresh the connection settings: %s", exc)
+        return False
+    if creds.apply_connection(data):
+        log.info("connection settings updated from the platform")
+    return True
+
+
+def register_runtime(
+    api_url: str, api_key: str, body: dict[str, Any]
+) -> dict[str, Any]:
+    """``POST /api/v1/runtimes``; returns the response body.
+
+    Raises ``httpx.HTTPStatusError`` with the API's message on refusal.
+    """
+    with _client() as client:
+        resp = client.post(
+            f"{api_url.rstrip('/')}/api/v1/runtimes",
+            json=body,
+            headers={"Authorization": f"Bearer {api_key}"},
+        )
+    if resp.status_code >= 400:
+        raise httpx.HTTPStatusError(
+            _error_message(resp), request=resp.request, response=resp
+        )
+    data = resp.json()
+    if not isinstance(data, dict) or "credentials" not in data:
+        raise ValueError("unexpected response from the platform")
+    return data
+
+
+def _error_message(resp: httpx.Response) -> str:
+    try:
+        detail = resp.json().get("detail")
+    except ValueError:
+        detail = None
+    if isinstance(detail, dict):
+        detail = detail.get("message") or detail.get("reason")
+    return str(detail or f"HTTP {resp.status_code}")

levain_runtime/process.py

@@ -0,0 +1,172 @@
+"""Subprocess streaming shared by the direct and docker backends."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import os
+import signal
+import sys
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import Any
+
+_CHUNK = 65536
+_DROPPED_ENV = ("CLAUDECODE", "CLAUDE_CODE_ENTRYPOINT")
+
+
+def host_env() -> dict[str, str]:
+    """The host environment minus the Claude CLI's nested-session markers.
+
+    The platform runs each node as ``python /sandbox/node_runner.py``,
+    so ``python`` must resolve to an interpreter that has the SDK
+    installed. That is this process's own interpreter (``uv tool
+    install`` puts only the ``levain-runtime`` entrypoint on PATH),
+    so its bin directory goes first.
+    """
+    env = {k: v for k, v in os.environ.items() if k not in _DROPPED_ENV}
+    env["PYTHONUNBUFFERED"] = "1"
+    own_bin = str(Path(sys.executable).parent)
+    path = env.get("PATH", "")
+    if own_bin not in path.split(os.pathsep):
+        env["PATH"] = own_bin + (os.pathsep + path if path else "")
+    return env
+
+
+def _kill_group(proc: asyncio.subprocess.Process) -> None:
+    if proc.returncode is not None:
+        return
+    with contextlib.suppress(ProcessLookupError, PermissionError):
+        os.killpg(proc.pid, signal.SIGKILL)
+
+
+async def _pump(
+    stream: asyncio.StreamReader | None,
+    key: str,
+    queue: asyncio.Queue[dict[str, Any] | None],
+) -> None:
+    buf = b""
+    try:
+        while stream is not None:
+            chunk = await stream.read(_CHUNK)
+            if not chunk:
+                break
+            buf += chunk
+            while b"\n" in buf:
+                line, buf = buf.split(b"\n", 1)
+                await queue.put({key: line.decode("utf-8", "replace")})
+        if buf:
+            await queue.put({key: buf.decode("utf-8", "replace")})
+    finally:
+        await queue.put(None)
+
+
+async def _feed_stdin(proc: asyncio.subprocess.Process, data: bytes) -> None:
+    assert proc.stdin is not None
+    with contextlib.suppress(BrokenPipeError, ConnectionResetError):
+        proc.stdin.write(data)
+        await proc.stdin.drain()
+    with contextlib.suppress(Exception):
+        proc.stdin.close()
+
+
+class _Watchdog:
+    def __init__(
+        self, proc: asyncio.subprocess.Process, timeout: float
+    ) -> None:
+        self.fired = False
+        self._proc = proc
+        self._timeout = timeout
+
+    async def run(self) -> None:
+        await asyncio.sleep(self._timeout)
+        self.fired = True
+        _kill_group(self._proc)
+
+
+async def stream_process(
+    argv: list[str],
+    *,
+    env: dict[str, str],
+    stdin: bytes | None = None,
+    timeout: float = 0,
+    cwd: str | None = None,
+) -> AsyncIterator[dict[str, Any]]:
+    """Run ``argv`` and yield ``{"out"|"err": line}`` then ``{"exit", "timed_out"}``.
+
+    Lines are split by hand from fixed-size reads so a very long line
+    (an SDK message dump) never trips a reader limit. The process gets
+    its own session so a timeout can kill the whole group.
+    """
+    proc = await asyncio.create_subprocess_exec(
+        *argv,
+        env=env,
+        cwd=cwd,
+        stdin=asyncio.subprocess.PIPE
+        if stdin is not None
+        else asyncio.subprocess.DEVNULL,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+        start_new_session=True,
+    )
+    queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue()
+    tasks = [
+        asyncio.create_task(_pump(proc.stdout, "out", queue)),
+        asyncio.create_task(_pump(proc.stderr, "err", queue)),
+    ]
+    if stdin is not None:
+        tasks.append(asyncio.create_task(_feed_stdin(proc, stdin)))
+    watchdog = _Watchdog(proc, timeout) if timeout and timeout > 0 else None
+    if watchdog is not None:
+        tasks.append(asyncio.create_task(watchdog.run()))
+    try:
+        closed = 0
+        while closed < 2:
+            item = await queue.get()
+            if item is None:
+                closed += 1
+                continue
+            yield item
+        rc = await proc.wait()
+        timed_out = watchdog is not None and watchdog.fired
+        yield {"exit": -1 if timed_out else rc, "timed_out": timed_out}
+    finally:
+        # A consumer that stops reading (client gone) must not leave
+        # the node running.
+        _kill_group(proc)
+        for t in tasks:
+            t.cancel()
+        for t in tasks:
+            with contextlib.suppress(BaseException):
+                await t
+
+
+async def run_process(
+    argv: list[str],
+    *,
+    env: dict[str, str] | None = None,
+    stdin: bytes | None = None,
+    timeout: float = 0,
+    cwd: str | None = None,
+) -> tuple[int, bytes, bytes]:
+    """Run ``argv`` to completion; returns ``(exit, stdout, stderr)``."""
+    proc = await asyncio.create_subprocess_exec(
+        *argv,
+        env=env if env is not None else host_env(),
+        cwd=cwd,
+        stdin=asyncio.subprocess.PIPE
+        if stdin is not None
+        else asyncio.subprocess.DEVNULL,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+        start_new_session=True,
+    )
+    try:
+        out, err = await asyncio.wait_for(
+            proc.communicate(stdin), timeout=timeout or None
+        )
+    except TimeoutError:
+        _kill_group(proc)
+        await proc.wait()
+        return -1, b"", b"timed out"
+    return proc.returncode or 0, out, err

levain_runtime/protocol.py

@@ -0,0 +1,9 @@
+"""Wire constants shared with the platform side. Change in lockstep."""
+
+# Loopback listener the platform reaches through the reverse tunnel.
+RUNNER_HTTP_PORT = 17777
+# Loopback ports the tunnel forwards to platform services.
+FORWARD_SANDBOX_API_PORT = 18010
+FORWARD_PROXY_PORT = 14100
+
+CHISEL_VERSION = "1.11.3"

levain_runtime/sandbox/node_runner.py.vendored

@@ -0,0 +1,304 @@
+r"""Node execution entrypoint for sandbox containers.
+
+Invoked as::
+
+    python /sandbox/node_runner.py \
+        --function claude_agent \
+        --params /io/params.json \
+        --state /io/state.json \
+        --output /io/result.json
+
+Reads node parameters and graph state from JSON files,
+dispatches to the built-in handler or resolves a
+``module:function`` reference, and writes a result JSON file.
+The result is a state-delta dict consumed by the orchestrator.
+
+Lives next to the sandbox Dockerfile (not in the SDK) because
+it's container plumbing.  The SDK is for recipe authors, who
+never call this script.  Handler logic (``run_claude_node``,
+``run_shell_node``) lives in ``levain_sdk.handlers`` and is
+usable from other contexts.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import importlib
+import inspect
+import json
+import logging
+import sys
+from pathlib import Path
+from typing import Annotated, Any
+
+# Make ``run_logging`` (sibling module) importable when this file
+# is loaded as a script (``python /sandbox/node_runner.py`` — the
+# container path) and when it's imported under a package name by
+# tests (``from services.sandbox.node_runner import ...`` — pytest
+# rootdir).  The script path auto-prepends sys.path[0]; the test
+# path doesn't, so we do it ourselves.
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import run_logging
+import typer
+from levain_sdk.handlers import HANDLERS
+from levain_sdk.mcp_registry import resolve_mcp_servers
+from levain_sdk.messages import from_jsonable_state, to_jsonable
+
+log = logging.getLogger("knead.node_runner")
+
+app = typer.Typer(add_completion=False)
+
+
+def _error_envelope(exc: BaseException) -> dict[str, Any]:
+    """Serialise a node-failure exception into the result-dict envelope.
+
+    The harness reads ``_error_status`` to classify wake retries
+    deterministically — a 4xx-not-429 from any provider is
+    terminal and must drop the wake on first attempt instead of
+    burning four more sandbox provisions.  ``status_code`` is a
+    typed attribute on every ``litellm`` / ``openai`` /
+    ``anthropic`` exception class, so we walk the ``__cause__``
+    chain and grab the first integer we find — robust to wrappers
+    like the SDK's ``AgentError(str(litellm_exc))`` pattern.
+
+    ``_error_class`` carries the underlying exception class name
+    so callers can match deterministic-but-statusless types
+    (notably ``BudgetExceededError``, which is a plain
+    ``Exception`` with no ``status_code``).
+
+    ``_error`` keeps the original 500-char string for backward
+    compat with code paths that haven't been rewired to read the
+    structured fields yet.
+    """
+    return {
+        "_error": str(exc)[:500],
+        "_error_status": _extract_status_code(exc),
+        "_error_class": _root_cause_class(exc),
+    }
+
+
+def _extract_status_code(exc: BaseException) -> int | None:
+    """Walk ``__cause__`` / ``__context__`` for an HTTP status code.
+
+    Stops on the first integer ``status_code`` attribute — typed
+    on ``openai.APIStatusError`` and its anthropic / litellm
+    subclasses, so any wrapped provider exception surfaces the
+    real status.  Returns ``None`` when no link in the chain
+    carries one (e.g. plain ``RuntimeError`` from infra code, or
+    a stringified error coming back through claude_agent_sdk's
+    ``ResultMessage`` which throws away the typed exception).
+    """
+    seen: set[int] = set()
+    current: BaseException | None = exc
+    while current is not None and id(current) not in seen:
+        seen.add(id(current))
+        status = getattr(current, "status_code", None)
+        if isinstance(status, int):
+            return status
+        if current.__cause__ is not None:
+            current = current.__cause__
+        elif (
+            not current.__suppress_context__
+            and current.__context__ is not None
+        ):
+            current = current.__context__
+        else:
+            current = None
+    return None
+
+
+def _root_cause_class(exc: BaseException) -> str:
+    """Return the deepest ``__cause__`` class name, or the immediate one.
+
+    The harness uses this as a fallback signal for typed
+    failures that don't carry a status code — chiefly
+    ``BudgetExceededError``, which LiteLLM raises as a plain
+    ``Exception`` with no HTTP envelope but is unambiguously
+    terminal.
+    """
+    seen: set[int] = set()
+    current: BaseException = exc
+    while True:
+        seen.add(id(current))
+        nxt = current.__cause__
+        if nxt is None and not current.__suppress_context__:
+            nxt = current.__context__
+        if nxt is None or id(nxt) in seen:
+            break
+        current = nxt
+    return type(current).__name__
+
+
+def _resolve(function: str) -> Any:
+    """Resolve a node function name to a callable.
+
+    Built-in names (``claude_agent``, ``shell``) come from
+    ``HANDLERS``.  Recipe-defined ``module:function`` refs
+    are imported via ``importlib``.  ``ImportError`` and
+    ``AttributeError`` are re-raised so the caller can
+    distinguish "unknown built-in" from "broken import" vs
+    "missing function in module" — otherwise typos produce
+    a misleading "Unknown function" with no hint where the
+    actual failure is.
+    """
+    if function in HANDLERS:
+        return HANDLERS[function]
+    if ":" in function:
+        mod_path, _, fn_name = function.partition(":")
+        mod = importlib.import_module(mod_path)
+        try:
+            return getattr(mod, fn_name)
+        except AttributeError as exc:
+            raise AttributeError(
+                f"Module {mod_path!r} has no attribute {fn_name!r}"
+            ) from exc
+    return None
+
+
+def _register_platform_mcp_servers() -> None:
+    """Register platform MCP servers via platform_sdk's import side-effect."""
+    try:
+        import platform_sdk  # noqa: F401
+    except ModuleNotFoundError as exc:
+        # Only swallow platform_sdk-absent (customer sandbox); a broken
+        # install — missing transitive dep — must surface, not be masked.
+        if exc.name != "platform_sdk":
+            raise
+
+
+async def execute_node(
+    function: str,
+    params_path: Path,
+    state_path: Path,
+    output_path: Path,
+) -> None:
+    """Read inputs, dispatch to handler, write result."""
+    try:
+        params = json.loads(params_path.read_text())
+        # Reducers don't run in the sandbox; rehydrate Message dicts.
+        state = from_jsonable_state(json.loads(state_path.read_text()))
+    except (OSError, json.JSONDecodeError) as exc:
+        result: dict[str, Any] = {
+            "_error": f"Failed to read IO files: {exc}",
+        }
+        output_path.write_text(json.dumps(result))
+        return
+
+    try:
+        handler = _resolve(function)
+    except (ImportError, AttributeError) as exc:
+        result = {
+            "_error": f"Cannot resolve {function!r}: {exc}",
+        }
+        output_path.write_text(json.dumps(result))
+        return
+
+    # Resolve ``@mcp_servers`` names → SDK configs once, in the
+    # sandbox process where ``LEVAIN_INTEGRATIONS_B64`` is in scope
+    # and stdio platform servers (``knead-review``) can be spawned
+    # in-process.  Resolution feeds both call paths: built-in
+    # ``claude_agent`` reads ``params["mcp_servers"]`` directly,
+    # recipe nodes read ``config["configurable"]["_resolved_mcp_servers"]``
+    # via ``levain_sdk.claude.resolved_mcp_servers``.
+    requested_mcp = params.get("mcp_servers")
+    resolved_mcp: dict[str, Any] = {}
+    if isinstance(requested_mcp, list) and requested_mcp:
+        # Catch everything (not just ValueError): a platform MCP server's
+        # init() can raise other types — e.g. knead-review touching git on
+        # a not-yet-cloned /workspace. A clean error envelope keeps it from
+        # escaping node_runner as a bare "Pod exec exited with 1".
+        try:
+            _register_platform_mcp_servers()
+            resolved_mcp = resolve_mcp_servers(
+                [n for n in requested_mcp if isinstance(n, str)],
+                repo_root=str(state.get("cwd") or "/workspace"),
+            )
+        except Exception as exc:
+            output_path.write_text(json.dumps(_error_envelope(exc)))
+            return
+        params["mcp_servers"] = resolved_mcp
+    elif isinstance(requested_mcp, dict):
+        resolved_mcp = requested_mcp
+
+    if handler is None:
+        result = {
+            "_error": f"Unknown function: {function}",
+        }
+    elif function in HANDLERS:
+        # Built-in handler: (params, state) signature.
+        try:
+            result = await handler(params, state)
+        except Exception as exc:
+            log.error(
+                "Node execution failed: %s",
+                exc,
+                exc_info=True,
+            )
+            result = _error_envelope(exc)
+    else:
+        # Recipe function: (state, config) LangGraph signature.
+        # ``params["_configurable"]`` (if present) is the
+        # serialized subset of ``RunnableConfig["configurable"]``
+        # that the orchestrator forwarded into the container —
+        # recipe nodes read values like ``task`` or
+        # ``message_hint`` from it via
+        # ``config["configurable"]``.
+        configurable = dict(params.get("_configurable") or {})
+        if resolved_mcp:
+            configurable["_resolved_mcp_servers"] = resolved_mcp
+        config: dict[str, Any] = {"configurable": configurable}
+        try:
+            result = handler(state, config)
+            if inspect.isawaitable(result):
+                result = await result
+        except Exception as exc:
+            log.error(
+                "Node execution failed: %s",
+                exc,
+                exc_info=True,
+            )
+            result = _error_envelope(exc)
+
+    if not isinstance(result, dict):
+        result = {
+            "_error": (
+                f"Node returned {type(result).__name__}, expected dict"
+            ),
+        }
+
+    # Recipes return ``Message`` dataclasses (and other typed
+    # values like ``datetime``) in state-delta dicts; the default
+    # JSON encoder can't serialise those, so convert via
+    # ``to_jsonable`` before write. The harness side coerces dict
+    # shapes back into ``Message`` via ``add_messages``.
+    try:
+        payload = json.dumps(to_jsonable(result))
+    except (TypeError, ValueError) as exc:
+        payload = json.dumps(_error_envelope(exc))
+    output_path.write_text(payload)
+
+
+@app.command()
+def main(
+    function: Annotated[str, typer.Option("--function")],
+    params: Annotated[Path, typer.Option("--params")],
+    state: Annotated[Path, typer.Option("--state")],
+    output: Annotated[Path, typer.Option("--output")],
+) -> None:
+    """Execute a single graph node inside a sandbox."""
+    # JSON Lines on stdout so the harness consolidator can ingest
+    # every record this subprocess emits.  Identity (run / session
+    # / agent ids) is read from ``LEVAIN_*`` env vars set by the
+    # harness's ``_env_prefix`` on the dispatch.  DEBUG so the
+    # SDK's raw ``sdk_message`` dumps (``log.debug(...)`` in
+    # ``levain_sdk.claude``) reach the consolidator — completeness
+    # over economy here since one subprocess handles one node.
+    run_logging.configure(level=logging.DEBUG)
+    run_logging.emit_node_boot("knead.node_runner", function)
+    asyncio.run(execute_node(function, params, state, output))
+
+
+if __name__ == "__main__":
+    # Allow running as a plain script (not just typer app)
+    app()