feat(v1): replay taskset — re-enter saved rollouts as continue/recheck tasks

[faresobeid]

What

A built-in v1 taskset, replay, that re-enters saved rollout records (Trace.to_record() JSONL lines) as fresh tasks — plus the small framework seams a wrapping taskset needs. This is the verifiers half of prime-rl's Continue/Recheck training (prime-rl PR to follow); everything is driven by env config, no orchestrator changes.

taskset = { id = "replay", records = "<glob>", mode = "continue" | "recheck", source = { id = "<original taskset>", ... } }

Each replayed task is the original typed task with its prompt swapped for the seed (model_copy), and every hook — tools, user, setup/finalize, score/score_group, stops, state — delegates to the source taskset loaded by id. New completions are therefore judged by the original env's own verifier, with its real toolkit.

Modes

• continue, anchor = "compaction" (default): one task per recorded context restart. Detection is structural, not textual — a restart is a branch that forks off shared history on unsampled, system/user-only launch context before its first sampled node — so records from different harnesses (each with its own compaction wording/shape) mix freely. The canonical [system, user(summary)] restart lowers to a plain string prompt, so it runs under any harness, including agent CLIs like rlm (a fresh launch reproduces the post-compaction context exactly). A guard drops candidates equal to the task's own launch context (branch order follows leaf order, so an auxiliary same-system branch can sort first), and retokenization splits are excluded because their forked context carries assistant/tool copies.
• continue, anchor = "tool-call": one deterministically-drawn resume point per rollout (Random(f"{seed}:{trace.id}")), right after a complete tool-result run (every tool_call id of the issuing assistant answered). Messages prompt → needs a message-seeding harness (default/null).
• recheck: the attempt's final branch, truncation artifacts stripped (everything from the first partially-answered tool run; trailing unanswered user turns), plus a configurable verification turn. last_reply-style rewards score the revision — seeded context is never sampled.

Other recycling schemes subclass ReplayTaskset and override seeds(trace); record loading, delegation, and snapshot plumbing stay inherited.

Sandbox snapshots (seam only)

No runtime can capture/restore snapshots yet, but the contract is in place so producers can start recording: trace.info["snapshots"] maps node index → opaque ref, keyed at the node a resume point re-enters at. A record carrying any refs offers only snapshotted anchors (an unsnapshotted mid-trajectory resume would be unfaithful once state is captured), and setup restores via restore_snapshot(), which fails loudly until support lands (endgame documented there: promote the ref to a base Task field alongside image/workdir).

Framework seams (shared surface — the part to review closely)

• Taskset.defines_tools / defines_user / needs_container / stops() / defines_group_rewards() / state_type() — instance-level seams replacing class-identity checks (type(ts).tools is not Taskset.tools) and direct discover_decorated calls in Environment.__init__, Rollout.run, EnvServer, Episode, and cli/validate. Defaults preserve today's behavior exactly; a wrapping taskset overrides them to report its source. Note defines_group_rewards is deliberately a method: hook discovery walks instances with inspect.getmembers, which evaluates properties — a discovery-running property recurses (pinned by a test).
• Harness.write_message_prompt() + shared INITIAL_MESSAGES_FILE: a Messages prompt now ships to the default/null programs as a file in the runtime instead of the INITIAL_MESSAGES env var — replayed histories are exactly the long contexts that overflow MAX_ARG_STRLEN (~128 KiB per env string).

Operational details

• ReplayConfig validation freezes records globs to the concrete matched files, so every pool worker (elastic pools spawn workers mid-run) builds the identical, positionally-indexed task list.
• Mixed-env record files: lines are filtered by validating the task payload against the source's task type (cheap, before full-trace validation) and counted as skipped. The filter is structural — sources whose task type has no distinguishing required fields should use single-env files (documented).
• max_seed_tokens drops over-long seeds at load (recorded token count when present, chars/4 estimate otherwise).

Verification

• tests/v1/test_replay_taskset.py (18 tests): seed construction across compaction shapes (rlm shape, kept-task-message shape, multi-compaction), retokenization/subagent/aux-branch exclusions, complete-vs-partial tool runs, recheck cleanup, snapshot gating, config narrowing + dump→re-validate round trip (the pool-worker path), glob freezing, capability/state delegation against the echo fixtures.
• Full local suite: tests/v1/ + tests/test_imports.py + tests/test_environment.py + tests/test_env_server.py → 87 passed (excluding the pre-existing wiki_search_compact.toml env-dependent failure, which fails identically on the base commit).
• End-to-end smoke: synthesized records → load_tasks in all three modes → correct prompts/typed task fields; Environment construction under the rlm harness passes the capability gates.

🤖 Generated with Claude Code

Update: online self-replay (records are always followed)

replay can now follow a growing records source — the current run's own rollouts/step_*/train_rollouts.jsonl — instead of requiring a finished prior run. There is no freeze/knob — load_tasks is always incremental + append-only (numeric-aware step ordering, only new files parsed, indices never reassigned); the env server refreshes through a new Taskset.reload_tasks() seam on info and on out-of-range task indices, so pool workers refreshing at different times stay index-aligned. The replay env's own output (continue:/recheck: task names) is skipped so seeds don't compound. Covered by a follow-mode test (append across reloads, step_9/step_10 ordering, self-replay guard).

Update: max_anchors + source-reward filters + per-env records

max_anchors: int | None = 1 on continue/tool-call draws that many resume points per source rollout (None = every complete tool-run boundary becomes a task). min/max_source_reward seed only from selected-outcome rollouts (e.g. recheck only incorrect attempts). prime-rl now writes per-env train record files (train_rollouts_.jsonl), so mixed-file filtering is a fallback rather than the mechanism. Validated by three live gsm8k debug runs (recheck, tool-call continue, incorrect-only recheck: repair rate 31%→37% over 100 steps).

Update: rollout record index

Record files can now ship with a derived sibling index (prime-rl writes index_<env>.jsonl beside each train_rollouts_<env>.jsonl): one row per record line with selection fields (task_name, reward, branches, ...) and the line's byte span. New pure helpers in records.py — read_index / select_index_rows / iter_indexed_records — let load_tasks filter records index-side (already-replayed names, out-of-range source rewards, and un-compacted rollouts for continue/compaction) and then seek/read/parse only the surviving spans, opening the record file once. A file without an index full-parses exactly as before, and the parsed-record checks (foreign task validation, empty nodes) still run on every survivor. Covered by 3 new tests: index-selected loading matches full parse, filtered records are provably never parsed (invalid JSON at their spans), and branch-count prefiltering for compaction; plus explicit missing-index fallback asserts.

Update: cleanup pass

Consolidated the index API: read_index / select_index_rows / iter_indexed_records collapsed into one iter_selected_records generator (index-side filtering + span-selective reads, full parse when no index exists), with index_path / index_row kept as the small writer-side contract (prime-rl builds its index rows through them). Shared REPLAYED_PREFIXES constant for both replayed-name filters, max_anchors bounds via Annotated[int, Field(ge=1)] | None, and tool_call_seeds now uses per-branch token-count prefix sums (O(1) estimates per anchor) plus a Counter ordinal for name disambiguation. Behavior unchanged; tests reorganized around the shared helpers (one redundant reward-filter test folded into the index-vs-full-parse test).

[macroscopeapp]

🟡 Medium replay/taskset.py:215

In follow mode, load_tasks permanently ignores records appended to a JSONL file it has already read. It tracks only file paths in _files_done and excludes any previously seen path from new_paths, so when rollouts/step_N/train_rollouts.jsonl grows after the first read, the new lines are never parsed — online self-replay silently drops tasks from the current step. Consider tracking a per-file byte offset (or line count) and re-reading from that position for files already seen, instead of skipping them outright.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/tasksets/replay/taskset.py around line 215:

In `follow` mode, `load_tasks` permanently ignores records appended to a JSONL file it has already read. It tracks only file paths in `_files_done` and excludes any previously seen path from `new_paths`, so when `rollouts/step_N/train_rollouts.jsonl` grows after the first read, the new lines are never parsed — online self-replay silently drops tasks from the current step. Consider tracking a per-file byte offset (or line count) and re-reading from that position for files already seen, instead of skipping them outright.

[macroscopeapp]

🟠 High replay/taskset.py:146

ReplayConfig._validate documents that it freezes the records globs to the matched files so every pool worker sees the same file set, but it never actually replaces self.records with the matched list — the method just returns self unchanged. As a result, load_tasks re-expands the globs on every call. A worker started later, after new record files have appeared, sees extra files and assigns different tasks to the same numeric idx values than an earlier worker. Requests routed by task index then execute or score the wrong replay task across pool workers. The validator should replace self.records with the resolved file list before returning, or freeze the expansion in a separate field that load_tasks reads instead of the raw globs.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/tasksets/replay/taskset.py around line 146:

`ReplayConfig._validate` documents that it freezes the `records` globs to the matched files so every pool worker sees the same file set, but it never actually replaces `self.records` with the matched list — the method just returns `self` unchanged. As a result, `load_tasks` re-expands the globs on every call. A worker started later, after new record files have appeared, sees extra files and assigns different tasks to the same numeric `idx` values than an earlier worker. Requests routed by task index then execute or score the wrong replay task across pool workers. The validator should replace `self.records` with the resolved file list before returning, or freeze the expansion in a separate field that `load_tasks` reads instead of the raw globs.

[macroscopeapp]

🟡 Medium serve/server.py:126

_refresh_tasks() replaces self.tasks with a new list, but self.serving() was already entered at server startup with the original list object. When the taskset grows (e.g. starts empty and gains tasks later), Environment.shared_tools keeps the stale list reference, so newly discovered tasks never get their shared tool servers. Replay tasksets that rely on shared MCP tools will run those later tasks without the intended shared resources, changing rollout behavior. Consider re-entering serving() (or rebuilding shared tools) when the task list actually changes in _refresh_tasks(), or otherwise ensuring shared serving resources track the refreshed task list.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/serve/server.py around line 126:

`_refresh_tasks()` replaces `self.tasks` with a new list, but `self.serving()` was already entered at server startup with the original list object. When the taskset grows (e.g. starts empty and gains tasks later), `Environment.shared_tools` keeps the stale list reference, so newly discovered tasks never get their shared tool servers. Replay tasksets that rely on `shared` MCP tools will run those later tasks without the intended shared resources, changing rollout behavior. Consider re-entering `serving()` (or rebuilding shared tools) when the task list actually changes in `_refresh_tasks()`, or otherwise ensuring shared serving resources track the refreshed task list.

[macroscopeapp]

🟡 Medium replay/taskset.py:212

load_tasks silently discards any source record whose task name begins with "continue:" or "recheck:", treating it as this taskset's own recycled output. Because Task.name is arbitrary user data from the source environment, a legitimate source task named e.g. "recheck: algebra" or "continue: step 2" is dropped and can never be replayed, even though it is valid source data. The name-prefix heuristic is an unreliable way to detect self-output. Consider tagging replayed output with a dedicated marker field rather than overloading the user-facing name.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/tasksets/replay/taskset.py around line 212:

`load_tasks` silently discards any source record whose task name begins with `"continue:"` or `"recheck:"`, treating it as this taskset's own recycled output. Because `Task.name` is arbitrary user data from the source environment, a legitimate source task named e.g. `"recheck: algebra"` or `"continue: step 2"` is dropped and can never be replayed, even though it is valid source data. The name-prefix heuristic is an unreliable way to detect self-output. Consider tagging replayed output with a dedicated marker field rather than overloading the user-facing `name`.

[macroscopeapp]

🟡 Medium replay/taskset.py:259

load_tasks() only updates _files_done at the very end of the batch. If any later file/line in new_paths raises (e.g. malformed JSON in iter_records(), or trace_cls.model_validate(record) rejects a bad trace), every earlier valid file in that batch is left unmarked. The next reload_tasks() reparses those already-consumed files and appends duplicate tasks, so task indices stop being stable across pool workers. Marking each file done as it's parsed prevents the duplicates.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/tasksets/replay/taskset.py around line 259:

`load_tasks()` only updates `_files_done` at the very end of the batch. If any later file/line in `new_paths` raises (e.g. malformed JSON in `iter_records()`, or `trace_cls.model_validate(record)` rejects a bad trace), every earlier valid file in that batch is left unmarked. The next `reload_tasks()` reparses those already-consumed files and appends duplicate tasks, so task indices stop being stable across pool workers. Marking each file done as it's parsed prevents the duplicates.

[macroscopeapp]

🟡 Medium replay/taskset.py:202

load_tasks() validates each record as Trace[task_cls] (the base Trace type), which forbids unknown fields. Source tasksets can subclass Trace to persist extra trace fields, and their rollout records contain those extra keys — so trace_cls.model_validate(record) raises ValidationError and the record is skipped. Replay therefore drops all records from any taskset that uses a custom Trace subclass. Use the source taskset's own Trace subclass type (obtainable via self.source.trace_type() or equivalent) for validation instead of the generic base type.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/tasksets/replay/taskset.py around line 202:

`load_tasks()` validates each record as `Trace[task_cls]` (the base `Trace` type), which forbids unknown fields. Source tasksets can subclass `Trace` to persist extra trace fields, and their rollout records contain those extra keys — so `trace_cls.model_validate(record)` raises `ValidationError` and the record is skipped. Replay therefore drops all records from any taskset that uses a custom `Trace` subclass. Use the source taskset's own `Trace` subclass type (obtainable via `self.source.trace_type()` or equivalent) for validation instead of the generic base type.

[macroscopeapp]

🟠 High replay/taskset.py:270

load_tasks() overwrites every replayed task's prompt with seed.prompt, even when the source task's original prompt was None. The server uses task.prompt is None as the signal to let the user() simulator open the conversation on the initial respond("") call, so replay tasks lose that opening turn. In environments like alphabet_sort_v1, this leaves the simulator queue one turn behind and injects the initial prompt again after the first assistant reply, corrupting the replayed conversation and scoring the wrong trajectory.

This happens at trace.task.model_copy(update={... "prompt": seed.prompt ...}) — it unconditionally replaces the original prompt. Consider preserving the None semantics when the source task had no prompt, or signaling the replay-seeded context differently so the server's opening-turn detection is not disrupted.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/tasksets/replay/taskset.py around line 270:

`load_tasks()` overwrites every replayed task's `prompt` with `seed.prompt`, even when the source task's original `prompt` was `None`. The server uses `task.prompt is None` as the signal to let the `user()` simulator open the conversation on the initial `respond("")` call, so replay tasks lose that opening turn. In environments like `alphabet_sort_v1`, this leaves the simulator queue one turn behind and injects the initial prompt again after the first assistant reply, corrupting the replayed conversation and scoring the wrong trajectory.

This happens at `trace.task.model_copy(update={... "prompt": seed.prompt ...})` — it unconditionally replaces the original prompt. Consider preserving the `None` semantics when the source task had no prompt, or signaling the replay-seeded context differently so the server's opening-turn detection is not disrupted.

[macroscopeapp]

🟡 Medium replay/records.py:298

tool_call_seeds always strips the leading system message from the seeded history. When the source task's system prompt lives in-band in Task.prompt (as Environment builds by prepending it into the message list instead of populating Task.system_prompt), the seed loses it entirely — the harness re-emits Task.system_prompt, which is empty, so the replayed rollout resumes with different instructions than the original. Consider preserving the system message in the seed so the replayed context matches what the original rollout actually had.

     for nodes, end, tokens in candidates:
         prefix = nodes[: end + 1]
         messages = [node.message for node in prefix]
-        if messages[0].role == "system":
-            messages = messages[1:]
         base = f"continue:{trace.id[:8]}:tool-call{end}"

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/tasksets/replay/records.py around lines 298-301:

`tool_call_seeds` always strips the leading system message from the seeded history. When the source task's system prompt lives in-band in `Task.prompt` (as `Environment` builds by prepending it into the message list instead of populating `Task.system_prompt`), the seed loses it entirely — the harness re-emits `Task.system_prompt`, which is empty, so the replayed rollout resumes with different instructions than the original. Consider preserving the system message in the seed so the replayed context matches what the original rollout actually had.

[macroscopeapp]

🟡 Medium replay/records.py:274

tool_call_seeds can produce replay seeds from fabricated context that was never part of the original rollout. When a branch carries copied tool messages (e.g. from a retokenization split or compaction), those nodes have sampled=False, but the function only checks node.message.role == "tool" and never checks node.sampled or issuer.sampled. This creates continue:...:tool-call seeds from context the source model never produced, polluting the replay taskset with invalid resume points. Consider filtering candidates so the issuing assistant node and tool result nodes must all have sampled=True.

-        for end, node in enumerate(nodes):
-            if node.message.role != "tool" or id(node) in seen:
+        for end, node in enumerate(nodes):
+            if node.message.role != "tool" or not node.sampled or id(node) in seen:

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/tasksets/replay/records.py around lines 274-276:

`tool_call_seeds` can produce replay seeds from fabricated context that was never part of the original rollout. When a branch carries copied `tool` messages (e.g. from a retokenization split or compaction), those nodes have `sampled=False`, but the function only checks `node.message.role == "tool"` and never checks `node.sampled` or `issuer.sampled`. This creates `continue:...:tool-call` seeds from context the source model never produced, polluting the replay taskset with invalid resume points. Consider filtering candidates so the issuing assistant node and tool result nodes must all have `sampled=True`.