Record GitLab Duo Non-Agentic Chat token usage per agent turn

What does this MR do and why?

GitLab Duo Non-Agentic Chat answers in agent turns, and the AI Gateway ends each completed turn with a usage frame of tokens per model. Rails parses it but stores no usage event, so administrators cannot query each user's Chat tokens.

This MR records two chat AI usage events, which the instance aiUsageData query returns:

  • report_duo_chat_token_usage, per agent turn that reports usage: the frame's per-model map, as sent, in extras.usage, with request_id.
  • complete_duo_chat_response, once per question: request_id and, if the response failed, its error_code.

usage and request_id are keyword arguments, so only the AI usage tracker records them, and the track_events API cannot set them.

Like request_duo_chat_response, both events pass a namespace, so a question on a group work item writes the group's namespace_id.

Both events' aiUserMetrics count fields start counting, and Chat's totalEventCount triples for a one-turn question.

A new instance setting, duo_token_usage_tracking_enabled, on by default, turns the token event off. It is a key in ai_settings.feature_settings, so there is no migration.

This MR depends on !259818 (merged), which defines and registers events 54 and 55 and adds duo_token_usage_tracking_enabled to the ai_settings.feature_settings schema a release earlier, so a node on the previous release can drain, read and save what this MR writes. It merged on 7 October, and this MR should merge only after it is in a release candidate (its released::candidate label). With no feature flag, GitLab.com records both events from this MR's deploy, and GitLab Self-Managed and GitLab Dedicated from the release that carries it.

How a turn becomes a row

sequenceDiagram
    participant GW as AI Gateway /v2/chat/agent
    participant SE as StepExecutor step
    participant RE as ReactExecutor
    participant AT as Gitlab::Tracking::AiTracking
    participant DB as ai_usage_events
    loop Each agent turn, up to ten per question
        GW->>SE: usage frame, tokens per model
        SE->>RE: AgentEvents::Usage, to the step_forward block
        opt duo_token_usage_tracking_enabled
            RE->>AT: report_duo_chat_token_usage
        end
    end
    RE->>AT: complete_duo_chat_response from execute
    AT->>DB: one row per event

Limits

  • A self-hosted model that reports no usage gets an AI Gateway LiteLLM estimate, unmarked.
  • A tool's own model call, such as the issue reader's or the documentation tool's, is not metered. Slash commands (/explain, /refactor, /tests, /fix and the others) record neither event, because Completions::Chat hands them to the command's tool executor (chat.rb:148, chat.rb:178) and ReactExecutor never runs. Both are left for follow-up in #466072.
  • DuoCodeReviewChatWorker answers @GitLabDuo mentions with this executor, writing both events, except review requests and, in :dap mode, assistant_flow ones.
  • A Sidekiq re-run of a Chat job reuses the request_id, and no field marks the re-run.
  • A retried turn records only the retry's tokens, because the AI Gateway sends no frame after an error.
  • Group-level aiUsageData does not expose extras.
  • Chat questions that IDEs send through the GitLab Language Server record both events with no namespace, as request_duo_chat_response does, because the language server sends no resource with them (gitlab_chat_api.ts:349). The instance aiUsageData query returns these rows and group queries do not. #572546 tracks Chat's missing namespace.
  • With AI logs on (enabled_expanded_logging), llm.log still records the usage frame when duo_token_usage_tracking_enabled is off.

References

Screenshots or screen recordings

Admin > GitLab Duo > Change configuration:

Before After
Data and privacy group on master Data and privacy group with the new checkbox

How to set up and validate locally

With ClickHouse running:

  1. bundle exec rspec ee/spec/lib/gitlab/tracking/ai_tracking_spec.rb ee/spec/lib/ee/gitlab/tracking/event_eligibility_checker_spec.rb ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb ee/spec/lib/gitlab/duo/chat/step_executor_spec.rb ee/spec/graphql/types/analytics/ai_usage/ai_usage_event_type_enum_spec.rb spec/lib/gitlab/tracking/event_definition_spec.rb spec/lib/gitlab/usage/metric_definition_validate_all_spec.rb ee/spec/config/metrics/every_metric_definition_spec.rb ee/spec/requests/api/graphql/analytics/ai_analytics/ai_user_metrics_spec.rb ee/spec/models/analytics/aggregation_engines/ai_usage_events_spec.rb ee/spec/models/ai/setting_spec.rb ee/spec/models/application_setting_spec.rb ee/spec/lib/ee/api/entities/application_setting_spec.rb ee/spec/presenters/admin/ai_configuration_presenter_spec.rb ee/spec/requests/api/settings_spec.rb passes.
  2. Jest on ee/spec/frontend/ai/settings/components/duo_token_usage_tracking_form_spec.js, ee/spec/frontend/ai/settings/pages/ai_admin_settings_spec.js and ee/spec/frontend/ai/settings/index_spec.js passes.
  3. In a development Rails console, after firing each event once, require_relative 'spec/support/helpers/service_ping_helpers.rb'; ServicePingHelpers.get_current_usage_metric_value('counts.count_total_report_duo_chat_token_usage_monthly') and the same call for counts.count_total_complete_duo_chat_response_monthly each return 1.

On a self-managed GitLab 19.4.1 instance on the cloud native chart, with GitLab Duo Self-Hosted and this MR's backend changes applied (which then included the registration and event definitions now in !259818 (merged)), each turn started from the GraphQL aiAction mutation the chat panel sends and asked Chat to summarize an issue. A stub model reported a different count on every call, so each token row matches one call:

  • On 19.4.1 as released, the turn wrote only request_duo_chat_response.
  • With this MR, it wrote two token rows, each map equal to its agent call, such as {"usage":{"<model>":{"input_tokens":5518,"output_tokens":382}},"request_id":"..."}, and one completion row, all on the mutation's request_id.
  • With duo_token_usage_tracking_enabled set to false through PUT /api/v4/application/settings, the turn answered and wrote its completion row and no token row.
  • The instance aiUsageData query returned each map unchanged in extras, with the project's namespacePath, and the group query returned the same counts per event.

Database

duo_token_usage_tracking_enabled= makes the same lookup and update as ai_audit_events_streaming_enabled=. It loads the default organization, finds the organization's ai_settings row through the unique organization_id index, and updates it by primary key. If that row does not exist, for_organization creates it first, as it does for the other setters. The plans are from Database Lab on gitlab-production-main, where the update set feature_settings to this MR's key alone. Each link opens a rerun of its statement.

Update (update!(duo_token_usage_tracking_enabled: value)):

UPDATE "ai_settings"
SET "updated_at" = '2026-10-03 19:12:10.936306',
    "feature_settings" = '{"duo_cli_enabled":true,"duo_agent_platform_enabled":true,"duo_token_usage_tracking_enabled":false,"ai_audit_events_streaming_enabled":false}'
WHERE "ai_settings"."id" = 1
Plan: one row through ai_settings_pkey, 5.665 ms
 Update on public.ai_settings  (cost=0.13..3.15 rows=0 width=0) (actual time=5.593..5.593 rows=0 loops=1)
   Buffers: shared hit=31 read=9 dirtied=2
   WAL: records=2 fpi=1 bytes=1191
   I/O Timings: read=5.242 write=0.000
   ->  Index Scan using ai_settings_pkey on public.ai_settings  (cost=0.13..3.15 rows=1 width=46) (actual time=0.768..0.770 rows=1 loops=1)
         Index Cond: (ai_settings.id = 1)
         Buffers: shared hit=4 read=1
         I/O Timings: read=0.732 write=0.000

Database Lab

Lookups, the same as for the other setters: 1.723 ms and 1.828 ms
SELECT "organizations".* FROM "organizations" WHERE "organizations"."id" = 1 LIMIT 1
 Limit  (cost=0.14..3.16 rows=1 width=78) (actual time=1.664..1.665 rows=1 loops=1)
   Buffers: shared hit=3 read=2
   I/O Timings: read=1.624 write=0.000
   ->  Index Scan using organizations_pkey on public.organizations  (cost=0.14..3.16 rows=1 width=78) (actual time=1.663..1.663 rows=1 loops=1)
         Index Cond: (organizations.id = 1)
         Buffers: shared hit=3 read=2
         I/O Timings: read=1.624 write=0.000

Database Lab

SELECT "ai_settings"."id", "ai_settings"."created_at", "ai_settings"."updated_at", "ai_settings"."amazon_q_oauth_application_id", "ai_settings"."amazon_q_service_account_user_id", "ai_settings"."amazon_q_ready", "ai_settings"."amazon_q_role_arn", "ai_settings"."duo_workflow_service_account_user_id", "ai_settings"."duo_workflow_oauth_application_id", "ai_settings"."duo_core_features_enabled", "ai_settings"."foundational_agents_default_enabled", "ai_settings"."minimum_access_level_execute", "ai_settings"."minimum_access_level_manage", "ai_settings"."minimum_access_level_enable_on_projects", "ai_settings"."minimum_access_level_execute_async", "ai_settings"."feature_settings", "ai_settings"."include_recommended_allowed", "ai_settings"."allow_all_unix_sockets", "ai_settings"."enforce_on_local_clients", "ai_settings"."allow_project_extension", "ai_settings"."allowed_domains", "ai_settings"."denied_domains", "ai_settings"."organization_id" FROM "ai_settings" WHERE "ai_settings"."organization_id" = 1 LIMIT 1
 Limit  (cost=0.13..3.15 rows=1 width=183) (actual time=1.756..1.758 rows=1 loops=1)
   Buffers: shared hit=3 read=2
   I/O Timings: read=1.705 write=0.000
   ->  Index Scan using index_ai_settings_on_organization_id on public.ai_settings  (cost=0.13..3.15 rows=1 width=183) (actual time=1.755..1.755 rows=1 loops=1)
         Index Cond: (ai_settings.organization_id = 1)
         Buffers: shared hit=3 read=2
         I/O Timings: read=1.705 write=0.000

Database Lab

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

An AI coding assistant wrote this change, ran the RSpec and Jest specs above, and tested the behavior on a self-managed GitLab 19.4.1 instance on the cloud native chart.

Edited by Andrew Dunn

Merge request reports

Loading
Loading