fix(chat): keep every message a turn appends after the reply - #237

Problem

A chat continuation turn is the previous request, then the model's reply, then the new messages. _run_turn sliced that turn's delta as if exactly one message followed the reply: it rendered chat[:-1] as the prefix and chat[:-2] as the history the reply was generated after. Any other count goes wrong silently, because the request still succeeds:

• Two or more messages after the reply (parallel tool results, an observation plus a reminder): only the last one reaches the generation context and the trained tokens. The others are dropped, and so is the template's close of the reply (the \n after <|im_end|> in ChatML).
• No message after the reply (the agent asks the model to keep going): the reply is appended again as feedback, so the model sees it twice.

Minimal reproduction: a ChatML template, a reply that makes two tool calls, and the two tool responses sent back as two user messages (the <tool_response> convention chat_geo3k.py uses). This is the context turn 2 generates over on main:

<|im_start|>user
Q<|im_end|>
<|im_start|>assistant
<tool_call>a</tool_call><tool_call>b</tool_call><|im_end|><|im_start|>user
<tool_response>
result b
</tool_response><|im_end|>
<|im_start|>assistant

result a never reaches the model, and the trajectory trains on this context.

Fix

Index the slices by the previous request's length n. messages[n] is the reply, so the prefix is chat[:n + 1], the history is chat[:n], and the delta is everything after the reply. New images are read from all of messages[n + 1:], not only the last message.

A request that extends the previous one without echoing the reply (messages[n] is not an assistant message) rewrites the history the segment holds. It now starts a fresh segment, like a compaction, instead of being sliced into the old one.

A turn that appends exactly one message after the reply, which is what chat_geo3k.py sends, renders exactly as before.

Tests

In tests/unit/test_chat_server.py, using the existing drifty template and engine fakes:

• test_continuation_delta_covers_every_message_after_the_reply, parametrized over two tool results and the reply alone. Turn 2's generation context and the segment must equal the template's render, with each new message exactly once. On main, <tool>OBS1</tool> is missing in the first case, and the second case has <assistant>A1</assistant> duplicated.
• test_turn_without_the_reply_starts_fresh_segment. On main, the request is sliced into the old segment.

Validation

• pytest tests/unit/test_chat_server.py: 33 passed. The 3 new tests fail on main.
• Full unit suite on CPU (macOS, Python 3.12, no vLLM or AutoModel installed): 325 passed and 3 skipped, which is the 322 that pass on main plus the 3 new tests. test_refit_lora_merge.py needs nemo_automodel, and three modules import vllm, so these four do not collect here. With a local vllm stub, test_prompts_dataset.py and test_collect_rollout_rewards.py pass (17 tests). test_gather_group_lineage.py needs a local Ray cluster, which this sandbox cannot start.
• python -m compileall -q molt examples/python tests and git diff --check. On the changed files: ruff check, ruff format --check (0.9.9), isort 8.0.1 and autoflake 2.3.3, at the versions pinned in .pre-commit-config.yaml.
• The ChatML reproduction above, rerun with the fix: turn 2's context equals the template's render of the request (result a, result b and the \n close are all present).

🤖 Generated with Claude Code