Merge pull request #4384 from wp-a/fix/openai-model-scoped-transient-cooldown

[codex] scope OpenAI transient cooldowns by model
This commit is contained in:
Wesley Liddick
2026-07-16 09:33:09 +08:00
committed by GitHub
26 changed files with 1003 additions and 122 deletions
+3 -3
View File
@@ -273,7 +273,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
return
}
if failoverErr.ShouldReportAccountScheduleFailure() {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil)
}
if c.Writer.Size() != writerSizeBeforeForward {
h.handleFailoverExhausted(c, failoverErr, true)
@@ -321,7 +321,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
)
continue
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil)
if c.Writer.Size() == writerSizeBeforeForward {
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
}
@@ -332,7 +332,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
return
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), true, nil)
if endpoint.IsGenerationRequest() && strings.TrimSpace(result.ResponseID) != "" {
if err := h.gatewayService.BindGrokMediaVideoRequestAccount(requestCtx, apiKey.GroupID, result.ResponseID, account.ID); err != nil {
reqLog.Warn("grok_media.bind_video_request_account_failed",
@@ -163,7 +163,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
service.SetOpsLatencyMs(c, service.OpsResponseLatencyMsKey, time.Since(forwardStart).Milliseconds())
if err == nil {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestedModel), true, nil)
if result != nil {
h.recordAlphaSearchUsage(c, apiKey, account, subscription, channelMapping, requestedModel, body, result, subject.UserID)
}
@@ -172,7 +172,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
var failoverErr *service.UpstreamFailoverError
if !errors.As(err, &failoverErr) {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestedModel), false, nil)
if c.Writer.Size() == writerSizeBeforeForward {
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
}
@@ -180,7 +180,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
return
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestedModel), false, nil)
if c.Writer.Size() != writerSizeBeforeForward {
h.handleFailoverExhausted(c, failoverErr, true)
return
@@ -255,7 +255,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
return
}
if failoverErr.ShouldReportAccountScheduleFailure() {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil)
}
if !failoverErr.ShouldRetryNextAccount() {
h.handleFailoverExhausted(c, failoverErr, streamStarted)
@@ -300,7 +300,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
)
continue
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil)
upstreamErrorAlreadyCommunicated := openAIForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err)
wroteFallback := false
if !upstreamErrorAlreadyCommunicated {
@@ -316,9 +316,9 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
}
}
if result != nil {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, result.FirstTokenMs)
} else {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, nil)
}
userAgent := c.GetHeader("User-Agent")
@@ -192,7 +192,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
h.handleFailoverExhausted(c, failoverErr, true)
return
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil)
if failoverClientGone(c) {
reqLog.Info("openai_embeddings.failover_aborted_client_disconnected",
zap.Int64("account_id", account.ID),
@@ -216,7 +216,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
)
continue
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil)
if c.Writer.Size() == writerSizeBeforeForward {
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
}
@@ -227,7 +227,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
return
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, nil)
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
inboundEndpoint := GetInboundEndpoint(c)
@@ -42,6 +42,10 @@ type OpenAIGatewayHandler struct {
const maxOpenAIFirstOutputTimeoutSwitches = 1
func openAIForwardSucceededForScheduling(result *service.OpenAIForwardResult) bool {
return result.SucceededForScheduling()
}
func resolveOpenAIMessagesDispatchMappedModel(apiKey *service.APIKey, requestedModel string) string {
if apiKey == nil || apiKey.Group == nil {
return ""
@@ -481,7 +485,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
streamStarted = true
}
if failoverErr.ShouldReportAccountScheduleFailure() {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil)
}
if !failoverErr.ShouldRetryNextAccount() {
h.handleFailoverExhausted(c, failoverErr, streamStarted)
@@ -541,7 +545,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
reqLog.Warn("openai.upstream_failover_switching", failoverSwitchFields...)
continue
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil)
upstreamErrorAlreadyCommunicated := openAIForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err)
wroteFallback := false
if !upstreamErrorAlreadyCommunicated {
@@ -566,9 +570,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
if account.Type == service.AccountTypeOAuth && !account.IsShadow() {
h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(c.Request.Context(), account.ID, result.ResponseHeaders)
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), result.FirstTokenMs)
} else {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), nil)
}
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context
@@ -1009,7 +1013,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
return
}
if failoverErr.ShouldReportAccountScheduleFailure() {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), false, nil)
}
if !failoverErr.ShouldRetryNextAccount() {
h.handleAnthropicFailoverExhausted(c, failoverErr, streamStarted)
@@ -1061,7 +1065,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
)
return
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), false, nil)
wroteFallback := h.ensureAnthropicErrorResponse(c, streamStarted)
reqLog.Warn("openai_messages.forward_failed",
zap.Int64("account_id", account.ID),
@@ -1072,9 +1076,9 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
}
}
if result != nil {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), true, result.FirstTokenMs)
} else {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), true, nil)
}
userAgent := c.GetHeader("User-Agent")
@@ -1558,7 +1562,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
return false
}
if failoverErr.ShouldReportAccountScheduleFailure() {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil)
}
releaseAccountSlot()
if !failoverErr.ShouldRetryNextAccount() {
@@ -1777,7 +1781,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
if account.Type == service.AccountTypeOAuth && !account.IsShadow() {
h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(ctx, account.ID, result.ResponseHeaders)
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), result.FirstTokenMs)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
@@ -1857,7 +1861,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
return
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil)
closeStatus, closeReason := summarizeWSCloseErrorForLog(err)
proxyFailedFields := []zap.Field{
zap.Int64("account_id", account.ID),
@@ -96,6 +96,19 @@ func TestOpenAIHandleStreamingAwareError_JSONEscaping(t *testing.T) {
}
}
func TestOpenAIForwardSucceededForScheduling(t *testing.T) {
require.True(t, openAIForwardSucceededForScheduling(nil))
require.True(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{}))
require.True(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{
OpenAIWSMode: true,
UpstreamTerminalEvent: "response.completed",
}))
require.False(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{
OpenAIWSMode: true,
UpstreamTerminalEvent: "response.failed",
}))
}
func TestResolveOpenAIMessagesMetadataSession_DoesNotDerivePromptCacheKey(t *testing.T) {
body := []byte(`{"model":"claude-sonnet-4-5","metadata":{"user_id":"claude-code-session"},"messages":[{"role":"user","content":"hello"}]}`)
+5 -5
View File
@@ -253,7 +253,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
var imageUpstreamErr *service.OpenAIImagesUpstreamError
if errors.As(err, &imageUpstreamErr) {
retryableServerError := service.IsOpenAIImagesRetryableUpstreamError(imageUpstreamErr)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, !retryableServerError, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), !retryableServerError, nil)
logEvent := "openai.images.upstream_user_error"
if retryableServerError {
logEvent = "openai.images.upstream_server_error_after_flush"
@@ -269,7 +269,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
}
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil)
if service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeForward {
reqLog.Warn("openai.images.upstream_failover_skipped_after_flush",
zap.Int64("account_id", account.ID),
@@ -323,7 +323,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
)
continue
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil)
upstreamErrorAlreadyCommunicated := openAIForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err)
wroteFallback := false
if !upstreamErrorAlreadyCommunicated {
@@ -348,9 +348,9 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
if account.Type == service.AccountTypeOAuth && !account.IsShadow() {
h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(c.Request.Context(), account.ID, result.ResponseHeaders)
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), true, result.FirstTokenMs)
} else {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), true, nil)
}
userAgent := c.GetHeader("User-Agent")
@@ -0,0 +1,177 @@
package service
import (
"strings"
"sync"
"time"
)
const (
openAIModelTransientFailureWindow = time.Minute
openAIModelTransientShortCooldown = 10 * time.Second
openAIModelTransientLongCooldown = 45 * time.Second
openAIModelTransientDefaultMax = 4096
openAIModelTransientMaxModelBytes = 512
)
type openAIAccountModelKey struct {
AccountID int64
Model string
}
type openAIAccountModelTransientEntry struct {
failureStreak int
lastFailure time.Time
blockUntil time.Time
lastTouched time.Time
}
type openAIAccountModelTransientDecision struct {
FailureStreak int
Cooldown time.Duration
BlockUntil time.Time
}
type openAIAccountModelTransientState struct {
mu sync.Mutex
entries map[openAIAccountModelKey]openAIAccountModelTransientEntry
maxEntries int
}
func newOpenAIAccountModelTransientState(maxEntries int) *openAIAccountModelTransientState {
if maxEntries <= 0 {
maxEntries = openAIModelTransientDefaultMax
}
return &openAIAccountModelTransientState{
entries: make(map[openAIAccountModelKey]openAIAccountModelTransientEntry),
maxEntries: maxEntries,
}
}
func normalizeOpenAIAccountModelTransientModel(model string) string {
model = strings.TrimSpace(model)
if len(model) > openAIModelTransientMaxModelBytes {
return ""
}
return strings.ToLower(model)
}
func openAIAccountModelTransientKey(accountID int64, model string) (openAIAccountModelKey, bool) {
model = normalizeOpenAIAccountModelTransientModel(model)
if accountID <= 0 || model == "" {
return openAIAccountModelKey{}, false
}
return openAIAccountModelKey{AccountID: accountID, Model: model}, true
}
func (s *openAIAccountModelTransientState) recordFailure(accountID int64, model string, now time.Time) openAIAccountModelTransientDecision {
key, ok := openAIAccountModelTransientKey(accountID, model)
if s == nil || !ok {
return openAIAccountModelTransientDecision{}
}
if now.IsZero() {
now = time.Now()
}
s.mu.Lock()
defer s.mu.Unlock()
if s.entries == nil {
s.entries = make(map[openAIAccountModelKey]openAIAccountModelTransientEntry)
}
if s.maxEntries <= 0 {
s.maxEntries = openAIModelTransientDefaultMax
}
entry, exists := s.entries[key]
if !exists {
s.evictOldestLocked()
}
if !exists || entry.lastFailure.IsZero() || now.Sub(entry.lastFailure) > openAIModelTransientFailureWindow || now.Before(entry.lastFailure) {
entry.failureStreak = 0
entry.blockUntil = time.Time{}
}
entry.failureStreak++
entry.lastFailure = now
entry.lastTouched = now
cooldown := time.Duration(0)
switch {
case entry.failureStreak >= 3:
cooldown = openAIModelTransientLongCooldown
case entry.failureStreak == 2:
cooldown = openAIModelTransientShortCooldown
}
if cooldown > 0 {
entry.blockUntil = now.Add(cooldown)
} else {
entry.blockUntil = time.Time{}
}
s.entries[key] = entry
return openAIAccountModelTransientDecision{
FailureStreak: entry.failureStreak,
Cooldown: cooldown,
BlockUntil: entry.blockUntil,
}
}
func (s *openAIAccountModelTransientState) recordSuccess(accountID int64, model string) {
key, ok := openAIAccountModelTransientKey(accountID, model)
if s == nil || !ok {
return
}
s.mu.Lock()
delete(s.entries, key)
s.mu.Unlock()
}
func (s *openAIAccountModelTransientState) isBlocked(accountID int64, model string, now time.Time) bool {
key, ok := openAIAccountModelTransientKey(accountID, model)
if s == nil || !ok {
return false
}
if now.IsZero() {
now = time.Now()
}
s.mu.Lock()
defer s.mu.Unlock()
entry, exists := s.entries[key]
if !exists {
return false
}
if !entry.lastFailure.IsZero() && now.Sub(entry.lastFailure) > openAIModelTransientFailureWindow {
delete(s.entries, key)
return false
}
entry.lastTouched = now
s.entries[key] = entry
return !entry.blockUntil.IsZero() && now.Before(entry.blockUntil)
}
func (s *openAIAccountModelTransientState) size() int {
if s == nil {
return 0
}
s.mu.Lock()
defer s.mu.Unlock()
return len(s.entries)
}
func (s *openAIAccountModelTransientState) evictOldestLocked() {
if len(s.entries) < s.maxEntries {
return
}
var oldestKey openAIAccountModelKey
var oldestTime time.Time
found := false
for key, entry := range s.entries {
if !found || entry.lastTouched.Before(oldestTime) {
oldestKey = key
oldestTime = entry.lastTouched
found = true
}
}
if found {
delete(s.entries, oldestKey)
}
}
@@ -0,0 +1,130 @@
package service
import (
"fmt"
"strings"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestOpenAIModelTransient_FirstFailureDoesNotCreateLongBlock(t *testing.T) {
state := newOpenAIAccountModelTransientState(128)
now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC)
decision := state.recordFailure(35, "gpt-5.5", now)
assert.Equal(t, 1, decision.FailureStreak)
assert.Zero(t, decision.Cooldown)
assert.False(t, state.isBlocked(35, "gpt-5.5", now))
}
func TestOpenAIModelTransient_SecondFailureCreatesShortModelBlock(t *testing.T) {
state := newOpenAIAccountModelTransientState(128)
now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC)
state.recordFailure(35, "gpt-5.5", now)
decision := state.recordFailure(35, "gpt-5.5", now.Add(time.Second))
assert.Equal(t, 2, decision.FailureStreak)
assert.Equal(t, openAIModelTransientShortCooldown, decision.Cooldown)
assert.True(t, state.isBlocked(35, "gpt-5.5", now.Add(2*time.Second)))
assert.False(t, state.isBlocked(35, "gpt-5.5", now.Add(openAIModelTransientShortCooldown+2*time.Second)))
}
func TestOpenAIModelTransient_ThirdFailureCreatesFortyFiveSecondModelBlock(t *testing.T) {
state := newOpenAIAccountModelTransientState(128)
now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC)
state.recordFailure(35, "gpt-5.5", now)
state.recordFailure(35, "gpt-5.5", now.Add(time.Second))
decision := state.recordFailure(35, "gpt-5.5", now.Add(2*time.Second))
assert.Equal(t, 3, decision.FailureStreak)
assert.Equal(t, 45*time.Second, decision.Cooldown)
assert.True(t, state.isBlocked(35, "gpt-5.5", now.Add(40*time.Second)))
assert.False(t, state.isBlocked(35, "gpt-5.5", now.Add(48*time.Second)))
}
func TestOpenAIModelTransient_BlockIsIsolatedByModel(t *testing.T) {
state := newOpenAIAccountModelTransientState(128)
now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC)
state.recordFailure(35, "gpt-5.6-terra", now)
state.recordFailure(35, "GPT-5.6-TERRA", now.Add(time.Second))
assert.True(t, state.isBlocked(35, "gpt-5.6-terra", now.Add(2*time.Second)))
assert.False(t, state.isBlocked(35, "gpt-5.5", now.Add(2*time.Second)))
assert.False(t, state.isBlocked(47, "gpt-5.6-terra", now.Add(2*time.Second)))
}
func TestOpenAIModelTransient_SuccessClearsStreakAndBlock(t *testing.T) {
state := newOpenAIAccountModelTransientState(128)
now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC)
state.recordFailure(35, "gpt-5.5", now)
state.recordFailure(35, "gpt-5.5", now.Add(time.Second))
require.True(t, state.isBlocked(35, "gpt-5.5", now.Add(2*time.Second)))
state.recordSuccess(35, "gpt-5.5")
assert.False(t, state.isBlocked(35, "gpt-5.5", now.Add(2*time.Second)))
decision := state.recordFailure(35, "gpt-5.5", now.Add(3*time.Second))
assert.Equal(t, 1, decision.FailureStreak)
assert.Zero(t, decision.Cooldown)
}
func TestOpenAIModelTransient_StaleStreakExpires(t *testing.T) {
state := newOpenAIAccountModelTransientState(128)
now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC)
state.recordFailure(35, "gpt-5.5", now)
decision := state.recordFailure(35, "gpt-5.5", now.Add(openAIModelTransientFailureWindow+time.Second))
assert.Equal(t, 1, decision.FailureStreak)
assert.Zero(t, decision.Cooldown)
}
func TestOpenAIModelTransient_IgnoresInvalidKeys(t *testing.T) {
state := newOpenAIAccountModelTransientState(128)
now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC)
assert.Zero(t, state.recordFailure(0, "gpt-5.5", now).FailureStreak)
assert.Zero(t, state.recordFailure(35, " ", now).FailureStreak)
assert.False(t, state.isBlocked(0, "gpt-5.5", now))
assert.False(t, state.isBlocked(35, "", now))
assert.Equal(t, 0, state.size())
}
func TestOpenAIModelTransient_IgnoresOversizedModelKey(t *testing.T) {
state := newOpenAIAccountModelTransientState(128)
now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC)
model := strings.Repeat("m", openAIModelTransientMaxModelBytes+1)
decision := state.recordFailure(35, model, now)
assert.Zero(t, decision.FailureStreak)
assert.False(t, state.isBlocked(35, model, now))
assert.Equal(t, 0, state.size())
}
func TestOpenAIModelTransient_StateIsBoundedAndConcurrencySafe(t *testing.T) {
const maxEntries = 16
state := newOpenAIAccountModelTransientState(maxEntries)
now := time.Date(2026, 7, 10, 10, 0, 0, 0, time.UTC)
var wg sync.WaitGroup
for i := 0; i < 128; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
model := fmt.Sprintf("gpt-test-%d", i)
state.recordFailure(int64(i+1), model, now.Add(time.Duration(i)*time.Millisecond))
_ = state.isBlocked(int64(i+1), model, now.Add(time.Second))
}(i)
}
wg.Wait()
assert.LessOrEqual(t, state.size(), maxEntries)
}
@@ -2,7 +2,9 @@ package service
import (
"context"
"log/slog"
"net/http"
"strings"
"sync"
"time"
)
@@ -43,7 +45,9 @@ func isOpenAIAccount(account *Account) bool {
return account != nil && (account.Platform == PlatformOpenAI || account.Platform == PlatformGrok)
}
func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, requestedModel ...string) bool {
// handleOpenAIAccountUpstreamError expects canonicalModel to be the model used
// for scheduling after applying account mapping exactly once.
func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, canonicalModel ...string) bool {
stateCtx, cancel := openAIAccountStateContext(ctx)
defer cancel()
@@ -64,16 +68,43 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont
if s == nil || account == nil || s.rateLimitService == nil {
return false
}
if len(requestedModel) > 0 && s.rateLimitService.HandleUpstreamModelNotFound(stateCtx, account, requestedModel[0], statusCode, responseBody) {
if len(canonicalModel) > 0 && s.rateLimitService.HandleUpstreamModelNotFound(stateCtx, account, canonicalModel[0], statusCode, responseBody) {
return true
}
shouldDisable := s.rateLimitService.HandleUpstreamError(stateCtx, account, statusCode, headers, responseBody)
if shouldDisable {
s.BlockAccountScheduling(account, time.Time{}, "upstream_disable")
}
if !shouldDisable && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey && shouldCooldownOpenAITransientUpstreamError(statusCode, responseBody) {
model := ""
if len(canonicalModel) > 0 {
model = canonicalModel[0]
}
decision := s.recordOpenAIAccountModelTransientFailure(account, model, time.Now())
if decision.FailureStreak > 0 {
slog.Warn("openai_model_transient_state",
"account_id", account.ID,
"model", openAIAccountModelTransientModel(model),
"failure_streak", decision.FailureStreak,
"cooldown_ms", decision.Cooldown.Milliseconds(),
"block_scope", "account_model",
)
}
}
return shouldDisable
}
func shouldCooldownOpenAITransientUpstreamError(statusCode int, responseBody []byte) bool {
switch statusCode {
case http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout, 520, 521, 522, 523, 524:
return true
case http.StatusBadRequest:
return isOpenAITransientProcessingError(statusCode, "", responseBody)
default:
return false
}
}
func (s *OpenAIGatewayService) markOpenAIOAuth429RateLimited(ctx context.Context, account *Account, headers http.Header, responseBody []byte) {
if s == nil || !isOpenAIOAuthAccount(account) {
return
@@ -191,6 +222,68 @@ func (s *OpenAIGatewayService) isOpenAIAccountRuntimeBlocked(account *Account) b
return false
}
func (s *OpenAIGatewayService) getOpenAIAccountModelTransientState() *openAIAccountModelTransientState {
if s == nil {
return nil
}
s.openaiModelTransientOnce.Do(func() {
if s.openaiModelTransient == nil {
s.openaiModelTransient = newOpenAIAccountModelTransientState(openAIModelTransientDefaultMax)
}
})
return s.openaiModelTransient
}
func canonicalOpenAIAccountSchedulingModel(account *Account, requestedModel string) string {
model := strings.TrimSpace(requestedModel)
if account == nil || model == "" {
return model
}
if mapped := strings.TrimSpace(account.GetMappedModel(model)); mapped != "" {
return mapped
}
return model
}
func openAIAccountModelTransientModel(canonicalModel string) string {
return normalizeOpenAIAccountModelTransientModel(canonicalModel)
}
func (s *OpenAIGatewayService) recordOpenAIAccountModelTransientFailure(account *Account, canonicalModel string, now time.Time) openAIAccountModelTransientDecision {
if s == nil || account == nil {
return openAIAccountModelTransientDecision{}
}
state := s.getOpenAIAccountModelTransientState()
if state == nil {
return openAIAccountModelTransientDecision{}
}
return state.recordFailure(account.ID, openAIAccountModelTransientModel(canonicalModel), now)
}
func (s *OpenAIGatewayService) clearOpenAIAccountModelTransientState(accountID int64, model string) {
state := s.getOpenAIAccountModelTransientState()
if state == nil {
return
}
state.recordSuccess(accountID, model)
}
func (s *OpenAIGatewayService) isOpenAIAccountModelRuntimeBlocked(account *Account, requestedModel string) bool {
if s == nil || account == nil {
return false
}
state := s.getOpenAIAccountModelTransientState()
if state == nil {
return false
}
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, requestedModel)
return state.isBlocked(account.ID, openAIAccountModelTransientModel(canonicalModel), time.Now())
}
func (s *OpenAIGatewayService) isOpenAIAccountRequestRuntimeBlocked(account *Account, requestedModel string) bool {
return s != nil && (s.isOpenAIAccountRuntimeBlocked(account) || s.isOpenAIAccountModelRuntimeBlocked(account, requestedModel))
}
func (s *OpenAIGatewayService) recordOpenAIOAuth429() {
if s == nil {
return
@@ -0,0 +1,116 @@
package service
import (
"context"
"net/http"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
type transientCooldownAccountRepo struct {
AccountRepository
}
func (transientCooldownAccountRepo) SetOverloaded(context.Context, int64, time.Time) error {
return nil
}
func TestHandleOpenAITransientError_BlocksOnlyRequestedModel(t *testing.T) {
svc := &OpenAIGatewayService{}
svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil)
account := &Account{
ID: 5105,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
}
firstShouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadGateway, http.Header{}, []byte(`{"error":{"message":"Upstream request failed","type":"upstream_error"}}`), "gpt-5.5")
secondShouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadGateway, http.Header{}, []byte(`{"error":{"message":"Upstream request failed","type":"upstream_error"}}`), "gpt-5.5")
require.False(t, firstShouldDisable)
require.False(t, secondShouldDisable)
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5"))
require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.6-terra"))
}
func TestHandleOpenAITransientError_TransientStatusesUseModelScope(t *testing.T) {
for _, statusCode := range []int{http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout, 520, 521, 522, 523, 524} {
t.Run(http.StatusText(statusCode), func(t *testing.T) {
svc := &OpenAIGatewayService{}
svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil)
account := &Account{
ID: int64(5100 + statusCode),
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
}
firstShouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, statusCode, http.Header{}, []byte(`{"error":{"message":"temporary upstream failure"}}`), "gpt-5.5")
secondShouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, statusCode, http.Header{}, []byte(`{"error":{"message":"temporary upstream failure"}}`), "gpt-5.5")
require.False(t, firstShouldDisable)
require.False(t, secondShouldDisable)
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account), "status %d must not block the whole account", statusCode)
require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5"), "status %d should block the failing model", statusCode)
})
}
}
func TestHandleOpenAITransientError_529RemainsOverloadOnly(t *testing.T) {
require.False(t, shouldCooldownOpenAITransientUpstreamError(529, []byte(`{"error":{"message":"overloaded"}}`)))
}
func TestHandleOpenAITransientError_CanonicalModelIsNotMappedTwice(t *testing.T) {
svc := &OpenAIGatewayService{}
svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil)
account := &Account{
ID: 5107,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"model_mapping": map[string]any{
"public-alias": "upstream-a",
"upstream-a": "upstream-b",
},
},
}
canonicalModel := account.GetMappedModel("public-alias")
require.Equal(t, "upstream-a", canonicalModel)
for range 2 {
svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadGateway, http.Header{}, []byte(`{"error":{"message":"temporary upstream failure"}}`), canonicalModel)
}
require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "public-alias"))
svc.ReportOpenAIAccountScheduleResult(account.ID, canonicalModel, true, nil)
require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "public-alias"))
}
func TestHandleOpenAITransientError_DoesNotBlockParameter400(t *testing.T) {
svc := &OpenAIGatewayService{}
svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil)
account := &Account{
ID: 5103,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
}
shouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadRequest, http.Header{}, []byte(`{"error":{"message":"Invalid type for input[0].arguments"}}`), "gpt-5.5")
require.False(t, shouldDisable)
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5"))
}
func TestHandleOpenAITransientError_HardDisableStillBlocksWholeAccount(t *testing.T) {
svc := &OpenAIGatewayService{}
account := &Account{ID: 5106, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
svc.BlockAccountScheduling(account, time.Now().Add(time.Minute), "upstream_disable")
require.True(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.5"))
require.True(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.6-sol"))
}
@@ -1108,7 +1108,7 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance(
if !account.IsSchedulable() || account.Platform != normalizeOpenAICompatiblePlatform(req.Platform) || !account.IsOpenAICompatible() {
continue
}
if s.service.isOpenAIAccountRuntimeBlocked(account) {
if s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) {
continue
}
// require_privacy_set: 跳过 privacy 未设置的账号并标记异常
@@ -1366,7 +1366,7 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatible(ctx context.C
if account == nil {
return false
}
if s != nil && s.service != nil && s.service.isOpenAIAccountRuntimeBlocked(account) {
if s != nil && s.service != nil && s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) {
return false
}
// Quota auto-pause must be evaluated during the initial filter too. Without it the
@@ -1841,7 +1841,10 @@ func (s *OpenAIGatewayService) isOpenAIAccountTransportCompatible(account *Accou
return s.getOpenAIWSProtocolResolver().Resolve(account).Transport == requiredTransport
}
func (s *OpenAIGatewayService) ReportOpenAIAccountScheduleResult(accountID int64, success bool, firstTokenMs *int) {
func (s *OpenAIGatewayService) ReportOpenAIAccountScheduleResult(accountID int64, model string, success bool, firstTokenMs *int) {
if success {
s.clearOpenAIAccountModelTransientState(accountID, normalizeOpenAIAccountModelTransientModel(model))
}
scheduler := s.getOpenAIAccountScheduler(context.Background())
if scheduler == nil {
return
@@ -1032,7 +1032,7 @@ func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics_DisabledNoOp(t *test
svc := &OpenAIGatewayService{}
ttft := 120
svc.ReportOpenAIAccountScheduleResult(10, true, &ttft)
svc.ReportOpenAIAccountScheduleResult(10, "", true, &ttft)
svc.RecordOpenAIAccountSwitch()
snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics()
@@ -2079,6 +2079,30 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_UsesAccountPriorityWith
}
}
func TestOpenAIAccountScheduler_SkipsAccountBlockedForRequestedModel(t *testing.T) {
now := time.Now()
account := &Account{ID: 21633, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
svc := &OpenAIGatewayService{openaiModelTransient: newOpenAIAccountModelTransientState(128)}
svc.openaiModelTransient.recordFailure(account.ID, "gpt-5.5", now)
svc.openaiModelTransient.recordFailure(account.ID, "gpt-5.5", now.Add(time.Millisecond))
scheduler := &defaultOpenAIAccountScheduler{service: svc}
require.False(t, scheduler.isAccountRequestCompatible(context.Background(), account, OpenAIAccountScheduleRequest{RequestedModel: "gpt-5.5"}))
require.True(t, scheduler.isAccountRequestCompatible(context.Background(), account, OpenAIAccountScheduleRequest{RequestedModel: "gpt-5.6-sol"}))
}
func TestReportOpenAIAccountScheduleResult_SuccessClearsModelTransientState(t *testing.T) {
svc := &OpenAIGatewayService{openaiModelTransient: newOpenAIAccountModelTransientState(128)}
now := time.Now()
svc.openaiModelTransient.recordFailure(21636, "gpt-5.5", now)
svc.openaiModelTransient.recordFailure(21636, "gpt-5.5", now.Add(time.Millisecond))
require.True(t, svc.openaiModelTransient.isBlocked(21636, "gpt-5.5", now.Add(2*time.Millisecond)))
svc.ReportOpenAIAccountScheduleResult(21636, "gpt-5.5", true, nil)
require.False(t, svc.openaiModelTransient.isBlocked(21636, "gpt-5.5", now.Add(2*time.Millisecond)))
}
func TestDefaultOpenAIAccountScheduler_ShouldEscapeStickyAccount_ThresholdBoundary(t *testing.T) {
stats := newOpenAIAccountRuntimeStats()
accountID := int64(21501)
@@ -2534,7 +2558,7 @@ func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics(t *testing.T) {
selection, _, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_metrics", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
require.NoError(t, err)
require.NotNil(t, selection)
svc.ReportOpenAIAccountScheduleResult(account.ID, true, intPtrForTest(120))
svc.ReportOpenAIAccountScheduleResult(account.ID, "", true, intPtrForTest(120))
svc.RecordOpenAIAccountSwitch()
snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics()
@@ -2871,7 +2895,7 @@ func TestOpenAIGatewayService_SchedulerWrappersAndDefaults(t *testing.T) {
svc := &OpenAIGatewayService{}
ttft := 120
svc.ReportOpenAIAccountScheduleResult(10, true, &ttft)
svc.ReportOpenAIAccountScheduleResult(10, "", true, &ttft)
svc.RecordOpenAIAccountSwitch()
snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics()
require.Equal(t, OpenAIAccountSchedulerMetricsSnapshot{}, snapshot)
@@ -576,7 +576,8 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough(
setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail)
logOpenAIInstructionsRequiredDebug(ctx, c, account, resp.StatusCode, upstreamMsg, requestBody, body)
reqModel, _, _ := extractOpenAIRequestMetaFromBody(requestBody)
_ = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, reqModel)
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, reqModel)
_ = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, canonicalModel)
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
@@ -637,7 +638,8 @@ func (s *OpenAIGatewayService) handleErrorResponsePassthrough(
// 刚被限流的账号。cyber 例外:不冷却账号。
if !cyberHit {
reqModel, _, _ := extractOpenAIRequestMetaFromBody(requestBody)
_ = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, reqModel)
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, reqModel)
_ = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, canonicalModel)
}
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
@@ -659,7 +659,7 @@ func (s *OpenAIGatewayService) tryStickySessionHit(ctx context.Context, groupID
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
return nil
}
if s.isOpenAIAccountRuntimeBlocked(account) {
if s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) {
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
return nil
}
@@ -864,7 +864,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
} else if !openAIStickyAccountMatchesGroup(account, groupID) {
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
} else if s.isOpenAIAccountRuntimeBlocked(account) {
} else if s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) {
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
} else if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel, requireCompact) {
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
@@ -927,7 +927,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex
if !parentHealthyForShadow(acc, parentLookupL2) {
continue
}
if s.isOpenAIAccountRuntimeBlocked(acc) {
if s.isOpenAIAccountRequestRuntimeBlocked(acc, requestedModel) {
continue
}
if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, acc, requestedModel, requireCompact) {
@@ -1162,7 +1162,7 @@ func (s *OpenAIGatewayService) resolveFreshSchedulableOpenAIAccount(ctx context.
if !parentHealthyForShadow(fresh, s.parentAccountLookup(ctx)) {
return nil
}
if s.isOpenAIAccountRuntimeBlocked(fresh) {
if s.isOpenAIAccountRequestRuntimeBlocked(fresh, requestedModel) {
return nil
}
return fresh
@@ -1207,7 +1207,7 @@ func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDB(ctx context.Co
if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) {
return nil
}
if s.isOpenAIAccountRuntimeBlocked(latest) {
if s.isOpenAIAccountRequestRuntimeBlocked(latest, requestedModel) {
return nil
}
return latest
@@ -236,22 +236,25 @@ type OpenAIForwardResult struct {
ServiceTier *string
// ReasoningEffort is extracted from request body (reasoning.effort) or derived from model suffix.
// Stored for usage records display; nil means not provided / not applicable.
ReasoningEffort *string
Stream bool
OpenAIWSMode bool
ResponseHeaders http.Header
Duration time.Duration
FirstTokenMs *int
ClientDisconnect bool
ImageCount int
ImageSize string
ImageInputSize string
ImageOutputSize string
ImageOutputSizes []string
ImageSizeSource string
ImageSizeBreakdown map[string]int
VideoCount int
VideoResolution string
ReasoningEffort *string
Stream bool
OpenAIWSMode bool
// UpstreamTerminalEvent is the normalized terminal event observed on an
// upstream Responses WebSocket turn. Empty preserves legacy/non-WS success.
UpstreamTerminalEvent string
ResponseHeaders http.Header
Duration time.Duration
FirstTokenMs *int
ClientDisconnect bool
ImageCount int
ImageSize string
ImageInputSize string
ImageOutputSize string
ImageOutputSizes []string
ImageSizeSource string
ImageSizeBreakdown map[string]int
VideoCount int
VideoResolution string
// VideoDurationSeconds 是提交时请求的生成时长(xAI 按输出秒数计费),已归一化到 1-15 秒。
VideoDurationSeconds int
// WebSearchCalls 是 Codex alpha/search 网页搜索调用次数(每次成功请求为 1)。
@@ -262,6 +265,21 @@ type OpenAIForwardResult struct {
wsReplayInputExists bool
}
// SucceededForScheduling reports whether this result is an upstream success
// that may clear model-scoped transient state. The zero value remains a success
// for existing non-WS callers.
func (r *OpenAIForwardResult) SucceededForScheduling() bool {
if r == nil || !r.OpenAIWSMode || r.UpstreamTerminalEvent == "" {
return true
}
switch r.UpstreamTerminalEvent {
case "response.completed", "response.done":
return true
default:
return false
}
}
// SetActualOpenAIUpstreamEndpoint records the endpoint selected by the current
// forwarding attempt. It covers error paths where no OpenAIForwardResult is
// available for usage and operations logging.
@@ -393,12 +411,14 @@ type OpenAIGatewayService struct {
openaiWSStateStoreOnce sync.Once
openaiSchedulerOnce sync.Once
openaiWSPassthroughDialerOnce sync.Once
openaiModelTransientOnce sync.Once
agentIdentityTaskMu sync.Mutex
openaiWSPool *openAIWSConnPool
openaiWSStateStore OpenAIWSStateStore
openaiScheduler OpenAIAccountScheduler
openaiWSPassthroughDialer openAIWSClientDialer
openaiAccountStats *openAIAccountRuntimeStats
openaiModelTransient *openAIAccountModelTransientState
openaiWSFallbackUntil sync.Map // key: int64(accountID), value: time.Time
openaiAccountRuntimeBlockUntil sync.Map // key: int64(accountID), value: time.Time
@@ -475,6 +495,7 @@ func NewOpenAIGatewayService(
userPlatformQuotaRepo: userPlatformQuotaRepo,
responseHeaderFilter: compileResponseHeaderFilter(cfg),
codexSnapshotThrottle: newAccountWriteThrottle(openAICodexSnapshotPersistMinInterval),
openaiModelTransient: newOpenAIAccountModelTransientState(openAIModelTransientDefaultMax),
}
if rateLimitService != nil {
rateLimitService.SetAccountRuntimeBlocker(svc)
@@ -262,9 +262,9 @@ func (s *OpenAIGatewayService) readUpstreamErrorBody(resp *http.Response) []byte
return body
}
func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, responseBody []byte, requestedModel ...string) {
if len(requestedModel) > 0 {
s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody, requestedModel[0])
func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, responseBody []byte, canonicalModel ...string) {
if len(canonicalModel) > 0 {
s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody, canonicalModel[0])
return
}
s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody)
@@ -385,6 +385,7 @@ func (s *OpenAIGatewayService) handleErrorResponse(
}
if reqModel == "" {
reqModel, _, _ = extractOpenAIRequestMetaFromBody(requestBody)
reqModel = canonicalOpenAIAccountSchedulingModel(account, reqModel)
}
shouldDisable := s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, reqModel)
kind := "http_error"
@@ -670,6 +670,8 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
return acquireTurnLease(turn, preferred, forcePreferredConn)
}
if acquireErr != nil {
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, ingressSessionOriginalModel)
s.handleOpenAIWSDialTransientFailure(ctx, account, canonicalModel, acquireErr)
dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(acquireErr)
logOpenAIWSModeInfo(
"ingress_ws_upstream_acquire_fail account_id=%d turn=%d reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_preferred_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v",
@@ -814,6 +816,8 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
lastEventType = eventType
}
if eventType == "error" {
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel)
s.handleOpenAIWSErrorEventTransientFailure(ctx, account, canonicalModel, lease.HandshakeHeaders(), upstreamMessage)
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage)
s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw)
fallbackReason, _ := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw)
@@ -946,6 +950,8 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
}
}
if isTerminalEvent {
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel)
terminalEvent := s.handleOpenAIWSTerminalTransientFailure(ctx, account, canonicalModel, lease.HandshakeHeaders(), upstreamMessage)
// 客户端已断连时,上游连接的 session 状态不可信,标记 broken 避免回池复用。
if clientDisconnected {
lease.MarkBroken()
@@ -973,17 +979,18 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
}
imageCount := imageCounter.Count()
result := &OpenAIForwardResult{
RequestID: responseID,
Usage: usage,
Model: originalModel,
UpstreamModel: mappedModel,
ServiceTier: extractOpenAIServiceTierFromBody(payload),
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(payload, mappedModel, originalModel), payload, mappedModel),
Stream: reqStream,
OpenAIWSMode: true,
ResponseHeaders: lease.HandshakeHeaders(),
Duration: time.Since(turnStart),
FirstTokenMs: firstTokenMs,
RequestID: responseID,
Usage: usage,
Model: originalModel,
UpstreamModel: mappedModel,
ServiceTier: extractOpenAIServiceTierFromBody(payload),
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(payload, mappedModel, originalModel), payload, mappedModel),
Stream: reqStream,
OpenAIWSMode: true,
UpstreamTerminalEvent: terminalEvent,
ResponseHeaders: lease.HandshakeHeaders(),
Duration: time.Since(turnStart),
FirstTokenMs: firstTokenMs,
}
if replayInput := replayCollector.Items(); len(replayInput) > 0 {
result.wsReplayInput = replayInput
@@ -72,11 +72,11 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT
}
serverErrCh := make(chan error, 1)
turnWSModeCh := make(chan bool, 2)
turnTerminalCh := make(chan string, 2)
hooks := &OpenAIWSIngressHooks{
AfterTurn: func(_ int, result *OpenAIForwardResult, turnErr error) {
if turnErr == nil && result != nil {
turnWSModeCh <- result.OpenAIWSMode
turnTerminalCh <- result.UpstreamTerminalEvent
}
},
}
@@ -146,8 +146,8 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT
secondTurnEvent := readMessage()
require.Equal(t, "response.completed", gjson.GetBytes(secondTurnEvent, "type").String())
require.Equal(t, "resp_ingress_turn_2", gjson.GetBytes(secondTurnEvent, "response.id").String())
require.True(t, <-turnWSModeCh, "首轮 turn 应标记为 WS 模式")
require.True(t, <-turnWSModeCh, "第二轮 turn 应标记为 WS 模式")
require.Equal(t, "response.completed", <-turnTerminalCh, "首轮 turn 应保留成功终态")
require.Equal(t, "response.completed", <-turnTerminalCh, "第二轮 turn 应保留成功终态")
_ = clientConn.Close(coderws.StatusNormalClosure, "done")
@@ -464,10 +464,61 @@ func TestOpenAIGatewayService_Forward_WSv2_RewriteModelAndToolCallsOnCompletedEv
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "resp_model_tool_1", result.RequestID)
require.Equal(t, "response.completed", result.UpstreamTerminalEvent)
require.True(t, result.SucceededForScheduling())
require.Equal(t, "custom-original-model", gjson.GetBytes(rec.Body.Bytes(), "model").String(), "响应模型应回写为原始请求模型")
require.Equal(t, "edit", gjson.GetBytes(rec.Body.Bytes(), "tool_calls.0.function.name").String(), "工具名称应被修正为 OpenCode 规范")
}
func TestOpenAIGatewayService_Forward_WSv2_ResponseFailedIsNotSchedulingSuccess(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
cfg := newOpenAIWSV2TestConfig()
cfg.Security.URLAllowlist.Enabled = false
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
captureConn := &openAIWSCaptureConn{events: [][]byte{
[]byte(`{"type":"response.failed","response":{"id":"resp_failed_1","model":"gpt-5.5","error":{"code":"server_error","message":"Internal error"}}}`),
}}
pool := newOpenAIWSConnPool(cfg)
pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn})
svc := &OpenAIGatewayService{
cfg: cfg,
rateLimitService: NewRateLimitService(transientCooldownAccountRepo{}, nil, cfg, nil, nil),
httpUpstream: &httpUpstreamRecorder{},
cache: &stubGatewayCache{},
openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
toolCorrector: NewCodexToolCorrector(),
openaiWSPool: pool,
}
account := &Account{
ID: 1302,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{"api_key": "sk-test"},
Extra: map[string]any{"responses_websockets_v2_enabled": true},
}
svc.recordOpenAIAccountModelTransientFailure(account, "gpt-5.5", time.Now())
result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.5","stream":false,"input":"hello"}`))
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "response.failed", result.UpstreamTerminalEvent)
require.False(t, result.SucceededForScheduling())
require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5"))
}
func TestOpenAIWSPayloadString_OnlyAcceptsStringValues(t *testing.T) {
payload := map[string]any{
"type": nil,
@@ -204,6 +204,98 @@ func isOpenAIWSTerminalEvent(eventType string) bool {
}
}
func normalizeOpenAIWSTerminalEvent(eventType string) string {
switch strings.TrimSpace(eventType) {
case "response.completed":
return "response.completed"
case "response.done":
return "response.done"
case "response.failed":
return "response.failed"
case "response.incomplete":
return "response.incomplete"
case "response.cancelled", "response.canceled":
return "response.cancelled"
default:
return ""
}
}
func openAIWSPayloadTransientStatus(payload []byte) int {
if len(payload) == 0 {
return 0
}
status := int(gjson.GetBytes(payload, "response.error.status_code").Int())
if status == 0 {
status = int(gjson.GetBytes(payload, "response.error.status").Int())
}
if status == 0 {
status = int(gjson.GetBytes(payload, "error.status_code").Int())
}
if status == 0 {
status = int(gjson.GetBytes(payload, "error.status").Int())
}
if shouldCooldownOpenAITransientUpstreamError(status, payload) {
return status
}
if status != 0 {
return 0
}
code := strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "response.error.code").String()))
errType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "response.error.type").String()))
if code == "" {
code = strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "error.code").String()))
}
if errType == "" {
errType = strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "error.type").String()))
}
switch {
case code == "server_is_overloaded", code == "slow_down":
return http.StatusServiceUnavailable
case strings.Contains(code, "server_error"),
strings.Contains(code, "internal_error"),
strings.Contains(code, "upstream_error"),
strings.Contains(errType, "server_error"),
strings.Contains(errType, "internal_error"),
strings.Contains(errType, "upstream_error"):
return http.StatusInternalServerError
default:
return 0
}
}
func (s *OpenAIGatewayService) handleOpenAIWSTerminalTransientFailure(ctx context.Context, account *Account, canonicalModel string, headers http.Header, payload []byte) string {
eventType, _, _ := parseOpenAIWSEventEnvelope(payload)
terminalEvent := normalizeOpenAIWSTerminalEvent(eventType)
if terminalEvent != "response.failed" {
return terminalEvent
}
status := openAIWSPayloadTransientStatus(payload)
if status != 0 {
s.handleOpenAIAccountUpstreamError(ctx, account, status, headers, payload, canonicalModel)
}
return terminalEvent
}
func (s *OpenAIGatewayService) handleOpenAIWSErrorEventTransientFailure(ctx context.Context, account *Account, canonicalModel string, headers http.Header, payload []byte) {
eventType, _, _ := parseOpenAIWSEventEnvelope(payload)
if eventType != "error" {
return
}
status := openAIWSPayloadTransientStatus(payload)
if status != 0 {
s.handleOpenAIAccountUpstreamError(ctx, account, status, headers, payload, canonicalModel)
}
}
func (s *OpenAIGatewayService) handleOpenAIWSDialTransientFailure(ctx context.Context, account *Account, canonicalModel string, err error) {
var dialErr *openAIWSDialError
if !errors.As(err, &dialErr) || dialErr == nil || !shouldCooldownOpenAITransientUpstreamError(dialErr.StatusCode, dialErr.ResponseBody) {
return
}
s.handleOpenAIAccountUpstreamError(ctx, account, dialErr.StatusCode, dialErr.ResponseHeaders, dialErr.ResponseBody, canonicalModel)
}
func isOpenAIWSTokenEvent(eventType string) bool {
eventType = strings.TrimSpace(eventType)
if eventType == "" {
@@ -440,7 +532,7 @@ func (s *OpenAIGatewayService) resolveAccountByPreviousResponseIDForCapability(
if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, latest); paused {
return 0, nil, "", nil
}
if s.isOpenAIAccountRuntimeBlocked(latest) {
if s.isOpenAIAccountRequestRuntimeBlocked(latest, requestedModel) {
_ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
return 0, nil, "", nil
}
@@ -1,9 +1,13 @@
package service
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
@@ -107,6 +111,81 @@ func TestOpenAIWSCyberPolicyMark_NonCyberPayload(t *testing.T) {
require.False(t, hit, "detectOpenAICyberPolicy should return false for non-cyber_policy error code")
}
func TestOpenAIForwardResultSucceededForScheduling_TerminalEvents(t *testing.T) {
tests := []struct {
name string
result *OpenAIForwardResult
expected bool
}{
{name: "nil legacy result", result: nil, expected: true},
{name: "non websocket zero value", result: &OpenAIForwardResult{}, expected: true},
{name: "websocket legacy empty terminal", result: &OpenAIForwardResult{OpenAIWSMode: true}, expected: true},
{name: "completed", result: &OpenAIForwardResult{OpenAIWSMode: true, UpstreamTerminalEvent: "response.completed"}, expected: true},
{name: "done", result: &OpenAIForwardResult{OpenAIWSMode: true, UpstreamTerminalEvent: "response.done"}, expected: true},
{name: "failed", result: &OpenAIForwardResult{OpenAIWSMode: true, UpstreamTerminalEvent: "response.failed"}, expected: false},
{name: "incomplete", result: &OpenAIForwardResult{OpenAIWSMode: true, UpstreamTerminalEvent: "response.incomplete"}, expected: false},
{name: "cancelled", result: &OpenAIForwardResult{OpenAIWSMode: true, UpstreamTerminalEvent: "response.cancelled"}, expected: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.expected, tt.result.SucceededForScheduling())
})
}
}
func TestOpenAIWSTerminalEvent_ResponseFailedRecordsModelTransient(t *testing.T) {
svc := &OpenAIGatewayService{}
svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil)
account := &Account{ID: 5201, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
payload := []byte(`{"type":"response.failed","response":{"error":{"code":"server_error","message":"Internal error"}}}`)
for range 2 {
terminalEvent := svc.handleOpenAIWSTerminalTransientFailure(context.Background(), account, "gpt-5.5", http.Header{}, payload)
require.Equal(t, "response.failed", terminalEvent)
}
require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5"))
}
func TestOpenAIWSErrorEvent_ServerErrorRecordsModelTransient(t *testing.T) {
svc := &OpenAIGatewayService{}
svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil)
account := &Account{ID: 5203, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
payload := []byte(`{"type":"error","error":{"code":"server_error","type":"server_error","message":"Internal error"}}`)
for range 2 {
svc.handleOpenAIWSErrorEventTransientFailure(context.Background(), account, "gpt-5.5", http.Header{}, payload)
}
require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5"))
}
func TestOpenAIWSPayloadTransientStatus_Explicit529IsNotModelTransient(t *testing.T) {
payload := []byte(`{"type":"response.failed","response":{"error":{"status_code":529,"code":"server_error","message":"overloaded"}}}`)
require.Zero(t, openAIWSPayloadTransientStatus(payload))
}
func TestOpenAIWSDial5xxRecordsModelTransient(t *testing.T) {
svc := &OpenAIGatewayService{}
svc.rateLimitService = NewRateLimitService(transientCooldownAccountRepo{}, nil, &config.Config{}, nil, nil)
account := &Account{ID: 5202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
dialErr := &openAIWSDialError{
StatusCode: http.StatusBadGateway,
ResponseHeaders: http.Header{"X-Request-Id": []string{"req-ws-502"}},
ResponseBody: []byte(`{"error":{"message":"bad gateway"}}`),
}
for range 2 {
svc.handleOpenAIWSDialTransientFailure(context.Background(), account, "gpt-5.5", dialErr)
}
require.Eventually(t, func() bool {
return svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5")
}, time.Second, 10*time.Millisecond)
}
// TestIsOpenAIWSTokenEvent_DisjointWithTerminal 守护「token 事件集合与终止事件集合互斥」的不变量。
// firstTokenMs 的计算依赖于 isTokenEvent && !isTerminalEvent
// 若两者再次出现交集,则 issue #2651 描述的 latency 误报会重现。
@@ -200,6 +200,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
}
return nil, &agentIdentityTaskRecoveredError{}
}
s.handleOpenAIWSDialTransientFailure(ctx, account, mappedModel, err)
dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(err)
logOpenAIWSModeInfo(
"acquire_fail account_id=%d account_type=%s transport=%s reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_new_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v",
@@ -352,6 +353,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
flushedBufferedEventCount := 0
firstEventType := ""
lastEventType := ""
upstreamTerminalEvent := ""
var flusher http.Flusher
if reqStream {
@@ -572,6 +574,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
}
if eventType == "error" {
s.handleOpenAIWSErrorEventTransientFailure(ctx, account, mappedModel, lease.HandshakeHeaders(), message)
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message)
s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw)
errMsg := strings.TrimSpace(errMsgRaw)
@@ -670,6 +673,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
}
if isTerminalEvent {
upstreamTerminalEvent = s.handleOpenAIWSTerminalTransientFailure(ctx, account, mappedModel, lease.HandshakeHeaders(), message)
// A terminal event must be the final JSON document in its WS message.
// Ignore any tail for the completed client turn, but never reuse the
// ambiguous upstream connection for another request.
@@ -741,19 +745,20 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
)
return &OpenAIForwardResult{
RequestID: responseID,
Usage: *usage,
Model: originalModel,
UpstreamModel: mappedModel,
ImageCount: imageCounter.Count(),
ImageOutputSizes: imageCounter.Sizes(),
ServiceTier: extractOpenAIServiceTier(reqBody),
ReasoningEffort: extractOpenAIReasoningEffort(reqBody, mappedModel, originalModel),
Stream: reqStream,
OpenAIWSMode: true,
ResponseHeaders: lease.HandshakeHeaders(),
Duration: time.Since(startTime),
FirstTokenMs: firstTokenMs,
RequestID: responseID,
Usage: *usage,
Model: originalModel,
UpstreamModel: mappedModel,
ImageCount: imageCounter.Count(),
ImageOutputSizes: imageCounter.Sizes(),
ServiceTier: extractOpenAIServiceTier(reqBody),
ReasoningEffort: extractOpenAIReasoningEffort(reqBody, mappedModel, originalModel),
Stream: reqStream,
OpenAIWSMode: true,
UpstreamTerminalEvent: upstreamTerminalEvent,
ResponseHeaders: lease.HandshakeHeaders(),
Duration: time.Since(startTime),
FirstTokenMs: firstTokenMs,
}, nil
}
@@ -235,6 +235,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes))
if account.Platform == PlatformGrok {
s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
} else if shouldCooldownOpenAITransientUpstreamError(resp.StatusCode, respBody) {
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel)
s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, canonicalModel)
}
upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
if upstreamMsg == "" {
@@ -258,6 +261,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
replayCollector := &openAIWSToolCallReplayCollector{}
firstEventType := ""
lastEventType := ""
upstreamTerminalEvent := ""
sawDone := false
wroteDownstream := false
clientDisconnected := false
@@ -275,17 +279,18 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
resultWithUsage := func() *OpenAIForwardResult {
imageCount := imageCounter.Count()
result := &OpenAIForwardResult{
RequestID: responseID,
Usage: usage,
Model: originalModel,
UpstreamModel: mappedModel,
ServiceTier: extractOpenAIServiceTierFromBody(body),
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(body, mappedModel, originalModel), body, mappedModel),
Stream: reqStream,
OpenAIWSMode: true,
ResponseHeaders: cloneHeader(resp.Header),
Duration: time.Since(turnStart),
FirstTokenMs: firstTokenMs,
RequestID: responseID,
Usage: usage,
Model: originalModel,
UpstreamModel: mappedModel,
ServiceTier: extractOpenAIServiceTierFromBody(body),
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(body, mappedModel, originalModel), body, mappedModel),
Stream: reqStream,
OpenAIWSMode: true,
UpstreamTerminalEvent: upstreamTerminalEvent,
ResponseHeaders: cloneHeader(resp.Header),
Duration: time.Since(turnStart),
FirstTokenMs: firstTokenMs,
}
if replayInput := replayCollector.Items(); len(replayInput) > 0 {
result.wsReplayInput = replayInput
@@ -384,6 +389,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
}
if eventType == "error" {
s.handleOpenAIWSErrorEventTransientFailure(ctx, account, canonicalOpenAIAccountSchedulingModel(account, originalModel), resp.Header, upstreamMessage)
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage)
s.persistOpenAIWSRateLimitSignal(ctx, account, resp.Header, upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw)
errMessage := strings.TrimSpace(errMsgRaw)
@@ -393,6 +399,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
return resultWithUsage(), errors.New(errMessage)
}
if isOpenAIWSTerminalEvent(eventType) {
upstreamTerminalEvent = s.handleOpenAIWSTerminalTransientFailure(ctx, account, canonicalOpenAIAccountSchedulingModel(account, originalModel), resp.Header, upstreamMessage)
terminalEventCount++
firstTokenMsValue := -1
if firstTokenMs != nil {
@@ -239,6 +239,52 @@ func TestOpenAIGatewayService_Forward_WSv2Handshake429PersistsRateLimit(t *testi
require.Contains(t, repo.updateExtra[0], "codex_usage_updated_at")
}
func TestOpenAIGatewayService_Forward_WSv2Handshake502RecordsModelTransient(t *testing.T) {
gin.SetMode(gin.TestMode)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("x-request-id", "req-ws-502")
w.WriteHeader(http.StatusBadGateway)
_, _ = w.Write([]byte(`{"error":{"type":"server_error","message":"bad gateway"}}`))
}))
defer server.Close()
cfg := newOpenAIWSV2TestConfig()
cfg.Security.URLAllowlist.Enabled = false
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
account := Account{
ID: 504,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{"api_key": "sk-test", "base_url": server.URL},
Extra: map[string]any{"responses_websockets_v2_enabled": true},
}
svc := &OpenAIGatewayService{
cfg: cfg,
rateLimitService: NewRateLimitService(transientCooldownAccountRepo{}, nil, cfg, nil, nil),
httpUpstream: &httpUpstreamRecorder{},
cache: &stubGatewayCache{},
openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
toolCorrector: NewCodexToolCorrector(),
}
body := []byte(`{"model":"gpt-5.5","stream":false,"input":"hello"}`)
for range 2 {
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
result, err := svc.Forward(context.Background(), c, &account, body)
require.Error(t, err)
require.Nil(t, result)
}
require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(&account, "gpt-5.5"))
}
func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_ErrorEventUsageLimitPersistsRateLimit(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -386,7 +386,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
if errors.As(err, &handshakeErr) && handshakeErr != nil {
responseBody = handshakeErr.Body
}
dialErr := &openAIWSDialError{StatusCode: statusCode, ResponseBody: responseBody, Err: err}
dialErr := &openAIWSDialError{StatusCode: statusCode, ResponseHeaders: cloneHeader(handshakeHeaders), ResponseBody: responseBody, Err: err}
if s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidWSDialError(dialErr) && !agentTaskRecoveryTried {
agentTaskRecoveryTried = true
if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil {
@@ -400,6 +400,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
statusCode,
truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen),
)
s.handleOpenAIWSDialTransientFailure(ctx, account, capturedSessionModel, dialErr)
if statusCode == http.StatusTooManyRequests {
s.persistOpenAIWSRateLimitSignal(ctx, account, handshakeHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error()))
return &UpstreamFailoverError{
@@ -571,14 +572,15 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
CacheReadInputTokens: turn.Usage.CacheReadInputTokens,
ImageOutputTokens: turn.Usage.ImageOutputTokens,
},
Model: turn.RequestModel,
ServiceTier: usageMeta.serviceTier.Load(),
ReasoningEffort: usageMeta.reasoningEffort.Load(),
Stream: true,
OpenAIWSMode: true,
ResponseHeaders: cloneHeader(handshakeHeaders),
Duration: turn.Duration,
FirstTokenMs: turn.FirstTokenMs,
Model: turn.RequestModel,
ServiceTier: usageMeta.serviceTier.Load(),
ReasoningEffort: usageMeta.reasoningEffort.Load(),
Stream: true,
OpenAIWSMode: true,
UpstreamTerminalEvent: normalizeOpenAIWSTerminalEvent(turn.TerminalEventType),
ResponseHeaders: cloneHeader(handshakeHeaders),
Duration: turn.Duration,
FirstTokenMs: turn.FirstTokenMs,
}
logOpenAIWSV2Passthrough(
"relay_turn_completed account_id=%d turn=%d request_id=%s terminal_event=%s duration_ms=%d first_token_ms=%d input_tokens=%d output_tokens=%d cache_read_tokens=%d",
@@ -597,10 +599,17 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
}
},
BeforeWriteClient: func(msgType coderws.MessageType, payload []byte, wroteDownstream bool) error {
if msgType != coderws.MessageText || wroteDownstream {
if msgType != coderws.MessageText {
return nil
}
if eventType, _, _ := parseOpenAIWSEventEnvelope(payload); eventType != "error" {
eventType, _, _ := parseOpenAIWSEventEnvelope(payload)
if isOpenAIWSTerminalEvent(eventType) {
s.handleOpenAIWSTerminalTransientFailure(ctx, account, capturedSessionModel, handshakeHeaders, payload)
}
if eventType == "error" {
s.handleOpenAIWSErrorEventTransientFailure(ctx, account, capturedSessionModel, handshakeHeaders, payload)
}
if wroteDownstream || eventType != "error" {
return nil
}
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(payload)
@@ -646,14 +655,15 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
CacheReadInputTokens: relayResult.Usage.CacheReadInputTokens,
ImageOutputTokens: relayResult.Usage.ImageOutputTokens,
},
Model: relayResult.RequestModel,
ServiceTier: usageMeta.serviceTier.Load(),
ReasoningEffort: usageMeta.reasoningEffort.Load(),
Stream: true,
OpenAIWSMode: true,
ResponseHeaders: cloneHeader(handshakeHeaders),
Duration: relayResult.Duration,
FirstTokenMs: relayResult.FirstTokenMs,
Model: relayResult.RequestModel,
ServiceTier: usageMeta.serviceTier.Load(),
ReasoningEffort: usageMeta.reasoningEffort.Load(),
Stream: true,
OpenAIWSMode: true,
UpstreamTerminalEvent: normalizeOpenAIWSTerminalEvent(relayResult.TerminalEventType),
ResponseHeaders: cloneHeader(handshakeHeaders),
Duration: relayResult.Duration,
FirstTokenMs: relayResult.FirstTokenMs,
}
turnCount := int(completedTurns.Load())