mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-08-31 01:13:06 +08:00
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:
@@ -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"}]}`)
|
||||
|
||||
|
||||
@@ -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())
|
||||
|
||||
Reference in New Issue
Block a user