github/sub2api
4
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-30 16:23:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): forward received chat stream usage

Browse Source
This commit is contained in:
wucm667
2026-09-21 12:58:56 +08:00
parent 7c700729c2
commit b13200d7a1
5 changed files with 138 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/anthropic_chat_stream_usage_test.go
+111
View File
@@ -0,0 +1,111 @@
//go:build unit
package service
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestAnthropicChatStreamAuthoritativeUsage(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, adapter := range []string{"anthropic", "native"} {
for _, options := range []string{"", `,"stream_options":{"include_usage":true}`, `,"stream_options":{"include_usage":false}`} {
for _, tc := range []struct {
name, start, delta string
want bool
prompt, output, cached, created int
}{
{"cache", `,"usage":{"input_tokens":308,"cache_read_input_tokens":241,"cache_creation_input_tokens":17}`, `,"usage":{"output_tokens":49}`, true, 566, 49, 241, 17},
{"zero_start", `,"usage":{"input_tokens":0,"output_tokens":0}`, "", true, 0, 0, 0, 0},
{"zero_delta", "", `,"usage":{"input_tokens":0,"output_tokens":0}`, true, 0, 0, 0, 0},
{"absent", "", "", false, 0, 0, 0, 0},
{"null", `,"usage":null`, `,"usage":null`, false, 0, 0, 0, 0},
} {
for _, terminal := range []string{"stop", "eof"} {
t.Run(adapter+"/"+options+"/"+tc.name+"/"+terminal, func(t *testing.T) {
sse := fmt.Sprintf("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"k3\",\"content\":[]%s}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}%s}\n\n", tc.start, tc.delta)
if terminal == "stop" {
// Repeated terminal events and EOF must not duplicate usage.
sse += "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(sse)),
}}
body := []byte(`{"model":"k3","messages":[{"role":"user","content":"hi"}],"stream":true` + options + "}")
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(string(body)))
if adapter == "native" {
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
result, err := svc.forwardChatCompletionsViaNativeAnthropic(context.Background(), c, nativeAnthropicTestAccount(), body, "")
require.NoError(t, err)
require.Equal(t, tc.output, result.Usage.OutputTokens)
require.Equal(t, tc.cached, result.Usage.CacheReadInputTokens)
require.Equal(t, tc.created, result.Usage.CacheCreationInputTokens)
} else {
svc := &GatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
account := &Account{ID: 1, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "sk-test"}}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, nil)
require.NoError(t, err)
require.Equal(t, tc.output, result.Usage.OutputTokens)
require.Equal(t, tc.cached, result.Usage.CacheReadInputTokens)
require.Equal(t, tc.created, result.Usage.CacheCreationInputTokens)
}
count, done := 0, 0
var last *apicompat.ChatCompletionsChunk
for _, line := range strings.Split(rec.Body.String(), "\n") {
if !strings.HasPrefix(line, "data: ") {
continue
}
payload := strings.TrimPrefix(line, "data: ")
if payload == "[DONE]" {
done++
if tc.want {
require.NotNil(t, last)
require.NotNil(t, last.Usage)
}
continue
}
require.Zero(t, done, "chunk after DONE")
var chunk apicompat.ChatCompletionsChunk
require.NoError(t, json.Unmarshal([]byte(payload), &chunk))
last = &chunk
if chunk.Usage == nil {
continue
}
count++
require.Empty(t, chunk.Choices)
require.Equal(t, tc.prompt, chunk.Usage.PromptTokens)
require.Equal(t, tc.output, chunk.Usage.CompletionTokens)
require.Equal(t, tc.prompt+tc.output, chunk.Usage.TotalTokens)
if tc.cached > 0 {
require.NotNil(t, chunk.Usage.PromptTokensDetails)
require.Equal(t, tc.cached, chunk.Usage.PromptTokensDetails.CachedTokens)
require.Equal(t, tc.created, chunk.Usage.PromptTokensDetails.CacheCreationTokens)
}
}
require.Equal(t, 1, done)
if tc.want {
require.Equal(t, 1, count)
} else {
require.Zero(t, count)
}
})
}
}
}
}
}
backend/internal/service/gateway_forward_as_chat_completions.go
+18 -4
View File
@@ -42,7 +42,6 @@ func (s *GatewayService) ForwardAsChatCompletions(
}
originalModel := ccReq.Model
clientStream := ccReq.Stream
includeUsage := ccReq.StreamOptions != nil && ccReq.StreamOptions.IncludeUsage
// 2. Convert CC → Responses → Anthropic (chained conversion)
responsesReq, err := apicompat.ChatCompletionsToResponses(&ccReq)
@@ -185,7 +184,7 @@ func (s *GatewayService) ForwardAsChatCompletions(
var result *ForwardResult
var handleErr error
if clientStream {
result, handleErr = s.handleCCStreamingFromAnthropic(resp, c, originalModel, mappedModel, reasoningEffort, startTime, includeUsage)
result, handleErr = s.handleCCStreamingFromAnthropic(resp, c, originalModel, mappedModel, reasoningEffort, startTime)
} else {
result, handleErr = s.handleCCBufferedFromAnthropic(resp, c, originalModel, mappedModel, reasoningEffort, startTime)
}
@@ -358,7 +357,6 @@ func (s *GatewayService) handleCCStreamingFromAnthropic(
mappedModel string,
reasoningEffort *string,
startTime time.Time,
includeUsage bool,
) (*ForwardResult, error) {
requestID := resp.Header.Get("x-request-id")
@@ -376,7 +374,6 @@ func (s *GatewayService) handleCCStreamingFromAnthropic(
anthState.Model = originalModel
ccState := apicompat.NewResponsesEventToChatState()
ccState.Model = originalModel
ccState.IncludeUsage = includeUsage
var usage ClaudeUsage
var firstTokenMs *int
@@ -467,6 +464,10 @@ func (s *GatewayService) handleCCStreamingFromAnthropic(
continue
}
// Forward received usage regardless of the client stream_options.
// The intermediate Responses converter synthesizes usage even when absent.
ccState.IncludeUsage = ccState.IncludeUsage || anthropicChatStreamHasUsage(&event, payload)
if processAnthropicEvent(&event) {
return resultWithUsage(), nil
}
@@ -512,3 +513,16 @@ func writeGatewayCCError(c *gin.Context, statusCode int, errType, message string
},
})
}
// anthropicChatStreamHasUsage distinguishes an explicit zero-valued usage object
// from omitted/null usage before the Anthropic→Responses conversion loses that distinction.
func anthropicChatStreamHasUsage(event *apicompat.AnthropicStreamEvent, payload string) bool {
switch event.Type {
case "message_start":
return event.Message != nil && gjson.Get(payload, "message.usage").IsObject()
case "message_delta":
return event.Usage != nil
default:
return false
}
}
backend/internal/service/gateway_forward_as_chat_completions_test.go
+2 -2
View File
@@ -199,7 +199,7 @@ func TestHandleCCStreamingFromAnthropic_CompactSSEFormat(t *testing.T) {
}
svc := &GatewayService{}
result, err := svc.handleCCStreamingFromAnthropic(resp, c, "k3", "k3", nil, time.Now(), true)
result, err := svc.handleCCStreamingFromAnthropic(resp, c, "k3", "k3", nil, time.Now())
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, 21, result.Usage.InputTokens)
@@ -236,7 +236,7 @@ func TestHandleCCStreamingFromAnthropic_PreservesMessageStartCacheUsageAndReason
}
svc := &GatewayService{}
result, err := svc.handleCCStreamingFromAnthropic(resp, c, "gpt-5", "claude-sonnet-4.5", &reasoningEffort, time.Now(), true)
result, err := svc.handleCCStreamingFromAnthropic(resp, c, "gpt-5", "claude-sonnet-4.5", &reasoningEffort, time.Now())
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, 20, result.Usage.InputTokens)
backend/internal/service/openai_gateway_anthropic_native_pump_test.go
+2 -2
View File
@@ -140,7 +140,7 @@ func TestCCStreamingFromNativeAnthropic_HangTimesOut(t *testing.T) {
resp, pr, pw := newHangingUpstreamResponse()
start := time.Now()
res, err := svc.handleCCStreamingFromNativeAnthropic(resp, c, "glm-4.7", "glm-4.7", "glm-4.7", nil, start, true)
res, err := svc.handleCCStreamingFromNativeAnthropic(resp, c, "glm-4.7", "glm-4.7", "glm-4.7", nil, start)
_ = pw.Close()
_ = pr.Close()
@@ -220,7 +220,7 @@ func TestCCStreamingFromNativeAnthropic_HappyPathStillConverts(t *testing.T) {
}()
defer func() { _ = pr.Close() }()
res, err := svc.handleCCStreamingFromNativeAnthropic(resp, c, "glm-4.7", "glm-4.7", "glm-4.7", nil, time.Now(), true)
res, err := svc.handleCCStreamingFromNativeAnthropic(resp, c, "glm-4.7", "glm-4.7", "glm-4.7", nil, time.Now())
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
backend/internal/service/openai_gateway_chat_completions_anthropic_native.go
+5 -4
View File
@@ -54,7 +54,6 @@ func (s *OpenAIGatewayService) forwardChatCompletionsViaNativeAnthropic(
return nil, fmt.Errorf("missing model in request")
}
clientStream := ccReq.Stream
includeUsage := ccReq.StreamOptions != nil && ccReq.StreamOptions.IncludeUsage
// 2. Convert CC → Responses → Anthropic (chained conversion)
responsesReq, err := apicompat.ChatCompletionsToResponses(&ccReq)
@@ -136,7 +135,7 @@ func (s *OpenAIGatewayService) forwardChatCompletionsViaNativeAnthropic(
reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, billingModel)
if clientStream {
return s.handleCCStreamingFromNativeAnthropic(resp, c, originalModel, billingModel, upstreamModel, reasoningEffort, startTime, includeUsage)
return s.handleCCStreamingFromNativeAnthropic(resp, c, originalModel, billingModel, upstreamModel, reasoningEffort, startTime)
}
return s.handleCCBufferedFromNativeAnthropic(resp, c, originalModel, billingModel, upstreamModel, reasoningEffort, startTime)
}
@@ -303,7 +302,6 @@ func (s *OpenAIGatewayService) handleCCStreamingFromNativeAnthropic(
upstreamModel string,
reasoningEffort *string,
startTime time.Time,
includeUsage bool,
) (*OpenAIForwardResult, error) {
requestID := resp.Header.Get("x-request-id")
@@ -320,7 +318,6 @@ func (s *OpenAIGatewayService) handleCCStreamingFromNativeAnthropic(
anthState.Model = originalModel
ccState := apicompat.NewResponsesEventToChatState()
ccState.Model = originalModel
ccState.IncludeUsage = includeUsage
var usage ClaudeUsage
var firstTokenMs *int
@@ -461,6 +458,10 @@ func (s *OpenAIGatewayService) handleCCStreamingFromNativeAnthropic(
continue
}
// Forward received usage regardless of the client stream_options.
// The intermediate Responses converter synthesizes usage even when absent.
ccState.IncludeUsage = ccState.IncludeUsage || anthropicChatStreamHasUsage(&event, payload)
if processAnthropicEvent(&event) {
return resultWithUsage(), nil
}