From bbd97446c26092f2e7250af429096064b9e0f899 Mon Sep 17 00:00:00 2001 From: Calcium-Ion Date: Thu, 3 Sep 2026 10:40:05 +0800 Subject: [PATCH] fix(relay): follow-up billing integrity and conversion completions (#7170) Deferred follow-ups from the relaykit-tools review cycle, verified by live end-to-end billing tests: - billing: normalize Gemini modality keys consistently between stream merge and settlement (case/whitespace variants no longer drop independent audio/image pricing) and sum duplicate modality entries on both paths - billing: sync legacy flat Claude cache-creation fields from the CacheCreation sub-object (including zeroing) and fall back to flat fields only when the snapshot never carried a sub-object, closing a stale 1h-cache overcharge path in cascaded deployments - relay: move Chat-to-Claude and Chat-to-Gemini stream conversion state from gin.Context onto RelayInfo and reset it with SendResponseCount in InitChannelMeta, so channel retries start clean while per-request state (stream error collection, conversion diagnostics, channel chain, billing accumulators) survives - relay: Claude channel now serves Gemini-format clients (request via registry conversion, response and stream composed through the Chat pivot), removing the last unimplemented conversion direction - relaykit: recognize legacy pseudo tool names (googleSearch, codeExecution, urlContext) in the toolconv decode stage and drop the string-matching bypass in the Chat-to-Gemini converter; native Gemini tool output is restored and non-Gemini targets follow standard loss diagnostics - relaykit: attach upstream Gemini usage (with billing_usage sidecar) to intermediate stream chunks so converted Claude streams report upstream truth from message_start, and preserve the sidecar through Claude stream usage merges; billing settlement unchanged - billing: clamp negative Total-Prompt completion derivation, OR the Estimated flag across cross-dialect snapshot replacement, and fill canonical OpenAI prompt details via field-wise merge --- relay/channel/claude/adaptor.go | 12 +- relay/channel/claude/adaptor_test.go | 97 +++++++++ relay/channel/claude/relay-claude.go | 99 +++++++++- relay/channel/gemini/relay-gemini.go | 3 + .../channel/gemini/relay_gemini_usage_test.go | 125 ++++++++++++ relay/channel/openai/helper.go | 18 +- relay/common/relay_info.go | 16 +- relay/common/relay_info_test.go | 93 +++++++++ relaykit/dto/billing_usage.go | 45 ++--- relaykit/dto/billing_usage_test.go | 32 +++ relaykit/dto/usage_merge.go | 38 +++- relaykit/dto/usage_merge_test.go | 169 +++++++++++++++- .../internal/gemini_chat/to_oai_chat_resp.go | 5 + .../oai_chat/to_claude_messages_resp.go | 51 +++-- .../oai_chat/to_claude_messages_resp_test.go | 186 ++++++++++++++++++ .../internal/oai_chat/to_gemini_chat_req.go | 30 --- .../relayconvert/internal/toolconv/decode.go | 38 ++++ .../relayconvert/internal/toolconv/encode.go | 4 +- 18 files changed, 958 insertions(+), 103 deletions(-) diff --git a/relay/channel/claude/adaptor.go b/relay/channel/claude/adaptor.go index 8d3583f09a..2dc7d2a1bf 100644 --- a/relay/channel/claude/adaptor.go +++ b/relay/channel/claude/adaptor.go @@ -21,9 +21,15 @@ import ( type Adaptor struct { } -func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) { - //TODO implement me - return nil, errors.New("not implemented") +func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + result, err := service.ConvertRequest(c, info, types.RelayFormatClaude, request) + if err != nil { + return nil, err + } + return result.Value, nil } func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) { diff --git a/relay/channel/claude/adaptor_test.go b/relay/channel/claude/adaptor_test.go index 01c035638e..885ea19a5d 100644 --- a/relay/channel/claude/adaptor_test.go +++ b/relay/channel/claude/adaptor_test.go @@ -88,3 +88,100 @@ func TestConvertClaudeRequestDoesNotOverwriteTrimmedUpstreamModelName(t *testing require.NoError(t, err) assert.Equal(t, "claude-3-7-sonnet", info.UpstreamModelName) } + +func geminiToClaudeInfo() *relaycommon.RelayInfo { + return &relaycommon.RelayInfo{ + OriginModelName: "claude-3-7-sonnet", + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "claude-3-7-sonnet", + }, + } +} + +func TestConvertGeminiRequestMapsSystemInstructionToolsAndMultimodal(t *testing.T) { + req := &dto.GeminiChatRequest{ + Contents: []dto.GeminiChatContent{ + { + Role: "user", + Parts: []dto.GeminiPart{ + {Text: "What is in this image?"}, + {InlineData: &dto.GeminiInlineData{MimeType: "image/png", Data: "aGVsbG8="}}, + }, + }, + }, + SystemInstructions: &dto.GeminiChatContent{ + Parts: []dto.GeminiPart{{Text: "You are a helpful assistant."}}, + }, + } + req.SetTools([]dto.GeminiChatTool{ + { + FunctionDeclarations: []dto.FunctionRequest{ + { + Name: "lookup", + Description: "Lookup data", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{"q": map[string]any{"type": "string"}}, + }, + }, + }, + }, + }) + + out, err := (&Adaptor{}).ConvertGeminiRequest(nil, geminiToClaudeInfo(), req) + require.NoError(t, err) + converted, ok := out.(*dto.ClaudeRequest) + require.True(t, ok) + + system := converted.ParseSystem() + require.NotEmpty(t, system) + assert.Contains(t, system[0].GetText(), "You are a helpful assistant.") + require.NotEmpty(t, converted.Messages) + assert.Equal(t, "user", converted.Messages[0].Role) + + blocks, parseErr := converted.Messages[0].ParseContent() + require.NoError(t, parseErr) + var foundImage bool + for _, block := range blocks { + if block.Type == "image" || (block.Source != nil && block.Source.Type == "base64") { + foundImage = true + break + } + } + assert.True(t, foundImage) + + require.NotNil(t, converted.Tools) + tools, err := common.Marshal(converted.Tools) + require.NoError(t, err) + assert.Contains(t, string(tools), `"lookup"`) + require.NotNil(t, converted.MaxTokens) + assert.Greater(t, *converted.MaxTokens, uint(0)) +} + +func TestConvertGeminiRequestThinkingConfigUsesReasoningIntent(t *testing.T) { + budget := 1024 + maxTokens := uint(4096) + req := &dto.GeminiChatRequest{ + Contents: []dto.GeminiChatContent{ + {Role: "user", Parts: []dto.GeminiPart{{Text: "think"}}}, + }, + GenerationConfig: dto.GeminiChatGenerationConfig{ + MaxOutputTokens: &maxTokens, + ThinkingConfig: &dto.GeminiThinkingConfig{ThinkingBudget: &budget}, + }, + } + + out, err := (&Adaptor{}).ConvertGeminiRequest(nil, geminiToClaudeInfo(), req) + require.NoError(t, err) + converted, ok := out.(*dto.ClaudeRequest) + require.True(t, ok) + require.NotNil(t, converted.Thinking) + assert.Equal(t, "enabled", converted.Thinking.Type) + require.NotNil(t, converted.Thinking.BudgetTokens) + assert.Equal(t, 1024, *converted.Thinking.BudgetTokens) +} + +func TestConvertGeminiRequestNilRequest(t *testing.T) { + _, err := (&Adaptor{}).ConvertGeminiRequest(nil, geminiToClaudeInfo(), nil) + require.Error(t, err) +} diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index 511c7eb79a..ac1456dfd2 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -20,8 +20,6 @@ import ( "github.com/gin-gonic/gin" ) -const claudeToChatStreamStateKey = "relaykit.claude_to_chat_stream_state" - func stopReasonClaude2OpenAI(reason string) string { return relayconvert.StopReasonClaudeToOpenAI(reason) } @@ -120,7 +118,7 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud countClaudeStreamBillableTools(c, info, &claudeResponse) helper.ClaudeChunkData(c, claudeResponse, data) } else if info.RelayFormat == types.RelayFormatOpenAI { - state, err := claudeToChatStreamState(c) + state, err := claudeToChatStreamState(info) if err != nil { return types.NewError(err, types.ErrorCodeBadResponseBody) } @@ -142,24 +140,80 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud if err != nil { logger.LogError(c, "send_stream_response_failed: "+err.Error()) } + } else if info.RelayFormat == types.RelayFormatGemini { + state, err := claudeToGeminiStreamState(info) + if err != nil { + return types.NewError(err, types.ErrorCodeBadResponseBody) + } + results, err := service.ConvertStreamResponseChunk(c, info, state, &claudeResponse) + if err != nil { + return types.NewError(err, types.ErrorCodeBadResponseBody) + } + if !FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo) { + return nil + } + countClaudeStreamBillableTools(c, info, &claudeResponse) + if sendErr := sendGeminiStreamResults(c, results); sendErr != nil { + return sendErr + } } return nil } -func claudeToChatStreamState(c *gin.Context) (*relayconvert.ClaudeToChatStreamState, error) { - if value, ok := c.Get(claudeToChatStreamStateKey); ok { - state, ok := value.(*relayconvert.ClaudeToChatStreamState) +func claudeToChatStreamState(info *relaycommon.RelayInfo) (*relayconvert.ClaudeToChatStreamState, error) { + if info != nil && info.ClaudeToChatStreamState != nil { + state, ok := info.ClaudeToChatStreamState.(*relayconvert.ClaudeToChatStreamState) if !ok || state == nil { - return nil, fmt.Errorf("invalid Claude-to-Chat stream state %T", value) + return nil, fmt.Errorf("invalid Claude-to-Chat stream state %T", info.ClaudeToChatStreamState) } return state, nil } state := relayconvert.NewClaudeToChatStreamState() - c.Set(claudeToChatStreamStateKey, state) + if info != nil { + info.ClaudeToChatStreamState = state + } return state, nil } +func claudeToGeminiStreamState(info *relaycommon.RelayInfo) (*relayconvert.ResponseStreamState, error) { + if info != nil && info.ChatToGeminiStreamState != nil { + state, ok := info.ChatToGeminiStreamState.(*relayconvert.ResponseStreamState) + if !ok || state == nil { + return nil, fmt.Errorf("invalid Claude-to-Gemini stream state %T", info.ChatToGeminiStreamState) + } + return state, nil + } + + state, err := relayconvert.NewResponseStreamState(types.RelayFormatClaude, types.RelayFormatGemini, relayconvert.ResponseStreamOptions{}) + if err != nil { + return nil, err + } + if info != nil { + info.ChatToGeminiStreamState = state + } + return state, nil +} + +func sendGeminiStreamResults(c *gin.Context, results []relayconvert.ResponseResult) *types.NewAPIError { + for _, result := range results { + geminiResponse, ok := result.Value.(*dto.GeminiChatResponse) + if !ok { + return types.NewError(fmt.Errorf("expected Gemini stream response, got %T", result.Value), types.ErrorCodeBadResponseBody) + } + if geminiResponse == nil { + continue + } + data, err := common.Marshal(geminiResponse) + if err != nil { + return types.NewError(err, types.ErrorCodeBadResponseBody) + } + c.Render(-1, common.CustomEvent{Data: "data: " + string(data)}) + _ = helper.FlushWriter(c) + } + return nil +} + func countClaudeStreamBillableTools(c *gin.Context, info *relaycommon.RelayInfo, claudeResponse *dto.ClaudeResponse) { if claudeResponse == nil { return @@ -213,6 +267,20 @@ func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, clau } } helper.Done(c) + } else if info.RelayFormat == types.RelayFormatGemini { + state, err := claudeToGeminiStreamState(info) + if err != nil { + common.SysLog("error creating Gemini stream state: " + err.Error()) + return + } + results, err := service.FinalizeStreamResponse(c, info, state) + if err != nil { + common.SysLog("error finalizing Gemini stream response: " + err.Error()) + return + } + if sendErr := sendGeminiStreamResults(c, results); sendErr != nil { + common.SysLog("send final Gemini stream response failed: " + sendErr.Error()) + } } } @@ -293,6 +361,21 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud } case types.RelayFormatClaude: responseData = data + case types.RelayFormatGemini: + { + convertResult, convertErr := service.ConvertResponse(c, info, types.RelayFormatGemini, &claudeResponse) + if convertErr != nil { + return types.NewError(convertErr, types.ErrorCodeBadResponseBody) + } + geminiResponse, ok := convertResult.Value.(*dto.GeminiChatResponse) + if !ok { + return types.NewError(fmt.Errorf("expected Gemini generateContent response, got %T", convertResult.Value), types.ErrorCodeBadResponseBody) + } + responseData, err = common.Marshal(geminiResponse) + if err != nil { + return types.NewError(err, types.ErrorCodeBadResponseBody) + } + } } if claudeResponse.Usage != nil && claudeResponse.Usage.ServerToolUse != nil && claudeResponse.Usage.ServerToolUse.WebSearchRequests > 0 { diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index e437f1277b..81b4542257 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -299,6 +299,9 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * if info.SendResponseCount == 0 { // send first response emptyResponse := helper.GenerateStartEmptyResponse(id, createAt, info.UpstreamModelName, nil) + // Claude message_start is emitted from this first OpenAI chunk. + // Carry upstream usage when the current Gemini frame provided it. + emptyResponse.Usage = response.Usage if response.IsToolCall() { if len(emptyResponse.Choices) > 0 && len(response.Choices) > 0 { toolCalls := response.Choices[0].Delta.ToolCalls diff --git a/relay/channel/gemini/relay_gemini_usage_test.go b/relay/channel/gemini/relay_gemini_usage_test.go index 1ae34c13ca..1fd9f1c3a0 100644 --- a/relay/channel/gemini/relay_gemini_usage_test.go +++ b/relay/channel/gemini/relay_gemini_usage_test.go @@ -5,6 +5,7 @@ import ( "io" "net/http" "net/http/httptest" + "strings" "testing" "github.com/QuantumNous/new-api/common" @@ -16,6 +17,130 @@ import ( "github.com/stretchr/testify/require" ) +func TestStreamResponseGeminiChat2OpenAIAttachesUsageMetadata(t *testing.T) { + t.Parallel() + + withUsage, isStop := streamResponseGeminiChat2OpenAI(&dto.GeminiChatResponse{ + Candidates: []dto.GeminiChatCandidate{{ + Content: dto.GeminiChatContent{ + Role: "model", + Parts: []dto.GeminiPart{{Text: "hello"}}, + }, + }}, + UsageMetadata: dto.GeminiUsageMetadata{ + PromptTokenCount: 3868, + CandidatesTokenCount: 0, + TotalTokenCount: 3868, + }, + }) + require.False(t, isStop) + require.NotNil(t, withUsage) + require.NotNil(t, withUsage.Usage) + require.Equal(t, 3868, withUsage.Usage.PromptTokens) + require.Equal(t, 3868, withUsage.Usage.TotalTokens) + require.NotNil(t, withUsage.Usage.BillingUsage) + require.Equal(t, dto.BillingUsageSourceGeminiChat, withUsage.Usage.BillingUsage.Source) + require.Equal(t, dto.BillingUsageSemanticGemini, withUsage.Usage.BillingUsage.Semantic) + require.NotNil(t, withUsage.Usage.BillingUsage.GeminiUsageMetadata) + require.Equal(t, 3868, withUsage.Usage.BillingUsage.GeminiUsageMetadata.PromptTokenCount) + require.False(t, withUsage.Usage.BillingUsage.Estimated) + + withoutUsage, _ := streamResponseGeminiChat2OpenAI(&dto.GeminiChatResponse{ + Candidates: []dto.GeminiChatCandidate{{ + Content: dto.GeminiChatContent{ + Role: "model", + Parts: []dto.GeminiPart{{Text: "hello"}}, + }, + }}, + }) + require.NotNil(t, withoutUsage) + require.Nil(t, withoutUsage.Usage) +} + +func TestGeminiChatStreamHandlerClaudeFirstFrameUsesUpstreamUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + oldStreamingTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 300 + t.Cleanup(func() { + constant.StreamingTimeout = oldStreamingTimeout + }) + + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatClaude, + OriginModelName: "gemini-2.5-flash", + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "gemini-2.5-flash", + }, + ClaudeConvertInfo: &relaycommon.ClaudeConvertInfo{ + LastMessagesType: relaycommon.LastMessageTypeNone, + }, + } + info.SetEstimatePromptTokens(4994) + + chunkData, err := common.Marshal(dto.GeminiChatResponse{ + Candidates: []dto.GeminiChatCandidate{{ + Content: dto.GeminiChatContent{ + Role: "model", + Parts: []dto.GeminiPart{{Text: "hello"}}, + }, + }}, + UsageMetadata: dto.GeminiUsageMetadata{ + PromptTokenCount: 3868, + TotalTokenCount: 3868, + }, + }) + require.NoError(t, err) + resp := &http.Response{ + Body: io.NopCloser(bytes.NewReader([]byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n"))), + } + + usage, newAPIError := GeminiChatStreamHandler(c, info, resp) + require.Nil(t, newAPIError) + require.NotNil(t, usage) + require.Equal(t, 3868, usage.PromptTokens) + + var startUsage, deltaUsage *dto.ClaudeUsage + for _, line := range strings.Split(recorder.Body.String(), "\n") { + payload, ok := strings.CutPrefix(strings.TrimSpace(line), "data: ") + if !ok { + continue + } + var event dto.ClaudeResponse + if err := common.UnmarshalJsonStr(payload, &event); err != nil { + continue + } + switch event.Type { + case "message_start": + if event.Message != nil { + startUsage = event.Message.Usage + } + case "message_delta": + deltaUsage = event.Usage + } + } + + require.NotNil(t, startUsage) + require.Equal(t, 3868, startUsage.InputTokens) + require.NotNil(t, startUsage.BillingUsage) + require.Equal(t, dto.BillingUsageSourceGeminiChat, startUsage.BillingUsage.Source) + require.Equal(t, dto.BillingUsageSemanticGemini, startUsage.BillingUsage.Semantic) + require.NotNil(t, startUsage.BillingUsage.GeminiUsageMetadata) + require.Equal(t, 3868, startUsage.BillingUsage.GeminiUsageMetadata.PromptTokenCount) + require.False(t, startUsage.BillingUsage.Estimated) + + require.NotNil(t, deltaUsage) + require.Equal(t, 3868, deltaUsage.InputTokens) + require.NotNil(t, deltaUsage.BillingUsage) + require.Equal(t, dto.BillingUsageSourceGeminiChat, deltaUsage.BillingUsage.Source) + require.Equal(t, dto.BillingUsageSemanticGemini, deltaUsage.BillingUsage.Semantic) + require.NotNil(t, deltaUsage.BillingUsage.GeminiUsageMetadata) + require.Equal(t, 3868, deltaUsage.BillingUsage.GeminiUsageMetadata.PromptTokenCount) +} + func TestGeminiChatHandlerCompletionTokensExcludeToolUsePromptTokens(t *testing.T) { t.Parallel() diff --git a/relay/channel/openai/helper.go b/relay/channel/openai/helper.go index 3999ba560f..a3cbd115d1 100644 --- a/relay/channel/openai/helper.go +++ b/relay/channel/openai/helper.go @@ -19,8 +19,6 @@ import ( "github.com/gin-gonic/gin" ) -const chatToGeminiStreamStateKey = "relaykit.chat_to_gemini_stream_state" - // 辅助函数 func HandleStreamFormat(c *gin.Context, info *relaycommon.RelayInfo, data string, forceFormat bool, thinkToContent bool) error { switch info.RelayFormat { @@ -68,7 +66,7 @@ func handleGeminiFormat(c *gin.Context, data string, info *relaycommon.RelayInfo return err } - state, err := chatToGeminiStreamState(c, &streamResponse) + state, err := chatToGeminiStreamState(info, &streamResponse) if err != nil { return err } @@ -79,11 +77,11 @@ func handleGeminiFormat(c *gin.Context, data string, info *relaycommon.RelayInfo return sendGeminiStreamResults(c, results) } -func chatToGeminiStreamState(c *gin.Context, streamResponse *dto.ChatCompletionsStreamResponse) (*relayconvert.ResponseStreamState, error) { - if value, ok := c.Get(chatToGeminiStreamStateKey); ok { - state, ok := value.(*relayconvert.ResponseStreamState) +func chatToGeminiStreamState(info *relaycommon.RelayInfo, streamResponse *dto.ChatCompletionsStreamResponse) (*relayconvert.ResponseStreamState, error) { + if info != nil && info.ChatToGeminiStreamState != nil { + state, ok := info.ChatToGeminiStreamState.(*relayconvert.ResponseStreamState) if !ok || state == nil { - return nil, fmt.Errorf("invalid Chat-to-Gemini stream state %T", value) + return nil, fmt.Errorf("invalid Chat-to-Gemini stream state %T", info.ChatToGeminiStreamState) } return state, nil } @@ -96,7 +94,9 @@ func chatToGeminiStreamState(c *gin.Context, streamResponse *dto.ChatCompletions if err != nil { return nil, err } - c.Set(chatToGeminiStreamStateKey, state) + if info != nil { + info.ChatToGeminiStreamState = state + } return state, nil } @@ -233,7 +233,7 @@ func HandleFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, lastStream return } - state, err := chatToGeminiStreamState(c, &streamResponse) + state, err := chatToGeminiStreamState(info, &streamResponse) if err != nil { common.SysLog("error creating Gemini stream state: " + err.Error()) return diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 727476b1ec..ca0150b969 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -123,9 +123,14 @@ type RelayInfo struct { UserSetting dto.UserSetting UserEmail string UserQuota int - RelayFormat types.RelayFormat - SendResponseCount int - ReceivedResponseCount int + RelayFormat types.RelayFormat + SendResponseCount int + // ClaudeToChatStreamState / ChatToGeminiStreamState hold per-attempt + // stream converters. InitChannelMeta nils them so a retry cannot resume a + // dirty converter (advanced tool index / finalized). + ClaudeToChatStreamState any + ChatToGeminiStreamState any + ReceivedResponseCount int FinalPreConsumedQuota int // 最终预消耗的配额 // ForcePreConsume 为 true 时禁用 BillingSession 的信任额度旁路， // 强制预扣全额。用于异步任务（视频/音乐生成等），因为请求返回后任务仍在运行， @@ -203,6 +208,11 @@ func (info *RelayInfo) InitChannelMeta(c *gin.Context) { info.FinalRequestRelayFormat = "" info.RequestConversionChain = nil info.InitRequestConversionChain() + // Per-attempt only. Do not clear StreamStatus, conversion diagnostics, + // LastError, or billing accumulators — those are request-scoped. + info.SendResponseCount = 0 + info.ClaudeToChatStreamState = nil + info.ChatToGeminiStreamState = nil channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType) paramOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelParamOverride) headerOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelHeaderOverride) diff --git a/relay/common/relay_info_test.go b/relay/common/relay_info_test.go index 5142c5414b..125df738fb 100644 --- a/relay/common/relay_info_test.go +++ b/relay/common/relay_info_test.go @@ -1,11 +1,13 @@ package common import ( + "context" "encoding/json" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/relaykit/relayconvert" "github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta" "github.com/QuantumNous/new-api/relaykit/types" "github.com/gin-gonic/gin" @@ -178,3 +180,94 @@ func TestInitChannelMetaRestoresRequestReasoningEffortForRetry(t *testing.T) { info.InitChannelMeta(ctx) assert.Equal(t, "max", info.ReasoningEffort) } + +func TestInitChannelMetaResetsPerAttemptStreamStateAndPreservesRequestState(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil) + + info, err := GenRelayInfo(ctx, types.RelayFormatOpenAI, &dto.GeneralOpenAIRequest{Model: "gpt-test"}, nil) + require.NoError(t, err) + + claudeState := relayconvert.NewClaudeToChatStreamState() + _, err = claudeState.ConvertChunk(&dto.ClaudeResponse{ + Type: "content_block_start", + Index: ptr(7), + ContentBlock: &dto.ClaudeMediaMessage{ + Type: "tool_use", + Id: "toolu_1", + Name: "lookup", + }, + }) + require.NoError(t, err) + _, err = claudeState.ConvertChunk(&dto.ClaudeResponse{ + Type: "content_block_delta", + Index: ptr(7), + Delta: &dto.ClaudeMediaMessage{ + Type: "input_json_delta", + PartialJson: ptr(`{"q":"x"}`), + }, + }) + require.NoError(t, err) + + geminiState, err := relayconvert.NewResponseStreamState(types.RelayFormatOpenAI, types.RelayFormatGemini, relayconvert.ResponseStreamOptions{ + ID: "chatcmpl_1", + Model: "gpt-test", + }) + require.NoError(t, err) + + info.SendResponseCount = 3 + info.ClaudeToChatStreamState = claudeState + info.ChatToGeminiStreamState = geminiState + info.LastError = types.NewError(assert.AnError, types.ErrorCodeBadResponseBody) + info.StreamStatus = NewStreamStatus() + info.StreamStatus.RecordError("attempt 1 soft error") + info.RecordConversionDiagnostics(context.Background(), []types.ConversionDiagnostic{{ + Code: "test.loss", + Message: "attempt 1 conversion loss", + Severity: types.ConversionDiagnosticWarning, + From: types.RelayFormatClaude, + To: types.RelayFormatOpenAI, + }}) + + info.InitChannelMeta(ctx) + + assert.Zero(t, info.SendResponseCount) + assert.Nil(t, info.ClaudeToChatStreamState) + assert.Nil(t, info.ChatToGeminiStreamState) + + require.NotNil(t, info.StreamStatus) + assert.True(t, info.StreamStatus.HasErrors()) + assert.Equal(t, 1, info.StreamStatus.TotalErrorCount()) + diagnostics := info.ConversionDiagnostics() + require.Len(t, diagnostics, 1) + assert.Equal(t, "test.loss", diagnostics[0].Code) + require.NotNil(t, info.LastError) + + freshClaude := relayconvert.NewClaudeToChatStreamState() + _, err = freshClaude.ConvertChunk(&dto.ClaudeResponse{ + Type: "content_block_delta", + Index: ptr(7), + Delta: &dto.ClaudeMediaMessage{ + Type: "input_json_delta", + PartialJson: ptr(`{"q":"x"}`), + }, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "unknown content block index") + + info.IncrSendResponseCount() + responses := relayconvert.StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Id: "chatcmpl_retry", + Model: "gpt-test", + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Content: ptr("hello")}, + }}, + }, info) + require.NotEmpty(t, responses) + assert.Equal(t, "message_start", responses[0].Type) +} + +func ptr[T any](value T) *T { + return &value +} diff --git a/relaykit/dto/billing_usage.go b/relaykit/dto/billing_usage.go index ac3535e8bf..f8cb491ed7 100644 --- a/relaykit/dto/billing_usage.go +++ b/relaykit/dto/billing_usage.go @@ -270,25 +270,12 @@ func (usage *BillingUsage) CanonicalUsage() (*Usage, bool) { func (usage *BillingUsage) canonicalOpenAIUsage() *Usage { canonical := cloneOpenAIUsage(usage.OpenAIUsage) - if inputDetails := canonical.InputTokensDetails; inputDetails != nil { - if canonical.PromptTokensDetails.CachedTokens == 0 && inputDetails.CachedTokens > 0 { - canonical.PromptTokensDetails.CachedTokens = inputDetails.CachedTokens - } - if canonical.PromptTokensDetails.CachedCreationTokens == 0 && inputDetails.CachedCreationTokens > 0 { - canonical.PromptTokensDetails.CachedCreationTokens = inputDetails.CachedCreationTokens - } - if canonical.PromptTokensDetails.CacheWriteTokens == 0 && inputDetails.CacheWriteTokens > 0 { - canonical.PromptTokensDetails.CacheWriteTokens = inputDetails.CacheWriteTokens - } - if canonical.PromptTokensDetails.TextTokens == 0 && inputDetails.TextTokens > 0 { - canonical.PromptTokensDetails.TextTokens = inputDetails.TextTokens - } - if canonical.PromptTokensDetails.ImageTokens == 0 && inputDetails.ImageTokens > 0 { - canonical.PromptTokensDetails.ImageTokens = inputDetails.ImageTokens - } - if canonical.PromptTokensDetails.AudioTokens == 0 && inputDetails.AudioTokens > 0 { - canonical.PromptTokensDetails.AudioTokens = inputDetails.AudioTokens - } + if canonical.InputTokensDetails != nil { + // InputTokensDetails fills fields that PromptTokensDetails omitted; + // existing PromptTokensDetails values stay canonical on overlap. + filled := *canonical.InputTokensDetails + mergeInputTokenDetails(&filled, canonical.PromptTokensDetails) + canonical.PromptTokensDetails = filled } if canonical.PromptTokensDetails.CachedTokens == 0 && canonical.PromptCacheHitTokens > 0 { canonical.PromptTokensDetails.CachedTokens = canonical.PromptCacheHitTokens @@ -316,12 +303,15 @@ func (usage *BillingUsage) canonicalOpenAIUsage() *Usage { func (usage *BillingUsage) canonicalClaudeUsage() *Usage { claudeUsage := usage.ClaudeUsage - cacheCreation5m := claudeUsage.GetCacheCreation5mTokens() - if cacheCreation5m == 0 { + // Flat legacy fields are a fallback only when this snapshot never carried + // a CacheCreation sub-object. Presence (non-nil), not zero vs non-zero, + // is the discriminator — a later sub-object that zeros 1h must win. + var cacheCreation5m, cacheCreation1h int + if claudeUsage.CacheCreation != nil { + cacheCreation5m = claudeUsage.GetCacheCreation5mTokens() + cacheCreation1h = claudeUsage.GetCacheCreation1hTokens() + } else { cacheCreation5m = claudeUsage.ClaudeCacheCreation5mTokens - } - cacheCreation1h := claudeUsage.GetCacheCreation1hTokens() - if cacheCreation1h == 0 { cacheCreation1h = claudeUsage.ClaudeCacheCreation1hTokens } @@ -363,7 +353,7 @@ func (usage *BillingUsage) canonicalGeminiUsage() *Usage { addGeminiInputTokenDetail(&canonical.PromptTokensDetails, detail) } for _, detail := range metadata.CandidatesTokensDetails { - switch detail.Modality { + switch normalizeGeminiModality(detail.Modality) { case "IMAGE": canonical.CompletionTokenDetails.ImageTokens += detail.TokenCount case "AUDIO": @@ -377,6 +367,9 @@ func (usage *BillingUsage) canonicalGeminiUsage() *Usage { canonical.TotalTokens = canonical.PromptTokens + canonical.CompletionTokens } else if canonical.CompletionTokens <= 0 { canonical.CompletionTokens = canonical.TotalTokens - canonical.PromptTokens + if canonical.CompletionTokens < 0 { + canonical.CompletionTokens = 0 + } } if canonical.PromptTokens > 0 && canonical.PromptTokensDetails.TextTokens == 0 && canonical.PromptTokensDetails.AudioTokens == 0 { canonical.PromptTokensDetails.TextTokens = canonical.PromptTokens @@ -385,7 +378,7 @@ func (usage *BillingUsage) canonicalGeminiUsage() *Usage { } func addGeminiInputTokenDetail(details *InputTokenDetails, detail GeminiPromptTokensDetails) { - switch detail.Modality { + switch normalizeGeminiModality(detail.Modality) { case "AUDIO": details.AudioTokens += detail.TokenCount case "IMAGE": diff --git a/relaykit/dto/billing_usage_test.go b/relaykit/dto/billing_usage_test.go index a09e2ac59c..261d5180b3 100644 --- a/relaykit/dto/billing_usage_test.go +++ b/relaykit/dto/billing_usage_test.go @@ -63,6 +63,38 @@ func TestNewEstimatedGeminiChatBillingUsage(t *testing.T) { assert.Equal(t, 18, billingUsage.GeminiUsageMetadata.TotalTokenCount) } +func TestCanonicalGeminiUsageClampsNegativeCompletionFromTotalMinusPrompt(t *testing.T) { + usage, ok := NewGeminiChatBillingUsage(&GeminiUsageMetadata{ + PromptTokenCount: 50, + TotalTokenCount: 30, + }).CanonicalUsage() + require.True(t, ok) + assert.Equal(t, 0, usage.CompletionTokens) +} + +func TestCanonicalOpenAIUsageMergesInputTokenDetailsFieldwise(t *testing.T) { + usage, ok := NewOpenAIResponsesBillingUsage(&Usage{ + PromptTokens: 10, + PromptTokensDetails: InputTokenDetails{ + CachedTokens: 8, + TextTokens: 12, + ImageTokens: 4, + AudioTokens: 3, + }, + InputTokensDetails: &InputTokenDetails{ + CachedTokens: 5, + CachedCreationTokens: 7, + TextTokens: 2, + }, + }).CanonicalUsage() + require.True(t, ok) + assert.Equal(t, 8, usage.PromptTokensDetails.CachedTokens) + assert.Equal(t, 12, usage.PromptTokensDetails.TextTokens) + assert.Equal(t, 4, usage.PromptTokensDetails.ImageTokens) + assert.Equal(t, 3, usage.PromptTokensDetails.AudioTokens) + assert.Equal(t, 7, usage.PromptTokensDetails.CachedCreationTokens) +} + func TestBillingUsageJSONUsesProtocolNamedFields(t *testing.T) { billingUsage := &BillingUsage{ OpenAIUsage: &Usage{PromptTokens: 1, BillingUsage: NewClaudeMessagesBillingUsage(&ClaudeUsage{InputTokens: 9})}, diff --git a/relaykit/dto/usage_merge.go b/relaykit/dto/usage_merge.go index 39383b1b77..0bb2fd5e22 100644 --- a/relaykit/dto/usage_merge.go +++ b/relaykit/dto/usage_merge.go @@ -101,7 +101,13 @@ func MergeBillingUsageNonZero(current *BillingUsage, incoming *BillingUsage) *Bi return CloneBillingUsage(current) } if current == nil || !sameBillingUsageDialect(current, incoming) { - return CloneBillingUsage(incoming) + replaced := CloneBillingUsage(incoming) + if current != nil && replaced != nil { + // Replacement carries the incoming payload; Estimated carries the + // history of any local synthesis on either side. + replaced.Estimated = current.Estimated || incoming.Estimated + } + return replaced } merged := CloneBillingUsage(current) @@ -120,7 +126,7 @@ func MergeBillingUsageNonZero(current *BillingUsage, incoming *BillingUsage) *Bi cloneOpenAIUsage(incoming.OpenAIUsage), ) case current.ClaudeUsage != nil && incoming.ClaudeUsage != nil: - merged.ClaudeUsage = mergeClaudeUsageNonZero(current.ClaudeUsage, incoming.ClaudeUsage) + merged.ClaudeUsage = MergeClaudeUsageNonZero(current.ClaudeUsage, incoming.ClaudeUsage) case current.GeminiUsageMetadata != nil && incoming.GeminiUsageMetadata != nil: merged.GeminiUsageMetadata = MergeGeminiUsageMetadataNonZero(current.GeminiUsageMetadata, incoming.GeminiUsageMetadata) } @@ -140,12 +146,15 @@ func sameBillingUsageDialect(current *BillingUsage, incoming *BillingUsage) bool current.GeminiUsageMetadata != nil && incoming.GeminiUsageMetadata != nil } -func mergeClaudeUsageNonZero(current *ClaudeUsage, incoming *ClaudeUsage) *ClaudeUsage { +func MergeClaudeUsageNonZero(current *ClaudeUsage, incoming *ClaudeUsage) *ClaudeUsage { merged := cloneClaudeUsage(current) if merged == nil { merged = &ClaudeUsage{} } if incoming == nil { + if current != nil { + merged.BillingUsage = CloneBillingUsage(current.BillingUsage) + } return merged } if incoming.InputTokens > 0 { @@ -169,6 +178,11 @@ func mergeClaudeUsageNonZero(current *ClaudeUsage, incoming *ClaudeUsage) *Claud if incoming.CacheCreation != nil { cacheCreation := *incoming.CacheCreation merged.CacheCreation = &cacheCreation + // Flat legacy fields are the same information as the sub-object. + // Sync them as a whole overwrite, including explicit zeros, so a + // later correction cannot leave a stale high-watermark behind. + merged.ClaudeCacheCreation5mTokens = cacheCreation.Ephemeral5mInputTokens + merged.ClaudeCacheCreation1hTokens = cacheCreation.Ephemeral1hInputTokens } if incoming.ServerToolUse != nil { if merged.ServerToolUse == nil { @@ -187,6 +201,14 @@ func mergeClaudeUsageNonZero(current *ClaudeUsage, incoming *ClaudeUsage) *Claud merged.ServerToolUse.ToolSearchRequests = incoming.ServerToolUse.ToolSearchRequests } } + // cloneClaudeUsage strips BillingUsage so a nested Claude snapshot cannot + // recurse. Restore the client-visible sidecar here: incoming wins when + // present (authoritative/upstream), otherwise keep current's. + if incoming.BillingUsage != nil { + merged.BillingUsage = CloneBillingUsage(incoming.BillingUsage) + } else if current != nil { + merged.BillingUsage = CloneBillingUsage(current.BillingUsage) + } return merged } @@ -238,19 +260,23 @@ func MergeGeminiUsageMetadataNonZero(current *GeminiUsageMetadata, incoming *Gem return &merged } +func normalizeGeminiModality(modality string) string { + return strings.ToUpper(strings.TrimSpace(modality)) +} + func mergeGeminiTokenDetails(current []GeminiPromptTokensDetails, incoming []GeminiPromptTokensDetails) []GeminiPromptTokensDetails { merged := append([]GeminiPromptTokensDetails{}, current...) indexes := make(map[string]int, len(merged)) for index, detail := range merged { - indexes[strings.ToUpper(strings.TrimSpace(detail.Modality))] = index + indexes[normalizeGeminiModality(detail.Modality)] = index } for _, detail := range incoming { if detail.TokenCount <= 0 { continue } - key := strings.ToUpper(strings.TrimSpace(detail.Modality)) + key := normalizeGeminiModality(detail.Modality) if index, ok := indexes[key]; ok { - merged[index] = detail + merged[index].TokenCount += detail.TokenCount continue } indexes[key] = len(merged) diff --git a/relaykit/dto/usage_merge_test.go b/relaykit/dto/usage_merge_test.go index 088ff02405..738a4bcafe 100644 --- a/relaykit/dto/usage_merge_test.go +++ b/relaykit/dto/usage_merge_test.go @@ -10,7 +10,7 @@ import ( func TestMergeClaudeUsageCacheCreationReplacesWholeObject(t *testing.T) { t.Parallel() - merged := mergeClaudeUsageNonZero( + merged := MergeClaudeUsageNonZero( &ClaudeUsage{ CacheCreation: &ClaudeCacheCreationUsage{Ephemeral1hInputTokens: 1000}, }, @@ -52,6 +52,173 @@ func TestMergeGeminiUsageMetadataCandidatesAndThoughtsReplacedAsPair(t *testing. assert.Equal(t, 150, usage.CompletionTokens) } +func TestGeminiModalityKeysSettleConsistentlyAndDuplicateEntriesSum(t *testing.T) { + t.Parallel() + + for _, modality := range []string{"audio", " AUDIO ", "AUDIO"} { + t.Run("settle_"+modality, func(t *testing.T) { + t.Parallel() + billing := NewGeminiChatBillingUsage(&GeminiUsageMetadata{ + PromptTokenCount: 100, + PromptTokensDetails: []GeminiPromptTokensDetails{ + {Modality: modality, TokenCount: 40}, + {Modality: "TEXT", TokenCount: 60}, + }, + }) + usage, ok := billing.CanonicalUsage() + require.True(t, ok) + assert.Equal(t, 40, usage.PromptTokensDetails.AudioTokens) + assert.Equal(t, 60, usage.PromptTokensDetails.TextTokens) + }) + } + + mergedDetails := mergeGeminiTokenDetails( + []GeminiPromptTokensDetails{{Modality: "AUDIO", TokenCount: 10}}, + []GeminiPromptTokensDetails{{Modality: "audio", TokenCount: 15}}, + ) + require.Len(t, mergedDetails, 1) + assert.Equal(t, 25, mergedDetails[0].TokenCount) + + streamMerged := MergeGeminiUsageMetadataNonZero( + &GeminiUsageMetadata{ + PromptTokenCount: 10, + PromptTokensDetails: []GeminiPromptTokensDetails{{Modality: "AUDIO", TokenCount: 10}}, + }, + &GeminiUsageMetadata{ + PromptTokenCount: 25, + PromptTokensDetails: []GeminiPromptTokensDetails{{Modality: "audio", TokenCount: 15}}, + }, + ) + require.NotNil(t, streamMerged) + streamUsage, ok := NewGeminiChatBillingUsage(streamMerged).CanonicalUsage() + require.True(t, ok) + + decodedUsage, ok := NewGeminiChatBillingUsage(&GeminiUsageMetadata{ + PromptTokenCount: 25, + PromptTokensDetails: []GeminiPromptTokensDetails{ + {Modality: "AUDIO", TokenCount: 10}, + {Modality: "audio", TokenCount: 15}, + }, + }).CanonicalUsage() + require.True(t, ok) + assert.Equal(t, decodedUsage.PromptTokensDetails.AudioTokens, streamUsage.PromptTokensDetails.AudioTokens) + assert.Equal(t, 25, decodedUsage.PromptTokensDetails.AudioTokens) +} + +func TestClaudeCacheCreationSubObjectZeroDoesNotReviveFlatLegacyFields(t *testing.T) { + t.Parallel() + + merged := MergeClaudeUsageNonZero( + &ClaudeUsage{ClaudeCacheCreation1hTokens: 1000}, + &ClaudeUsage{ + InputTokens: 10, + CacheCreation: &ClaudeCacheCreationUsage{ + Ephemeral5mInputTokens: 1000, + Ephemeral1hInputTokens: 0, + }, + }, + ) + require.NotNil(t, merged.CacheCreation) + assert.Equal(t, 1000, merged.CacheCreation.Ephemeral5mInputTokens) + assert.Equal(t, 0, merged.CacheCreation.Ephemeral1hInputTokens) + assert.Equal(t, 1000, merged.ClaudeCacheCreation5mTokens) + assert.Equal(t, 0, merged.ClaudeCacheCreation1hTokens) + + usage, ok := NewClaudeMessagesBillingUsage(merged).CanonicalUsage() + require.True(t, ok) + assert.Equal(t, 1000, usage.ClaudeCacheCreation5mTokens) + assert.Equal(t, 0, usage.ClaudeCacheCreation1hTokens) +} + +func TestClaudeCacheCreationFlatFieldsStillSettleWhenSnapshotNeverHadSubObject(t *testing.T) { + t.Parallel() + + usage, ok := NewClaudeMessagesBillingUsage(&ClaudeUsage{ + InputTokens: 10, + ClaudeCacheCreation1hTokens: 1000, + }).CanonicalUsage() + require.True(t, ok) + assert.Equal(t, 0, usage.ClaudeCacheCreation5mTokens) + assert.Equal(t, 1000, usage.ClaudeCacheCreation1hTokens) +} + +func TestMergeBillingUsageORsEstimatedOnSameAndCrossDialect(t *testing.T) { + t.Parallel() + + estimated := NewEstimatedGeminiChatBillingUsage(&Usage{PromptTokens: 10, CompletionTokens: 2}) + require.NotNil(t, estimated) + require.True(t, estimated.Estimated) + + sameDialect := MergeBillingUsageNonZero(estimated, NewGeminiChatBillingUsage(&GeminiUsageMetadata{ + PromptTokenCount: 11, + CandidatesTokenCount: 3, + TotalTokenCount: 14, + })) + require.NotNil(t, sameDialect) + assert.True(t, sameDialect.Estimated) + + crossDialect := MergeBillingUsageNonZero(estimated, NewOpenAIChatBillingUsage(&Usage{ + PromptTokens: 12, + CompletionTokens: 4, + TotalTokens: 16, + })) + require.NotNil(t, crossDialect) + assert.True(t, crossDialect.Estimated) + require.NotNil(t, crossDialect.OpenAIUsage) + assert.Equal(t, 12, crossDialect.OpenAIUsage.PromptTokens) +} + +func TestMergeClaudeUsageNonZeroPreservesBillingUsage(t *testing.T) { + t.Parallel() + + currentSidecar := NewGeminiChatBillingUsage(&GeminiUsageMetadata{ + PromptTokenCount: 3868, + TotalTokenCount: 3868, + CachedContentTokenCount: 20, + }) + incomingSidecar := NewGeminiChatBillingUsage(&GeminiUsageMetadata{ + PromptTokenCount: 3868, + CandidatesTokenCount: 12, + TotalTokenCount: 3880, + }) + require.NotNil(t, currentSidecar) + require.NotNil(t, incomingSidecar) + + withIncoming := MergeClaudeUsageNonZero( + &ClaudeUsage{ + InputTokens: 3868, + CacheReadInputTokens: 20, + BillingUsage: currentSidecar, + }, + &ClaudeUsage{ + InputTokens: 3868, + OutputTokens: 12, + BillingUsage: incomingSidecar, + }, + ) + require.NotNil(t, withIncoming.BillingUsage) + assert.Equal(t, BillingUsageSourceGeminiChat, withIncoming.BillingUsage.Source) + assert.Equal(t, BillingUsageSemanticGemini, withIncoming.BillingUsage.Semantic) + require.NotNil(t, withIncoming.BillingUsage.GeminiUsageMetadata) + assert.Equal(t, 12, withIncoming.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount) + assert.Equal(t, 20, withIncoming.CacheReadInputTokens) + assert.NotSame(t, incomingSidecar, withIncoming.BillingUsage) + + keepCurrent := MergeClaudeUsageNonZero( + &ClaudeUsage{ + InputTokens: 3868, + CacheReadInputTokens: 20, + BillingUsage: currentSidecar, + }, + &ClaudeUsage{InputTokens: 3868, OutputTokens: 12}, + ) + require.NotNil(t, keepCurrent.BillingUsage) + require.NotNil(t, keepCurrent.BillingUsage.GeminiUsageMetadata) + assert.Equal(t, 20, keepCurrent.BillingUsage.GeminiUsageMetadata.CachedContentTokenCount) + assert.Equal(t, 0, keepCurrent.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount) + assert.Equal(t, 20, keepCurrent.CacheReadInputTokens) +} + func TestMergeUsageNonZeroKeepsPositiveValuesAndTakesMaxTotal(t *testing.T) { t.Parallel() diff --git a/relaykit/relayconvert/internal/gemini_chat/to_oai_chat_resp.go b/relaykit/relayconvert/internal/gemini_chat/to_oai_chat_resp.go index c74af6d577..58116d7dd5 100644 --- a/relaykit/relayconvert/internal/gemini_chat/to_oai_chat_resp.go +++ b/relaykit/relayconvert/internal/gemini_chat/to_oai_chat_resp.go @@ -283,6 +283,11 @@ func StreamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d Object: "chat.completion.chunk", Choices: choices, } + // Only attach usage the chunk actually reported. Do not fall back to a + // local prompt estimate — converters treat this as first-frame truth. + if metadata := geminiResponse.GetUsageMetadata(); dto.HasGeminiUsageMetadataTokens(metadata) { + response.Usage = UsageFromGeminiMetadata(metadata, 0) + } return &response, isStop } diff --git a/relaykit/relayconvert/internal/oai_chat/to_claude_messages_resp.go b/relaykit/relayconvert/internal/oai_chat/to_claude_messages_resp.go index 78fb875bc3..06e11356fd 100644 --- a/relaykit/relayconvert/internal/oai_chat/to_claude_messages_resp.go +++ b/relaykit/relayconvert/internal/oai_chat/to_claude_messages_resp.go @@ -89,6 +89,21 @@ func buildClaudeUsageFromOpenAIUsage(oaiUsage *dto.Usage) *dto.ClaudeUsage { return sharedclaude.UsageFromOpenAI(oaiUsage) } +func clientVisibleClaudeStreamUsage(state *convmeta.ClaudeConvertInfo, incoming *dto.Usage) *dto.ClaudeUsage { + prior := buildClaudeUsageFromOpenAIUsage(state.Usage) + converted := buildClaudeUsageFromOpenAIUsage(incoming) + if incoming != nil { + state.Usage = dto.MergeUsageNonZero(state.Usage, incoming) + } + if prior == nil { + return converted + } + if converted == nil { + return prior + } + return dto.MergeClaudeUsageNonZero(prior, converted) +} + func NormalizeCacheCreationSplit(totalTokens int, tokens5m int, tokens1h int) (int, int) { return sharedclaude.NormalizeCacheCreationSplit(totalTokens, tokens5m, tokens1h) } @@ -168,15 +183,27 @@ func StreamResponseOpenAI2Claude(openAIResponse *dto.ChatCompletionsStreamRespon } } if info.GetSendResponseCount() == 1 { + // Client-visible Claude stream usage matches billing merge: first + // frame records first, a later non-zero field overrides, and a later + // zero/missing field never erases a first-frame positive. Anthropic + // clients treat message_delta as authoritative for the fields it + // carries, including a corrected input_tokens. + startUsage := &dto.ClaudeUsage{ + InputTokens: info.GetEstimatePromptTokens(), + OutputTokens: 0, + } + if openAIResponse.Usage != nil && dto.HasOpenAIUsageTokens(openAIResponse.Usage) { + if real := buildClaudeUsageFromOpenAIUsage(openAIResponse.Usage); real != nil { + startUsage = real + } + state.Usage = dto.MergeUsageNonZero(state.Usage, openAIResponse.Usage) + } msg := &dto.ClaudeMediaMessage{ Id: openAIResponse.Id, Model: openAIResponse.Model, Type: "message", Role: "assistant", - Usage: &dto.ClaudeUsage{ - InputTokens: info.GetEstimatePromptTokens(), - OutputTokens: 0, - }, + Usage: startUsage, } msg.SetContent(make([]any, 0)) claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ @@ -187,10 +214,7 @@ func StreamResponseOpenAI2Claude(openAIResponse *dto.ChatCompletionsStreamRespon if len(openAIResponse.Choices) == 0 { // Some OpenAI-compatible upstreams end with a usage-only SSE chunk. - oaiUsage := openAIResponse.Usage - if oaiUsage == nil { - oaiUsage = state.Usage - } + oaiUsage := clientVisibleClaudeStreamUsage(state, openAIResponse.Usage) if oaiUsage != nil { appendStopOpenBlocks() stopReason := stopReasonOpenAI2Claude(state.FinishReason) @@ -199,7 +223,7 @@ func StreamResponseOpenAI2Claude(openAIResponse *dto.ChatCompletionsStreamRespon } claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ Type: "message_delta", - Usage: buildClaudeUsageFromOpenAIUsage(oaiUsage), + Usage: oaiUsage, Delta: &dto.ClaudeMediaMessage{ StopReason: kitutil.GetPointer[string](stopReason), }, @@ -367,10 +391,7 @@ func StreamResponseOpenAI2Claude(openAIResponse *dto.ChatCompletionsStreamRespon appendCitationDeltas(chosenChoice.Delta.Annotations) if doneChunk || state.Done { - oaiUsage := openAIResponse.Usage - if oaiUsage == nil { - oaiUsage = state.Usage - } + oaiUsage := clientVisibleClaudeStreamUsage(state, openAIResponse.Usage) if oaiUsage == nil { // Some upstreams emit finish_reason first, then send a final usage-only chunk. // Keep content blocks open until usage is available so the terminal message_delta @@ -380,7 +401,7 @@ func StreamResponseOpenAI2Claude(openAIResponse *dto.ChatCompletionsStreamRespon appendStopOpenBlocks() claudeResponses = append(claudeResponses, &dto.ClaudeResponse{ Type: "message_delta", - Usage: buildClaudeUsageFromOpenAIUsage(oaiUsage), + Usage: oaiUsage, Delta: &dto.ClaudeMediaMessage{ StopReason: kitutil.GetPointer[string](stopReasonOpenAI2Claude(state.FinishReason)), }, @@ -414,7 +435,7 @@ func FinalizeStreamResponseOpenAI2Claude(info convmeta.Meta) []*dto.ClaudeRespon responses = append(responses, &dto.ClaudeResponse{ Type: "message_delta", - Usage: buildClaudeUsageFromOpenAIUsage(state.Usage), + Usage: clientVisibleClaudeStreamUsage(state, nil), Delta: &dto.ClaudeMediaMessage{ StopReason: kitutil.GetPointer[string](stopReason), }, diff --git a/relaykit/relayconvert/internal/oai_chat/to_claude_messages_resp_test.go b/relaykit/relayconvert/internal/oai_chat/to_claude_messages_resp_test.go index 976e725919..cbb38813da 100644 --- a/relaykit/relayconvert/internal/oai_chat/to_claude_messages_resp_test.go +++ b/relaykit/relayconvert/internal/oai_chat/to_claude_messages_resp_test.go @@ -223,6 +223,192 @@ func TestStreamResponseOpenAI2ClaudeClosesTextThinkingAndToolBlocks(t *testing.T assert.Equal(t, "message_stop", finishResponses[2].Type) } +func TestStreamResponseOpenAI2ClaudeFirstFrameUsesUpstreamUsageWhenPresent(t *testing.T) { + info := &convmeta.Values{ + EstimatePromptTokens: 32, + SendResponseCount: 1, + ClaudeConvertInfo: &convmeta.ClaudeConvertInfo{LastMessagesType: convmeta.LastMessageTypeNone}, + } + + responses := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Id: "chatcmpl_1", + Model: "gpt-test", + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Content: ptr("hello")}, + }}, + Usage: &dto.Usage{PromptTokens: 29, CompletionTokens: 0, TotalTokens: 29}, + }, info) + require.NotEmpty(t, responses) + require.Equal(t, "message_start", responses[0].Type) + require.NotNil(t, responses[0].Message) + require.NotNil(t, responses[0].Message.Usage) + assert.Equal(t, 29, responses[0].Message.Usage.InputTokens) +} + +func TestStreamResponseOpenAI2ClaudeMessageDeltaCorrectsEstimatedFirstFrame(t *testing.T) { + info := &convmeta.Values{ + EstimatePromptTokens: 32, + SendResponseCount: 1, + ClaudeConvertInfo: &convmeta.ClaudeConvertInfo{LastMessagesType: convmeta.LastMessageTypeNone}, + } + + first := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Id: "chatcmpl_1", + Model: "gpt-test", + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Content: ptr("hello")}, + }}, + }, info) + require.NotEmpty(t, first) + require.Equal(t, "message_start", first[0].Type) + require.NotNil(t, first[0].Message.Usage) + assert.Equal(t, 32, first[0].Message.Usage.InputTokens) + + info.SendResponseCount = 2 + finish := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Id: "chatcmpl_1", + Model: "gpt-test", + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + FinishReason: ptr("stop"), + }}, + Usage: &dto.Usage{PromptTokens: 29, CompletionTokens: 4, TotalTokens: 33}, + }, info) + var delta *dto.ClaudeResponse + for _, resp := range finish { + if resp.Type == "message_delta" { + delta = resp + break + } + } + require.NotNil(t, delta) + require.NotNil(t, delta.Usage) + assert.Equal(t, 29, delta.Usage.InputTokens) + assert.Equal(t, 4, delta.Usage.OutputTokens) +} + +func TestStreamResponseOpenAI2ClaudeMessageDeltaDoesNotZeroFirstFrameCache(t *testing.T) { + info := &convmeta.Values{ + EstimatePromptTokens: 8, + SendResponseCount: 1, + ClaudeConvertInfo: &convmeta.ClaudeConvertInfo{LastMessagesType: convmeta.LastMessageTypeNone}, + } + + first := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Id: "chatcmpl_1", + Model: "gpt-test", + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Content: ptr("hello")}, + }}, + Usage: &dto.Usage{ + PromptTokens: 40, + CompletionTokens: 0, + TotalTokens: 40, + PromptTokensDetails: dto.InputTokenDetails{ + CachedTokens: 20, + CachedCreationTokens: 10, + }, + }, + }, info) + require.Equal(t, "message_start", first[0].Type) + require.NotNil(t, first[0].Message.Usage) + assert.Equal(t, 20, first[0].Message.Usage.CacheReadInputTokens) + assert.Equal(t, 10, first[0].Message.Usage.CacheCreationInputTokens) + + info.SendResponseCount = 2 + finish := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Id: "chatcmpl_1", + Model: "gpt-test", + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + FinishReason: ptr("stop"), + }}, + Usage: &dto.Usage{PromptTokens: 29, CompletionTokens: 4, TotalTokens: 33}, + }, info) + var delta *dto.ClaudeResponse + for _, resp := range finish { + if resp.Type == "message_delta" { + delta = resp + break + } + } + require.NotNil(t, delta) + require.NotNil(t, delta.Usage) + assert.Equal(t, 29, delta.Usage.InputTokens) + assert.Equal(t, 20, delta.Usage.CacheReadInputTokens) + assert.Equal(t, 10, delta.Usage.CacheCreationInputTokens) +} + +func TestStreamResponseOpenAI2ClaudeGeminiBillingUsageOnStartAndDelta(t *testing.T) { + info := &convmeta.Values{ + EstimatePromptTokens: 4994, + SendResponseCount: 1, + ClaudeConvertInfo: &convmeta.ClaudeConvertInfo{LastMessagesType: convmeta.LastMessageTypeNone}, + } + + firstUsage := &dto.Usage{ + PromptTokens: 3868, + CompletionTokens: 0, + TotalTokens: 3868, + BillingUsage: dto.NewGeminiChatBillingUsage(&dto.GeminiUsageMetadata{ + PromptTokenCount: 3868, + TotalTokenCount: 3868, + }), + } + first := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Id: "chatcmpl_1", + Model: "gpt-test", + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Content: ptr("hello")}, + }}, + Usage: firstUsage, + }, info) + require.NotEmpty(t, first) + require.Equal(t, "message_start", first[0].Type) + require.NotNil(t, first[0].Message) + require.NotNil(t, first[0].Message.Usage) + assert.Equal(t, 3868, first[0].Message.Usage.InputTokens) + require.NotNil(t, first[0].Message.Usage.BillingUsage) + assert.Equal(t, dto.BillingUsageSourceGeminiChat, first[0].Message.Usage.BillingUsage.Source) + assert.Equal(t, dto.BillingUsageSemanticGemini, first[0].Message.Usage.BillingUsage.Semantic) + require.NotNil(t, first[0].Message.Usage.BillingUsage.GeminiUsageMetadata) + assert.Equal(t, 3868, first[0].Message.Usage.BillingUsage.GeminiUsageMetadata.PromptTokenCount) + + info.SendResponseCount = 2 + finish := StreamResponseOpenAI2Claude(&dto.ChatCompletionsStreamResponse{ + Id: "chatcmpl_1", + Model: "gpt-test", + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + FinishReason: ptr("stop"), + }}, + Usage: &dto.Usage{ + PromptTokens: 3868, + CompletionTokens: 12, + TotalTokens: 3880, + BillingUsage: dto.NewGeminiChatBillingUsage(&dto.GeminiUsageMetadata{ + PromptTokenCount: 3868, + CandidatesTokenCount: 12, + TotalTokenCount: 3880, + }), + }, + }, info) + var delta *dto.ClaudeResponse + for _, resp := range finish { + if resp.Type == "message_delta" { + delta = resp + break + } + } + require.NotNil(t, delta) + require.NotNil(t, delta.Usage) + assert.Equal(t, 3868, delta.Usage.InputTokens) + assert.Equal(t, 12, delta.Usage.OutputTokens) + require.NotNil(t, delta.Usage.BillingUsage) + assert.Equal(t, dto.BillingUsageSourceGeminiChat, delta.Usage.BillingUsage.Source) + assert.Equal(t, dto.BillingUsageSemanticGemini, delta.Usage.BillingUsage.Semantic) + require.NotNil(t, delta.Usage.BillingUsage.GeminiUsageMetadata) + assert.Equal(t, 3868, delta.Usage.BillingUsage.GeminiUsageMetadata.PromptTokenCount) + assert.Equal(t, 12, delta.Usage.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount) +} + func TestNormalizeCacheCreationSplit(t *testing.T) { cache5m, cache1h := NormalizeCacheCreationSplit(10, 3, 2) assert.Equal(t, 8, cache5m) diff --git a/relaykit/relayconvert/internal/oai_chat/to_gemini_chat_req.go b/relaykit/relayconvert/internal/oai_chat/to_gemini_chat_req.go index 3828b5fa03..7e721632bc 100644 --- a/relaykit/relayconvert/internal/oai_chat/to_gemini_chat_req.go +++ b/relaykit/relayconvert/internal/oai_chat/to_gemini_chat_req.go @@ -157,22 +157,7 @@ func OpenAIChatRequestToGeminiGenerateContent(c context.Context, textRequest dto if textRequest.Tools != nil { functions := make([]dto.FunctionRequest, 0, len(textRequest.Tools)) - googleSearch := false - codeExecution := false - urlContext := false for _, tool := range textRequest.Tools { - if tool.Function.Name == "googleSearch" { - googleSearch = true - continue - } - if tool.Function.Name == "codeExecution" { - codeExecution = true - continue - } - if tool.Function.Name == "urlContext" { - urlContext = true - continue - } if tool.Function.Parameters != nil { if params, ok := tool.Function.Parameters.(map[string]interface{}); ok { if props, hasProps := params["properties"].(map[string]interface{}); hasProps && len(props) == 0 { @@ -184,21 +169,6 @@ func OpenAIChatRequestToGeminiGenerateContent(c context.Context, textRequest dto functions = append(functions, tool.Function) } geminiTools := geminiRequest.GetTools() - if codeExecution { - geminiTools = append(geminiTools, dto.GeminiChatTool{ - CodeExecution: make(map[string]string), - }) - } - if googleSearch { - geminiTools = append(geminiTools, dto.GeminiChatTool{ - GoogleSearch: make(map[string]string), - }) - } - if urlContext { - geminiTools = append(geminiTools, dto.GeminiChatTool{ - URLContext: make(map[string]string), - }) - } if len(functions) > 0 { geminiTools = append(geminiTools, dto.GeminiChatTool{ FunctionDeclarations: functions, diff --git a/relaykit/relayconvert/internal/toolconv/decode.go b/relaykit/relayconvert/internal/toolconv/decode.go index 748fed041d..ba64370bbc 100644 --- a/relaykit/relayconvert/internal/toolconv/decode.go +++ b/relaykit/relayconvert/internal/toolconv/decode.go @@ -56,6 +56,10 @@ func extractOpenAIChatRequest(request any) (any, Set, error) { } for index, tool := range source.Tools { if tool.Type == "function" || tool.Type == "" { + if definition, ok := decodeOpenAIChatPseudoHostedTool(tool.Function.Name); ok { + set.Definitions = append(set.Definitions, definition) + continue + } set.Definitions = append(set.Definitions, Definition{ Kind: KindFunction, Execution: ExecutionClient, @@ -511,6 +515,40 @@ func rawBoolPointer(raw json.RawMessage) *bool { return &value } +// decodeOpenAIChatPseudoHostedTool recognizes the OpenAI Chat dialect that +// declares Gemini hosted tools as function definitions named googleSearch, +// codeExecution, or urlContext. The names are the historical public contract; +// recognition lives here so every target format goes through the same hosted +// ToolDefinition pipeline. +func decodeOpenAIChatPseudoHostedTool(name string) (Definition, bool) { + switch name { + case "googleSearch": + return Definition{ + Kind: KindWebSearch, + Execution: ExecutionServer, + NativeType: "googleSearch", + Name: "googleSearch", + WebSearch: &WebSearch{}, + }, true + case "codeExecution": + return Definition{ + Kind: KindCodeExecution, + Execution: ExecutionServer, + NativeType: "codeExecution", + Name: "codeExecution", + }, true + case "urlContext": + return Definition{ + Kind: KindURLContext, + Execution: ExecutionServer, + NativeType: "urlContext", + Name: "urlContext", + }, true + default: + return Definition{}, false + } +} + func decodeOpenAIChatLocation(raw json.RawMessage) (*ApproximateLocation, error) { if len(raw) == 0 { return nil, nil diff --git a/relaykit/relayconvert/internal/toolconv/encode.go b/relaykit/relayconvert/internal/toolconv/encode.go index 224901f478..b2e60132ab 100644 --- a/relaykit/relayconvert/internal/toolconv/encode.go +++ b/relaykit/relayconvert/internal/toolconv/encode.go @@ -381,13 +381,13 @@ func attachGeminiRequest(request any, set Set) (any, []types.ConversionDiagnosti diagnostics = append(diagnostics, geminiWebSearchDiagnostics(index, definition.WebSearch)...) } case KindCodeExecution: - if set.Source == types.RelayFormatGemini { + if set.Source == types.RelayFormatGemini || definition.NativeType == "codeExecution" { tools = append(tools, map[string]any{"codeExecution": map[string]any{}}) continue } diagnostics = append(diagnostics, semanticLoss(fmt.Sprintf("tools[%d]", index), "unverified_tool_mapping", "code execution semantics differ across providers")) case KindURLContext: - if set.Source == types.RelayFormatGemini { + if set.Source == types.RelayFormatGemini || definition.NativeType == "urlContext" { tools = append(tools, map[string]any{"urlContext": map[string]any{}}) continue }