From b1e60ba4535ea1fc2cc57a181a92bab03a2e0782 Mon Sep 17 00:00:00 2001 From: myy_321 <1948281933@qq.com> Date: Sun, 16 Aug 2026 08:29:12 +0800 Subject: [PATCH] =?UTF-8?q?fix(gateway):=20=E4=BF=AE=E5=A4=8D=E6=B1=A0?= =?UTF-8?q?=E6=A8=A1=E5=BC=8F=E5=90=8C=E8=B4=A6=E5=8F=B7=E9=94=99=E8=AF=AF?= =?UTF-8?q?=E9=87=8D=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 透传上游错误处理结果，保留池模式重试判定\n- 对可重试状态码设置同账号重试标记\n- 补充 Chat Completions 与 Responses 的 429 回归测试 --- .../gateway_forward_as_chat_completions.go | 8 +- .../service/gateway_forward_as_responses.go | 8 +- .../service/gateway_pool_mode_retry_test.go | 78 +++++++++++++++++++ 3 files changed, 88 insertions(+), 6 deletions(-) create mode 100644 backend/internal/service/gateway_pool_mode_retry_test.go diff --git a/backend/internal/service/gateway_forward_as_chat_completions.go b/backend/internal/service/gateway_forward_as_chat_completions.go index 93da85a7b2..24c4b39b56 100644 --- a/backend/internal/service/gateway_forward_as_chat_completions.go +++ b/backend/internal/service/gateway_forward_as_chat_completions.go @@ -165,12 +165,14 @@ func (s *GatewayService) ForwardAsChatCompletions( Kind: "failover", Message: upstreamMsg, }) + shouldDisable := false if s.rateLimitService != nil { - s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, mappedModel) + shouldDisable = s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, mappedModel) } return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, + StatusCode: resp.StatusCode, + ResponseBody: respBody, + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } diff --git a/backend/internal/service/gateway_forward_as_responses.go b/backend/internal/service/gateway_forward_as_responses.go index 2aa5255c29..a4c4f0670b 100644 --- a/backend/internal/service/gateway_forward_as_responses.go +++ b/backend/internal/service/gateway_forward_as_responses.go @@ -170,12 +170,14 @@ func (s *GatewayService) ForwardAsResponses( Kind: "failover", Message: upstreamMsg, }) + shouldDisable := false if s.rateLimitService != nil { - s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, mappedModel) + shouldDisable = s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, mappedModel) } return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, + StatusCode: resp.StatusCode, + ResponseBody: respBody, + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } diff --git a/backend/internal/service/gateway_pool_mode_retry_test.go b/backend/internal/service/gateway_pool_mode_retry_test.go new file mode 100644 index 0000000000..83455e5b95 --- /dev/null +++ b/backend/internal/service/gateway_pool_mode_retry_test.go @@ -0,0 +1,78 @@ +package service + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestGatewayCompatPoolMode429AllowsSameAccountRetry(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + path string + body []byte + call func(*GatewayService, context.Context, *gin.Context, *Account, []byte) (*ForwardResult, error) + }{ + { + name: "chat completions", + path: "/v1/chat/completions", + body: []byte(`{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}]}`), + call: func(svc *GatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) { + return svc.ForwardAsChatCompletions(ctx, c, account, body, nil) + }, + }, + { + name: "responses", + path: "/v1/responses", + body: []byte(`{"model":"claude-sonnet-4-5","input":"hello"}`), + call: func(svc *GatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) { + return svc.ForwardAsResponses(ctx, c, account, body, nil) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + upstream := &queuedHTTPUpstreamStub{responses: []*http.Response{{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"X-Request-Id": []string{"pool-429"}}, + Body: io.NopCloser(http.NoBody), + }}} + svc := &GatewayService{ + cfg: &config.Config{}, + httpUpstream: upstream, + tlsFPProfileService: &TLSFingerprintProfileService{}, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, tt.path, nil) + account := &Account{ + ID: 1, + Name: "pool-account", + Platform: PlatformAnthropic, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "test-key", + "pool_mode": true, + }, + } + + result, err := tt.call(svc, context.Background(), c, account, tt.body) + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) + require.True(t, failoverErr.RetryableOnSameAccount) + require.Equal(t, 1, upstream.callCount) + require.Empty(t, recorder.Body.String()) + }) + } +}