Wave 21: close both Wave 20 debt items — chat-template alignment + structural is_error

Two architectural-debt items flagged "Wave 20 candidate" in the framework doc,
both CPU-only, both now fixed and regression-guarded.

## 1. SDPO mask chat-template drift (was ~67% aligned, now ~100%)

`ComposerDataCollator._build_segment_mask` built the sdpo_loss_mask (and the
aligned-student response_mask) by tokenizing each content segment in isolation
and concatenating. That ignored the scaffolding tokens apply_chat_template
inserts around every message (<|im_start|>{role}\n ... <|im_end|>\n, BOS), so
mask bits drifted left of the real content tokens — the residual ~33%
contamination documented in the Wave 19 production audit.

Fix: new `_build_chat_aligned_mask` derives the mask from per-message
apply_chat_template prefix deltas. For message k it computes the token span
len(template(msgs[:k+1])) - len(template(msgs[:k])), then locates the content
run inside that span by subsequence match and marks ONLY those positions as
loss. Falls back to whole-span marking if the content run can't be located
(tokenizer merge across the boundary) so SDPO signal is never silently dropped.
Degenerates exactly to the old concat behavior on stub tokenizers (no template),
so the 15 stub collator tests stay green.

Validated against the real Qwen2.5-0.5B-Instruct chat template: alignment ratio
67% -> 100%, and the in-loss tokens decode to exactly the recovery turn's
content with zero <|im_start|>/<|im_end|> leakage.

## 2. TOOL_ERROR_TAG string-coupling

ClaudeCodeIngester wrote the literal "[TOOL_RESULT (ERROR)]" string for
is_error:true tool_results, then the trace_examples adapter grepped that same
string back out to detect SDPO error sites. Brittle: any serialization drift
silently darkened the SDPO channel.

Fix: ingester now surfaces a structural `tool_error: True` boolean on user
messages (the is_error bool was already known, just discarded into the tag).
The adapter's new `_user_turn_has_error` reads the boolean first and falls back
to the string tag only when the structural flag is absent (backward-compat for
old traces / third-party producers). Structural flag wins both ways — a
producer can set tool_error:False to suppress a tag that appears in quoted text.

## Tests
- composer_replication/trainer/tests/ (new): 3 real-tokenizer alignment guards
(skip cleanly when transformers/model cache absent)
- test_trace_examples_adapter.py: +5 structural-flag tests incl. drift-resilience
(tag absent but flag present -> still detected) and inverse (flag False wins)
- Full package suite: 146 passed, 16 skipped, 0 failed

## Docs
- production SDPO example audit note updated: drift is fixed, sub-95% now means
a NEW regression rather than a known-residual bug

Also adds uv.lock (uv-pinned deps; pyproject already present, lock was untracked).

composer_replication/ingestion/claude_code.py

@@ -135,9 +135,16 @@ class ClaudeCodeIngester:
                     # Either text blocks (a real human prompt) or tool_result
                     # blocks (an observation). Both go into history as user
                     # messages, but we serialize them differently.
-                    flat = self._flatten_user_content(content)
+                    flat, had_tool_error = self._flatten_user_content(content)
                     if flat:
-                        history.append({"role": "user", "content": flat})
+                        user_msg: dict[str, Any] = {"role": "user", "content": flat}
+                        if had_tool_error:
+                            # Structural error flag — the SDPO source of truth.
+                            # The [TOOL_RESULT (ERROR)] string tag still lives in
+                            # `content` for readability, but downstream detection
+                            # reads THIS boolean (see trace_examples adapter).
+                            user_msg["tool_error"] = True
+                        history.append(user_msg)
 
             elif rec_type == "assistant":
                 msg = rec.get("message", {})
@@ -215,9 +222,20 @@ class ClaudeCodeIngester:
                     logger.debug("Truncated/malformed line in %s: %s", path, e)
                     continue
 
-    def _flatten_user_content(self, content: list[Any]) -> str:
-        """Convert a user record's content list to a single string."""
+    def _flatten_user_content(self, content: list[Any]) -> tuple[str, bool]:
+        """Convert a user record's content list to a single string.
+
+        Returns ``(flattened_text, had_tool_error)`` where ``had_tool_error``
+        is True iff any ``tool_result`` block in this user content carried
+        ``is_error: true``. The boolean is the STRUCTURAL source of truth for
+        SDPO error-site detection; the ``[TOOL_RESULT (ERROR)]`` string tag in
+        the text is kept only for human-readability and ``apply_chat_template``
+        rendering. Downstream consumers (the trace_examples adapter) should
+        read the structural flag, never grep the tag — see Wave 20 design note
+        on TOOL_ERROR_TAG string-coupling debt.
+        """
         parts: list[str] = []
+        had_tool_error = False
         for block in content:
             if not isinstance(block, dict):
                 continue
@@ -237,11 +255,13 @@ class ClaudeCodeIngester:
                     tc = "\n".join(sub)
                 tu_id = block.get("tool_use_id", "<unknown>")
                 is_err = block.get("is_error", False)
+                if is_err:
+                    had_tool_error = True
                 tag = "[TOOL_RESULT (ERROR)]" if is_err else "[TOOL_RESULT]"
                 parts.append(f"{tag} (id={tu_id})\n{tc}")
             elif bt == "image":
                 parts.append("[IMAGE OMITTED]")
-        return "\n\n".join(parts)
+        return "\n\n".join(parts), had_tool_error
 
     def _serialize_assistant_content(
         self, content: list[Any], *, strip_thinking: bool,

composer_replication/ingestion/tests/test_trace_examples_adapter.py

@@ -174,6 +174,126 @@ def test_tool_error_tag_matches_ingester_output():
     )
 
 
+# ----------------------------------------------------------------------
+# Structural error flag (Wave 20 — eliminate TOOL_ERROR_TAG coupling)
+# ----------------------------------------------------------------------
+
+
+def test_ingester_sets_structural_tool_error_flag():
+    """The ingester must set a STRUCTURAL `tool_error: True` boolean on
+    user messages whose source JSONL had `is_error: true`, independent of
+    the rendered string tag."""
+    ingester = ClaudeCodeIngester(skip_sidechain=True, strip_thinking=True)
+    states = list(ingester.ingest(ERROR_FIXTURE))
+    flagged = [
+        m for s in states for m in s["messages"]
+        if m.get("role") == "user" and m.get("tool_error") is True
+    ]
+    assert flagged, (
+        "Expected ≥1 user message with structural tool_error=True flag; "
+        "the ingester is not surfacing is_error structurally."
+    )
+    # And every structurally-flagged message must also render the tag
+    # (the tag is kept for readability — both should co-occur on the fixture).
+    for m in flagged:
+        assert TOOL_ERROR_TAG in m["content"], (
+            "Structural flag set but string tag missing — the two views "
+            "of the same error have diverged within the ingester."
+        )
+
+
+def test_clean_fixture_has_no_structural_flag():
+    """No user message on the clean fixture should carry tool_error=True."""
+    ingester = ClaudeCodeIngester(skip_sidechain=True, strip_thinking=True)
+    states = list(ingester.ingest(OK_FIXTURE))
+    flagged = [
+        m for s in states for m in s["messages"]
+        if m.get("role") == "user" and m.get("tool_error")
+    ]
+    assert not flagged, f"Clean fixture should have 0 structural flags; got {len(flagged)}"
+
+
+def test_structural_flag_survives_tag_drift():
+    """THE drift-resilience guarantee: if the rendered string tag drifts
+    (e.g. a future serialization change strips or renames it) but the
+    structural `tool_error: True` flag is present, the adapter MUST still
+    detect the error site. This is the entire point of the Wave 20 fix —
+    detection no longer depends on grepping a human-readable string."""
+    # Hand-build a state where the tag is ABSENT from content but the
+    # structural flag is set — simulating ingester serialization drift.
+    states = [{
+        "state_id": "drift-0",
+        "messages": [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "run the build"},
+            {"role": "assistant", "content": "[TOOL_USE] name=Bash input={}"},
+            # Tag DELIBERATELY absent from content; only the structural flag.
+            {"role": "user", "content": "build failed: missing target",
+             "tool_error": True},
+            {"role": "assistant", "content": "Let me fix the target."},
+        ],
+    }]
+    examples = claude_states_to_trace_examples(states)
+    assert len(examples) == 1
+    err_turns = [t for t in examples[0]["turns"] if t.get("tool_error")]
+    assert len(err_turns) == 1, (
+        "Structural flag present but adapter failed to detect the error "
+        "site without the string tag — the coupling fix is broken."
+    )
+    # The recovery turn is the assistant immediately after the flagged user turn.
+    assert err_turns[0]["content"] == "Let me fix the target."
+
+
+def test_structural_false_suppresses_tag_match():
+    """Inverse drift case: a producer sets `tool_error: False` to assert
+    'this is NOT an error' even though the rendered content happens to
+    contain the tag string. The structural flag must WIN over the string."""
+    states = [{
+        "state_id": "false-0",
+        "messages": [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "look at this log"},
+            {"role": "assistant", "content": "[TOOL_USE] name=Read input={}"},
+            # Content contains the tag verbatim (e.g. quoting a prior log)
+            # but the producer asserts it's not a live error site.
+            {"role": "user",
+             "content": f"the docs mention {TOOL_ERROR_TAG} as an example",
+             "tool_error": False},
+            {"role": "assistant", "content": "I see, that's just documentation."},
+        ],
+    }]
+    examples = claude_states_to_trace_examples(states)
+    err_turns = [t for t in examples[0]["turns"] if t.get("tool_error")]
+    assert not err_turns, (
+        "tool_error=False should suppress detection even when the string "
+        "tag is present in content; structural flag must take precedence."
+    )
+
+
+def test_string_tag_fallback_when_no_structural_flag():
+    """Backward-compat: an OLD trace (no structural flag anywhere) with the
+    tag in content must STILL be detected via the string fallback path."""
+    states = [{
+        "state_id": "legacy-0",
+        "messages": [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "run it"},
+            {"role": "assistant", "content": "[TOOL_USE] name=Bash input={}"},
+            # No tool_error key at all — pure legacy serialization.
+            {"role": "user",
+             "content": f"{TOOL_ERROR_TAG} (id=x)\nno such file or directory"},
+            {"role": "assistant", "content": "Creating the file first."},
+        ],
+    }]
+    examples = claude_states_to_trace_examples(states)
+    err_turns = [t for t in examples[0]["turns"] if t.get("tool_error")]
+    assert len(err_turns) == 1, (
+        "Legacy trace without structural flag must fall back to the string "
+        "tag match; backward compatibility broken."
+    )
+    assert err_turns[0]["tool_error"] == "file_not_found"
+
+
 # ----------------------------------------------------------------------
 # Empty input
 # ----------------------------------------------------------------------

composer_replication/ingestion/trace_examples.py

@@ -88,6 +88,30 @@ def default_classify_error(content: str) -> str:
     return "tool_error"
 
 
+def _user_turn_has_error(msg: Mapping[str, Any], flat_content: str) -> bool:
+    """Decide whether a user-role turn is a tool-error site.
+
+    Precedence (Wave 20 — eliminate TOOL_ERROR_TAG string-coupling):
+
+    1. **Structural flag** — if the message dict carries an explicit
+       ``tool_error`` key, trust it as the source of truth. The ingester sets
+       ``tool_error: True`` whenever the source JSONL had ``is_error: true``.
+       A third-party producer can set ``tool_error: False`` to assert "no
+       error here" even if the rendered text happens to contain the tag.
+    2. **String-tag fallback** — only when no structural flag is present
+       (older serialized traces, or producers that never learned the boolean
+       contract) do we fall back to matching ``TOOL_ERROR_TAG`` in the
+       rendered content. This keeps backward compatibility without making the
+       brittle string match the primary path.
+
+    Returns True iff the turn should trigger SDPO error-site handling.
+    """
+    structural = msg.get("tool_error")
+    if structural is not None:
+        return bool(structural)
+    return TOOL_ERROR_TAG in flat_content
+
+
 # ---------------------------------------------------------------------------
 # Adapter
 # ---------------------------------------------------------------------------
@@ -163,7 +187,14 @@ def claude_states_to_trace_examples(
                             str(c.get("text", c)) if isinstance(c, dict) else str(c)
                             for c in prev_content
                         )
-                    if TOOL_ERROR_TAG in prev_content:
+                    # STRUCTURAL detection (Wave 20): the ingester sets a
+                    # `tool_error: True` boolean on user messages whose source
+                    # JSONL had `is_error: true`. This is the source of truth.
+                    # We fall back to string-matching the TOOL_ERROR_TAG only
+                    # for messages that lack the structural flag (older traces
+                    # or third-party producers that didn't set it) — see
+                    # `_user_turn_has_error`.
+                    if _user_turn_has_error(prev, prev_content):
                         error_kind_found = error_kind_fn(prev_content)
                         error_content_found = prev_content
                         break

composer_replication/trainer/data_collator.py

@@ -318,8 +318,18 @@ class ComposerDataCollator:
 
         # Tokenize the full teacher conversation
         teacher_ids = self._tokenize_messages(teacher_messages)
-        # Build the per-token loss mask by tokenizing each segment and concatenating
-        sdpo_mask = self._build_segment_mask(teacher_loss_segments)
+        # Build the per-token loss mask ALIGNED to the chat-template tokenization
+        # (Wave 20 fix). The old path tokenized each segment's raw text in
+        # isolation and concatenated; that ignored the scaffolding tokens
+        # (<|im_start|>{role}\n ... <|im_end|>\n, BOS, etc.) that
+        # apply_chat_template inserts, so mask positions drifted left of the
+        # real content tokens — the residual ~33% misalignment documented in
+        # the Wave 19 production audit. `_build_chat_aligned_mask` derives the
+        # mask from per-message apply_chat_template deltas instead, so loss
+        # bits land exactly on content tokens regardless of template markers.
+        sdpo_mask = self._build_chat_aligned_mask(
+            teacher_messages, teacher_loss_segments, teacher_ids
+        )
         # Truncate mask to teacher_ids length if tokenization round-tripped slightly differently
         sdpo_mask = sdpo_mask[: len(teacher_ids)]
         if len(sdpo_mask) < len(teacher_ids):
@@ -465,14 +475,17 @@ class ComposerDataCollator:
         # Tokenize the full student conversation via apply_chat_template
         # (mirrors teacher's path so chat-template markers are identical).
         student_ids = self._tokenize_messages(student_messages)
-        # Build response mask via the same segment-tokenization helper used
-        # for sdpo_mask, then reinterpret 1=in-response, 0=not-in-response.
-        # We can't reuse _build_segment_mask (which uses ignore_index for
-        # non-loss); inline a 0/1 variant.
-        resp_mask: list[int] = []
-        for is_resp, text in student_loss_segments:
-            seg_ids = self._tokenize_text(text)
-            resp_mask.extend([1 if is_resp else 0] * len(seg_ids))
+        # Build response mask ALIGNED to the chat-template tokenization (Wave 20
+        # fix — same drift bug as the teacher sdpo_mask path). We derive the
+        # mask from per-message apply_chat_template deltas so 1-bits land on
+        # the assistant content tokens exactly, not shifted by the template
+        # scaffolding. `_build_chat_aligned_mask` emits 1 for loss segments and
+        # ignore_index for the rest; we remap ignore_index -> 0 because the
+        # response_mask convention here is 1=in-response, 0=not.
+        raw_mask = self._build_chat_aligned_mask(
+            student_messages, student_loss_segments, student_ids
+        )
+        resp_mask = [1 if v == 1 else 0 for v in raw_mask]
         # Pad/truncate response_mask to student_ids length (same as teacher path).
         resp_mask = resp_mask[: len(student_ids)]
         if len(resp_mask) < len(student_ids):
@@ -486,6 +499,15 @@ class ComposerDataCollator:
         """For each (is_loss, text) segment, tokenize and emit per-token mask values.
 
         Loss-active tokens get 1; non-loss tokens get -100 (ignore_index).
+
+        NOTE (Wave 20): this naive per-segment concatenation IGNORES the
+        chat-template scaffolding that `apply_chat_template` inserts around
+        each message, so the resulting mask drifts out of alignment with a
+        sequence produced via `_tokenize_messages`. It is retained only for
+        the degenerate fallback inside `_build_chat_aligned_mask` and for
+        callers that build sequences via raw segment concatenation (no chat
+        template). The SDPO/response-mask paths now use
+        `_build_chat_aligned_mask` instead.
         """
         out: list[int] = []
         for is_loss, text in segments:
@@ -494,6 +516,94 @@ class ComposerDataCollator:
             out.extend([mask_value] * len(seg_ids))
         return out
 
+    @staticmethod
+    def _find_subseq(haystack: list[int], needle: list[int], start: int = 0) -> int:
+        """Return the index where ``needle`` first occurs in ``haystack`` at or
+        after ``start``, or -1 if absent. Linear scan (spans are short)."""
+        if not needle:
+            return start
+        n, m = len(haystack), len(needle)
+        for i in range(start, n - m + 1):
+            if haystack[i:i + m] == needle:
+                return i
+        return -1
+
+    def _build_chat_aligned_mask(
+        self,
+        messages: Sequence[dict],
+        segments: Sequence[tuple[bool, str]],
+        full_ids: list[int],
+    ) -> list[int]:
+        """Build a per-token loss mask aligned to a chat-template tokenization.
+
+        The caller builds ``messages`` and ``segments`` in lockstep — element
+        ``k`` of each describes the same logical chunk, where ``segments[k] =
+        (is_loss, content_text)`` and ``messages[k] = {role, content}``.
+
+        We need a mask over ``full_ids = apply_chat_template(messages)`` whose
+        1-bits sit exactly on the content tokens of loss segments. The hard
+        part is that ``apply_chat_template`` inserts role/BOS/EOS scaffolding
+        between and around messages, so the naive ``_build_segment_mask``
+        (which tokenizes each content string in isolation and concatenates)
+        drifts: its k-th block of mask bits lands at the wrong offset because
+        all the preceding scaffolding tokens are unaccounted for.
+
+        Algorithm — per-message prefix deltas:
+
+          prev_len = len(apply_chat_template(messages[:k]))
+          cur_len  = len(apply_chat_template(messages[:k+1]))
+          # message k occupies full_ids[prev_len : cur_len] (content + its
+          # own scaffolding). Locate the content token run inside that span
+          # by subsequence match against the content's standalone
+          # tokenization, mark THOSE positions with the segment value and
+          # leave the scaffolding as ignore_index.
+
+        Falls back gracefully:
+          * If the tokenizer has no usable chat template (stub / no template),
+            ``_tokenize_messages`` returns a plain concatenation and the prefix
+            deltas equal the raw content token counts — so the content
+            subsequence match is trivially the whole span and the result
+            matches ``_build_segment_mask`` exactly (stub tests stay green).
+          * If a content run can't be located inside its span (rare tokenizer
+            merges across the content/scaffolding boundary), we mark the whole
+            message span with the segment value when it is a loss segment, so
+            we never silently drop SDPO signal — we over-include by at most a
+            couple scaffolding tokens rather than misalign.
+        """
+        mask = [self.config.ignore_index] * len(full_ids)
+        prev_len = 0
+        search_from = 0
+        for k, msg in enumerate(messages):
+            prefix_ids = self._tokenize_messages(list(messages[: k + 1]))
+            cur_len = len(prefix_ids)
+            span_start, span_end = prev_len, cur_len
+            prev_len = cur_len
+            if span_end <= span_start:
+                continue
+            is_loss = segments[k][0] if k < len(segments) else False
+            content = segments[k][1] if k < len(segments) else msg.get("content", "")
+            if not is_loss:
+                search_from = span_end
+                continue
+            # Loss segment: mark only the content tokens within the span.
+            content_ids = self._tokenize_text(content)
+            # Search for the content run inside this message's span. Anchor the
+            # search at span_start so we don't match content from a later msg.
+            idx = self._find_subseq(full_ids[:span_end], content_ids, start=max(span_start, search_from))
+            if idx != -1 and idx >= span_start:
+                for p in range(idx, min(idx + len(content_ids), span_end)):
+                    mask[p] = 1
+                search_from = idx + len(content_ids)
+            else:
+                # Fallback: couldn't locate the content run (tokenizer merged
+                # the content/scaffolding boundary). Mark the whole span as
+                # loss rather than drop the SDPO signal entirely. Over-includes
+                # at most the message's own scaffolding tokens.
+                for p in range(span_start, span_end):
+                    mask[p] = 1
+                search_from = span_end
+        return mask
+
     # ----------------------------------------------------------------------
     # Channel 3: trace-replay DPO inputs
     # ----------------------------------------------------------------------

composer_replication/trainer/tests/test_chat_template_alignment.py

@@ -0,0 +1,146 @@
+"""Wave 20 — chat-template alignment regression guard for the PACKAGE collator.
+
+`composer_replication.trainer.data_collator.ComposerDataCollator` builds the
+SDPO `sdpo_loss_mask` (and the aligned-student `response_mask`) so that in-loss
+positions sit exactly on content tokens. The hard part is that
+`apply_chat_template` inserts role/BOS/EOS scaffolding around each message; the
+old `_build_segment_mask` tokenized each content string in isolation and
+concatenated, so the mask drifted left of the real content tokens. The Wave 19
+production audit measured this drift at ~67% aligned. Wave 20's
+`_build_chat_aligned_mask` derives the mask from per-message
+`apply_chat_template` prefix deltas instead, restoring ~100% alignment.
+
+These tests use a REAL chat-template tokenizer (the stub used by
+spikes/005 cannot expose the drift — its `apply_chat_template` adds no
+scaffolding). They skip cleanly when transformers / the model cache is absent.
+"""
+from __future__ import annotations
+
+import pytest
+
+from composer_replication.trainer.data_collator import (
+    CollatorConfig,
+    ComposerDataCollator,
+)
+
+
+def _load_real_chat_tokenizer():
+    """Return a real tokenizer with a chat template, or None to skip."""
+    try:
+        import os
+
+        os.environ.setdefault("HF_HUB_OFFLINE", "1")
+        os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
+        from transformers import AutoTokenizer
+    except Exception:
+        return None
+    for model in ("Qwen/Qwen2.5-0.5B-Instruct", "Qwen/Qwen2.5-1.5B-Instruct"):
+        try:
+            t = AutoTokenizer.from_pretrained(model)
+            if getattr(t, "chat_template", None):
+                return t
+        except Exception:
+            continue
+    return None
+
+
+_REAL_TOK = _load_real_chat_tokenizer()
+_SKIP_REASON = "real chat-template tokenizer not available (offline / not cached)"
+
+
+@pytest.fixture
+def real_chat_tok():
+    if _REAL_TOK is None:
+        pytest.skip(_SKIP_REASON)
+    return _REAL_TOK
+
+
+@pytest.fixture
+def multiturn_error_trace():
+    """Multi-turn trace with an error site after several turns, so the
+    chat-template scaffolding drift compounds (what exposed the old 33%)."""
+    return {
+        "trace_id": "real-align-1",
+        "turns": [
+            {"role": "user", "content": "Read /etc/app/config.yaml and summarize it."},
+            {"role": "assistant", "content": '[TOOL_USE] name=Read input={"path":"/etc/app/config.yaml"}'},
+            {"role": "user", "content": "[TOOL_RESULT (ERROR)] (id=t1)\nError: no such file or directory"},
+            {
+                "role": "assistant",
+                "content": "The file does not exist there. Let me search for it instead.",
+                "tool_error": "file_not_found",
+                "error_meta": {"source_role": "user"},
+            },
+            {"role": "user", "content": "[TOOL_RESULT] (id=t2)\nFound /opt/app/config.yaml"},
+            {"role": "assistant", "content": "Found it at /opt/app/config.yaml. Reading now."},
+        ],
+        "final_reward": 0.0,
+    }
+
+
+def _hint_gen(kind, _meta):
+    return f"The path was wrong (kind: {kind}). Search with Glob before reading."
+
+
+def test_real_chat_template_sdpo_mask_fully_aligned(real_chat_tok, multiturn_error_trace):
+    """THE Wave 20 guarantee: with a REAL chat template, every in-loss
+    sdpo_loss_mask position must have student==teacher token id. Before the
+    fix this drifted to ~67% because the mask was built from per-segment
+    tokenization that ignored apply_chat_template scaffolding."""
+    cfg = CollatorConfig(hint_generator=_hint_gen, enable_replay_dpo=False)
+    collator = ComposerDataCollator(tokenizer=real_chat_tok, config=cfg)
+    batch = collator([multiturn_error_trace])
+
+    assert "sdpo_loss_mask" in batch, "SDPO channel did not fire on the error trace"
+    s_in = batch["input_ids"]
+    t_in = batch["ctx_teacher_input_ids"]
+    m_in = batch["sdpo_loss_mask"]
+    assert s_in.shape == t_in.shape == m_in.shape
+
+    n_aligned = n_total = 0
+    for row in range(s_in.shape[0]):
+        in_loss = m_in[row] == 1
+        if int(in_loss.sum()) == 0:
+            continue
+        s_at = s_in[row][in_loss]
+        t_at = t_in[row][in_loss]
+        n_aligned += int((s_at == t_at).sum().item())
+        n_total += int(in_loss.sum().item())
+
+    assert n_total > 0, "No in-loss positions — SDPO mask is empty"
+    ratio = n_aligned / n_total
+    assert ratio >= 0.95, (
+        f"SDPO mask alignment is only {100 * ratio:.1f}% ({n_aligned}/{n_total}); "
+        f"the chat-template drift fix has regressed. Expected ~100%."
+    )
+
+
+def test_real_chat_template_in_loss_tokens_are_content_not_scaffolding(
+    real_chat_tok, multiturn_error_trace
+):
+    """The in-loss teacher tokens must decode to the recovery turn's CONTENT,
+    not chat-template markers (<|im_start|>, role strings, etc.)."""
+    cfg = CollatorConfig(hint_generator=_hint_gen, enable_replay_dpo=False)
+    collator = ComposerDataCollator(tokenizer=real_chat_tok, config=cfg)
+    batch = collator([multiturn_error_trace])
+
+    t_in = batch["ctx_teacher_input_ids"][0]
+    m_in = batch["sdpo_loss_mask"][0]
+    in_loss = m_in == 1
+    decoded = real_chat_tok.decode(t_in[in_loss].tolist())
+    assert "does not exist" in decoded, (
+        f"In-loss tokens don't contain the recovery content; got: {decoded!r}"
+    )
+    for marker in ("<|im_start|>", "<|im_end|>", "<|endoftext|>"):
+        assert marker not in decoded, (
+            f"Chat-template marker {marker!r} leaked into the in-loss span: {decoded!r}"
+        )
+
+
+def test_real_chat_template_student_teacher_shapes_match(real_chat_tok, multiturn_error_trace):
+    """The SDPO gate requires student_logits.shape == teacher_logits.shape;
+    verify the aligned-student path produces matching sequence lengths."""
+    cfg = CollatorConfig(hint_generator=_hint_gen, enable_replay_dpo=False)
+    collator = ComposerDataCollator(tokenizer=real_chat_tok, config=cfg)
+    batch = collator([multiturn_error_trace])
+    assert batch["input_ids"].shape == batch["ctx_teacher_input_ids"].shape

examples/sdpo_with_real_traces_production/run.py

@@ -308,16 +308,22 @@ def main() -> int:
         ratio = n_aligned / n_total_in_loss
         log.info("  alignment audit: %d / %d in-loss positions match student==teacher (%.1f%%)",
                  n_aligned, n_total_in_loss, 100 * ratio)
-        if ratio < 1.0:
+        if ratio < 0.95:
             log.warning(
                 "  NOTE: %d positions (%.1f%%) of the SDPO mask cover non-aligned "
-                "tokens. This is a residual segment-vs-chat-template drift bug "
-                "in the existing _build_segment_mask: the segment-tokenizer "
-                "doesn't account for chat-template markers added by "
-                "apply_chat_template. Tracked for Wave 20.",
+                "tokens. As of Wave 20 the chat-template drift was fixed via "
+                "ComposerDataCollator._build_chat_aligned_mask (per-message "
+                "apply_chat_template prefix deltas). A ratio below ~100%% now "
+                "indicates a NEW regression — investigate the collator, not a "
+                "known-residual bug.",
                 n_total_in_loss - n_aligned,
                 100 * (1 - ratio),
             )
+        else:
+            log.info(
+                "  ✓ Wave 20 chat-template alignment holding (%.1f%% — was ~67%% "
+                "before the _build_chat_aligned_mask fix).", 100 * ratio,
+            )
 
     log.info("=" * 64)
     log.info("Summary")