mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-07 13:08:03 +08:00
fix(openai): 非流式生图脱钩上游 context，客户端断开不再导致图已出却不扣费
forwardOpenAIImagesAPIKey 走的是 detachStreamUpstreamContext(ctx, parsed.Stream)， 该函数在非流式时原样返回请求 context（gateway_usage_billing.go:488）。于是客户端 中途断开会连带取消已经在出图的上游调用：上游那边图已生成并计费，网关这边拿到 context canceled，记 502，result 为 nil，handler 的 result.ImageCount > 0 兜底 无法命中，本次请求不产生扣费。 生图是长耗时（数十秒）且上游侧已产生实际成本的操作，与普通非流式 chat 不同： 提前取消并不能省下上游开销，只会丢掉已付费的产出。 同一端点的 OAuth 分支 forwardOpenAIImagesOAuth（openai_images_responses.go:1690） 以及同属媒体生成的 grok_media.go 本来就用无条件脱钩的 detachUpstreamContext， OpenAI 网关侧 13 个转发点里只有这一处是例外。这里对齐。 Fixes #5411
This commit is contained in:
li
@@ -601,7 +601,12 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey(
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, parsed.Stream)
|
||||
// 生图是长耗时、上游侧已产生实际成本的操作：客户端中途断开不应连带取消上游请求。
|
||||
// detachStreamUpstreamContext 在非流式时原样返回请求 context，于是客户端一断开
|
||||
// 就把已经在出图的上游调用打断成 context canceled，网关记 502、不扣费，而上游那边
|
||||
// 图已经生成并计费。同一端点的 OAuth 分支 forwardOpenAIImagesOAuth 以及 Grok 媒体
|
||||
// 路径本来就无条件脱钩，这里对齐；上游侧仍由 ResponseHeaderTimeout 兜底。
|
||||
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
||||
defer releaseUpstreamCtx()
|
||||
|
||||
token, _, err := s.GetAccessToken(upstreamCtx, account)
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
)
|
||||
|
||||
func newOpenAIImagesTestContext(t *testing.T, body []byte) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
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
|
||||
return c, rec
|
||||
}
|
||||
|
||||
func newOpenAIImagesTestService(upstream HTTPUpstream) *OpenAIGatewayService {
|
||||
return &OpenAIGatewayService{
|
||||
httpUpstream: upstream,
|
||||
cfg: &config.Config{
|
||||
Security: config.SecurityConfig{
|
||||
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newOpenAIImagesAPIKeyAccount() *Account {
|
||||
return &Account{
|
||||
ID: 31,
|
||||
Name: "openai-apikey-images",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "sk-test",
|
||||
"base_url": "https://api.openai.com/v1",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func openAIImagesJSONResponse() *http.Response {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"application/json"},
|
||||
"X-Request-Id": []string{"req_img_ctx"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
`{"created":1710000000,"data":[{"b64_json":"aGVsbG8="}],"usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}`,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
// issue #5411：生图是长耗时、上游侧已经产生实际成本的操作。客户端中途断开时，
|
||||
// 如果连带取消上游请求，就会出现「上游已出图并计费、网关记 502 context canceled、
|
||||
// 用户不扣费」。非流式路径以前走 detachStreamUpstreamContext(ctx, false)，
|
||||
// 该函数在非流式时原样返回请求 context，因此不脱钩。
|
||||
func TestForwardOpenAIImagesAPIKey_NonStreamDetachesUpstreamContext(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","response_format":"b64_json"}`)
|
||||
c, _ := newOpenAIImagesTestContext(t, body)
|
||||
|
||||
recorder := &httpUpstreamRecorder{resp: openAIImagesJSONResponse()}
|
||||
svc := newOpenAIImagesTestService(recorder)
|
||||
|
||||
parsed, err := svc.ParseOpenAIImagesRequest(c, body)
|
||||
require.NoError(t, err)
|
||||
require.False(t, parsed.Stream, "本用例覆盖非流式生图")
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // 客户端已断开
|
||||
|
||||
result, err := svc.ForwardImages(ctx, c, newOpenAIImagesAPIKeyAccount(), body, parsed, "")
|
||||
|
||||
require.NoError(t, err, "客户端断开不应把已在出图的上游调用打断成 context canceled")
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 1, result.ImageCount, "图片已产出，必须带回结果供计费")
|
||||
|
||||
require.NotNil(t, recorder.lastReq)
|
||||
require.NoError(t, recorder.lastReq.Context().Err(),
|
||||
"交给上游的请求 context 必须已脱钩，不随客户端断开取消")
|
||||
}
|
||||
|
||||
// 流式路径本来就脱钩，这条守卫防止对齐时把它改坏。
|
||||
func TestForwardOpenAIImagesAPIKey_StreamKeepsDetachedUpstreamContext(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true,"response_format":"b64_json"}`)
|
||||
c, _ := newOpenAIImagesTestContext(t, body)
|
||||
|
||||
recorder := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"text/event-stream"},
|
||||
"X-Request-Id": []string{"req_img_ctx_stream"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000000}}\n\n" +
|
||||
"data: {\"type\":\"response.image_generation_call.completed\",\"result\":\"aGVsbG8=\"}\n\n" +
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":10,\"output_tokens\":20}}}\n\n",
|
||||
)),
|
||||
}}
|
||||
svc := newOpenAIImagesTestService(recorder)
|
||||
|
||||
parsed, err := svc.ParseOpenAIImagesRequest(c, body)
|
||||
require.NoError(t, err)
|
||||
require.True(t, parsed.Stream)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
_, _ = svc.ForwardImages(ctx, c, newOpenAIImagesAPIKeyAccount(), body, parsed, "")
|
||||
|
||||
require.NotNil(t, recorder.lastReq)
|
||||
require.NoError(t, recorder.lastReq.Context().Err(),
|
||||
"流式路径原本就脱钩，不能被改回随客户端取消")
|
||||
}
|
||||
|
||||
// 两个 detach 辅助函数的语义差异是本次修复的根据，锁死它们防止被悄悄改动。
|
||||
func TestDetachUpstreamContextSemantics(t *testing.T) {
|
||||
t.Run("detachUpstreamContext_always_detaches", func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
detached, release := detachUpstreamContext(ctx)
|
||||
defer release()
|
||||
require.NoError(t, detached.Err())
|
||||
})
|
||||
|
||||
t.Run("detachStreamUpstreamContext_keeps_cancel_when_not_streaming", func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
same, release := detachStreamUpstreamContext(ctx, false)
|
||||
defer release()
|
||||
require.ErrorIs(t, same.Err(), context.Canceled,
|
||||
"非流式时该函数原样返回请求 context —— 生图路径不能用它")
|
||||
})
|
||||
|
||||
t.Run("detachStreamUpstreamContext_detaches_when_streaming", func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
detached, release := detachStreamUpstreamContext(ctx, true)
|
||||
defer release()
|
||||
require.NoError(t, detached.Err())
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user