From b4a6233cd5ba1c761727fa5130e74639c2c0aeab Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 10 Sep 2026 04:34:41 +0000 Subject: [PATCH 1/4] Preserve captured context across mixed stop boundaries --- src/art/trajectories/_tokenize.py | 49 ++++++++++++++++++------ tests/unit/trajectories/test_tokenize.py | 41 +++++++++++++++++++- 2 files changed, 78 insertions(+), 12 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 081e3bcd8..c40305e60 100644 --- a/src/art/trajectories/_tokenize.py +++ b/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,9 +5254,25 @@ 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 @@ -5310,7 +5337,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 diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 8b0711546..10a14c9b3 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -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,44 @@ 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: str, +) -> None: + history, tokenizer, captured = _character_template_history( + terminal_sampled_stop=False + ) + source = history.message_sources[-1] + assert source is not None + 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]) + + def test_public_exact_chain_probes_multi_part_length_response() -> None: history, tokenizer, expected = _character_template_history( length_reasoning="thinking" From 708b7d590548f2e05eb761b2c9c849a709db5530 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 10 Sep 2026 04:51:06 +0000 Subject: [PATCH 2/4] test: narrow terminal stop fixture response types --- tests/unit/trajectories/test_tokenize.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 10a14c9b3..1e158e4f6 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/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 @@ -702,13 +702,14 @@ def test_public_exact_chain_preserves_raw_drift_across_proven_length_boundary() @pytest.mark.parametrize("finish_reason", ["stop", "tool_calls"]) def test_length_chain_retains_exact_prefix_with_terminal_synthetic_stop( - finish_reason: str, + 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) From 3afdc48d7f9201fda897a866acd80b0da575b8c8 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 10 Sep 2026 04:54:16 +0000 Subject: [PATCH 3/4] Prove length boundaries before the next assistant tool prefix --- src/art/trajectories/_tokenize.py | 8 ++- tests/unit/trajectories/test_tokenize.py | 75 ++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 3 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index c40305e60..94f14f4ef 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5312,11 +5312,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: diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 1e158e4f6..969823c97 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -739,6 +739,81 @@ def test_terminal_synthetic_stop_does_not_relax_nonterminal_length_proof( assert not all(flag & tr.TokenFlag.EXACT for flag in tokenized.flags[:-1]) +@pytest.mark.parametrize("mismatch", [False, True]) +def test_length_boundary_ends_before_next_assistant_tool_prefix(mismatch: bool) -> None: + 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 += ( + "" + function["name"] + function["arguments"] + "" + ) + 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("lookup{}") + 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, *tool_output, *tokenizer._encode("§")] + 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) + + def test_public_exact_chain_probes_multi_part_length_response() -> None: history, tokenizer, expected = _character_template_history( length_reasoning="thinking" From 5de63beed7cb155938310a895b1da33c37ad1b7d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 10 Sep 2026 05:26:07 +0000 Subject: [PATCH 4/4] Prove complete sampled terminal span before retaining renderer tail --- src/art/trajectories/_tokenize.py | 12 ++++++++ tests/unit/trajectories/test_tokenize.py | 37 +++++++++++++++++++++--- 2 files changed, 45 insertions(+), 4 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 94f14f4ef..6a7b586e8 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5279,6 +5279,18 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: 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 ): diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 969823c97..7dae26ecc 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -740,7 +740,16 @@ def test_terminal_synthetic_stop_does_not_relax_nonterminal_length_proof( @pytest.mark.parametrize("mismatch", [False, True]) -def test_length_boundary_ends_before_next_assistant_tool_prefix(mismatch: bool) -> None: +@pytest.mark.parametrize( + "sampled_closer", ["", "<", "", ""] +) +def test_length_boundary_ends_before_next_assistant_tool_prefix( + mismatch: bool, sampled_closer: str +) -> None: + closing_markup = ( + "" if sampled_closer == "" else "" + ) + class ToolTokenizer(_CharacterTemplateTokenizer): def apply_chat_template( self, @@ -756,7 +765,10 @@ def apply_chat_template( for call in message.get("tool_calls") or []: function = call["function"] text += ( - "" + function["name"] + function["arguments"] + "" + "" + + function["name"] + + function["arguments"] + + closing_markup ) if message.get("role") == "assistant": text += "§" @@ -770,7 +782,7 @@ def apply_chat_template( next_prompt = [*prompt, *output, 9, *tokenizer._encode("turn 1")] if mismatch: next_prompt[len(prompt) + len(output)] = 1000 - tool_output = tokenizer._encode("lookup{}") + tool_output = tokenizer._encode("lookup{}" + sampled_closer) second = _chat_exchange(next_prompt, tool_output, offset=1) data = second.response.model_dump(mode="python") data["choices"][0].update( @@ -792,7 +804,10 @@ def apply_chat_template( exchanges=TrajectoryExchanges(chat_completions=[first, second]) ).chat_completions_history() tokenized = history.tokenize(tokenizer=tokenizer) - expected = [*next_prompt, *tool_output, *tokenizer._encode("§")] + expected = [ + *next_prompt, + *tokenizer._encode("lookup{}" + closing_markup + "§"), + ] if mismatch: assert tokenized.tokens != expected return @@ -814,6 +829,20 @@ def apply_chat_template( 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"