feat(v1): integrate nano-RLM training over ACP

verifiers/v1/harnesses/rlm/harness.py

@@ -1,20 +1,18 @@
 """RLM over ACP, with MCP tools exposed as pre-imported IPython skills."""
 
-import json
 import logging
 import random
 import shlex
 from typing import Literal
 
-from pydantic import Field, PositiveInt, model_validator
+from pydantic import BaseModel, ConfigDict, Field, PositiveInt, model_validator
 
-from verifiers.v1.acp import ACPConfig, ACPHarness
+from verifiers.v1.acp import ACPConfig, ACPHarness, ACPTurnResult, JsonObject
 from verifiers.v1.clients import ModelContext
 from verifiers.v1.configs.harness import HarnessConfig
 from verifiers.v1.runtimes import Runtime
 from verifiers.v1.task import TaskData
 from verifiers.v1.trace import Trace
-from verifiers.v1.utils.decorators import metric
 
 logger = logging.getLogger(__name__)
 
@@ -25,15 +23,24 @@
 RLM_BIN = f"{RLM_DIR}/bin/rlm"
 SKILLS_DIR = "/task/rlm-skills"
 RLM_STATE_DIR = ".vf-rlm"
+RLM_RUNTIME_METADATA_KEY = "ai.prime.rlm/runtime-v1"
+RLM_SESSION_METADATA_KEY = "ai.prime.rlm/session-v1"
+
+
+class _SessionSnapshot(BaseModel):
+    model_config = ConfigDict(extra="ignore", strict=True)
+
+    session_id: str = Field(pattern=r"^[A-Za-z0-9._:-]{1,128}$")
+    metrics: dict[str, int | float]
 
 
 class RLMHarnessConfig(HarnessConfig):
     version: str = Field(
-        default="83ef01f7a6c97328919387343bd30cf4edaac20d", min_length=1
+        default="5ee1c34024a183bbbd3a38a6129995f5b982631d", min_length=1
     )
     """Git ref (branch, tag, or commit) of nano-rlm to install."""
     max_depth: int = 0
-    """Recursion depth rlm may spawn sub-harnesses to (RLM_MAX_DEPTH)."""
+    """Recursion depth RLM may spawn sub-harnesses to."""
     builtin_skills: list[BuiltinSkill] = Field(default_factory=list)
     """Built-in rlm skills to enable (RLM_SKILLS), e.g. `["edit"]`; empty enables none.
     The tool set is fixed (ipython); the base `skills` field takes SKILL.md paths."""
@@ -54,7 +61,6 @@ def validate_range(self) -> "RLMHarnessConfig":
 
     @model_validator(mode="after")
     def reject_disabled_tools(self) -> "RLMHarnessConfig":
-        # rlm's only tool is ipython, which must stay enabled, so there's nothing to disable.
         if self.disabled_tools:
             raise ValueError(
                 "the rlm harness has a fixed tool set (ipython) and does not support "
@@ -91,41 +97,45 @@ async def setup(self, runtime: Runtime) -> None:
             raise RuntimeError(f"rlm install failed: {result.stderr.strip()[-500:]}")
         await super().setup(runtime)
 
-    def summarize_threshold(self, task_idx: int | None) -> str:
-        """The `RLM_SUMMARIZE_AT_TOKENS` value: a range draws per-group (seeded by task index —
-        0 when unset — so a task's rollouts share one threshold). Always set — "" when disabled —
-        so the typed field, not a host var the subprocess runtime would inherit, wins."""
+    def summarize_threshold(self, task_idx: int | None) -> int | None:
+        """Resolve a fixed or per-task compaction threshold."""
         value = self.config.summarize_at_tokens
         if value is None:
-            return ""
+            return None
         if isinstance(value, tuple):
             lo, hi = value
-            return str(random.Random(task_idx or 0).randint(lo, hi))
-        return str(value)
+            return random.Random(task_idx or 0).randint(lo, hi)
+        return value
 
-    def _env(
+    def _runtime_metadata(
         self,
         ctx: ModelContext,
         trace: Trace,
+        runtime: Runtime,
         endpoint: str,
         secret: str,
         data: TaskData,
         system_prompt: str | None,
-    ) -> dict[str, str]:
-        env = {
-            **self.config.resolved_env,
-            "RLM_BASE_URL": endpoint,
-            "RLM_API_KEY": secret,
-            "RLM_MODEL": ctx.model,
-            "RLM_MAX_DEPTH": str(self.config.max_depth),
-            "RLM_HOME": self._home(trace),
-            "RLM_SUMMARIZE_AT_TOKENS": self.summarize_threshold(data.idx),
+    ) -> JsonObject:
+        payload = {
+            "session_id": trace.id,
+            "model": ctx.model,
+            "provider": {
+                "base_url": endpoint,
+                "api_key": secret,
+            },
+            "policy": {
+                "max_depth": self.config.max_depth,
+                "summarize_at_tokens": self.summarize_threshold(data.idx),
+                "max_concurrent_subagents": max(4, self.config.max_depth),
+            },
+            "system_prompt_path": None,
+            "append_to_system_prompt": system_prompt,
+            "skills": list(self.config.builtin_skills),
+            "kernel_env": runtime.env,
+            "search_api_key": self.config.resolved_env.get("SERPER_API_KEY"),
         }
-        if system_prompt is not None:
-            env["RLM_APPEND_TO_SYSTEM_PROMPT"] = system_prompt
-        if self.config.builtin_skills:
-            env["RLM_SKILLS"] = ",".join(self.config.builtin_skills)
-        return env
+        return {RLM_RUNTIME_METADATA_KEY: payload}
 
     async def prepare_acp(
         self,
@@ -139,29 +149,21 @@ async def prepare_acp(
     ) -> ACPConfig:
         system_prompt, prompt = self.resolve_prompt(data)
         return ACPConfig(
-            env=self._env(ctx, trace, endpoint, secret, data, system_prompt),
+            env={**self.config.resolved_env, "RLM_HOME": self._home(trace)},
             command=[RLM_BIN, "--acp"],
             prompt=prompt,
+            session_meta=self._runtime_metadata(
+                ctx, trace, runtime, endpoint, secret, data, system_prompt
+            ),
         )
 
-    @metric
-    async def rlm(self, trace: Trace, runtime: Runtime) -> dict[str, float]:
-        # RolloutRun closes the harness session before metrics, which finalizes
-        # RLM's meta.json while leaving the harness-owned state available here.
-        home = shlex.quote(self._home(trace))
-        latest = f'cat "$(ls -t {home}/sessions/*/meta.json | head -1)"'
-        result = await runtime.run(["sh", "-c", latest], {})
-        if result.exit_code != 0 or not result.stdout.strip():
-            return {}
-        try:
-            meta = json.loads(result.stdout)
-        except json.JSONDecodeError:
-            return {}
-        return {
-            key: float(value)
-            for key, value in meta.get("metrics", {}).items()
-            if isinstance(value, (int, float)) and not isinstance(value, bool)
-        }
+    def acp_turn_result(self, trace: Trace, result: ACPTurnResult) -> None:
+        snapshot = _SessionSnapshot.model_validate(
+            result.response_metadata.get(RLM_SESSION_METADATA_KEY)
+        )
+        if snapshot.session_id != trace.id:
+            raise ValueError("RLM session snapshot does not match the rollout")
+        trace.record_metrics(snapshot.metrics)
 
     async def cleanup(self, trace: Trace, runtime: Runtime) -> None:
         await runtime.run(["rm", "-rf", f"{RLM_STATE_DIR}/{trace.id}"], {})