Draft: fix: resolves sending billing event for proxy requests on response body

What does this merge request do and why?

In https://gitlab.com/gitlab-org/gitlab/-/work_items/574417, we added instrumentation to ensure our v1/proxy/<provider> requests appeared in raw_usage_billing events in Clickhouse. However, we noticed that these billing events didn't actually show up in staging with a valid event_type="ai_gateway_proxy_use".

The reason we didn't see this information is because we were checking for response.body existence. However, the usage data from the provider was actually parsed earlier, and response.body doesn't exist at that point.

        # Extract token usage from response for billing tracking
        # Only extract for non-streaming responses
        if not stream and response_from_upstream.status_code == 200:
            try:
                response_json = response_from_upstream.json()
                token_usage = self._extract_token_usage(upstream_path, response_json)
                # Store in request state for billing event tracking
                request.state.proxy_token_usage = token_usage
                request.state.proxy_model_name = model_name
            except Exception:
                # If token extraction fails, continue without it
                # The billing event will still be tracked without token details
                pass

To ensure the billing client gets called, I've added additional logging to verify the billing client is actually invoked.

Kibana logs supporting that response.body is not present: https://log.gprd.gitlab.net/app/discover#/?_g=h@3d4aa0d&_a=h@6fe9f77

AI Generated Code Flow Explanation

track_billing_event decorator
  └─> calls func() (proxy endpoint)
       └─> calls BaseProxyClient.proxy()
            └─> sends request to upstream
            └─> receives response
            └─> extracts tokens  stores in request.state 
            └─> returns Response object
  └─> receives Response object 
  └─> reads request.state.proxy_token_usage 
  └─> tracks billing event
  └─> returns Response to client

Kibana logs supporting this: https://log.gprd.gitlab.net/app/discover#/?_g=h@3d4aa0d&_a=h@6fe9f77

How to set up and validate locally

  1. Enable Snowplow Micro.
  2. Enable billing events with the following environment variables in .env in AI Gateway:
# Billing Events
AIGW_BILLING_EVENT__ENABLED=true
AIGW_BILLING_EVENT__ENDPOINT=http://localhost:9091 # snowplow micro url
AIGW_BILLING_EVENT__NAMESPACE=gl
AIGW_BILLING_EVENT__BATCH_SIZE=1
AIGW_BILLING_EVENT__THREAD_COUNT=1
  1. Make a proxy request:
curl -X POST "http://gdk.test:5052/v1/proxy/anthropic/v1/messages" \
  -H "Content-Type: application/json" \
  -H 'bypass-auth: true' \
  -H "X-Gitlab-Realm: saas" \
  -H "X-Gitlab-Unit-Primitive: ai_gateway_model_provider_proxy" \
  -H "X-Gitlab-Enabled-Feature-Flags: duo_use_billing_endpoint" \
  -d '{
    "model": "claude-sonnet-4-5-20250929", 
    "max_tokens": 1024,
    "messages": [
      {
        "role": "user",
        "content": "Hello, Claude!"
      }
    ]
  }'
  1. Check http://localhost:9091/micro/ui to verify that the billable event was registered.

Merge request checklist

  • Tests added for new functionality. If not, please raise an issue to follow up.
  • Documentation added/updated, if needed.
  • If this change requires executor implementation: verified that issues/MRs exist for both Go executor and Node executor, or confirmed that changes are backward-compatible and don't break existing executor functionality.
Edited by Nathan Weinshenker

Merge request reports

Loading