agent: optional exact token-ID capture for token-level training

[akhatua2]

Adds an opt-in capture_token_ids flag to LitellmModelConfig. Off by default — existing runs are byte-identical.

Why

Token-level training (on-policy distillation, TITO-style SFT) needs the tokens the model actually saw. Reconstructing them afterwards by re-tokenizing a finished trajectory is not equivalent:

• BPE is non-injective (HAVING can tokenize as H+AVING or HAV+ING)
• tool-call serialization can differ between inference and training (whitespace, JSON key order)
• under context compaction, the prompt a turn saw no longer exists in the final message list

Capturing at request time makes the training input identical to the inference input by construction.

Concretely: a scan of a 10-pair flash_10 run found 1942 assistant turns and zero usable states, because there was no way to record ids.

What it does

With capture_token_ids: true, requests carry extra_body={"return_token_ids": true} and the returned ids land on each assistant message as:

message["extra"]["token_capture"] = {"prompt_token_ids": [...], "output_token_ids": [...]}

which persists into the saved trajectory.

Details worth noting:

• _with_return_token_ids merges into any existing extra_body instead of replacing it.
• _extract_token_capture handles both placements seen in practice — vLLM puts prompt_token_ids on the response and token_ids on the choice; some versions nest both under the choice — and returns None rather than partial data.
• A server that ignores the flag logs a warning. Silently returning no ids otherwise surfaces much later as a training set with zero usable rows.
• The summarizer path is deliberately not captured; those turns are excluded from training anyway.

Testing

tests/agents: 121 passed. The 2 failures in test_team_wiring.py are a missing optional openhands.tools.task_tracker.coop_definition module and reproduce on clean main.

Verified directly:

capture_token_ids default: False
settable via config      : True
extra_body merge         : {'extra_body': {'a': 1, 'return_token_ids': True}}
capture extraction       : {'prompt_token_ids': [1, 2], 'output_token_ids': [3]}

plus both response placements, missing ids, and empty-list ids (all correctly rejected).

Requires vLLM >= 0.10.2 or SGLang with return_token_ids support. Bumps version to 0.0.20.