fix: allow TITO session rollback to the empty checkpoint

[Shi-Dong]

Summary

A TITO session permanently rejects every request once a client re-sends the very first turn. The session becomes unusable and the whole trajectory scores reward 0.

The failure

_try_detect_and_rollback_to_assistant_checkpoint rolls a session back to the last assistant message inside the prefix that the incoming request still shares with the stored history. Retrying the first turn is the one case where that prefix cannot contain an assistant:

	messages
stored after turn 1	[user, assistant₁]
request re-sent	[user]

match_len is 1 and stored[0] is a user message, so there was no rollback target and the request was rejected with:

rollback failed: no assistant message found in the first 1 matched messages
(stored has 2 messages, request has 1 messages)

From that point the session is poisoned. The stored history still holds an assistant the client does not know about, so every subsequent request diverges the same way and 400s again. In the agentic run that surfaced this, 1446 of 1504 /sessions/*/v1/chat/completions calls returned 400 and every trial ended with reward 0.

What re-sends the first turn

Tracing one session end-to-end in that run gives an exact account. Its access log is one 200, then nine 400s — and the 200 is on a different connection (source port 46994) from every one of the 400s (all on 47040):

1. terminus-2 issues its first LLM call. The session server answers 200 and stores [user, assistant₁].
2. The client never successfully reads that response — the connection it arrived on is never used again.
3. The OpenAI SDK's own max_retries re-sends the request byte-for-byte on a fresh connection, which is why match_len == 1. It gets the 400.
4. BadRequestError is not retryable, so it propagates up into two nested tenacity decorators — harbor's LiteLLM.call (3 attempts) and terminus-2's _query_llm (3 attempts, no backoff). Those 3 × 3 = 9 attempts are the nine 400s.

The counts across the whole run match that structure exactly: 164 distinct sessions, each with precisely nine 400s. That uniformity is what rules out flakiness — this is a deterministic retry cascade, not an intermittent one. The 200 plus nine 400s also means ten HTTP requests against nine tenacity slots, which is the extra SDK-level re-send in step 3.

The trial artifacts corroborate step 2 from the agent's side: trajectory.json contains a single step with source: "user" and no assistant at all, final_metrics reports total_completion_tokens: 0, and trial.log logs Unknown Error in LLM interaction: ... 400 ... exactly three times — once per _query_llm attempt. The agent never saw a usable completion, even though the server had produced and stored one.

So the trigger is a client-side re-send of a request the server considered successful. Why the client fails to read that first response is a separate defect, still under investigation, and not fixed here. What is fixed here is that the re-send is unrecoverable.

Ruled out so far, for whoever picks that up:

• Response-header passthrough. _DROP_RESPONSE_HEADERS already strips content-length / transfer-encoding / content-encoding, and the deployed build has it.
• Turn-1 output truncation. finish_reason == "length" would make LiteLLM.call raise OutputLengthExceededError, which is a genuine verbatim-re-send path (and one that _query_llm's tenacity retries, since it only excludes ContextLengthExceededError). But no trial artifact or server log contains any truncation signal, so it is not what fired here.
• Streaming / the fake-SSE response path. harbor never sets stream, and LiteLLM.call raises on CustomStreamWrapper, so the plain-JSON branch is the one in use.
• Client timeouts and upstream 5xx. Zero timeouts anywhere, and no 5xx precedes any rollout 400.

Two corrections to an earlier revision of this description, which got the trigger wrong:

• It claimed an upstream 5xx caused an SDK retry. It did not. The Retrying request to /chat/completions in 0.409547 seconds line and the 500s I quoted came from a synthetic probe I ran myself at 06:51:56, sixteen minutes after the failure, against an engine that had already crashed — not from the rollouts. No 5xx and no timeout precedes any rollout 400.
• response_length_exceeded_policy is abort in this run (HARBOR_RESPONSE_LENGTH_POLICY=abort), so the regenerate path was not involved either.

That second path is still worth closing, though: response_length_exceeded_policy="regenerate" re-sends the history without the truncated assistant turn and relies on exactly this rollback to discard it (the comment in harbor's terminus_2.py says so). A response-length overflow on turn 1 walks into the same hole by design.

The fix

Treat "no assistant in the matched prefix" as a rollback to the empty checkpoint instead of an error. Session state resets, and prepare_pretokenized re-renders the prompt from scratch — exactly what it already does for a genuine first turn. Turn 1 now regenerates the same way any later turn does.

That required moving the rollback check in prepare_pretokenized above the "nothing stored yet" early return, because the rollback itself is what empties the session.

MAX_ASSISTANT_ROLLBACK_STEPS still bounds the discard, so this only succeeds while exactly one assistant has been stored. A session holding two or more stored assistants is still rejected, with state left untouched.

The docs line describing the append-only contract is updated to match.

Scope

This is a robustness fix, not a root-cause fix. A session has to survive a duplicate turn-1 request whatever produced it — at-least-once delivery is normal HTTP client behaviour, and today the first such retry on a session is unrecoverable and costs the entire trajectory.

The guard is also not a recent regression: it arrived together with _try_detect_and_rollback_to_assistant_checkpoint in #783 (2026-03-25), and MAX_ASSISTANT_ROLLBACK_STEPS in #857 (2026-04-02). The hole has been latent since March; it only became fatal once a harness started retrying turn 1.

Test plan

• tests/fast/router/test_linear_trajectory.py — 42 passed. Replaced the test that asserted the old rejection, and added coverage for a verbatim first-turn re-send (regenerates and commits cleanly) plus the >1-assistant case still raising with state intact.
• tests/fast/router tests/fast/utils — 2585 passed, 1 skipped. No regressions.
• black, isort and ruff check clean.
• End-to-end on a live agentic run — the patch is staged on the devbox but the GLM-4.7-Flash / terminus-2 job still needs relaunching to confirm non-zero rollout/raw_reward.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[guapisolo]

Good fix. LGTM.