Merge pull request #2505 from is7Qin/fix/openai-compat-usage-parsing

修复 Claude 映射 GPT 后被记为 0 token 的计费漏洞
This commit is contained in:
Wesley Liddick 2026-05-19 09:53:50 +08:00 committed by GitHub
commit f9fec78b70
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 152 additions and 24 deletions

View File

@ -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"`

View File

@ -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)

View File

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

View File

@ -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) {

View File

@ -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) {

View File

@ -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) {