Add projects

backend/internal/runtime/executor/claude_executor_tokens.go

@@ -0,0 +1,298 @@
+package executor
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
+	"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
+	"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
+	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
+	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
+	sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
+	log "github.com/sirupsen/logrus"
+	"github.com/tidwall/gjson"
+	"github.com/tidwall/sjson"
+)
+
+func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
+	apiKey, baseURL := claudeCreds(auth)
+	if baseURL == "" {
+		baseURL = "https://api.anthropic.com"
+	}
+	// Only Anthropic's first-party origin has the measured native count_tokens
+	// contract. Every custom/third-party base URL keeps local estimation,
+	// regardless of whether the credential is OAuth or an API key.
+	if shouldUseClaudeUpstreamTokenCount(apiKey, baseURL) {
+		return e.countTokensUpstream(ctx, auth, req, opts)
+	}
+
+	baseModel := thinking.ParseSuffix(req.Model).ModelName
+	from := opts.SourceFormat
+	responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts)
+	to := sdktranslator.FromString("claude")
+
+	// Use streaming translation to preserve function calling, except for claude.
+	stream := from != to
+	body := helps.TranslateRequestWithAPIKeyModelCompatibility(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, stream, helps.APIKeyModelIsCompat(req))
+	var errThinking error
+	body, errThinking = helps.ApplyRequestThinking(body, req, opts, from.String(), to.String(), e.Identifier())
+	if errThinking != nil {
+		return cliproxyexecutor.Response{}, errThinking
+	}
+	if rebuildMidSystemMessageEnabled(e.cfg, auth) {
+		body = rebuildMidSystemMessagesToTopLevel(body)
+	}
+	body = sanitizeClaudeMessagesForClaudeUpstreamWithDebug(ctx, body, baseModel, helps.APIKeyModelIsCompat(req))
+	if errValidate := validateClaudeTokenCountRequest(body); errValidate != nil {
+		return cliproxyexecutor.Response{}, errValidate
+	}
+
+	// Custom API-key gateways without a native count_tokens contract continue to
+	// use the local estimator without injecting generation-only CLI instructions.
+	count, err := helps.CountClaudeInputTokens(body)
+	if err != nil {
+		return cliproxyexecutor.Response{}, fmt.Errorf("claude executor: token counting failed: %w", err)
+	}
+
+	usageJSON := []byte(fmt.Sprintf(`{"input_tokens":%d}`, count))
+	out := sdktranslator.TranslateTokenCount(ctx, to, responseFormat, count, usageJSON)
+	return cliproxyexecutor.Response{Payload: out}, nil
+}
+
+type claudeTokenCountValidationError struct {
+	statusErr
+}
+
+func (claudeTokenCountValidationError) IsRequestScoped() bool {
+	return true
+}
+
+func newClaudeTokenCountValidationError(message string) error {
+	return claudeTokenCountValidationError{statusErr{code: http.StatusBadRequest, msg: message}}
+}
+
+func validateClaudeTokenCountRequest(body []byte) error {
+	if !gjson.ValidBytes(body) {
+		return newClaudeTokenCountValidationError("invalid Claude token count request JSON")
+	}
+	root := gjson.ParseBytes(body)
+	if !root.IsObject() {
+		return newClaudeTokenCountValidationError("Claude token count request must be a JSON object")
+	}
+	messages := root.Get("messages")
+	if !messages.IsArray() || len(messages.Array()) == 0 {
+		return newClaudeTokenCountValidationError("Claude token count request messages must be a non-empty array")
+	}
+	for _, message := range messages.Array() {
+		if !message.IsObject() {
+			return newClaudeTokenCountValidationError("Claude token count request messages must contain objects")
+		}
+		role := message.Get("role").String()
+		if role != "user" && role != "assistant" {
+			return newClaudeTokenCountValidationError("Claude token count request message role must be user or assistant")
+		}
+		content := message.Get("content")
+		if content.Type == gjson.String {
+			continue
+		}
+		if !content.IsArray() {
+			return newClaudeTokenCountValidationError("Claude token count request message content must be a string or array")
+		}
+		for _, block := range content.Array() {
+			if !block.IsObject() || block.Get("type").Type != gjson.String || block.Get("type").String() == "" {
+				return newClaudeTokenCountValidationError("Claude token count request content blocks must be typed objects")
+			}
+		}
+	}
+	return nil
+}
+
+func shouldUseClaudeUpstreamTokenCount(apiKey, baseURL string) bool {
+	return strings.TrimSpace(apiKey) != "" && isAnthropicUpstreamBase(baseURL)
+}
+
+// countTokensUpstream preserves Anthropic's native token-counting contract.
+func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
+	baseModel := thinking.ParseSuffix(req.Model).ModelName
+	upstreamModel := e.upstreamModel(baseModel)
+
+	apiKey, baseURL := claudeCreds(auth)
+	if baseURL == "" {
+		baseURL = "https://api.anthropic.com"
+	}
+	url := fmt.Sprintf("%s/v1/messages/count_tokens?beta=true", baseURL)
+	fp := resolveClaudeFingerprintPolicy(e.cfg, auth, apiKey)
+
+	from := opts.SourceFormat
+	responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts)
+	to := sdktranslator.FromString("claude")
+	originalPayload := req.Payload
+	if len(opts.OriginalRequest) > 0 {
+		originalPayload = opts.OriginalRequest
+	}
+	incomingHeaders, claudeCodeDetection := detectIncomingClaudeCodeRequest(ctx, opts.Headers, originalPayload, true, e.cfg)
+	confirmedClaudeCode := claudeCodeDetection.Confirmed
+	claudeSessionID := ""
+	if fp.ProfileClaudeCodeCLI {
+		claudeSessionID = helps.ClaudeAgentSessionUUIDForRequest(incomingHeaders, originalPayload, req.Payload, confirmedClaudeCode, opts.Metadata, req.Metadata)
+	}
+	// Use streaming translation to preserve function calling, except for claude.
+	stream := from != to
+	body := helps.TranslateRequestWithAPIKeyModelCompatibility(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, stream, helps.APIKeyModelIsCompat(req))
+	body = helps.SetStringIfDifferent(body, "model", upstreamModel)
+	var errThinking error
+	body, errThinking = helps.ApplyRequestThinking(body, req, opts, from.String(), to.String(), e.Identifier())
+	if errThinking != nil {
+		return cliproxyexecutor.Response{}, errThinking
+	}
+	if rebuildMidSystemMessageEnabled(e.cfg, auth) {
+		body = rebuildMidSystemMessagesToTopLevel(body)
+	}
+
+	directAnthropic := isAnthropicUpstreamBase(baseURL)
+	// Claude Code's count_tokens carries only model, messages and tools, so the
+	// full Messages cloaking must not run here for any origin. Apply the parts
+	// that still have to hold: relocate the caller's system prompt into messages
+	// so its tokens stay counted, and obfuscate sensitive words exactly like the
+	// Messages path. Kimi opt-in uses the same contract.
+	policy, settings := resolveClaudeWirePolicy(e.cfg, auth, apiKey, confirmedClaudeCode)
+	cloaked := policy.Cloak
+	if cloaked {
+		if !settings.strictMode {
+			if errSystem := validateClaudeCallerSystemBlocks(gjson.GetBytes(body, "system")); errSystem != nil {
+				return cliproxyexecutor.Response{}, errSystem
+			}
+		}
+		body = relocateClaudeSystemPromptForCountTokens(body, settings.strictMode)
+		if len(settings.sensitiveWords) > 0 {
+			body = helps.ObfuscateSensitiveWords(body, helps.BuildSensitiveWordMatcher(settings.sensitiveWords))
+		}
+	}
+
+	// Keep count_tokens requests compatible with Anthropic cache-control constraints too.
+	body = enforceCacheControlLimit(body, 4)
+	body = normalizeCacheControlTTL(body)
+
+	// Extract betas from body and convert to header (for count_tokens too)
+	var extraBetas []string
+	extraBetas, body = extractAndRemoveBetas(body)
+	// Claude Code 2.1.220's beta.messages.countTokens() always appends this beta.
+	extraBetas = append(extraBetas, claudeTokenCountingBeta)
+	if fp.MCPAlias && cloaked {
+		mcpAliases := resolveClaudeMCPAliasOptions(ctx)
+		body, _ = prepareClaudeOAuthToolNamesForUpstream(body, mcpAliases)
+	}
+	body = sanitizeClaudeMessagesForClaudeUpstreamWithDebug(ctx, body, baseModel, helps.APIKeyModelIsCompat(req))
+	// Two different reasons converge on the same deletions, and they must stay
+	// separable.
+	//
+	// api.anthropic.com rejects these fields on count_tokens outright ("metadata:
+	// Extra inputs are not permitted"), so they have to go for every credential
+	// that lands there, opted in or not. That is upstream compatibility, not
+	// fingerprinting.
+	//
+	// Elsewhere (Kimi, delegated Anthropic Messages providers) the caller owns its
+	// body by default: a caller that deliberately sends context_management expects
+	// the token count to reflect it, so CPA must not silently rewrite the request.
+	// Only an explicit claude-code-cli profile aligns the shape, and then it aligns
+	// to the measured one: Claude Code 2.1.220 count_tokens carries exactly model,
+	// messages and tools, never a system block.
+	alignCLICountTokensShape := fp.ProfileClaudeCodeCLI
+	if directAnthropic || alignCLICountTokensShape {
+		body, _ = sjson.DeleteBytes(body, "metadata")
+		body, _ = sjson.DeleteBytes(body, "context_management")
+		body, _ = sjson.DeleteBytes(body, "diagnostics")
+	}
+	if alignCLICountTokensShape {
+		body = util.StripClaudeCodeAttributionSystem(body)
+	}
+	// Runs on the finished body: payload rules can rewrite model and messages
+	// long after translation, so an earlier check would not describe the request
+	// that is about to be sent.
+	if errMidSystem := validateClaudeMidSystemMessageModel(body, confirmedClaudeCode, directAnthropic); errMidSystem != nil {
+		return cliproxyexecutor.Response{}, errMidSystem
+	}
+	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
+	if err != nil {
+		return cliproxyexecutor.Response{}, err
+	}
+	if errHeaders := applyClaudeHeaders(httpReq, auth, apiKey, false, extraBetas, body, e.cfg, incomingHeaders, confirmedClaudeCode && !cloaked, claudeSessionID); errHeaders != nil {
+		return cliproxyexecutor.Response{}, errHeaders
+	}
+	var authID, authLabel, authType, authValue string
+	if auth != nil {
+		authID = auth.ID
+		authLabel = auth.Label
+		authType, authValue = auth.AccountInfo()
+	}
+	helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{
+		URL:       url,
+		Method:    http.MethodPost,
+		Headers:   httpReq.Header.Clone(),
+		Body:      body,
+		Provider:  e.upstreamRequestLogProvider(),
+		AuthID:    authID,
+		AuthLabel: authLabel,
+		AuthType:  authType,
+		AuthValue: authValue,
+	})
+
+	httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0)
+	resp, err := doClaudeUpstreamRequest(httpClient, httpReq)
+	if err != nil {
+		helps.RecordAPIResponseError(ctx, e.cfg, err)
+		return cliproxyexecutor.Response{}, err
+	}
+	helps.RecordAPIResponseMetadata(ctx, e.cfg, resp.StatusCode, resp.Header.Clone())
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		// Decompress error responses — pass the Content-Encoding value (may be empty)
+		// and let decodeResponseBody handle both header-declared and magic-byte-detected
+		// compression.  This keeps error-path behaviour consistent with the success path.
+		errBody, decErr := decodeResponseBody(resp.Body, claudeResponseContentEncoding(resp.Header))
+		if decErr != nil {
+			helps.RecordAPIResponseError(ctx, e.cfg, decErr)
+			msg := fmt.Sprintf("failed to decode error response body: %v", decErr)
+			helps.LogWithRequestID(ctx).Warn(msg)
+			return cliproxyexecutor.Response{}, classifyClaudeUpstreamError(resp.StatusCode, resp.Header, []byte(msg))
+		}
+		b, readErr := io.ReadAll(errBody)
+		if readErr != nil {
+			helps.RecordAPIResponseError(ctx, e.cfg, readErr)
+			msg := fmt.Sprintf("failed to read error response body: %v", readErr)
+			helps.LogWithRequestID(ctx).Warn(msg)
+			b = []byte(msg)
+		}
+		helps.AppendAPIResponseChunk(ctx, e.cfg, b)
+		if errClose := errBody.Close(); errClose != nil {
+			log.Errorf("response body close error: %v", errClose)
+		}
+		return cliproxyexecutor.Response{}, classifyClaudeUpstreamError(resp.StatusCode, resp.Header, b)
+	}
+	decodedBody, err := decodeResponseBody(resp.Body, claudeResponseContentEncoding(resp.Header))
+	if err != nil {
+		helps.RecordAPIResponseError(ctx, e.cfg, err)
+		if errClose := resp.Body.Close(); errClose != nil {
+			log.Errorf("response body close error: %v", errClose)
+		}
+		return cliproxyexecutor.Response{}, err
+	}
+	defer func() {
+		if errClose := decodedBody.Close(); errClose != nil {
+			log.Errorf("response body close error: %v", errClose)
+		}
+	}()
+	data, err := io.ReadAll(decodedBody)
+	if err != nil {
+		helps.RecordAPIResponseError(ctx, e.cfg, err)
+		return cliproxyexecutor.Response{}, err
+	}
+	helps.AppendAPIResponseChunk(ctx, e.cfg, data)
+	count := gjson.GetBytes(data, "input_tokens").Int()
+	out := sdktranslator.TranslateTokenCount(ctx, to, responseFormat, count, data)
+	return cliproxyexecutor.Response{Payload: out, Headers: resp.Header.Clone()}, nil
+}