Feature: Add feature flag for token usage tracking#44
Merged
SantiagoDePolonia merged 1 commit intofeature/usage-token-trackingfrom Jan 21, 2026
Merged
Conversation
Add a new configuration flag that automatically adds stream_options:
{"include_usage": true} to streaming requests for OpenAI-compatible
providers. This ensures usage data is always returned in streaming
responses when the flag is enabled (default: true).
Changes:
- Add StreamOptions struct to core/types.go for stream_options param
- Add StreamOptions field to ChatRequest and ResponsesRequest
- Update WithStreaming() methods to preserve StreamOptions
- Add EnforceReturningUsageData to usage.Config and config.UsageConfig
- Update handlers to set StreamOptions when flag is enabled
- Document ENFORCE_RETURNING_USAGE_DATA in .env.template
Contributor
📝 WalkthroughWalkthroughThis PR introduces functionality to control usage data inclusion in streaming responses through a new configuration option. A Changes
Estimated Code Review Effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In @.env.template:
- Around line 87-90: Rename the environment variable
ENFORCE_RETURNING_USAGE_DATA to USAGE_ENFORCE_RETURNING_USAGE_DATA to match the
other usage-prefixed settings and the Viper key
usage.enforce_returning_usage_data; update the .env.template entry text and any
code/config that reads ENFORCE_RETURNING_USAGE_DATA (search for that symbol) so
it reads USAGE_ENFORCE_RETURNING_USAGE_DATA and ensure Viper/env binding or
defaults still map to usage.enforce_returning_usage_data; also update any
documentation or tests that reference the old name.
In `@config/config.go`:
- Around line 420-422: Update the environment override in the init logic to use
the new variable name: replace the os.Getenv lookup for
"ENFORCE_RETURNING_USAGE_DATA" with "USAGE_ENFORCE_RETURNING_USAGE_DATA" so
cfg.Usage.EnforceReturningUsageData is set from the renamed env var (keeping the
existing strings.EqualFold(..., "true") || value == "1" logic); adjust any
surrounding comment/text referencing the old name if present and ensure the
change is made where the code sets cfg.Usage.EnforceReturningUsageData.
- Around line 296-302: The UsageConfig initialization uses a mismatched env var
name for EnforceReturningUsageData; update the call in the UsageConfig block to
read getEnvBoolOrDefault("USAGE_ENFORCE_RETURNING_USAGE_DATA", true) instead of
"ENFORCE_RETURNING_USAGE_DATA" so the field name and other usage env vars share
the USAGE_ prefix; ensure the change is applied where UsageConfig is constructed
and any docs/tests referencing the old env var are updated accordingly.
In `@internal/server/handlers.go`:
- Around line 98-104: Extract a helper on the Handler type to avoid duplicated
enforcement logic: add a method like shouldEnforceUsageData() bool that returns
h.usageLogger != nil && h.usageLogger.Config().EnforceReturningUsageData, then
replace the duplicated blocks in the current handler and the Responses handler
with a call to that method and the existing code that ensures req.StreamOptions
is non-nil and sets req.StreamOptions.IncludeUsage = true; ensure both sites
preserve the nil check for req.StreamOptions before setting IncludeUsage.
Comment on lines
+87
to
+90
| # Enforce returning usage data in streaming responses (default: true) | ||
| # When true, stream_options: {"include_usage": true} is automatically added | ||
| # to streaming requests for OpenAI-compatible providers | ||
| # ENFORCE_RETURNING_USAGE_DATA=true |
Contributor
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Environment variable naming inconsistency.
Other usage-related environment variables use the USAGE_ prefix (e.g., USAGE_ENABLED, USAGE_BUFFER_SIZE), but ENFORCE_RETURNING_USAGE_DATA does not. Consider renaming to USAGE_ENFORCE_RETURNING_USAGE_DATA for consistency.
Note: The Viper default in config/config.go already uses usage.enforce_returning_usage_data, suggesting the USAGE_ prefix would be more appropriate.
Suggested naming
-# ENFORCE_RETURNING_USAGE_DATA=true
+# USAGE_ENFORCE_RETURNING_USAGE_DATA=true🤖 Prompt for AI Agents
In @.env.template around lines 87 - 90, Rename the environment variable
ENFORCE_RETURNING_USAGE_DATA to USAGE_ENFORCE_RETURNING_USAGE_DATA to match the
other usage-prefixed settings and the Viper key
usage.enforce_returning_usage_data; update the .env.template entry text and any
code/config that reads ENFORCE_RETURNING_USAGE_DATA (search for that symbol) so
it reads USAGE_ENFORCE_RETURNING_USAGE_DATA and ensure Viper/env binding or
defaults still map to usage.enforce_returning_usage_data; also update any
documentation or tests that reference the old name.
Comment on lines
296
to
302
| Usage: UsageConfig{ | ||
| Enabled: getEnvBoolOrDefault("USAGE_ENABLED", true), | ||
| BufferSize: getEnvIntOrDefault("USAGE_BUFFER_SIZE", 1000), | ||
| FlushInterval: getEnvIntOrDefault("USAGE_FLUSH_INTERVAL", 5), | ||
| RetentionDays: getEnvIntOrDefault("USAGE_RETENTION_DAYS", 90), | ||
| Enabled: getEnvBoolOrDefault("USAGE_ENABLED", true), | ||
| EnforceReturningUsageData: getEnvBoolOrDefault("ENFORCE_RETURNING_USAGE_DATA", true), | ||
| BufferSize: getEnvIntOrDefault("USAGE_BUFFER_SIZE", 1000), | ||
| FlushInterval: getEnvIntOrDefault("USAGE_FLUSH_INTERVAL", 5), | ||
| RetentionDays: getEnvIntOrDefault("USAGE_RETENTION_DAYS", 90), | ||
| }, |
Contributor
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Environment variable naming inconsistency with legacy path.
The legacy environment-based initialization uses ENFORCE_RETURNING_USAGE_DATA while other usage-related variables use the USAGE_ prefix (e.g., USAGE_ENABLED, USAGE_BUFFER_SIZE). This creates an inconsistency:
| Field | Env Var |
|---|---|
| Enabled | USAGE_ENABLED |
| EnforceReturningUsageData | ENFORCE_RETURNING_USAGE_DATA |
| BufferSize | USAGE_BUFFER_SIZE |
Consider using USAGE_ENFORCE_RETURNING_USAGE_DATA for consistency.
Suggested fix
Usage: UsageConfig{
Enabled: getEnvBoolOrDefault("USAGE_ENABLED", true),
- EnforceReturningUsageData: getEnvBoolOrDefault("ENFORCE_RETURNING_USAGE_DATA", true),
+ EnforceReturningUsageData: getEnvBoolOrDefault("USAGE_ENFORCE_RETURNING_USAGE_DATA", true),
BufferSize: getEnvIntOrDefault("USAGE_BUFFER_SIZE", 1000),📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| Usage: UsageConfig{ | |
| Enabled: getEnvBoolOrDefault("USAGE_ENABLED", true), | |
| BufferSize: getEnvIntOrDefault("USAGE_BUFFER_SIZE", 1000), | |
| FlushInterval: getEnvIntOrDefault("USAGE_FLUSH_INTERVAL", 5), | |
| RetentionDays: getEnvIntOrDefault("USAGE_RETENTION_DAYS", 90), | |
| Enabled: getEnvBoolOrDefault("USAGE_ENABLED", true), | |
| EnforceReturningUsageData: getEnvBoolOrDefault("ENFORCE_RETURNING_USAGE_DATA", true), | |
| BufferSize: getEnvIntOrDefault("USAGE_BUFFER_SIZE", 1000), | |
| FlushInterval: getEnvIntOrDefault("USAGE_FLUSH_INTERVAL", 5), | |
| RetentionDays: getEnvIntOrDefault("USAGE_RETENTION_DAYS", 90), | |
| }, | |
| Usage: UsageConfig{ | |
| Enabled: getEnvBoolOrDefault("USAGE_ENABLED", true), | |
| EnforceReturningUsageData: getEnvBoolOrDefault("USAGE_ENFORCE_RETURNING_USAGE_DATA", true), | |
| BufferSize: getEnvIntOrDefault("USAGE_BUFFER_SIZE", 1000), | |
| FlushInterval: getEnvIntOrDefault("USAGE_FLUSH_INTERVAL", 5), | |
| RetentionDays: getEnvIntOrDefault("USAGE_RETENTION_DAYS", 90), | |
| }, |
🤖 Prompt for AI Agents
In `@config/config.go` around lines 296 - 302, The UsageConfig initialization uses
a mismatched env var name for EnforceReturningUsageData; update the call in the
UsageConfig block to read
getEnvBoolOrDefault("USAGE_ENFORCE_RETURNING_USAGE_DATA", true) instead of
"ENFORCE_RETURNING_USAGE_DATA" so the field name and other usage env vars share
the USAGE_ prefix; ensure the change is applied where UsageConfig is constructed
and any docs/tests referencing the old env var are updated accordingly.
Comment on lines
+420
to
+422
| if enforceUsage := os.Getenv("ENFORCE_RETURNING_USAGE_DATA"); enforceUsage != "" { | ||
| cfg.Usage.EnforceReturningUsageData = strings.EqualFold(enforceUsage, "true") || enforceUsage == "1" | ||
| } |
Contributor
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Environment variable override also needs renaming.
If you rename the environment variable to USAGE_ENFORCE_RETURNING_USAGE_DATA as suggested, update this override block as well.
Suggested fix
-if enforceUsage := os.Getenv("ENFORCE_RETURNING_USAGE_DATA"); enforceUsage != "" {
+if enforceUsage := os.Getenv("USAGE_ENFORCE_RETURNING_USAGE_DATA"); enforceUsage != "" {
cfg.Usage.EnforceReturningUsageData = strings.EqualFold(enforceUsage, "true") || enforceUsage == "1"
}🤖 Prompt for AI Agents
In `@config/config.go` around lines 420 - 422, Update the environment override in
the init logic to use the new variable name: replace the os.Getenv lookup for
"ENFORCE_RETURNING_USAGE_DATA" with "USAGE_ENFORCE_RETURNING_USAGE_DATA" so
cfg.Usage.EnforceReturningUsageData is set from the renamed env var (keeping the
existing strings.EqualFold(..., "true") || value == "1" logic); adjust any
surrounding comment/text referencing the old name if present and ensure the
change is made where the code sets cfg.Usage.EnforceReturningUsageData.
Comment on lines
+98
to
+104
| // Enforce returning usage data in streaming responses if configured | ||
| if h.usageLogger != nil && h.usageLogger.Config().EnforceReturningUsageData { | ||
| if req.StreamOptions == nil { | ||
| req.StreamOptions = &core.StreamOptions{} | ||
| } | ||
| req.StreamOptions.IncludeUsage = true | ||
| } |
Contributor
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
LGTM! Consider extracting helper to reduce duplication.
The logic correctly ensures StreamOptions.IncludeUsage is set when EnforceReturningUsageData is enabled. The null check on h.usageLogger is proper defensive programming.
Note: This logic is duplicated in the Responses handler (lines 164-170). Consider extracting a helper method if more endpoints require this pattern.
Optional: Extract helper method
// ensureStreamOptionsIncludeUsage sets IncludeUsage if EnforceReturningUsageData is configured.
func (h *Handler) shouldEnforceUsageData() bool {
return h.usageLogger != nil && h.usageLogger.Config().EnforceReturningUsageData
}Then use in handlers:
if req.Stream {
- if h.usageLogger != nil && h.usageLogger.Config().EnforceReturningUsageData {
+ if h.shouldEnforceUsageData() {
if req.StreamOptions == nil {
req.StreamOptions = &core.StreamOptions{}
}
req.StreamOptions.IncludeUsage = true
}🤖 Prompt for AI Agents
In `@internal/server/handlers.go` around lines 98 - 104, Extract a helper on the
Handler type to avoid duplicated enforcement logic: add a method like
shouldEnforceUsageData() bool that returns h.usageLogger != nil &&
h.usageLogger.Config().EnforceReturningUsageData, then replace the duplicated
blocks in the current handler and the Responses handler with a call to that
method and the existing code that ensures req.StreamOptions is non-nil and sets
req.StreamOptions.IncludeUsage = true; ensure both sites preserve the nil check
for req.StreamOptions before setting IncludeUsage.
5d26044
into
feature/usage-token-tracking
1 check passed
SantiagoDePolonia
added a commit
that referenced
this pull request
Jan 21, 2026
* feat: added token usage tracking * docs: saved executed plan for implementing tokens usage * fix: enabled by default - usage tracking feature * fix: actual anthropic models added * fix: get anthropic model dynamically * fix: pass X-Request-ID * feat: add ENFORCE_RETURNING_USAGE_DATA flag for streaming responses (#44) Add a new configuration flag that automatically adds stream_options: {"include_usage": true} to streaming requests for OpenAI-compatible providers. This ensures usage data is always returned in streaming responses when the flag is enabled (default: true). Changes: - Add StreamOptions struct to core/types.go for stream_options param - Add StreamOptions field to ChatRequest and ResponsesRequest - Update WithStreaming() methods to preserve StreamOptions - Add EnforceReturningUsageData to usage.Config and config.UsageConfig - Update handlers to set StreamOptions when flag is enabled - Document ENFORCE_RETURNING_USAGE_DATA in .env.template Co-authored-by: Claude <noreply@anthropic.com> * fix: replace fixed sleep with polling loop in TestLogger Use time.After for timeout and time.Ticker for polling instead of a fixed time.Sleep, making the test more reliable and faster. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * feat: added common storage type setting * chore: added appriopriate env variables to the docker compose file and CLOUD.md upgraded * fix: fixed the issue with no logged tokens for streaming * fix: preserve originalErr in 404 error handling The NotFound case in ParseProviderError was dropping the original error, losing debugging context. Other error branches preserve this information. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * feat: forward X-Request-ID to upstream LLM providers - Add context-based request ID passing via core.WithRequestID/GetRequestID - Fix middleware to set X-Request-ID on both request and response headers - Update handlers to enrich context with request ID before calling providers - Forward request ID to all providers with provider-specific header names: - OpenAI, Groq, xAI: X-Request-ID - Anthropic: X-Request-Id - Gemini: X-Goog-Request-Id This enables end-to-end request tracing from client through GOModel to upstream providers. Clients can override the auto-generated ID by passing their own X-Request-ID header. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * chore: reduced the claude permissions * docs: claude upgraded * docs: condense CLAUDE.md and fix outdated information - Reduce from 272 to 84 lines (~69% reduction) - Fix incorrect claim about init() registration (now explicit factory.Register()) - Fix incorrect 5-minute refresh interval (actually 1 hour default) - Add clarification that Docker Compose is for manual testing only - Update config section to reference .env.template and config/config.yaml Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * fix: error related fix * fix: logger updated to cover some rare scenarios * fix: fixed google and openai Request ID headers * fix: clone RawData maps in usage extractors to prevent races Entries are queued asynchronously; if the source map is mutated after Write, the logged data can be corrupted. Added defensive copy via cloneRawData helper for all three extractor functions. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add a new configuration flag that automatically adds stream_options: {"include_usage": true} to streaming requests for OpenAI-compatible providers. This ensures usage data is always returned in streaming responses when the flag is enabled (default: true).
Changes:
Summary by CodeRabbit
ENFORCE_RETURNING_USAGE_DATAenvironment variable (defaults to enabled) to configure whether usage data is returned when streaming.✏️ Tip: You can customize this high-level summary in your review settings.