Fix tau-bench token deltas for reasoning templates

[DongzhuoranZhou]

Summary

• preserve complete user-message token deltas when a reasoning-aware chat template rewrites earlier assistant history
• keep the existing fast prefix-slice path for append-only assistant/tool turns
• fail closed when an unexpected non-user history rewrite cannot be represented safely
• add a self-contained regression test for Qwen3-style historical-reasoning removal

Root cause

The tau-bench example computed each newly appended message with:

new_text = curr[len(prev):]

This assumes the newly rendered conversation always starts with the previous rendering. Qwen3 reasoning templates violate that assumption: after a new real user message arrives, the template removes the preceding assistant reasoning from curr. Slicing curr at len(prev) therefore starts inside, or after, the new user message. The resulting GRPO sample.tokens can omit user context while its remaining mask lengths still align.

Fix

When curr.startswith(prev), behavior is unchanged. When the latest message is a user message and history was rewritten, render that user message independently and verify that it is exactly the suffix of curr before tokenizing it. Other unexpected rewrite patterns raise an error instead of silently producing corrupt training data.

This does not change whether prior reasoning is retained in the accumulated GRPO training context; it only prevents new user messages from being truncated.

Validation

• pytest -q tests/test_tau_bench_token_delta.py — 2 passed
• pre-commit hooks on all changed files — passed
• py_compile on the helper and tau-bench agent — passed
• local probe with the real Qwen3-4B-Thinking-2507 tokenizer over five assistant/user turns:
  • every later user turn triggered the historical rewrite
  • every user message remained complete and appeared exactly once
  • response token and loss-mask lengths matched (255/255)