levain-runtime 0.1.2__py3-none-any.whl

levain_runtime/sandbox/run_logging.py.vendored

@@ -0,0 +1,153 @@
+"""Sandbox-side logging: JSON Lines to stdout for harness ingestion.
+
+The harness streams this subprocess's stdout into the per-run log
+file consolidator.  Records are JSON Lines (one object per line)
+matching the harness's ``_JsonFormatter`` field shape so the
+consolidator can deserialize without sniffing.
+
+Run identity is provided as ``LEVAIN_*`` env vars set by
+``_env_prefix`` on the dispatch — see
+:mod:`levain.agent.harness.run_context`.  These values are baked
+into every record so the consolidator can correlate without
+re-parsing the env on each line.
+
+This module lives next to ``node_runner.py`` because it ships
+inside the sandbox image and must not import from
+``levain.agent`` (which is not on the sandbox's PYTHONPATH).
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import sys
+from typing import Any
+
+# Mirror of ``levain.agent.harness.run_context.ENV_VAR_NAMES``.
+# Duplicated here because the sandbox image doesn't ship the
+# harness package; any change must touch both copies.
+_ENV_TO_FIELD: dict[str, str] = {
+    "LEVAIN_RUN_ID": "run_id",
+    "LEVAIN_SESSION_ID": "session_id",
+    "LEVAIN_WORKSPACE_ID": "workspace_id",
+    "LEVAIN_AGENT_ID": "agent_id",
+    "LEVAIN_AGENT_VERSION_ID": "agent_version_id",
+    "LEVAIN_TURN_ID": "turn_id",
+    "LEVAIN_NODE": "node",
+}
+
+
+def _snapshot_identity() -> dict[str, str]:
+    """Read run-scoped identifiers from env at configure time.
+
+    A sandbox subprocess handles exactly one node dispatch, so
+    these fields are constants for the process lifetime.  Empty
+    values are dropped to keep records terse when the harness
+    didn't supply a particular id (e.g. platform runs without an
+    agent).
+    """
+    out: dict[str, str] = {}
+    for env_name, field in _ENV_TO_FIELD.items():
+        value = os.environ.get(env_name)
+        if value:
+            out[field] = value
+    return out
+
+
+class _JsonFormatter(logging.Formatter):
+    """One JSON object per line — same shape as the harness side.
+
+    Bakes the env-supplied run identity into every record so the
+    consolidator can route lines into the right per-run file
+    without inspecting subprocess env separately.
+    """
+
+    def __init__(self, identity: dict[str, str]) -> None:
+        super().__init__()
+        self._identity = identity
+
+    def format(self, record: logging.LogRecord) -> str:
+        entry: dict[str, object] = {
+            "ts": record.created,
+            "level": record.levelname,
+            "logger": record.name,
+            "msg": record.getMessage(),
+        }
+        entry.update(self._identity)
+        for key in ("label", "cost_usd", "tokens", "duration_ms", "event"):
+            val = getattr(record, key, None)
+            if val is not None:
+                entry[key] = val
+        if record.exc_info and record.exc_info[1] is not None:
+            entry["exc"] = self.formatException(record.exc_info)
+        return json.dumps(entry, default=str)
+
+
+_NOISY_DEBUG_LOGGERS: tuple[str, ...] = (
+    "httpx",
+    "httpcore",
+    "urllib3",
+    "anthropic",
+    "openai",
+    "asyncio",
+    "langchain",
+    "langgraph",
+    "litellm",
+    "boto3",
+    "botocore",
+    "s3transfer",
+    # Content-extraction libs used by the crawler — per-page DEBUG
+    # (element discarding, readability scoring) drowns the run log.
+    "trafilatura",
+    "htmldate",
+    "tzlocal",
+)
+
+
+def configure(level: int = logging.INFO) -> dict[str, str]:
+    """Install the JSON formatter on the root logger.
+
+    Replaces any existing handlers so ``node_runner.py``'s prior
+    plain-text ``basicConfig`` doesn't double-print.  Returns the
+    identity snapshot so the caller can include it in any
+    bootstrap log line.
+
+    When ``level`` is ``DEBUG`` (the sandbox default), known-noisy
+    third-party loggers are clamped to ``INFO`` so transport-layer
+    detail from the LLM SDKs doesn't drown out the curated SDK /
+    recipe records we actually want in the run log.
+    """
+    identity = _snapshot_identity()
+    root = logging.getLogger()
+    root.setLevel(level)
+    # Drop any handlers a prior ``logging.basicConfig`` left
+    # behind so we don't emit both plaintext and JSON for the
+    # same record.
+    for h in root.handlers[:]:
+        root.removeHandler(h)
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(level)
+    handler.setFormatter(_JsonFormatter(identity))
+    root.addHandler(handler)
+    if level <= logging.DEBUG:
+        for name in _NOISY_DEBUG_LOGGERS:
+            logging.getLogger(name).setLevel(logging.INFO)
+    return identity
+
+
+def emit_node_boot(
+    logger_name: str,
+    function: str,
+    extra: dict[str, Any] | None = None,
+) -> None:
+    """Emit a single ``node_boot`` record summarizing the dispatch.
+
+    Lets the consolidator anchor the subprocess's activity even
+    when the recipe handler logs nothing useful on its own.
+    """
+    log = logging.getLogger(logger_name)
+    event: dict[str, Any] = {"kind": "node_boot", "function": function}
+    if extra:
+        event.update(extra)
+    log.info("NODE  %s starting", function, extra={"event": event})

levain_runtime/server.py

@@ -0,0 +1,226 @@
+"""Loopback HTTP server the platform drives through the reverse tunnel."""
+
+from __future__ import annotations
+
+import json
+import tempfile
+import uuid
+from collections.abc import AsyncIterator
+from typing import Any
+
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response, StreamingResponse
+from starlette.routing import Mount, Route
+
+from .config import RuntimeConfig
+from .sessions import (
+    MissingTarget,
+    RuntimeBusy,
+    SessionManager,
+    UnknownSession,
+)
+
+_SPOOL_BYTES = 8 * 1024 * 1024
+
+
+def _manager(request: Request) -> SessionManager:
+    manager: SessionManager = request.app.state.manager
+    return manager
+
+
+def _sid(request: Request) -> str:
+    raw = request.path_params["sid"]
+    try:
+        return str(uuid.UUID(raw))
+    except ValueError as exc:
+        raise ValueError("sid must be a UUID") from exc
+
+
+async def _json_body(request: Request) -> dict[str, Any]:
+    try:
+        body = await request.json()
+    except json.JSONDecodeError as exc:
+        raise ValueError("body must be JSON") from exc
+    if not isinstance(body, dict):
+        raise ValueError("body must be a JSON object")
+    return body
+
+
+def _env_of(body: dict[str, Any]) -> dict[str, str]:
+    env = body.get("env") or {}
+    if not isinstance(env, dict) or not all(
+        isinstance(k, str) and isinstance(v, str) for k, v in env.items()
+    ):
+        raise ValueError("env must map strings to strings")
+    return env
+
+
+async def health(request: Request) -> Response:
+    m = _manager(request)
+    return JSONResponse(
+        {"ok": True, "mode": m.mode, "active_sessions": m.active_sessions()}
+    )
+
+
+async def session(request: Request) -> Response:
+    m = _manager(request)
+    sid = _sid(request)
+    if request.method == "GET":
+        return JSONResponse({"status": await m.status(sid)})
+    if request.method == "DELETE":
+        keep = request.query_params.get("keep_workspace", "0") == "1"
+        await m.delete(sid, keep_workspace=keep)
+        return Response(status_code=204)
+    body = await _json_body(request)
+    name, created = await m.provision(
+        sid,
+        workspace_id=str(body.get("workspace_id") or ""),
+        agent_id=body.get("agent_id") or None,
+        workspace_scope=body.get("workspace_scope") or None,
+        workspace_user=body.get("workspace_user") or None,
+        env=_env_of(body),
+    )
+    return JSONResponse({"name": name, "created": created})
+
+
+async def pause(request: Request) -> Response:
+    await _manager(request).pause(_sid(request))
+    return Response(status_code=204)
+
+
+async def resume(request: Request) -> Response:
+    await _manager(request).resume(_sid(request))
+    return Response(status_code=204)
+
+
+async def exec_(request: Request) -> Response:
+    m = _manager(request)
+    sid = _sid(request)
+    body = await _json_body(request)
+    cmd = body.get("cmd")
+    if not isinstance(cmd, str) or not cmd:
+        raise ValueError("cmd is required")
+    stdin = body.get("stdin")
+    if stdin is not None and not isinstance(stdin, str):
+        raise ValueError("stdin must be a string")
+    timeout = float(body.get("timeout") or 0)
+    stream = m.exec(sid, cmd, _env_of(body), stdin=stdin, timeout=timeout)
+
+    async def lines() -> AsyncIterator[bytes]:
+        async for item in stream:
+            yield (json.dumps(item) + "\n").encode()
+
+    return StreamingResponse(lines(), media_type="application/x-ndjson")
+
+
+async def capture(request: Request) -> Response:
+    m = _manager(request)
+    sid = _sid(request)
+    body = await _json_body(request)
+    cmd = body.get("cmd")
+    if not isinstance(cmd, str) or not cmd:
+        raise ValueError("cmd is required")
+    out, rc = await m.capture(
+        sid, cmd, _env_of(body), timeout=float(body.get("timeout") or 0)
+    )
+    return JSONResponse({"stdout": out, "exit": rc})
+
+
+def _path_param(request: Request) -> str:
+    path = request.query_params.get("path")
+    if not path:
+        raise ValueError("path is required")
+    return path
+
+
+async def files(request: Request) -> Response:
+    m = _manager(request)
+    sid = _sid(request)
+    path = _path_param(request)
+    if request.method == "GET":
+        data = await m.get_file(sid, path)
+        if data is None:
+            return JSONResponse({"detail": "not found"}, status_code=404)
+        return Response(data, media_type="application/octet-stream")
+    await m.put_file(sid, path, await request.body())
+    return Response(status_code=204)
+
+
+async def tree(request: Request) -> Response:
+    m = _manager(request)
+    sid = _sid(request)
+    path = _path_param(request)
+    with tempfile.SpooledTemporaryFile(max_size=_SPOOL_BYTES) as spooled:
+        async for chunk in request.stream():
+            spooled.write(chunk)
+        spooled.seek(0)
+        await m.put_tree(sid, path, spooled)
+    return Response(status_code=204)
+
+
+async def agent_workspace(request: Request) -> Response:
+    m = _manager(request)
+    user = request.query_params.get("user") or None
+    await m.reset_agent_workspace(request.path_params["agent_id"], user=user)
+    return Response(status_code=204)
+
+
+async def link(request: Request) -> Response:
+    m = _manager(request)
+    sid = _sid(request)
+    body = await _json_body(request)
+    path, target = body.get("path"), body.get("target")
+    if not isinstance(path, str) or not isinstance(target, str):
+        raise ValueError("path and target are required")
+    await m.link(sid, path, target)
+    return Response(status_code=204)
+
+
+def _error(status: int) -> Any:
+    async def handler(request: Request, exc: Exception) -> Response:
+        return JSONResponse({"detail": str(exc)}, status_code=status)
+
+    return handler
+
+
+def create_app(config: RuntimeConfig) -> Starlette:
+    app = Starlette(
+        routes=[
+            Mount(
+                "/v1",
+                routes=[
+                    Route("/health", health, methods=["GET"]),
+                    Route(
+                        "/sessions/{sid}",
+                        session,
+                        methods=["POST", "GET", "DELETE"],
+                    ),
+                    Route("/sessions/{sid}/pause", pause, methods=["POST"]),
+                    Route("/sessions/{sid}/resume", resume, methods=["POST"]),
+                    Route("/sessions/{sid}/exec", exec_, methods=["POST"]),
+                    Route(
+                        "/sessions/{sid}/capture", capture, methods=["POST"]
+                    ),
+                    Route(
+                        "/sessions/{sid}/files", files, methods=["PUT", "GET"]
+                    ),
+                    Route("/sessions/{sid}/tree", tree, methods=["PUT"]),
+                    Route("/sessions/{sid}/link", link, methods=["PUT"]),
+                    Route(
+                        "/agents/{agent_id}/workspace",
+                        agent_workspace,
+                        methods=["DELETE"],
+                    ),
+                ],
+            )
+        ],
+        exception_handlers={
+            ValueError: _error(400),
+            UnknownSession: _error(404),
+            MissingTarget: _error(404),
+            RuntimeBusy: _error(409),
+        },
+    )
+    app.state.manager = SessionManager(config)
+    return app