diff --git a/docs/SR_ARCHITECTURE.md b/docs/SR_ARCHITECTURE.md
index 0e3c3b92..e605b37f 100644
--- a/docs/SR_ARCHITECTURE.md
+++ b/docs/SR_ARCHITECTURE.md
@@ -144,6 +144,47 @@ No manual configuration is needed — pass SR adapter paths to the composer and
`validate_cross_stream_population` then asserts every adapter in a `dual_stream` checkpoint
has non-zero `cross_stream` weights, catching a silent load failure.
+## Keeping Control Tokens in History
+
+A text-based multi-turn caller re-renders history from `messages`, which carry no
+control tokens, so an SR turn's header renders with its site token instead of the
+control token that was sent. For SR this changes no KV and no used output, but the
+ids no longer match what was sent and the prefix cache misses from there on.
+
+Pass `history_adapters`, exactly one entry per assistant message (the adapter
+that produced it, or `None`), and the template puts the tokens back. The control
+token replaces the site token; nothing is added:
+
+```python
+tokenizer.apply_chat_template(messages, add_generation_prompt=True,
+ adapter_name="sr2", history_adapters=[None, "sr1"])
+# vLLM /v1/chat/completions and mellea: same keys inside chat_template_kwargs
+```
+
+```
+4.1 assistant<|end_of_role|>A1 -> assistant<|sr1|>A1
+4.2 enable_thinking=F assistant\nA1 -> assistant\n<|sr1|>A1
+```
+
+A list of the wrong length, or an entry that is not an SR adapter, raises. No
+effect when the current turn is base.
+
+For the cache to hit, the caller must:
+
+* store the exact answer text the model produced (no post-processing), and send
+ `skip_special_tokens: False` so a control token generated mid-answer is not
+ stripped from it;
+* on Granite 4.2 with `enable_thinking=True`, store `"\n"` + the returned
+ text and pass `truncate_history_thinking=False`; otherwise history thinking is
+ rewritten and the match breaks right after the control token;
+* serve Granite 4.2 with `--chat-template-content-format string`, or vLLM rejects
+ every chat request with a 400.
+
+Even exact text can re-encode to different ids (the model may emit `"\n"` + `"\n"`
+where the tokenizer picks `"\n\n"`), so the cache misses from that point on that
+turn. Only a caller that keeps the generated ids and sends them to `/v1/completions`
+avoids this.
+
## Config Parameters
| Parameter | Type | Default | Description |
diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py
index 29be363b..5b3815f4 100644
--- a/src/granite_switch/composer/tokenizer_setup.py
+++ b/src/granite_switch/composer/tokenizer_setup.py
@@ -9,6 +9,8 @@
import re
from dataclasses import dataclass
+import jinja2
+
from .adapter_loader import load_adapter_config
# Name of the optional return-to-base control token (``<|base_reset|>``). It is
@@ -552,6 +554,185 @@ def _rewrite(emission: "re.Match") -> str:
return head + region, site, count
+#: A ``for`` / ``endfor`` statement tag, for depth-matching the message loop.
+#: Accepts every whitespace-control spelling: ``{%``, ``{%-`` and ``{%+``.
+_FOR_TAG_RE = re.compile(r"\{%[-+]?\s*(for|endfor)\b")
+
+
+def _jinja_string_literal(text: str) -> str:
+ """*text* as a single-quoted Jinja string literal that renders back to it."""
+ escaped = text.replace("\\", "\\\\").replace("'", "\\'").replace("\n", "\\n")
+ return "'" + escaped + "'"
+
+
+def _sr_history_header(template: str, site: str) -> str:
+ """The rendered assistant header, through *site*, that history turns repeat.
+
+ Taken from the base template's generation prompt: every branch opens with
+ the same assistant header, and a history assistant turn renders that same
+ header before its content (``<|start_of_role|>assistant<|end_of_role|>`` on
+ 4.1, ``<|im_start|>assistant\\n`` on 4.2, where the history turn
+ continues ````).
+
+ Raises:
+ ValueError: If the generation prompt's branches share no prefix that
+ contains *site*.
+ """
+ _, region = _split_at_generation_prompt(template)
+ literals = [_unescape_literal(body) for body in _generation_prompt_literals(region)]
+ prefix = os.path.commonprefix(literals) if literals else ""
+ cut = prefix.rfind(site)
+ if cut == -1:
+ raise ValueError(
+ f"The generation prompt's branches share no prefix containing the "
+ f"Shadow Residual substitution site {site!r} ({literals!r}), so the "
+ f"history_adapters rewrite has no header to match."
+ )
+ return prefix[: cut + len(site)]
+
+
+def _wrap_message_loop_for_history_adapters(
+ template: str, loop_var_source: str, header: str, site: str
+) -> str:
+ """Let ``history_adapters`` put SR control tokens back into history turns.
+
+ Rendering history from ``messages`` emits every assistant header with the
+ SR *site* (the token the control token replaced at generation time), so the
+ re-rendered ids diverge from the ones that were sent and the prefix cache
+ misses. ``history_adapters[k]`` names the adapter that produced the k-th
+ assistant message; this rewrites that message's header back so its *site*
+ is the control token.
+
+ The match is on the whole *header* (e.g. ``<|im_start|>assistant\\n``),
+ which ends in *site*, not on *site* alone. Granite 4.2 writes its own
+ ```` only when the answer has no think tags, so for an answer
+ that mentions ```` the first bare ```` is the answer's text; the
+ header text cannot occur there.
+
+ The main loop body is wrapped in ``{% filter replace(header, new, 1) %}``,
+ so the replacement applies to the body's rendered text, rather than
+ rewriting each assistant emission in the source: the 4.2 template emits the
+ assistant header from several branches and reads ```` in its own
+ history-truncation logic, so a token placed in the source would be seen
+ (and stripped, or misread) by that logic. Rendered-text replacement runs
+ after it and is indifferent to which branch ran.
+
+ Only applied when the current turn is SR (``ns.adapter_type == 'sr'``): a
+ base turn after a kept SR token would generate under SR (latest-wins).
+
+ ``history_adapters`` must have exactly one entry per assistant message
+ (``None`` for a turn that had no adapter); any other length raises.
+
+ Raises:
+ ValueError: If *header* does not end in *site*, the template has no
+ main message loop with a matching ``endfor``, or the wrapped
+ template no longer parses (e.g. an ``endfor`` inside a comment
+ fooled the tag scan).
+ """
+ if not header.endswith(site):
+ raise ValueError(
+ f"history_adapters header {header!r} does not end in the Shadow "
+ f"Residual substitution site {site!r}."
+ )
+ open_tag = re.search(
+ r"\{%[-+]?\s*for message in " + re.escape(loop_var_source) + r"\s*[-+]?%\}",
+ template,
+ )
+ if open_tag is None:
+ raise ValueError(
+ f"No '{{%- for message in {loop_var_source} %}}' loop in the chat "
+ f"template, so history_adapters has nothing to wrap."
+ )
+ depth = 1
+ for tag in _FOR_TAG_RE.finditer(template, open_tag.end()):
+ depth += 1 if tag.group(1) == "for" else -1
+ if depth == 0:
+ close_at = tag.start()
+ break
+ else:
+ raise ValueError(
+ "The chat template's message loop has no matching endfor, so "
+ "history_adapters has nothing to wrap."
+ )
+
+ # Checked once, before the loop: a list that does not line up one-to-one
+ # with the assistant messages is a caller bug, and silently padding or
+ # truncating it would put tokens on the wrong turns.
+ precheck = (
+ "{%- if ns.adapter_type == '" + ANCHOR_MODE_SR + "'"
+ " and history_adapters is defined and history_adapters is not none %}\n"
+ "{%- if history_adapters is string or history_adapters is mapping"
+ " or history_adapters is not iterable %}\n"
+ "{{- raise_exception('history_adapters must be a list of adapter names "
+ "(or none), one per assistant message') }}\n"
+ "{%- endif %}\n"
+ "{%- set _gs_expected = " + loop_var_source + " | selectattr('role', "
+ "'equalto', 'assistant') | list | length %}\n"
+ "{%- if history_adapters | length != _gs_expected %}\n"
+ "{{- raise_exception('history_adapters has ' ~ (history_adapters | length)"
+ " ~ ' entries but the conversation has ' ~ _gs_expected ~ ' assistant "
+ "messages; pass exactly one entry (an adapter name or none) per "
+ "assistant message') }}\n"
+ "{%- endif %}\n"
+ "{%- for _gs_entry in history_adapters %}\n"
+ # ``is not string`` first: ``or`` short-circuits, so an unhashable entry
+ # never reaches the ``in`` lookup and gets this message, not a TypeError.
+ "{%- if _gs_entry is not none and (_gs_entry is not string"
+ " or _gs_entry not in adapter_map"
+ " or adapter_map[_gs_entry]['type'] != '" + ANCHOR_MODE_SR + "') %}\n"
+ "{{- raise_exception('history_adapters names ' ~ _gs_entry ~ ', which is "
+ "not a Shadow Residual adapter of this checkpoint') }}\n"
+ "{%- endif %}\n"
+ "{%- endfor %}\n"
+ "{%- endif %}\n"
+ )
+ # Picks this message's replacement before its body renders; the default
+ # (header -> header) leaves the body untouched.
+ choose = (
+ "\n{%- set _gs_new = " + _jinja_string_literal(header) + " %}\n"
+ "{%- if message.role == 'assistant' %}\n"
+ "{%- if ns.adapter_type == '" + ANCHOR_MODE_SR + "'"
+ " and history_adapters is defined and history_adapters is not none"
+ " and history_adapters[ns.assistant_idx] is not none %}\n"
+ "{%- set _gs_new = "
+ + _jinja_string_literal(header[: len(header) - len(site)])
+ + " ~ adapter_map[history_adapters[ns.assistant_idx]]['token'] %}\n"
+ "{%- endif %}\n"
+ "{%- set ns.assistant_idx = ns.assistant_idx + 1 %}\n"
+ "{%- endif %}\n"
+ "{%- filter replace(" + _jinja_string_literal(header) + ", _gs_new, 1) %}"
+ )
+ wrapped = (
+ template[: open_tag.start()]
+ + precheck
+ + template[open_tag.start() : open_tag.end()]
+ + choose
+ + template[open_tag.end() : close_at]
+ + "{%- endfilter %}\n"
+ + template[close_at:]
+ )
+
+ # The tag scan is textual, so a template it misreads composes fine and then
+ # fails at render time. Parse both versions to catch that here instead; a
+ # template that does not parse even before the wrap (an extension tag this
+ # environment lacks) is left for the real renderer to judge.
+ env = jinja2.Environment(extensions=["jinja2.ext.loopcontrols"])
+ try:
+ env.parse(template)
+ except jinja2.TemplateSyntaxError:
+ return wrapped
+ try:
+ env.parse(wrapped)
+ except jinja2.TemplateSyntaxError as exc:
+ raise ValueError(
+ f"Wrapping the chat template's message loop for history_adapters "
+ f"produced a template that does not parse ({exc}). The loop's "
+ f"endfor was probably mis-identified, e.g. one inside a comment or "
+ f"a string."
+ ) from exc
+ return wrapped
+
+
def add_control_tokens(
tokenizer,
discovered_adapters: list[tuple[str | None, str, str, str | None]],
@@ -1165,7 +1346,9 @@ def configure_chat_template(
"\n adapter_invocation_text=adapter_invocation_text,"
"\n adapter_invocation_drop_chars=adapter_invocation_drop_chars,"
"\n alora_target_idx=-1,"
- "\n skip_next_role_marker=false"
+ # history_adapters state (SR only): the assistant message being visited.
+ + ("\n assistant_idx=0," if sr_anchors else "")
+ + "\n skip_next_role_marker=false"
"\n )"
)
modified_chat_template = (
@@ -1318,6 +1501,12 @@ def configure_chat_template(
for name, info in adapter_mapping.items()
if info["type"] == ANCHOR_MODE_SR
}
+ modified_chat_template = _wrap_message_loop_for_history_adapters(
+ modified_chat_template,
+ fmt.loop_var_source,
+ _sr_history_header(base_chat_template, sr_site),
+ sr_site,
+ )
# Skip-once wrapper for every role-open marker emission in the template.
# ns.skip_next_role_marker is set to true immediately after a LoRA or
@@ -1412,4 +1601,8 @@ def _split_merged(match: "re.Match") -> str:
f" - SR tokens: substituted for {sr_site!r} in all "
f"{sr_replacements} generation-prompt render path(s)"
)
+ print(
+ f" - SR history: history_adapters=[...] restores {sr_site!r} -> "
+ f"control token in earlier assistant turns"
+ )
return fmt.name, sr_substitute_token_ids
diff --git a/tests/composer/test_chat_template.py b/tests/composer/test_chat_template.py
index 5cbabc0c..32744bde 100644
--- a/tests/composer/test_chat_template.py
+++ b/tests/composer/test_chat_template.py
@@ -23,6 +23,7 @@
- End-to-end: adapter_config.json → render (no patching)
"""
+import json
import os
from types import SimpleNamespace
from unittest.mock import patch
@@ -31,6 +32,8 @@
from jinja2 import Environment
from granite_switch.composer.tokenizer_setup import (
+ _wrap_message_loop_for_history_adapters,
+ add_control_tokens,
configure_audio_chat_template,
configure_chat_template,
detect_template_format,
@@ -1423,3 +1426,428 @@ def test_sr_anchor_id_must_match_the_tokenizer(self):
configure_chat_template(
tokenizer, [(self._SR_GRANITE, "answerability", "sr")]
)
+
+
+class TestShadowResidualHistoryAdapters:
+ """``history_adapters`` restores SR control tokens in earlier assistant turns.
+
+ For SR the token at a history position changes no KV (K/V come from the base
+ stream) and no used output; what it buys is id equality with what was sent,
+ which is what vLLM's prefix cache keys on. So the central property is a prefix
+ relation: the turn-N prompt plus the model's answer must be a prefix of the
+ turn-N+1 render.
+
+ Renders use HF's whitespace flags (``trim_blocks``/``lstrip_blocks``), as
+ ``apply_chat_template`` and vLLM do: under plain Jinja the 4.2 template's own
+ history-truncation path leaks whitespace, so the prefix relation fails there
+ even without this feature.
+ """
+
+ _ENV = Environment(trim_blocks=True, lstrip_blocks=True)
+
+ _SR_GRANITE = os.path.join(_FIXTURES, "sr_answerability_adapter")
+ _SR_CHATML = os.path.join(_FIXTURES, "sr_answerability_adapter_chatml")
+
+ _MESSAGES = [
+ {"role": "user", "content": "Q1"},
+ {"role": "assistant", "content": "A1"},
+ {"role": "user", "content": "Q2"},
+ {"role": "assistant", "content": "A2"},
+ {"role": "user", "content": "Q3"},
+ ]
+
+ @staticmethod
+ def _raise(message):
+ raise ValueError(message)
+
+ def _granite(self):
+ tokenizer = _SRTokenizer(_GRANITE_TEMPLATE, {"<|end_of_role|>": [49153]})
+ configure_chat_template(
+ tokenizer,
+ [(self._SR_GRANITE, "sr1", "sr"), (self._SR_GRANITE, "sr2", "sr")],
+ )
+ return tokenizer
+
+ def _chatml(self):
+ tokenizer = _SRTokenizer(_CHATML_TEMPLATE, {"": [49154]})
+ configure_chat_template(
+ tokenizer,
+ [(self._SR_CHATML, "sr1", "sr"), (self._SR_CHATML, "sr2", "sr")],
+ )
+ return tokenizer
+
+ def _render(self, tokenizer, messages, **kwargs):
+ return self._ENV.from_string(tokenizer.chat_template).render(
+ messages=messages,
+ add_generation_prompt=True,
+ raise_exception=self._raise,
+ **kwargs,
+ )
+
+ def test_granite_history_headers_carry_their_control_tokens(self):
+ result = self._render(
+ self._granite(),
+ self._MESSAGES,
+ adapter_name="sr1",
+ history_adapters=["sr1", "sr2"],
+ )
+ assert "<|start_of_role|>assistant<|sr1|>A1<|end_of_text|>" in result
+ assert "<|start_of_role|>assistant<|sr2|>A2<|end_of_text|>" in result
+ assert result.endswith("<|start_of_role|>assistant<|sr1|>")
+
+ def test_granite_sent_prompt_is_prefix_of_next_turn(self):
+ tokenizer = self._granite()
+ sent = self._render(
+ tokenizer, self._MESSAGES[:3], adapter_name="sr2", history_adapters=["sr1"]
+ )
+ nxt = self._render(
+ tokenizer,
+ self._MESSAGES,
+ adapter_name="sr1",
+ history_adapters=["sr1", "sr2"],
+ )
+ assert nxt.startswith(sent + "A2<|end_of_text|>\n")
+
+ def test_chatml_no_thinking_sent_prompt_is_prefix_of_next_turn(self):
+ tokenizer = self._chatml()
+ sent = self._render(
+ tokenizer,
+ self._MESSAGES[:3],
+ enable_thinking=False,
+ adapter_name="sr2",
+ history_adapters=["sr1"],
+ )
+ nxt = self._render(
+ tokenizer,
+ self._MESSAGES,
+ enable_thinking=False,
+ adapter_name="sr1",
+ history_adapters=["sr1", "sr2"],
+ )
+ assert "<|im_start|>assistant\n<|sr1|>A1" in nxt
+ assert nxt.startswith(sent + "A2<|im_end|>\n")
+
+ def test_chatml_thinking_places_tokens_but_breaks_the_prefix(self):
+ """With thinking on and the answer stored as returned, the sent prompt
+ ends ``<|sr2|>\\n`` but history re-renders as ``<|sr2|>A``, so no
+ prefix relation holds. The next test shows the caller-side fix.
+ """
+ tokenizer = self._chatml()
+ sent = self._render(
+ tokenizer,
+ self._MESSAGES[:3],
+ enable_thinking=True,
+ adapter_name="sr2",
+ history_adapters=["sr1"],
+ )
+ nxt = self._render(
+ tokenizer,
+ self._MESSAGES,
+ enable_thinking=True,
+ adapter_name="sr1",
+ history_adapters=["sr1", "sr2"],
+ )
+ assert "<|im_start|>assistant\n<|sr2|>A2" in nxt
+ assert not nxt.startswith(sent + "A2")
+
+ @pytest.mark.parametrize("truncate", [False, True])
+ def test_chatml_thinking_prefix_holds_with_think_restored(self, truncate):
+ """With thinking on, the control token takes ````'s place in the sent
+ prompt, so the model's text starts after it. Storing ``"\\n"`` + that
+ text gives the template its header back and the prefix holds -- but only
+ with ``truncate_history_thinking=False``: truncation strips the reasoning
+ out of every answer before the last user turn.
+ """
+ tokenizer = self._chatml()
+ answers = ["R1\n\nA1", "R2\n\nA2"]
+ messages = [
+ {"role": "user", "content": "Q1"},
+ {"role": "assistant", "content": "\n" + answers[0]},
+ {"role": "user", "content": "Q2"},
+ {"role": "assistant", "content": "\n" + answers[1]},
+ {"role": "user", "content": "Q3"},
+ ]
+ kwargs = {"enable_thinking": True, "truncate_history_thinking": truncate}
+ sent = self._render(
+ tokenizer,
+ messages[:3],
+ adapter_name="sr2",
+ history_adapters=["sr1"],
+ **kwargs,
+ )
+ nxt = self._render(
+ tokenizer,
+ messages,
+ adapter_name="sr1",
+ history_adapters=["sr1", "sr2"],
+ **kwargs,
+ )
+ assert sent.endswith("<|im_start|>assistant\n<|sr2|>\n")
+ assert nxt.startswith(sent + answers[1] + "<|im_end|>\n") is not truncate
+
+ @pytest.mark.parametrize("fmt", ["granite", "chatml"])
+ def test_base_turn_emits_no_history_tokens(self, fmt):
+ tokenizer = self._granite() if fmt == "granite" else self._chatml()
+ result = self._render(
+ tokenizer, self._MESSAGES, history_adapters=["sr1", "sr2"]
+ )
+ assert "<|sr" not in result
+
+ @pytest.mark.parametrize("fmt", ["granite", "chatml"])
+ def test_history_adapters_none_is_a_no_op(self, fmt):
+ tokenizer = self._granite() if fmt == "granite" else self._chatml()
+ absent = self._render(tokenizer, self._MESSAGES, adapter_name="sr1")
+ explicit_none = self._render(
+ tokenizer, self._MESSAGES, adapter_name="sr1", history_adapters=None
+ )
+ all_none = self._render(
+ tokenizer,
+ self._MESSAGES,
+ adapter_name="sr1",
+ history_adapters=[None, None],
+ )
+ assert absent == explicit_none == all_none
+ assert absent.count("<|sr1|>") == 1
+
+ def test_absent_history_adapters_matches_the_unpatched_template(self):
+ """The wrap is invisible when unused: undoing the generation-prompt
+ substitution reproduces the original template's render exactly."""
+ for enable_thinking in (False, True):
+ base = self._ENV.from_string(_CHATML_TEMPLATE).render(
+ messages=self._MESSAGES,
+ add_generation_prompt=True,
+ enable_thinking=enable_thinking,
+ )
+ patched = self._render(
+ self._chatml(),
+ self._MESSAGES,
+ enable_thinking=enable_thinking,
+ adapter_name="sr1",
+ )
+ assert patched.replace("<|sr1|>", "") == base
+ base = self._ENV.from_string(_GRANITE_TEMPLATE).render(
+ messages=self._MESSAGES, add_generation_prompt=True
+ )
+ patched = self._render(self._granite(), self._MESSAGES, adapter_name="sr1")
+ assert patched.replace("<|sr1|>", "<|end_of_role|>") == base
+
+ def test_render_identical_under_plain_jinja(self):
+ """The class renders with HF's flags; the filter block must not shift
+ whitespace under plain Jinja either."""
+ env = Environment()
+ for tokenizer, site, base_template in (
+ (self._granite(), "<|end_of_role|>", _GRANITE_TEMPLATE),
+ (self._chatml(), "", _CHATML_TEMPLATE),
+ ):
+ base = env.from_string(base_template).render(
+ messages=self._MESSAGES, add_generation_prompt=True
+ )
+ patched = env.from_string(tokenizer.chat_template).render(
+ messages=self._MESSAGES,
+ add_generation_prompt=True,
+ adapter_name="sr1",
+ raise_exception=self._raise,
+ )
+ assert patched.replace("<|sr1|>", site) == base
+
+ def test_index_counts_assistant_messages_only(self):
+ messages = [
+ {"role": "system", "content": "S"},
+ {"role": "user", "content": "Q1"},
+ {"role": "assistant", "content": "A1"},
+ {"role": "tool", "content": "T1"},
+ {"role": "user", "content": "Q2"},
+ {"role": "assistant", "content": "A2"},
+ {"role": "user", "content": "Q3"},
+ ]
+ result = self._render(
+ self._granite(),
+ messages,
+ adapter_name="sr1",
+ history_adapters=[None, "sr2"],
+ )
+ assert "assistant<|end_of_role|>A1" in result
+ assert "assistant<|sr2|>A2" in result
+
+ @pytest.mark.parametrize(
+ "history_adapters",
+ [["sr2"], ["sr2", "sr2", "sr2"], []],
+ ids=["too-short", "too-long", "empty"],
+ )
+ def test_length_mismatch_raises(self, history_adapters):
+ """One entry per assistant message, no more and no fewer: a list that
+ does not line up would put tokens on the wrong turns."""
+ with pytest.raises(ValueError, match="has 2 assistant messages"):
+ self._render(
+ self._granite(),
+ self._MESSAGES,
+ adapter_name="sr1",
+ history_adapters=history_adapters,
+ )
+
+ def test_length_is_not_checked_on_a_base_turn(self):
+ """A base turn does not read the list, so a wrong length is not an error."""
+ result = self._render(self._granite(), self._MESSAGES, history_adapters=["sr2"])
+ assert "<|sr" not in result
+
+ def test_unknown_adapter_raises(self):
+ with pytest.raises(ValueError, match="not a Shadow Residual adapter"):
+ self._render(
+ self._granite(),
+ self._MESSAGES,
+ adapter_name="sr1",
+ history_adapters=["nope", None],
+ )
+
+ def test_non_sr_checkpoint_is_unaffected(self):
+ with patch(_PATCH_TARGET, return_value=""):
+ tokenizer = _make_tokenizer()
+ configure_chat_template(tokenizer, [("/path/a", "req_check", "alora")])
+ assert "history_adapters" not in tokenizer.chat_template
+ assert "assistant_idx" not in tokenizer.chat_template
+
+ def test_chatml_answer_containing_the_site_is_not_rewritten(self):
+ """4.2 writes its own ```` header only when the answer has
+ no think tags, so for this answer the first ```` is the answer's
+ own text. Matching the full header leaves it alone."""
+ messages = [
+ {"role": "user", "content": "Q1"},
+ {"role": "assistant", "content": "See tag"},
+ {"role": "user", "content": "Q2"},
+ ]
+ result = self._render(
+ self._chatml(),
+ messages,
+ enable_thinking=False,
+ adapter_name="sr1",
+ history_adapters=["sr2"],
+ )
+ assert "<|im_start|>assistant\nSee tag<|im_end|>" in result
+ assert "<|sr2|>" not in result
+
+ def test_bare_string_raises(self):
+ """A string is a sequence of characters, not a list of names."""
+ with pytest.raises(ValueError, match="must be a list"):
+ self._render(
+ self._granite(),
+ self._MESSAGES,
+ adapter_name="sr1",
+ history_adapters="sr1",
+ )
+
+ def test_non_string_entry_raises_the_named_error(self):
+ with pytest.raises(ValueError, match="not a Shadow Residual adapter"):
+ self._render(
+ self._granite(),
+ self._MESSAGES,
+ adapter_name="sr1",
+ history_adapters=[["sr1"], None],
+ )
+
+ def test_wrap_rejects_a_result_that_does_not_parse(self):
+ """An ``endfor`` inside a comment fools the tag scan; the wrapped
+ template must fail at compose time, not at render time."""
+ template = (
+ "{%- set ns = namespace(adapter_type='') %}"
+ "{%- for message in messages %}{# a {% endfor %} b #}"
+ "{{ message.content }}{%- endfor %}"
+ )
+ with pytest.raises(ValueError, match="does not parse"):
+ _wrap_message_loop_for_history_adapters(template, "messages", "", "")
+
+ def test_wrap_counts_plus_whitespace_control_tags(self):
+ """``{%+ for`` opens a loop like ``{%- for`` does."""
+ template = (
+ "{%- set ns = namespace(adapter_type='', assistant_idx=0) %}"
+ "{%- for message in messages %}"
+ "{%+ for x in [1] %}{{ message.content }}{% endfor %}"
+ "{%- endfor %}"
+ )
+ wrapped = _wrap_message_loop_for_history_adapters(
+ template, "messages", "", ""
+ )
+ rendered = self._ENV.from_string(wrapped).render(
+ messages=[{"role": "user", "content": "a"}], adapter_map={}
+ )
+ assert rendered == "a"
+
+
+class TestShadowResidualHistoryAdaptersTokenIds:
+ """The prefix relation on token ids, with the real Granite tokenizers.
+
+ vLLM's prefix cache keys on ids, not text; matching strings can still
+ tokenize differently at a boundary such as the one right after a control
+ token. Skipped when the tokenizer cannot be fetched.
+ """
+
+ _CASES = [
+ (
+ "ibm-granite/granite-4.1-3b",
+ "sr_answerability_adapter",
+ {},
+ "<|end_of_text|>",
+ ),
+ (
+ "ibm-granite/granite-4.2-3b",
+ "sr_answerability_adapter_chatml",
+ {"enable_thinking": False},
+ "<|im_end|>",
+ ),
+ ]
+
+ @staticmethod
+ def _ids(out):
+ # transformers 5.x returns a BatchEncoding here; 4.x returns a plain list.
+ return list(out["input_ids"]) if hasattr(out, "keys") else list(out)
+
+ @pytest.mark.parametrize(
+ "model,fixture,kwargs,turn_end", _CASES, ids=["granite-4.1", "granite-4.2"]
+ )
+ def test_sent_ids_are_prefix_of_next_turn(
+ self, model, fixture, kwargs, turn_end, tmp_path
+ ):
+ from transformers import AutoTokenizer
+
+ try:
+ tokenizer = AutoTokenizer.from_pretrained(model)
+ except Exception as e:
+ pytest.skip(f"could not fetch {model} tokenizer: {e}")
+ # The fixture pins the fake tokenizer's site id; drop it so the composer
+ # resolves the id against the real one.
+ with open(os.path.join(_FIXTURES, fixture, "adapter_config.json")) as f:
+ adapter_config = json.load(f)
+ adapter_config.pop("last_context_token_id", None)
+ adapter_path = str(tmp_path / fixture)
+ os.makedirs(adapter_path)
+ with open(os.path.join(adapter_path, "adapter_config.json"), "w") as f:
+ json.dump(adapter_config, f)
+ adapters = [
+ (adapter_path, "sr1", "sr", None),
+ (adapter_path, "sr2", "sr", None),
+ ]
+ add_control_tokens(tokenizer, adapters)
+ configure_chat_template(tokenizer, [a[:3] for a in adapters])
+
+ messages = TestShadowResidualHistoryAdapters._MESSAGES
+ sent = self._ids(
+ tokenizer.apply_chat_template(
+ messages[:3],
+ add_generation_prompt=True,
+ adapter_name="sr2",
+ history_adapters=[None],
+ **kwargs,
+ )
+ )
+ answer = tokenizer.encode(
+ messages[3]["content"] + turn_end, add_special_tokens=False
+ )
+ nxt = self._ids(
+ tokenizer.apply_chat_template(
+ messages,
+ add_generation_prompt=True,
+ adapter_name="sr1",
+ history_adapters=[None, "sr2"],
+ **kwargs,
+ )
+ )
+ assert nxt[: len(sent) + len(answer)] == sent + answer