Compare commits

..
Author SHA1 Message Date
shaw bd52e5d770 fix(gateway): record observed usage when anthropic stream is interrupted
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.
2026-08-01 11:29:58 +08:00
Wesley Liddick d4cada3b6b Merge pull request #5089 from feeeei/fix/openai_sse_rate_limit
fix(openai): retry SSE rate limits as HTTP 429
2026-08-01 10:50:57 +08:00
Wesley Liddick 8f5caef78e Merge pull request #5153 from Wei-Shaw/fix/issue-5152-classifier-multi-system-entries
fix(anthropic): recognize classifier requests with extra system entries
2026-08-01 10:50:27 +08:00
feeeei 85a27fae39 fix(openai): retry SSE rate limits as HTTP 429
OpenAI Responses may return code=rate_limit_exceeded in a response.failed SSE event while the HTTP status remains 200. Classify these failures as 429 so configured pool-mode retries and account failover are applied. No account-level rate-limit state is written on this path: the 200-stream response headers carry normal quota snapshots, and retry semantics stay owned by the failover engine.

Ignore compact keepalive bytes when determining whether semantic output has started, preserving safe retries before real output.
2026-08-01 10:01:47 +08:00
16 changed files with 747 additions and 77 deletions
+75 -56
View File
@@ -827,6 +827,65 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
if accountReleaseFunc != nil {
accountReleaseFunc()
}
// 提交 usage 记录。成功路径与"流中断但 Forward 已观测到 usage 的部分结果"
// 错误路径共用:后者若不入账,上游已计量的请求会完全漏记漏计费(#5148)。
submitForwardUsage := func(result *service.ForwardResult) {
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context)
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
// Forward 内部可能继续改写 body,usage 去重指纹必须使用最终上游接受的当前 body。
requestPayloadHash := service.HashUsageRequestPayload(attemptParsedReq.Body.Bytes())
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
if result.ReasoningEffort == nil {
result.ReasoningEffort = service.NormalizeClaudeOutputEffort(attemptParsedReq.OutputEffort)
}
// 同上(重试路径中的对称填充)。详见非重试路径同名注释。
if result.ReasoningEffort == nil && attemptParsedReq.ThinkingEnabled {
protocolModel := result.UpstreamModel
if protocolModel == "" {
protocolModel = result.Model
}
result.ReasoningEffort = service.DefaultEffortForThinkingEnabled(protocolModel)
}
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
forceCacheBilling := fs.ForceCacheBilling
quotaPlatform := service.QuotaPlatform(c.Request.Context(), currentAPIKey)
sessionID := service.ExtractClientSessionID(c)
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: result,
QuotaPlatform: quotaPlatform,
APIKey: currentAPIKey,
User: currentAPIKey.User,
Account: account,
Subscription: currentSubscription,
InboundEndpoint: inboundEndpoint,
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
SessionID: sessionID,
RequestPayloadHash: requestPayloadHash,
ForceCacheBilling: forceCacheBilling,
APIKeyService: h.apiKeyService,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
}); err != nil {
logger.L().With(
zap.String("component", "handler.gateway.messages"),
zap.Int64("user_id", subject.UserID),
zap.Int64("api_key_id", currentAPIKey.ID),
zap.Any("group_id", currentAPIKey.GroupID),
zap.String("model", reqModel),
zap.Int64("account_id", account.ID),
).Error("gateway.record_usage_failed", zap.Error(err))
}
})
}
if err != nil {
// Beta policy block: return 400 immediately, no failover
var betaBlockedErr *service.BetaBlockedError
@@ -925,6 +984,12 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
forwardFailedFields = append(forwardFailedFields, zap.Int64p("proxy_id", account.ProxyID))
}
reqLog.Error("gateway.forward_failed", forwardFailedFields...)
// Forward 与错误一起返回的部分结果:流中断前上游已计量的 usage 照常入账,
// 避免上游已产生消耗的请求完全漏记(#5148)。failover 错误恒定 result=nil,
// 不会走到这里重复计费。
if result != nil {
submitForwardUsage(result)
}
return
}
@@ -948,59 +1013,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
}
}
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context)
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
// Forward 内部可能继续改写 body,usage 去重指纹必须使用最终上游接受的当前 body。
requestPayloadHash := service.HashUsageRequestPayload(attemptParsedReq.Body.Bytes())
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
if result.ReasoningEffort == nil {
result.ReasoningEffort = service.NormalizeClaudeOutputEffort(attemptParsedReq.OutputEffort)
}
// 同上(重试路径中的对称填充)。详见非重试路径同名注释。
if result.ReasoningEffort == nil && attemptParsedReq.ThinkingEnabled {
protocolModel := result.UpstreamModel
if protocolModel == "" {
protocolModel = result.Model
}
result.ReasoningEffort = service.DefaultEffortForThinkingEnabled(protocolModel)
}
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
forceCacheBilling := fs.ForceCacheBilling
quotaPlatform := service.QuotaPlatform(c.Request.Context(), currentAPIKey)
sessionID := service.ExtractClientSessionID(c)
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: result,
QuotaPlatform: quotaPlatform,
APIKey: currentAPIKey,
User: currentAPIKey.User,
Account: account,
Subscription: currentSubscription,
InboundEndpoint: inboundEndpoint,
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
SessionID: sessionID,
RequestPayloadHash: requestPayloadHash,
ForceCacheBilling: forceCacheBilling,
APIKeyService: h.apiKeyService,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
}); err != nil {
logger.L().With(
zap.String("component", "handler.gateway.messages"),
zap.Int64("user_id", subject.UserID),
zap.Int64("api_key_id", currentAPIKey.ID),
zap.Any("group_id", currentAPIKey.GroupID),
zap.String("model", reqModel),
zap.Int64("account_id", account.ID),
).Error("gateway.record_usage_failed", zap.Error(err))
}
})
submitForwardUsage(result)
return
}
if !retryWithFallback {
@@ -2336,10 +2349,16 @@ func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task serv
}
task = wrapUsageRecordTaskContext(parent, task)
if h.usageRecordWorkerPool != nil {
h.usageRecordWorkerPool.Submit(task)
return
if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped {
return
}
// 池已停止(进程关停窗口):计费任务不能静默丢失,降级为内联同步执行。
// 显式配置的 drop/sample 溢出丢弃仍按配置语义保留。
logger.L().With(
zap.String("component", "handler.gateway.messages"),
).Warn("gateway.usage_record_task_stopped_sync_fallback")
}
// 回退路径:worker 池未注入时同步执行,避免退回到无界 goroutine 模式。
// 回退路径:worker 池未注入或已停止时同步执行,避免退回到无界 goroutine 模式。
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
defer func() {
@@ -2183,10 +2183,16 @@ func (h *OpenAIGatewayHandler) submitUsageRecordTask(parent context.Context, tas
}
task = wrapUsageRecordTaskContext(parent, task)
if h.usageRecordWorkerPool != nil {
h.usageRecordWorkerPool.Submit(task)
return
if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped {
return
}
// 池已停止(进程关停窗口):计费任务不能静默丢失,降级为内联同步执行。
// 显式配置的 drop/sample 溢出丢弃仍按配置语义保留。
logger.L().With(
zap.String("component", "handler.openai_gateway.responses"),
).Warn("openai.usage_record_task_stopped_sync_fallback")
}
// 回退路径:worker 池未注入时同步执行,避免退回到无界 goroutine 模式。
// 回退路径:worker 池未注入或已停止时同步执行,避免退回到无界 goroutine 模式。
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
defer func() {
@@ -2214,7 +2220,7 @@ func (h *OpenAIGatewayHandler) submitMandatoryUsageRecordTask(parent context.Con
}
task = wrapUsageRecordTaskContext(parent, task)
if h.usageRecordWorkerPool != nil {
if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDropped {
if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() {
return
}
logger.L().With(
@@ -1812,6 +1812,40 @@ func (u *openAIHTTPPassthroughFailoverUpstream) calls() []int64 {
return append([]int64(nil), u.accountIDs...)
}
type openAIHTTPPassthroughSSERateLimitUpstream struct {
service.HTTPUpstream
mu sync.Mutex
accountIDs []int64
}
func (u *openAIHTTPPassthroughSSERateLimitUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
u.mu.Lock()
u.accountIDs = append(u.accountIDs, accountID)
u.mu.Unlock()
body := strings.Join([]string{
"event: response.created",
`data: {"type":"response.created","response":{"id":"resp_rate_limited"}}`,
"",
"event: response.failed",
`data: {"type":"response.failed","response":{"id":"resp_rate_limited","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
"",
}, "\n")
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/event-stream"},
"Retry-After": []string{"1"},
},
Body: io.NopCloser(strings.NewReader(body)),
}, nil
}
func (u *openAIHTTPPassthroughSSERateLimitUpstream) calls() []int64 {
u.mu.Lock()
defer u.mu.Unlock()
return append([]int64(nil), u.accountIDs...)
}
func (s *openAIWSFailoverHandlerAccountRepoStub) ListSchedulableByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
out := make([]service.Account, 0, len(s.accounts))
for _, account := range s.accounts {
@@ -2012,6 +2046,88 @@ func TestOpenAIResponses_APIKeyPassthroughPool5xxRetriesThenExhaustsMaxSwitches(
require.Equal(t, "Upstream service temporarily unavailable", gjson.GetBytes(rec.Body.Bytes(), "error.message").String())
}
func TestOpenAIResponses_APIKeyPassthroughSSERateLimitUsesConfiguredPoolRetry(t *testing.T) {
gin.SetMode(gin.TestMode)
groupID := int64(4204)
accounts := []service.Account{
{
ID: 9912, Name: "pool-sse-rate-limit", Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 1,
Credentials: map[string]any{
"api_key": "sk-pool",
"base_url": "https://api.example.test",
"pool_mode": true,
"pool_mode_retry_count": float64(1),
"pool_mode_retry_status_codes": []any{float64(http.StatusTooManyRequests)},
},
Extra: map[string]any{"openai_passthrough": true},
},
}
cfg := &config.Config{RunMode: config.RunModeSimple}
cfg.Default.RateMultiplier = 1
cfg.Security.URLAllowlist.Enabled = false
cfg.Gateway.MaxAccountSwitches = 1
accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
upstream := &openAIHTTPPassthroughSSERateLimitUpstream{}
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
t.Cleanup(billingCacheSvc.Stop)
gatewaySvc := service.NewOpenAIGatewayService(
accountRepo,
nil,
nil,
nil,
nil,
nil,
nil,
cfg,
nil,
nil,
service.NewBillingService(cfg, nil),
nil,
billingCacheSvc,
upstream,
&service.DeferredService{},
nil,
nil,
nil,
nil,
nil,
nil,
nil,
)
h := NewOpenAIGatewayHandler(
gatewaySvc,
service.NewConcurrencyService(nil),
billingCacheSvc,
service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg),
nil,
nil,
nil,
nil,
cfg,
)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(`{"model":"gpt-5.6-sol","input":"hello","stream":true}`))
c.Request.Header.Set("Content-Type", "application/json")
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
ID: 1804, GroupID: &groupID,
User: &service.User{ID: 1704, Status: service.StatusActive},
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
})
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1704, Concurrency: 0})
h.Responses(c)
require.Equal(t, []int64{9912, 9912}, upstream.calls())
require.Equal(t, http.StatusTooManyRequests, rec.Code)
require.Equal(t, "1", rec.Header().Get("Retry-After"))
require.Equal(t, "rate_limit_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
require.Equal(t, "Upstream rate limit exceeded, please retry later", gjson.GetBytes(rec.Body.Bytes(), "error.message").String())
}
func TestOpenAIResponsesWebSocket_FailoverOnUpstreamUsageLimitEvent(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -0,0 +1,77 @@
package handler
import (
"context"
"sync/atomic"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
// 本文件覆盖:worker 池已停止(进程关停窗口)时,计费任务不得静默丢失,
// 必须降级为内联同步执行;显式配置的 drop/sample 溢出丢弃仍按配置语义保留。
func newStoppedUsageRecordPoolForTest() *service.UsageRecordWorkerPool {
pool := service.NewUsageRecordWorkerPoolWithOptions(service.UsageRecordWorkerPoolOptions{
WorkerCount: 1,
QueueSize: 1,
TaskTimeout: time.Second,
OverflowPolicy: config.UsageRecordOverflowPolicySync,
})
pool.Stop()
return pool
}
func TestGatewayHandlerSubmitUsageRecordTask_StoppedPoolFallsBackToSync(t *testing.T) {
h := &GatewayHandler{usageRecordWorkerPool: newStoppedUsageRecordPoolForTest()}
executed := false
h.submitUsageRecordTask(context.Background(), func(ctx context.Context) {
executed = true
})
require.True(t, executed, "池已停止时计费任务必须内联同步执行")
}
func TestOpenAIGatewayHandlerSubmitUsageRecordTask_StoppedPoolFallsBackToSync(t *testing.T) {
h := &OpenAIGatewayHandler{usageRecordWorkerPool: newStoppedUsageRecordPoolForTest()}
executed := false
h.submitUsageRecordTask(context.Background(), func(ctx context.Context) {
executed = true
})
require.True(t, executed, "池已停止时计费任务必须内联同步执行")
}
func TestGatewayHandlerSubmitUsageRecordTask_DropPolicyOverflowStillDrops(t *testing.T) {
pool := service.NewUsageRecordWorkerPoolWithOptions(service.UsageRecordWorkerPoolOptions{
WorkerCount: 1,
QueueSize: 1,
TaskTimeout: time.Minute,
OverflowPolicy: config.UsageRecordOverflowPolicyDrop,
})
t.Cleanup(pool.Stop)
h := &GatewayHandler{usageRecordWorkerPool: pool}
started := make(chan struct{})
block := make(chan struct{})
t.Cleanup(func() { close(block) })
// 占满 worker 槽位后再填满队列,保证第三个任务触发溢出。
require.Equal(t, service.UsageRecordSubmitModeEnqueued, pool.Submit(func(ctx context.Context) {
close(started)
<-block
}))
<-started
require.Equal(t, service.UsageRecordSubmitModeEnqueued, pool.Submit(func(ctx context.Context) {
<-block
}))
var executed atomic.Bool
h.submitUsageRecordTask(context.Background(), func(ctx context.Context) {
executed.Store(true)
})
time.Sleep(50 * time.Millisecond)
require.False(t, executed.Load(), "drop 溢出策略是运维显式配置的取舍,不应被同步兜底覆盖")
}
@@ -269,6 +269,11 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
if input.RequestStream {
streamResult, err := s.handleStreamingResponseAnthropicAPIKeyPassthrough(ctx, resp, c, account, input.StartTime, input.RequestModel)
if err != nil {
// 流中断时保留已观测到的 usage 与错误一起返回,避免上游已计量的请求
// 完全漏记漏计费(issue #5148)。
if partial := partialStreamUsageResult(resp, streamResult, input.OriginalModel, input.RequestModel, input.StartTime, err); partial != nil {
return partial, err
}
return nil, err
}
usage = streamResult.usage
@@ -843,6 +843,11 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
ResponseBody: body,
}
}
// 流中断(缺失 terminal 事件、读错误、数据间隔超时等)时保留已观测到的
// usage 与错误一起返回,handler 在错误处理完成后照常提交 usage 记录。
if partial := partialStreamUsageResult(resp, streamResult, originalModel, mappedModel, startTime, err); partial != nil {
return partial, err
}
return nil, err
}
usage = streamResult.usage
@@ -0,0 +1,226 @@
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)
}
@@ -647,6 +647,44 @@ type streamingResult struct {
clientDisconnect bool // 客户端是否在流式传输过程中断开
}
// hasObservedTokens 报告流式过程中是否已观测到任何上游计量的 token。
func (u *ClaudeUsage) hasObservedTokens() bool {
if u == nil {
return false
}
return u.InputTokens > 0 || u.OutputTokens > 0 ||
u.CacheCreationInputTokens > 0 || u.CacheReadInputTokens > 0 ||
u.CacheCreation5mTokens > 0 || u.CacheCreation1hTokens > 0 ||
u.ImageOutputTokens > 0
}
// partialStreamUsageResult 在流式转发中途出错时,把已观测到 usage 的部分结果包装为
// ForwardResult(与错误一起返回给 handler 记录)。上游一旦下发过 message_start,
// input/cache token 就已计量,直接丢弃会让请求完全漏记漏计费(issue #5148)。
// 无已观测 usage 时返回 nil。
//
// 不变式:UpstreamFailoverError 必须保持 result=nil——failover 重试成功后按成功请求
// 计费,若同时返回部分 usage 会造成双重计费,此处显式拦截兜底。
func partialStreamUsageResult(resp *http.Response, streamResult *streamingResult, model, upstreamModel string, startTime time.Time, err error) *ForwardResult {
if streamResult == nil || !streamResult.usage.hasObservedTokens() {
return nil
}
var failoverErr *UpstreamFailoverError
if errors.As(err, &failoverErr) {
return nil
}
return &ForwardResult{
RequestID: resp.Header.Get("x-request-id"),
Usage: *streamResult.usage,
Model: model,
UpstreamModel: upstreamModel,
Stream: true,
Duration: time.Since(startTime),
FirstTokenMs: streamResult.firstTokenMs,
ClientDisconnect: streamResult.clientDisconnect,
}
}
func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel string, mimicClaudeCode bool) (*streamingResult, error) {
// 更新5h窗口状态
s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header)
@@ -284,6 +284,20 @@ func TestOpenAICompactKeepaliveAdjustedWrittenSize_ExcludesHeartbeatBytes(t *tes
require.Contains(t, rec.Body.String(), ": keepalive\n\n")
}
func TestOpenAIStreamClientOutputStarted_IgnoresCompactKeepaliveBytes(t *testing.T) {
c, _ := newCompactBridgeTestContext(t, true)
stop := StartOpenAICompactSSEKeepalive(c, keepaliveTestInterval)
defer stop()
waitForKeepaliveBeats()
require.True(t, c.Writer.Written())
require.False(t, openAIStreamClientOutputStarted(c, false), "keepalive comments are not semantic output")
_, err := c.Writer.Write([]byte("real-output"))
require.NoError(t, err)
require.True(t, openAIStreamClientOutputStarted(c, false))
}
// fast policy block 在心跳未提交时保持 403 JSON 原语义。
func TestWriteOpenAIFastPolicyBlockedResponse_BeforeKeepaliveCommit(t *testing.T) {
c, rec := newCompactBridgeTestContext(t, true)
@@ -441,7 +441,7 @@ func (s *OpenAIGatewayService) handleChatBufferedStreamingResponse(
}
message := openAICompatFailedResponseMessage(finalResponse)
if openAIStreamFailedEventShouldFailover(payload, message) {
return nil, s.newOpenAIStreamFailoverError(c, account, false, requestID, payload, message)
return nil, s.newOpenAIStreamFailoverError(c, account, false, requestID, payload, message, resp.Header)
}
message = s.recordOpenAIStreamUpstreamError(c, account, false, requestID, "http_error", payload, message)
// response.failed 到达在 HTTP 200 SSE 流上,无真实 HTTP 错误码;统一走语义
@@ -608,7 +608,7 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse(
return true
}
if openAIStreamFailedEventShouldFailover(payloadBytes, message) {
streamFailoverErr = s.newOpenAIStreamFailoverError(c, account, false, requestID, payloadBytes, message)
streamFailoverErr = s.newOpenAIStreamFailoverError(c, account, false, requestID, payloadBytes, message, resp.Header)
return true
}
message = s.recordOpenAIStreamUpstreamError(c, account, false, requestID, "http_error", payloadBytes, message)
@@ -561,7 +561,7 @@ func (s *OpenAIGatewayService) handleAnthropicBufferedStreamingResponse(
}
message := openAICompatFailedResponseMessage(finalResponse)
if openAIStreamFailedEventShouldFailover(payload, message) {
return nil, s.newOpenAIStreamFailoverError(c, account, false, requestID, payload, message)
return nil, s.newOpenAIStreamFailoverError(c, account, false, requestID, payload, message, resp.Header)
}
message = s.recordOpenAIStreamUpstreamError(c, account, false, requestID, "http_error", payload, message)
// 统一走语义状态推断 + body 归一化(与 /v1/responses 路径一致),
@@ -917,7 +917,7 @@ func (s *OpenAIGatewayService) handleAnthropicStreamingResponse(
// two model streams together. Surface a proper Anthropic error event
// instead of returning a failover error that the handler cannot retry.
if !clientOutputStarted && openAIStreamFailedEventShouldFailover(payloadBytes, message) {
streamFailoverErr = s.newOpenAIStreamFailoverError(c, account, false, requestID, payloadBytes, message)
streamFailoverErr = s.newOpenAIStreamFailoverError(c, account, false, requestID, payloadBytes, message, resp.Header)
return true
}
message = s.recordOpenAIStreamUpstreamError(c, account, false, requestID, "http_error", payloadBytes, message)
@@ -738,7 +738,13 @@ func openAIStreamClientOutputStarted(c *gin.Context, localStarted bool) bool {
if localStarted {
return true
}
return c != nil && c.Writer != nil && c.Writer.Written()
if c == nil || c.Writer == nil {
return false
}
// compact keepalive comments commit the HTTP response as 200, but they are
// not semantic model output and therefore must not block a safe retry.
// Without a compact keepalive this is equivalent to checking Writer.Size().
return OpenAICompactKeepaliveAdjustedWrittenSize(c) >= 0
}
func openAIStreamEventIsPreamble(eventType string) bool {
@@ -776,10 +782,10 @@ func openAIStreamFailedEventSemanticStatus(payload []byte, message string) int {
}
combined := strings.TrimSpace(errType + " " + code + " " + strings.ToLower(strings.TrimSpace(message)))
switch {
case strings.Contains(errType, "invalid_request"):
return http.StatusBadRequest
case strings.Contains(combined, "rate_limit"):
return http.StatusTooManyRequests
case strings.Contains(errType, "invalid_request"):
return http.StatusBadRequest
case strings.Contains(combined, "authentication") || strings.Contains(combined, "unauthorized") || strings.Contains(combined, "invalid_api_key"):
return http.StatusUnauthorized
case strings.Contains(combined, "permission") || strings.Contains(combined, "forbidden") || strings.Contains(combined, "access denied"):
@@ -791,6 +797,19 @@ func openAIStreamFailedEventSemanticStatus(payload []byte, message string) int {
}
}
func openAIStreamFailureStatus(payload []byte, message string) int {
if len(bytes.TrimSpace(payload)) == 0 || !gjson.ValidBytes(payload) {
return http.StatusBadGateway
}
// Keep the existing 502 failover behavior for other response.failed events.
// Only rate limits need promotion because they participate in the account's
// configurable 429 same-account retry policy.
if openAIStreamFailedEventSemanticStatus(payload, message) == http.StatusTooManyRequests {
return http.StatusTooManyRequests
}
return http.StatusBadGateway
}
func openAIStreamFailedEventPassthroughBody(payload []byte, failedMessage string) []byte {
if len(payload) == 0 || !gjson.ValidBytes(payload) {
return payload
@@ -867,6 +886,12 @@ func openAIStreamFailedEventShouldFailover(payload []byte, message string) bool
if isOpenAIContextWindowError(message, payload) {
return false
}
// A response.failed event is transported over HTTP 200. Prefer its semantic
// rate-limit status over a generic/invalid_request error type so it can enter
// the same 429 retry policy as a regular upstream HTTP response.
if openAIStreamFailureStatus(payload, message) == http.StatusTooManyRequests {
return true
}
if isOpenAITransientProcessingError(http.StatusBadRequest, message, payload) {
return true
}
@@ -921,6 +946,7 @@ func (s *OpenAIGatewayService) recordOpenAIStreamUpstreamError(
if message == "" {
message = "OpenAI upstream response failed"
}
statusCode := openAIStreamFailureStatus(payload, message)
detail := ""
if len(payload) > 0 && s != nil && s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
@@ -930,10 +956,10 @@ func (s *OpenAIGatewayService) recordOpenAIStreamUpstreamError(
detail = truncateString(string(payload), maxBytes)
}
if c != nil {
setOpsUpstreamError(c, http.StatusBadGateway, message, detail)
setOpsUpstreamError(c, statusCode, message, detail)
event := OpsUpstreamErrorEvent{
Platform: PlatformOpenAI,
UpstreamStatusCode: http.StatusBadGateway,
UpstreamStatusCode: statusCode,
UpstreamRequestID: strings.TrimSpace(upstreamRequestID),
Passthrough: passthrough,
Kind: kind,
@@ -957,21 +983,35 @@ func (s *OpenAIGatewayService) newOpenAIStreamFailoverError(
upstreamRequestID string,
payload []byte,
message string,
responseHeaders ...http.Header,
) *UpstreamFailoverError {
message = sanitizeUpstreamErrorMessage(strings.TrimSpace(message))
if message == "" {
message = "OpenAI stream disconnected before completion"
}
statusCode := openAIStreamFailureStatus(payload, message)
var headers http.Header
if len(responseHeaders) > 0 && responseHeaders[0] != nil {
headers = responseHeaders[0].Clone()
}
// 流内 failed 事件承载于 HTTP 200,响应头是正常配额快照而非限流信号,
// 不写账号级限流/封禁状态;重试与切号由 failover 引擎按
// StatusCode/RetryableOnSameAccount 决定。
message = s.recordOpenAIStreamUpstreamError(c, account, passthrough, upstreamRequestID, "failover", payload, message)
errType := "upstream_error"
if statusCode == http.StatusTooManyRequests {
errType = "rate_limit_error"
}
body, _ := json.Marshal(gin.H{
"error": gin.H{
"type": "upstream_error",
"type": errType,
"message": message,
},
})
return &UpstreamFailoverError{
StatusCode: http.StatusBadGateway,
StatusCode: statusCode,
ResponseBody: body,
ResponseHeaders: headers,
RetryableOnSameAccount: openAIStreamFailedEventRetryableOnSameAccount(account, payload, message),
}
}
@@ -1127,7 +1167,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough(
}
if openAIStreamFailedEventShouldFailover(dataBytes, failedMessage) {
return resultWithUsage(),
s.newOpenAIStreamFailoverError(c, account, true, upstreamRequestID, dataBytes, failedMessage)
s.newOpenAIStreamFailoverError(c, account, true, upstreamRequestID, dataBytes, failedMessage, resp.Header)
}
}
forceFlushFailedEvent = true
@@ -449,7 +449,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context.
}
if openAIStreamFailedEventShouldFailover(dataBytes, failedMessage) {
sawFailedEvent = true
streamEarlyErr = s.newOpenAIStreamFailoverError(c, account, false, upstreamRequestID, dataBytes, failedMessage)
streamEarlyErr = s.newOpenAIStreamFailoverError(c, account, false, upstreamRequestID, dataBytes, failedMessage, resp.Header)
return
}
}
@@ -1577,6 +1577,117 @@ func TestOpenAIStreamingResponseFailedBeforeOutputServerOverloadedCodeReturnsFai
require.Empty(t, rec.Body.String())
}
func TestOpenAIStreamingResponseFailedBeforeOutputRateLimitUsesPoolRetryPolicy(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{
Gateway: config.GatewayConfig{
StreamDataIntervalTimeout: 0,
StreamKeepaliveInterval: 0,
MaxLineSize: defaultMaxLineSize,
},
}
svc := &OpenAIGatewayService{cfg: cfg}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
"event: response.created",
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
"",
"event: response.failed",
`data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
"",
}, "\n"))),
Header: http.Header{
"X-Request-Id": []string{"rid-rate-limit-failed"},
"Retry-After": []string{"1"},
},
}
account := &Account{
ID: 1,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Name: "pool-account",
Credentials: map[string]any{
"pool_mode": true,
"pool_mode_retry_count": float64(1),
"pool_mode_retry_status_codes": []any{float64(http.StatusTooManyRequests)},
},
}
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
require.Error(t, err)
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
require.True(t, failoverErr.RetryableOnSameAccount)
require.Equal(t, "1", failoverErr.ResponseHeaders.Get("Retry-After"))
require.Equal(t, "rate_limit_error", gjson.GetBytes(failoverErr.ResponseBody, "error.type").String())
require.Contains(t, string(failoverErr.ResponseBody), "Concurrency limit exceeded")
require.False(t, c.Writer.Written())
require.Empty(t, rec.Body.String())
opsVal, ok := c.Get(OpsUpstreamErrorsKey)
require.True(t, ok)
opsEvents, ok := opsVal.([]*OpsUpstreamErrorEvent)
require.True(t, ok)
require.NotEmpty(t, opsEvents)
require.Equal(t, http.StatusTooManyRequests, opsEvents[len(opsEvents)-1].UpstreamStatusCode)
}
// 流内 rate limit 只产生 failover 错误,不写账号级限流/封禁状态:
// HTTP 200 流的 x-codex-* 头是正常配额快照,不能按 429 头驱动账号冷却。
func TestOpenAIStreamingResponseFailedRateLimitDoesNotBlockAccountScheduling(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{
Gateway: config.GatewayConfig{
StreamDataIntervalTimeout: 0,
StreamKeepaliveInterval: 0,
MaxLineSize: defaultMaxLineSize,
},
}
svc := &OpenAIGatewayService{cfg: cfg}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
"event: response.created",
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
"",
"event: response.failed",
`data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
"",
}, "\n"))),
Header: http.Header{
"X-Codex-Primary-Used-Percent": []string{"12"},
"X-Codex-Primary-Reset-After-Seconds": []string{"604800"},
"Retry-After": []string{"1"},
},
}
account := &Account{
ID: 11,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Name: "oauth-account",
}
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
require.Error(t, err)
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
require.False(t, failoverErr.RetryableOnSameAccount)
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
}
func TestOpenAIStreamingResponseFailedAfterOutputSanitizesVerboseResponseForClient(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{
@@ -43,9 +43,18 @@ type UsageRecordSubmitMode string
const (
UsageRecordSubmitModeEnqueued UsageRecordSubmitMode = "enqueued"
UsageRecordSubmitModeDropped UsageRecordSubmitMode = "dropped"
UsageRecordSubmitModeSync UsageRecordSubmitMode = "sync_fallback"
// UsageRecordSubmitModeDroppedStopped 表示任务因池已停止(进程关停窗口)被丢弃。
// 与显式 drop/sample 溢出策略的丢弃区分开:溢出丢弃是运维显式配置的取舍,
// 而关停窗口丢弃不是,计费关键任务应在调用侧降级为同步执行兜底。
UsageRecordSubmitModeDroppedStopped UsageRecordSubmitMode = "dropped_stopped"
UsageRecordSubmitModeSync UsageRecordSubmitMode = "sync_fallback"
)
// Dropped 报告任务是否未被执行(入队失败且未同步执行)。
func (m UsageRecordSubmitMode) Dropped() bool {
return m == UsageRecordSubmitModeDropped || m == UsageRecordSubmitModeDroppedStopped
}
// UsageRecordWorkerPoolOptions 使用量记录池配置。
type UsageRecordWorkerPoolOptions struct {
WorkerCount int
@@ -150,7 +159,7 @@ func (p *UsageRecordWorkerPool) Submit(task UsageRecordTask) UsageRecordSubmitMo
if p.pool == nil || p.pool.Stopped() {
p.droppedPoolStopped.Add(1)
p.logDrop("stopped")
return UsageRecordSubmitModeDropped
return UsageRecordSubmitModeDroppedStopped
}
_, ok := p.pool.TrySubmit(func() {
@@ -163,7 +172,7 @@ func (p *UsageRecordWorkerPool) Submit(task UsageRecordTask) UsageRecordSubmitMo
if p.pool.Stopped() {
p.droppedPoolStopped.Add(1)
p.logDrop("stopped")
return UsageRecordSubmitModeDropped
return UsageRecordSubmitModeDroppedStopped
}
switch p.overflowPolicy {
@@ -176,7 +176,11 @@ func TestUsageRecordWorkerPool_SubmitAfterStop(t *testing.T) {
pool.Stop()
mode := pool.Submit(func(ctx context.Context) {})
require.Equal(t, UsageRecordSubmitModeDropped, mode)
require.Equal(t, UsageRecordSubmitModeDroppedStopped, mode)
require.True(t, mode.Dropped())
require.True(t, UsageRecordSubmitModeDropped.Dropped())
require.False(t, UsageRecordSubmitModeEnqueued.Dropped())
require.False(t, UsageRecordSubmitModeSync.Dropped())
require.GreaterOrEqual(t, pool.Stats().DroppedPoolStopped, uint64(1))
}