Compare commits

...
Author SHA1 Message Date
shaw 21aacde0b3 fix(openai-ws): keep downstream writes off the relay cancellation context
coder/websocket arms a context.AfterFunc that hard-closes the connection
when a write context is canceled, and AfterFunc stop does not wait for a
callback that already started. An external cancellation (e.g. ingress
lease loss) landing inside the disarm window of an already-successful
downstream write could therefore kill the TCP connection before the
retry close frame (1013) was written, leaving the client with a bare
EOF. Mirror the read side: bound downstream writes with the write
timeout only, and rely on the explicit Close/CloseNow performed by every
relay exit path for teardown.

Fixes the flaky TestPassthroughLifecycle_LeaseLossSendsRetryClose.
2026-08-01 20:32:10 +08:00
Wesley Liddick b22f73e725 Merge pull request #5154 from Wei-Shaw/fix/issue-5148-stream-partial-usage-billing
fix(gateway): 流中断时保留已观测 usage 入账,修复 newapi 类上游大面积漏记(#5148)
2026-08-01 13:51:06 +08:00
shaw d6d53052f8 chore: update sponsors 2026-08-01 11:32:31 +08:00
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
20 changed files with 754 additions and 96 deletions
-6
View File
@@ -84,12 +84,6 @@ Please read the following carefully before using this project:
</td>
</tr>
<tr>
<td width="180"><a href="https://unity2.ai/register?source=sub2api"><img src="assets/partners/logos/unity2.png" alt="unity2" width="150"></a></td>
<td>Thanks to Unity2 for sponsoring this project! <a href="https://unity2.ai/register?source=sub2api">Unity2</a> is a high-performance AI model API relay for individuals, teams, and enterprises, handling 30B+ tokens/day with 5000 RPM concurrency. One API Key works across Claude Code, Codex, OpenAI models, IDE plugins, and Agent workflows, with balance billing, bundled subscriptions, enterprise invoicing, and 1-on-1 support. <a href="https://unity2.ai/register?source=sub2api">Register</a> to claim $2 in balance, plus $10 more by joining the official group — up to $12 in free credit.
</td>
</tr>
<tr>
<td width="180"><a href="https://veilx.io/#/hello/SJRBRVDV"><img src="assets/partners/logos/veilx.png" alt="veilx" width="150"></a></td>
<td>Thanks to Veilx for sponsoring this project! <a href="https://veilx.io/#/hello/SJRBRVDV">Veilx</a> CDN is purpose-built for large-scale AI API traffic, deeply optimized for relay services and call chains across OpenAI, Claude, Gemini, and scenarios like chat, image generation, embeddings, and streaming — delivering lower latency and higher stability under heavy concurrency. It also offers China three-network optimized return lines, making it ideal for global AI relay platforms, overseas AI SaaS, and cross-border high-concurrency deployments.
-6
View File
@@ -87,12 +87,6 @@
</td>
</tr>
<tr>
<td width="180"><a href="https://unity2.ai/register?source=sub2api"><img src="assets/partners/logos/unity2.png" alt="unity2" width="150"></a></td>
<td>感谢 Unity2 赞助本项目! <a href="https://unity2.ai/register?source=sub2api">Unity2</a> 是面向个人开发者、团队、企业的高性能 AI 模型 API 中转平台,长期服务国内头部企业,日均承载超 300 亿 token 调用,支持 5000 RPM 级高并发。一个 API Key 即可适配 Claude Code、Codex、OpenAI 模型、IDE 插件和 Agent 工作流等场景。具备企业级稳定供应能力,在高并发、持续调用和团队集中采购场景下依然保持低延迟、高可用。同时支持余额计费、组合订阅、首充优惠、企业开票、专属 1v1 对接,适合个人高频使用和企业长期接入。现在注册 Unity2.ai 可领取 $2 余额,加入官方群再送 $10 余额,合计最高可领 $12 免费额度,适合先体验后长期使用。<a href="https://unity2.ai/register?source=sub2api">注册链接</a>
</td>
</tr>
<tr>
<td width="180"><a href="https://veilx.io/#/hello/SJRBRVDV"><img src="assets/partners/logos/veilx.png" alt="veilx" width="150"></a></td>
<td>感谢 Veilx 赞助本项目! <a href="https://veilx.io/#/hello/SJRBRVDV">Veilx</a> CDN 专为超大规模 API 请求场景打造,针对 AI 中转站业务与 AI API 调用链路进行了深度优化,轻松应对高并发、高频请求与大流量传输,为开发者与企业提供更快、更稳、更低延迟的加速体验。无论是 OpenAI、Claude、Gemini 等 AI 接口中转,还是聊天、绘图、Embedding、流式输出等复杂场景,Veilx 都能显著提升响应速度与连接稳定性,有效降低网络波动带来的超时与失败问题。同时,Veilx 提供中国三网优化回国极速线路,大幅提升中国大陆地区访问海外 AI 服务的速度与稳定性,特别适合全球 AI 中转平台、海外 AI SaaS、跨境业务与高并发 API 系统部署。专为 AI API 而生,让你的 AI 中转服务更快、更稳、更省心。<a href="https://veilx.io/#/hello/SJRBRVDV">购买地址</a>
-6
View File
@@ -86,12 +86,6 @@
</td>
</tr>
<tr>
<td width="180"><a href="https://unity2.ai/register?source=sub2api"><img src="assets/partners/logos/unity2.png" alt="unity2" width="150"></a></td>
<td>Unity2 のご支援に感謝します!<a href="https://unity2.ai/register?source=sub2api">Unity2</a> は個人開発者、チーム、企業向けの高性能 AI モデル API 中継プラットフォームです。中国の大手企業に長期にわたりサービスを提供しており、1日あたり 300 億以上のトークン呼び出しを処理し、5000 RPM 級の高並列性をサポートします。1つの API キーで Claude Code、Codex、OpenAI モデル、IDE プラグイン、Agent ワークフローなど様々なシナリオに対応できます。エンタープライズグレードの安定供給能力を備え、高並列・継続的な呼び出し・チームの集中購入シーンでも低レイテンシと高可用性を維持します。残高課金、組み合わせサブスクリプション、初回チャージ特典、企業向け請求書発行、専属 1v1 サポートにも対応しており、個人の頻繁な利用にも企業の長期導入にも適しています。今 Unity2.ai に登録すると $2 の残高、公式グループに参加するとさらに $10 の残高がもらえ、合計最大 $12 の無料クレジットを獲得できます — 試用後に長期利用したい方に最適です。<a href="https://unity2.ai/register?source=sub2api">登録リンク</a>
</td>
</tr>
<tr>
<td width="180"><a href="https://veilx.io/#/hello/SJRBRVDV"><img src="assets/partners/logos/veilx.png" alt="veilx" width="150"></a></td>
<td>Veilx のご支援に感謝します!<a href="https://veilx.io/#/hello/SJRBRVDV">Veilx</a> CDN は超大規模 API リクエストシナリオ向けに設計されており、AI 中継サービスと AI API 呼び出しチェーンに対して深く最適化されています。高並列・高頻度リクエスト・大容量トラフィックに容易に対応し、開発者と企業により高速で安定した、低レイテンシの加速体験を提供します。OpenAI、Claude、Gemini などの AI インターフェース中継はもちろん、チャット、画像生成、Embedding、ストリーミング出力などの複雑なシナリオでも、Veilx は応答速度と接続安定性を大幅に向上させ、ネットワーク変動によるタイムアウトや失敗を効果的に削減します。さらに、Veilx は中国三大ネットワーク最適化の高速回線を提供しており、中国本土から海外 AI サービスへのアクセス速度と安定性を大幅に向上させます。グローバル AI 中継プラットフォーム、海外 AI SaaS、越境ビジネス、高並列 API システム展開に特に適しています。AI API のために生まれ、あなたの AI 中継サービスをより速く、より安定して、より安心に。<a href="https://veilx.io/#/hello/SJRBRVDV">購入リンク</a>
+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{
@@ -173,7 +173,13 @@ func Relay(
return writeUpstream(msgType, payload)
}
writeClient := func(msgType coderws.MessageType, payload []byte) error {
writeCtx, cancel := context.WithTimeout(relayCtx, writeTimeout)
// 下行写超时故意不挂在 relayCtx 上:coder/websocket 在已武装的 write
// ctx 被取消时会直接硬关连接(context.AfterFunc 的 stop 不等待执行中
// 的回调),外部取消若落在一次已成功写入的解除武装窗口内,会连同尚未
// 发出的 close 帧一起冲掉,客户端只能看到裸 EOF 而收不到关闭码。与读
// 侧 conn.Read(context.Background()) 同理,取消路径的连接回收由各退出
// 分支的显式 Close/CloseNow 兜底。
writeCtx, cancel := context.WithTimeout(context.Background(), writeTimeout)
defer cancel()
return clientConn.WriteFrame(writeCtx, msgType, payload)
}
@@ -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))
}