refactor: fix OCP violation in usage cost mapping#96
Conversation
Move provider-specific cost mapping data out of usage/cost.go and into each provider's Registration var, wired through the factory at startup. This ensures adding a new provider no longer requires editing the usage package. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review infoConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughAdds a provider-driven token cost-mapping system: new core types ( Changes
Sequence DiagramsequenceDiagram
actor App as App Initialization
participant Factory as ProviderFactory
participant Usage as usage package
participant Core as internal/core
App->>Factory: Add(registration) × N
Note over Factory: registrations store CostMappings & InformationalFields
App->>Factory: CostRegistry()
Factory->>Factory: Aggregate provider CostMappings
Factory->>Factory: Deduplicate & sort informational fields
Factory-->>App: (map[string][]TokenCostMapping, []string)
App->>Usage: RegisterCostMappings(mappings, informational)
Usage->>Core: publish registry (atomic pointer)
Note over Usage,Core: Later during processing
Usage->>Core: loadCostRegistry()
Usage->>Usage: CalculateGranularCost() uses registry mappings & pricing fields
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
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 |
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
CLAUDE.md (1)
65-65:⚠️ Potential issue | 🟡 MinorStartup sequence missing the
RegisterCostMappingsstep.The sequence on line 65 jumps from "Init providers" directly to "Init audit logging", but
app.gonow inserts a cost-mapping registration step between them. Stale documentation here could mislead contributors adding new providers.✏️ Proposed doc fix
-**Startup:** Config load (defaults → YAML → env vars) → Register providers with factory → Init providers (cache → async model load → background refresh → router) → Init audit logging → Init usage tracking (shares storage if same backend) → Build guardrails pipeline → Create server → Start listening +**Startup:** Config load (defaults → YAML → env vars) → Register providers with factory → Init providers (cache → async model load → background refresh → router) → Register cost mappings (`factory.CostRegistry` → `usage.RegisterCostMappings`) → Init audit logging → Init usage tracking (shares storage if same backend) → Build guardrails pipeline → Create server → Start listening🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@CLAUDE.md` at line 65, The startup sequence in CLAUDE.md is missing the cost-mapping registration step; update the sequence to include the RegisterCostMappings step between "Init providers (cache → async model load → background refresh → router)" and "Init audit logging" to match app.go where RegisterCostMappings is invoked; reference the RegisterCostMappings symbol (and app.go startup flow) to ensure the doc order mirrors the code.internal/app/app.go (1)
217-222:⚠️ Potential issue | 🟡 Minor
Shutdowndoc comment is missing the usage-tracking step.The comment enumerates only 3 teardown steps, but the implementation performs 4. This will mislead readers about the actual shutdown sequence.
✏️ Proposed doc fix
// Shutdown gracefully shuts down all components in the correct order. // It ensures proper cleanup of resources: // 1. HTTP server (stop accepting new requests) // 2. Background refresh goroutine and cache -// 3. Audit logging +// 3. Usage tracking (flush buffer) +// 4. Audit logging (flush buffer) // // Safe to call multiple times; subsequent calls are no-ops.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/app/app.go` around lines 217 - 222, Update the Shutdown method comment in internal/app/app.go to include the missing fourth step (usage-tracking teardown) so the docstring accurately reflects the implementation; explicitly list: (1) stop HTTP server, (2) stop background refresh goroutine and clear cache, (3) stop/flush usage-tracking, and (4) stop/flush audit logging, referencing the Shutdown function so readers know the order matches the implementation.internal/providers/anthropic/anthropic.go (1)
503-510:⚠️ Potential issue | 🟠 MajorStreaming converter drops cache token counts, so
CostMappingsforcache_read_input_tokens/cache_creation_input_tokenshave no effect in streaming mode.
event.Usage(*anthropicUsage) carriesCacheReadInputTokensandCacheCreationInputTokens, butconvertEventonly emits the three standard fields in the"usage"object. Whenstream_wrapper.goiterates overdefaultCostRegistry.extendedFieldSet, it looks for these keys in the parsed usage map and finds nothing, silently producing zero cache costs for every streaming request.The non-streaming path is correct (
buildAnthropicRawUsagepopulatesRawDataproperly). The streaming path needs the same treatment:🐛 Proposed fix — include cache fields in streaming usage emission
if event.Usage != nil { - chunk["usage"] = map[string]interface{}{ + usagePayload := map[string]interface{}{ "prompt_tokens": event.Usage.InputTokens, "completion_tokens": event.Usage.OutputTokens, "total_tokens": event.Usage.InputTokens + event.Usage.OutputTokens, } + if event.Usage.CacheReadInputTokens > 0 { + usagePayload["cache_read_input_tokens"] = event.Usage.CacheReadInputTokens + } + if event.Usage.CacheCreationInputTokens > 0 { + usagePayload["cache_creation_input_tokens"] = event.Usage.CacheCreationInputTokens + } + chunk["usage"] = usagePayload }Apply the same fix to
responsesStreamConverterin theresponse.completedevent at the EOF handling block (lines ~791-796), usingsc.cachedUsage.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/providers/anthropic/anthropic.go` around lines 503 - 510, convertEvent currently emits only prompt_tokens, completion_tokens and total_tokens from event.Usage (anthropicUsage), which omits CacheReadInputTokens and CacheCreationInputTokens and causes cache-related CostMappings to be ignored in streaming mode; update the usage map built in convertEvent (the block that checks if event.Usage != nil) to also include "cache_read_input_tokens" and "cache_creation_input_tokens" (and keep total_tokens calculation) using event.Usage.CacheReadInputTokens and event.Usage.CacheCreationInputTokens, and apply the same change in responsesStreamConverter where the EOF/response.completed handling emits sc.cachedUsage so that sc.cachedUsage includes those two cache fields as well.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/core/cost_mapping.go`:
- Around line 19-25: The TokenCostMapping struct and its exported fields
(TokenCostMapping, RawDataKey, PricingField, Side, Unit) lack Go doc comments;
add concise doc comments above the TokenCostMapping type and each exported field
describing their purpose and types (e.g. RawDataKey is the raw data key string,
PricingField returns a pointer to a float64 from *ModelPricing, Side indicates
the CostSide, Unit indicates the CostUnit) so the package complies with
exported-symbol documentation guidelines and tooling.
- Around line 6-17: Add explicit zero-value sentinels for both enums to avoid
silent defaults: introduce CostSideUnknown as the first CostSide value and shift
CostSideInput/CostSideOutput to follow, and similarly introduce CostUnitUnknown
as the first CostUnit value and shift CostUnitPerMtok/CostUnitPerItem. Update
any validation logic (e.g., wherever TokenCostMapping instances are
validated/constructed) to reject Unknown values and ensure newly created
mappings set an explicit non-unknown CostSide and CostUnit. Reference the enum
types CostSide and CostUnit and the constants CostSideInput, CostSideOutput,
CostUnitPerMtok, CostUnitPerItem when making these changes so callers are
updated to use the new non-unknown values.
In `@internal/providers/factory.go`:
- Around line 107-127: CostRegistry currently builds informationalFields from
f.registrations (map iteration order is nondeterministic); before returning from
ProviderFactory.CostRegistry, sort the informationalFields slice to make output
deterministic (use sort.Strings) and ensure the package imports "sort" if not
present; update the function that references informationalFields (and the seen
map logic stays the same) so the only change is calling
sort.Strings(informationalFields) just prior to the return.
In `@internal/usage/cost.go`:
- Around line 113-127: Add explicit default branches to both switches that
handle m.Unit and m.Side (the switches around CostUnitPerMtok/CostUnitPerItem
and CostSideInput/CostSideOutput) so unknown enum values don't silently skip
cost accounting; in each default case emit a clear warning or error (using the
module's logger or return an error/panic path consistent with surrounding error
handling) that includes the unexpected enum value and context (e.g., the
meter/metric id and count) to surface misconfiguration early and ensure costs
are not silently dropped.
- Around line 29-55: RegisterCostMappings currently does a plain pointer write
to defaultCostRegistry while CalculateGranularCost reads it unsynchronized;
change publication to use atomic.Pointer[costRegistry] (or a sync.Once) so
readers see a safe, ordered store. Replace the package-level defaultCostRegistry
with an atomic.Pointer[costRegistry], update RegisterCostMappings to allocate
the new costRegistry and publish it with ptr.Store(newReg), and update
CalculateGranularCost (and any other readers) to use ptr.Load() and nil-check
before use; add the necessary import for sync/atomic and keep the costRegistry
type name unchanged.
In `@internal/usage/setup_test.go`:
- Around line 10-44: TestMain currently hardcodes provider cost mappings via
RegisterCostMappings which duplicates openai/anthropic/gemini/xai
Registration.CostMappings and can drift; replace the hardcoded maps by importing
each provider package and passing their authoritative Registration.CostMappings
(e.g., openai.Registration.CostMappings, anthropic.Registration.CostMappings,
gemini.Registration.CostMappings, xai.Registration.CostMappings) into
RegisterCostMappings inside TestMain so tests use the single source of truth;
keep the final shared slice (the second argument) as-is or build it from
provider vars if available.
---
Outside diff comments:
In `@CLAUDE.md`:
- Line 65: The startup sequence in CLAUDE.md is missing the cost-mapping
registration step; update the sequence to include the RegisterCostMappings step
between "Init providers (cache → async model load → background refresh →
router)" and "Init audit logging" to match app.go where RegisterCostMappings is
invoked; reference the RegisterCostMappings symbol (and app.go startup flow) to
ensure the doc order mirrors the code.
In `@internal/app/app.go`:
- Around line 217-222: Update the Shutdown method comment in internal/app/app.go
to include the missing fourth step (usage-tracking teardown) so the docstring
accurately reflects the implementation; explicitly list: (1) stop HTTP server,
(2) stop background refresh goroutine and clear cache, (3) stop/flush
usage-tracking, and (4) stop/flush audit logging, referencing the Shutdown
function so readers know the order matches the implementation.
In `@internal/providers/anthropic/anthropic.go`:
- Around line 503-510: convertEvent currently emits only prompt_tokens,
completion_tokens and total_tokens from event.Usage (anthropicUsage), which
omits CacheReadInputTokens and CacheCreationInputTokens and causes cache-related
CostMappings to be ignored in streaming mode; update the usage map built in
convertEvent (the block that checks if event.Usage != nil) to also include
"cache_read_input_tokens" and "cache_creation_input_tokens" (and keep
total_tokens calculation) using event.Usage.CacheReadInputTokens and
event.Usage.CacheCreationInputTokens, and apply the same change in
responsesStreamConverter where the EOF/response.completed handling emits
sc.cachedUsage so that sc.cachedUsage includes those two cache fields as well.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (12)
CLAUDE.mdinternal/app/app.gointernal/core/cost_mapping.gointernal/providers/anthropic/anthropic.gointernal/providers/factory.gointernal/providers/factory_test.gointernal/providers/gemini/gemini.gointernal/providers/openai/openai.gointernal/providers/xai/xai.gointernal/usage/cost.gointernal/usage/setup_test.gointernal/usage/stream_wrapper.go
| const ( | ||
| CostSideInput CostSide = iota | ||
| CostSideOutput | ||
| ) | ||
|
|
||
| // CostUnit indicates how the pricing field is applied. | ||
| type CostUnit int | ||
|
|
||
| const ( | ||
| CostUnitPerMtok CostUnit = iota // divide token count by 1M, multiply by rate | ||
| CostUnitPerItem // multiply count directly by rate | ||
| ) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Zero-value ambiguity on CostSide and CostUnit iota enums.
Both CostSideInput = iota (0) and CostUnitPerMtok = iota (0) are the zero values of their types. An accidentally zero-initialized or partially filled TokenCostMapping will silently adopt input/per-mtok semantics with no compile-time signal. Consider shifting the iota by introducing an explicit CostSideUnknown/CostUnitUnknown sentinel at 0 to make unintentional usage detectable at validation time.
💡 Suggested sentinel approach
const (
- CostSideInput CostSide = iota
+ CostSideUnknown CostSide = iota
+ CostSideInput
CostSideOutput
)
const (
- CostUnitPerMtok CostUnit = iota // divide token count by 1M, multiply by rate
+ CostUnitUnknown CostUnit = iota
+ CostUnitPerMtok // divide token count by 1M, multiply by rate
CostUnitPerItem // multiply count directly by rate
)📝 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.
| const ( | |
| CostSideInput CostSide = iota | |
| CostSideOutput | |
| ) | |
| // CostUnit indicates how the pricing field is applied. | |
| type CostUnit int | |
| const ( | |
| CostUnitPerMtok CostUnit = iota // divide token count by 1M, multiply by rate | |
| CostUnitPerItem // multiply count directly by rate | |
| ) | |
| const ( | |
| CostSideUnknown CostSide = iota | |
| CostSideInput | |
| CostSideOutput | |
| ) | |
| // CostUnit indicates how the pricing field is applied. | |
| type CostUnit int | |
| const ( | |
| CostUnitUnknown CostUnit = iota | |
| CostUnitPerMtok // divide token count by 1M, multiply by rate | |
| CostUnitPerItem // multiply count directly by rate | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@internal/core/cost_mapping.go` around lines 6 - 17, Add explicit zero-value
sentinels for both enums to avoid silent defaults: introduce CostSideUnknown as
the first CostSide value and shift CostSideInput/CostSideOutput to follow, and
similarly introduce CostUnitUnknown as the first CostUnit value and shift
CostUnitPerMtok/CostUnitPerItem. Update any validation logic (e.g., wherever
TokenCostMapping instances are validated/constructed) to reject Unknown values
and ensure newly created mappings set an explicit non-unknown CostSide and
CostUnit. Reference the enum types CostSide and CostUnit and the constants
CostSideInput, CostSideOutput, CostUnitPerMtok, CostUnitPerItem when making
these changes so callers are updated to use the new non-unknown values.
- Add CostSideUnknown/CostUnitUnknown zero-value sentinels and field doc comments to TokenCostMapping (core/cost_mapping.go) - Sort informationalFields in CostRegistry for deterministic output - Add default branches to Unit/Side switches to surface unknown enums - Use atomic.Pointer for defaultCostRegistry to prevent data races - Replace hardcoded test data with provider Registration vars - Add cache_read_input_tokens and cache_creation_input_tokens to Anthropic streaming usage (both ChatCompletion and Responses) - Fix Shutdown docstring and startup sequence in CLAUDE.md Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/providers/anthropic/anthropic.go`:
- Around line 798-809: Extract the duplicated conditional cache-field injection
into a small helper (e.g., appendCacheFields(usageMap map[string]interface{}, u
*anthropicUsage) or reuse buildAnthropicRawUsage) and call it from both
streaming code paths where usageMap is constructed (the blocks that reference
usageMap, sc.cachedUsage.CacheReadInputTokens,
sc.cachedUsage.CacheCreationInputTokens and set responseData["usage"]); the
helper should add "cache_read_input_tokens" and "cache_creation_input_tokens"
only when the corresponding sc.cachedUsage fields are > 0 so you can replace the
duplicated if statements in both ChatCompletion and Responses streaming
locations with a single call to the helper.
In `@internal/usage/setup_test.go`:
- Around line 14-25: The test hard-codes
openai.Registration.InformationalFields, which diverges from production
(Factory.CostRegistry) that aggregates/deduplicates fields from all providers;
update TestMain to build the InformationalFields the same way production does by
importing the providers package and retrieving/merging the providers'
InformationalFields (e.g., via providers.Factory().CostRegistry() or iterating
providers.Registry) and pass that aggregated/deduped slice into
RegisterCostMappings instead of openai.Registration.InformationalFields; also
add the import "gomodel/internal/providers".
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (8)
CLAUDE.mdinternal/app/app.gointernal/core/cost_mapping.gointernal/providers/anthropic/anthropic.gointernal/providers/factory.gointernal/usage/cost.gointernal/usage/setup_test.gointernal/usage/stream_wrapper.go
- Extract duplicated cache field injection in Anthropic streaming into appendCacheFields helper, called from both ChatCompletion and Responses stream converters - Use ProviderFactory.CostRegistry() in usage TestMain to aggregate informational fields the same way production does, instead of hardcoding openai.Registration.InformationalFields Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This reverts commit b65c319.
Move provider-specific cost mapping data out of usage/cost.go and into each provider's Registration var, wired through the factory at startup. This ensures adding a new provider no longer requires editing the usage package.
Summary by CodeRabbit
New Features
Refactor
Tests