[llm_server] Fix transcript replay and reasoning response handling

.github/workflows/_llm_server.yml

@@ -36,12 +36,15 @@ jobs:
         python -m pip install --progress-bar off \
           -r examples/llm_server/python/requirements.txt \
           httpx \
-          pytest
+          pytest \
+          tokenizers
 
         export PYTHONPATH="$(dirname "${PWD}"):${PYTHONPATH:-}"
         export PYTHONDONTWRITEBYTECODE=1
 
-        python -m pytest -q examples/llm_server/python/tests
+        python -m pytest -q \
+          examples/llm_server/python/tests \
+          examples/models/muse-glimmer/tests/test_serve.py
 
         cmake -S . -B cmake-out \
           -DCMAKE_BUILD_TYPE=Release \

.github/workflows/pull.yml

@@ -752,6 +752,28 @@ jobs:
       id-token: write
       contents: read
 
+  llm-server:
+    needs: [changed-files, run-decision]
+    if: |
+      github.event_name == 'pull_request' && (
+        contains(needs.changed-files.outputs.changed-files, 'examples/llm_server') ||
+        contains(needs.changed-files.outputs.changed-files, 'examples/models/muse-glimmer/serving') ||
+        contains(needs.changed-files.outputs.changed-files, 'examples/models/muse-glimmer/tests') ||
+        contains(needs.changed-files.outputs.changed-files, 'examples/models/muse_glimmer') ||
+        contains(needs.changed-files.outputs.changed-files, 'extension/llm/runner') ||
+        contains(needs.changed-files.outputs.changed-files, 'extension/llm/tokenizers') ||
+        contains(needs.changed-files.outputs.changed-files, '.github/workflows/_get-changed-files.yml') ||
+        contains(needs.changed-files.outputs.changed-files, '.github/workflows/_ci-run-decision.yml') ||
+        contains(needs.changed-files.outputs.changed-files, '.github/workflows/_llm_server.yml') ||
+        contains(needs.changed-files.outputs.changed-files, '.github/workflows/pull.yml') ||
+        contains(needs.changed-files.outputs.changed-files, '.github/workflows/trunk.yml') ||
+        needs.run-decision.outputs.is-full-run == 'true'
+      )
+    uses: ./.github/workflows/_llm_server.yml
+    permissions:
+      id-token: write
+      contents: read
+
   unittest:
     uses: ./.github/workflows/_unittest.yml
     permissions:

.github/workflows/trunk.yml

@@ -994,12 +994,20 @@ jobs:
   llm-server:
     needs: [changed-files, run-decision]
     if: |
-      contains(needs.changed-files.outputs.changed-files, 'examples/llm_server') ||
-      contains(needs.changed-files.outputs.changed-files, 'extension/llm/runner') ||
-      contains(needs.changed-files.outputs.changed-files, 'extension/llm/tokenizers') ||
-      contains(needs.changed-files.outputs.changed-files, '.github/workflows/_llm_server.yml') ||
-      contains(needs.changed-files.outputs.changed-files, '.github/workflows/trunk.yml') ||
-      needs.run-decision.outputs.is-full-run == 'true'
+      github.event_name != 'pull_request' && (
+        contains(needs.changed-files.outputs.changed-files, 'examples/llm_server') ||
+        contains(needs.changed-files.outputs.changed-files, 'examples/models/muse-glimmer/serving') ||
+        contains(needs.changed-files.outputs.changed-files, 'examples/models/muse-glimmer/tests') ||
+        contains(needs.changed-files.outputs.changed-files, 'examples/models/muse_glimmer') ||
+        contains(needs.changed-files.outputs.changed-files, 'extension/llm/runner') ||
+        contains(needs.changed-files.outputs.changed-files, 'extension/llm/tokenizers') ||
+        contains(needs.changed-files.outputs.changed-files, '.github/workflows/_get-changed-files.yml') ||
+        contains(needs.changed-files.outputs.changed-files, '.github/workflows/_ci-run-decision.yml') ||
+        contains(needs.changed-files.outputs.changed-files, '.github/workflows/_llm_server.yml') ||
+        contains(needs.changed-files.outputs.changed-files, '.github/workflows/pull.yml') ||
+        contains(needs.changed-files.outputs.changed-files, '.github/workflows/trunk.yml') ||
+        needs.run-decision.outputs.is-full-run == 'true'
+      )
     uses: ./.github/workflows/_llm_server.yml
     permissions:
       id-token: write

examples/llm_server/python/README.md

@@ -60,12 +60,20 @@ Key flags:
 | Flag | Effect |
 |------|--------|
 | `--hf-tokenizer` | model's HF chat template (required unless fallback) |
+| `--assistant-header` | exact assistant generation header, including trailing whitespace (default: ChatML) |
 | `--allow-chatml-fallback` | opt into approximate ChatML when no HF tokenizer |
 | `--no-think` | default `enable_thinking=False` (e.g. Qwen3) |
 | `--max-context N` | reject over-long prompts with 400 instead of failing mid-gen |
 | `--num-runners N` | Worker processes — **1 only** (one worker hosts many isolated sessions on one weight load; more would duplicate weights) |
 | `--worker-bin PATH` | path to a model worker binary that speaks the llm_server JSONL protocol |
 
+Set `--assistant-header` to the model template's exact generation boundary when
+it differs from ChatML. For Llama 3 templates, add
+`--assistant-header $'<|start_header_id|>assistant<|end_header_id|>\n\n'` in Bash;
+the `$'...'` quoting supplies literal newlines. The launcher warns once at startup
+if the configured header is absent from a rendered probe. Unverified boundaries
+use the rendered text, which can reduce KV reuse.
+
 ## Smoke test
 
 ```bash
@@ -114,7 +122,7 @@ Two layers, both contract-focused (assert on the wire, not internals):
 
 ```bash
 # 1. Model-free tests — unit coverage plus loopback disconnect integration.
-pip install pytest httpx
+pip install pytest httpx tokenizers
 pytest tests/
 
 # 2. Conformance — black-box, against a LIVE server (real model, or llama.cpp/mlx-lm).
@@ -126,6 +134,9 @@ real server/protocol/streaming code is tested over HTTP without a `.pte`. The
 worker JSONL protocol is covered separately by `tests/test_worker_client.py`,
 and `tests/test_stream_disconnect.py` uses real loopback Uvicorn/TCP plus a
 model-free subprocess to verify disconnect cancellation end to end.
+The BPE splice tests use an in-memory tokenizer with no model downloads. Optional
+integration tests use local tokenizer directories set with `QWEN_HF_DIR`,
+`GEMMA_HF_DIR`, or `MUSE_GLIMMER_HF_DIR` and require `transformers`.
 
 ## Architecture
 

examples/llm_server/python/chat_template.py

@@ -106,6 +106,7 @@ def __init__(
         # chat_template_kwargs override these.
         self._defaults = default_template_kwargs or {}
         self._assistant_header = assistant_header
+        self._warned_assistant_header = False
         self._strip_rendered_prefix = strip_rendered_prefix
         self._append_generation_prompt_after_tool_response = (
             append_generation_prompt_after_tool_response
@@ -229,8 +230,6 @@ def generation_preamble(
         the resident one (for Qwen3 the scaffold is tool-independent -> same key).
         Returns ``""`` for the fallback / no-scaffold templates (fix is a no-op).
         """
-        if self._hf is None:
-            return ""
         merged = {**self._defaults, **(template_kwargs or {})}
         if tools:
             try:
@@ -251,7 +250,15 @@ def generation_preamble(
             template_kwargs=template_kwargs,
         )
         marker = self._assistant_header
-        idx = rendered.rfind(marker)
+        idx = rendered.rfind(marker) if marker else -1
+        if idx == -1 and not self._warned_assistant_header:
+            logger.warning(
+                "Assistant header %r was not found in the chat template's generation "
+                "prompt. Stored-token replay may fall back to rendered text; configure "
+                "assistant_header (--assistant-header for the generic server).",
+                marker,
+            )
+            self._warned_assistant_header = True
         preamble = rendered[idx + len(marker) :] if idx != -1 else ""
         self._preamble_cache[key] = preamble
         return preamble

examples/llm_server/python/openai_transcript.py

@@ -16,10 +16,10 @@
 token ids and a fingerprint of the response. On the next request each prior
 assistant turn is replaced with a sentinel, the conversation is rendered once,
 and the rendered text is split on the sentinels with the stored ids spliced back
-in -- only for turns whose fingerprint matches the incoming message (an edited,
-branched, or reused history is never substituted with stale ids) and whose ids
-are present (a stop-trimmed turn is left as text). The worker's exact-token
-prefix check is the final backstop.
+in -- only for turns whose content/tool fingerprint and any supplied reasoning string
+match the recorded response, and whose ids are present (a stop-trimmed turn is
+left as text). An edit invalidates that turn and all later records. The worker's
+exact-token prefix check separately protects KV reuse.
 """
 
 import hashlib
@@ -111,9 +111,8 @@ def __init__(self, template: ChatTemplate):
         # boundary verification.
         _header_fn = getattr(template, "assistant_header", None)
         self._assist_hdr = _header_fn() if _header_fn else _ASSIST_HDR
-        # session_id -> [{"fp": str, "ids": list[int] | None}, ...] (one per
-        # assistant turn we produced, in order). Cleared on reset/close.
-        self._turns: dict[str, list[dict]] = {}
+        # Keyed by assistant-turn index, including gaps after invalidation.
+        self._turns: dict[str, dict[int, dict]] = {}
 
     @staticmethod
     def _assistant_fingerprint(content, tool_calls) -> str:
@@ -135,19 +134,24 @@ def _assistant_fingerprint(content, tool_calls) -> str:
         blob = json.dumps([content or "", norm], sort_keys=True, ensure_ascii=False)
         return hashlib.sha1(blob.encode("utf-8")).hexdigest()
 
+    @staticmethod
+    def _reasoning_fingerprint(reasoning_content: Optional[str]) -> Optional[bytes]:
+        if reasoning_content is None:
+            return None
+        return hashlib.sha256(reasoning_content.encode("utf-8")).digest()
+
     def _normalize_scaffold(self, text_chunk: str, preamble: str) -> Optional[str]:
         """Force the scaffold region (between the last assistant header in
         `text_chunk` and its end) to equal `preamble`, so the worker re-tokenizes
         the exact resident scaffold. The region is empty (history stripped it ->
         insert) or a think scaffold (history preserved it -> replace). Returns the
         adjusted text, or None if it isn't a recognized scaffold (-> text fallback)."""
+        if not self._assist_hdr:
+            return None
         h = text_chunk.rfind(self._assist_hdr)
         if h == -1:
-            # No assistant header: with a scaffold to reproduce this is
-            # ambiguous (-> text fallback); without one there is nothing to
-            # normalize, so splicing still works for templates with a different
-            # assistant header.
-            return None if preamble else text_chunk
+            # Without a verified boundary, splicing can duplicate template framing.
+            return None
         base = h + len(self._assist_hdr)
         if not preamble:
             # No generation scaffold: the worker prefills nothing ahead of the
@@ -266,40 +270,37 @@ def build_prompt_input(
         """Return a PromptInput: token-ID segments when this session has faithful
         stored ids for matching prior assistant turns, else the plain rendered
         text. Each incoming assistant turn is matched IN ORDER against the stored
-        records and only spliced when (a) its fingerprint matches what we returned
-        (else the history diverged -> stop, splice nothing further) and (b) we
-        kept faithful ids for it (a stop-trimmed turn's None -> rendered as text).
+        records and only spliced when its content/tool calls and any supplied
+        reasoning string match what we returned, and we kept faithful ids for it.
+        Omitted or null reasoning permits reuse; a string edit invalidates the tail.
         Falls back to text on a sentinel collision or a render that
         dropped/duplicated a sentinel."""
         stored = self._turns.get(session_id or "")
         if not stored:
             return PromptInput(text=rendered_prompt)
-        # Positional: stored[k] is the k-th assistant turn WE generated, matched
-        # against the k-th assistant message in the request. A client-injected
-        # turn (few-shot exemplar, pre-seeded turn, reused session) shifts that
-        # alignment -> fingerprint mismatch at k -> stop splicing. Always safe
-        # (text fallback + worker prefix backstop); just a lower hit rate.
+        # Missing records render as text without shifting later turn indices.
         positions = [i for i, m in enumerate(messages) if m.role == "assistant"]
         splice: dict[int, dict] = {}  # message index -> {"ids", "preamble"}
-        diverged_at = None
         for k, pos in enumerate(positions):
-            if k >= len(stored):
-                break
+            record = stored.get(k)
+            if record is None:
+                continue
             m = messages[pos]
-            if self._assistant_fingerprint(m.content, m.tool_calls) != stored[k]["fp"]:
-                diverged_at = k  # this stored turn and every later one are stale
+            if self._assistant_fingerprint(m.content, m.tool_calls) != record["fp"] or (
+                m.reasoning_content is not None
+                and self._reasoning_fingerprint(m.reasoning_content)
+                != record["reasoning_fp"]
+            ):
+                # Discard the stale tail without shifting subsequent turn indices.
+                self._turns[session_id or ""] = {
+                    index: record for index, record in stored.items() if index < k
+                }
                 break
-            if stored[k]["ids"] is not None:
+            if record["ids"] is not None:
                 splice[pos] = {
-                    "ids": stored[k]["ids"],
-                    "preamble": stored[k].get("preamble", ""),
+                    "ids": record["ids"],
+                    "preamble": record.get("preamble", ""),
                 }
-        if diverged_at is not None:
-            # Drop the stale tail from the first mismatch so an edited/branched
-            # earlier turn can't shadow future requests; the matched prefix still
-            # splices, the rest stays text until reset/close. Safe either way:
-            # stale ids are never spliced and the worker's prefix check backstops.
-            del stored[diverged_at:]
         if not splice:
             return PromptInput(text=rendered_prompt)
         tool_splice = {
@@ -351,25 +352,28 @@ def record_assistant_turn(
         generated_token_ids: list,
         prior_turns: int,
         preamble: str = "",
+        reasoning_content: Optional[str] = None,
     ) -> None:
         """Record this turn's {fingerprint, generated ids, generation preamble} at
         `prior_turns` (the assistant-turn count of the request it answers).
         Records at/after that index are dropped first, so a regenerated/branched
         turn replaces stale records rather than shadowing later hits. ids is None
         when the worker omitted them (stop-trimmed -> non-resumable), kept for
-        positional alignment. `preamble` is the generation scaffold (e.g. the
-        Qwen3 `<think>` block) reproduced ahead of the spliced ids next request."""
+        positional alignment. `reasoning_content` is the client-visible value,
+        including None when the client opted out. `preamble` is the generation
+        scaffold (e.g. the Qwen3 `<think>` block) reproduced ahead of the spliced
+        ids next request."""
         if not session_id:
             return
-        turns = self._turns.setdefault(session_id, [])
-        del turns[prior_turns:]
-        turns.append(
-            {
-                "fp": self._assistant_fingerprint(content, tool_calls),
-                "ids": list(generated_token_ids) if generated_token_ids else None,
-                "preamble": preamble,
-            }
-        )
+        turns = self._turns.setdefault(session_id, {})
+        for index in [index for index in turns if index >= prior_turns]:
+            del turns[index]
+        turns[prior_turns] = {
+            "fp": self._assistant_fingerprint(content, tool_calls),
+            "reasoning_fp": self._reasoning_fingerprint(reasoning_content),
+            "ids": list(generated_token_ids) if generated_token_ids else None,
+            "preamble": preamble,
+        }
 
     def reset(self, session_id: str) -> None:
         self._turns.pop(session_id, None)

examples/llm_server/python/protocol.py

@@ -42,6 +42,7 @@ class ChatMessage(BaseModel):
     # ResponseMessage.reasoning_content so a multi-turn client can echo an assistant
     # turn's reasoning back into the request; a chat template that renders prior
     # reasoning needs it here, and without the field it is dropped at parse.
+    # Omission/null allows stored-token replay; a string edit invalidates that turn.
     reasoning_content: Optional[str] = None
     tool_calls: Optional[list[ToolCall]] = None
     tool_call_id: Optional[str] = None

examples/llm_server/python/server.py

@@ -171,6 +171,12 @@ def main() -> None:
         help="Allow approximate generic ChatML templating when --hf-tokenizer is absent. "
         "Off by default: the fallback can't reproduce model-specific controls.",
     )
+    p.add_argument(
+        "--assistant-header",
+        default="<|im_start|>assistant\n",
+        help="Exact template text ending the assistant generation header, including "
+        "any trailing whitespace. Defaults to the ChatML header.",
+    )
     p.add_argument(
         "--model-id", default="executorch", help="Model id reported on /v1/models"
     )
@@ -218,7 +224,9 @@ def main() -> None:
         args.hf_tokenizer,
         default_template_kwargs=default_template_kwargs,
         allow_fallback=args.allow_chatml_fallback,
+        assistant_header=args.assistant_header,
     )
+    template.generation_preamble()
     worker = _spawn(args)  # one worker hosting many isolated sessions
     runtime = SessionRuntime(worker)
     serving = ServingChat(