mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 03:52:26 +08:00
Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
21aacde0b3 | ||
|
|
b22f73e725 | ||
|
|
d6d53052f8 | ||
|
|
bd52e5d770 | ||
|
|
d4cada3b6b | ||
|
|
8f5caef78e | ||
|
|
85a27fae39 |
@@ -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.
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user