Preserve proved boundaries in mixed-stop captured histories

src/art/trajectories/_tokenize.py

@@ -4046,11 +4046,9 @@ def _tokenize_exact_projected_chat_history(
         return None
     final_key = _sampled_source_key(final_source)
     final_stop_reason = _source_stop_evidence(final_source, final_key)[0]
-    terminal_boundary = (
-        (length_stop_boundaries or {}).get(final_key)
-        if final_stop_reason == "length"
-        else None
-    )
+    # A terminal synthetic stop can accompany an earlier length-stop boundary;
+    # neither tail is sampled, and both must retain their renderer proof.
+    terminal_boundary = (length_stop_boundaries or {}).get(final_key)
     # Unlike a nonterminal truncation, the final sampled output needs no
     # synthetic boundary to connect it to a later prompt.
     if terminal_boundary is not None and not terminal_boundary.tail:
@@ -4065,7 +4063,17 @@ def _tokenize_exact_projected_chat_history(
     )
     terminal_flags = (
         [
-            *([TokenFlag.ASSISTANT] * len(terminal_boundary.tail)),
+            *(
+                [
+                    TokenFlag.ASSISTANT
+                    | (
+                        TokenFlag.OUTPUT
+                        if final_stop_reason == "stop"
+                        else TokenFlag(0)
+                    )
+                ]
+                * len(terminal_boundary.tail)
+            ),
             *([TokenFlag(0)] * len(terminal_boundary.following)),
         ]
         if terminal_boundary is not None
@@ -4097,7 +4105,11 @@ def _tokenize_exact_projected_chat_history(
     if terminal_boundary is not None:
         flags[
             len(final_prompt) + len(final_output) + len(terminal_boundary.tail) - 1
-        ] = TokenFlag.STOP
+        ] = TokenFlag.STOP | (
+            TokenFlag.ASSISTANT | TokenFlag.OUTPUT
+            if final_stop_reason == "stop"
+            else TokenFlag(0)
+        )
     source_keys: list[_SampledSourceKey | None] = [
         *([None] * len(final_prompt)),
         *([final_key] * len(final_output)),
@@ -5214,7 +5226,6 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None:
         _projection_matches is True
         and chat_template is None
         and chat_template_kwargs is None
-        and not _history_needs_synthetic_stop(history, resolved_tokenizer)
     ):
         sampled_message_indices: list[int] = []
         seen_signatures: set[tuple[object, ...]] = set()
@@ -5243,15 +5254,43 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None:
             source = history.message_sources[message_index]
             assert source is not None
             source_key = _sampled_source_key(source)
-            if _source_stop_evidence(source, source_key)[0] != "length":
+            stop_reason = _source_stop_evidence(source, source_key)[0]
+            output = _source_output_tokens(source, source_key)
+            synthetic_stop = (
+                stop_reason == "stop"
+                and bool(_terminator_ids(resolved_tokenizer))
+                and output is not None
+                and not _sampled_stop_suffix(
+                    output,
+                    source=source,
+                    source_key=source_key,
+                    tokenizer=resolved_tokenizer,
+                )
+            )
+            if synthetic_stop and position + 1 < len(sampled_message_indices):
+                length_stop_boundaries_complete = False
+                break
+            if stop_reason != "length" and not synthetic_stop:
                 continue
-            length_stop_count += 1
+            length_stop_count += stop_reason == "length"
             bounds = (
                 direct_bounds[message_index]
                 if direct_bounds
                 else marked_bounds.get(message_index)
                 or probed_bounds.get(message_index)
             )
+            if synthetic_stop and bounds is not None:
+                # Tool part bounds can omit sampled closing markup. Prove the
+                # complete sampled prefix before appending only its remainder.
+                assert output is not None
+                rendered_start = bounds[0]
+                while rendered_start and assistant_mask[rendered_start - 1]:
+                    rendered_start -= 1
+                bounds = _prove_exact_length_stopped_assistant_prefix(
+                    locations(output, rendered_start),
+                    assistant_mask,
+                    expected_start=rendered_start,
+                )
             if position + 1 < len(sampled_message_indices) and source_matches_context(
                 source
             ):
@@ -5285,11 +5324,13 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None:
                     else marked_bounds.get(next_message_index)
                     or probed_bounds.get(next_message_index)
                 )
+                # Part bounds can start after sampled tool-call markup; the
+                # next generation boundary is the assistant span's start.
                 next_prompt_end = (
-                    next_bounds[0]
-                    if next_bounds is not None
-                    else _next_assistant_span_start(assistant_mask, after=bounds[1])
+                    _next_assistant_span_start(assistant_mask, after=bounds[1])
                     if bounds is not None
+                    else next_bounds[0]
+                    if next_bounds is not None
                     else None
                 )
             else:
@@ -5310,7 +5351,7 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None:
                 else None
             )
             if boundary is None:
-                if position + 1 < len(sampled_message_indices):
+                if synthetic_stop or position + 1 < len(sampled_message_indices):
                     length_stop_boundaries_complete = False
                     break
                 # A terminal length stop needs no renderer-owned tail: the

tests/unit/trajectories/test_tokenize.py

@@ -9,7 +9,7 @@
 import sys
 from time import perf_counter
 from types import ModuleType, SimpleNamespace
-from typing import Any, Never, cast
+from typing import Any, Literal, Never, cast
 
 from anthropic.types import ImageBlockParam, Message, MessageParam
 from openai.types import Completion
@@ -346,6 +346,7 @@ def _character_template_history(
     following_user: str = "turn 2",
     omit_length_tail: bool = False,
     length_reasoning: str | None = None,
+    terminal_sampled_stop: bool = True,
 ) -> tuple[ChatCompletionsHistory, _CharacterTemplateTokenizer, list[int]]:
     tokenizer = _CharacterTemplateTokenizer()
     answer = tokenizer._encode("answer")
@@ -359,7 +360,7 @@ def _character_template_history(
         *([] if omit_length_tail else [9]),
         *tokenizer._encode(following_user),
     ]
-    third_output = [*answer, 9]
+    third_output = [*answer, *([9] if terminal_sampled_stop else [])]
 
     first = _chat_exchange(first_prompt, first_output)
     second = _chat_exchange(second_prompt, second_output, offset=1)
@@ -699,6 +700,149 @@ def test_public_exact_chain_preserves_raw_drift_across_proven_length_boundary()
     assert not tokenized.flags[tail] & tr.TokenFlag.SAMPLED
 
 
+@pytest.mark.parametrize("finish_reason", ["stop", "tool_calls"])
+def test_length_chain_retains_exact_prefix_with_terminal_synthetic_stop(
+    finish_reason: Literal["stop", "tool_calls"],
+) -> None:
+    history, tokenizer, captured = _character_template_history(
+        terminal_sampled_stop=False
+    )
+    source = history.message_sources[-1]
+    assert source is not None
+    assert isinstance(source.exchange.response, ChatCompletion)
+    source.exchange.response.choices[0].finish_reason = finish_reason
+
+    tokenized = history.tokenize(tokenizer=tokenizer)
+
+    assert tokenized.tokens == [*captured, 9]
+    assert all(flag & tr.TokenFlag.EXACT for flag in tokenized.flags[:-1])
+    assert tokenized.flags[-1] == (
+        tr.TokenFlag.STOP | tr.TokenFlag.ASSISTANT | tr.TokenFlag.OUTPUT
+    )
+    assert sum(bool(flag & tr.TokenFlag.SAMPLED) for flag in tokenized.flags) == 19
+
+
+def test_terminal_synthetic_stop_does_not_relax_nonterminal_length_proof(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False
+    )
+    history, tokenizer, captured = _character_template_history(
+        terminal_sampled_stop=False,
+        following_user="§turn 2",
+        omit_length_tail=True,
+    )
+    with pytest.warns(UserWarning, match="retokenized an earlier sampled response"):
+        tokenized = history.tokenize(tokenizer=tokenizer)
+    assert tokenized.tokens != [*captured, 9]
+    assert not all(flag & tr.TokenFlag.EXACT for flag in tokenized.flags[:-1])
+
+
+@pytest.mark.parametrize("mismatch", [False, True])
+@pytest.mark.parametrize(
+    "sampled_closer", ["", "<", "</too", "</tool>", "</tool></tool>"]
+)
+def test_length_boundary_ends_before_next_assistant_tool_prefix(
+    mismatch: bool, sampled_closer: str
+) -> None:
+    closing_markup = (
+        "</tool></tool>" if sampled_closer == "</tool></tool>" else "</tool>"
+    )
+
+    class ToolTokenizer(_CharacterTemplateTokenizer):
+        def apply_chat_template(
+            self,
+            messages: list[dict[str, Any]],
+            *,
+            tokenize: bool = True,
+            add_generation_prompt: bool,
+            **kwargs: object,
+        ) -> str | list[int]:
+            text = ""
+            for message in messages:
+                text += str(message.get("content") or "")
+                for call in message.get("tool_calls") or []:
+                    function = call["function"]
+                    text += (
+                        "<tool>"
+                        + function["name"]
+                        + function["arguments"]
+                        + closing_markup
+                    )
+                if message.get("role") == "assistant":
+                    text += "§"
+            return self._encode(text) if tokenize else text
+
+    tokenizer = ToolTokenizer()
+    prompt = tokenizer._encode("turn 0")
+    output = tokenizer._encode("answer")
+    first = _chat_exchange(prompt, output)
+    first.response.choices[0].finish_reason = "length"
+    next_prompt = [*prompt, *output, 9, *tokenizer._encode("turn 1")]
+    if mismatch:
+        next_prompt[len(prompt) + len(output)] = 1000
+    tool_output = tokenizer._encode("<tool>lookup{}" + sampled_closer)
+    second = _chat_exchange(next_prompt, tool_output, offset=1)
+    data = second.response.model_dump(mode="python")
+    data["choices"][0].update(
+        finish_reason="tool_calls",
+        message={
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "call-1",
+                    "type": "function",
+                    "function": {"name": "lookup", "arguments": "{}"},
+                }
+            ],
+        },
+    )
+    second.response = ChatCompletion.model_validate(data)
+    history = art.Trajectory(
+        exchanges=TrajectoryExchanges(chat_completions=[first, second])
+    ).chat_completions_history()
+    tokenized = history.tokenize(tokenizer=tokenizer)
+    expected = [
+        *next_prompt,
+        *tokenizer._encode("<tool>lookup{}" + closing_markup + "§"),
+    ]
+    if mismatch:
+        assert tokenized.tokens != expected
+        return
+    assert tokenized.tokens == expected
+    assert all(
+        flag & tr.TokenFlag.EXACT
+        for flag in tokenized.flags[: len(next_prompt) + len(tool_output)]
+    )
+    assert (
+        tokenized.flags[-1]
+        == tr.TokenFlag.ASSISTANT | tr.TokenFlag.OUTPUT | tr.TokenFlag.STOP
+    )
+    sampled = [
+        index
+        for index, flag in enumerate(tokenized.flags)
+        if flag & tr.TokenFlag.SAMPLED
+    ]
+    assert len(sampled) == len(output) + len(tool_output)
+    assert all(math.isfinite(tokenized.logprobs[index]) for index in sampled)
+
+
+@pytest.mark.parametrize("matches", [[], [(1, 4), (4, 7)], [(4, 7)]])
+def test_terminal_sampled_prefix_requires_one_match_at_assistant_start(
+    matches: list[tuple[int, int]],
+) -> None:
+    from art.trajectories._tokenize import _prove_exact_length_stopped_assistant_prefix
+
+    assert (
+        _prove_exact_length_stopped_assistant_prefix(
+            matches, [False, *([True] * 6)], expected_start=1
+        )
+        is None
+    )
+
+
 def test_public_exact_chain_probes_multi_part_length_response() -> None:
     history, tokenizer, expected = _character_template_history(
         length_reasoning="thinking"