diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index ea9169e0b7..1262845cea 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -819,6 +819,8 @@ type GatewayConfig struct { ImageStreamDataIntervalTimeout int `mapstructure:"image_stream_data_interval_timeout"` // ImageStreamKeepaliveInterval: 图片流式 keepalive 间隔（秒），0表示禁用 ImageStreamKeepaliveInterval int `mapstructure:"image_stream_keepalive_interval"` + // ImageNonstreamKeepaliveInterval: 图片非流式 JSON keepalive 间隔（秒），0表示禁用 + ImageNonstreamKeepaliveInterval int `mapstructure:"image_nonstream_keepalive_interval"` // MaxLineSize: 上游 SSE 单行最大字节数（0使用默认值） MaxLineSize int `mapstructure:"max_line_size"` @@ -2037,6 +2039,7 @@ func setDefaults() { viper.SetDefault("gateway.stream_keepalive_interval", 10) viper.SetDefault("gateway.image_stream_data_interval_timeout", 900) viper.SetDefault("gateway.image_stream_keepalive_interval", 10) + viper.SetDefault("gateway.image_nonstream_keepalive_interval", 0) viper.SetDefault("gateway.max_line_size", 500*1024*1024) viper.SetDefault("gateway.scheduling.sticky_session_max_waiting", 3) viper.SetDefault("gateway.scheduling.sticky_session_wait_timeout", 120*time.Second) @@ -2723,6 +2726,13 @@ func (c *Config) Validate() error { (c.Gateway.ImageStreamKeepaliveInterval < 5 || c.Gateway.ImageStreamKeepaliveInterval > 60) { return fmt.Errorf("gateway.image_stream_keepalive_interval must be 0 or between 5-60 seconds") } + if c.Gateway.ImageNonstreamKeepaliveInterval < 0 { + return fmt.Errorf("gateway.image_nonstream_keepalive_interval must be non-negative") + } + if c.Gateway.ImageNonstreamKeepaliveInterval != 0 && + (c.Gateway.ImageNonstreamKeepaliveInterval < 5 || c.Gateway.ImageNonstreamKeepaliveInterval > 60) { + return fmt.Errorf("gateway.image_nonstream_keepalive_interval must be 0 or between 5-60 seconds") + } // 兼容旧键 sticky_previous_response_ttl_seconds if c.Gateway.OpenAIWS.StickyResponseIDTTLSeconds <= 0 && c.Gateway.OpenAIWS.StickyPreviousResponseTTLSeconds > 0 { c.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = c.Gateway.OpenAIWS.StickyPreviousResponseTTLSeconds diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index 4eea3a2840..2f9defb80e 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -259,6 +259,15 @@ func TestLoadOpenAIResponseHeaderTimeoutFromEnv(t *testing.T) { require.Equal(t, 1800, cfg.Gateway.OpenAIResponseHeaderTimeout) } +func TestLoadImageNonstreamKeepaliveFromEnv(t *testing.T) { + resetViperWithJWTSecret(t) + t.Setenv("GATEWAY_IMAGE_NONSTREAM_KEEPALIVE_INTERVAL", "15") + + cfg, err := Load() + require.NoError(t, err) + require.Equal(t, 15, cfg.Gateway.ImageNonstreamKeepaliveInterval) +} + func TestLoadOpenAIWSStickyTTLCompatibility(t *testing.T) { resetViperWithJWTSecret(t) t.Setenv("GATEWAY_OPENAI_WS_STICKY_RESPONSE_ID_TTL_SECONDS", "0") @@ -1429,6 +1438,16 @@ func TestValidateConfigErrors(t *testing.T) { mutate: func(c *Config) { c.Gateway.ImageStreamKeepaliveInterval = -1 }, wantErr: "gateway.image_stream_keepalive_interval must be non-negative", }, + { + name: "gateway image nonstream keepalive range", + mutate: func(c *Config) { c.Gateway.ImageNonstreamKeepaliveInterval = 4 }, + wantErr: "gateway.image_nonstream_keepalive_interval", + }, + { + name: "gateway image nonstream keepalive negative", + mutate: func(c *Config) { c.Gateway.ImageNonstreamKeepaliveInterval = -1 }, + wantErr: "gateway.image_nonstream_keepalive_interval must be non-negative", + }, { name: "gateway image stream data interval range", mutate: func(c *Config) { c.Gateway.ImageStreamDataIntervalTimeout = 30 }, @@ -1997,6 +2016,9 @@ func TestLoad_DefaultGatewayImageStreamConfig(t *testing.T) { if cfg.Gateway.ImageStreamKeepaliveInterval != 10 { t.Fatalf("image_stream_keepalive_interval = %d, want 10", cfg.Gateway.ImageStreamKeepaliveInterval) } + if cfg.Gateway.ImageNonstreamKeepaliveInterval != 0 { + t.Fatalf("image_nonstream_keepalive_interval = %d, want 0", cfg.Gateway.ImageNonstreamKeepaliveInterval) + } if cfg.Gateway.ImageConcurrency.Enabled { t.Fatalf("image_concurrency.enabled = true, want false") } diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 7d6dc2c17a..231363a8a3 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -2091,7 +2091,8 @@ func openAIForwardErrorAlreadyCommunicated(c *gin.Context, writerSizeBeforeForwa } // 与快照同口径：排除 compact 心跳字节，避免"仅心跳写出"被误判为 // 响应已写出（#3887）。 - if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward { + if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward || + service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward { return false } diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 5868f7f35b..c5982fb7d1 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -142,6 +142,9 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { failedAccountIDs := make(map[int64]struct{}) sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError + stopJSONKeepalive := func() {} + jsonKeepaliveStarted := false + defer func() { stopJSONKeepalive() }() for { reqLog.Debug("openai.images.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs))) @@ -210,8 +213,12 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { } service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds()) + if !parsed.Stream && !jsonKeepaliveStarted { + stopJSONKeepalive = service.StartOpenAIImagesJSONKeepalive(c, h.openAIImagesJSONKeepaliveInterval()) + jsonKeepaliveStarted = true + } forwardStart := time.Now() - writerSizeBeforeForward := c.Writer.Size() + writerSizeBeforeForward := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) result, err := func() (*service.OpenAIForwardResult, error) { defer func() { if accountReleaseFunc != nil { @@ -258,7 +265,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) - if c.Writer.Size() != writerSizeBeforeForward { + if service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeForward { reqLog.Warn("openai.images.upstream_failover_skipped_after_flush", zap.Int64("account_id", account.ID), zap.Int("upstream_status", failoverErr.StatusCode), @@ -383,6 +390,13 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { } } +func (h *OpenAIGatewayHandler) openAIImagesJSONKeepaliveInterval() time.Duration { + if h.cfg == nil || h.cfg.Gateway.ImageNonstreamKeepaliveInterval <= 0 { + return 0 + } + return time.Duration(h.cfg.Gateway.ImageNonstreamKeepaliveInterval) * time.Second +} + func isMultipartImagesContentType(contentType string) bool { return strings.HasPrefix(strings.ToLower(strings.TrimSpace(contentType)), "multipart/form-data") } diff --git a/backend/internal/service/openai_images_json_keepalive.go b/backend/internal/service/openai_images_json_keepalive.go new file mode 100644 index 0000000000..0d8b24fb0a --- /dev/null +++ b/backend/internal/service/openai_images_json_keepalive.go @@ -0,0 +1,268 @@ +package service + +import ( + "bufio" + "errors" + "net" + "net/http" + "sync" + "time" + + "github.com/gin-gonic/gin" +) + +const openAIImagesJSONKeepaliveKey = "openai_images_json_keepalive" + +// openAIImagesJSONKeepalive keeps non-streaming Images API requests alive while +// an OAuth upstream is producing SSE internally. JSON permits leading +// whitespace, so each heartbeat remains compatible with clients expecting one +// final JSON document. +// +// Once the first heartbeat is sent, the HTTP status is committed as 200. Late +// upstream errors are still returned as an OpenAI-compatible JSON error body, +// matching the status tradeoff used by the compact SSE keepalive path. +type openAIImagesJSONKeepalive struct { + mu sync.Mutex + writer gin.ResponseWriter + started bool + stopped bool + bytes int + stop chan struct{} +} + +// StartOpenAIImagesJSONKeepalive starts whitespace heartbeats for a +// non-streaming Images request. A non-positive interval disables the feature. +func StartOpenAIImagesJSONKeepalive(c *gin.Context, interval time.Duration) func() { + if c == nil || c.Writer == nil || interval <= 0 { + return func() {} + } + originalWriter := c.Writer + k := &openAIImagesJSONKeepalive{ + writer: originalWriter, + stop: make(chan struct{}), + } + c.Set(openAIImagesJSONKeepaliveKey, k) + wrappedWriter := &openAIImagesJSONKeepaliveWriter{ResponseWriter: originalWriter, k: k} + c.Writer = wrappedWriter + + var reqDone <-chan struct{} + if c.Request != nil { + reqDone = c.Request.Context().Done() + } + go func() { + timer := time.NewTimer(interval) + defer timer.Stop() + for { + select { + case <-k.stop: + return + case <-reqDone: + return + case <-timer.C: + } + if !k.beat() { + return + } + timer.Reset(interval) + } + }() + + return func() { + k.Stop() + if current, ok := c.Writer.(*openAIImagesJSONKeepaliveWriter); ok && current == wrappedWriter { + c.Writer = originalWriter + } + } +} + +func (k *openAIImagesJSONKeepalive) beat() bool { + k.mu.Lock() + defer k.mu.Unlock() + if k.stopped { + return false + } + if !k.started { + header := k.writer.Header() + header.Set("Content-Type", "application/json; charset=utf-8") + header.Set("Cache-Control", "no-cache") + header.Set("X-Accel-Buffering", "no") + k.writer.WriteHeader(http.StatusOK) + k.started = true + } + n, err := k.writer.Write([]byte(" \n")) + k.bytes += n + if err != nil { + k.stopped = true + return false + } + k.writer.Flush() + return true +} + +func (k *openAIImagesJSONKeepalive) Stop() { + k.mu.Lock() + k.markStoppedLocked() + k.mu.Unlock() +} + +func (k *openAIImagesJSONKeepalive) markStoppedLocked() { + if k.stopped { + return + } + k.stopped = true + close(k.stop) +} + +// StopOpenAIImagesJSONKeepaliveCommitted stops heartbeats and reports whether +// they already committed a 200 response. +func StopOpenAIImagesJSONKeepaliveCommitted(c *gin.Context) bool { + k := openAIImagesJSONKeepaliveFromContext(c) + if k == nil { + return false + } + k.mu.Lock() + k.markStoppedLocked() + committed := k.started + k.mu.Unlock() + return committed +} + +// OpenAIImagesJSONKeepaliveAdjustedWrittenSize excludes heartbeat whitespace +// from response-size checks so account retry and failover remain available. +func OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c *gin.Context) int { + if c == nil || c.Writer == nil { + return -1 + } + k := openAIImagesJSONKeepaliveFromContext(c) + if k == nil { + return c.Writer.Size() + } + k.mu.Lock() + defer k.mu.Unlock() + size := k.writer.Size() + if size < 0 { + return size + } + if real := size - k.bytes; real > 0 { + return real + } + return -1 +} + +func openAIImagesJSONKeepaliveFromContext(c *gin.Context) *openAIImagesJSONKeepalive { + if c == nil { + return nil + } + value, ok := c.Get(openAIImagesJSONKeepaliveKey) + if !ok { + return nil + } + k, _ := value.(*openAIImagesJSONKeepalive) + return k +} + +type openAIImagesJSONKeepaliveWriter struct { + gin.ResponseWriter + k *openAIImagesJSONKeepalive +} + +func (w *openAIImagesJSONKeepaliveWriter) suspend() { + if w.k != nil { + w.k.Stop() + } +} + +func (w *openAIImagesJSONKeepaliveWriter) Header() http.Header { + w.suspend() + if w.ResponseWriter == nil { + return http.Header{} + } + return w.ResponseWriter.Header() +} + +func (w *openAIImagesJSONKeepaliveWriter) Write(data []byte) (int, error) { + w.suspend() + if w.ResponseWriter == nil { + return 0, nil + } + return w.ResponseWriter.Write(data) +} + +func (w *openAIImagesJSONKeepaliveWriter) WriteString(s string) (int, error) { + w.suspend() + if w.ResponseWriter == nil { + return 0, nil + } + return w.ResponseWriter.WriteString(s) +} + +func (w *openAIImagesJSONKeepaliveWriter) WriteHeader(code int) { + w.suspend() + if w.ResponseWriter != nil { + w.ResponseWriter.WriteHeader(code) + } +} + +func (w *openAIImagesJSONKeepaliveWriter) WriteHeaderNow() { + w.suspend() + if w.ResponseWriter != nil { + w.ResponseWriter.WriteHeaderNow() + } +} + +func (w *openAIImagesJSONKeepaliveWriter) Flush() { + w.suspend() + if w.ResponseWriter != nil { + w.ResponseWriter.Flush() + } +} + +func (w *openAIImagesJSONKeepaliveWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + if w.ResponseWriter == nil { + return nil, nil, errors.New("response writer released") + } + return w.ResponseWriter.Hijack() +} + +func (w *openAIImagesJSONKeepaliveWriter) CloseNotify() <-chan bool { + if w.ResponseWriter == nil { + ch := make(chan bool) + close(ch) + return ch + } + return w.ResponseWriter.CloseNotify() +} + +func (w *openAIImagesJSONKeepaliveWriter) Pusher() http.Pusher { + if w.ResponseWriter == nil { + return nil + } + return w.ResponseWriter.Pusher() +} + +func (w *openAIImagesJSONKeepaliveWriter) Status() int { + if w.k == nil || w.ResponseWriter == nil { + return 0 + } + w.k.mu.Lock() + defer w.k.mu.Unlock() + return w.ResponseWriter.Status() +} + +func (w *openAIImagesJSONKeepaliveWriter) Size() int { + if w.k == nil || w.ResponseWriter == nil { + return 0 + } + w.k.mu.Lock() + defer w.k.mu.Unlock() + return w.ResponseWriter.Size() +} + +func (w *openAIImagesJSONKeepaliveWriter) Written() bool { + if w.k == nil || w.ResponseWriter == nil { + return false + } + w.k.mu.Lock() + defer w.k.mu.Unlock() + return w.ResponseWriter.Written() +} diff --git a/backend/internal/service/openai_images_json_keepalive_test.go b/backend/internal/service/openai_images_json_keepalive_test.go new file mode 100644 index 0000000000..a7207b3adb --- /dev/null +++ b/backend/internal/service/openai_images_json_keepalive_test.go @@ -0,0 +1,252 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestOpenAIImagesJSONKeepalive_PreservesValidJSONResponse(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + originalWriter := c.Writer + + stop := StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + waitForOpenAIImagesJSONKeepalive(t, c) + require.Equal(t, -1, OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)) + + c.JSON(http.StatusOK, gin.H{"data": []gin.H{{"b64_json": "aW1hZ2U="}}}) + stop() + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, "application/json; charset=utf-8", rec.Header().Get("Content-Type")) + require.Equal(t, "no", rec.Header().Get("X-Accel-Buffering")) + require.True(t, rec.Flushed) + require.True(t, json.Valid(rec.Body.Bytes()), rec.Body.String()) + require.Equal(t, "aW1hZ2U=", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) + require.Greater(t, OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c), 0) + require.Same(t, originalWriter, c.Writer) +} + +func TestOpenAIImagesJSONKeepalive_DisabledIsNoop(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + originalWriter := c.Writer + + stop := StartOpenAIImagesJSONKeepalive(c, 0) + stop() + c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"message": "invalid request"}}) + + require.Same(t, originalWriter, c.Writer) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Equal(t, "invalid request", gjson.Get(rec.Body.String(), "error.message").String()) +} + +func TestOpenAIImagesJSONKeepalive_FastErrorPreservesStatus(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + stop := StartOpenAIImagesJSONKeepalive(c, time.Second) + wrote := writeOpenAIImagesUpstreamErrorResponse(c, &OpenAIImagesUpstreamError{ + StatusCode: http.StatusBadRequest, + ErrorType: "invalid_request_error", + Message: "invalid size", + }) + stop() + + require.True(t, wrote) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.False(t, strings.HasPrefix(rec.Body.String(), " \n")) + require.Equal(t, "invalid size", gjson.Get(rec.Body.String(), "error.message").String()) +} + +func TestOpenAIImagesJSONKeepalive_LateErrorRemainsJSON(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + stop := StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + defer stop() + waitForOpenAIImagesJSONKeepalive(t, c) + + wrote := writeOpenAIImagesUpstreamErrorResponse(c, &OpenAIImagesUpstreamError{ + StatusCode: http.StatusBadRequest, + ErrorType: "image_generation_user_error", + Code: "moderation_blocked", + Message: "request rejected", + }) + + require.True(t, wrote) + require.Equal(t, http.StatusOK, rec.Code, "heartbeat already committed the status") + require.True(t, json.Valid(rec.Body.Bytes()), rec.Body.String()) + require.Equal(t, "moderation_blocked", gjson.Get(rec.Body.String(), "error.code").String()) + require.Equal(t, "request rejected", gjson.Get(rec.Body.String(), "error.message").String()) +} + +func TestOpenAIImagesJSONKeepalive_DoesNotBlockFailoverDetection(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + stop := StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + waitForOpenAIImagesJSONKeepalive(t, c) + + before := OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) + require.Equal(t, -1, before) + require.True(t, c.Writer.Written()) + require.Equal(t, before, OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)) + stop() + require.True(t, strings.TrimSpace(rec.Body.String()) == "") +} + +func TestOpenAIImagesJSONKeepalive_KeepsOAuthNonStreamResponseValid(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + reader, writer := io.Pipe() + go func() { + time.Sleep(20 * time.Millisecond) + _, _ = io.WriteString(writer, + "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000000,\"output\":[{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2U=\",\"output_format\":\"png\"}]}}\n\n"+ + "data: [DONE]\n\n", + ) + _ = writer.Close() + }() + + stop := StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: reader, + } + svc := &OpenAIGatewayService{} + _, imageCount, _, err := svc.handleOpenAIImagesOAuthNonStreamingResponse(resp, c, "b64_json", "gpt-image-2") + stop() + + require.NoError(t, err) + require.Equal(t, 1, imageCount) + require.True(t, rec.Flushed) + require.True(t, strings.HasPrefix(rec.Body.String(), " \n"), rec.Body.String()) + require.True(t, json.Valid(rec.Body.Bytes()), rec.Body.String()) + require.Equal(t, "aW1hZ2U=", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) +} + +func TestOpenAIImagesJSONKeepaliveWriter_NilGuards(t *testing.T) { + w := &openAIImagesJSONKeepaliveWriter{} + require.NotPanics(t, func() { + require.NotNil(t, w.Header()) + _, _ = w.Write([]byte("test")) + _, _ = w.WriteString("test") + w.WriteHeader(http.StatusOK) + w.WriteHeaderNow() + w.Flush() + require.Equal(t, 0, w.Status()) + require.Equal(t, 0, w.Size()) + require.False(t, w.Written()) + require.Nil(t, w.Pusher()) + }) + + conn, _, err := w.Hijack() + require.Error(t, err) + require.Nil(t, conn) + select { + case <-w.CloseNotify(): + default: + t.Fatal("nil writer CloseNotify channel should be closed") + } +} + +// 回归：failover 第 2+ 轮时，上一轮心跳残留的空白字节不得被误判为"已写响应"， +// 可重试上游错误必须仍转换为 UpstreamFailoverError（而非裸错误吞掉换号）。 +func TestOpenAIImagesJSONKeepalive_HeartbeatBeforeForwardStillFailsOver(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","response_format":"b64_json"}`) + + req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = req + + svc := &OpenAIGatewayService{ + httpUpstream: &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "X-Request-Id": []string{"req_img_heartbeat_failover"}, + }, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000021}}\n\n" + + "data: {\"type\":\"error\",\"error\":{\"type\":\"server_error\",\"code\":\"server_error\",\"message\":\"The image service is temporarily unavailable.\"}}\n\n", + )), + }, + }, + } + parsed, err := svc.ParseOpenAIImagesRequest(c, body) + require.NoError(t, err) + + // 模拟上一轮 failover 已发生：心跳已提交 200 并写出空白字节。 + stop := StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + defer stop() + waitForOpenAIImagesJSONKeepalive(t, c) + + account := &Account{ + ID: 22, + Name: "openai-oauth-heartbeat-failover", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "access_token": "token-123", + }, + } + + result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode) + require.Contains(t, string(failoverErr.ResponseBody), "temporarily unavailable") + require.Empty(t, strings.TrimSpace(rec.Body.String()), "only heartbeat whitespace may reach the client") + + rawEvents, ok := c.Get(OpsUpstreamErrorsKey) + require.True(t, ok) + events, ok := rawEvents.([]*OpsUpstreamErrorEvent) + require.True(t, ok) + require.Len(t, events, 1) + require.Equal(t, "failover", events[0].Kind) + require.Equal(t, account.ID, events[0].AccountID) + require.Equal(t, http.StatusBadGateway, events[0].UpstreamStatusCode) +} + +func waitForOpenAIImagesJSONKeepalive(t *testing.T, c *gin.Context) { + t.Helper() + k := openAIImagesJSONKeepaliveFromContext(c) + require.NotNil(t, k) + require.Eventually(t, func() bool { + k.mu.Lock() + defer k.mu.Unlock() + return k.started + }, time.Second, time.Millisecond) +} diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index 226cbf8e88..363ba9f591 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -1010,9 +1010,13 @@ func buildOpenAIImagesStreamErrorBodyFromUpstream(err *OpenAIImagesUpstreamError } func writeOpenAIImagesUpstreamErrorResponse(c *gin.Context, err *OpenAIImagesUpstreamError) bool { - if c == nil || c.Writer == nil || c.Writer.Written() || err == nil { + if c == nil || c.Writer == nil || err == nil { return false } + if c.Writer.Written() && OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) >= 0 { + return false + } + StopOpenAIImagesJSONKeepaliveCommitted(c) errorObj := gin.H{ "type": err.clientErrorType(), "message": err.clientMessage(), @@ -1176,7 +1180,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( var sseData openAISSEDataAccumulator var processDataErr error processDataDone := false - writerSizeBeforeResponse := c.Writer.Size() + writerSizeBeforeResponse := OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) processData := func(dataBytes []byte) { if processDataDone || processDataErr != nil { @@ -1591,7 +1595,9 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( imageOutputSizes []string firstTokenMs *int ) - writerSizeBeforeResponse := c.Writer.Size() + // 与 handleOpenAIImagesOAuthResponseError 的比较端同口径：排除非流式 JSON + // keepalive 心跳字节，避免 failover 第 2 轮起把上一轮心跳残留误判为已写响应。 + writerSizeBeforeResponse := OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) if parsed.Stream { usage, imageCount, imageOutputSizes, firstTokenMs, err = s.handleOpenAIImagesOAuthStreamingResponse(resp, c, startTime, parsed.ResponseFormat, openAIImagesStreamPrefix(parsed), requestModel) if err != nil { @@ -1672,7 +1678,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthResponseError( } retryable := IsOpenAIImagesRetryableUpstreamError(upstreamErr) - responseWritten := c != nil && c.Writer != nil && c.Writer.Size() != writerSizeBeforeResponse + responseWritten := c != nil && c.Writer != nil && OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeResponse kind := "http_error" if retryable { kind = "failover" diff --git a/deploy/.env.example b/deploy/.env.example index 1d1a6ec881..57056b4907 100644 --- a/deploy/.env.example +++ b/deploy/.env.example @@ -319,13 +319,15 @@ GATEWAY_SCHEDULING_OUTBOX_BACKLOG_REBUILD_ROWS=10000 GATEWAY_SCHEDULING_FULL_REBUILD_INTERVAL_SECONDS=300 # ----------------------------------------------------------------------------- -# Image Generation Stream & Concurrency (Optional) -# 图片生成流式与并发隔离配置（可选） +# Image Generation Keepalive & Concurrency (Optional) +# 图片生成保活与并发隔离配置（可选） # ----------------------------------------------------------------------------- # 图片流式上游数据间隔超时（秒）。0 表示禁用；非 0 时必须为 60-1800。 GATEWAY_IMAGE_STREAM_DATA_INTERVAL_TIMEOUT=900 # 图片流式 keepalive 间隔（秒）。0 表示禁用；非 0 时必须为 5-60。 GATEWAY_IMAGE_STREAM_KEEPALIVE_INTERVAL=10 +# 图片非流式 JSON keepalive 间隔（秒）。默认 0 禁用；首个心跳后 HTTP 状态会固化为 200。 +GATEWAY_IMAGE_NONSTREAM_KEEPALIVE_INTERVAL=0 # 是否启用进程级图片生成并发限制。默认 false，保持历史行为。 GATEWAY_IMAGE_CONCURRENCY_ENABLED=false # 当前进程允许同时处理的图片生成请求数。0 表示不限制。 diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index dfc584a59a..4cf759497f 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -390,6 +390,9 @@ gateway: # Image stream keepalive interval (seconds), 0=disable; independent from ordinary text streams # 图片流式 keepalive 间隔（秒），0=禁用；独立于普通文本流式 image_stream_keepalive_interval: 10 + # Non-streaming Images JSON keepalive interval (seconds), 0=disable; commits HTTP 200 after the first heartbeat + # 图片非流式 JSON keepalive 间隔（秒），0=禁用；首个心跳后 HTTP 状态会固化为 200 + image_nonstream_keepalive_interval: 0 # Image generation independent concurrency limiter (process-local, default disabled) # 图片生成独立并发限制（进程级，默认关闭；多实例总上限约为实例数×该值） image_concurrency: