Refactor/open closed principle#98
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>
- 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>
- 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>
📝 WalkthroughWalkthroughThis PR introduces a dynamic provider cost mapping framework by adding a new Changes
Sequence Diagram(s)sequenceDiagram
participant App as App Init
participant Factory as ProviderFactory
participant Usage as Usage Module
participant Calc as Cost Calculator
App->>Factory: CostRegistry()
Factory->>Factory: Aggregate CostMappings<br/>from all registrations
Factory-->>App: mappings, informationalFields
App->>Usage: RegisterCostMappings(mappings,<br/>informationalFields)
Usage->>Usage: Store in atomic<br/>costRegistry
rect rgba(100, 150, 255, 0.5)
Note over Calc: Later at runtime
Calc->>Usage: loadCostRegistry()
Usage-->>Calc: costRegistry{providerMappings,<br/>informationalFields}
Calc->>Calc: Look up RawDataKey mappings<br/>and pricing fields
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 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 |
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
internal/usage/stream_wrapper.go (1)
236-240: 🧹 Nitpick | 🔵 Trivial
loadCostRegistry()called on every SSE data line — consider caching the registry reference at wrapper creation.
loadCostRegistry()is invoked insideextractUsageFromJSON, which is called for everydata:line in every SSE event. Although anatomic.Pointerload is cheap, the map lookup loop overextendedFieldSetruns repeatedly with the same registry for the lifetime of the wrapper. Caching theextendedFieldSetmap reference (or the whole*costRegistry) inStreamUsageWrapperat construction time would be cleaner and avoids repeated atomic loads on the hot path.♻️ Proposed refactor
type StreamUsageWrapper struct { io.ReadCloser logger LoggerInterface pricingResolver PricingResolver eventBuffer bytes.Buffer cachedEntry *UsageEntry model string provider string requestID string endpoint string + extendedFields map[string]struct{} closed bool } func NewStreamUsageWrapper(...) *StreamUsageWrapper { return &StreamUsageWrapper{ ... + extendedFields: loadCostRegistry().extendedFieldSet, } }Then replace Line 236 with:
- for field := range loadCostRegistry().extendedFieldSet { + for field := range w.extendedFields {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/usage/stream_wrapper.go` around lines 236 - 240, Cache the registry reference when the wrapper is created instead of calling loadCostRegistry() on every SSE line: add a field to StreamUsageWrapper (e.g., cachedRegistry or cachedExtendedFieldSet of type map[string]struct{} or *costRegistry) and initialize it in the constructor, update extractUsageFromJSON to use that cachedExtendedFieldSet (or cachedRegistry.extendedFieldSet) instead of calling loadCostRegistry(), and remove the loadCostRegistry() call inside the loop that iterates over extendedFieldSet so the map lookup uses the pre-stored reference.
🤖 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 7-18: The CostSideUnknown and CostUnitUnknown zero-value sentinels
can slip into TokenCostMapping when providers omit Side or Unit, so add
validation to reject mappings with those zero-values: in the factory.Add (and/or
RegisterCostMappings) code path where TokenCostMapping entries are registered
(the same place that already panics on empty Type or nil New), check each
mapping's Side against CostSideUnknown and Unit against CostUnitUnknown and
panic or return an error if found; ensure the validation references
TokenCostMapping.Side and TokenCostMapping.Unit to make failures explicit and
prevent silent incorrect mappings.
In `@internal/usage/cost.go`:
- Around line 112-113: The code calls m.PricingField(pricing) without ensuring
the PricingField callback exists, which can panic if a TokenCostMapping is
zero-valued; update the logic in the block that currently invokes PricingField
(using variables m and pricing) to first check if m.PricingField != nil, and if
nil emit a caveat/log (e.g., note missing PricingField for that
TokenCostMapping) and skip calling it, otherwise call m.PricingField(pricing) as
before.
- Around line 47-65: RegisterCostMappings currently stores the caller-owned
mappings map directly into the costRegistry, allowing callers to mutate
slices/entries after registration and causing concurrent readers (e.g.,
CalculateGranularCost) to observe mutations; fix by defensively cloning the
entire mappings structure before publishing: allocate a new
map[string][]core.TokenCostMapping, iterate each key in mappings, create a new
slice for each value and copy its TokenCostMapping elements (so slices and
elements are not shared), then assign that new map to
costRegistry.providerMappings (and keep existing cloning of informational into
informationalFields) before calling costRegistryPtr.Store(reg); reference
RegisterCostMappings, costRegistry.providerMappings, core.TokenCostMapping,
RawDataKey, costRegistryPtr, and CalculateGranularCost when locating where to
implement the clone.
In `@internal/usage/setup_test.go`:
- Around line 14-27: Add a short maintenance comment above TestMain explaining
that the provider list (calls to providers.NewProviderFactory() and
factory.Add(...)) must be updated whenever a new provider is introduced so
CostRegistry() and RegisterCostMappings() remain complete; reference TestMain,
providers.NewProviderFactory, factory.Add and RegisterCostMappings in the
comment so future contributors know to add the new provider registrations to
this test setup.
---
Outside diff comments:
In `@internal/usage/stream_wrapper.go`:
- Around line 236-240: Cache the registry reference when the wrapper is created
instead of calling loadCostRegistry() on every SSE line: add a field to
StreamUsageWrapper (e.g., cachedRegistry or cachedExtendedFieldSet of type
map[string]struct{} or *costRegistry) and initialize it in the constructor,
update extractUsageFromJSON to use that cachedExtendedFieldSet (or
cachedRegistry.extendedFieldSet) instead of calling loadCostRegistry(), and
remove the loadCostRegistry() call inside the loop that iterates over
extendedFieldSet so the map lookup uses the pre-stored reference.
ℹ️ 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
| CostSideUnknown CostSide = iota // zero value; must not be used in mappings | ||
| CostSideInput | ||
| CostSideOutput | ||
| ) | ||
|
|
||
| // CostUnit indicates how the pricing field is applied. | ||
| type CostUnit int | ||
|
|
||
| const ( | ||
| CostUnitUnknown CostUnit = iota // zero value; must not be used in mappings | ||
| CostUnitPerMtok // divide token count by 1M, multiply by rate | ||
| CostUnitPerItem // multiply count directly by rate |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
CostSideUnknown/CostUnitUnknown zero-value sentinels are unvalidated — silently incorrect mappings possible.
The doc comments say these values "must not be used in mappings," but nothing enforces this. A provider that accidentally omits Side or Unit in a TokenCostMapping literal will compile and run without error, producing a mapping with a zero-value enum that the cost-calculation logic likely cannot handle correctly. Consider adding a validation guard in factory.Add() (which already panics on empty Type/nil New) or in RegisterCostMappings:
🛡️ Suggested validation in factory.Add (illustrative)
// Inside factory.Add, after existing guards:
+for i, m := range reg.CostMappings {
+ if m.RawDataKey == "" {
+ panic(fmt.Sprintf("provider %q: CostMappings[%d] has empty RawDataKey", reg.Type, i))
+ }
+ if m.Side == core.CostSideUnknown {
+ panic(fmt.Sprintf("provider %q: CostMappings[%d] has unknown Side", reg.Type, i))
+ }
+ if m.Unit == core.CostUnitUnknown {
+ panic(fmt.Sprintf("provider %q: CostMappings[%d] has unknown Unit", reg.Type, i))
+ }
+}🤖 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 7 - 18, The CostSideUnknown and
CostUnitUnknown zero-value sentinels can slip into TokenCostMapping when
providers omit Side or Unit, so add validation to reject mappings with those
zero-values: in the factory.Add (and/or RegisterCostMappings) code path where
TokenCostMapping entries are registered (the same place that already panics on
empty Type or nil New), check each mapping's Side against CostSideUnknown and
Unit against CostUnitUnknown and panic or return an error if found; ensure the
validation references TokenCostMapping.Side and TokenCostMapping.Unit to make
failures explicit and prevent silent incorrect mappings.
| func RegisterCostMappings(mappings map[string][]core.TokenCostMapping, informational []string) { | ||
| reg := &costRegistry{ | ||
| providerMappings: mappings, | ||
| informationalFields: make(map[string]struct{}, len(informational)), | ||
| extendedFieldSet: make(map[string]struct{}), | ||
| } | ||
|
|
||
| // extendedFieldSet is derived from providerMappings and contains all RawData keys | ||
| // that providers may report. Used by stream_wrapper.go to extract extended fields | ||
| // from SSE usage data without maintaining a separate hard-coded list. | ||
| var extendedFieldSet = func() map[string]struct{} { | ||
| set := make(map[string]struct{}) | ||
| for _, mappings := range providerMappings { | ||
| for _, m := range mappings { | ||
| set[m.rawDataKey] = struct{}{} | ||
| for _, f := range informational { | ||
| reg.informationalFields[f] = struct{}{} | ||
| } | ||
|
|
||
| for _, ms := range mappings { | ||
| for _, m := range ms { | ||
| reg.extendedFieldSet[m.RawDataKey] = struct{}{} | ||
| } | ||
| } | ||
| return set | ||
| }() | ||
|
|
||
| costRegistryPtr.Store(reg) | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cd /repo && find internal/usage -type f -name "*.go" | head -20Repository: ENTERPILOT/GOModel
Length of output: 117
🏁 Script executed:
find . -path "*/internal/usage/*.go" -type fRepository: ENTERPILOT/GOModel
Length of output: 758
🏁 Script executed:
cat -n ./internal/usage/cost.go | head -100Repository: ENTERPILOT/GOModel
Length of output: 3761
🏁 Script executed:
cat -n ./internal/usage/cost.go | sed -n '100,200p'Repository: ENTERPILOT/GOModel
Length of output: 3354
🏁 Script executed:
grep -rn "RegisterCostMappings" --include="*.go" .Repository: ENTERPILOT/GOModel
Length of output: 655
🏁 Script executed:
cat -n ./internal/app/app.go | sed -n '70,85p'Repository: ENTERPILOT/GOModel
Length of output: 751
🏁 Script executed:
grep -rn "CostRegistry" --include="*.go" . | head -20Repository: ENTERPILOT/GOModel
Length of output: 1140
🏁 Script executed:
cat -n ./internal/providers/factory.go | sed -n '106,140p'Repository: ENTERPILOT/GOModel
Length of output: 1005
🏁 Script executed:
cat -n ./internal/usage/cost_test.go | head -80Repository: ENTERPILOT/GOModel
Length of output: 3135
🏁 Script executed:
cat -n ./internal/providers/factory.go | sed -n '1,50p'Repository: ENTERPILOT/GOModel
Length of output: 1918
🏁 Script executed:
cat -n ./internal/providers/factory.go | sed -n '115,120p'Repository: ENTERPILOT/GOModel
Length of output: 304
🏁 Script executed:
cat -n ./internal/usage/setup_test.goRepository: ENTERPILOT/GOModel
Length of output: 978
🏁 Script executed:
rg -n "costRegistryPtr|loadCostRegistry" --include="*.go" -B 2 -A 2Repository: ENTERPILOT/GOModel
Length of output: 503
🏁 Script executed:
rg -n "costRegistryPtr|loadCostRegistry" -t go -B 2 -A 2Repository: ENTERPILOT/GOModel
Length of output: 1867
🏁 Script executed:
rg -n "costMappings\|RegisterCostMappings" -t go -B 3 -A 3 ./internal/app/app.goRepository: ENTERPILOT/GOModel
Length of output: 44
🏁 Script executed:
cat -n ./internal/app/app.go | sed -n '60,90p'Repository: ENTERPILOT/GOModel
Length of output: 1180
Defensively copy mappings before publishing the registry.
Line 49 stores the caller-owned map by reference. While current usage calls this once at startup, if the input map or slices are modified after registration, concurrent readers in CalculateGranularCost could observe mutations and panic. Clone map and slices before Store().
Suggested fix
func RegisterCostMappings(mappings map[string][]core.TokenCostMapping, informational []string) {
+ copiedMappings := make(map[string][]core.TokenCostMapping, len(mappings))
+ for providerType, providerMappings := range mappings {
+ mappingsCopy := make([]core.TokenCostMapping, len(providerMappings))
+ copy(mappingsCopy, providerMappings)
+ copiedMappings[providerType] = mappingsCopy
+ }
+
reg := &costRegistry{
- providerMappings: mappings,
+ providerMappings: copiedMappings,
informationalFields: make(map[string]struct{}, len(informational)),
extendedFieldSet: make(map[string]struct{}),
}
@@
- for _, ms := range mappings {
+ for _, ms := range copiedMappings {
for _, m := range ms {
reg.extendedFieldSet[m.RawDataKey] = struct{}{}
}
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@internal/usage/cost.go` around lines 47 - 65, RegisterCostMappings currently
stores the caller-owned mappings map directly into the costRegistry, allowing
callers to mutate slices/entries after registration and causing concurrent
readers (e.g., CalculateGranularCost) to observe mutations; fix by defensively
cloning the entire mappings structure before publishing: allocate a new
map[string][]core.TokenCostMapping, iterate each key in mappings, create a new
slice for each value and copy its TokenCostMapping elements (so slices and
elements are not shared), then assign that new map to
costRegistry.providerMappings (and keep existing cloning of informational into
informationalFields) before calling costRegistryPtr.Store(reg); reference
RegisterCostMappings, costRegistry.providerMappings, core.TokenCostMapping,
RawDataKey, costRegistryPtr, and CalculateGranularCost when locating where to
implement the clone.
| rate := m.PricingField(pricing) | ||
| if rate == nil { |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's find and read the cost.go file
find . -path "*internal/usage/cost.go" -type fRepository: ENTERPILOT/GOModel
Length of output: 86
🏁 Script executed:
# Read the file to see the context around lines 112-113
cat -n internal/usage/cost.go | head -130 | tail -50Repository: ENTERPILOT/GOModel
Length of output: 1938
🏁 Script executed:
# Let's also check the type definition of TokenCostMapping
fd "\.go$" -x grep -l "type TokenCostMapping" {}Repository: ENTERPILOT/GOModel
Length of output: 93
🏁 Script executed:
# Search for PricingField usage patterns in the codebase
rg "PricingField" -A 3 -B 3 internal/usage/Repository: ENTERPILOT/GOModel
Length of output: 811
🏁 Script executed:
# Let's examine the type definition of TokenCostMapping
cat -n internal/core/cost_mapping.goRepository: ENTERPILOT/GOModel
Length of output: 1389
🏁 Script executed:
# Let's check the test that mentions NilPricingFieldNoCaveat
rg "NilPricingFieldNoCaveat" -A 30 internal/usage/cost_test.goRepository: ENTERPILOT/GOModel
Length of output: 1169
🏁 Script executed:
# Let's check what initializes the providerMappings and how TokenCostMapping is constructed
rg "providerMappings\[" -B 5 -A 5 internal/Repository: ENTERPILOT/GOModel
Length of output: 778
🏁 Script executed:
# Search for all TokenCostMapping struct literals to see how they're initialized
rg "TokenCostMapping\{" -A 10 internal/Repository: ENTERPILOT/GOModel
Length of output: 6606
🏁 Script executed:
# Check if there are any other usages of m.PricingField that show nil checking patterns
rg "\.PricingField" -B 3 -A 3 internal/Repository: ENTERPILOT/GOModel
Length of output: 394
Guard nil PricingField callbacks before invocation.
Line 112 calls m.PricingField(pricing) unconditionally. If a TokenCostMapping is zero-valued or malformed, the function pointer will be nil and cause a panic. Add a nil check and emit a caveat instead.
Suggested fix
- rate := m.PricingField(pricing)
+ if m.PricingField == nil {
+ caveats = append(caveats, fmt.Sprintf("missing pricing field mapper for field %s", m.RawDataKey))
+ continue
+ }
+ rate := m.PricingField(pricing)
if rate == nil {
continue // Base rate covers this token type; no adjustment needed
}📝 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.
| rate := m.PricingField(pricing) | |
| if rate == nil { | |
| if m.PricingField == nil { | |
| caveats = append(caveats, fmt.Sprintf("missing pricing field mapper for field %s", m.RawDataKey)) | |
| continue | |
| } | |
| rate := m.PricingField(pricing) | |
| if rate == nil { | |
| continue // Base rate covers this token type; no adjustment needed | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@internal/usage/cost.go` around lines 112 - 113, The code calls
m.PricingField(pricing) without ensuring the PricingField callback exists, which
can panic if a TokenCostMapping is zero-valued; update the logic in the block
that currently invokes PricingField (using variables m and pricing) to first
check if m.PricingField != nil, and if nil emit a caveat/log (e.g., note missing
PricingField for that TokenCostMapping) and skip calling it, otherwise call
m.PricingField(pricing) as before.
| func TestMain(m *testing.M) { | ||
| // Build cost mappings and informational fields the same way production does: | ||
| // register all providers into a factory and use CostRegistry to aggregate. | ||
| factory := providers.NewProviderFactory() | ||
| factory.Add(openai.Registration) | ||
| factory.Add(anthropic.Registration) | ||
| factory.Add(gemini.Registration) | ||
| factory.Add(xai.Registration) | ||
|
|
||
| costMappings, informationalFields := factory.CostRegistry() | ||
| RegisterCostMappings(costMappings, informationalFields) | ||
|
|
||
| os.Exit(m.Run()) | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
LGTM — mirrors the production startup sequence faithfully.
The TestMain calls factory.CostRegistry() + RegisterCostMappings before m.Run(), exactly mirroring app.go's startup path. The os.Exit(m.Run()) idiom is correct.
One maintenance note: if a new provider is added to the codebase, this TestMain must be updated to include it, otherwise usage package tests will silently run with an incomplete registry (missing the new provider's extendedFieldSet keys). Consider adding a comment to this effect.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@internal/usage/setup_test.go` around lines 14 - 27, Add a short maintenance
comment above TestMain explaining that the provider list (calls to
providers.NewProviderFactory() and factory.Add(...)) must be updated whenever a
new provider is introduced so CostRegistry() and RegisterCostMappings() remain
complete; reference TestMain, providers.NewProviderFactory, factory.Add and
RegisterCostMappings in the comment so future contributors know to add the new
provider registrations to this test setup.
1 participant
Summary by CodeRabbit
New Features
Tests
Refactor