[noaakl]

Add history_adapters option to keep SR control tokens in history

Text callers re-render history without control tokens, so a past SR turn
shows its site token instead of the control token that was sent. The ids
stop matching and the prefix cache misses from that point on.

history_adapters lists the adapter behind each assistant message, and the
chat template puts each control token back in place of the site token.
Works for both the Granite 4.1 and 4.2 chat templates.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[freunda]

Thanks, this is a useful feature. I checked out a5fbe0ff and the new tests pass (26 template tests + 2 real-tokenizer tests against Granite 4.1/4.2). I then probed edge cases by rendering directly.

Inline comments below. The first three change output in ways the docs don't mention:

1. On 4.2 with enable_thinking=False, an answer containing <think> has that text rewritten into the control token.
2. history_adapters entries are silently ignored when the rendered turn doesn't contain the header (enable_thinking=True with stored answers, empty content).
3. The {% filter %} wrap drops output on {% continue %}/{% break %}, and adds newlines when the loop tags don't use -.

The rest are robustness items (undeclared jinja2 dependency, skipped validation for {% generation %} templates, textual endfor matching) plus a duplication nit.

[freunda]

Rewrites answer text on 4.2 when enable_thinking=False.

The docstring at L606–610 says the header text can't occur inside the answer, but it can. If a previous assistant answer is "<think> is the tag", the template doesn't add its own <think></think> (the answer already has a think tag), so the turn renders as <|im_start|>assistant\n<think> is the tag. That is exactly the header, so the replace matches the answer's <think>:

<|im_start|>assistant\n<|sr2|> is the tag<|im_end|>

The user's text is corrupted, and the ids still don't match what was sent (<|sr2|></think><think> is…), so the prefix-cache goal isn't met either. One option is to match header only when it's followed by the template's own </think>, or to place the token from the template branch that emits the empty think block.

[freunda]

Non-None entries can be silently dropped.

The length and name checks pass, but nothing checks that the replacement actually happened. Two cases render without the header:

• 4.2 with enable_thinking=True and the answer stored exactly as returned ('R1\n</think>\nA1') renders as assistant\nR1…, so replace matches nothing and history_adapters=['sr2'] has no effect.
• Empty assistant content on the truncation path (assistant\n<|im_end|>).

The caller gets no error and the prefix cache keeps missing. It would help to fail loudly here, e.g. render the body into a {% set %} block, check that header in body whenever _gs_new != header, and raise_exception otherwise.

[freunda]

{% filter %} around the loop body breaks {% continue %} / {% break %}.

loopcontrols is enabled both here and in HF, so templates that use them are accepted. But a continue/break inside the filter block throws away whatever that iteration already rendered:

{% for m in [1,2,3] %}{{m}}a{% if m==2 %}{% continue %}{% endif %}x{% endfor %}

renders 1ax2a3ax unwrapped and 1bx3bx wrapped (message 2's output is gone). With {% break %}: 1bx instead of 1ax2a. A chat template that skips messages this way would lose whole turns. It's worth rejecting templates whose loop body contains continue/break, or using a different mechanism.

It also adds whitespace when the loop tags don't use -. The inserted {%- endfilter %}\n (and the precheck's trailing {%- endif %}\n) adds a newline per message when the loop tags are plain {% %}, which _FOR_TAG_RE explicitly accepts. {% for message in messages %}{{ message.content }}{% endfor %} renders ab before the wrap and \na\nb\n after (plain Jinja, no trim_blocks). That contradicts the "invisible when unused" claim, and test_render_identical_under_plain_jinja only covers templates whose tags use -. Dropping the trailing \ns, or using -%} on the inserted tags, should fix it.

[freunda]

jinja2 isn't a declared dependency: pyproject lists only torch and transformers, and transformers pulls jinja2 in only through its chat-template / all extras. In a clean [compose] install without those, importing this module fails, which breaks every composer path, including LoRA/aLoRA-only builds that never use history_adapters. Options: declare it, or import it lazily inside _wrap_message_loop_for_history_adapters.

[freunda]

When the original template doesn't parse in this bare environment, wrapped is returned without any validation. Templates using HF's {% generation %}…{% endgeneration %} (assistant masks) hit this path, so a misread endfor would compose fine and then fail at render time, which is exactly what this check is meant to catch. Registering a no-op generation extension, or parsing with HF's own chat-template environment, would cover them.

[freunda]

Finding the matching endfor by regex depth-counting ignores {% raw %} blocks, {# comments #} and string literals. A for/endfor inside any of them shifts the depth, so compose either rejects a valid template or wraps the wrong span. The re-parse below only detects a bad match, and only when the result fails to parse. Walking env.lex(template) (block_begin → name for/endfor) would locate it exactly and handle every whitespace-control spelling for free.

[freunda]

Nit/drift risk: this recomputes the generation-prompt literals and their common prefix, which _resolve_sr_substitution_site already derives. If the site-resolution rules change (escape handling, branch selection), the two can drift apart and the history rewrite quietly stops matching. Could _replace_anchor_in_generation_prompt / _resolve_sr_substitution_site return the header (prefix through the site) along with the site?