fix(openai): complete request-scoped capacity recovery

This commit is contained in:
Perfecto
2026-08-15 22:54:21 +08:00
parent c3063e01a4
commit 5390647989
8 changed files with 250 additions and 36 deletions
@@ -55,6 +55,11 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont
if s != nil {
scheduleOllamaCloudUsageActivity(s.deferredService, account)
}
// Capacity shedding describes this request, not account health. Keep the
// account schedulable while the request-local retry budget handles recovery.
if account != nil && account.Platform == PlatformOpenAI && isOpenAIRequestScopedCapacityShed("", responseBody) {
return false
}
stateCtx, cancel := openAIAccountStateContext(ctx)
defer cancel()
@@ -77,6 +77,37 @@ func TestStreamFailedEventCapacityShedRetriesOnSameAccount(t *testing.T) {
require.False(t, openAIStreamFailedEventRetryableOnSameAccount(nonPool, other, "boom"))
}
func TestOpenAIHTTPCapacityShedIsRequestScopedForOAuthAccounts(t *testing.T) {
payload := []byte(`{"error":{"type":"server_error","message":"Our servers are currently overloaded. Please try again later."}}`)
failoverErr := newOpenAIUpstreamFailoverError(
http.StatusBadRequest,
http.Header{"X-Request-Id": []string{"rid-http-capacity"}},
payload,
"Our servers are currently overloaded. Please try again later.",
false,
)
require.True(t, failoverErr.RetryableOnSameAccount)
require.True(t, failoverErr.RequestScopedTransient)
repo := &capacityShedAccountRepoStub{}
(&GatewayService{accountRepo: repo}).TempUnscheduleRetryableError(context.Background(), 1, failoverErr)
require.Zero(t, repo.tempUnschedCalls)
rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
gateway := &OpenAIGatewayService{rateLimitService: rateLimitService}
account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
require.False(t, gateway.handleOpenAIAccountUpstreamError(
context.Background(),
account,
http.StatusBadRequest,
nil,
payload,
"gpt-5",
))
require.Zero(t, repo.tempUnschedCalls)
}
// 上游降载的真实序列是「event: error → event: response.failed」。error 帧不算
// 客户端输出:若把它当首输出 flush,clientOutputStarted 被固化,随后的 failed
// 事件就进不了 pre-output failover 分支,只能把致命错误原样转发给客户端。
@@ -217,6 +248,8 @@ func TestOpenAIStreamCapacityShedErrorFramePrecedingFailedStillFailsOver(t *test
// 并终止会话,对其余错误码执行内置退避重试。消息原样保留。
func TestOpenAIStreamCapacityShedAfterOutputRewritesCodeForClient(t *testing.T) {
gin.SetMode(gin.TestMode)
logSink, restore := captureStructuredLog(t)
defer restore()
cfg := &config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
}
@@ -256,6 +289,9 @@ func TestOpenAIStreamCapacityShedAfterOutputRewritesCodeForClient(t *testing.T)
require.Contains(t, body, `"code":"server_error"`)
require.NotContains(t, body, "server_is_overloaded")
require.Contains(t, body, "Our servers are currently overloaded")
require.True(t, logSink.ContainsMessage("gateway.failover_suppressed_after_semantic_output"))
require.True(t, logSink.ContainsFieldValue("path", "native_sse"))
require.True(t, logSink.ContainsFieldValue("upstream_request_id", "rid-shed-after-output"))
}
// helper 单测:只有降载码被改写,其余错误码(尤其 rate_limit_exceeded,客户端
@@ -1028,6 +1028,27 @@ func isOpenAIUpstreamCapacityShedEvent(payload []byte) bool {
return false
}
func logOpenAICapacityFailoverSuppressed(
ctx context.Context,
account *Account,
path string,
upstreamRequestID string,
eventType string,
) {
fields := []zap.Field{
zap.String("path", path),
zap.String("event_type", strings.TrimSpace(eventType)),
zap.String("upstream_request_id", strings.TrimSpace(upstreamRequestID)),
}
if account != nil {
fields = append(fields,
zap.Int64("account_id", account.ID),
zap.String("platform", account.Platform),
)
}
logger.FromContext(ctx).Warn("gateway.failover_suppressed_after_semantic_output", fields...)
}
// openAICapacityShedRetryableClientCode 是把上游容量降载错误转发给客户端时改写
// 使用的错误码。Codex CLI 按闭集对错误码分类:server_is_overloaded / slow_down
// 被判为致命错误(客户端提示 "Selected model is at capacity. Please try a
@@ -1372,6 +1393,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough(
sawTerminalEvent := false
sawFailedEvent := false
semanticOutputSeen := false
capacityFailoverSuppressedLogged := false
failedMessage := ""
clientOutputStarted := false
upstreamRequestID := strings.TrimSpace(resp.Header.Get("x-request-id"))
@@ -1458,6 +1480,13 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough(
}
}
eventType := strings.TrimSpace(gjson.Get(trimmedData, "type").String())
if !capacityFailoverSuppressedLogged && account != nil && account.Platform == PlatformOpenAI &&
(eventType == "error" || eventType == "response.failed") &&
openAIStreamClientOutputStarted(c, clientOutputStarted) &&
isOpenAIUpstreamCapacityShedEvent(dataBytes) {
logOpenAICapacityFailoverSuppressed(ctx, account, "passthrough_sse", upstreamRequestID, eventType)
capacityFailoverSuppressedLogged = true
}
if eventType == "error" && !openAIStreamClientOutputStarted(c, clientOutputStarted) {
errorMessage := extractOpenAISSEErrorMessage(dataBytes)
if status, errType, errMsg, matched := applyOpenAIStreamFailedErrorPassthroughRule(c, account.Platform, dataBytes, errorMessage); matched {
@@ -242,6 +242,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context.
sawTerminalEvent := false
sawFailedEvent := false
responsesSemanticOutputSeen := false
capacityFailoverSuppressedLogged := false
failedMessage := ""
clientOutputStarted := false
upstreamRequestID := strings.TrimSpace(resp.Header.Get("x-request-id"))
@@ -456,6 +457,13 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context.
responseID = extractOpenAIResponseIDFromJSONBytes(dataBytes)
}
forceFlushFailedEvent := false
if !capacityFailoverSuppressedLogged && account != nil && account.Platform == PlatformOpenAI &&
(eventType == "error" || eventType == "response.failed") &&
openAIStreamClientOutputStarted(c, clientOutputStarted) &&
isOpenAIUpstreamCapacityShedEvent(dataBytes) {
logOpenAICapacityFailoverSuppressed(ctx, account, "native_sse", upstreamRequestID, eventType)
capacityFailoverSuppressedLogged = true
}
if eventType == "error" && !openAIStreamClientOutputStarted(c, clientOutputStarted) {
errorMessage := extractOpenAISSEErrorMessage(dataBytes)
if status, errType, errMsg, matched := applyOpenAIStreamFailedErrorPassthroughRule(c, account.Platform, dataBytes, errorMessage); matched {
@@ -115,7 +115,7 @@ func (r stubOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Co
return r.ListSchedulableByPlatform(ctx, platform)
}
func TestOpenAIGatewayService_ForwardAsAnthropic_TempUnschedulableReturnsFailoverWithoutCommit(t *testing.T) {
func TestOpenAIGatewayService_ForwardAsAnthropic_CapacityShedReturnsRequestScopedFailoverWithoutCommit(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"gpt-5.4","max_tokens":32,"messages":[{"role":"user","content":"hello"}],"stream":false}`)
@@ -173,8 +173,10 @@ func TestOpenAIGatewayService_ForwardAsAnthropic_TempUnschedulableReturnsFailove
require.ErrorAs(t, err, &failoverErr)
require.Equal(t, http.StatusBadRequest, failoverErr.StatusCode)
require.True(t, failoverErr.ShouldRetryNextAccount())
require.Equal(t, account.ID, repo.modelRateLimitAccountID, "temporary unschedulability should exclude this account from reselection")
require.Equal(t, "gpt-5.4", repo.modelRateLimitKey)
require.True(t, failoverErr.RetryableOnSameAccount)
require.True(t, failoverErr.RequestScopedTransient)
require.Zero(t, repo.modelRateLimitAccountID, "request-scoped capacity shedding must not change account health")
require.Empty(t, repo.modelRateLimitKey)
require.False(t, IsResponseCommitted(c))
require.Equal(t, http.StatusOK, rec.Code)
require.Empty(t, rec.Body.String())
@@ -180,6 +180,12 @@ func isOpenAICapacityShedMessage(text string) bool {
strings.Contains(lower, "servers are currently overloaded")
}
func isOpenAIRequestScopedCapacityShed(upstreamMsg string, upstreamBody []byte) bool {
return isOpenAIUpstreamCapacityShedEvent(upstreamBody) ||
isOpenAICapacityShedMessage(upstreamMsg) ||
isOpenAICapacityShedMessage(string(upstreamBody))
}
func isOpenAIContextWindowError(upstreamMsg string, upstreamBody []byte) bool {
match := func(text string) bool {
lower := strings.ToLower(strings.TrimSpace(text))
@@ -264,14 +270,17 @@ func newOpenAIUpstreamFailoverError(
upstreamMsg string,
retryableOnSameAccount bool,
) *UpstreamFailoverError {
requestScopedCapacity := isOpenAIRequestScopedCapacityShed(upstreamMsg, responseBody)
failoverErr := &UpstreamFailoverError{
StatusCode: statusCode,
ResponseBody: responseBody,
ResponseHeaders: responseHeaders.Clone(),
RetryableOnSameAccount: retryableOnSameAccount,
RetryableOnSameAccount: retryableOnSameAccount || requestScopedCapacity,
RequestScopedTransient: requestScopedCapacity,
}
if isOpenAIRequestBodyTooLargeError(statusCode, upstreamMsg, responseBody) {
failoverErr.RetryableOnSameAccount = false
failoverErr.RequestScopedTransient = false
failoverErr.Scope = GatewayFailureScopeAccount
failoverErr.Reason = openAIRequestBodyTooLargeReason
failoverErr.NextAccountAction = NextAccountRetry
@@ -282,6 +282,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
upstreamTerminalEvent := ""
sawDone := false
wroteDownstream := false
pendingClientMessages := make([][]byte, 0, 4)
pendingClientMessageBytes := int64(0)
capacityFailoverSuppressedLogged := false
clientDisconnected := false
mappedModel := ""
needModelReplace := false
@@ -392,15 +395,20 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
replayCollector.AddEvent(eventType, upstreamMessage)
var upstreamEventErr error
if eventType == "error" {
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage)
errMessage := strings.TrimSpace(errMsgRaw)
if eventType == "error" || eventType == "response.failed" {
errMessage := extractOpenAISSEErrorMessage(upstreamMessage)
if errMessage == "" {
errMessage = "upstream error event"
}
statusCode := openAIWSErrorHTTPStatusFromRaw(errCodeRaw, errTypeRaw)
shouldFailover := s.shouldFailoverOpenAIUpstreamResponse(statusCode, errMessage, upstreamMessage)
if account.Platform == PlatformGrok {
statusCode := openAIStreamFailureStatus(upstreamMessage, errMessage)
shouldFailover := openAIStreamFailedEventShouldFailover(upstreamMessage, errMessage)
if eventType == "error" {
errCodeRaw, errTypeRaw, _ := parseOpenAIWSErrorEventFields(upstreamMessage)
statusCode = openAIWSErrorHTTPStatusFromRaw(errCodeRaw, errTypeRaw)
shouldFailover = s.shouldFailoverOpenAIUpstreamResponse(statusCode, errMessage, upstreamMessage)
}
requestScopedCapacity := isOpenAIUpstreamCapacityShedEvent(upstreamMessage)
if account.Platform == PlatformGrok && eventType == "error" {
// SSE error events do not carry an HTTP status. The local status
// mapper therefore defaults unknown xAI codes (for example
// new_sensitive) to 502; classify the body as a request-scoped
@@ -411,7 +419,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
shouldFailover = s.shouldFailoverGrokUpstreamError(statusCode, upstreamMessage)
s.handleGrokAccountUpstreamError(ctx, account, statusCode, resp.Header, upstreamMessage)
}
} else if shouldFailover {
} else if eventType == "error" && shouldFailover && !requestScopedCapacity {
accountStatus := statusCode
if transientStatus := openAIWSPayloadTransientStatus(upstreamMessage); transientStatus != 0 {
accountStatus = transientStatus
@@ -420,9 +428,18 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
s.handleOpenAIAccountUpstreamError(ctx, account, accountStatus, resp.Header, upstreamMessage, canonicalModel)
}
if turn == 1 && !wroteDownstream && shouldFailover {
return nil, newOpenAIUpstreamFailoverError(statusCode, resp.Header, upstreamMessage, errMessage, false)
if account.Platform == PlatformGrok {
return nil, newOpenAIUpstreamFailoverError(statusCode, resp.Header, upstreamMessage, errMessage, false)
}
return nil, s.newOpenAIStreamFailoverError(c, account, true, resp.Header.Get("x-request-id"), upstreamMessage, errMessage, resp.Header)
}
if wroteDownstream && requestScopedCapacity && !capacityFailoverSuppressedLogged {
logOpenAICapacityFailoverSuppressed(ctx, account, "ws_http_bridge", resp.Header.Get("x-request-id"), eventType)
capacityFailoverSuppressedLogged = true
}
if eventType == "error" {
upstreamEventErr = errors.New(errMessage)
}
upstreamEventErr = errors.New(errMessage)
}
// 客户端写出副本改写容量降载码:Codex 对 error/response.failed 中的
@@ -436,26 +453,50 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
}
}
if !clientDisconnected {
if err := writeClientMessage(clientMessage); err != nil {
if isOpenAIWSClientDisconnectError(err) {
clientDisconnected = true
closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err)
logOpenAIWSModeInfo(
"ingress_ws_http_bridge_client_disconnected_drain account_id=%d turn=%d close_status=%s close_reason=%s",
account.ID,
turn,
closeStatus,
truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen),
)
} else {
return nil, wrapOpenAIWSIngressTurnError(
"write_client",
fmt.Errorf("write client websocket event: %w", err),
wroteDownstream,
stageBeforeSemanticOutput := turn == 1 && account.Platform == PlatformOpenAI && !wroteDownstream
commitStagedMessages := !stageBeforeSemanticOutput ||
openAIStreamDataStartsClientOutput(string(clientMessage), eventType) ||
isOpenAIWSTerminalEvent(eventType)
if stageBeforeSemanticOutput && !commitStagedMessages {
if pendingClientMessageBytes+int64(len(clientMessage)) > openAIFirstOutputStageMaxBytes {
return nil, s.newOpenAIStreamFailoverError(
c,
account,
true,
resp.Header.Get("x-request-id"),
nil,
"OpenAI WS HTTP bridge first-output staging limit exceeded",
resp.Header,
)
}
pendingClientMessages = append(pendingClientMessages, append([]byte(nil), clientMessage...))
pendingClientMessageBytes += int64(len(clientMessage))
} else {
wroteDownstream = true
messages := append(pendingClientMessages, clientMessage)
pendingClientMessages = nil
pendingClientMessageBytes = 0
for _, message := range messages {
if err := writeClientMessage(message); err != nil {
if isOpenAIWSClientDisconnectError(err) {
clientDisconnected = true
closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err)
logOpenAIWSModeInfo(
"ingress_ws_http_bridge_client_disconnected_drain account_id=%d turn=%d close_status=%s close_reason=%s",
account.ID,
turn,
closeStatus,
truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen),
)
break
}
return nil, wrapOpenAIWSIngressTurnError(
"write_client",
fmt.Errorf("write client websocket event: %w", err),
wroteDownstream,
)
}
wroteDownstream = true
}
}
}
@@ -239,10 +239,9 @@ func TestProxyOpenAIWSHTTPBridgeTurnRewritesCapacityShedCodeForClient(t *testing
wantErr: true,
},
{
// response.failed 不走 error 事件分支:即便 turn 1 也会被当终止事件
// 原样转发(不 failover),因此改写必须在这里同样生效。
name: "turn1_bare_response_failed",
turn: 1,
// 后续 turn 不允许 replay,容量错误必须改写后交给客户端重试。
name: "turn2_bare_response_failed",
turn: 2,
body: "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_shed\",\"status\":\"failed\",\"error\":{\"code\":\"server_is_overloaded\",\"message\":\"Our servers are currently overloaded. Please try again later.\"}}}\n\n",
},
}
@@ -283,6 +282,91 @@ func TestProxyOpenAIWSHTTPBridgeTurnRewritesCapacityShedCodeForClient(t *testing
}
}
func TestProxyOpenAIWSHTTPBridgeTurnStagesMetadataBeforeCapacityFailover(t *testing.T) {
gin.SetMode(gin.TestMode)
body := strings.Join([]string{
`data: {"type":"response.created","response":{"id":"resp_shed"}}`,
"",
`data: {"type":"response.in_progress","response":{"id":"resp_shed"}}`,
"",
`data: {"type":"response.failed","response":{"id":"resp_shed","status":"failed","error":{"message":"Our servers are currently overloaded. Please try again later."}}}`,
"",
}, "\n")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"X-Request-Id": []string{"rid-ws-bridge-capacity"}},
Body: io.NopCloser(strings.NewReader(body)),
}}
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
account := &Account{ID: 12, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
var writes [][]byte
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
context.Background(), c, account, "sk-test", payload, len(payload),
"gpt-5", "", "", "", "", 1,
func(message []byte) error {
writes = append(writes, append([]byte(nil), message...))
return nil
},
)
require.Nil(t, result)
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.True(t, failoverErr.RetryableOnSameAccount)
require.True(t, failoverErr.RequestScopedTransient)
require.Empty(t, writes)
}
func TestProxyOpenAIWSHTTPBridgeTurnDoesNotReplayCapacityAfterSemanticOutput(t *testing.T) {
gin.SetMode(gin.TestMode)
logSink, restore := captureStructuredLog(t)
defer restore()
body := strings.Join([]string{
`data: {"type":"response.created","response":{"id":"resp_partial"}}`,
"",
`data: {"type":"response.output_text.delta","delta":"partial"}`,
"",
`data: {"type":"response.failed","response":{"id":"resp_partial","status":"failed","error":{"code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}}}`,
"",
}, "\n")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"X-Request-Id": []string{"rid-ws-bridge-post-output"}},
Body: io.NopCloser(strings.NewReader(body)),
}}
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
account := &Account{ID: 13, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
var writes [][]byte
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
context.Background(), c, account, "sk-test", payload, len(payload),
"gpt-5", "", "", "", "", 1,
func(message []byte) error {
writes = append(writes, append([]byte(nil), message...))
return nil
},
)
require.NotNil(t, result)
require.NoError(t, err)
require.Len(t, writes, 3)
var failoverErr *UpstreamFailoverError
require.False(t, errors.As(err, &failoverErr))
require.Contains(t, string(writes[2]), `"code":"server_error"`)
require.NotContains(t, string(writes[2]), "server_is_overloaded")
require.True(t, logSink.ContainsMessage("gateway.failover_suppressed_after_semantic_output"))
require.True(t, logSink.ContainsFieldValue("path", "ws_http_bridge"))
}
func TestProxyOpenAIWSHTTPBridgeTurnRequiresTerminalEvent(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -294,10 +378,10 @@ func TestProxyOpenAIWSHTTPBridgeTurnRequiresTerminalEvent(t *testing.T) {
}{
{name: "done_without_events_fails_over", body: "data: [DONE]\n\n", wantFailover: true},
{
name: "created_then_done_is_truncated_not_success",
name: "created_then_done_fails_over_before_semantic_output",
body: "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_truncated\"}}\n\n" +
"data: [DONE]\n\n",
wantWrites: 1,
wantFailover: true,
},
}
for _, tt := range tests {