shaw
bd52e5d770
Fixes #5148: with aggregator upstreams (e.g. newapi) that end SSE streams without a proper terminal event, every such request was silently missing from usage logs and billing. Root cause (tracked via the nested audit issue): the low-level Anthropic SSE readers already return the partially collected usage together with the stream error (missing terminal event, read error, interval timeout), but Forward converted every such result to (nil, err) and the handler returned before submitting RecordUsage. Changes: - Add partialStreamUsageResult: on stream errors, wrap observed usage into a ForwardResult and return it alongside the error, for both the regular Anthropic path and the API-key passthrough path. Invariants: UpstreamFailoverError always keeps result=nil (failover retries are billed as the successful attempt, never twice), and zero observed usage returns no partial result (no phantom zero-usage records). - Messages handler: hoist the usage submission block into a closure shared by the success path and the new partial-result error path. - Usage record worker pool: distinguish pool-stopped drops (dropped_stopped) from operator-configured drop/sample overflow drops; billing tasks now fall back to inline synchronous execution only during the shutdown window, while explicit drop/sample overflow semantics are preserved. Image usage keeps its mandatory fallback for both drop kinds via the new mode.Dropped() helper. Tests: Forward-level regressions for missing-terminal / read-error / no-usage / failover-invariant on both paths, plus handler-level stopped-pool sync fallback and drop-policy preservation tests.
227 lines
8.9 KiB
Go
227 lines
8.9 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"io"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"strings"
|
||
"testing"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/stretchr/testify/require"
|
||
)
|
||
|
||
// 本文件覆盖 issue #5148：流式转发中途出错（缺失 terminal 事件、读错误等）时，
|
||
// 已观测到的上游 usage 不得随错误一起被丢弃，Forward 必须把部分结果与错误一同
|
||
// 返回，供 handler 照常提交 usage 记录。
|
||
|
||
func newForwardPartialUsageServiceForTest(upstream *anthropicHTTPUpstreamRecorder) *GatewayService {
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
return &GatewayService{
|
||
cfg: cfg,
|
||
responseHeaderFilter: compileResponseHeaderFilter(cfg),
|
||
httpUpstream: upstream,
|
||
rateLimitService: &RateLimitService{},
|
||
deferredService: &DeferredService{},
|
||
}
|
||
}
|
||
|
||
func newAnthropicOAuthAccountForPartialUsageTest() *Account {
|
||
return &Account{
|
||
ID: 501,
|
||
Name: "anthropic-oauth-partial-usage",
|
||
Platform: PlatformAnthropic,
|
||
Type: AccountTypeOAuth,
|
||
Concurrency: 1,
|
||
Credentials: map[string]any{
|
||
"access_token": "oauth-token",
|
||
},
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
}
|
||
}
|
||
|
||
func TestGatewayService_Forward_StreamMissingTerminalPreservesPartialUsage(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
|
||
body := []byte(`{"model":"claude-3-5-sonnet-latest","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`)
|
||
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), PlatformAnthropic)
|
||
require.NoError(t, err)
|
||
|
||
// newapi 等聚合上游的典型失败形态：message_start/message_delta 携带 usage，
|
||
// 但流在 message_stop 前直接结束。
|
||
upstreamSSE := strings.Join([]string{
|
||
`event: message_start`,
|
||
`data: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-3-5-sonnet-latest","content":[],"usage":{"input_tokens":11,"cache_read_input_tokens":7}}}`,
|
||
"",
|
||
`event: content_block_delta`,
|
||
`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}`,
|
||
"",
|
||
`event: message_delta`,
|
||
`data: {"type":"message_delta","delta":{"stop_reason":null},"usage":{"output_tokens":5}}`,
|
||
"",
|
||
"",
|
||
}, "\n")
|
||
upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{
|
||
"Content-Type": []string{"text/event-stream"},
|
||
"X-Request-Id": []string{"rid-partial"},
|
||
},
|
||
Body: io.NopCloser(strings.NewReader(upstreamSSE)),
|
||
}}
|
||
svc := newForwardPartialUsageServiceForTest(upstream)
|
||
account := newAnthropicOAuthAccountForPartialUsageTest()
|
||
|
||
result, err := svc.Forward(context.Background(), c, account, parsed)
|
||
require.Error(t, err)
|
||
require.Contains(t, err.Error(), "missing terminal event")
|
||
require.NotNil(t, result, "流中断但已观测到 usage 时必须返回部分结果用于计费")
|
||
require.True(t, result.Stream)
|
||
require.Equal(t, 11, result.Usage.InputTokens)
|
||
require.Equal(t, 7, result.Usage.CacheReadInputTokens)
|
||
require.Equal(t, 5, result.Usage.OutputTokens)
|
||
require.Equal(t, "rid-partial", result.RequestID)
|
||
require.NotNil(t, result.FirstTokenMs)
|
||
}
|
||
|
||
func TestGatewayService_Forward_StreamReadErrorAfterOutputPreservesPartialUsage(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
|
||
body := []byte(`{"model":"claude-3-5-sonnet-latest","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`)
|
||
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), PlatformAnthropic)
|
||
require.NoError(t, err)
|
||
|
||
// message_start 已写出（含 usage），随后上游连接异常中断。
|
||
upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||
Body: &streamReadCloser{
|
||
payload: []byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":9,\"cache_creation_input_tokens\":4}}}\n\n"),
|
||
err: io.ErrUnexpectedEOF,
|
||
},
|
||
}}
|
||
svc := newForwardPartialUsageServiceForTest(upstream)
|
||
account := newAnthropicOAuthAccountForPartialUsageTest()
|
||
|
||
result, err := svc.Forward(context.Background(), c, account, parsed)
|
||
require.Error(t, err)
|
||
require.Contains(t, err.Error(), "stream read error")
|
||
require.NotNil(t, result, "已写出内容后的读错误必须保留部分 usage")
|
||
require.Equal(t, 9, result.Usage.InputTokens)
|
||
require.Equal(t, 4, result.Usage.CacheCreationInputTokens)
|
||
}
|
||
|
||
func TestGatewayService_Forward_StreamErrorWithoutUsageReturnsNilResult(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
|
||
body := []byte(`{"model":"claude-3-5-sonnet-latest","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`)
|
||
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), PlatformAnthropic)
|
||
require.NoError(t, err)
|
||
|
||
// 只有 ping、没有任何 usage 的流中断：不应产生零 usage 的幽灵账单记录。
|
||
upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||
Body: io.NopCloser(strings.NewReader("event: ping\ndata: {\"type\": \"ping\"}\n\n")),
|
||
}}
|
||
svc := newForwardPartialUsageServiceForTest(upstream)
|
||
account := newAnthropicOAuthAccountForPartialUsageTest()
|
||
|
||
result, err := svc.Forward(context.Background(), c, account, parsed)
|
||
require.Error(t, err)
|
||
require.Contains(t, err.Error(), "missing terminal event")
|
||
require.Nil(t, result, "无已观测 usage 时不应返回部分结果")
|
||
}
|
||
|
||
func TestGatewayService_Forward_FailoverErrorKeepsNilResult(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
|
||
body := []byte(`{"model":"claude-3-5-sonnet-latest","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`)
|
||
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), PlatformAnthropic)
|
||
require.NoError(t, err)
|
||
|
||
// 未向客户端写出任何字节前的读错误会包成 UpstreamFailoverError 走换号重试。
|
||
// 该路径必须保持 result=nil：failover 成功后按成功请求计费，双份结果会重复计费。
|
||
upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||
Body: &streamReadCloser{
|
||
err: errors.New("connection reset by peer"),
|
||
},
|
||
}}
|
||
svc := newForwardPartialUsageServiceForTest(upstream)
|
||
account := newAnthropicOAuthAccountForPartialUsageTest()
|
||
|
||
result, err := svc.Forward(context.Background(), c, account, parsed)
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.True(t, errors.As(err, &failoverErr))
|
||
require.Nil(t, result, "failover 错误必须保持 result=nil，防止重试成功后双重计费")
|
||
}
|
||
|
||
func TestGatewayService_AnthropicAPIKeyPassthrough_ForwardStreamMissingTerminalPreservesPartialUsage(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
|
||
body := []byte(`{"model":"claude-3-7-sonnet-20250219","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`)
|
||
parsed := &ParsedRequest{
|
||
Body: NewRequestBodyRef(body),
|
||
Model: "claude-3-7-sonnet-20250219",
|
||
Stream: true,
|
||
}
|
||
|
||
upstreamSSE := strings.Join([]string{
|
||
`data: {"type":"message_start","message":{"usage":{"input_tokens":9,"cache_read_input_tokens":2}}}`,
|
||
"",
|
||
`data: {"type":"message_delta","usage":{"output_tokens":3}}`,
|
||
"",
|
||
}, "\n")
|
||
upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{
|
||
"Content-Type": []string{"text/event-stream"},
|
||
"X-Request-Id": []string{"rid-pass-partial"},
|
||
},
|
||
Body: io.NopCloser(strings.NewReader(upstreamSSE)),
|
||
}}
|
||
svc := newForwardPartialUsageServiceForTest(upstream)
|
||
account := newAnthropicAPIKeyAccountForTest()
|
||
|
||
result, err := svc.Forward(context.Background(), c, account, parsed)
|
||
require.Error(t, err)
|
||
require.Contains(t, err.Error(), "missing terminal event")
|
||
require.NotNil(t, result, "透传流中断但已观测到 usage 时必须返回部分结果用于计费")
|
||
require.True(t, result.Stream)
|
||
require.Equal(t, 9, result.Usage.InputTokens)
|
||
require.Equal(t, 2, result.Usage.CacheReadInputTokens)
|
||
require.Equal(t, 3, result.Usage.OutputTokens)
|
||
require.Equal(t, "claude-3-7-sonnet-20250219", result.Model)
|
||
}
|