From 2c13ed32e748d845788acf66830aca61d7b3309c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=B9=8F?= <2829624376@qq.com> Date: Wed, 15 Jul 2026 02:54:22 +0800 Subject: [PATCH] fix(images): preserve JSON completion boundaries --- .../handler/openai_gateway_handler.go | 16 +++- .../handler/openai_gateway_handler_test.go | 92 +++++++++++++++++++ .../service/openai_images_json_keepalive.go | 6 ++ 3 files changed, 111 insertions(+), 3 deletions(-) diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 231363a8a3..4260f590fc 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -2051,13 +2051,23 @@ func (h *OpenAIGatewayHandler) ensureForwardErrorResponse(c *gin.Context, stream return false } // 先停 compact 心跳再读 Writer 状态，避免与心跳 goroutine 竞争。 - if service.StopOpenAICompactSSEKeepaliveCommitted(c) { + compactKeepaliveCommitted := service.StopOpenAICompactSSEKeepaliveCommitted(c) + if compactKeepaliveCommitted { streamStarted = true } - if service.IsResponseCommitted(c) { + imageKeepalivePresent := service.OpenAIImagesJSONKeepalivePresent(c) + service.StopOpenAIImagesJSONKeepaliveCommitted(c) + imageKeepalivePaddingOnly := false + imageKeepaliveResponseWritten := false + if imageKeepalivePresent { + adjustedSize := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) + imageKeepalivePaddingOnly = adjustedSize < 0 + imageKeepaliveResponseWritten = adjustedSize >= 0 + } + if service.IsResponseCommitted(c) || (!compactKeepaliveCommitted && imageKeepaliveResponseWritten) { return false } - if c.Writer.Written() { + if c.Writer.Written() && !imageKeepalivePaddingOnly { streamStarted = true } h.handleStreamingAwareError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed", streamStarted) diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index e4b594c0b8..808f652eed 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "errors" + "io" "net/http" "net/http/httptest" "strings" @@ -220,6 +221,97 @@ func TestOpenAIEnsureForwardErrorResponse_ResponsesRouteAfterWrittenEmitsRespons assert.Contains(t, body, "Upstream request failed") } +func TestOpenAIEnsureForwardErrorResponse_ImageJSONKeepaliveWritesSingleJSONFallback(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + stop := service.StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + defer stop() + before := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) + require.Eventually(t, c.Writer.Written, time.Second, time.Millisecond) + require.Equal(t, before, service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)) + require.False(t, openAIForwardErrorAlreadyCommunicated(c, before, errors.New("read upstream response: unexpected EOF"))) + + h := &OpenAIGatewayHandler{} + wrote := h.ensureForwardErrorResponse(c, false) + + require.True(t, wrote) + require.Equal(t, http.StatusOK, w.Code, "heartbeat already committed the status") + require.True(t, json.Valid(w.Body.Bytes()), w.Body.String()) + require.NotContains(t, w.Body.String(), "event:") + require.NotContains(t, w.Body.String(), "data:") + + decoder := json.NewDecoder(strings.NewReader(w.Body.String())) + var payload map[string]any + require.NoError(t, decoder.Decode(&payload)) + require.ErrorIs(t, decoder.Decode(&payload), io.EOF) + require.Equal(t, "upstream_error", gjson.Get(w.Body.String(), "error.type").String()) + require.Equal(t, "Upstream request failed", gjson.Get(w.Body.String(), "error.message").String()) +} + +func TestOpenAIEnsureForwardErrorResponse_ImageJSONKeepalivePreservesCompletedJSON(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + stop := service.StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond) + defer stop() + before := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) + require.Eventually(t, c.Writer.Written, time.Second, time.Millisecond) + c.JSON(http.StatusOK, gin.H{"data": []gin.H{{"b64_json": "aW1hZ2U="}}}) + completedBody := w.Body.String() + require.True(t, json.Valid([]byte(completedBody)), completedBody) + require.Greater(t, service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c), before) + require.False(t, openAIForwardErrorAlreadyCommunicated(c, before, errors.New("read upstream trailer: unexpected EOF"))) + + h := &OpenAIGatewayHandler{} + wrote := h.ensureForwardErrorResponse(c, false) + + require.False(t, wrote, "the completed Images JSON already communicated the response") + require.Equal(t, completedBody, w.Body.String()) + require.NotContains(t, w.Body.String(), "event:") + require.NotContains(t, w.Body.String(), "data:") + + decoder := json.NewDecoder(strings.NewReader(w.Body.String())) + var payload map[string]any + require.NoError(t, decoder.Decode(&payload)) + require.ErrorIs(t, decoder.Decode(&payload), io.EOF) + require.Equal(t, "aW1hZ2U=", gjson.Get(w.Body.String(), "data.0.b64_json").String()) +} + +func TestOpenAIEnsureForwardErrorResponse_FastImageJSONKeepalivePreservesCompletedJSON(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + stop := service.StartOpenAIImagesJSONKeepalive(c, time.Hour) + defer stop() + before := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) + c.JSON(http.StatusOK, gin.H{"data": []gin.H{{"b64_json": "ZmFzdC1pbWFnZQ=="}}}) + completedBody := w.Body.String() + require.True(t, json.Valid([]byte(completedBody)), completedBody) + require.Greater(t, service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c), before) + require.False(t, openAIForwardErrorAlreadyCommunicated(c, before, errors.New("read upstream trailer: unexpected EOF"))) + + h := &OpenAIGatewayHandler{} + wrote := h.ensureForwardErrorResponse(c, false) + + require.False(t, wrote, "fast completed Images JSON already communicated the response") + require.Equal(t, completedBody, w.Body.String()) + require.NotContains(t, w.Body.String(), "event:") + require.NotContains(t, w.Body.String(), "data:") + + decoder := json.NewDecoder(strings.NewReader(w.Body.String())) + var payload map[string]any + require.NoError(t, decoder.Decode(&payload)) + require.ErrorIs(t, decoder.Decode(&payload), io.EOF) + require.Equal(t, "ZmFzdC1pbWFnZQ==", gjson.Get(w.Body.String(), "data.0.b64_json").String()) +} + func TestShouldLogOpenAIForwardFailureAsWarn(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_images_json_keepalive.go b/backend/internal/service/openai_images_json_keepalive.go index 0d8b24fb0a..6ad79b1ae5 100644 --- a/backend/internal/service/openai_images_json_keepalive.go +++ b/backend/internal/service/openai_images_json_keepalive.go @@ -127,6 +127,12 @@ func StopOpenAIImagesJSONKeepaliveCommitted(c *gin.Context) bool { return committed } +// OpenAIImagesJSONKeepalivePresent reports whether the response writer belongs +// to an Images JSON request, including fast responses before the first beat. +func OpenAIImagesJSONKeepalivePresent(c *gin.Context) bool { + return openAIImagesJSONKeepaliveFromContext(c) != nil +} + // OpenAIImagesJSONKeepaliveAdjustedWrittenSize excludes heartbeat whitespace // from response-size checks so account retry and failover remain available. func OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c *gin.Context) int {