fix(images): align OAuth pre-response size snapshot with keepalive-adjusted metric

The snapshot taken before dispatching the OAuth response handler still used
raw c.Writer.Size(), while handleOpenAIImagesOAuthResponseError compares it
against the keepalive-adjusted size. From the second failover round on, the
previous round's heartbeat bytes made the raw snapshot positive while the
adjusted size stayed -1, so a retryable upstream error was misjudged as
'response already written': the failover was swallowed and the client got a
committed 200 with a whitespace-only body.

Snapshot now uses OpenAIImagesJSONKeepaliveAdjustedWrittenSize, matching the
comparison side. Regression test simulates a prior-round heartbeat before
ForwardImages and asserts the retryable error still converts to
UpstreamFailoverError (verified red without the fix, green with it).
This commit is contained in:
shaw
2026-07-14 11:24:44 +08:00
parent 002c0b9fda
commit ac7a141a24
2 changed files with 69 additions and 1 deletions
@@ -1,6 +1,8 @@
package service
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
@@ -174,6 +176,70 @@ func TestOpenAIImagesJSONKeepaliveWriter_NilGuards(t *testing.T) {
}
}
// 回归: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)
@@ -1595,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 {