Add AgentEvents::Usage for the Duo Chat usage event

What does this MR do and why?

Adds Gitlab::Duo::Chat::AgentEvents::Usage so Rails parses the terminal usage event the AI gateway emits on /v2/chat/agent, and updates the two places that reason about the event stream as a whole so a non-answer event arriving last is handled correctly.

The gateway started emitting usage in ai-assist!6185. Rails had no matching class, so AgentEventParser could not resolve it and logged Failed to find the event class in GitLab-Rails. once per round trip to the gateway. Chat responses themselves were unaffected: the parser returned nil and the executors ignored the event.

Closes #606921.

Parsing the event is the easy half

Until now every event that reached the events array was part of the answer. usage is the first that is not, and it is terminal, so it arrives last. Two consumers read the array as a whole and both needed updating. Neither failure would have been loud.

The streaming cut-off warning. add_cut_off_warning_event_if_response_cut_off read events.last twice, behind a guard matching only FinalAnswerDelta. A trailing usage means the guard never matches, so the truncation warning stops firing for every cut-off response with nothing logged. That warning tells someone their answer was cut short by a max-tokens limit or an ingress timeout, and it sits behind the duo_non_agentic_chat_ai_message_cut_off_warning flag.

The empty-stream check. ReactExecutor#execute raised EmptyEventsError on an empty array, producing error A1002 immediately. A stream carrying only usage is no longer empty, so it would pass that check, return nil from all three processors, and fall through to next unless answer. The identical request would then be reissued up to MAX_ITERATIONS (10) times before failing with ExhaustedLoopError and a misleading "reached the limit" message, at up to ten times the token cost and latency.

Both now exclude metadata through one predicate rather than repeating a class check:

# base_event.rb
def metadata?
  false
end

# usage.rb
def metadata?
  true
end
# step_executor.rb
answer_events = events.reject(&:metadata?)

# react_executor.rb
raise EmptyEventsError if events.all?(&:metadata?)

all? is correct on the empty array, so the empty-stream behaviour is unchanged and a usage-only stream returns to raising as it did before.

For the cut-off guard I used reject rather than searching backwards for the last FinalAnswerDelta. The latter reads better but changes behaviour in an unrelated case, since a stream ending in an Error produces no warning today and would begin producing one. reject is behaviour-preserving for every stream that carries no metadata event.

How to set up and validate locally

  1. Run the four spec files that cover the change:

    bundle exec rspec \
      ee/spec/lib/gitlab/duo/chat/agent_events/base_event_spec.rb \
      ee/spec/lib/gitlab/duo/chat/agent_events/usage_spec.rb \
      ee/spec/lib/gitlab/duo/chat/agent_event_parser_spec.rb \
      ee/spec/lib/gitlab/duo/chat/step_executor_spec.rb

    92 examples, 0 failures.

  2. To see the cut-off regression this prevents, change answer_events = events.reject(&:metadata?) in step_executor.rb back to answer_events = events, then run:

    bundle exec rspec ee/spec/lib/gitlab/duo/chat/step_executor_spec.rb -e "still appends the cut-off warning"

    It fails with expected #<Gitlab::Duo::Chat::AgentEvents::Usage ...> to be an instance of Gitlab::Duo::Chat::AgentEvents::FinalAnswerDelta. That is the trailing usage event pushing the answer out of the tail read, which is what silently suppresses the warning. Put the line back.

  3. To see the retry regression, change events.all?(&:metadata?) in react_executor.rb to events.empty?, then run:

    bundle exec rspec ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb -e "when only a metadata event is received"

    The example asserts error code A1002 and an EmptyEventsError reaching the tracker. With empty? restored, a usage-only stream is no longer empty, so the request is reissued instead and the example fails. Put the line back. This example needs Gitaly, so it first runs in CI here; see the note below.

Everything else that reads the events array

The remaining consumers filter by type rather than reading the tail (process_final_answer, process_tool_action and process_unknown in react_executor.rb, the Error and FinalAnswerDelta checks in its streaming block, and the Action check in StepExecutor#step), so a Usage in the array is inert for them. process_final_answer selects FinalAnswerDelta and joins in order, so whenever an answer is assembled its text is unchanged, with and without a usage event.

What I ran, for the record

Ran locally: 92 examples, 0 failures across base_event_spec.rb, usage_spec.rb, agent_event_parser_spec.rb and step_executor_spec.rb, with RuboCop clean on all four. That was a disposable source checkout with Postgres and Redis but no Gitaly, which is enough for these four because none of them touch git.

ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb builds projects with real repositories, so it needs Gitaly and its assertions first execute in this pipeline. Step 3 above is derived from the code path rather than from a run, and is the one claim here I have not executed myself.

Because a spec cannot easily show the absence of a warning being a regression, I also drove add_cut_off_warning_event_if_response_cut_off directly with rails runner, patching Feature.enabled? to force duo_non_agentic_chat_ai_message_cut_off_warning on (it defaults off) so the guarded path is reachable:

Variant Result
master plus the new event class only Cut-off warning silently dropped. The warning never enters the array, no exception is raised and nothing is logged.
Guard read moved off events.last, warning-text read left on it undefined method `finish_reason' for an instance of Gitlab::Duo::Chat::AgentEvents::Usage (NoMethodError)
This MR All checks pass, including the metadata predicate and the empty-stream guard across empty, usage-only and mixed arrays

The middle row is why both reads change and not just the guard.

Why this event is worth having in Rails

usage carries the per-request, per-model token count. Having it available in Rails is a prerequisite for storing it on Ai::UsageEvent#extras, one route to the per-user token attribution self-managed instances ask for in #466072. That issue names the need rather than the mechanism, so treat the storage route as a proposal, not a settled plan.

Counts are kept per model rather than summed, because tokens are not equivalent across models and a total would not be a meaningful figure.

Thanks @alejandro for flagging the gap and naming the file location when merging the gateway side.

Edited by Andrew Dunn

Merge request reports

Loading