rlmness-cli 0.2.0__py3-none-any.whl

rlmagent/kernel/wasm.py

@@ -0,0 +1,169 @@
+"""The python tool's kernel inside a WebAssembly sandbox (Pyodide, through rlmness).
+
+Same surface as KernelSession, so the python tool and the agent tree do not care
+which one they have. The sandbox has an empty filesystem, no network and no host
+environment; host functions (bridges) are the only way out.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from collections.abc import Awaitable, Callable, Mapping
+
+from rlmagent.kernel.session import CellResult, clip
+
+Bridge = Callable[..., Awaitable[object]]
+
+LOST = ("[The sandbox was replaced ({why}): variables from before are gone. "
+        "Re-run any setup you still need.]")
+
+
+def available() -> bool:
+    try:
+        from rlmness.wasm_runtime import wasm_available
+    except ImportError:
+        return False
+    return wasm_available()
+
+
+class WasmKernel:
+    def __init__(
+        self,
+        cwd: str,
+        bridges: Mapping[str, Bridge],
+        startup_code: str = "",
+        timeout: float = 300.0,
+        on_final: Callable[[object], Awaitable[None]] | None = None,
+        output_limit: int = 20_000,
+        restore: bool = False,
+    ) -> None:
+        self.cwd = cwd
+        self.env: dict[str, str] = {}
+        self.bridges = dict(bridges)
+        self.startup_code = startup_code
+        self.timeout = timeout
+        self.on_final = on_final
+        self.output_limit = output_limit
+        self.restore = restore
+        self._checkpoint: dict = {}
+        self._owned: set[str] | None = None
+        self._recovery_note = ""
+        self._runtime = None
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._lock = asyncio.Lock()
+
+    @contextlib.contextmanager
+    def waiting_outside(self):
+        # The sandbox does not time a cell out while one of its bridges is running.
+        yield
+
+    @property
+    def started(self) -> bool:
+        return self._runtime is not None
+
+    def _blocking(self, bridge: Bridge) -> Callable[..., object]:
+        """rlmness serves bridges on worker threads; the work belongs on this agent's loop."""
+        def call(*args, **kwargs):
+            future = asyncio.run_coroutine_threadsafe(bridge(*args, **kwargs), self._loop)
+            return future.result()
+        return call
+
+    async def _start(self) -> str | None:
+        from rlmness.wasm_runtime import WasmRuntime
+
+        self._loop = asyncio.get_running_loop()
+        bridges = {name: self._blocking(fn) for name, fn in self.bridges.items()}
+        options = {"session": self._checkpoint} if self.restore else {}
+        # Loading Pyodide on a cold host has a separate deadline from user code.
+        self._runtime = await asyncio.to_thread(
+            WasmRuntime, None, bridges, max(30.0, self.timeout), **options
+        )
+        if self.startup_code:
+            outcome = await asyncio.to_thread(self._runtime.execute, self.startup_code)
+            if outcome.error:
+                await self._drop()
+                return outcome.error
+        if self.restore:
+            if self._owned is None:
+                initial = await asyncio.to_thread(self._runtime.sweep)
+                self._owned = {name for entries in initial.values() for name in entries}
+                self._owned.update(("CONTEXT", "PROMPT", "FINAL", "commit"))
+            if self._checkpoint:
+                failed = self._runtime.restore_failed
+                saved = set().union(*(self._checkpoint.get(k, {}) for k in
+                                      ("variables", "functions", "modules")))
+                restored = sorted(saved - failed)
+                missing = sorted(set(self._checkpoint.get("dropped", {})) | failed)
+                self._recovery_note = ("[Sandbox checkpoint restored: "
+                                       + (", ".join(restored) or "no values")
+                                       + "; unavailable: " + (", ".join(missing) or "none")
+                                       + ". Changes in the interrupted cell were not recovered.]\n")
+        self._runtime.timeout = self.timeout
+        return None
+
+    async def _save_checkpoint(self, code: str) -> None:
+        state = await asyncio.to_thread(self._runtime.sweep, code)
+        if state:
+            self._checkpoint = {kind: {name: value for name, value in entries.items()
+                                      if name not in self._owned}
+                                for kind, entries in state.items()}
+
+    def _lost(self, why: str) -> str:
+        if not self.restore:
+            return LOST.format(why=why)
+        return (f"[The sandbox was replaced ({why}). The next cell will restore the last "
+                "completed checkpoint; changes from the interrupted cell are lost. "
+                "External side effects may already have happened.]")
+
+    async def _drop(self) -> None:
+        runtime, self._runtime = self._runtime, None
+        if runtime is not None:
+            await asyncio.to_thread(runtime.close)
+
+    async def run(self, code: str, timeout: float | None = None) -> CellResult:
+        from rlmness.runtime import CellTimeout, RuntimeGone
+
+        async with self._lock:
+            if self._runtime is None:
+                failure = await self._start()
+                if failure:
+                    return CellResult(output="", error=f"sandbox start-up failed:\n{failure}")
+            cell_timeout = self.timeout if timeout is None else timeout
+            self._runtime.timeout = cell_timeout
+            try:
+                outcome = await asyncio.to_thread(self._runtime.execute, code)
+            except CellTimeout:
+                from rlmness.wasm_runtime import WasmRuntime, reset_host
+
+                self._runtime = None
+                if not getattr(WasmRuntime, "ISOLATED", False):
+                    await asyncio.to_thread(reset_host)
+                return CellResult(output=self._lost(why=f"no answer within {cell_timeout:.0f}s"),
+                                  timed_out=True, restarted=True)
+            except RuntimeGone:
+                await self._drop()
+                return CellResult(output=self._lost(why="it stopped"), restarted=True)
+            finally:
+                if self._runtime is not None:
+                    self._runtime.timeout = self.timeout
+            if self.restore and not outcome.error:
+                await self._save_checkpoint(code)
+            recovery_note, self._recovery_note = self._recovery_note, ""
+        output = recovery_note + outcome.stdout
+        if outcome.final_given and self.on_final is not None:
+            await self.on_final(outcome.final)
+            output += "FINAL recorded.\n"
+        return CellResult(output=clip(output, self.output_limit),
+                          error=clip(outcome.error, self.output_limit) if outcome.error else None)
+
+    async def interrupt(self) -> None:
+        # Pyodide cannot be interrupted from outside; stopping means replacing it.
+        await self._drop()
+
+    async def kill(self) -> None:
+        await self._drop()
+
+    async def shutdown(self) -> None:
+        await self._drop()
+        self._checkpoint = {}

rlmagent/logs.py

@@ -0,0 +1,208 @@
+"""Run logging: trajectory, raw events, and stdout captured under ``.logs/``."""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+from rlmagent.refresh import _scrub
+from rlmagent.harness.contracts.stream import (
+    AgentEvent,
+    MessageEndEvent,
+    RunEndEvent,
+    RunStartEvent,
+    ToolRunEndEvent,
+    ToolRunStartEvent,
+)
+from rlmagent.harness.contracts.transcript import HumanEntry, ModelEntry
+
+_DIR_ENV = "RLM_AGENT_LOG_DIR"
+_DISABLE_ENV = "RLM_AGENT_NO_LOGS"
+_DEFAULT_DIR = ".logs"
+
+
+def _now() -> str:
+    return datetime.now(UTC).isoformat(timespec="milliseconds")
+
+
+def logs_enabled() -> bool:
+    """Whether run logging is switched on."""
+    return os.environ.get(_DISABLE_ENV, "").strip().lower() not in {"1", "true", "yes"}
+
+
+def log_root(cwd: str | Path | None = None) -> Path:
+    """Directory holding run logs — ``RLM_AGENT_LOG_DIR`` or ``<cwd>/.logs``."""
+    override = os.environ.get(_DIR_ENV)
+    if override:
+        return Path(override)
+    return Path(cwd or os.getcwd()) / _DEFAULT_DIR
+
+
+def _clean(value: Any) -> Any:
+    """Recursively scrub strings inside a JSON-compatible value."""
+    if isinstance(value, str):
+        return _scrub(value)
+    if isinstance(value, dict):
+        return {k: _clean(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_clean(v) for v in value]
+    return value
+
+
+@dataclass(slots=True)
+class RunLogger:
+    """Writes one session's logs into ``<root>/<session_id>/``."""
+
+    session_id: str
+    root: Path
+    provider: str = ""
+    model: str = ""
+    steps: list[dict[str, Any]] = field(default_factory=list)
+    #: Optional zero-arg callable returning the session's usage stats.
+    stats_source: Any = None
+
+    # -- construction -------------------------------------------------------
+
+    @classmethod
+    def create(
+        cls,
+        session_id: str,
+        *,
+        provider: str = "",
+        model: str = "",
+        cwd: str | Path | None = None,
+        root: Path | None = None,
+    ) -> RunLogger | None:
+        """Prepare the log directory, or return ``None`` when disabled."""
+        if not logs_enabled():
+            return None
+        base = (root or log_root(cwd)) / session_id
+        try:
+            base.mkdir(parents=True, exist_ok=True)
+            # The logs sit in the user's project; keep them out of its git status.
+            ignore = base.parent / ".gitignore"
+            if not ignore.exists():
+                ignore.write_text("*\n", encoding="utf-8")
+        except OSError:
+            return None
+        logger = cls(session_id=session_id, root=base, provider=provider, model=model)
+        logger.write_line(f"session {session_id} started ({provider} · {model})")
+        return logger
+
+    # -- paths --------------------------------------------------------------
+
+    @property
+    def trajectory_path(self) -> Path:
+        return self.root / "trajectory.json"
+
+    @property
+    def events_path(self) -> Path:
+        return self.root / "events.jsonl"
+
+    @property
+    def stdout_path(self) -> Path:
+        return self.root / "stdout.log"
+
+    # -- writing ------------------------------------------------------------
+
+    def write_line(self, text: str) -> None:
+        """Append a timestamped line to ``stdout.log``."""
+        self._append(self.stdout_path, f"{_now()}  {_scrub(text)}\n")
+
+    def record_event(self, event: AgentEvent) -> None:
+        """Append one event to ``events.jsonl`` and update the trajectory."""
+        try:
+            payload = _clean(event.model_dump(mode="json", by_alias=True))
+        except Exception:  # noqa: BLE001 - logging must never break a run
+            return
+        self._append(self.events_path, json.dumps(payload, default=str) + "\n")
+        self._track(event)
+
+    def finish(self, stats: Any = None) -> None:
+        """Write ``trajectory.json``."""
+        if stats is None and self.stats_source is not None:
+            try:
+                stats = self.stats_source()
+            except Exception:  # noqa: BLE001
+                stats = None
+        doc: dict[str, Any] = {
+            "session_id": self.session_id,
+            "provider": self.provider,
+            "model": self.model,
+            "finished_at": _now(),
+            "steps": self.steps,
+        }
+        if stats is not None:
+            doc["usage"] = {
+                "turns": getattr(stats, "turn_count", 0),
+                "messages": getattr(stats, "message_count", 0),
+                "input_tokens": getattr(stats, "total_input_tokens", 0),
+                "output_tokens": getattr(stats, "total_output_tokens", 0),
+                "cost_usd": getattr(stats, "total_cost_usd", 0.0),
+                "context_tokens": getattr(stats, "estimated_context_tokens", 0),
+            }
+        # Scrub the whole document, not just the steps built here — any
+        # caller-supplied content must be covered too.
+        self._write(
+            self.trajectory_path, json.dumps(_clean(doc), indent=2, default=str),
+        )
+        self.write_line("session finished")
+
+    # -- internals ----------------------------------------------------------
+
+    def _track(self, event: AgentEvent) -> None:
+        """Fold an event into the human-readable trajectory."""
+        if isinstance(event, RunStartEvent):
+            self.write_line("run started")
+        elif isinstance(event, ToolRunStartEvent):
+            self.steps.append({
+                "at": _now(), "kind": "tool_call",
+                "tool": event.tool_name,
+                "args": _clean(dict(event.args)),
+            })
+            self.write_line(f"tool {event.tool_name} started")
+        elif isinstance(event, ToolRunEndEvent):
+            self.steps.append({
+                "at": _now(), "kind": "tool_result",
+                "tool": event.tool_name,
+                "is_error": event.is_error,
+                "output": _scrub(event.result.text)[:4000],
+            })
+            self.write_line(
+                f"tool {event.tool_name} {'failed' if event.is_error else 'ok'}"
+            )
+        elif isinstance(event, MessageEndEvent):
+            message = event.message
+            if isinstance(message, ModelEntry) and message.text:
+                self.steps.append({
+                    "at": _now(), "kind": "assistant",
+                    "text": _scrub(message.text),
+                    "stop_reason": message.stop_reason,
+                })
+            elif isinstance(message, HumanEntry) and message.text:
+                self.steps.append({
+                    "at": _now(), "kind": "prompt", "text": _scrub(message.text),
+                })
+        elif isinstance(event, RunEndEvent):
+            self.write_line("run ended")
+            self.finish()
+
+    def _append(self, path: Path, text: str) -> None:
+        try:
+            with path.open("a", encoding="utf-8") as fh:
+                fh.write(text)
+        except OSError:
+            pass  # logging is best effort — never break the run
+
+    def _write(self, path: Path, text: str) -> None:
+        try:
+            path.write_text(text, encoding="utf-8")
+        except OSError:
+            pass
+
+
+__all__ = ["RunLogger", "log_root", "logs_enabled"]