warn on ambiguous prefix match in interleave_rollout

When more than one active prefix matches a step's prompt, log a warning
with the example id, step index, set of matching prefix lengths, total
active prefixes, and the prompt length. Longest-match still picks the
correct extension; the warning just surfaces the rare ambiguous case so
it's debuggable if it starts showing up in real rollouts (e.g. from
compaction/rollback turns).

Co-authored-by: Cursor <cursoragent@cursor.com>
(cherry picked from commit ca38614)

Author: 2 people

src/prime_rl/orchestrator/trajectories.py

@@ -413,11 +413,26 @@ def extend_sample(
         # silently absorb the longer sample's generated tokens as user input.
         matched_idx = None
         matched_len = -1
+        matching_prefix_lens: list[int] = []
         for idx, (prefix_tokens, _, _) in enumerate(active_samples):
             pl = len(prefix_tokens)
-            if pl > matched_len and step_prompt_ids[:pl] == prefix_tokens:
-                matched_idx = idx
-                matched_len = pl
+            if step_prompt_ids[:pl] == prefix_tokens:
+                matching_prefix_lens.append(pl)
+                if pl > matched_len:
+                    matched_idx = idx
+                    matched_len = pl
+
+        if len(matching_prefix_lens) > 1:
+            # Ambiguous extension: rare, but reachable via compaction/rollback
+            # where a new sample's prefix happens to start with an older
+            # sample's prefix. Longest-match is the correct choice; surface
+            # the ambiguity so we can audit if it shows up in real rollouts.
+            logger.warning(
+                f"Ambiguous prefix match at step {step_idx} for example {output['example_id']}: "
+                f"{len(matching_prefix_lens)} of {len(active_samples)} active prefixes match "
+                f"(lens={sorted(matching_prefix_lens)}, step_prompt_len={len(step_prompt_ids)}). "
+                f"Extending the longest (len={matched_len})."
+            )
 
         if matched_idx is not None:
             # Extension holds - merge into matched sample