diff --git a/controller/claude_count_tokens.go b/controller/claude_count_tokens.go new file mode 100644 index 000000000000..4162bb01519c --- /dev/null +++ b/controller/claude_count_tokens.go @@ -0,0 +1,35 @@ +package controller + +import ( + "net/http" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" +) + +// ClaudeCountTokens implements POST /v1/messages/count_tokens. +// +// The estimate is computed locally — no upstream channel is selected +// and no quota is consumed. This is what makes the endpoint suitable +// for the SDK / CLI to poll before every chat. +func ClaudeCountTokens(c *gin.Context) { + var req dto.ClaudeRequest + if err := common.UnmarshalBodyReusable(c, &req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "type": "error", + "error": types.ClaudeError{ + Type: "invalid_request_error", + Message: "invalid JSON body: " + err.Error(), + }, + }) + return + } + + c.JSON(http.StatusOK, gin.H{ + "input_tokens": service.EstimateClaudeInputTokens(&req), + }) +} diff --git a/router/relay-router.go b/router/relay-router.go index 17a13cad7fd6..23c66ac8aac1 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -71,6 +71,16 @@ func SetRelayRouter(router *gin.Engine) { relayV1Router.Use(middleware.SystemPerformanceCheck()) relayV1Router.Use(middleware.TokenAuth()) relayV1Router.Use(middleware.ModelRequestRateLimit()) + + // Anthropic's /v1/messages/count_tokens — local estimate only. Mounted + // directly on relayV1Router so it inherits TokenAuth + RouteTag + + // SystemPerformanceCheck + ModelRequestRateLimit, but NOT + // middleware.Distribute(). No channel is selected, no quota is consumed. + // Anthropic's spec defines this endpoint as token-counting only; the + // project's Distribute pipeline would otherwise pick a channel + run + // PreConsume on a request that never reaches an upstream. + relayV1Router.POST("/messages/count_tokens", controller.ClaudeCountTokens) + { // WebSocket 路由(统一到 Relay) wsRouter := relayV1Router.Group("") diff --git a/service/claude_token_estimator.go b/service/claude_token_estimator.go new file mode 100644 index 000000000000..6443713aa25b --- /dev/null +++ b/service/claude_token_estimator.go @@ -0,0 +1,106 @@ +package service + +import ( + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" +) + +// EstimateClaudeInputTokens returns a free, fast local estimate of +// input_tokens for an Anthropic /v1/messages/count_tokens request. +// +// Reuses ClaudeRequest.GetTokenCountMeta to do the canonical flattening +// of system / messages (text, tool_use, tool_result) / tools, and then +// the project's existing Claude-tuned tokenizer EstimateTokenByModel to +// turn that into a token count. This means the value matches whatever +// /v1/messages itself would report on the same body, so callers can +// reason about the two numbers consistently. +// +// Image tokens are intentionally NOT added: getImageToken requires a +// RelayInfo + http.Request context, which this route deliberately +// avoids by bypassing the channel pipeline. The Claude CLI context-bar +// probe (the failure mode this endpoint exists to mitigate) never +// carries images, so the gap doesn't matter for that case. Requests +// that do contain images will under-estimate, which is documented and +// acceptable for an estimate endpoint. +func EstimateClaudeInputTokens(req *dto.ClaudeRequest) int { + if req == nil { + return 0 + } + normalizeRequestTools(req) + meta := req.GetTokenCountMeta() + if meta == nil { + return 0 + } + return EstimateTokenByModel(req.Model, meta.CombineText) +} + +// normalizeRequestTools converts raw map[string]any entries (which is what +// json.Unmarshal of `tools` produces when the field is declared as `any`) +// into the typed *dto.Tool / *dto.ClaudeWebSearchTool values that +// dto.ProcessTools (called from GetTokenCountMeta) accepts. Without this, +// every tool entry on a count_tokens request is silently dropped on the +// `default: continue` arm of ProcessTools — see dto/claude.go:439-442. +// +// Mutates req.Tools in place. Safe because the request is parsed in this +// handler and not shared with anything else (count_tokens does not enter +// the channel pipeline). +// +// Web-search tools are distinguished from regular user-defined tools by a +// top-level "type" field whose value begins with "web_search" (e.g. +// "web_search_20250305", "web_search_20250604"). Other Anthropic server tools +// — computer_*, bash_*, text_editor_*, code_execution_*, mcp_* — also carry +// a top-level "type" but have a completely different shape from +// ClaudeWebSearchTool; if we unmarshalled them into that struct we would lose +// the rest of their schema from the estimate (ProcessTools would then only +// see Name + UserLocation). Prefix-matching the web-search family is a +// deliberate narrow allowlist: unknown "type" values fall through to the +// dto.Tool path so ProcessTools still sees Name / Description / InputSchema. +// (dto.Tool doesn't model every server-tool field, so undercount can still +// happen for exotic server tools — that's an upstream schema gap outside +// this PR's scope; the fallthrough at least stops the regression from +// misrouting.) +func normalizeRequestTools(req *dto.ClaudeRequest) { + if req == nil || req.Tools == nil { + return + } + rawTools, ok := req.Tools.([]any) + if !ok { + return + } + normalized := make([]any, 0, len(rawTools)) + for _, t := range rawTools { + switch t.(type) { + case *dto.Tool, dto.Tool, *dto.ClaudeWebSearchTool, dto.ClaudeWebSearchTool: + normalized = append(normalized, t) + continue + } + m, ok := t.(map[string]any) + if !ok { + continue + } + b, err := common.Marshal(m) + if err != nil { + continue + } + if typeVal, hasType := m["type"]; hasType { + if typeStr, ok := typeVal.(string); ok && strings.HasPrefix(typeStr, "web_search") { + var ws dto.ClaudeWebSearchTool + if err := common.Unmarshal(b, &ws); err == nil && ws.Type != "" { + normalized = append(normalized, &ws) + } + continue + } + // Non-web-search server tools (computer_*, bash_*, ...): fall + // through to the dto.Tool path below so ProcessTools still sees + // Name / Description / InputSchema instead of truncating the + // tool to just Name + UserLocation. + } + var tool dto.Tool + if err := common.Unmarshal(b, &tool); err == nil && tool.Name != "" { + normalized = append(normalized, &tool) + } + } + req.Tools = normalized +} diff --git a/service/claude_token_estimator_test.go b/service/claude_token_estimator_test.go new file mode 100644 index 000000000000..ec5674018868 --- /dev/null +++ b/service/claude_token_estimator_test.go @@ -0,0 +1,190 @@ +package service + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" +) + +func TestEstimateClaudeInputTokens(t *testing.T) { + cases := []struct { + name string + body string + want int + }{ + { + name: "empty messages", + body: `{"model":"claude-haiku-4-5","max_tokens":1024,"messages":[]}`, + want: 0, + }, + { + name: "single short user message", + body: `{"model":"claude-haiku-4-5","max_tokens":1024,"messages":[{"role":"user","content":"count"}]}`, + want: 4, + }, + { + name: "string system + user", + body: `{"model":"claude-haiku-4-5","max_tokens":1024,"system":"You are helpful.","messages":[{"role":"user","content":"hello"}]}`, + want: 9, + }, + { + name: "CLI-probe shape: trivial message + bash tool schema", + body: `{"model":"claude-haiku-4-5","max_tokens":1,"messages":[{"role":"user","content":"count"}],"tools":[{"name":"Bash","description":"Execute a shell command on the host","input_schema":{"type":"object","properties":{"command":{"type":"string"}},"required":["command"]}}]}`, + want: 43, + }, + { + name: "system as array of text blocks", + body: `{"model":"claude-haiku-4-5","max_tokens":1024,"system":[{"type":"text","text":"You are a helpful assistant."}],"messages":[{"role":"user","content":"hi"}]}`, + want: 12, + }, + { + name: "CJK content", + body: `{"model":"claude-haiku-4-5","max_tokens":1024,"messages":[{"role":"user","content":"你好世界,请用中文回答我的问题"}]}`, + want: 20, + }, + { + name: "tool_use block on assistant turn", + body: `{"model":"claude-haiku-4-5","max_tokens":1024,"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"Bash","input":{"command":"ls"}}]}]}`, + want: 10, + }, + { + name: "tool_result block on user turn", + body: `{"model":"claude-haiku-4-5","max_tokens":1024,"messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"file1\nfile2"}]}]}`, + want: 9, + }, + { + name: "web-search tool (different schema, distinguished by type field)", + body: `{"model":"claude-haiku-4-5","max_tokens":1024,"messages":[{"role":"user","content":"news"}],"tools":[{"type":"web_search_20250305","name":"web_search","max_uses":3}]}`, + want: 7, + }, + { + name: "tools array with already-typed entries (defensive: caller pre-normalized)", + // Synthesised in-test below via direct ClaudeRequest construction. + body: "", + want: 18, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var req dto.ClaudeRequest + if tc.body != "" { + if err := common.Unmarshal([]byte(tc.body), &req); err != nil { + t.Fatalf("unmarshal: %v", err) + } + } else { + // pre-typed tool entry — exercises the early-return arm of + // normalizeRequestTools. + req = dto.ClaudeRequest{ + Model: "claude-haiku-4-5", + Messages: []dto.ClaudeMessage{ + {Role: "user", Content: "list"}, + }, + Tools: []any{ + &dto.Tool{ + Name: "ListFiles", + Description: "List directory contents", + InputSchema: map[string]any{"type": "object"}, + }, + }, + } + } + got := EstimateClaudeInputTokens(&req) + if got != tc.want { + t.Errorf("EstimateClaudeInputTokens() = %d, want %d", got, tc.want) + } + }) + } +} + +func TestEstimateClaudeInputTokens_NilSafety(t *testing.T) { + if got := EstimateClaudeInputTokens(nil); got != 0 { + t.Errorf("nil request should return 0, got %d", got) + } + + // Empty model string + no content — should not panic. + if got := EstimateClaudeInputTokens(&dto.ClaudeRequest{}); got != 0 { + t.Errorf("zero-value request should return 0, got %d", got) + } +} + +func TestNormalizeRequestTools(t *testing.T) { + t.Run("nil tools", func(t *testing.T) { + req := &dto.ClaudeRequest{} + normalizeRequestTools(req) + if req.Tools != nil { + t.Errorf("expected Tools to remain nil, got %v", req.Tools) + } + }) + + t.Run("non-array tools field is left alone", func(t *testing.T) { + req := &dto.ClaudeRequest{Tools: "not-an-array"} + normalizeRequestTools(req) + if got, ok := req.Tools.(string); !ok || got != "not-an-array" { + t.Errorf("expected non-array Tools to be preserved, got %v", req.Tools) + } + }) + + t.Run("malformed tool entry is dropped, others survive", func(t *testing.T) { + req := &dto.ClaudeRequest{Tools: []any{ + map[string]any{"name": "Good", "description": "ok"}, + map[string]any{"description": "missing name"}, // no Name → dropped + "a string entry", // wrong shape → dropped + }} + normalizeRequestTools(req) + got, ok := req.Tools.([]any) + if !ok { + t.Fatalf("expected []any after normalize, got %T", req.Tools) + } + if len(got) != 1 { + t.Errorf("expected 1 surviving tool, got %d (%v)", len(got), got) + } + }) + + t.Run("non-web-search server tool routes to dto.Tool, not ClaudeWebSearchTool", func(t *testing.T) { + // Anthropic's built-in server tools (computer_*, bash_*, text_editor_*, + // code_execution_*, mcp_*) also carry a top-level "type" field, but + // their shape is different from ClaudeWebSearchTool. Previously they + // were routed into ClaudeWebSearchTool and ProcessTools then only saw + // Name + UserLocation, dropping the rest of the tool schema from the + // estimate. The fix narrows web-search routing to a "web_search" + // prefix allowlist; unknown types must fall through to dto.Tool. + req := &dto.ClaudeRequest{Tools: []any{ + map[string]any{ + "type": "computer_20250124", + "name": "computer", + "description": "Operate a virtual desktop", + "input_schema": map[string]any{"type": "object"}, + }, + }} + normalizeRequestTools(req) + got, ok := req.Tools.([]any) + if !ok || len(got) != 1 { + t.Fatalf("expected 1 tool, got %v (%T)", req.Tools, req.Tools) + } + if _, isWS := got[0].(*dto.ClaudeWebSearchTool); isWS { + t.Errorf("computer_20250124 must NOT be routed to *ClaudeWebSearchTool") + } + if _, isTool := got[0].(*dto.Tool); !isTool { + t.Errorf("expected *dto.Tool, got %T", got[0]) + } + }) + + t.Run("web_search variant (future version suffix) still routes to ClaudeWebSearchTool via prefix rule", func(t *testing.T) { + req := &dto.ClaudeRequest{Tools: []any{ + map[string]any{ + "type": "web_search_20250604", + "name": "web_search", + }, + }} + normalizeRequestTools(req) + got, ok := req.Tools.([]any) + if !ok || len(got) != 1 { + t.Fatalf("expected 1 tool, got %v (%T)", req.Tools, req.Tools) + } + if _, isWS := got[0].(*dto.ClaudeWebSearchTool); !isWS { + t.Errorf("web_search_20250604 should route to *ClaudeWebSearchTool, got %T", got[0]) + } + }) +}