feat(claude): support /v1/messages/count_tokens with local estimation

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),
+	})
+}

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("")

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
+}

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])
+		}
+	})
+}