diff --git a/backend/internal/pkg/apicompat/types.go b/backend/internal/pkg/apicompat/types.go index f9cd5a1c..df75ce50 100644 --- a/backend/internal/pkg/apicompat/types.go +++ b/backend/internal/pkg/apicompat/types.go @@ -306,6 +306,37 @@ type ResponsesUsage struct { OutputTokensDetails *ResponsesOutputTokensDetails `json:"output_tokens_details,omitempty"` } +func (u *ResponsesUsage) UnmarshalJSON(data []byte) error { + type responsesUsageAlias ResponsesUsage + var aux struct { + responsesUsageAlias + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + PromptTokensDetails *ResponsesInputTokensDetails `json:"prompt_tokens_details,omitempty"` + CompletionTokensDetails *ResponsesOutputTokensDetails `json:"completion_tokens_details,omitempty"` + } + if err := json.Unmarshal(data, &aux); err != nil { + return err + } + *u = ResponsesUsage(aux.responsesUsageAlias) + if u.InputTokens == 0 && aux.PromptTokens != 0 { + u.InputTokens = aux.PromptTokens + } + if u.OutputTokens == 0 && aux.CompletionTokens != 0 { + u.OutputTokens = aux.CompletionTokens + } + if u.InputTokensDetails == nil && aux.PromptTokensDetails != nil { + u.InputTokensDetails = aux.PromptTokensDetails + } + if u.OutputTokensDetails == nil && aux.CompletionTokensDetails != nil { + u.OutputTokensDetails = aux.CompletionTokensDetails + } + if u.TotalTokens == 0 && (u.InputTokens != 0 || u.OutputTokens != 0) { + u.TotalTokens = u.InputTokens + u.OutputTokens + } + return nil +} + // ResponsesInputTokensDetails breaks down input token usage. type ResponsesInputTokensDetails struct { CachedTokens int `json:"cached_tokens,omitempty"` diff --git a/backend/internal/service/openai_compat_model_test.go b/backend/internal/service/openai_compat_model_test.go index 0ba2a63f..f8b9d360 100644 --- a/backend/internal/service/openai_compat_model_test.go +++ b/backend/internal/service/openai_compat_model_test.go @@ -183,6 +183,63 @@ func TestForwardAsAnthropic_NormalizesRoutingAndEffortForGpt54XHigh(t *testing.T t.Logf("response body: %s", rec.Body.String()) } +func TestForwardAsAnthropic_MappedClaudeModelAcceptsChatUsageShape(t *testing.T) { + t.Parallel() + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := []byte(`{"model":"claude-opus-4-7","max_tokens":16,"messages":[{"role":"user","content":"compact this"}],"stream":true}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + upstreamBody := strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp_compact","model":"gpt-5.5","status":"in_progress","output":[]}}`, + "", + `data: {"type":"response.output_text.delta","delta":"ok"}`, + "", + `data: {"type":"response.completed","response":{"id":"resp_compact","object":"response","model":"gpt-5.5","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"prompt_tokens":31,"completion_tokens":9,"total_tokens":40,"prompt_tokens_details":{"cached_tokens":11}}}}`, + "", + "data: [DONE]", + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_compact_usage"}}, + Body: io.NopCloser(strings.NewReader(upstreamBody)), + }} + + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}}, + } + account := &Account{ + ID: 1, + Name: "openai-apikey", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://api.openai.com/v1", + "model_mapping": map[string]any{ + "gpt-5.5": "gpt-5.5", + }, + }, + } + + result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "gpt-5.5") + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "claude-opus-4-7", result.Model) + require.Equal(t, "gpt-5.5", result.BillingModel) + require.Equal(t, "gpt-5.5", result.UpstreamModel) + require.Equal(t, 31, result.Usage.InputTokens) + require.Equal(t, 9, result.Usage.OutputTokens) + require.Equal(t, 11, result.Usage.CacheReadInputTokens) + require.Equal(t, "gpt-5.5", gjson.GetBytes(upstream.lastBody, "model").String()) +} + func TestForwardAsAnthropic_InjectsPromptCacheKeyForAPIKeyMessagesDispatch(t *testing.T) { t.Parallel() gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index a2276353..97f88222 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -4709,28 +4709,47 @@ func (s *OpenAIGatewayService) parseSSEUsageBytes(data []byte, usage *OpenAIUsag return } - usage.InputTokens = int(gjson.GetBytes(data, "response.usage.input_tokens").Int()) - usage.OutputTokens = int(gjson.GetBytes(data, "response.usage.output_tokens").Int()) - usage.CacheReadInputTokens = int(gjson.GetBytes(data, "response.usage.input_tokens_details.cached_tokens").Int()) - usage.ImageOutputTokens = int(gjson.GetBytes(data, "response.usage.output_tokens_details.image_tokens").Int()) + if parsedUsage, ok := extractOpenAIUsageFromJSONBytes(data); ok { + *usage = parsedUsage + } } func extractOpenAIUsageFromJSONBytes(body []byte) (OpenAIUsage, bool) { if len(body) == 0 || !gjson.ValidBytes(body) { return OpenAIUsage{}, false } - values := gjson.GetManyBytes( - body, - "usage.input_tokens", - "usage.output_tokens", - "usage.input_tokens_details.cached_tokens", - "usage.output_tokens_details.image_tokens", - ) + if usage, ok := openAIUsageFromGJSON(gjson.GetBytes(body, "usage")); ok { + return usage, true + } + return openAIUsageFromGJSON(gjson.GetBytes(body, "response.usage")) +} + +func openAIUsageFromGJSON(value gjson.Result) (OpenAIUsage, bool) { + if !value.Exists() || !value.IsObject() { + return OpenAIUsage{}, false + } + inputTokens := value.Get("input_tokens").Int() + if inputTokens == 0 { + inputTokens = value.Get("prompt_tokens").Int() + } + outputTokens := value.Get("output_tokens").Int() + if outputTokens == 0 { + outputTokens = value.Get("completion_tokens").Int() + } + cacheReadTokens := value.Get("input_tokens_details.cached_tokens").Int() + if cacheReadTokens == 0 { + cacheReadTokens = value.Get("prompt_tokens_details.cached_tokens").Int() + } + imageOutputTokens := value.Get("output_tokens_details.image_tokens").Int() + if imageOutputTokens == 0 { + imageOutputTokens = value.Get("completion_tokens_details.image_tokens").Int() + } return OpenAIUsage{ - InputTokens: int(values[0].Int()), - OutputTokens: int(values[1].Int()), - CacheReadInputTokens: int(values[2].Int()), - ImageOutputTokens: int(values[3].Int()), + InputTokens: int(inputTokens), + OutputTokens: int(outputTokens), + CacheCreationInputTokens: int(value.Get("cache_creation_input_tokens").Int()), + CacheReadInputTokens: int(cacheReadTokens), + ImageOutputTokens: int(imageOutputTokens), }, true } diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index d636cf27..7013a477 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -2174,6 +2174,25 @@ func TestParseSSEUsage_SelectiveParsing(t *testing.T) { require.Equal(t, 13, usage.InputTokens) require.Equal(t, 15, usage.OutputTokens) require.Equal(t, 4, usage.CacheReadInputTokens) + + svc.parseSSEUsage(`{"type":"response.completed","response":{"usage":{"prompt_tokens":21,"completion_tokens":8,"prompt_tokens_details":{"cached_tokens":6}}}}`, usage) + require.Equal(t, 21, usage.InputTokens) + require.Equal(t, 8, usage.OutputTokens) + require.Equal(t, 6, usage.CacheReadInputTokens) +} + +func TestExtractOpenAIUsageFromJSONBytes_AcceptsResponseAndChatUsageShapes(t *testing.T) { + usage, ok := extractOpenAIUsageFromJSONBytes([]byte(`{"id":"resp_1","usage":{"input_tokens":3,"output_tokens":5,"input_tokens_details":{"cached_tokens":2}}}`)) + require.True(t, ok) + require.Equal(t, 3, usage.InputTokens) + require.Equal(t, 5, usage.OutputTokens) + require.Equal(t, 2, usage.CacheReadInputTokens) + + usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"type":"response.completed","response":{"usage":{"prompt_tokens":13,"completion_tokens":7,"prompt_tokens_details":{"cached_tokens":4}}}}`)) + require.True(t, ok) + require.Equal(t, 13, usage.InputTokens) + require.Equal(t, 7, usage.OutputTokens) + require.Equal(t, 4, usage.CacheReadInputTokens) } func TestExtractCodexFinalResponse_SampleReplay(t *testing.T) { diff --git a/backend/internal/service/openai_ws_forwarder.go b/backend/internal/service/openai_ws_forwarder.go index 77cf7d95..192ff90a 100644 --- a/backend/internal/service/openai_ws_forwarder.go +++ b/backend/internal/service/openai_ws_forwarder.go @@ -399,15 +399,9 @@ func parseOpenAIWSResponseUsageFromCompletedEvent(message []byte, usage *OpenAIU if usage == nil || len(message) == 0 { return } - values := gjson.GetManyBytes( - message, - "response.usage.input_tokens", - "response.usage.output_tokens", - "response.usage.input_tokens_details.cached_tokens", - ) - usage.InputTokens = int(values[0].Int()) - usage.OutputTokens = int(values[1].Int()) - usage.CacheReadInputTokens = int(values[2].Int()) + if parsedUsage, ok := extractOpenAIUsageFromJSONBytes(message); ok { + *usage = parsedUsage + } } func parseOpenAIWSErrorEventFields(message []byte) (code string, errType string, errMessage string) { diff --git a/backend/internal/service/openai_ws_forwarder_hotpath_optimization_test.go b/backend/internal/service/openai_ws_forwarder_hotpath_optimization_test.go index 76167603..0350bde9 100644 --- a/backend/internal/service/openai_ws_forwarder_hotpath_optimization_test.go +++ b/backend/internal/service/openai_ws_forwarder_hotpath_optimization_test.go @@ -29,6 +29,14 @@ func TestParseOpenAIWSResponseUsageFromCompletedEvent(t *testing.T) { require.Equal(t, 11, usage.InputTokens) require.Equal(t, 7, usage.OutputTokens) require.Equal(t, 3, usage.CacheReadInputTokens) + + parseOpenAIWSResponseUsageFromCompletedEvent( + []byte(`{"type":"response.completed","response":{"usage":{"prompt_tokens":19,"completion_tokens":5,"prompt_tokens_details":{"cached_tokens":4}}}}`), + usage, + ) + require.Equal(t, 19, usage.InputTokens) + require.Equal(t, 5, usage.OutputTokens) + require.Equal(t, 4, usage.CacheReadInputTokens) } func TestOpenAIWSErrorEventHelpers_ConsistentWithWrapper(t *testing.T) {