mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-07 13:18:18 +08:00
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:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user