Log AI Gateway token usage per model in llm.log

What does this MR do and why?

The Received response from AI Gateway entry in llm.log reads the AI Gateway's token counts from usage.input_tokens and usage.output_tokens as flat keys, but /v1/prompts answers with a string and /v2/prompts nests the counts under each model's name, so those fields are never filled.

This MR logs one token_usage field instead, with one entry for each model:

"token_usage": [{ "model_name": "claude-sonnet-4-5", "input_tokens": 100, "output_tokens": 42 }]

The keys are the same whichever models answer, as the Log Field Standards review rule asks of new fields, and match the names the AI Gateway logs for these counts. The rule also asks for new fields in Labkit::Fields, which has no token or model field yet. Adding one is a change to the LabKit spec that the field standardisation design describes, outside this repository.

A streamed /v2/prompts/chat answer arrives as newline-delimited frames, and several frames do not parse as one document. PromptResponseParser now keeps the usage the last frame reports, and Chain::Requests::AiGateway logs it. The client reads usage only from a response that is not streamed, so each answer's usage is logged once.

flowchart LR
  A["POST /v2/prompts/:prompt_id"] --> B["Client#35;log_response_received"]
  C["POST /v2/prompts/chat, streamed"] --> B
  C --> D["PromptResponseParser#35;usage"]
  D --> E["Chain::Requests::AiGateway#35;request"]
  B -->|"not streamed"| F["token_usage_fields"]
  E --> F
  F --> G["llm.log response_received<br/>token_usage, one entry per model"]

Danger's AI logging warning: both changed entries use log_conditional_info, which drops extra fields unless expanded AI logging is on, and token_usage holds only model names and token counts.

No changelog: only /v2/prompts responses carry usage, and ai_prompts_v2 is disabled by default (https://docs.gitlab.com/development/feature_flags/#changelog).

References

How to set up and validate locally

  1. In a Rails console, Feature.enable(:expanded_ai_logging) and Feature.enable(:ai_prompts_v2, User.find_by_username('root')).
  2. In Duo Chat, with the Agentic toggle off, /explain a code snippet. In log/llm.log, Made request to AI Client carries token_usage and that request's Received response from AI Gateway does not.
  3. Generate a commit message in a merge request. Its Received response from AI Gateway entry carries token_usage.

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist.

An AI coding assistant wrote this change and ran the changed specs in a disposable development environment.

Edited by Andrew Dunn

Merge request reports

Loading
Loading