fix(account): 计数、Bedrock 与重试耗尽入口把上游模型交给错误处理

OpenAI 与 Anthropic 的 count_tokens / input_tokens 入口、Bedrock 故障转移入口
以及 OAuth 403 重试耗尽入口调 HandleUpstreamError 时没有传入实际发往上游的模型,
显式 scope=model 的 401 规则在这些入口看到空模型,仍落到 handleAuthError 整号置 error。

- 四处计数入口传入映射后的上游模型;Anthropic 透传计数函数新增 upstreamModel 参数
- handleRetryExhaustedSideEffects / handleRetryExhaustedError 增加模型变参,
  Forward、Anthropic 透传与 Bedrock 三处调用方传入实际模型
- Bedrock 入口传调度器的模型键(account.GetMappedModel),不是发往上游的
  Bedrock 模型 ID,模型级冷却才会被调度器读取
- 这些入口的 404 模型未找到与缺省 scope 的非 401 规则随之按主转发路径语义处理
  (模型已知走模型级)
This commit is contained in:
feeeei
2026-10-13 00:53:39 +08:00
parent fc82e83102
commit 2023c290e1
7 changed files with 205 additions and 21 deletions
@@ -184,7 +184,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
logger.LegacyPrintf("service.gateway", "[Anthropic Passthrough] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d RequestID=%s Body=%s",
account.ID, account.Name, resp.StatusCode, resp.Header.Get("x-request-id"), truncateString(string(respBody), 1000))
s.handleRetryExhaustedSideEffects(ctx, resp, account)
s.handleRetryExhaustedSideEffects(ctx, resp, account, input.RequestModel)
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
ProxyID: opsUpstreamProxyID(account),
ProxyName: opsUpstreamProxyName(account),
@@ -209,7 +209,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
return s.handleRetryExhaustedError(ctx, resp, c, account)
return s.handleRetryExhaustedError(ctx, resp, c, account, input.RequestModel)
}
if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) {
+10 -6
View File
@@ -130,7 +130,9 @@ func (s *GatewayService) forwardBedrock(
// 错误/failover 处理
if resp.StatusCode >= 400 {
return s.handleBedrockUpstreamErrors(ctx, resp, c, account)
// 账号状态按模型隔离时用调度器的模型键(账号映射后的模型名),不是发往上游的 Bedrock 模型 ID;
// 两者不一致会让模型级冷却写到调度器不会读取的键上。
return s.handleBedrockUpstreamErrors(ctx, resp, c, account, account.GetMappedModel(reqModel))
}
// Bedrock 分支绕过通用 Forward 成功路径,这里保持上游接受回调语义一致。
@@ -263,12 +265,14 @@ func (s *GatewayService) executeBedrockUpstream(
return resp, nil
}
// handleBedrockUpstreamErrors 处理 Bedrock 上游 4xx/5xx 错误(failover + 错误响应)
// handleBedrockUpstreamErrors 处理 Bedrock 上游 4xx/5xx 错误(failover + 错误响应);
// modelKey 是调度器按模型隔离账号状态时使用的模型键(account.GetMappedModel 的结果)。
func (s *GatewayService) handleBedrockUpstreamErrors(
ctx context.Context,
resp *http.Response,
c *gin.Context,
account *Account,
modelKey string,
) (*ForwardResult, error) {
// retry exhausted + failover
if s.shouldRetryUpstreamError(account, resp.StatusCode) {
@@ -280,7 +284,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors(
logger.LegacyPrintf("service.gateway", "[Bedrock] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d Body=%s",
account.ID, account.Name, resp.StatusCode, truncateString(string(respBody), 1000))
s.handleRetryExhaustedSideEffects(ctx, resp, account)
s.handleRetryExhaustedSideEffects(ctx, resp, account, modelKey)
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
ProxyID: opsUpstreamProxyID(account),
ProxyName: opsUpstreamProxyName(account),
@@ -297,7 +301,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors(
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
return s.handleRetryExhaustedError(ctx, resp, c, account)
return s.handleRetryExhaustedError(ctx, resp, c, account, modelKey)
}
// non-retryable failover
@@ -306,7 +310,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors(
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
s.handleFailoverSideEffects(ctx, resp, account)
s.handleFailoverSideEffects(ctx, resp, account, modelKey)
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
ProxyID: opsUpstreamProxyID(account),
ProxyName: opsUpstreamProxyName(account),
@@ -325,7 +329,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors(
}
// other errors
return s.handleErrorResponse(ctx, resp, c, account)
return s.handleErrorResponse(ctx, resp, c, account, modelKey)
}
// buildUpstreamRequestBedrock 构建 Bedrock 上游请求
@@ -42,13 +42,15 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() {
passthroughBody := parsed.Body.Bytes()
passthroughModel := parsed.Model
if reqModel := parsed.Model; reqModel != "" {
if mappedModel := account.GetMappedModel(reqModel); mappedModel != reqModel {
passthroughBody = s.replaceModelInBody(passthroughBody, mappedModel)
passthroughModel = mappedModel
logger.LegacyPrintf("service.gateway", "CountTokens passthrough model mapping: %s -> %s (account: %s)", reqModel, mappedModel, account.Name)
}
}
return s.forwardCountTokensAnthropicAPIKeyPassthrough(ctx, c, account, passthroughBody)
return s.forwardCountTokensAnthropicAPIKeyPassthrough(ctx, c, account, passthroughBody, passthroughModel)
}
// Bedrock 不支持 count_tokens 端点
@@ -220,8 +222,8 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
// 处理错误响应
if resp.StatusCode >= 400 {
// 标记账号状态(429/529等)
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
// 标记账号状态(429/529等);传入实际发往上游的模型,临时不可调度规则才能按模型隔离
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, reqModel)
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
@@ -267,7 +269,9 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
return nil
}
func (s *GatewayService) forwardCountTokensAnthropicAPIKeyPassthrough(ctx context.Context, c *gin.Context, account *Account, body []byte) error {
// forwardCountTokensAnthropicAPIKeyPassthrough 透传 count_tokens 请求;upstreamModel 是 body 中实际发往上游的模型,
// 仅用于上游错误处理按模型隔离。
func (s *GatewayService) forwardCountTokensAnthropicAPIKeyPassthrough(ctx context.Context, c *gin.Context, account *Account, body []byte, upstreamModel string) error {
token, tokenType, err := s.GetAccessToken(ctx, account)
if err != nil {
s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to get access token")
@@ -322,7 +326,7 @@ func (s *GatewayService) forwardCountTokensAnthropicAPIKeyPassthrough(ctx contex
if resp.StatusCode >= 400 {
if s.rateLimitService != nil {
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel)
}
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
+2 -2
View File
@@ -689,7 +689,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d RequestID=%s Body=%s",
account.ID, account.Name, resp.StatusCode, resp.Header.Get("x-request-id"), truncateString(string(respBody), 1000))
s.handleRetryExhaustedSideEffects(ctx, resp, account)
s.handleRetryExhaustedSideEffects(ctx, resp, account, reqModel)
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
ProxyID: opsUpstreamProxyID(account),
ProxyName: opsUpstreamProxyName(account),
@@ -713,7 +713,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
return s.handleRetryExhaustedError(ctx, resp, c, account)
return s.handleRetryExhaustedError(ctx, resp, c, account, reqModel)
}
// 处理可切换账号的错误
@@ -527,13 +527,19 @@ func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Res
return nil, fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, upstreamMsg)
}
func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, resp *http.Response, account *Account) {
// handleRetryExhaustedSideEffects 在重试耗尽后标记账号状态;requestedModel 是实际发往上游的模型,
// 临时不可调度规则据此按模型隔离。
func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) {
body, _ := s.readUpstreamErrorBody(resp)
statusCode := resp.StatusCode
// OAuth/Setup Token 账号的 403:标记账号异常
if account.IsOAuth() && statusCode == 403 {
s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, resp.Header, body)
if len(requestedModel) > 0 {
s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, resp.Header, body, requestedModel[0])
} else {
s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, resp.Header, body)
}
logger.LegacyPrintf("service.gateway", "Account %d: marked as error after %d retries for status %d", account.ID, maxRetryAttempts, statusCode)
} else {
// API Key 未配置错误码:不标记账号状态
@@ -553,14 +559,14 @@ func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *ht
// handleRetryExhaustedError 处理重试耗尽后的错误
// OAuth 403:标记账号异常
// API Key 未配置错误码:仅返回错误,不标记账号
func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account) (*ForwardResult, error) {
func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, requestedModel ...string) (*ForwardResult, error) {
MarkResponseCommitted(c)
// Capture upstream error body before side-effects consume the stream.
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
s.handleRetryExhaustedSideEffects(ctx, resp, account)
s.handleRetryExhaustedSideEffects(ctx, resp, account, requestedModel...)
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
@@ -104,7 +104,7 @@ func (s *OpenAIGatewayService) ForwardResponsesInputTokens(
return nil
}
if s.rateLimitService != nil {
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, prepared.UpstreamModel)
}
upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, "")
@@ -346,7 +346,7 @@ func (s *OpenAIGatewayService) ForwardCountTokensAsAnthropic(
}
if s.rateLimitService != nil {
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, prepared.UpstreamModel)
}
if isOpenAIInputTokensUnsupported(resp.StatusCode, respBody) {
@@ -3,14 +3,20 @@
package service
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
type tempUnschedScopeRepoStub struct {
@@ -321,3 +327,167 @@ func TestCheckErrorPolicy_401FollowsTempUnschedScope(t *testing.T) {
})
}
}
// 计数入口把实际发往上游的模型交给错误处理:显式模型级 401 规则只冷却该模型,不把整号置为 error。
func TestCountTokensEntries_ExplicitModelScope401OnlyCoolsModel(t *testing.T) {
gin.SetMode(gin.TestMode)
unauthorized := func() *http.Response {
return &http.Response{
StatusCode: http.StatusUnauthorized,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"unauthorized"}}`)),
}
}
cfg := &config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false, AllowInsecureHTTP: true}},
}
newContext := func(path string, body []byte) (*gin.Context, *httptest.ResponseRecorder) {
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, path, bytes.NewReader(body))
return c, rec
}
t.Run("openai count_tokens as anthropic", func(t *testing.T) {
repo := &tempUnschedScopeRepoStub{}
upstream := &httpUpstreamRecorder{resp: unauthorized()}
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream, rateLimitService: &RateLimitService{accountRepo: repo, cfg: cfg}}
account := newTempUnschedScope401Account(PlatformOpenAI, TempUnschedScopeModel)
account.Credentials["api_key"] = "sk-test"
account.Credentials["base_url"] = "http://upstream.example"
account.Concurrency = 1
body := []byte(`{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}]}`)
c, _ := newContext("/v1/messages/count_tokens", body)
require.Error(t, svc.ForwardCountTokensAsAnthropic(context.Background(), c, account, body, "gpt-5.4"))
require.Equal(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String())
require.Zero(t, repo.errorCalls)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope)
})
t.Run("openai responses input_tokens", func(t *testing.T) {
repo := &tempUnschedScopeRepoStub{}
upstream := &httpUpstreamRecorder{resp: unauthorized()}
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream, rateLimitService: &RateLimitService{accountRepo: repo, cfg: cfg}}
account := newTempUnschedScope401Account(PlatformOpenAI, TempUnschedScopeModel)
account.Credentials["api_key"] = "sk-test"
account.Credentials["model_mapping"] = map[string]any{"gpt-5.4": "gpt-5.4-upstream"}
account.Concurrency = 1
body := []byte(`{"model":"gpt-5.4","input":"hello"}`)
c, _ := newContext("/v1/responses/input_tokens", body)
require.Error(t, svc.ForwardResponsesInputTokens(context.Background(), c, account, body))
require.NotNil(t, upstream.lastReq)
require.Equal(t, "gpt-5.4-upstream", gjson.GetBytes(upstream.lastBody, "model").String())
require.Zero(t, repo.errorCalls)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "gpt-5.4-upstream", repo.modelRateLimitCalls[0].scope)
})
for _, passthrough := range []bool{false, true} {
name := "anthropic count_tokens"
if passthrough {
name = "anthropic count_tokens api key passthrough"
}
t.Run(name, func(t *testing.T) {
repo := &tempUnschedScopeRepoStub{}
upstream := &anthropicHTTPUpstreamRecorder{resp: unauthorized()}
svc := &GatewayService{
cfg: cfg,
responseHeaderFilter: compileResponseHeaderFilter(cfg),
httpUpstream: upstream,
rateLimitService: &RateLimitService{accountRepo: repo, cfg: cfg},
}
account := newTempUnschedScope401Account(PlatformAnthropic, TempUnschedScopeModel)
account.Credentials["api_key"] = "upstream-anthropic-key"
account.Credentials["base_url"] = "https://api.anthropic.com"
account.Credentials["model_mapping"] = map[string]any{"claude-sonnet-4-5": "claude-sonnet-4-5-upstream"}
account.Concurrency = 1
if passthrough {
account.Extra = map[string]any{"anthropic_passthrough": true}
}
body := []byte(`{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}]}`)
c, _ := newContext("/v1/messages/count_tokens", body)
parsed := &ParsedRequest{Body: NewRequestBodyRef(body), Model: "claude-sonnet-4-5"}
require.Error(t, svc.ForwardCountTokens(context.Background(), c, account, parsed))
require.NotNil(t, upstream.lastReq)
require.Equal(t, "claude-sonnet-4-5-upstream", gjson.GetBytes(upstream.lastBody, "model").String())
require.Zero(t, repo.errorCalls)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "claude-sonnet-4-5-upstream", repo.modelRateLimitCalls[0].scope)
})
}
}
// Bedrock 故障转移入口与重试耗尽入口把实际发往上游的模型交给错误处理:显式模型级规则只冷却该模型。
func TestAnthropicFailoverEntries_ExplicitModelScopeOnlyCoolsModel(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}
newResponse := func(status int, body string) *http.Response {
return &http.Response{
StatusCode: status,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(body)),
}
}
t.Run("bedrock failover 401", func(t *testing.T) {
repo := &tempUnschedScopeRepoStub{}
upstream := &anthropicHTTPUpstreamRecorder{resp: newResponse(http.StatusUnauthorized, `{"message":"unauthorized"}`)}
svc := &GatewayService{cfg: cfg, httpUpstream: upstream, rateLimitService: &RateLimitService{accountRepo: repo, cfg: cfg}}
account := newTempUnschedScope401Account(PlatformAnthropic, TempUnschedScopeModel)
account.Type = AccountTypeBedrock
account.Status = StatusActive
account.Schedulable = true
account.Credentials["auth_mode"] = "apikey"
account.Credentials["api_key"] = "bedrock-key"
body := []byte(`{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}]}`)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
parsed := &ParsedRequest{Body: NewRequestBodyRef(body), Model: "claude-sonnet-4-5"}
_, err := svc.forwardBedrock(context.Background(), c, account, parsed, time.Now())
var failover *UpstreamFailoverError
require.ErrorAs(t, err, &failover)
require.NotNil(t, upstream.lastReq)
require.Contains(t, upstream.lastReq.URL.Path, "anthropic.claude-sonnet-4-5-", "上游收到 Bedrock 模型 ID")
require.Zero(t, repo.errorCalls)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
// 冷却键必须是调度器查找的键(账号映射后的模型名),而不是 Bedrock 模型 ID。
scope := repo.modelRateLimitCalls[0].scope
require.Equal(t, "claude-sonnet-4-5", scope)
require.True(t, account.IsSchedulableForModelWithContext(context.Background(), "claude-sonnet-4-5"))
setAccountModelRateLimitSnapshot(account, scope, time.Now().Add(10*time.Minute), "temp_unschedulable", time.Now())
require.False(t, account.IsSchedulableForModelWithContext(context.Background(), "claude-sonnet-4-5"), "调度器按同一个键排除该模型")
})
t.Run("oauth retry exhausted 403", func(t *testing.T) {
repo := &tempUnschedScopeRepoStub{}
svc := &GatewayService{cfg: cfg, rateLimitService: &RateLimitService{accountRepo: repo, cfg: cfg}}
account := newTempUnschedScopeAccount(TempUnschedScopeModel, map[string]any{
"error_code": float64(http.StatusForbidden),
"keywords": []any{"forbidden"},
})
account.Platform = PlatformAnthropic
account.Type = AccountTypeOAuth
svc.handleRetryExhaustedSideEffects(context.Background(), newResponse(http.StatusForbidden, `{"error":{"message":"forbidden"}}`), account, "claude-sonnet-4-5")
require.Zero(t, repo.errorCalls)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].scope)
})
}