Option B: typed MessageNode.kind for replay resume points

environments/compact/compact/annotate.py

@@ -0,0 +1,42 @@
+"""Producer-side helpers: locate the compaction resume points a harness should tag.
+
+A compacting rollout splits into branches (one per context rewrite). Each compaction exposes
+two replay resume points:
+
+- ``compaction_after``  — the post-compaction branch start (the rewritten ``[system, user(notes)]``).
+  Resuming here, the model continues solving *from* the compaction message.
+- ``compaction_before`` — the leaf of the branch that compaction summarized (the prior turn's
+  response). Resuming here, the model is back in the pre-compaction context and its continuation
+  *writes* the compaction itself (then keeps solving).
+
+This module only *finds* the nodes; where the tag is stored is the A/B decision (Option A:
+``trace.info``; Option B: ``MessageNode.kind``) and lives at the harness write site.
+"""
+
+from __future__ import annotations
+
+from verifiers.v1 import graph
+from verifiers.v1.trace import Trace
+
+
+def compaction_after_nodes(trace: Trace) -> list[int]:
+    """Post-compaction branch starts: the first node of each forked branch. A node with >1
+    child is a fork point; ``children[0]`` is the original line, ``children[1:]`` are the
+    rewritten (post-compaction) branches. The compacting harness rewrites every turn, so each
+    is a compaction boundary."""
+    children: dict[int | None, list[int]] = {}
+    for nid, node in enumerate(trace.nodes):
+        children.setdefault(node.parent, []).append(nid)
+    starts: list[int] = []
+    for kids in children.values():
+        if len(kids) > 1:
+            starts.extend(kids[1:])
+    return starts
+
+
+def compaction_before_nodes(trace: Trace) -> list[int]:
+    """Pre-compaction points: every branch leaf except the final-answer branch. Each such leaf
+    is the turn whose output was summarized into the next branch's compaction message, so
+    resuming there puts the model right before it writes a compaction."""
+    leaves = sorted(graph.leaves(trace))
+    return leaves[:-1] if len(leaves) > 1 else []

environments/compact/compact/harness.py

@@ -16,6 +16,8 @@
 from verifiers.v1.runtimes import ProgramResult, Runtime
 from verifiers.v1.trace import Trace
 
+from compact.annotate import compaction_after_nodes, compaction_before_nodes
+
 PROGRAM_SOURCE = (Path(__file__).resolve().parent / "program.py").read_text()
 
 
@@ -51,4 +53,13 @@ async def launch(
                 {"mcpServers": {name: {"url": url} for name, url in mcp_urls.items()}}
             )
         program = await runtime.prepare_uv_script(PROGRAM_SOURCE, self.config.env)
-        return await runtime.run_program([*program, trace.task.prompt], env)
+        result = await runtime.run_program([*program, trace.task.prompt], env)
+
+        # Tag compaction resume points on the finished graph (Option B: typed MessageNode.kind).
+        # Two points per compaction: the post-compaction branch start (resume to continue from
+        # the compaction message) and the pre-compaction branch leaf (resume to regenerate it).
+        for nid in compaction_after_nodes(trace):
+            trace.nodes[nid].kind = "compaction_after"
+        for nid in compaction_before_nodes(trace):
+            trace.nodes[nid].kind = "compaction_before"
+        return result

environments/replay/pyproject.toml

@@ -0,0 +1,13 @@
+[project]
+name = "replay"
+version = "0.1.0"
+description = "replay — a replay-buffer taskset that resumes old rollouts from compaction points."
+requires-python = ">=3.10"
+dependencies = ["verifiers"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["replay"]

environments/replay/replay/__init__.py

@@ -0,0 +1,18 @@
+"""replay — resume old rollouts from tagged compaction points (a replay buffer).
+
+Offline (``ReplayTaskset`` materializes tasks from the buffer) or online (``ReplayHarness``
+samples the live buffer per rollout). Scoring reuses the original env's verifier.
+
+Option B: compaction tags + snapshot refs are read from typed ``MessageNode`` fields (see
+``replay/selector.py``)."""
+
+from replay.harness import ReplayHarness, ReplayHarnessConfig
+from replay.taskset import ReplayTask, ReplayTaskset, ReplayTasksetConfig
+
+__all__ = [
+    "ReplayTaskset",
+    "ReplayTasksetConfig",
+    "ReplayTask",
+    "ReplayHarness",
+    "ReplayHarnessConfig",
+]

environments/replay/replay/harness.py

@@ -0,0 +1,122 @@
+"""ReplayHarness — rollout-time sampling for an online (growing) replay buffer.
+
+The offline path materializes tasks in ``ReplayTaskset.load_tasks``, which runs once at
+env-server start and so can't see rollouts written later. This harness instead samples a stored
+trace + resume point from the *live* buffer on each rollout, so this run's own rollouts are
+replayed as the buffer fills. It restores the resume point's sandbox snapshot, seeds the default
+chat loop with the replay prefix (``root->node``), and stashes provenance in ``trace.info`` so
+``ReplayTaskset.score`` can reuse the original verifier.
+
+It reuses the default harness's program (a growing-message-list chat loop) and seeds it via the
+same ``INITIAL_MESSAGES`` channel the default harness uses for a Messages prompt.
+"""
+
+from __future__ import annotations
+
+import glob
+import json
+import random
+from pathlib import Path
+
+from verifiers.v1.clients import RolloutContext
+from verifiers.v1.dialects.chat import message_to_wire
+from verifiers.v1.harnesses.default.harness import (
+    PROGRAM_SOURCE,
+    DefaultHarness,
+    DefaultHarnessConfig,
+)
+from verifiers.v1.runtimes import ProgramResult, Runtime
+from verifiers.v1.trace import Trace, WireTrace
+
+from replay.selector import (
+    DEFAULT_FOLLOWUP,
+    DEFAULT_KINDS,
+    build_seed,
+    resume_points,
+    snapshot_ref_of,
+)
+
+
+class ReplayHarnessConfig(DefaultHarnessConfig):
+    buffer_glob: str = ""
+    """Glob of stored-rollout JSONL files to sample from (the live, possibly growing, buffer)."""
+    kinds: list[str] = DEFAULT_KINDS
+    """Which replay kinds to sample (see ReplayTaskset; add ``"judge"`` to opt in)."""
+    followup: str = DEFAULT_FOLLOWUP
+    """The user turn appended for ``recheck`` points."""
+
+
+class ReplayHarness(DefaultHarness):
+    """Subclasses the default harness (its chat-loop program) but seeds from a sampled buffer
+    rollout instead of ``task.prompt``."""
+
+    SUPPORTS_MESSAGE_PROMPT = True
+
+    async def launch(
+        self,
+        ctx: RolloutContext,
+        trace: Trace,
+        runtime: Runtime,
+        endpoint: str,
+        secret: str,
+        mcp_urls: dict[str, str],
+    ) -> ProgramResult:
+        rng = random.Random(trace.id)  # deterministic per rollout, varies across rollouts
+        sample = self._sample(rng)
+        if sample is None:  # buffer empty (warmup) or no matching resume points yet
+            trace.stop("replay_buffer_empty")
+            return ProgramResult(exit_code=0, stdout="", stderr="")
+        src, point = sample
+
+        ref = snapshot_ref_of(src, point["node"])
+        if ref is not None:  # exec/sandbox replay; skeleton refs are None -> skip
+            await runtime.restore(ref)
+
+        # Stash provenance so ReplayTaskset.score can reuse the original verifier.
+        trace.info["replay"] = {
+            "source_id": src.id,
+            "resume_node": point["node"],
+            "kind": point["kind"],
+            "original_task": src.task.model_dump(),
+            "original_reward": src.reward,
+        }
+
+        # Seed the default chat loop with the replay prefix (mirror DefaultHarness.launch).
+        seed = build_seed(src, point, self.config.followup)
+        env = {**self.config.env}
+        env["INITIAL_MESSAGES"] = json.dumps([message_to_wire(m) for m in seed])
+        args = [f"--base-url={endpoint}", f"--api-key={secret}", f"--model={ctx.model}"]
+        if mcp_urls:
+            args.append(
+                "--mcp-config="
+                + json.dumps(
+                    {"mcpServers": {name: {"url": url} for name, url in mcp_urls.items()}}
+                )
+            )
+        program = await runtime.prepare_uv_script(PROGRAM_SOURCE, self.config.env)
+        return await runtime.run_program([*program, *args], env)
+
+    def _sample(self, rng: random.Random) -> tuple[Trace, dict] | None:
+        """Scan the live buffer in random order; return the first (trace, resume point) found.
+        Re-globs every rollout, so files written after env-server start are included."""
+        kinds = set(self.config.kinds)
+        files = sorted(glob.glob(self.config.buffer_glob))
+        rng.shuffle(files)
+        for path in files:
+            try:
+                lines = Path(path).read_text().splitlines()
+            except OSError:
+                continue
+            rng.shuffle(lines)
+            for line in lines:
+                line = line.strip()
+                if not line:
+                    continue
+                src = WireTrace.model_validate(json.loads(line))
+                points = resume_points(src, kinds=kinds)
+                if points:
+                    return src, rng.choice(points)
+        return None
+
+
+__all__ = ["ReplayHarness", "ReplayHarnessConfig"]

environments/replay/replay/selector.py

@@ -0,0 +1,100 @@
+"""Consumer-side buffer reading + resume-point selection (Option B: ``MessageNode`` fields).
+
+Turns the compaction tags the producing harness stamped on ``MessageNode.kind`` into resume
+points the ReplayTaskset/ReplayHarness sample. ``build_seed`` produces the seed conversation per
+mode; ``snapshot_ref_of`` returns the durable sandbox handle to restore (None until per-turn
+snapshot capture is wired). Only ``get_tag``/``snapshot_ref_of`` differ from Option A — everything
+else (buffer reading, resume-point logic, seed building) is shared.
+"""
+
+from __future__ import annotations
+
+import glob
+import json
+
+from verifiers.v1 import graph
+from verifiers.v1.trace import Trace, WireTrace
+from verifiers.v1.types import Messages, UserMessage
+
+DEFAULT_KINDS = ["recheck", "compaction_after", "compaction_before"]
+DEFAULT_FOLLOWUP = "Check your work. If anything is wrong, fix it and give the corrected final answer."
+
+
+def get_tag(trace: Trace, node_id: int) -> str | None:
+    """Read a node's replay tag (Option B: from the typed ``MessageNode.kind`` field)."""
+    return trace.nodes[node_id].kind
+
+
+def snapshot_ref_of(trace: Trace, node_id: int) -> str | None:
+    """Durable sandbox snapshot ref for a node (Option B: from ``MessageNode.snapshot_ref``).
+    None when snapshotting was off/unsupported (the skeleton never captures one yet)."""
+    return trace.nodes[node_id].snapshot_ref
+
+
+def iter_traces(buffer_glob: str):
+    """Yield each stored rollout (``WireTrace``) from the buffer glob, in file then line order."""
+    for path in sorted(glob.glob(buffer_glob)):
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    yield WireTrace.model_validate(json.loads(line))
+
+
+def seed_messages(trace: Trace, node_id: int) -> Messages:
+    """The replay prefix: messages along root->node_id, in order."""
+    path: Messages = []
+    nid: int | None = node_id
+    while nid is not None:
+        path.append(trace.nodes[nid].message)
+        nid = trace.nodes[nid].parent
+    path.reverse()
+    return path
+
+
+def resume_points(trace: Trace, *, kinds: set[str]) -> list[dict]:
+    """Resume points whose ``kind`` is in ``kinds``. ``compaction_before``/``compaction_after``
+    come from the harness tags; ``recheck`` and ``judge`` are the structural final-answer leaf
+    (re-roll vs. judge-the-attempt). Each: ``node`` id and ``kind``."""
+    points: list[dict] = [
+        {"node": nid, "kind": get_tag(trace, nid)}
+        for nid in range(len(trace.nodes))
+        if get_tag(trace, nid) in kinds
+    ]
+    leaves = graph.leaves(trace)
+    if leaves:
+        final = max(leaves)  # the rollout's final-answer leaf
+        points += [{"node": final, "kind": k} for k in ("recheck", "judge") if k in kinds]
+    return points
+
+
+def render_transcript(trace: Trace, node_id: int) -> str:
+    """The conversation along root->node_id as plain text, for a judge prompt."""
+    lines = []
+    for m in seed_messages(trace, node_id):
+        content = m.content if isinstance(m.content, str) else (m.content or "")
+        lines.append(f"{m.role}: {content}")
+    return "\n".join(lines)
+
+
+def judge_prompt(trace: Trace, node_id: int) -> str:
+    """A 'was this attempt correct?' prompt presenting the rollout's transcript."""
+    return (
+        "You are judging whether a previous attempt solved its task correctly.\n\n"
+        f"--- transcript ---\n{render_transcript(trace, node_id)}\n--- end transcript ---\n\n"
+        "Was the final answer correct? Reply with exactly 'yes' or 'no'."
+    )
+
+
+def build_seed(trace: Trace, point: dict, followup: str) -> Messages:
+    """The seed conversation for a resume point:
+    - ``judge``   -> a single user turn presenting the rollout to be graded;
+    - ``recheck`` -> the full rollout prefix + an appended check-your-work user turn;
+    - ``compaction_*`` -> the plain ``root->node`` prefix.
+    """
+    if point["kind"] == "judge":
+        return [UserMessage(content=judge_prompt(trace, point["node"]))]
+    msgs = seed_messages(trace, point["node"])
+    if point["kind"] == "recheck":
+        msgs = [*msgs, UserMessage(content=followup)]
+    return msgs