levain-runtime 0.1.2__py3-none-any.whl

levain_runtime/sessions.py

@@ -0,0 +1,721 @@
+"""Session registry and the two execution backends (direct host, docker)."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+import shutil
+import tarfile
+import uuid
+from collections.abc import AsyncIterator, Awaitable, Callable
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import IO, Any, Literal
+
+from .config import RuntimeConfig
+from .process import host_env, run_process, stream_process
+
+log = logging.getLogger("levain_runtime")
+
+State = Literal["active", "paused"]
+_ID_RE = re.compile(r"^[A-Za-z0-9-]{1,64}$")
+_SESSION_FILE = "session.json"
+# Trees under these roots are host bind mounts in docker mode, so they
+# are extracted on the host for both backends.
+_HOST_TREE_ROOTS = ("/recipes", "/sdk")
+_CLAUDE_READY_SENTINEL = "/tmp/.levain-runtime-ready"
+
+Runner = Callable[..., Awaitable[tuple[int, bytes, bytes]]]
+
+
+class RuntimeBusy(Exception):
+    pass
+
+
+class UnknownSession(Exception):
+    pass
+
+
+class MissingTarget(Exception):
+    """``link`` asked for a target tree that is not on this runtime."""
+
+
+@dataclass
+class SessionMeta:
+    sid: str
+    workspace_id: str
+    agent_id: str | None
+    workspace_scope: str | None
+    workspace_user: str | None
+    env: dict[str, str]
+    workspace_dir: str
+    state: State = "paused"
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, raw: dict[str, Any]) -> SessionMeta:
+        return cls(**raw)
+
+
+def validate_abs_path(path: str) -> str:
+    if not path.startswith("/") or ".." in Path(path).parts:
+        raise ValueError("path must be absolute and contain no '..'")
+    return path
+
+
+def _check_id(value: str, what: str) -> str:
+    if not _ID_RE.match(value):
+        raise ValueError(f"invalid {what}")
+    return value
+
+
+def extract_tar(src: IO[bytes], dest: Path) -> None:
+    """Unpack ``src`` at ``dest`` atomically; keep ``dest`` if it exists.
+
+    Trees are content-addressed, so an existing ``dest`` already holds
+    these bytes. Staging next to it means a failed extraction never
+    leaves a half tree that a later probe would mistake for a full one.
+    """
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.parent / f".{dest.name}.{uuid.uuid4().hex[:8]}"
+    tmp.mkdir()
+    try:
+        with tarfile.open(fileobj=src, mode="r:") as tar:
+            # The data filter rejects absolute members and ``..`` escapes.
+            tar.extractall(tmp, filter="data")
+        try:
+            os.rename(tmp, dest)
+        except OSError:
+            if not dest.is_dir():
+                raise
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)
+
+
+def flip_link(link: Path, target: str, *, exists_at: Path) -> None:
+    """Point ``link`` at ``target`` in one ``rename(2)``.
+
+    ``target`` is the string stored in the symlink (the path as nodes
+    see it); ``exists_at`` is where that tree lives from here, which
+    differs in docker mode where ``/recipes`` is a bind mount.
+    """
+    if not exists_at.is_dir():
+        raise MissingTarget(f"{target} is not on this runtime")
+    link.parent.mkdir(parents=True, exist_ok=True)
+    tmp = link.parent / f".{link.name}.{uuid.uuid4().hex[:8]}"
+    os.symlink(target, tmp)
+    try:
+        os.replace(tmp, link)
+    except OSError:
+        tmp.unlink(missing_ok=True)
+        raise
+
+
+class SessionStore:
+    """On-disk session metadata under ``<root>/sessions/<sid>/``."""
+
+    def __init__(self, root: Path) -> None:
+        self.root = root
+        self._cache: dict[str, SessionMeta] = {}
+
+    def dir(self, sid: str) -> Path:
+        return self.root / "sessions" / sid
+
+    def get(self, sid: str) -> SessionMeta | None:
+        if sid in self._cache:
+            return self._cache[sid]
+        path = self.dir(sid) / _SESSION_FILE
+        if not path.is_file():
+            return None
+        meta = SessionMeta.from_dict(json.loads(path.read_text()))
+        self._cache[sid] = meta
+        return meta
+
+    def save(self, meta: SessionMeta) -> None:
+        self.dir(meta.sid).mkdir(parents=True, exist_ok=True)
+        path = self.dir(meta.sid) / _SESSION_FILE
+        tmp = path.with_suffix(".tmp")
+        tmp.write_text(json.dumps(meta.to_dict()))
+        os.replace(tmp, path)
+        self._cache[meta.sid] = meta
+
+    def delete(self, sid: str, *, keep_workspace: bool) -> None:
+        self._cache.pop(sid, None)
+        sdir = self.dir(sid)
+        if not sdir.exists():
+            return
+        if keep_workspace:
+            (sdir / _SESSION_FILE).unlink(missing_ok=True)
+            return
+        shutil.rmtree(sdir, ignore_errors=True)
+
+    def all(self) -> list[SessionMeta]:
+        base = self.root / "sessions"
+        if not base.is_dir():
+            return []
+        out: list[SessionMeta] = []
+        for child in base.iterdir():
+            if child.name.startswith("."):
+                continue
+            meta = self.get(child.name)
+            if meta is not None:
+                out.append(meta)
+        return out
+
+    def workspace_dir(
+        self,
+        sid: str,
+        *,
+        agent_id: str | None,
+        scope: str | None,
+        user: str | None,
+    ) -> Path:
+        if scope == "agent" and agent_id:
+            return (
+                self.root
+                / "agents"
+                / _check_id(agent_id, "agent_id")
+                / "workspace"
+            )
+        if scope == "user" and agent_id and user:
+            return (
+                self.root
+                / "agents"
+                / _check_id(agent_id, "agent_id")
+                / "users"
+                / _check_id(user, "workspace_user")
+                / "workspace"
+            )
+        return self.dir(sid) / "workspace"
+
+
+class DirectBackend:
+    """Nodes run as host processes; one active session owns ``/workspace``."""
+
+    def __init__(self, root: Path) -> None:
+        self.root = root
+        self.active: str | None = None
+
+    @property
+    def idle_workspace(self) -> Path:
+        return self.root / "sessions" / ".idle" / "workspace"
+
+    def _point_workspace(self, target: Path) -> None:
+        target.mkdir(parents=True, exist_ok=True)
+        link = self.root / "workspace"
+        tmp = self.root / f".workspace.{uuid.uuid4().hex[:8]}"
+        os.symlink(target, tmp)
+        os.replace(tmp, link)
+
+    async def activate(self, meta: SessionMeta) -> None:
+        if self.active not in (None, meta.sid):
+            raise RuntimeBusy("runtime busy")
+        self._point_workspace(Path(meta.workspace_dir))
+        self.active = meta.sid
+
+    async def release(self, meta: SessionMeta) -> None:
+        if self.active == meta.sid:
+            self.active = None
+            self._point_workspace(self.idle_workspace)
+
+    async def destroy(self, meta: SessionMeta) -> None:
+        await self.release(meta)
+
+    def status(self, meta: SessionMeta) -> State:
+        return "active" if self.active == meta.sid else "paused"
+
+    def _env(self, meta: SessionMeta, extra: dict[str, str]) -> dict[str, str]:
+        return {**host_env(), **meta.env, **extra}
+
+    def exec(
+        self,
+        meta: SessionMeta,
+        cmd: str,
+        env: dict[str, str],
+        *,
+        stdin: str | None,
+        timeout: float,
+    ) -> AsyncIterator[dict[str, Any]]:
+        return stream_process(
+            ["/bin/sh", "-c", cmd],
+            env=self._env(meta, env),
+            stdin=stdin.encode() if stdin is not None else None,
+            timeout=timeout,
+            cwd=meta.workspace_dir
+            if Path(meta.workspace_dir).is_dir()
+            else None,
+        )
+
+    async def capture(
+        self,
+        meta: SessionMeta,
+        cmd: str,
+        env: dict[str, str],
+        *,
+        timeout: float,
+    ) -> tuple[str, int]:
+        rc, out, _ = await run_process(
+            ["/bin/sh", "-c", cmd],
+            env=self._env(meta, env),
+            timeout=timeout,
+            cwd=meta.workspace_dir
+            if Path(meta.workspace_dir).is_dir()
+            else None,
+        )
+        return out.decode("utf-8", "replace"), rc
+
+    async def put_file(
+        self, meta: SessionMeta, path: str, data: bytes
+    ) -> None:
+        p = Path(path)
+        p.parent.mkdir(parents=True, exist_ok=True)
+        p.write_bytes(data)
+
+    async def get_file(self, meta: SessionMeta, path: str) -> bytes | None:
+        p = Path(path)
+        return p.read_bytes() if p.is_file() else None
+
+    async def put_tree(
+        self, meta: SessionMeta, path: str, src: IO[bytes]
+    ) -> None:
+        extract_tar(src, Path(path))
+
+    async def link(self, meta: SessionMeta, path: str, target: str) -> None:
+        flip_link(Path(path), target, exists_at=Path(target))
+
+
+class DockerBackend:
+    """One container per session; the runtime home is bind-mounted in."""
+
+    def __init__(
+        self,
+        root: Path,
+        image: str,
+        *,
+        share_claude_login: bool = True,
+        runner: Runner = run_process,
+        streamer: Callable[
+            ..., AsyncIterator[dict[str, Any]]
+        ] = stream_process,
+    ) -> None:
+        self.root = root
+        self.image = image
+        self.share_claude_login = share_claude_login
+        self._run = runner
+        self._stream = streamer
+
+    @staticmethod
+    def container_name(sid: str) -> str:
+        return "levain-rt-" + sid.replace("-", "")[:12]
+
+    def _claude_binds(self) -> list[str]:
+        if not self.share_claude_login:
+            return []
+        binds: list[str] = []
+        cfg = Path(
+            os.environ.get("CLAUDE_CONFIG_DIR") or "~/.claude"
+        ).expanduser()
+        if cfg.is_dir():
+            binds += ["-v", f"{cfg}:/mnt/host-claude:ro"]
+        cfg_json = Path("~/.claude.json").expanduser()
+        if cfg_json.is_file():
+            binds += ["-v", f"{cfg_json}:/mnt/host-claude.json:ro"]
+        return binds
+
+    @staticmethod
+    def _startup_script() -> str:
+        # Copy the host login into a writable $HOME; the CLI refreshes
+        # tokens and writes session state there. Skip the bulky,
+        # host-specific directories.
+        return "\n".join(
+            [
+                "if [ -d /mnt/host-claude ]; then",
+                '  mkdir -p "$HOME/.claude"',
+                "  for entry in /mnt/host-claude/* /mnt/host-claude/.[!.]*; do",
+                '    [ -e "$entry" ] || continue',
+                '    case "$(basename "$entry")" in '
+                "sessions|projects|cache|debug|backups|statsig|"
+                "shell-snapshots|todos|history) continue ;; esac",
+                '    cp -r "$entry" "$HOME/.claude/" || echo "warning: copy failed" >&2',
+                "  done",
+                "fi",
+                '[ -f /mnt/host-claude.json ] && cp /mnt/host-claude.json "$HOME/.claude.json"',
+                f"touch {_CLAUDE_READY_SENTINEL}",
+                "exec sleep infinity",
+            ]
+        )
+
+    def create_argv(self, meta: SessionMeta) -> list[str]:
+        argv = [
+            "docker",
+            "run",
+            "-d",
+            "--name",
+            self.container_name(meta.sid),
+            "-v",
+            f"{meta.workspace_dir}:/workspace",
+            "-v",
+            f"{self.root / 'recipes'}:/recipes",
+            "-v",
+            f"{self.root / 'sdk'}:/sdk",
+            "-v",
+            f"{self.root / 'sandbox'}:/sandbox:ro",
+            "-e",
+            "PYTHONUNBUFFERED=1",
+        ]
+        for k, v in meta.env.items():
+            argv += ["-e", f"{k}={v}"]
+        argv += self._claude_binds()
+        argv += [self.image, "/bin/sh", "-c", self._startup_script()]
+        return argv
+
+    async def _inspect(self, name: str) -> str | None:
+        rc, out, _ = await self._run(
+            ["docker", "inspect", "-f", "{{.State.Running}}", name],
+            timeout=30,
+        )
+        if rc != 0:
+            return None
+        return "running" if out.strip() == b"true" else "stopped"
+
+    async def _wait_ready(self, name: str) -> None:
+        rc, _, _ = await self._run(
+            [
+                "docker",
+                "exec",
+                name,
+                "/bin/sh",
+                "-c",
+                f"i=0; until [ -f {_CLAUDE_READY_SENTINEL} ]; do"
+                " i=$((i+1)); [ $i -ge 300 ] && exit 1; sleep 0.1; done",
+            ],
+            timeout=45,
+        )
+        if rc != 0:
+            raise RuntimeError(f"container {name} did not become ready")
+
+    async def activate(self, meta: SessionMeta) -> None:
+        Path(meta.workspace_dir).mkdir(parents=True, exist_ok=True)
+        name = self.container_name(meta.sid)
+        state = await self._inspect(name)
+        if state == "running":
+            return
+        if state == "stopped":
+            rc, _, err = await self._run(["docker", "start", name], timeout=60)
+            if rc == 0:
+                return
+            log.warning("docker start %s failed: %s; recreating", name, err)
+            await self._run(["docker", "rm", "-f", name], timeout=60)
+        rc, _, err = await self._run(self.create_argv(meta), timeout=300)
+        if rc != 0:
+            raise RuntimeError(
+                f"docker run failed: {err.decode(errors='replace')}"
+            )
+        await self._wait_ready(name)
+
+    async def release(self, meta: SessionMeta) -> None:
+        await self._run(
+            ["docker", "stop", self.container_name(meta.sid)], timeout=60
+        )
+
+    async def destroy(self, meta: SessionMeta) -> None:
+        await self._run(
+            ["docker", "rm", "-f", self.container_name(meta.sid)], timeout=60
+        )
+
+    async def status_async(self, meta: SessionMeta) -> State:
+        state = await self._inspect(self.container_name(meta.sid))
+        return "active" if state == "running" else "paused"
+
+    def exec(
+        self,
+        meta: SessionMeta,
+        cmd: str,
+        env: dict[str, str],
+        *,
+        stdin: str | None,
+        timeout: float,
+    ) -> AsyncIterator[dict[str, Any]]:
+        argv = ["docker", "exec", "-i"]
+        for k, v in env.items():
+            argv += ["-e", f"{k}={v}"]
+        argv.append(self.container_name(meta.sid))
+        if timeout and timeout > 0:
+            # Kill inside the container too; the host watchdog only
+            # reaches the docker client process.
+            argv += ["timeout", "-s", "KILL", str(int(timeout) + 1)]
+        argv += ["/bin/sh", "-c", cmd]
+        return self._stream(
+            argv,
+            env=host_env(),
+            stdin=stdin.encode() if stdin is not None else None,
+            timeout=timeout,
+        )
+
+    async def capture(
+        self,
+        meta: SessionMeta,
+        cmd: str,
+        env: dict[str, str],
+        *,
+        timeout: float,
+    ) -> tuple[str, int]:
+        argv = ["docker", "exec"]
+        for k, v in env.items():
+            argv += ["-e", f"{k}={v}"]
+        argv += [self.container_name(meta.sid), "/bin/sh", "-c", cmd]
+        rc, out, _ = await self._run(argv, timeout=timeout or 600)
+        return out.decode("utf-8", "replace"), rc
+
+    def put_file_argv(self, meta: SessionMeta, path: str) -> list[str]:
+        return [
+            "docker",
+            "exec",
+            "-i",
+            self.container_name(meta.sid),
+            "/bin/sh",
+            "-c",
+            'mkdir -p "$(dirname -- "$1")" && cat > "$1"',
+            "sh",
+            path,
+        ]
+
+    async def put_file(
+        self, meta: SessionMeta, path: str, data: bytes
+    ) -> None:
+        rc, _, err = await self._run(
+            self.put_file_argv(meta, path), stdin=data, timeout=120
+        )
+        if rc != 0:
+            raise RuntimeError(
+                f"write {path} failed: {err.decode(errors='replace')}"
+            )
+
+    async def get_file(self, meta: SessionMeta, path: str) -> bytes | None:
+        rc, out, _ = await self._run(
+            ["docker", "exec", self.container_name(meta.sid), "cat", path],
+            timeout=120,
+        )
+        return out if rc == 0 else None
+
+    def host_path(self, path: str) -> Path | None:
+        """Where a container path lives on the host, for bind-mounted roots."""
+        for prefix in _HOST_TREE_ROOTS:
+            if path == prefix or path.startswith(prefix + "/"):
+                return self.root / path.lstrip("/")
+        return None
+
+    async def put_tree(
+        self, meta: SessionMeta, path: str, src: IO[bytes]
+    ) -> None:
+        host = self.host_path(path)
+        if host is not None:
+            extract_tar(src, host)
+            return
+        rc, _, err = await self._run(
+            [
+                "docker",
+                "exec",
+                "-i",
+                self.container_name(meta.sid),
+                "/bin/sh",
+                "-c",
+                'mkdir -p "$1" && tar -x -C "$1"',
+                "sh",
+                path,
+            ],
+            stdin=src.read(),
+            timeout=600,
+        )
+        if rc != 0:
+            raise RuntimeError(
+                f"extract into {path} failed: {err.decode(errors='replace')}"
+            )
+
+    async def link(self, meta: SessionMeta, path: str, target: str) -> None:
+        # The symlink is created on the host but read inside the
+        # container, so it stores the container path and the existence
+        # check uses the host one.
+        host_link = self.host_path(path)
+        host_target = self.host_path(target)
+        if host_link is None or host_target is None:
+            raise ValueError(
+                f"link is limited to {', '.join(_HOST_TREE_ROOTS)} in docker mode"
+            )
+        flip_link(host_link, target, exists_at=host_target)
+
+
+class SessionManager:
+    """Routes session operations to the configured backend."""
+
+    def __init__(self, config: RuntimeConfig) -> None:
+        self.config = config
+        self.store = SessionStore(config.root)
+        self.backend: DirectBackend | DockerBackend
+        if config.mode == "docker":
+            if not config.image:
+                raise ValueError("docker mode requires an image")
+            self.backend = DockerBackend(
+                config.root,
+                config.image,
+                share_claude_login=config.share_claude_login,
+            )
+        else:
+            self.backend = DirectBackend(config.root)
+            self._restore_direct_slot()
+
+    def _restore_direct_slot(self) -> None:
+        assert isinstance(self.backend, DirectBackend)
+        active = [m for m in self.store.all() if m.state == "active"]
+        # Only one session may own /workspace; after a restart, any
+        # recorded holder is stale until the platform resumes it.
+        for meta in active:
+            meta.state = "paused"
+            self.store.save(meta)
+
+    @property
+    def mode(self) -> str:
+        return self.config.mode
+
+    def require(self, sid: str) -> SessionMeta:
+        meta = self.store.get(sid)
+        if meta is None:
+            raise UnknownSession(sid)
+        return meta
+
+    def active_sessions(self) -> list[str]:
+        return [m.sid for m in self.store.all() if m.state == "active"]
+
+    async def provision(
+        self,
+        sid: str,
+        *,
+        workspace_id: str,
+        agent_id: str | None,
+        workspace_scope: str | None,
+        workspace_user: str | None,
+        env: dict[str, str],
+    ) -> tuple[str, bool]:
+        meta = self.store.get(sid)
+        created = meta is None
+        if meta is None:
+            ws = self.store.workspace_dir(
+                sid,
+                agent_id=agent_id,
+                scope=workspace_scope,
+                user=workspace_user,
+            )
+            meta = SessionMeta(
+                sid=sid,
+                workspace_id=workspace_id,
+                agent_id=agent_id,
+                workspace_scope=workspace_scope,
+                workspace_user=workspace_user,
+                env=dict(env),
+                workspace_dir=str(ws),
+            )
+        else:
+            meta.env = dict(env)
+        Path(meta.workspace_dir).mkdir(parents=True, exist_ok=True)
+        await self.backend.activate(meta)
+        meta.state = "active"
+        self.store.save(meta)
+        return meta.sid, created
+
+    async def status(self, sid: str) -> State:
+        meta = self.require(sid)
+        if isinstance(self.backend, DockerBackend):
+            return await self.backend.status_async(meta)
+        return self.backend.status(meta)
+
+    async def pause(self, sid: str) -> None:
+        meta = self.require(sid)
+        await self.backend.release(meta)
+        meta.state = "paused"
+        self.store.save(meta)
+
+    async def resume(self, sid: str) -> None:
+        meta = self.require(sid)
+        await self.backend.activate(meta)
+        meta.state = "active"
+        self.store.save(meta)
+
+    async def delete(self, sid: str, *, keep_workspace: bool) -> None:
+        meta = self.store.get(sid)
+        if meta is None:
+            return
+        await self.backend.destroy(meta)
+        self.store.delete(sid, keep_workspace=keep_workspace)
+
+    def exec(
+        self,
+        sid: str,
+        cmd: str,
+        env: dict[str, str],
+        *,
+        stdin: str | None,
+        timeout: float,
+    ) -> AsyncIterator[dict[str, Any]]:
+        meta = self.require(sid)
+        return self.backend.exec(meta, cmd, env, stdin=stdin, timeout=timeout)
+
+    async def capture(
+        self, sid: str, cmd: str, env: dict[str, str], *, timeout: float
+    ) -> tuple[str, int]:
+        meta = self.require(sid)
+        return await self.backend.capture(meta, cmd, env, timeout=timeout)
+
+    async def put_file(self, sid: str, path: str, data: bytes) -> None:
+        await self.backend.put_file(
+            self.require(sid), validate_abs_path(path), data
+        )
+
+    async def get_file(self, sid: str, path: str) -> bytes | None:
+        return await self.backend.get_file(
+            self.require(sid), validate_abs_path(path)
+        )
+
+    async def put_tree(self, sid: str, path: str, src: IO[bytes]) -> None:
+        await self.backend.put_tree(
+            self.require(sid), validate_abs_path(path), src
+        )
+
+    async def link(self, sid: str, path: str, target: str) -> None:
+        await self.backend.link(
+            self.require(sid),
+            validate_abs_path(path),
+            validate_abs_path(target),
+        )
+
+    async def reset_agent_workspace(
+        self, agent_id: str, *, user: str | None
+    ) -> None:
+        """Empty a durable workspace; refused while a session is using it."""
+        _check_id(agent_id, "agent_id")
+        target = self.store.workspace_dir(
+            "",
+            agent_id=agent_id,
+            scope="user" if user else "agent",
+            user=user,
+        )
+        for meta in self.store.all():
+            if meta.workspace_dir == str(target) and meta.state == "active":
+                raise RuntimeBusy(f"workspace in use by session {meta.sid}")
+        shutil.rmtree(target, ignore_errors=True)
+        target.mkdir(parents=True, exist_ok=True)
+
+
+__all__ = [
+    "DirectBackend",
+    "DockerBackend",
+    "MissingTarget",
+    "RuntimeBusy",
+    "SessionManager",
+    "SessionMeta",
+    "UnknownSession",
+    "validate_abs_path",
+]