fix(rollout/session): reseed instead of erroring when no assistant persisted

[DavidBellamy]

When the agent retries with a divergent user message before the first assistant turn has been persisted, _try_detect_and_rollback_to_assistant_checkpoint raises MessageValidationError because the matched prefix contains no assistant checkpoint to roll back to. This kills trials that would otherwise recover on retry.

Repro

Observed on two Harbor trials (battery-charging-optimization #6 and #7) aborted with:

litellm.BadRequestError: OpenAIException - Error code: 400 - {error: rollback failed: no assistant message found in the first 1 matched messages (stored has 2 messages, request has 3 messages)}

Chain of events:

1. Agent emits the first assistant turn, which hits max_tokens truncation and errors out before the response can be persisted.
2. Session state is therefore just [system, user] with num_assistant=0.
3. Agent retry path sends a new request whose user content diverges from the stored user message.
4. match_len=1 (system matches, user diverges); no assistant in stored[:1], so checkpoint_index=-1, and the function raises.

Fix

In the checkpoint_index < 0 branch, if num_assistant == 0 the stored state is a prompt-only seed whose first assistant turn never persisted. Clear the session and return True so the caller treats this as a fresh session. The existing raise path is preserved for cases where at least one assistant has been stored (rollback past an assistant turn remains disallowed).

The caller (prepare_pretokenized) early-returns None on the reseed signal, equivalent to the existing empty-session branch.

Diff shape

• _try_detect_and_rollback_to_assistant_checkpoint: return type None -> bool
• New num_assistant == 0 branch inside the existing if checkpoint_index < 0:
• Existing early returns get explicit return False
• prepare_pretokenized propagates the reseed signal

No behavior change for any case where at least one assistant has been persisted.

[gemini-code-assist]

Code Review

This pull request introduces a 'pre-assistant re-seed' mechanism to handle session divergences that occur before the first assistant turn is persisted. By resetting the session state instead of raising a validation error, the system can recover during retries. Feedback indicates that the new logic branch might be unreachable due to existing early return guards in 'prepare_pretokenized' and the rollback method itself, which may need adjustment to allow the re-seed logic to execute.

[gemini-code-assist]

The logic in this new branch appears to be unreachable from prepare_pretokenized given the current implementation:

1. In prepare_pretokenized (line 76), the function returns None if self.token_ids is empty. Since self.token_ids is derived from self.trajectory_token_ids, and self.num_assistant is kept in sync with the length of self.trajectory_token_ids (see update_pretokenized_state at line 145), num_assistant == 0 implies token_ids is empty. Thus, prepare_pretokenized would have already returned None before calling this method.
2. Even if called, _try_detect_and_rollback_to_assistant_checkpoint returns early at line 213 if self.trajectory_token_ids is empty (which is the case when num_assistant == 0).

If the intention is to handle reseeding when a divergence occurs before the first assistant turn is persisted, the check at line 76 in prepare_pretokenized (and line 212 here) might need to be adjusted. Additionally, when implementing logic to reset or rollback session state, ensure all generated outputs and metadata fields are cleared to their default state to prevent carrying over stale data from previous attempts.

References

1. When implementing a function to reset objects for retry, ensure all generated outputs and metadata fields are cleared to their default state to prevent carrying over stale data from previous attempts.

[guapisolo]

Thanks for contribution. I understood that this PR intended to handle "The agent retries with a divergent user message" case. But we think that current session server design should only handle append-only trajectory... and diverge should be limited to one step