feat(usage): track prompt tokens and cost saved by request rewriters#481
Merged
Conversation
Request rewriters (ext extensions) can now report how many prompt tokens their body rewrite removed via ext.Result.TokensSaved. Core aggregates this into usage tracking and the dashboard: - The rewrite middleware sums applied rewriters' estimates and carries the total on the request context. - Both usage write paths fold it into the request's usage entry: the orchestrator's non-streaming logUsage and the SSE stream observer (wired at both construction sites). - usage.ApplyRewriteSavings prices the removed tokens as the input-cost delta between the request as sent and as forwarded, via the existing granular cost engine — so tiered, batch, and endpoint-adjusted rates stay honest (a tier crossing re-rates the whole input). - New usage columns rewrite_tokens_saved / rewrite_cost_saved in the SQLite and PostgreSQL stores (idempotent migrations; MongoDB persists via bson tags), summed into GetSummary on all three backends. - Pricing recalculation refreshes rewrite_cost_saved from the stored token estimate alongside the other cost fields. - Usage page gains "Rewrite Saved" (cost) and "Tokens Saved" cards next to Cache Saved, shown only once a rewriter reports savings — zero visual change for gateways without rewriters. Savings ride on provider usage rows only; local response-cache hits never carry them (no provider call, nothing saved). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughThis PR adds rewrite token and cost savings tracking from request rewriting through usage storage, aggregation, recalculation, dashboard display, and API schemas. ChangesRewrite Savings Tracking
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Rewriter
participant RequestRewriteMiddleware
participant RequestContext
participant StreamUsageObserver
participant UsageStore
Rewriter->>RequestRewriteMiddleware: Result{TokensSaved}
RequestRewriteMiddleware->>RequestContext: WithRewriteTokensSaved(tokensSaved)
RequestContext->>StreamUsageObserver: RewriteTokensSavedFromContext
StreamUsageObserver->>StreamUsageObserver: ApplyRewriteSavings(entry, tokensSaved, pricing)
StreamUsageObserver->>UsageStore: UsageEntry{RewriteTokensSaved, RewriteCostSaved}
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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 |
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
internal/usage/recalculate_pricing_postgresql.go (1)
68-94: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd PostgreSQL coverage for rewrite savings recalculation. The
rewrite_tokens_savedpath is only exercised by the SQLite recalculation test; add an equivalent PostgreSQL test that seedsrewrite_tokens_saved, runsRecalculatePricing, and assertsrewrite_cost_savedis refreshed.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/usage/recalculate_pricing_postgresql.go` around lines 68 - 94, Add PostgreSQL test coverage for the rewrite savings recalculation path in RecalculatePricing. Create an equivalent PostgreSQL test that seeds usage rows with rewrite_tokens_saved, runs the recalculation flow, and verifies rewrite_cost_saved is updated/refreshed. Use the recalculation logic around recalculationEntry and the postgres usage query/scan path in internal/usage/recalculate_pricing_postgresql.go as the target behavior to exercise.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cmd/gomodel/docs/docs.go`:
- Around line 7833-7839: Add the nullable marker to the source field for
RewriteCostSaved in internal/usage/usage.go, since it is a pointer-backed float
like the other cost fields but currently generates a non-nullable schema. Update
the struct tag on RewriteCostSaved to include extensions:"x-nullable", then
regenerate the docs so the schema for rewrite_cost_saved reflects nullability
consistently with the related usage fields.
In `@internal/server/request_rewrite_test.go`:
- Around line 439-496: Add a regression test in RequestRewriteMiddleware for a
rewriter that returns TokensSaved > 0 but leaves Result.Body nil, since the
middleware should not accumulate savings unless a body rewrite is actually
applied. Use the existing test patterns in
TestRequestRewriteMiddlewareStoresTokensSavedInContext and
TestRequestRewriteMiddlewareNoSavingsLeavesContextZero, and assert
RewriteTokensSavedFromContext stays 0 when a stubRewriter reports savings
without changing the body.
---
Outside diff comments:
In `@internal/usage/recalculate_pricing_postgresql.go`:
- Around line 68-94: Add PostgreSQL test coverage for the rewrite savings
recalculation path in RecalculatePricing. Create an equivalent PostgreSQL test
that seeds usage rows with rewrite_tokens_saved, runs the recalculation flow,
and verifies rewrite_cost_saved is updated/refreshed. Use the recalculation
logic around recalculationEntry and the postgres usage query/scan path in
internal/usage/recalculate_pricing_postgresql.go as the target behavior to
exercise.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 8e35f419-ab9e-4135-9a4e-863db7fac1f3
📒 Files selected for processing (28)
cmd/gomodel/docs/docs.godocs/openapi.jsonext/ext.gointernal/admin/dashboard/static/js/modules/usage.jsinternal/admin/dashboard/static/js/modules/usage.test.cjsinternal/admin/dashboard/templates/page-usage.htmlinternal/core/context.gointernal/gateway/usage.gointernal/server/passthrough_support.gointernal/server/request_rewrite.gointernal/server/request_rewrite_test.gointernal/server/translated_inference_service.gointernal/usage/reader.gointernal/usage/reader_mongodb.gointernal/usage/reader_postgresql.gointernal/usage/reader_sqlite.gointernal/usage/recalculate_pricing.gointernal/usage/recalculate_pricing_mongodb.gointernal/usage/recalculate_pricing_postgresql.gointernal/usage/recalculate_pricing_sqlite.gointernal/usage/savings.gointernal/usage/savings_test.gointernal/usage/store_postgresql.gointernal/usage/store_postgresql_test.gointernal/usage/store_sqlite.gointernal/usage/stream_observer.gointernal/usage/stream_observer_test.gointernal/usage/usage.go
Comment on lines
+7833
to
+7839
| "rewrite_cost_saved": { | ||
| "type": "number" | ||
| }, | ||
| "rewrite_tokens_saved": { | ||
| "description": "Rewrite savings: prompt tokens request rewriters removed before the\nprovider call, and the estimated input cost avoided (nil when no\nmatched row had a priced savings estimate).", | ||
| "type": "integer" | ||
| }, |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Confirm whether RewriteCostSaved lacks the x-nullable extension tag, and compare with a sibling pointer field that already has it.
rg -n -B2 -A2 'RewriteCostSaved' internal/usage/usage.go
rg -n -B2 -A2 'InputCost \*float64' internal/usage/usage.goRepository: ENTERPILOT/GoModel
Length of output: 558
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## usage.go\n'
rg -n -B3 -A3 'RewriteCostSaved|InputCost \*float64' internal/usage/usage.go || true
printf '\n## docs.go rewrite_cost_saved context\n'
rg -n -B4 -A4 '"rewrite_cost_saved"|x-nullable' cmd/gomodel/docs/docs.go | sed -n '1,120p'
printf '\n## other pointer-backed cost fields in docs.go\n'
rg -n -B2 -A4 '"input_cost"|x-nullable' cmd/gomodel/docs/docs.go | sed -n '1,160p'Repository: ENTERPILOT/GoModel
Length of output: 7850
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the exact schema block around rewrite_cost_saved.
sed -n '7826,7845p' cmd/gomodel/docs/docs.go
# Show the struct definition with surrounding comments in internal/usage/usage.go.
sed -n '80,98p' internal/usage/usage.goRepository: ENTERPILOT/GoModel
Length of output: 1975
Mark RewriteCostSaved as nullable in the source struct.
internal/usage/usage.go defines it as *float64 but omits extensions:"x-nullable", so the generated schema shows "type": "number" without nullability while sibling pointer-backed cost fields include it. Add the tag there and regenerate docs.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cmd/gomodel/docs/docs.go` around lines 7833 - 7839, Add the nullable marker
to the source field for RewriteCostSaved in internal/usage/usage.go, since it is
a pointer-backed float like the other cost fields but currently generates a
non-nullable schema. Update the struct tag on RewriteCostSaved to include
extensions:"x-nullable", then regenerate the docs so the schema for
rewrite_cost_saved reflects nullability consistently with the related usage
fields.
…ecalculation Review follow-ups on #481: - Regression tests that a rewriter claiming TokensSaved without an applied body rewrite contributes nothing — alone (context stays zero) and mid-chain (excluded from the accumulated sum). - PostgreSQL integration test exercising the savings recalculation path end to end: a usage row with rewrite_tokens_saved and a stale rewrite_cost_saved is repriced to the exact input-cost delta by POST /admin/usage/recalculate-pricing. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Jul 7, 2026
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.
What
Request rewriters registered through the
extAPI can now report how many prompt tokens their rewrite removed (ext.Result.TokensSaved), and the gateway turns that into dashboard-visible savings: total tokens saved and the estimated money saved, aggregated in usage tracking.New cards on the Usage page (shown only once a rewriter reports savings — zero visual change otherwise):
$), same period/filters as Estimated CostHow
Result.TokensSaved int— a rewriter's estimate of prompt tokens its applied body change removed.RequestRewriteMiddlewaresums the estimates of rewriters whose body was applied and stores the total on the request context (core.WithRewriteTokensSaved).logUsageandStreamUsageObserver(SetRewriteTokensSaved, wired at both construction sites). Local response-cache hits never carry savings (no provider call was made).usage.ApplyRewriteSavingsprices removed tokens as the input-cost delta between the request as the client sent it (input + saved tokens) and as forwarded, via the existingCalculateGranularCost— so endpoint-adjusted (batch) rates and tiered pricing stay honest: when the un-rewritten prompt would have landed in a higher tier, the delta includes re-rating the whole input. Unknown pricing → cost staysnil, tokens still counted.rewrite_tokens_saved/rewrite_cost_savedcolumns in SQLite and PostgreSQL (idempotent migrations, fresh-DB DDL); MongoDB persists via bson tags.GetSummarysums them on all three backends with the standard filters./admin/usage/recalculate-pricing) refreshesrewrite_cost_savedfrom the stored token estimate alongside the other cost fields on all three backends.Tests
ApplyRewriteSavingsunit table: flat rate, nil pricing, tier crossing, batch endpoint.null, not0).rewrite_cost_saved.go test ./...,make test-dashboard,make test-e2e,make lint,make build— all green. Swagger/openapi regenerated.Notes for reviewers
🤖 Generated with Claude Code
Summary by CodeRabbit