distinct_agent/rlm.py

"""Optional DSPy RLM harness with a network-denied Deno/Pyodide interpreter.

The interpreter receives no filesystem, environment, process, FFI, or network
grants.  Explicit Distinct tools remain host-side broker callbacks, so a job
can use only the exact tool versions selected by the user.
"""

from __future__ import annotations

import importlib
import json
import os
import re
import shutil
import subprocess
import sys
import threading
from collections.abc import Callable, Mapping
from contextlib import nullcontext
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from typing import Any

from distinct_protocol import (
    PHASE_CALLING_TOOL,
    PHASE_GENERATING,
    PHASE_PLANNING,
    STEP_PHASE,
    STEP_TOOL_CALL,
    STEP_TOOL_RESULT,
    JobSpec,
)

from .harness import HarnessResult, emit
from .models import DiscoveredModel
from .runners import InferenceRunner, RunnerCancelled
from .tools import ToolBrokerAdapter


class RlmUnavailable(RuntimeError):
    """Raised rather than silently dropping the requested RLM sandbox."""


@dataclass(frozen=True)
class RlmLimits:
    max_iters: int = 8
    max_llm_calls: int = 12
    max_output_chars: int = 10_000

    def __post_init__(self) -> None:
        if not 1 <= self.max_iters <= 20:
            raise ValueError("max_iters must be between 1 and 20")
        if not 1 <= self.max_llm_calls <= 50:
            raise ValueError("max_llm_calls must be between 1 and 50")
        if not 1_000 <= self.max_output_chars <= 65_536:
            raise ValueError("max_output_chars must be between 1,000 and 65,536")


class DspyRlmHarness:
    """Explore conversation context through DSPy's local WASM Python REPL."""

    name = "dspy-rlm-deno"
    handles_plain_jobs = True

    def __init__(
        self,
        *,
        limits: RlmLimits | None = None,
        dspy_module: Any = None,
        interpreter_factory: Callable[[], Any] | None = None,
    ) -> None:
        self.limits = limits or RlmLimits()
        self._dspy_module = dspy_module
        self._interpreter_factory = interpreter_factory

    def run(
        self,
        *,
        job: JobSpec,
        model: DiscoveredModel,
        runner: InferenceRunner,
        broker: ToolBrokerAdapter,
        prompt: str,
        cancel_event: threading.Event,
        progress: Callable[[float, str], None],
    ) -> HarnessResult:
        dspy = self._load_dspy()
        emit(
            progress,
            0.03,
            "Starting network-denied DSPy RLM sandbox",
            phase=PHASE_PLANNING,
            step={"kind": STEP_PHASE, "text": "Starting the sandboxed agent"},
        )
        model_usage: list[Mapping[str, Any]] = []
        tool_events: list[Mapping[str, Any]] = []
        artifacts: list[Mapping[str, Any]] = []
        local_lm = _make_local_lm(
            dspy,
            model=model,
            runner=runner,
            job=job,
            cancel_event=cancel_event,
            progress=progress,
            usage_sink=model_usage,
        )
        tools = _brokered_rlm_tools(
            broker,
            tool_events,
            cancel_event,
            artifact_sink=artifacts,
            progress=progress,
        )

        # An empty grant list is intentional.  sync_files=False also prevents
        # even explicitly mounted files (none here) from being copied back.
        interpreter_factory = self._interpreter_factory or partial(
            _network_denied_interpreter, dspy
        )
        try:
            rlm = dspy.RLM(
                "context, query -> answer",
                max_iters=self.limits.max_iters,
                max_llm_calls=self.limits.max_llm_calls,
                max_output_chars=self.limits.max_output_chars,
                verbose=False,
                tools=tools,
                sub_lm=local_lm,
                interpreter_factory=interpreter_factory,
            )
            context_factory = getattr(dspy, "context", None)
            context = context_factory(lm=local_lm) if callable(context_factory) else nullcontext()
            with context:
                prediction = rlm(context=prompt, query=job.prompt)
        except RunnerCancelled:
            raise
        except Exception as exc:
            raise RlmUnavailable(
                f"DSPy RLM/Deno execution failed: {type(exc).__name__}: {exc}"
            ) from exc
        if cancel_event.is_set():
            raise RunnerCancelled("job was cancelled during RLM execution")
        answer = getattr(prediction, "answer", None)
        if not isinstance(answer, str):
            raise RlmUnavailable("DSPy RLM returned no text answer")
        if len(answer) > 1_000_000:
            raise RlmUnavailable("DSPy RLM answer exceeds the result limit")
        trajectory = getattr(prediction, "trajectory", ())
        iterations = len(trajectory) if isinstance(trajectory, list | tuple) else None
        emit(progress, 1.0, "DSPy RLM inference complete", phase=PHASE_GENERATING)
        usage: dict[str, Any] = {
            "harness": self.name,
            "model_calls": len(model_usage),
            "tool_calls": len(tool_events),
            "model_usage": [dict(value) for value in model_usage],
            "sandbox": {
                "runtime": "Deno/Pyodide/WASM",
                "network": "denied",
                "filesystem": "denied",
                "environment": "denied",
                "sync_files": False,
            },
        }
        if iterations is not None:
            usage["rlm_iterations"] = iterations
        return HarnessResult(answer, usage, tuple(tool_events), tuple(artifacts))

    def _load_dspy(self) -> Any:
        if self._dspy_module is not None:
            return self._dspy_module
        try:
            module = importlib.import_module("dspy")
        except ImportError as exc:
            raise RlmUnavailable(
                'DSPy RLM is unavailable; install the agent extra "dspy[deno]==3.3.0"'
            ) from exc
        for attribute in ("BaseLM", "LMResponse", "PythonInterpreter", "RLM"):
            if not hasattr(module, attribute):
                raise RlmUnavailable(f"installed DSPy has no {attribute}")
        return module

    def check_ready(self) -> None:
        """Start and close a sandbox before the worker advertises availability."""

        dspy = self._load_dspy()
        factory = self._interpreter_factory or partial(_network_denied_interpreter, dspy)
        interpreter = None
        try:
            interpreter = factory()
            start = getattr(interpreter, "start", None)
            if callable(start):
                start()
            result = interpreter("print(40 + 2)")
            if str(result).strip() != "42":
                raise RlmUnavailable("Deno/Pyodide readiness probe returned an unexpected value")
        except RlmUnavailable:
            raise
        except Exception as exc:
            raise RlmUnavailable(
                f"Deno/Pyodide readiness probe failed: {type(exc).__name__}: {exc}"
            ) from exc
        finally:
            if interpreter is not None:
                shutdown = getattr(interpreter, "shutdown", None)
                if callable(shutdown):
                    shutdown()


def _network_denied_interpreter(dspy: Any) -> Any:
    """Create a cached-only Deno child with no ambient runtime grants."""

    executable = _deno_executable()
    cache_dir = _deno_cache_directory(executable)
    runner_path = _dspy_runner_path(dspy)
    lock_path = _lock_path()
    # Deno's comma-separated permission syntax cannot represent a comma in a
    # path without broadening access.  Reject that uncommon case explicitly.
    if any("," in str(path) for path in (runner_path, cache_dir, lock_path)):
        raise RlmUnavailable("Deno runtime paths containing commas are unsupported")
    command = [
        executable,
        "run",
        "--no-config",
        "--node-modules-dir=false",
        "--cached-only",
        f"--lock={lock_path}",
        "--frozen",
        f"--allow-read={runner_path},{cache_dir}",
        str(runner_path),
    ]
    return dspy.PythonInterpreter(
        deno_command=command,
        enable_read_paths=[],
        enable_write_paths=[],
        enable_env_vars=[],
        enable_network_access=[],
        sync_files=False,
    )


def prepare_rlm_runtime(*, dspy_module: Any = None, timeout_seconds: float = 300.0) -> None:
    """Explicitly fetch lock-pinned Deno dependencies, then verify offline use.

    This is an installation/build operation and may access the network. Normal
    inference always uses ``--cached-only`` and cannot fetch dependencies.
    """

    dspy = dspy_module or importlib.import_module("dspy")
    executable = _deno_executable()
    runner_path = _dspy_runner_path(dspy)
    lock_path = _lock_path()
    try:
        completed = subprocess.run(
            [
                executable,
                "cache",
                "--no-config",
                "--node-modules-dir=false",
                f"--lock={lock_path}",
                "--frozen",
                str(runner_path),
            ],
            stdin=subprocess.DEVNULL,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=timeout_seconds,
            check=False,
        )
    except (OSError, subprocess.SubprocessError) as exc:
        raise RlmUnavailable("failed to prepare the lock-pinned Deno cache") from exc
    if completed.returncode != 0:
        detail = " ".join(completed.stderr.split())[-1000:]
        raise RlmUnavailable(f"Deno dependency preparation failed: {detail or 'no detail'}")
    DspyRlmHarness(dspy_module=dspy).check_ready()


def _dspy_runner_path(dspy: Any) -> Path:
    module_path = getattr(dspy, "__file__", None)
    if not isinstance(module_path, str):
        raise RlmUnavailable("cannot locate the installed DSPy package")
    runner_path = (Path(module_path).resolve().parent / "primitives" / "runner.js").resolve()
    if not runner_path.is_file():
        raise RlmUnavailable("DSPy Deno runner.js is missing")
    return runner_path


def _lock_path() -> Path:
    path = Path(__file__).resolve().with_name("deno.lock")
    if not path.is_file():
        raise RlmUnavailable("the pinned Deno dependency lock is missing")
    return path


def _deno_executable() -> str:
    located = shutil.which("deno")
    if located:
        return str(Path(located).resolve())
    try:
        deno = importlib.import_module("deno")
        candidate = deno.find_deno_bin()
    except (ImportError, AttributeError, OSError) as exc:
        raise RlmUnavailable(
            'managed Deno is unavailable; install the agent extra "dspy[deno]==3.3.0"'
        ) from exc
    path = Path(candidate).resolve()
    if not path.is_file():
        raise RlmUnavailable("managed Deno executable is missing")
    return str(path)


def _deno_cache_directory(executable: str) -> Path:
    bundle_root = getattr(sys, "_MEIPASS", None)
    if isinstance(bundle_root, str):
        bundled = (Path(bundle_root) / "deno_cache").resolve()
        if bundled.is_dir():
            return bundled
    configured = os.environ.get("DENO_DIR")
    if configured:
        path = Path(configured).expanduser().resolve()
    else:
        try:
            completed = subprocess.run(
                [executable, "info", "--json"],
                stdin=subprocess.DEVNULL,
                capture_output=True,
                text=True,
                encoding="utf-8",
                errors="replace",
                timeout=10,
                check=False,
            )
            value = json.loads(completed.stdout) if completed.returncode == 0 else {}
            path = Path(value.get("denoDir", "")).resolve()
        except (OSError, subprocess.SubprocessError, json.JSONDecodeError) as exc:
            raise RlmUnavailable("cannot locate Deno's dependency cache") from exc
    if not path.is_dir():
        raise RlmUnavailable(
            "Deno dependency cache is absent; prepare the pinned RLM runtime before serving jobs"
        )
    return path


def _make_local_lm(
    dspy: Any,
    *,
    model: DiscoveredModel,
    runner: InferenceRunner,
    job: JobSpec,
    cancel_event: threading.Event,
    progress: Callable[[float, str], None],
    usage_sink: list[Mapping[str, Any]],
) -> Any:
    class LocalRunnerLM(dspy.BaseLM):
        forward_contract = "typed_lm"

        def __init__(self) -> None:
            super().__init__(model=f"local/{model.manifest.id}", cache=False, num_retries=0)

        def forward(self, request: Any) -> Any:
            if cancel_event.is_set():
                raise RunnerCancelled("job was cancelled before an RLM model call")
            progress(0.1, "DSPy RLM local model call")
            inference = runner.run(
                model,
                _request_prompt(request),
                cancel_event=cancel_event,
                progress=None,
                limits=job.limits,
            )
            usage_sink.append(dict(inference.usage))
            return dspy.LMResponse.from_text(inference.text, model=self.model)

    return LocalRunnerLM()


def _request_prompt(request: Any) -> str:
    messages = getattr(request, "messages", None)
    if messages is None:
        prompt = getattr(request, "prompt", None)
        return prompt if isinstance(prompt, str) else str(request)
    parts: list[str] = []
    for message in messages:
        value = _message_mapping(message)
        role = str(value.get("role", "user"))
        content = value.get("content", "")
        if isinstance(content, str):
            rendered = content
        else:
            rendered = json.dumps(content, ensure_ascii=False, default=str)
        parts.append(f"{role}: {rendered}")
    return "\n".join(parts)


def _message_mapping(message: Any) -> Mapping[str, Any]:
    if isinstance(message, Mapping):
        return message
    for method_name in ("to_dict", "model_dump"):
        method = getattr(message, method_name, None)
        if callable(method):
            value = method()
            if isinstance(value, Mapping):
                return value
    return {
        "role": getattr(message, "role", "user"),
        "content": getattr(message, "content", str(message)),
    }


def _brokered_rlm_tools(
    broker: ToolBrokerAdapter,
    event_sink: list[Mapping[str, Any]],
    cancel_event: threading.Event,
    *,
    artifact_sink: list[Mapping[str, Any]] | None = None,
    progress: Callable[..., None] | None = None,
) -> list[Callable[..., str]]:
    """Expose exactly the job's granted tools to DSPy, and nothing else.

    The adapter is generic: it is derived from the broker's manifest, so it
    neither knows nor cares which tools exist.  With the registry shipping
    empty this returns an empty list, and the RLM runs with no tools at all —
    which is the correct behaviour, not a degraded one.
    """

    tools: list[Callable[..., str]] = []
    for item in broker.tool_manifest():
        tool_id = str(item.get("id", ""))
        version = str(item.get("version", ""))
        if not tool_id or not version:
            raise RlmUnavailable("tool manifest entry is missing an id or version")
        tools.append(
            _make_brokered_tool(
                broker=broker,
                tool_id=tool_id,
                version=version,
                description=str(item.get("description", ""))[:2_000],
                schema=item.get("input_schema"),
                event_sink=event_sink,
                cancel_event=cancel_event,
                artifact_sink=artifact_sink,
                progress=progress,
            )
        )
    return tools


def _make_brokered_tool(
    *,
    broker: ToolBrokerAdapter,
    tool_id: str,
    version: str,
    description: str,
    schema: Any,
    event_sink: list[Mapping[str, Any]],
    cancel_event: threading.Event,
    artifact_sink: list[Mapping[str, Any]] | None = None,
    progress: Callable[..., None] | None = None,
) -> Callable[..., str]:
    """Wrap one exact ``id@version`` as a single-argument DSPy callable.

    Arguments arrive as a JSON object so the wrapper never has to mirror a
    tool's parameter list.  Validation stays where it belongs: in the tool's
    own schema, enforced behind the broker.
    """

    schema_json = json.dumps(schema, ensure_ascii=False, default=str) if schema else "{}"

    def call_tool(arguments: Any = None, **keywords: Any) -> str:
        if cancel_event.is_set():
            raise RunnerCancelled("job was cancelled before a tool call")
        if arguments is None:
            payload: Mapping[str, Any] = keywords
        elif isinstance(arguments, Mapping):
            payload = {**arguments, **keywords}
        elif isinstance(arguments, str):
            try:
                decoded = json.loads(arguments)
            except (TypeError, ValueError) as exc:
                raise ValueError("tool arguments must be a JSON object") from exc
            if not isinstance(decoded, Mapping):
                raise ValueError("tool arguments must be a JSON object")
            payload = {**decoded, **keywords}
        else:
            raise ValueError("tool arguments must be a JSON object")

        if progress is not None:
            from .harness import _argument_preview

            emit(
                progress,
                0.5,
                f"Calling allowed tool {tool_id}",
                phase=PHASE_CALLING_TOOL,
                step={
                    "kind": STEP_TOOL_CALL,
                    "tool": f"{tool_id}@{version}",
                    "text": _argument_preview(dict(payload)),
                },
            )
        result = broker.invoke(tool_id, dict(payload), version=version)
        value = result.to_dict() if hasattr(result, "to_dict") else result
        if not isinstance(value, Mapping):
            raise RuntimeError("tool returned an invalid result envelope")
        if artifact_sink is not None:
            # Same lift as the structured harness: an artifact's bytes go to
            # the run result for the user's session, never back into the
            # model's context.
            from .harness import lift_artifact

            value = lift_artifact(value, artifact_sink)
        event_sink.append(_safe_tool_event(value))
        if progress is not None:
            from .harness import _result_preview

            emit(
                progress,
                0.55,
                f"{tool_id} returned",
                phase=PHASE_CALLING_TOOL,
                step={
                    "kind": STEP_TOOL_RESULT,
                    "tool": f"{tool_id}@{version}",
                    "ok": bool(value.get("ok", True)),
                    "text": _result_preview(value),
                },
            )
        return json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":"))

    call_tool.__name__ = re.sub(r"[^0-9a-zA-Z_]", "_", f"{tool_id}_{version}")
    call_tool.__doc__ = (
        f"{description}\n\nCall with one JSON object of arguments. "
        f"INPUT_SCHEMA={schema_json}"
    )
    return call_tool


def _safe_tool_event(value: Mapping[str, Any]) -> Mapping[str, Any]:
    event: dict[str, Any] = {
        "tool_id": str(value.get("tool_id", ""))[:128],
        "version": str(value.get("version", ""))[:32],
        "ok": value.get("ok") is True,
    }
    elapsed = value.get("elapsed_ms")
    if isinstance(elapsed, int) and not isinstance(elapsed, bool) and elapsed >= 0:
        event["elapsed_ms"] = elapsed
    error = value.get("error")
    if isinstance(error, Mapping):
        event["error_code"] = str(error.get("code", "tool_error"))[:80]
    output = value.get("output")
    if isinstance(output, Mapping):
        saved = output.get("artifact_saved")
        if isinstance(saved, Mapping) and isinstance(saved.get("name"), str):
            event["artifact"] = saved["name"][:120]
    return event


#: The two modes an operator chooses between, and their aliases.
#:
#: ``agent`` is the default: a DSPy RLM agent that reasons, acts and loops
#: inside the Deno sandbox. ``simple`` is one plan pass, the tools, and one
#: answer pass, for models too small to hold an agent loop together.
#:
#: ``structured`` is neither, and is not offered as a mode. It is the older
#: bounded read-eval loop, kept because the smoke tests and the deterministic
#: demo runner are written against it.
HARNESS_ALIASES: dict[str, str] = {
    "agent": "dspy-rlm",
    "dspy-rlm": "dspy-rlm",
    "simple": "simple",
    "structured": "structured",
}

AGENT_MODES: tuple[str, ...] = ("agent", "simple")


def build_inference_harness(name: str) -> DspyRlmHarness | Any:
    """Create a named harness without importing optional DSPy eagerly."""

    resolved = HARNESS_ALIASES.get(name, name)
    if resolved == "dspy-rlm":
        return DspyRlmHarness()
    if resolved == "simple":
        from .harness import SimpleToolHarness

        return SimpleToolHarness()
    if resolved == "structured":
        from .harness import StructuredToolHarness

        return StructuredToolHarness()
    raise ValueError(f"unknown inference harness {name!r}")