feat(claude): support /v1/messages/count_tokens with local estimation - #4441
Open
daymade wants to merge 3 commits into
Open
feat(claude): support /v1/messages/count_tokens with local estimation#4441daymade wants to merge 3 commits into
daymade wants to merge 3 commits into
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds a Claude-compatible POST /v1/messages/count_tokens endpoint with a Gin handler and relay route, a local estimator that normalizes request tools and computes input token estimates, and unit tests for estimator and normalization behavior. ChangesClaude count_tokens
Sequence DiagramsequenceDiagram
participant Client as Client
participant Handler as ClaudeCountTokens
participant Service as EstimateClaudeInputTokens
participant Estimator as EstimateTokenByModel
Client->>Handler: POST /v1/messages/count_tokens with body
Handler->>Handler: Unmarshal body -> dto.ClaudeRequest
alt Unmarshal fails
Handler-->>Client: 400 Bad Request (invalid_request_error)
else Success
Handler->>Service: EstimateClaudeInputTokens(&req)
Service->>Service: normalizeRequestTools(&req)
Service->>Service: meta := req.GetTokenCountMeta().CombineText
Service->>Estimator: EstimateTokenByModel(req.Model, meta)
Estimator-->>Service: token count
Service-->>Handler: input_tokens
Handler-->>Client: 200 OK { "input_tokens": N }
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (6)
.github/workflows/ghcr-publish.yml (3)
10-12: Redundant branch pattern.
feature/**already matchesfeature/count-tokens-local-estimate, so line 11 is dead configuration. Safe to remove unless you specifically want a visible marker.🧹 Proposed tidy-up
branches: - - feature/count-tokens-local-estimate - feature/**🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/ghcr-publish.yml around lines 10 - 12, The branches list contains a redundant pattern: "feature/**" already covers "feature/count-tokens-local-estimate"; remove the specific "feature/count-tokens-local-estimate" entry (or remove "feature/**" if you intended only the single branch) so only the necessary pattern ("feature/**" for all feature branches) remains; update the branches array accordingly to eliminate the dead configuration.
19-22: Consider adding aconcurrencygroup to avoid racing on the mutablecounttokenstag.Two rapid pushes to any
feature/**branch will run this job in parallel and race to overwriteghcr.io/daymade/new-api:counttokens. The per-SHA tag is safe, but the floating tag can end up pointing at an older commit if the jobs finish out of order. A branch-scoped concurrency group (withcancel-in-progress: true) keeps the latest push authoritative and also saves CI minutes.🛡️ Proposed addition
jobs: build-and-push: runs-on: ubuntu-latest + concurrency: + group: ghcr-publish-${{ github.ref }} + cancel-in-progress: true steps:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/ghcr-publish.yml around lines 19 - 22, The workflow's build-and-push job can race when two pushes to the same feature branch write the mutable ghcr tag "counttokens"; add a concurrency block on the job "build-and-push" using a branch-scoped group (e.g. use github.ref or github.ref_name with the feature/** pattern) and set cancel-in-progress: true so only the latest push wins and parallel runs are cancelled; update the job definition to include this concurrency stanza to prevent tag overwrite races.
57-59: Remove redundantTARGETOSandTARGETARCHbuild args.When
platforms: linux/amd64is set, BuildKit automatically providesTARGETOSandTARGETARCHas predefined global build args to every stage that declares them (which the Dockerfile does at lines 14-15). Explicitly passing them here is redundant and risks silently breaking cross-builds if you later expandplatformsto multi-architecture builds, since the hardcodedamd64would override the per-platform value.♻️ Proposed simplification
provenance: false sbom: false - build-args: | - TARGETOS=linux - TARGETARCH=amd64🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/ghcr-publish.yml around lines 57 - 59, The build workflow is passing redundant hardcoded build args TARGETOS and TARGETARCH which override BuildKit's per-platform values; remove the TARGETOS=linux and TARGETARCH=amd64 entries from the build-args block so that the platforms: linux/amd64 setting can allow BuildKit to inject the correct TARGETOS/TARGETARCH automatically (look for the build-args block containing TARGETOS and TARGETARCH in the ghcr-publish workflow).service/claude_token_estimator_test.go (1)
60-72: Consider exercising the handler path, not just the estimator.These tests validate the estimator in isolation, but the more interesting risk surface — the HTTP handler's body parsing and response shape — is not covered. A small
httptest-based test that posts raw JSON toClaudeCountTokensand asserts the{"input_tokens": N}envelope (plus the 400 error shape on malformed JSON) would catch regressions in the wire contract that the Anthropic SDK depends on.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@service/claude_token_estimator_test.go` around lines 60 - 72, Add an httptest-based integration test that exercises the HTTP handler ClaudeCountTokens rather than only calling EstimateClaudeInputTokens: construct raw JSON bodies (valid and malformed), POST them to the ClaudeCountTokens handler using httptest.NewRequest/httptest.NewRecorder, decode the JSON response and assert the envelope contains {"input_tokens": N} for valid input (N within the expected range from EstimateClaudeInputTokens) and that malformed JSON returns the 400 error shape; reuse dto.ClaudeRequest for building valid bodies and call EstimateClaudeInputTokens in the test to compute the expected token range to assert against the handler response.service/claude_token_estimator.go (2)
28-50: Estimator is non-deterministic across requests due to Go map ordering.
claudeToolsCharsrelies onjson.Marshal(tools). Whentoolscontains nestedmap[string]any(which is whatany-typed fields produce after generic unmarshaling — seedto.ClaudeRequest.Tools any),encoding/jsonsorts map keys, so this happens to be stable. However,ToolChoice,systemas[]any, and toolinput_schema.propertiesall decode tomap[string]anytoo — the "stable estimate" promise in the doc comment (lines 24–25) holds only because of that Go stdlib guarantee. Worth a one-line comment acknowledging the dependency so a future switch to a faster/non-sorting JSON encoder doesn't silently break the invariant.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@service/claude_token_estimator.go` around lines 28 - 50, The token estimator relies on json.Marshal's deterministic key-sorting to keep estimates stable; add a one-line comment in EstimateClaudeInputTokens (or directly above claudeToolsChars) noting that the stable-estimate guarantee depends on encoding/json's map key sorting for any map[string]any structures (e.g., dto.ClaudeRequest.Tools, ToolChoice, system as []any, and tool input_schema.properties), so switching to a non-sorting JSON encoder would break determinism.
9-27: Nice doc comment; minor: nothing prevents the "fast" claim from regressing.The header promises "fast, free local estimate" and calls out the 250 RPM fallback scenario. Given a sufficiently adversarial tool schema (deeply nested, MB-scale),
json.Marshalontoolscan still be CPU-visible. Consider capping the marshaled byte count you actually consume (e.g. stop at N KB) so the handler's latency is bounded even for pathological inputs — this aligns with the route's purpose of shedding traffic cheaply.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@service/claude_token_estimator.go` around lines 9 - 27, The comment warns that EstimateClaudeInputTokens currently json.Marshal(s) the full tools payload which can regress the "fast, free" promise for pathological inputs; modify the EstimateClaudeInputTokens implementation (the json.Marshal of tools / any full marshaling paths) to cap how many bytes are actually processed (e.g. stop after a configurable N KB like 64KB) by using a streaming/limited marshal approach (io.LimitReader/LimitWriter or write to a buffer and bail when length exceeds the cap) and then compute the chars/3 estimate from the truncated output (ensuring you still round up) so the handler's latency is bounded for large/deep tool schemas.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@controller/claude_count_tokens.go`:
- Line 33: The handler uses io.ReadAll(c.Request.Body) without any size limit,
allowing attackers to send arbitrarily large payloads; replace that call by
wrapping c.Request.Body with a size-limiting reader (e.g., http.MaxBytesReader
or io.LimitReader) using the project's configured max for Claude requests and
then read from that limited reader; update the code paths referencing body and
error handling in the claude_count_tokens handler to reject requests that exceed
the limit (return appropriate error) and ensure TokenAuth-protected endpoints
use the same cap.
- Line 4: Replace direct use of the standard library JSON package in
controller/claude_count_tokens.go with the project's JSON wrapper: remove the
"encoding/json" import and change all calls that currently use
json.Marshal/json.Unmarshal/json.NewDecoder or similar (e.g., in the functions
handling payload encoding/decoding around lines 46-55) to the corresponding
common package helpers such as common.Marshal(), common.Unmarshal(),
common.DecodeJson() or common.UnmarshalJsonStr() as appropriate, preserving the
same error handling and types; ensure imports include the common package and
update any variable names if needed to match the wrapper return signatures.
- Around line 32-43: Replace the manual body read/unmarshal in ClaudeCountTokens
with the project's helper: remove io and encoding/json imports, declare var req
dto.ClaudeRequest and call common.UnmarshalBodyReusable(c, &req); on error
return the same HTTP 400 JSON error payload but use the error message "invalid
JSON body: "+err.Error(); keep all other logic intact and reference the
ClaudeCountTokens handler and dto.ClaudeRequest when making the change.
In `@router/relay-router.go`:
- Around line 92-101: The ClaudeCountTokens route is registered on httpRouter
and therefore inherits relayV1Router middlewares (ModelRequestRateLimit,
Distribute) contrary to its docstring; create a sibling router group off
relayV1Router that only applies TokenAuth (and optionally RouteTag) and register
controller.ClaudeCountTokens on that group instead of httpRouter so it bypasses
ModelRequestRateLimit and Distribute; additionally, in
controller/claude_count_tokens.go replace the direct encoding/json.Unmarshal
call with common.Unmarshal to conform to project JSON helpers.
In `@service/claude_token_estimator.go`:
- Around line 81-89: claudeMessageChars currently uses
dto.ClaudeMessage.GetStringContent() which ignores non-text blocks (tool_result,
tool_use, image) and undercounts; update claudeMessageChars to fall back when
GetStringContent() returns empty but msg.Content != nil by computing the byte
length of json.Marshal(msg.Content) (or by explicitly extracting
tool_result.content and tool_use.input if you prefer) and return that length so
non-text blocks are accounted for; ensure you still return 0 for a nil msg and
preserve the existing behavior when GetStringContent() yields non-empty text.
- Line 4: Replace direct usage of the standard library JSON package with the
project's JSON wrapper: remove the "encoding/json" import and import the common
JSON wrapper (package common), then replace all json.Marshal and json.Unmarshal
calls in this file (the import at top and the calls around the areas flagged,
e.g., the json.Unmarshal call near line 74 and the json.Marshal call near line
95) with common.Unmarshal(...) and common.Marshal(...). Ensure the error
handling and variable types remain the same when switching to
common.Marshal/common.Unmarshal.
---
Nitpick comments:
In @.github/workflows/ghcr-publish.yml:
- Around line 10-12: The branches list contains a redundant pattern:
"feature/**" already covers "feature/count-tokens-local-estimate"; remove the
specific "feature/count-tokens-local-estimate" entry (or remove "feature/**" if
you intended only the single branch) so only the necessary pattern ("feature/**"
for all feature branches) remains; update the branches array accordingly to
eliminate the dead configuration.
- Around line 19-22: The workflow's build-and-push job can race when two pushes
to the same feature branch write the mutable ghcr tag "counttokens"; add a
concurrency block on the job "build-and-push" using a branch-scoped group (e.g.
use github.ref or github.ref_name with the feature/** pattern) and set
cancel-in-progress: true so only the latest push wins and parallel runs are
cancelled; update the job definition to include this concurrency stanza to
prevent tag overwrite races.
- Around line 57-59: The build workflow is passing redundant hardcoded build
args TARGETOS and TARGETARCH which override BuildKit's per-platform values;
remove the TARGETOS=linux and TARGETARCH=amd64 entries from the build-args block
so that the platforms: linux/amd64 setting can allow BuildKit to inject the
correct TARGETOS/TARGETARCH automatically (look for the build-args block
containing TARGETOS and TARGETARCH in the ghcr-publish workflow).
In `@service/claude_token_estimator_test.go`:
- Around line 60-72: Add an httptest-based integration test that exercises the
HTTP handler ClaudeCountTokens rather than only calling
EstimateClaudeInputTokens: construct raw JSON bodies (valid and malformed), POST
them to the ClaudeCountTokens handler using
httptest.NewRequest/httptest.NewRecorder, decode the JSON response and assert
the envelope contains {"input_tokens": N} for valid input (N within the expected
range from EstimateClaudeInputTokens) and that malformed JSON returns the 400
error shape; reuse dto.ClaudeRequest for building valid bodies and call
EstimateClaudeInputTokens in the test to compute the expected token range to
assert against the handler response.
In `@service/claude_token_estimator.go`:
- Around line 28-50: The token estimator relies on json.Marshal's deterministic
key-sorting to keep estimates stable; add a one-line comment in
EstimateClaudeInputTokens (or directly above claudeToolsChars) noting that the
stable-estimate guarantee depends on encoding/json's map key sorting for any
map[string]any structures (e.g., dto.ClaudeRequest.Tools, ToolChoice, system as
[]any, and tool input_schema.properties), so switching to a non-sorting JSON
encoder would break determinism.
- Around line 9-27: The comment warns that EstimateClaudeInputTokens currently
json.Marshal(s) the full tools payload which can regress the "fast, free"
promise for pathological inputs; modify the EstimateClaudeInputTokens
implementation (the json.Marshal of tools / any full marshaling paths) to cap
how many bytes are actually processed (e.g. stop after a configurable N KB like
64KB) by using a streaming/limited marshal approach (io.LimitReader/LimitWriter
or write to a buffer and bail when length exceeds the cap) and then compute the
chars/3 estimate from the truncated output (ensuring you still round up) so the
handler's latency is bounded for large/deep tool schemas.
🪄 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: CHILL
Plan: Pro
Run ID: f42381a1-157d-4c99-a5af-d170ba35439e
📒 Files selected for processing (5)
.github/workflows/ghcr-publish.ymlcontroller/claude_count_tokens.gorouter/relay-router.goservice/claude_token_estimator.goservice/claude_token_estimator_test.go
Comment on lines
+92
to
+101
| // Anthropic's /v1/messages/count_tokens — local estimate, does not | ||
| // hit any upstream channel. The Claude CLI / Anthropic SDK polls | ||
| // this endpoint frequently; serving a fast 200 here prevents the | ||
| // SDK from falling back to fake max_tokens=1 probes that flood | ||
| // upstream RPM quotas. See controller.ClaudeCountTokens for the | ||
| // rationale and the security boundary (TokenAuth still applies via | ||
| // the parent group; Distribute / PreConsume are intentionally | ||
| // skipped because no channel is selected). | ||
| httpRouter.POST("/messages/count_tokens", controller.ClaudeCountTokens) | ||
|
|
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm what Distribute() and ModelRequestRateLimit() do, to assess the impact.
rg -nP --type=go -C3 '\bfunc\s+Distribute\s*\('
rg -nP --type=go -C3 '\bfunc\s+ModelRequestRateLimit\s*\('
# Also verify no other route registrations bypass Distribute by using a sibling group under relayV1Router.
rg -nP --type=go -C2 'relayV1Router\.Group\b'Repository: QuantumNous/new-api
Length of output: 1452
🏁 Script executed:
sed -n '65,105p' router/relay-router.goRepository: QuantumNous/new-api
Length of output: 1640
🏁 Script executed:
# Check what ModelRequestRateLimit does and if it applies to count_tokens
sed -n '167,185p' middleware/model-rate-limit.go
# Check what Distribute does with channel selection
sed -n '30,50p' middleware/distributor.go
# Check the controller implementation
sed -n '1,50p' controller/claude.go | grep -A 30 "ClaudeCountTokens"Repository: QuantumNous/new-api
Length of output: 1548
🏁 Script executed:
# Find the ClaudeCountTokens controller
find . -type f -name "*.go" -exec grep -l "func ClaudeCountTokens" {} \;
# Search for ClaudeCountTokens definition
rg -n "func.*ClaudeCountTokens" --type=goRepository: QuantumNous/new-api
Length of output: 177
🏁 Script executed:
# Read the ClaudeCountTokens controller implementation
cat -n controller/claude_count_tokens.goRepository: QuantumNous/new-api
Length of output: 2246
Route placement contradicts the handler's documented design: move to a dedicated local-only group.
The ClaudeCountTokens handler's docstring (controller/claude_count_tokens.go:29–31) explicitly states it "is wired up under a dedicated router group that runs only TokenAuth," but the route is registered on httpRouter (router/relay-router.go:100), which inherits ModelRequestRateLimit from the parent relayV1Router (line 73) and applies Distribute() (line 85). This causes every request to:
- Hit
ModelRequestRateLimit, subjecting a local token estimate to per-model RPM limits, - Hit
Distribute, attempting channel selection for an endpoint that intentionally bypasses the channel pipeline.
Create a sibling group under relayV1Router with only TokenAuth middleware (and optionally RouteTag for observability):
🔧 Proposed fix
relayV1Router := router.Group("/v1")
relayV1Router.Use(middleware.RouteTag("relay"))
relayV1Router.Use(middleware.SystemPerformanceCheck())
relayV1Router.Use(middleware.TokenAuth())
relayV1Router.Use(middleware.ModelRequestRateLimit())
{
// WebSocket 路由（统一到 Relay）
wsRouter := relayV1Router.Group("")
wsRouter.Use(middleware.Distribute())
wsRouter.GET("/realtime", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatOpenAIRealtime)
})
}
+ {
+ // Local-only endpoints: no Distribute, no rate limit.
+ localRouter := relayV1Router.Group("")
+ localRouter.POST("/messages/count_tokens", controller.ClaudeCountTokens)
+ }
{
//http router
httpRouter := relayV1Router.Group("")
httpRouter.Use(middleware.Distribute())
httpRouter.POST("/messages", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatClaude)
})
- // Anthropic's /v1/messages/count_tokens — …
- httpRouter.POST("/messages/count_tokens", controller.ClaudeCountTokens)Also, replace the direct encoding/json.Unmarshal() call at controller/claude_count_tokens.go:46 with common.Unmarshal() per Go coding guidelines.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@router/relay-router.go` around lines 92 - 101, The ClaudeCountTokens route is
registered on httpRouter and therefore inherits relayV1Router middlewares
(ModelRequestRateLimit, Distribute) contrary to its docstring; create a sibling
router group off relayV1Router that only applies TokenAuth (and optionally
RouteTag) and register controller.ClaudeCountTokens on that group instead of
httpRouter so it bypasses ModelRequestRateLimit and Distribute; additionally, in
controller/claude_count_tokens.go replace the direct encoding/json.Unmarshal
call with common.Unmarshal to conform to project JSON helpers.
daymade
force-pushed
the
feature/count-tokens-local-estimate
branch
from
April 24, 2026 12:34
f91a002 to
40cc2c4
Compare
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (1)
service/claude_token_estimator.go (1)
3-7:⚠️ Potential issue | 🟠 MajorReplace
encoding/jsonwith the project's wrapper.Per coding guidelines, all JSON marshal/unmarshal operations must use
common.Marshal()/common.Unmarshal()—encoding/jsonmust not be imported in business code. This was flagged in the prior review and remains unaddressed in this revision.🔧 Proposed fix
import ( - "encoding/json" - + "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" ) @@ - b, err := json.Marshal(m) + b, err := common.Marshal(m) if err != nil { continue } if _, isWebSearch := m["type"]; isWebSearch { var ws dto.ClaudeWebSearchTool - if err := json.Unmarshal(b, &ws); err == nil && ws.Type != "" { + if err := common.Unmarshal(b, &ws); err == nil && ws.Type != "" { normalized = append(normalized, &ws) } continue } var tool dto.Tool - if err := json.Unmarshal(b, &tool); err == nil && tool.Name != "" { + if err := common.Unmarshal(b, &tool); err == nil && tool.Name != "" { normalized = append(normalized, &tool) }As per coding guidelines: "All JSON marshal/unmarshal operations MUST use wrapper functions from
common/json.go:common.Marshal(),common.Unmarshal(),common.UnmarshalJsonStr(),common.DecodeJson(),common.GetJsonType(). Do NOT directly import or callencoding/jsonin business code."Also applies to: 70-83
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@service/claude_token_estimator.go` around lines 3 - 7, The import of the standard "encoding/json" package must be removed and any JSON operations in this file (including the code referenced around lines 70-83) should be rewritten to use the project's JSON wrapper functions in the common package (e.g., common.Marshal(), common.Unmarshal(), common.UnmarshalJsonStr(), common.DecodeJson(), or common.GetJsonType()); update imports to include the common package instead of "encoding/json" and replace direct calls to json.Marshal/json.Unmarshal (or similar) in functions such as those handling DTOs (referenced by the dto usage) with the appropriate common.* wrapper so the file no longer imports or calls encoding/json directly.
🧹 Nitpick comments (1)
service/claude_token_estimator_test.go (1)
3-8: Usecommon.Unmarshalin tests for consistency with the project JSON policy.The JSON-wrapper rule applies to all
*.gofiles, not just production code. Please swapencoding/jsonfor the project wrapper here as well so this file does not become a precedent for re-introducing directencoding/jsonusage.🔧 Proposed fix
import ( - "encoding/json" "testing" + "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" ) @@ - if err := json.Unmarshal([]byte(tc.body), &req); err != nil { + if err := common.Unmarshal([]byte(tc.body), &req); err != nil { t.Fatalf("unmarshal: %v", err) }As per coding guidelines: "All JSON marshal/unmarshal operations MUST use wrapper functions from
common/json.go... Do NOT directly import or callencoding/jsonin business code."Also applies to: 70-76
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@service/claude_token_estimator_test.go` around lines 3 - 8, The test imports encoding/json directly and calls json.Unmarshal (lines ~70-76); replace that import with the project JSON wrapper (import common from "github.com/QuantumNous/new-api/common") and use common.Unmarshal in place of json.Unmarshal in claude_token_estimator_test.go so tests follow the project's JSON wrapper policy and remove the direct encoding/json usage.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@service/claude_token_estimator.go`:
- Around line 3-7: The import of the standard "encoding/json" package must be
removed and any JSON operations in this file (including the code referenced
around lines 70-83) should be rewritten to use the project's JSON wrapper
functions in the common package (e.g., common.Marshal(), common.Unmarshal(),
common.UnmarshalJsonStr(), common.DecodeJson(), or common.GetJsonType()); update
imports to include the common package instead of "encoding/json" and replace
direct calls to json.Marshal/json.Unmarshal (or similar) in functions such as
those handling DTOs (referenced by the dto usage) with the appropriate common.*
wrapper so the file no longer imports or calls encoding/json directly.
---
Nitpick comments:
In `@service/claude_token_estimator_test.go`:
- Around line 3-8: The test imports encoding/json directly and calls
json.Unmarshal (lines ~70-76); replace that import with the project JSON wrapper
(import common from "github.com/QuantumNous/new-api/common") and use
common.Unmarshal in place of json.Unmarshal in claude_token_estimator_test.go so
tests follow the project's JSON wrapper policy and remove the direct
encoding/json usage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8b982c07-be9d-4131-978b-6af54f4927a6
📒 Files selected for processing (4)
controller/claude_count_tokens.gorouter/relay-router.goservice/claude_token_estimator.goservice/claude_token_estimator_test.go
Author
|
Round 3 已处理 CodeRabbit Round 2 的
Round 2 已修的几条复核（相关 reviewer 可略过）：
所有 15 个单测仍然全过： ``` Image token 精算 + docstring coverage 如需要可另开 follow-up issue。 |
Author
|
@6639835 谢谢反馈。评估过 7 种替代方案，列一下结论供参考：
至于 "SDK 会不会 cache 上游不支持 endpoint、永远走 如果你有具体的其他方向，欢迎补充——尤其是：(1) 精度提升对哪个实际调用链有影响？(2) 外部 tokenizer 跟版本漂移怎么维护？ |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@service/claude_token_estimator.go`:
- Around line 73-83: The current heuristic treats any tool JSON with a top-level
"type" key as a ClaudeWebSearchTool which misclassifies many Anthropic tools;
update the detection in the block that reads m["type"] (in
claude_token_estimator.go) to read the type value as a string and only route to
ClaudeWebSearchTool (ws) when the type value indicates web-search (e.g.,
strings.HasPrefix(typeVal, "web_search") or exact "web_search"); otherwise allow
the payload to fall through to the dto.Tool path (tool) so ProcessTools /
CombineText see the full tool structure (Description, InputSchema, etc.).
Optionally mention in a comment near this logic that a whitelist of web-search
prefixes or checking for web-search-specific fields like UserLocation is
preferred for future robustness.
🪄 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: CHILL
Plan: Pro
Run ID: 069ae5fc-6432-42b6-9af5-868ccdae9b1e
📒 Files selected for processing (2)
service/claude_token_estimator.goservice/claude_token_estimator_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- service/claude_token_estimator_test.go
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
service/claude_token_estimator.go (1)
87-99: Minor: malformedweb_search_*tools are silently dropped.When a tool entry has
typestarting with"web_search"butcommon.UnmarshalintoClaudeWebSearchToolfails (or produces an emptyType), thecontinueon line 93 prevents fallthrough to thedto.Toolpath. That means a payload that looks like a web-search tool but is malformed contributes 0 chars to the estimate, while a non-web-search server tool with a similar issue would also be skipped (via the empty-Nameguard at line 101). This is consistent and fine for an estimate endpoint, but worth a one-line comment so future readers don't try to "fix" the unconditional continue.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@service/claude_token_estimator.go` around lines 87 - 99, The block that checks m["type"] and unmarshals into dto.ClaudeWebSearchTool (using common.Unmarshal into ws, then normalized = append(...), and the unconditional continue) can silently drop malformed web_search_* entries; add a one-line comment immediately above the continue explaining that failing/untyped web_search_* tools are intentionally not falling through to the dto.Tool path (so they contribute 0 chars) to prevent future maintainers from removing the continue thinking it's a bug—reference the variables/type checks (typeVal/typeStr), the common.Unmarshal into ws, the continue, and the fallback dto.Tool/ProcessTools behavior in the comment.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@service/claude_token_estimator.go`:
- Around line 87-99: The block that checks m["type"] and unmarshals into
dto.ClaudeWebSearchTool (using common.Unmarshal into ws, then normalized =
append(...), and the unconditional continue) can silently drop malformed
web_search_* entries; add a one-line comment immediately above the continue
explaining that failing/untyped web_search_* tools are intentionally not falling
through to the dto.Tool path (so they contribute 0 chars) to prevent future
maintainers from removing the continue thinking it's a bug—reference the
variables/type checks (typeVal/typeStr), the common.Unmarshal into ws, the
continue, and the fallback dto.Tool/ProcessTools behavior in the comment.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 3d79ac81-3f79-40b3-a29a-3b0b5beb7719
📒 Files selected for processing (2)
service/claude_token_estimator.goservice/claude_token_estimator_test.go
daymade
force-pushed
the
feature/count-tokens-local-estimate
branch
from
May 12, 2026 23:28
5a9421f to
7a65744
Compare
|
非常好 PR！我们也急需这个功能，希望可以尽快合并😊 |
11 tasks
|
现在是不是 /v1/messages/count_tokens现在都不支持？ |
我的分支 已经实现了这个接口，不过目前仅限 AWS bedrock 渠道，你可以试试，直接基于这份代码构建 docker 镜像就行了 |
我也写了，不过是透传上游。最近一直在排查相关非流请求循环的BUG，最后发现问题在这 |
11 tasks
11 tasks
Implements Anthropic's POST /v1/messages/count_tokens endpoint (https://docs.claude.com/en/api/messages-count-tokens) by estimating input_tokens locally instead of forwarding to any upstream channel. Why --- The Anthropic JS SDK and Claude CLI poll this endpoint before each chat to size the context window. When new-api returns 404 (which it does today, since the route is unregistered), the SDK falls back to sending fake max_tokens=1 messages with the entire tool schema attached. We observed bursts of 250 RPM of these probes from a single Claude Code Desktop session, exhausting the upstream provider's RPM quota and starving real traffic on the same channel for ~2 minutes. Returning a fast 200 here from new-api itself fixes the failure mode at the source: the SDK is satisfied, no upstream RPM is consumed, no billing entry is created. Closes QuantumNous#1694 Closes QuantumNous#2847 Closes QuantumNous#1979 Approach -------- - Mounted on relayV1Router directly, NOT under httpRouter — the former gives us TokenAuth + RouteTag + SystemPerformanceCheck + ModelRequestRateLimit; the latter would also pull in middleware.Distribute() which selects a channel and starts the PreConsume flow. Anthropic defines count_tokens as token-counting only with no quota impact, so the bypass is intentional. - Estimation reuses the project's existing pieces: * ClaudeRequest.GetTokenCountMeta() for the canonical flattening of system / messages (text, tool_use, tool_result) / tools * EstimateTokenByModel() for the Claude-tuned tokenizer that /v1/messages itself bills against This keeps the count consistent with what /v1/messages would report on the same body, so callers can reason about both numbers together. - One small helper, normalizeRequestTools(), converts raw map[string]any tool entries (what json.Unmarshal produces when the field is `any`) into the typed *dto.Tool / *dto.ClaudeWebSearchTool values that dto.ProcessTools requires. Without this every tool entry on a count_tokens request would be silently dropped on the `default: continue` arm of ProcessTools (dto/claude.go:439-442) and the tools text — usually 80%+ of a CLI probe body — would not be counted. - Body parsing uses common.UnmarshalBodyReusable so the request body remains available to any logging/observability that runs after this handler. - Image tokens are intentionally not added: getImageToken() needs a RelayInfo + http.Request context, which this route doesn't have. The CLI probe (the failure mode this PR exists to mitigate) never carries images. Documented in the estimator's godoc. How vs. PR QuantumNous#2384 ---------------- QuantumNous#2384 took the deeper path: a new RelayFormat enum value + adaptor changes + middleware.distributor changes + autoban bypass. That PR was closed without merge. This PR keeps the surface area small — one new route, one controller, one estimator, no enum/adaptor changes, no impact on any code path that doesn't hit the new route. Files ----- - router/relay-router.go + 9 lines (one route registration) - controller/claude_count_tokens.go new (~35 lines incl. doc) - service/claude_token_estimator.go new (~90 lines) - service/claude_token_estimator_test.go new (12 cases) Tests ----- $ go test ./service/ -run 'TestEstimateClaude|TestNormalizeRequestTools' -v === RUN TestEstimateClaudeInputTokens --- PASS (10 cases: empty / short user / system+user / CLI probe / system as array / CJK / tool_use / tool_result / web-search / pre-typed tools) === RUN TestEstimateClaudeInputTokens_NilSafety --- PASS (2 cases: nil request / zero-value request) === RUN TestNormalizeRequestTools --- PASS (3 cases: nil / non-array / malformed dropped) Manual verification of the failure mode (sample CLI probe body — trivial message + Bash tool schema, ~200 bytes): EstimateClaudeInputTokens = 43 input tokens Endpoint returns: 200 OK {"input_tokens": 43} Claude CLI accepts and stops sending fake max_tokens=1 fallback.
…ideline CodeRabbit round 2 flagged service/claude_token_estimator.go + claude_token_estimator_test.go as still importing encoding/json after round 1. Replaces all three json.Marshal/Unmarshal calls in normalizeRequestTools + one json.Unmarshal call in the table-driven test body with the project's common.Marshal / common.Unmarshal wrappers. All 15 estimator/normalize test cases still pass. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
CodeRabbit round 3 flagged that normalizeRequestTools was routing any tool JSON with a top-level "type" field into ClaudeWebSearchTool. But Anthropic's other built-in server tools (computer_*, bash_*, text_editor_*, code_execution_*, mcp_*) also carry a top-level "type" — their shape is different from ClaudeWebSearchTool, so computer_use and friends were being unmarshalled into the web-search struct. After that, ProcessTools only saw Name + UserLocation and dropped the rest of the tool payload from the estimate. Fix: narrow the web-search branch to type values whose string begins with "web_search" (matches current web_search_20250305 and any future version suffix). Unknown "type" values fall through to the generic dto.Tool path, so ProcessTools still sees Name + Description + InputSchema. Caveat: dto.Tool does not model every server-tool-specific field (display_width_px on computer_use, etc.), so exotic server tools can still be undercounted. That is an upstream schema gap outside this PR's scope — this fix just stops the regression from misrouting them all into ClaudeWebSearchTool. Two new test cases lock in the behavior: - computer_20250124 routes to *dto.Tool (not *ClaudeWebSearchTool) - web_search_20250604 (future version suffix) still routes to *ClaudeWebSearchTool via the prefix rule Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
daymade
force-pushed
the
feature/count-tokens-local-estimate
branch
from
July 18, 2026 11:03
7a65744 to
cb1d4c5
Compare
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
5 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.
📝 变更描述 / Description
实现 Anthropic
/v1/messages/count_tokens端点的本地估算，避免 Claude SDK / CLI 在该端点 404 时退化为发送max_tokens=1的伪请求。真实事故场景：单个 Claude Code Desktop session 在一分钟内对同一上游 channel 发出 250 个伪 probe 请求（每个请求体 ~53KB，全是 tool schema 但
max_tokens=1），打爆上游 RPM 配额，导致同 channel 上其他用户在 ~2 分钟内收到 429。本地估算端点解决根因：SDK 拿到 200 + 一个合理的input_tokens数字就不会 fallback。实现要点：
relayV1Router上而不是httpRouter，跳过middleware.Distribute()——这与 Anthropic 规范一致（count_tokens 不路由到上游、不消费配额）ClaudeRequest.GetTokenCountMeta()+EstimateTokenByModel()，与/v1/messages实际计费的 tokenizer 一致normalizeRequestTools()修了dto.ProcessTools在收到 rawmap[string]anytool 时静默丢弃的 latent bug——CLI probe 体里 80%+ 的字节都是 tool schema，不修这个 estimate 会严重偏低🚀 变更类型 / Type of change
🔗 关联任务 / Related Issue
参考 #2384（已 close 未 merge）。本 PR 与之相比：不引入新 RelayFormat enum、不改 adaptor、不改 distributor、不引入 autoban 旁路。改动面只有一条新路由 + 一个 controller + 一个 estimator，对所有未触达
/v1/messages/count_tokens的代码路径零影响。✅ 提交前检查项 / Checklist
httpRouter中间件链含Distribute()，刻意挂到relayV1Router上避开。go test ./service/ -run 'TestEstimateClaude|TestNormalizeRequestTools' -v全部通过；go build ./...通过；本地 curl 实测端点返回正确 shape。types.ClaudeError）+ body parser（common.UnmarshalBodyReusable）。📸 运行证明 / Proof of Work
1. 单元测试 (12 cases passed)
2. 端到端: 不带 PR 时 SDK 行为 (实测自我们生产环境的事故)
CLI 检测到
/v1/messages/count_tokens返 404，退化为发送伪请求到/v1/messages。截取一个真实的 53KB 伪请求体（已脱敏，源自我们的 OSS 归档）：{ "model": "claude-haiku-4-5-20251001", "max_tokens": 1, "messages": [{"role": "user", "content": "count"}], "tools": [ /* 18 个完整 tool schema，~52KB，含 Bash / AskUserQuestion / ... */ ] }请求头
user-agent: claude-cli/2.1.111 (external, claude-desktop-3p, agent-sdk/0.2.111)。同一 device_id + session_id 在 60 秒内连发 250 个，149 个被上游返429 rate limit reached for RPM。3. 带 PR 时端点行为 (实测自本地 build)
错误路径 (malformed body):
维护者可能的问题 / Anticipated reviewer questions
Q: 为什么不直接转发到上游？
A: 大部分 Claude 兼容上游（Qiniu QNAIGC / OpenRouter / 等）自己也没实现这个端点 (实测 Qiniu 返 404)。即使是 Anthropic 官方端点，虽免费但仍占 RPM 配额，pass-through 不解决我们这个 PR 要解的问题。
Q: 为什么不用真 tokenizer (tiktoken / anthropic-tokenizer)?
A: 项目已有
EstimateTokenByModel(model, text)带 Claude-tuned multipliers (Word 1.13, CJK 1.21, MathSymbol 4.52, ...)，用在service/token_counter.go等多处。本 PR 不引入新依赖，复用已有实现，估算与/v1/messages自身计费一致。Q:
normalizeRequestTools是否应该改dto.ProcessTools而不是在 estimator 自己 normalize？A: 改
ProcessTools是更彻底的 fix，但会影响所有 caller（distributor / handler / 其他 future code paths）的行为，scope 太大、回归风险也大。本 PR 只改自己 reach 到的代码，把 latent bug 在 estimator 边界处兜住。如果维护者倾向把 normalize 提到ProcessTools，欢迎反馈，我可以再开一个 PR 单独处理。Q: 图片 token 不算合理吗？
A: Anthropic 规范说 count_tokens 应当包含图片。我们权衡的是：(1)
getImageToken需要RelayInfo + http.Request，本路由刻意不走 channel pipeline；(2) 触发本 PR 的 CLI probe 失败模式不带图片；(3) 真要加图片可后续 PR，需要把 image-token 计算从service/token_counter.go中抽成 standalone helper。Summary by CodeRabbit