opensource/sub2api
3
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-09-24 16:05:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #4298 from wp-a/fix/openai-images-json-completion-boundary

Browse Source 
fix(images): preserve JSON completion boundaries
This commit is contained in:
Wesley Liddick
2026-07-15 11:08:47 +08:00
committed by GitHub
parent ab369c363f 2c13ed32e7
commit 4344f6afb0
3 changed files with 110 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/handler/openai_gateway_handler.go
+13 -3
View File
@@ -2209,13 +2209,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)
backend/internal/handler/openai_gateway_handler_test.go
+91
View File
@@ -221,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)
backend/internal/service/openai_images_json_keepalive.go
+6
View File
@@ -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 {