Skip to content

refactor: fix OCP violation in usage cost mapping#96

Merged
SantiagoDePolonia merged 3 commits intomainfrom
refactor/open-closed-principle
Feb 25, 2026
Merged

refactor: fix OCP violation in usage cost mapping#96
SantiagoDePolonia merged 3 commits intomainfrom
refactor/open-closed-principle

Conversation

@SantiagoDePolonia
Copy link
Copy Markdown
Contributor

@SantiagoDePolonia SantiagoDePolonia commented Feb 25, 2026

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

    • Provider-specific token cost mappings and informational token fields added for finer-grained cost attribution.
    • Usage payloads and stream responses now include additional cache and token breakdowns.
  • Refactor

    • Centralized provider registration and a startup-initialized cost registry to publish per-provider mappings and informational fields.
    • Startup/shutdown sequencing adjusted to register and clean up cost/usage components.
  • Tests

    • New tests and test startup wiring validating the cost registry and provider mappings.

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>
@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai bot commented Feb 25, 2026

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a0cb34f and 8f305eb.

📒 Files selected for processing (2)
  • internal/providers/anthropic/anthropic.go
  • internal/usage/setup_test.go

📝 Walkthrough

Walkthrough

Adds a provider-driven token cost-mapping system: new core types (CostSide, CostUnit, TokenCostMapping), provider Registration gains CostMappings and InformationalFields, ProviderFactory exposes CostRegistry(), and the usage subsystem is switched to a startup-registered registry consumed by CalculateGranularCost.

Changes

Cohort / File(s) Summary
Core Cost-Mapping Types
internal/core/cost_mapping.go
New exported types/constants (CostSide, CostUnit) and TokenCostMapping struct to represent provider token cost metadata and pricing field accessors.
Provider Factory & Tests
internal/providers/factory.go, internal/providers/factory_test.go
Replaces builders with registrations, stores full Registration values, adds CostRegistry() to aggregate per-provider TokenCostMappings and deduplicated informational fields; adds tests for CostRegistry.
Provider Registrations
internal/providers/.../openai.go, internal/providers/.../anthropic.go, internal/providers/.../gemini.go, internal/providers/.../xai.go
Registration literals extended with CostMappings []core.TokenCostMapping (OpenAI also adds InformationalFields); Anthropic also surfaces cache token fields into usage payloads.
App Initialization
internal/app/app.go
App startup calls factory.CostRegistry() and registers mappings via usage.RegisterCostMappings(); shutdown ordering notes adjusted.
Usage Registry & Calculation
internal/usage/cost.go, internal/usage/stream_wrapper.go, internal/usage/setup_test.go
Introduces costRegistry with atomic publish, RegisterCostMappings() and loadCostRegistry(); CalculateGranularCost() and stream extraction read from the registry instead of hard-coded mappings; test setup registers mappings in TestMain.
Provider Scaffolding Template
CLAUDE.md
Scaffold template updated to include optional CostMappings and InformationalFields in provider Registration.

Sequence Diagram

sequenceDiagram
    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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

Suggested labels

enhancement

Poem

🐇 I hop through mappings, counting every token,
From cached crumbs to reasoning spoken.
Registrations whisper prices and side,
A registry blooms where totals abide.
Hooray — costs tallied, neat and woken!

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'refactor: fix OCP violation in usage cost mapping' directly and clearly describes the main objective: refactoring to fix an Open/Closed Principle (OCP) violation related to cost mapping.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
  • 📝 Generate docstrings (stacked PR)
  • 📝 Generate docstrings (commit on current branch)
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch refactor/open-closed-principle

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 | 🟡 Minor

Startup sequence missing the RegisterCostMappings step.

The sequence on line 65 jumps from "Init providers" directly to "Init audit logging", but app.go now 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

Shutdown doc 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 | 🟠 Major

Streaming converter drops cache token counts, so CostMappings for cache_read_input_tokens / cache_creation_input_tokens have no effect in streaming mode.

event.Usage (*anthropicUsage) carries CacheReadInputTokens and CacheCreationInputTokens, but convertEvent only emits the three standard fields in the "usage" object. When stream_wrapper.go iterates over defaultCostRegistry.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 (buildAnthropicRawUsage populates RawData properly). 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 responsesStreamConverter in the response.completed event at the EOF handling block (lines ~791-796), using sc.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

📥 Commits

Reviewing files that changed from the base of the PR and between aaa5e24 and f9e62fd.

📒 Files selected for processing (12)
  • CLAUDE.md
  • internal/app/app.go
  • internal/core/cost_mapping.go
  • internal/providers/anthropic/anthropic.go
  • internal/providers/factory.go
  • internal/providers/factory_test.go
  • internal/providers/gemini/gemini.go
  • internal/providers/openai/openai.go
  • internal/providers/xai/xai.go
  • internal/usage/cost.go
  • internal/usage/setup_test.go
  • internal/usage/stream_wrapper.go

Comment on lines +6 to +17
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
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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.

Suggested change
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>
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between f9e62fd and a0cb34f.

📒 Files selected for processing (8)
  • CLAUDE.md
  • internal/app/app.go
  • internal/core/cost_mapping.go
  • internal/providers/anthropic/anthropic.go
  • internal/providers/factory.go
  • internal/usage/cost.go
  • internal/usage/setup_test.go
  • internal/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>
@SantiagoDePolonia SantiagoDePolonia merged commit b65c319 into main Feb 25, 2026
12 checks passed
SantiagoDePolonia added a commit that referenced this pull request Feb 25, 2026
@SantiagoDePolonia SantiagoDePolonia deleted the refactor/open-closed-principle branch March 22, 2026 14:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant