From f3e36f19c0882a76a46a565133a5c181f840ea4e Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Wed, 22 Jul 2026 05:21:14 +0800 Subject: [PATCH] feat(executor): add Claude input token state handling and token estimation logic - Introduced `ClaudeInputTokenState` to track and estimate input token usage for Claude message_start events. - Added `TranslateStreamWithClaudeInputTokens` for token-aware stream translation. - Updated executors (XAI, Kimi, Gemini) to leverage the new logic. - Included robust test cases for token counting, JSON validation, and concurrent tokenization scenarios. --- go.mod | 4 +- go.sum | 8 +- .../runtime/executor/aistudio_executor.go | 9 +- .../runtime/executor/antigravity_executor.go | 5 +- internal/runtime/executor/codex_executor.go | 3 +- .../executor/codex_websockets_executor.go | 3 +- internal/runtime/executor/gemini_executor.go | 12 +- .../executor/gemini_vertex_executor.go | 10 +- .../executor/helps/claude_input_tokens.go | 369 +++++++++++++++ .../helps/claude_input_tokens_test.go | 443 ++++++++++++++++++ internal/runtime/executor/kimi_executor.go | 5 +- .../executor/openai_compat_executor.go | 5 +- internal/runtime/executor/xai_executor.go | 3 +- .../executor/xai_websockets_executor.go | 6 +- 14 files changed, 858 insertions(+), 27 deletions(-) create mode 100644 internal/runtime/executor/helps/claude_input_tokens.go create mode 100644 internal/runtime/executor/helps/claude_input_tokens_test.go diff --git a/go.mod b/go.mod index 178184b8..7264ba4a 100644 --- a/go.mod +++ b/go.mod @@ -23,7 +23,7 @@ require ( github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 github.com/tidwall/gjson v1.18.0 github.com/tidwall/sjson v1.2.5 - github.com/tiktoken-go/tokenizer v0.7.0 + github.com/tiktoken-go/tokenizer v0.8.1 golang.org/x/crypto v0.54.0 golang.org/x/net v0.57.0 golang.org/x/oauth2 v0.30.0 @@ -35,6 +35,7 @@ require ( require ( github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/dlclark/regexp2/v2 v2.5.1 // indirect github.com/rogpeppe/go-internal v1.15.0 // indirect go.uber.org/atomic v1.11.0 // indirect ) @@ -56,7 +57,6 @@ require ( github.com/cloudflare/circl v1.6.3 // indirect github.com/cloudwego/base64x v0.1.4 // indirect github.com/cloudwego/iasm v0.2.0 // indirect - github.com/dlclark/regexp2 v1.11.5 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/emirpasic/gods v1.18.1 // indirect github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect diff --git a/go.sum b/go.sum index 352f166e..0f9a92e2 100644 --- a/go.sum +++ b/go.sum @@ -53,8 +53,8 @@ github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ= -github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dlclark/regexp2/v2 v2.5.1 h1:E5Ug7Dh264W1ymdySmiHNcDG7fmsR307APCE5R07a20= +github.com/dlclark/regexp2/v2 v2.5.1/go.mod h1:avUrQvPaLz2DrFNHJF0taWAFFX2C1GMSSoeiqFjcBmU= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc= @@ -197,8 +197,8 @@ github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs= github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= -github.com/tiktoken-go/tokenizer v0.7.0 h1:VMu6MPT0bXFDHr7UPh9uii7CNItVt3X9K90omxL54vw= -github.com/tiktoken-go/tokenizer v0.7.0/go.mod h1:6UCYI/DtOallbmL7sSy30p6YQv60qNyU/4aVigPOx6w= +github.com/tiktoken-go/tokenizer v0.8.1 h1:4obDoB6/dhdBt9xMweX4nww5cjdOq/nYF4ecwPq2+mg= +github.com/tiktoken-go/tokenizer v0.8.1/go.mod h1:eLA0t6nGvn9mDc7gt90qt7pMat+gE9ViqwQ6l9B+tA4= github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE= diff --git a/internal/runtime/executor/aistudio_executor.go b/internal/runtime/executor/aistudio_executor.go index ab588935..041baa02 100644 --- a/internal/runtime/executor/aistudio_executor.go +++ b/internal/runtime/executor/aistudio_executor.go @@ -291,6 +291,11 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth go func(first wsrelay.StreamEvent) { defer close(out) responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + originalRequest := opts.OriginalRequest + if len(originalRequest) == 0 { + originalRequest = req.Payload + } + claudeInputTokens := helps.NewClaudeInputTokenState(opts.SourceFormat, body.toFormat, responseFormat, originalRequest) var param any metadataLogged := false processEvent := func(event wsrelay.StreamEvent) bool { @@ -318,7 +323,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth if detail, ok := helps.ParseGeminiStreamUsage(filtered); ok { reporter.Publish(ctx, detail) } - lines := sdktranslator.TranslateStream(ctx, body.toFormat, responseFormat, req.Model, opts.OriginalRequest, translatedReq, filtered, ¶m) + lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, opts.OriginalRequest, translatedReq, filtered, ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}: @@ -340,7 +345,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth reporter.MarkFirstResponseByte() helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload) } - lines := sdktranslator.TranslateStream(ctx, body.toFormat, responseFormat, req.Model, opts.OriginalRequest, translatedReq, event.Payload, ¶m) + lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, opts.OriginalRequest, translatedReq, event.Payload, ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}: diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index 30b3b8d0..8fe024bb 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -1533,6 +1533,7 @@ attemptLoop: }() scanner := bufio.NewScanner(resp.Body) scanner.Buffer(nil, streamScannerBuffer) + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var param any for scanner.Scan() { line := scanner.Bytes() @@ -1555,7 +1556,7 @@ attemptLoop: } payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, payload) - chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bytes.Clone(payload), ¶m) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bytes.Clone(payload), ¶m, claudeInputTokens) for i := range chunks { select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: @@ -1564,7 +1565,7 @@ attemptLoop: } } } - tail := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, []byte("[DONE]"), ¶m) + tail := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, []byte("[DONE]"), ¶m, claudeInputTokens) for i := range tail { select { case out <- cliproxyexecutor.StreamChunk{Payload: tail[i]}: diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 56b01069..4847fbf0 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -1487,6 +1487,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au }() scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, 52_428_800) // 50MB + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var param any outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte @@ -1535,7 +1536,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au } translatedLine = applyCodexIdentityExposeResponsePayload(translatedLine, identityState) - chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, originalPayload, body, translatedLine, ¶m) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, body, translatedLine, ¶m, claudeInputTokens) for i := range chunks { select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index ae503fa7..266612cf 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -761,6 +761,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } } + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var param any outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte @@ -880,7 +881,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr eventType = gjson.GetBytes(payload, "type").String() clientPayload = applyCodexIdentityExposeResponsePayload(payload, identityState) line := encodeCodexWebsocketAsSSE(clientPayload) - chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, originalPayload, clientBody, line, ¶m) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, clientBody, line, ¶m, claudeInputTokens) for i := range chunks { if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i]}) { terminateReason = "context_done" diff --git a/internal/runtime/executor/gemini_executor.go b/internal/runtime/executor/gemini_executor.go index 0cd284dd..e4d76730 100644 --- a/internal/runtime/executor/gemini_executor.go +++ b/internal/runtime/executor/gemini_executor.go @@ -339,6 +339,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A }() scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, streamScannerBuffer) + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var param any for scanner.Scan() { line := scanner.Bytes() @@ -351,7 +352,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A if detail, ok := helps.ParseGeminiStreamUsage(payload); ok { reporter.Publish(ctx, detail) } - lines := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, bytes.Clone(payload), ¶m) + lines := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, bytes.Clone(payload), ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -360,7 +361,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A } } } - lines := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, []byte("[DONE]"), ¶m) + lines := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, []byte("[DONE]"), ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -530,6 +531,11 @@ func (e *GeminiExecutor) executeInteractionsStream(ctx context.Context, auth *cl }() scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, streamScannerBuffer) + originalRequest := opts.OriginalRequest + if len(originalRequest) == 0 { + originalRequest = req.Payload + } + claudeInputTokens := helps.NewClaudeInputTokenState(opts.SourceFormat, sdktranslator.FormatInteractions, responseFormat, originalRequest) var param any var frame []byte emitFrame := func() bool { @@ -564,7 +570,7 @@ func (e *GeminiExecutor) executeInteractionsStream(ctx context.Context, auth *cl return true } var lines [][]byte - lines = sdktranslator.TranslateStream(ctx, sdktranslator.FormatInteractions, responseFormat, req.Model, opts.OriginalRequest, body, payload, ¶m) + lines = helps.TranslateStreamWithClaudeInputTokens(ctx, sdktranslator.FormatInteractions, responseFormat, req.Model, opts.OriginalRequest, body, payload, ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: diff --git a/internal/runtime/executor/gemini_vertex_executor.go b/internal/runtime/executor/gemini_vertex_executor.go index f97eb84d..fa5f96be 100644 --- a/internal/runtime/executor/gemini_vertex_executor.go +++ b/internal/runtime/executor/gemini_vertex_executor.go @@ -661,6 +661,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte }() scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, streamScannerBuffer) + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var param any for scanner.Scan() { line := scanner.Bytes() @@ -668,7 +669,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte if detail, ok := helps.ParseGeminiStreamUsage(line); ok { reporter.Publish(ctx, detail) } - lines := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, bytes.Clone(line), ¶m) + lines := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, bytes.Clone(line), ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -677,7 +678,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte } } } - lines := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, []byte("[DONE]"), ¶m) + lines := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, []byte("[DONE]"), ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -806,6 +807,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth }() scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, streamScannerBuffer) + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var param any for scanner.Scan() { line := scanner.Bytes() @@ -813,7 +815,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth if detail, ok := helps.ParseGeminiStreamUsage(line); ok { reporter.Publish(ctx, detail) } - lines := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, bytes.Clone(line), ¶m) + lines := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, bytes.Clone(line), ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -822,7 +824,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth } } } - lines := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, []byte("[DONE]"), ¶m) + lines := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, []byte("[DONE]"), ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: diff --git a/internal/runtime/executor/helps/claude_input_tokens.go b/internal/runtime/executor/helps/claude_input_tokens.go new file mode 100644 index 00000000..b14429f4 --- /dev/null +++ b/internal/runtime/executor/helps/claude_input_tokens.go @@ -0,0 +1,369 @@ +package helps + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "strings" + "sync" + + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + "github.com/tiktoken-go/tokenizer" + + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" +) + +var ( + claudeInputTokenizerOnce sync.Once + claudeInputTokenizerCodec tokenizer.Codec + claudeInputTokenizerErr error +) + +// ClaudeInputTokenState tracks the one-time input token update for a translated Claude stream. +type ClaudeInputTokenState struct { + upstreamFormat sdktranslator.Format + responseFormat sdktranslator.Format + originalRequest []byte + codec tokenizer.Codec + handled bool +} + +// NewClaudeInputTokenState creates request-scoped state for translated Claude input token usage. +func NewClaudeInputTokenState(sourceFormat, upstreamFormat, responseFormat sdktranslator.Format, originalRequest []byte) *ClaudeInputTokenState { + enabled := sourceFormat == sdktranslator.FormatClaude && + upstreamFormat != sdktranslator.FormatClaude && + responseFormat == sdktranslator.FormatClaude + return &ClaudeInputTokenState{ + upstreamFormat: upstreamFormat, + responseFormat: responseFormat, + originalRequest: originalRequest, + handled: !enabled, + } +} + +// TranslateStreamWithClaudeInputTokens translates a stream chunk and estimates Claude message_start input usage once. +func TranslateStreamWithClaudeInputTokens( + ctx context.Context, + upstreamFormat, responseFormat sdktranslator.Format, + model string, + originalRequestRawJSON, requestRawJSON, rawJSON []byte, + param *any, + state *ClaudeInputTokenState, +) [][]byte { + chunks := sdktranslator.TranslateStream( + ctx, + upstreamFormat, + responseFormat, + model, + originalRequestRawJSON, + requestRawJSON, + rawJSON, + param, + ) + if state == nil { + return chunks + } + return state.apply(ctx, chunks) +} + +func claudeInputTokenizer() (tokenizer.Codec, error) { + claudeInputTokenizerOnce.Do(func() { + claudeInputTokenizerCodec, claudeInputTokenizerErr = tokenizer.Get(tokenizer.O200kBase) + }) + return claudeInputTokenizerCodec, claudeInputTokenizerErr +} + +func countClaudeInputTokens(enc tokenizer.Codec, payload []byte) (int64, error) { + if enc == nil { + return 0, fmt.Errorf("encoder is nil") + } + segments, err := collectClaudeInputTokenSegments(payload) + if err != nil { + return 0, err + } + if len(segments) == 0 { + return 0, nil + } + count, err := enc.Count(strings.Join(segments, "\n")) + if err != nil { + return 0, err + } + return int64(count), nil +} + +func collectClaudeInputTokenSegments(payload []byte) ([]string, error) { + if len(bytes.TrimSpace(payload)) == 0 { + return nil, nil + } + if !gjson.ValidBytes(payload) { + return nil, fmt.Errorf("invalid Claude request JSON") + } + + root := gjson.ParseBytes(payload) + segments := make([]string, 0, 32) + collectClaudeSystemTokenSegments(root.Get("system"), &segments) + collectClaudeMessageTokenSegments(root.Get("messages"), &segments) + collectClaudeToolTokenSegments(root.Get("tools"), &segments) + collectClaudeToolChoiceTokenSegments(root.Get("tool_choice"), &segments) + return segments, nil +} + +func collectClaudeSystemTokenSegments(system gjson.Result, segments *[]string) { + if system.Type == gjson.String { + appendClaudeTokenString(segments, system.String()) + return + } + if !system.IsArray() { + return + } + system.ForEach(func(_, part gjson.Result) bool { + if part.Type == gjson.String { + appendClaudeTokenString(segments, part.String()) + } else if part.Get("type").String() == "text" { + appendClaudeTokenString(segments, part.Get("text").String()) + } + return true + }) +} + +func collectClaudeMessageTokenSegments(messages gjson.Result, segments *[]string) { + if !messages.IsArray() { + return + } + messages.ForEach(func(_, message gjson.Result) bool { + appendClaudeTokenString(segments, message.Get("role").String()) + collectClaudeContentTokenSegments(message.Get("content"), segments) + return true + }) +} + +func collectClaudeContentTokenSegments(content gjson.Result, segments *[]string) { + if !content.Exists() { + return + } + if content.Type == gjson.String { + appendClaudeTokenString(segments, content.String()) + return + } + if content.IsArray() { + content.ForEach(func(_, part gjson.Result) bool { + collectClaudeContentTokenSegments(part, segments) + return true + }) + return + } + if !content.IsObject() { + return + } + + switch content.Get("type").String() { + case "text": + appendClaudeTokenString(segments, content.Get("text").String()) + case "thinking": + appendClaudeTokenString(segments, content.Get("thinking").String()) + case "document": + collectClaudeDocumentTokenSegments(content, segments) + case "tool_use", "server_tool_use", "mcp_tool_use": + appendClaudeTokenString(segments, content.Get("id").String()) + appendClaudeTokenString(segments, content.Get("name").String()) + appendClaudeTokenJSON(segments, content.Get("input")) + case "tool_result", "mcp_tool_result", "web_search_tool_result", "web_fetch_tool_result", "code_execution_tool_result", "bash_code_execution_tool_result", "text_editor_code_execution_tool_result": + appendClaudeTokenString(segments, content.Get("tool_use_id").String()) + appendClaudeTokenString(segments, content.Get("tool_call_id").String()) + collectClaudeContentTokenSegments(content.Get("content"), segments) + case "web_search_result", "search_result": + if source := content.Get("source"); source.Type == gjson.String { + appendClaudeTokenString(segments, source.String()) + } + appendClaudeTokenString(segments, content.Get("title").String()) + appendClaudeTokenString(segments, content.Get("url").String()) + appendClaudeTokenString(segments, content.Get("page_age").String()) + collectClaudeContentTokenSegments(content.Get("content"), segments) + case "web_fetch_result": + appendClaudeTokenString(segments, content.Get("url").String()) + appendClaudeTokenString(segments, content.Get("retrieved_at").String()) + collectClaudeContentTokenSegments(content.Get("content"), segments) + case "code_execution_result", "bash_code_execution_result", "text_editor_code_execution_result": + appendClaudeTokenString(segments, content.Get("stdout").String()) + appendClaudeTokenString(segments, content.Get("stderr").String()) + appendClaudeTokenString(segments, content.Get("return_code").String()) + collectClaudeContentTokenSegments(content.Get("content"), segments) + collectClaudeContentTokenSegments(content.Get("output"), segments) + case "tool_reference": + appendClaudeTokenString(segments, content.Get("tool_name").String()) + case "image", "input_audio", "audio", "video", "redacted_thinking": + return + case "": + appendClaudeTokenJSON(segments, content) + default: + appendClaudeTokenString(segments, content.Get("text").String()) + } +} + +func collectClaudeDocumentTokenSegments(document gjson.Result, segments *[]string) { + source := document.Get("source") + if source.Get("type").String() != "text" { + return + } + appendClaudeTokenString(segments, document.Get("title").String()) + appendClaudeTokenString(segments, document.Get("context").String()) + appendClaudeTokenString(segments, source.Get("data").String()) + appendClaudeTokenString(segments, source.Get("content").String()) +} + +func collectClaudeToolTokenSegments(tools gjson.Result, segments *[]string) { + if !tools.IsArray() { + return + } + tools.ForEach(func(_, tool gjson.Result) bool { + appendClaudeTokenString(segments, tool.Get("type").String()) + appendClaudeTokenString(segments, tool.Get("name").String()) + appendClaudeTokenString(segments, tool.Get("description").String()) + appendClaudeTokenJSON(segments, tool.Get("input_schema")) + return true + }) +} + +func collectClaudeToolChoiceTokenSegments(toolChoice gjson.Result, segments *[]string) { + if !toolChoice.Exists() { + return + } + if toolChoice.Type == gjson.String { + appendClaudeTokenString(segments, toolChoice.String()) + return + } + appendClaudeTokenString(segments, toolChoice.Get("type").String()) + appendClaudeTokenString(segments, toolChoice.Get("name").String()) +} + +func appendClaudeTokenString(segments *[]string, value string) { + if segments == nil { + return + } + if trimmed := strings.TrimSpace(value); trimmed != "" { + *segments = append(*segments, trimmed) + } +} + +func appendClaudeTokenJSON(segments *[]string, value gjson.Result) { + if !value.Exists() { + return + } + if value.Type == gjson.String { + appendClaudeTokenString(segments, value.String()) + return + } + raw := strings.TrimSpace(value.Raw) + if raw == "" { + return + } + var compact bytes.Buffer + if err := json.Compact(&compact, []byte(raw)); err == nil { + appendClaudeTokenString(segments, compact.String()) + return + } + appendClaudeTokenString(segments, raw) +} + +func (state *ClaudeInputTokenState) apply(ctx context.Context, chunks [][]byte) [][]byte { + if state == nil || state.handled { + return chunks + } + for i := range chunks { + updated, found := state.applyChunk(ctx, chunks[i]) + if !found { + continue + } + state.handled = true + chunks[i] = updated + break + } + return chunks +} + +func (state *ClaudeInputTokenState) applyChunk(ctx context.Context, chunk []byte) ([]byte, bool) { + for lineStart := 0; lineStart < len(chunk); { + lineEnd := bytes.IndexByte(chunk[lineStart:], '\n') + if lineEnd < 0 { + lineEnd = len(chunk) + } else { + lineEnd += lineStart + } + + contentEnd := lineEnd + if contentEnd > lineStart && chunk[contentEnd-1] == '\r' { + contentEnd-- + } + line := chunk[lineStart:contentEnd] + trimmedLeft := bytes.TrimLeft(line, " \t") + if bytes.HasPrefix(trimmedLeft, []byte("data:")) { + payloadOffset := len(line) - len(trimmedLeft) + len("data:") + for payloadOffset < len(line) && (line[payloadOffset] == ' ' || line[payloadOffset] == '\t') { + payloadOffset++ + } + payloadEnd := len(line) + for payloadEnd > payloadOffset && (line[payloadEnd-1] == ' ' || line[payloadEnd-1] == '\t') { + payloadEnd-- + } + payload := line[payloadOffset:payloadEnd] + if gjson.GetBytes(payload, "type").String() == "message_start" { + inputTokens := gjson.GetBytes(payload, "message.usage.input_tokens") + if inputTokens.Exists() && inputTokens.Int() != 0 { + return chunk, true + } + count, err := state.estimate() + if err != nil { + state.logEstimateError(ctx, err) + return chunk, true + } + if count == 0 { + return chunk, true + } + updatedPayload, errSet := sjson.SetBytes(payload, "message.usage.input_tokens", count) + if errSet != nil { + state.logEstimateError(ctx, fmt.Errorf("set message_start usage: %w", errSet)) + return chunk, true + } + payloadStart := lineStart + payloadOffset + payloadStop := lineStart + payloadEnd + updated := make([]byte, 0, len(chunk)+len(updatedPayload)-len(payload)) + updated = append(updated, chunk[:payloadStart]...) + updated = append(updated, updatedPayload...) + updated = append(updated, chunk[payloadStop:]...) + return updated, true + } + } + + if lineEnd == len(chunk) { + break + } + lineStart = lineEnd + 1 + } + return chunk, false +} + +func (state *ClaudeInputTokenState) estimate() (int64, error) { + enc := state.codec + if enc == nil { + var err error + enc, err = claudeInputTokenizer() + if err != nil { + return 0, fmt.Errorf("initialize O200kBase tokenizer: %w", err) + } + } + count, err := countClaudeInputTokens(enc, state.originalRequest) + if err != nil { + return 0, fmt.Errorf("count Claude input tokens: %w", err) + } + return count, nil +} + +func (state *ClaudeInputTokenState) logEstimateError(ctx context.Context, err error) { + LogWithRequestID(ctx).WithFields(log.Fields{ + "upstream_format": state.upstreamFormat.String(), + "response_format": state.responseFormat.String(), + }).WithError(err).Warn("failed to estimate Claude input tokens") +} diff --git a/internal/runtime/executor/helps/claude_input_tokens_test.go b/internal/runtime/executor/helps/claude_input_tokens_test.go new file mode 100644 index 00000000..dadea4b9 --- /dev/null +++ b/internal/runtime/executor/helps/claude_input_tokens_test.go @@ -0,0 +1,443 @@ +package helps + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "strings" + "sync" + "testing" + + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tiktoken-go/tokenizer" + + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" +) + +type failingClaudeInputCodec struct{} + +func (failingClaudeInputCodec) GetName() string { + return "failing" +} + +func (failingClaudeInputCodec) Count(string) (int, error) { + return 0, errors.New("count failed") +} + +func (failingClaudeInputCodec) Encode(string) ([]uint, []string, error) { + return nil, nil, errors.New("encode failed") +} + +func (failingClaudeInputCodec) Decode([]uint) (string, error) { + return "", errors.New("decode failed") +} + +func TestCollectClaudeInputTokenSegments(t *testing.T) { + payload := []byte(`{ + "model":"claude-test", + "system":[ + {"type":"text","text":"Follow repository rules.","cache_control":{"type":"ephemeral"}}, + {"type":"image","source":{"type":"base64","media_type":"image/png","data":"ignored-system-image"}} + ], + "messages":[ + {"role":"user","content":[ + {"type":"text","text":"Review the implementation."}, + {"type":"document","source":{"type":"text","data":"Reference document text."}}, + {"type":"image","source":{"type":"base64","media_type":"image/png","data":"ignored-image"}} + ]}, + {"role":"assistant","content":[ + {"type":"thinking","thinking":"Inspect the relevant files.","signature":"ignored-signature"}, + {"type":"tool_use","id":"toolu_1","name":"read_file","input":{"path":"main.go"}} + ]}, + {"role":"user","content":[ + {"type":"tool_result","tool_use_id":"toolu_1","content":[ + {"type":"text","text":"package main"}, + {"type":"image","source":{"type":"base64","data":"ignored-tool-image"}} + ]} + ]} + ], + "tools":[{ + "name":"read_file", + "description":"Reads a repository file.", + "input_schema":{"type":"object","properties":{"path":{"type":"string"}}}, + "cache_control":{"type":"ephemeral"} + }], + "tool_choice":{"type":"tool","name":"read_file"}, + "metadata":{"user_id":"ignored-metadata"}, + "max_tokens":4096, + "stream":true + }`) + + got, err := collectClaudeInputTokenSegments(payload) + if err != nil { + t.Fatalf("collectClaudeInputTokenSegments() error = %v", err) + } + want := []string{ + "Follow repository rules.", + "user", + "Review the implementation.", + "Reference document text.", + "assistant", + "Inspect the relevant files.", + "toolu_1", + "read_file", + `{"path":"main.go"}`, + "user", + "toolu_1", + "package main", + "read_file", + "Reads a repository file.", + `{"type":"object","properties":{"path":{"type":"string"}}}`, + "tool", + "read_file", + } + if fmt.Sprint(got) != fmt.Sprint(want) { + t.Fatalf("segments = %#v, want %#v", got, want) + } +} + +func TestCollectClaudeInputTokenSegmentsIncludesKnownToolResults(t *testing.T) { + payload := []byte(`{ + "messages":[{"role":"user","content":[ + {"type":"web_search_tool_result","tool_use_id":"ws_tool_1","content":[ + {"type":"web_search_result","source":"Search source","title":"Search result title","url":"https://search.example/result","page_age":"1 day","encrypted_content":"ignored-secret"} + ]}, + {"type":"web_fetch_tool_result","tool_use_id":"fetch_tool_1","content":{ + "type":"web_fetch_result","url":"https://docs.example/page","retrieved_at":"2026-07-22T00:00:00Z","content":{ + "type":"document","title":"Fetched document","source":{"type":"text","data":"Fetched body"} + } + }}, + {"type":"bash_code_execution_tool_result","tool_use_id":"bash_tool_1","content":{ + "type":"bash_code_execution_result","stdout":"command output","stderr":"command error","return_code":1, + "content":[{"type":"text","text":"additional output"}] + }}, + {"type":"tool_result","tool_use_id":"toolu_1","content":[ + {"type":"tool_reference","tool_name":"proxy_mcp__nia__manage_resource"} + ]} + ]}] + }`) + + segments, err := collectClaudeInputTokenSegments(payload) + if err != nil { + t.Fatalf("collectClaudeInputTokenSegments() error = %v", err) + } + joined := "\n" + strings.Join(segments, "\n") + "\n" + for _, want := range []string{ + "ws_tool_1", + "Search source", + "Search result title", + "https://search.example/result", + "1 day", + "fetch_tool_1", + "https://docs.example/page", + "2026-07-22T00:00:00Z", + "Fetched document", + "Fetched body", + "bash_tool_1", + "command output", + "command error", + "1", + "additional output", + "toolu_1", + "proxy_mcp__nia__manage_resource", + } { + if !strings.Contains(joined, "\n"+want+"\n") { + t.Errorf("segments do not contain %q: %#v", want, segments) + } + } + if strings.Contains(joined, "ignored-secret") { + t.Fatalf("segments contain encrypted content: %#v", segments) + } +} + +func TestCountClaudeInputTokensExcludesMultimediaAndControlFields(t *testing.T) { + enc, err := tokenizer.Get(tokenizer.O200kBase) + if err != nil { + t.Fatalf("tokenizer.Get() error = %v", err) + } + + base := []byte(`{ + "system":"System text.", + "messages":[{"role":"user","content":[{"type":"text","text":"User text."}]}], + "tools":[{"name":"lookup","description":"Looks up data.","input_schema":{"type":"object"}}] + }`) + withExcludedFields := []byte(`{ + "model":"claude-test", + "system":"System text.", + "messages":[{"role":"user","content":[ + {"type":"text","text":"User text."}, + {"type":"image","source":{"type":"base64","media_type":"image/png","data":"very-large-image-data"}}, + {"type":"input_audio","source":{"type":"base64","data":"very-large-audio-data"}}, + {"type":"video","source":{"type":"url","url":"https://example.com/video.mp4"}}, + {"type":"document","source":{"type":"base64","media_type":"application/pdf","data":"very-large-pdf-data"}} + ]}], + "tools":[{"name":"lookup","description":"Looks up data.","input_schema":{"type":"object"},"cache_control":{"type":"ephemeral"}}], + "metadata":{"large_wrapper":"ignored"}, + "max_tokens":8192, + "temperature":0.8, + "top_p":0.9, + "thinking":{"type":"enabled","budget_tokens":4096}, + "stream":true + }`) + + baseCount, errBase := countClaudeInputTokens(enc, base) + if errBase != nil { + t.Fatalf("countClaudeInputTokens(base) error = %v", errBase) + } + excludedCount, errExcluded := countClaudeInputTokens(enc, withExcludedFields) + if errExcluded != nil { + t.Fatalf("countClaudeInputTokens(withExcludedFields) error = %v", errExcluded) + } + if excludedCount != baseCount { + t.Fatalf("count with excluded fields = %d, want %d", excludedCount, baseCount) + } +} + +func TestTranslateStreamWithClaudeInputTokensPatchesMessageStartOnce(t *testing.T) { + upstreamFormat := sdktranslator.Format("claude-input-token-test-upstream") + sdktranslator.Register(sdktranslator.FormatClaude, upstreamFormat, nil, sdktranslator.ResponseTransform{ + Stream: func(_ context.Context, _ string, _, _, rawJSON []byte, _ *any) [][]byte { + return [][]byte{rawJSON} + }, + }) + + originalRequest := []byte(`{"system":"System text.","messages":[{"role":"user","content":"Hello."}]}`) + state := NewClaudeInputTokenState(sdktranslator.FormatClaude, upstreamFormat, sdktranslator.FormatClaude, originalRequest) + var param any + combined := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":0,\"output_tokens\":0}}}\n\n" + + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0}\n\n") + + got := TranslateStreamWithClaudeInputTokens( + context.Background(), + upstreamFormat, + sdktranslator.FormatClaude, + "claude-test", + originalRequest, + nil, + combined, + ¶m, + state, + ) + if tokens := messageStartInputTokens(got); tokens <= 0 { + t.Fatalf("message_start input_tokens = %d, want positive estimate; output = %q", tokens, joinClaudeInputChunks(got)) + } + if !state.handled { + t.Fatal("state.handled = false, want true after message_start") + } + if !strings.Contains(joinClaudeInputChunks(got), `"type":"content_block_start"`) { + t.Fatalf("combined non-target event was not preserved: %q", joinClaudeInputChunks(got)) + } + + secondStart := []byte(`event: message_start +data: {"type":"message_start","message":{"usage":{"input_tokens":0}}} + +`) + gotSecond := TranslateStreamWithClaudeInputTokens( + context.Background(), + upstreamFormat, + sdktranslator.FormatClaude, + "claude-test", + originalRequest, + nil, + secondStart, + ¶m, + state, + ) + if tokens := messageStartInputTokens(gotSecond); tokens != 0 { + t.Fatalf("second message_start input_tokens = %d, want 0 after state handled", tokens) + } +} + +func TestClaudeInputTokenStatePreservesCRLFAndNonTargetEvents(t *testing.T) { + originalRequest := []byte(`{"messages":[{"role":"user","content":"Hello."}]}`) + state := NewClaudeInputTokenState(sdktranslator.FormatClaude, sdktranslator.FormatOpenAI, sdktranslator.FormatClaude, originalRequest) + chunk := []byte("event: message_start\r\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":0,\"output_tokens\":0}}} \r\n\r\n" + + "event: ping\r\ndata: {\"type\":\"ping\",\"value\":\"keep\"}\r\n\r\n") + + got := state.apply(context.Background(), [][]byte{chunk}) + tokens := messageStartInputTokens(got) + if tokens <= 0 { + t.Fatalf("input_tokens = %d, want positive estimate", tokens) + } + want := fmt.Sprintf("event: message_start\r\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":%d,\"output_tokens\":0}}} \r\n\r\n"+ + "event: ping\r\ndata: {\"type\":\"ping\",\"value\":\"keep\"}\r\n\r\n", tokens) + if joined := joinClaudeInputChunks(got); joined != want { + t.Fatalf("output bytes changed unexpectedly:\n got: %q\nwant: %q", joined, want) + } +} + +func TestClaudeInputTokenStatePatchesMissingAndPreservesNonZero(t *testing.T) { + originalRequest := []byte(`{"messages":[{"role":"user","content":"Hello."}]}`) + + t.Run("missing", func(t *testing.T) { + state := NewClaudeInputTokenState(sdktranslator.FormatClaude, sdktranslator.FormatOpenAI, sdktranslator.FormatClaude, originalRequest) + chunks := [][]byte{[]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"output_tokens\":0}}}\n\n")} + got := state.apply(context.Background(), chunks) + if tokens := messageStartInputTokens(got); tokens <= 0 { + t.Fatalf("input_tokens = %d, want positive estimate", tokens) + } + }) + + t.Run("non-zero", func(t *testing.T) { + state := NewClaudeInputTokenState(sdktranslator.FormatClaude, sdktranslator.FormatOpenAI, sdktranslator.FormatClaude, []byte(`not valid json`)) + chunks := [][]byte{[]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":73}}}\n\n")} + got := state.apply(context.Background(), chunks) + if tokens := messageStartInputTokens(got); tokens != 73 { + t.Fatalf("input_tokens = %d, want preserved value 73", tokens) + } + if !state.handled { + t.Fatal("state.handled = false, want true") + } + }) +} + +func TestClaudeInputTokenStateSkipsUnsupportedFlows(t *testing.T) { + originalRequest := []byte(`{"messages":[{"role":"user","content":"Hello."}]}`) + testCases := []struct { + name string + sourceFormat sdktranslator.Format + upstreamFormat sdktranslator.Format + responseFormat sdktranslator.Format + }{ + {name: "non-Claude source", sourceFormat: sdktranslator.FormatOpenAI, upstreamFormat: sdktranslator.FormatGemini, responseFormat: sdktranslator.FormatClaude}, + {name: "Claude passthrough", sourceFormat: sdktranslator.FormatClaude, upstreamFormat: sdktranslator.FormatClaude, responseFormat: sdktranslator.FormatClaude}, + {name: "non-Claude response", sourceFormat: sdktranslator.FormatClaude, upstreamFormat: sdktranslator.FormatOpenAI, responseFormat: sdktranslator.FormatOpenAI}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + state := NewClaudeInputTokenState(tc.sourceFormat, tc.upstreamFormat, tc.responseFormat, originalRequest) + chunks := [][]byte{[]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":0}}}\n\n")} + got := state.apply(context.Background(), chunks) + if tokens := messageStartInputTokens(got); tokens != 0 { + t.Fatalf("input_tokens = %d, want unchanged 0", tokens) + } + if !state.handled { + t.Fatal("state.handled = false, want disabled flow handled at initialization") + } + }) + } +} + +func TestClaudeInputTokenStateCountErrorKeepsZero(t *testing.T) { + originalLogOutput := log.StandardLogger().Out + log.SetOutput(io.Discard) + defer log.SetOutput(originalLogOutput) + + state := NewClaudeInputTokenState( + sdktranslator.FormatClaude, + sdktranslator.FormatOpenAI, + sdktranslator.FormatClaude, + []byte(`{"messages":[{"role":"user","content":"Hello."}]}`), + ) + state.codec = failingClaudeInputCodec{} + chunks := [][]byte{[]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":0}}}\n\n")} + + got := state.apply(context.Background(), chunks) + if tokens := messageStartInputTokens(got); tokens != 0 { + t.Fatalf("input_tokens = %d, want fallback 0", tokens) + } + if !state.handled { + t.Fatal("state.handled = false, want true after failed estimate") + } +} + +func TestClaudeInputTokenStateInvalidJSONKeepsZeroWithoutLoggingRequest(t *testing.T) { + originalLogOutput := log.StandardLogger().Out + var logOutput bytes.Buffer + log.SetOutput(&logOutput) + defer log.SetOutput(originalLogOutput) + + const sensitiveRequest = `{"messages":["sensitive-original-request"` + state := NewClaudeInputTokenState( + sdktranslator.FormatClaude, + sdktranslator.FormatOpenAI, + sdktranslator.FormatClaude, + []byte(sensitiveRequest), + ) + chunks := [][]byte{[]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":0}}}\n\n")} + + got := state.apply(context.Background(), chunks) + if tokens := messageStartInputTokens(got); tokens != 0 { + t.Fatalf("input_tokens = %d, want fallback 0", tokens) + } + if !state.handled { + t.Fatal("state.handled = false, want true after invalid JSON") + } + if !strings.Contains(logOutput.String(), "failed to estimate Claude input tokens") { + t.Fatalf("warning not logged: %q", logOutput.String()) + } + if strings.Contains(logOutput.String(), "sensitive-original-request") { + t.Fatalf("warning leaked original request: %q", logOutput.String()) + } +} + +func TestClaudeInputTokenizerConcurrentCount(t *testing.T) { + first, errFirst := claudeInputTokenizer() + if errFirst != nil { + t.Fatalf("claudeInputTokenizer() error = %v", errFirst) + } + second, errSecond := claudeInputTokenizer() + if errSecond != nil { + t.Fatalf("claudeInputTokenizer() second error = %v", errSecond) + } + if first != second { + t.Fatal("claudeInputTokenizer() returned different codec instances") + } + + const workers = 32 + const iterations = 50 + var wg sync.WaitGroup + errs := make(chan error, workers) + for worker := 0; worker < workers; worker++ { + worker := worker + wg.Add(1) + go func() { + defer wg.Done() + for iteration := 0; iteration < iterations; iteration++ { + payload := []byte(fmt.Sprintf(`{"messages":[{"role":"user","content":"worker %d iteration %d 你好"}]}`, worker, iteration)) + count, errCount := countClaudeInputTokens(first, payload) + if errCount != nil { + errs <- errCount + return + } + if count <= 0 { + errs <- fmt.Errorf("non-positive count: %d", count) + return + } + } + }() + } + wg.Wait() + close(errs) + for err := range errs { + t.Error(err) + } +} + +func messageStartInputTokens(chunks [][]byte) int64 { + for _, chunk := range chunks { + for _, line := range strings.Split(string(chunk), "\n") { + trimmed := strings.TrimSpace(line) + if !strings.HasPrefix(trimmed, "data:") { + continue + } + payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + if gjson.Get(payload, "type").String() == "message_start" { + return gjson.Get(payload, "message.usage.input_tokens").Int() + } + } + } + return 0 +} + +func joinClaudeInputChunks(chunks [][]byte) string { + var builder strings.Builder + for _, chunk := range chunks { + builder.Write(chunk) + } + return builder.String() +} diff --git a/internal/runtime/executor/kimi_executor.go b/internal/runtime/executor/kimi_executor.go index 7f67d3d2..76434f54 100644 --- a/internal/runtime/executor/kimi_executor.go +++ b/internal/runtime/executor/kimi_executor.go @@ -297,6 +297,7 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut }() scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, 1_048_576) // 1MB + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var param any var streamUsage helps.StreamUsageBuffer defer streamUsage.Publish(ctx, reporter) @@ -304,7 +305,7 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut line := scanner.Bytes() helps.AppendAPIResponseChunk(ctx, e.cfg, line) streamUsage.ObserveOpenAIStream(line) - chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, bytes.Clone(line), ¶m) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, bytes.Clone(line), ¶m, claudeInputTokens) for i := range chunks { select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: @@ -313,7 +314,7 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut } } } - doneChunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, []byte("[DONE]"), ¶m) + doneChunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, []byte("[DONE]"), ¶m, claudeInputTokens) for i := range doneChunks { select { case out <- cliproxyexecutor.StreamChunk{Payload: doneChunks[i]}: diff --git a/internal/runtime/executor/openai_compat_executor.go b/internal/runtime/executor/openai_compat_executor.go index d18ab967..2288100e 100644 --- a/internal/runtime/executor/openai_compat_executor.go +++ b/internal/runtime/executor/openai_compat_executor.go @@ -392,6 +392,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy }() scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, 52_428_800) // 50MB + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var param any var streamUsage helps.StreamUsageBuffer defer streamUsage.Publish(ctx, reporter) @@ -423,7 +424,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy } // OpenAI-compatible streams must use SSE data lines. - chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bytes.Clone(trimmedLine), ¶m) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bytes.Clone(trimmedLine), ¶m, claudeInputTokens) for i := range chunks { select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: @@ -443,7 +444,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy // In case the upstream close the stream without a terminal [DONE] marker. // Feed a synthetic done marker through the translator so pending // response.completed events are still emitted exactly once. - chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, []byte("data: [DONE]"), ¶m) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, []byte("data: [DONE]"), ¶m, claudeInputTokens) for i := range chunks { select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 70b092e8..2cbb4184 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -669,13 +669,14 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth }() scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, 52_428_800) + claudeInputTokens := helps.NewClaudeInputTokenState(prepared.from, prepared.to, prepared.responseFormat, prepared.originalPayload) var param any outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools) var pendingEventLine []byte emitTranslatedLine := func(translatedLine []byte) bool { - chunks := sdktranslator.TranslateStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, translatedLine, ¶m) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, translatedLine, ¶m, claudeInputTokens) for i := range chunks { select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index 1fe240a1..967172e3 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -21,7 +21,6 @@ import ( "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" @@ -586,6 +585,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } } + claudeInputTokens := helps.NewClaudeInputTokenState(prepared.from, prepared.to, prepared.responseFormat, prepared.originalPayload) var param any outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte @@ -719,7 +719,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox payload = normalizeCodexWebsocketCompletion(payload) line := encodeCodexWebsocketAsSSE(payload) - chunks := sdktranslator.TranslateStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, line, ¶m) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, line, ¶m, claudeInputTokens) for i := range chunks { if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i]}) { terminateReason = "context_done" @@ -729,7 +729,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } if len(warmupCompletedPayload) > 0 { line = encodeCodexWebsocketAsSSE(warmupCompletedPayload) - chunks = sdktranslator.TranslateStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, line, ¶m) + chunks = helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, line, ¶m, claudeInputTokens) for i := range chunks { if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i]}) { terminateReason = "context_done"