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
- Enable Snowplow Micro.
- Enable billing events with the following environment variables in
.envin 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
- 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!"
}
]
}'
- 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.