diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index 646c886047..403752baf5 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -331,6 +331,11 @@ func (s *OpenAIGatewayService) forwardAsChatCompletions( // 6. Build upstream request upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) + cancelUpstream := func() {} + if clientStream { + upstreamCtx, cancelUpstream = context.WithCancel(upstreamCtx) + } + defer cancelUpstream() upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, responsesBody, token, true, promptCacheKey, false) releaseUpstreamCtx() if err != nil { @@ -355,7 +360,10 @@ func (s *OpenAIGatewayService) forwardAsChatCompletions( if err != nil { return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) } - defer func() { _ = resp.Body.Close() }() + defer func() { + cancelUpstream() + _ = resp.Body.Close() + }() // 8. Handle error response with failover if resp.StatusCode >= 400 { diff --git a/backend/internal/service/openai_gateway_chat_completions_cancellation_test.go b/backend/internal/service/openai_gateway_chat_completions_cancellation_test.go new file mode 100644 index 0000000000..4115dc65b4 --- /dev/null +++ b/backend/internal/service/openai_gateway_chat_completions_cancellation_test.go @@ -0,0 +1,116 @@ +//go:build unit + +package service + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type contextBoundBlockingReadCloser struct { + data []byte + offset int + ctx context.Context + forceClose chan struct{} + closeOnce sync.Once +} + +func newContextBoundBlockingReadCloser(data []byte) *contextBoundBlockingReadCloser { + return &contextBoundBlockingReadCloser{ + data: data, + forceClose: make(chan struct{}), + } +} + +func (r *contextBoundBlockingReadCloser) Read(p []byte) (int, error) { + if r.offset < len(r.data) { + n := copy(p, r.data[r.offset:]) + r.offset += n + return n, nil + } + select { + case <-r.ctx.Done(): + return 0, r.ctx.Err() + case <-r.forceClose: + return 0, io.EOF + } +} + +func (r *contextBoundBlockingReadCloser) Close() error { + select { + case <-r.ctx.Done(): + case <-r.forceClose: + } + return nil +} + +func (r *contextBoundBlockingReadCloser) forceUnblock() { + r.closeOnce.Do(func() { close(r.forceClose) }) +} + +type contextBoundHTTPUpstream struct { + body *contextBoundBlockingReadCloser +} + +func (u *contextBoundHTTPUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + u.body.ctx = req.Context() + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: u.body, + }, nil +} + +func (u *contextBoundHTTPUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) { + return u.Do(req, proxyURL, accountID, accountConcurrency) +} + +func TestForwardAsChatCompletions_CancelsUpstreamBeforeClosingBody(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":true}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(string(body))) + c.Request.Header.Set("Content-Type", "application/json") + + upstreamBody := []byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"id\":\"msg_1\",\"role\":\"assistant\",\"status\":\"completed\",\"content\":[{\"type\":\"output_text\",\"text\":\"ok\"}]}],\"usage\":{\"input_tokens\":17,\"output_tokens\":8,\"total_tokens\":25}}}\n\n") + stream := newContextBoundBlockingReadCloser(upstreamBody) + t.Cleanup(stream.forceUnblock) + + cfg := rawChatCompletionsTestConfig() + cfg.Gateway.StreamKeepaliveInterval = 1 + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &contextBoundHTTPUpstream{body: stream}, + } + + type forwardResult struct { + result *OpenAIForwardResult + err error + } + resultCh := make(chan forwardResult, 1) + go func() { + result, err := svc.ForwardAsChatCompletions(context.Background(), c, rawChatCompletionsTestAccount(), body, "", "gpt-5.1") + resultCh <- forwardResult{result: result, err: err} + }() + + select { + case got := <-resultCh: + require.NoError(t, got.err) + require.NotNil(t, got.result) + require.Equal(t, 17, got.result.Usage.InputTokens) + case <-time.After(time.Second): + t.Fatal("ForwardAsChatCompletions did not cancel upstream before closing the body") + } +}