From a10ff1255cbaf2ec61dc9dcbc647cbc5bf5d3535 Mon Sep 17 00:00:00 2001 From: Ray <19251440+swjturay@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:27:52 +0800 Subject: [PATCH] fix(openai-ws): classify canceled client streams --- .../handler/openai_client_disconnect_test.go | 25 +++ .../handler/openai_gateway_handler.go | 20 ++ .../openai_ws_forwarder_client_cancel_test.go | 184 ++++++++++++++++++ .../service/openai_ws_forwarder_v2.go | 115 ++++++++--- 4 files changed, 317 insertions(+), 27 deletions(-) create mode 100644 backend/internal/handler/openai_client_disconnect_test.go create mode 100644 backend/internal/service/openai_ws_forwarder_client_cancel_test.go diff --git a/backend/internal/handler/openai_client_disconnect_test.go b/backend/internal/handler/openai_client_disconnect_test.go new file mode 100644 index 0000000000..71ff3d14f7 --- /dev/null +++ b/backend/internal/handler/openai_client_disconnect_test.go @@ -0,0 +1,25 @@ +package handler + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestOpenAIEnsureForwardErrorResponse_SkipsCanceledClient(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + ctx, cancel := context.WithCancel(context.Background()) + c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil).WithContext(ctx) + cancel() + + h := &OpenAIGatewayHandler{} + require.False(t, h.ensureForwardErrorResponse(c, true)) + require.Equal(t, statusClientClosedRequest, c.Writer.Status()) + require.Empty(t, recorder.Body.String()) +} diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index e57cc759e6..de54001530 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -792,6 +792,22 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { }) } if err != nil { + if result != nil && result.ClientDisconnect { + reqLog.Info("openai.client_disconnected", + zap.Int64("account_id", account.ID), + zap.Error(err), + ) + submitResponsesUsage(result) + return + } + if failoverClientGone(c) { + reqLog.Info("openai.client_disconnected", + zap.Int64("account_id", account.ID), + zap.Error(err), + ) + submitResponsesUsage(result) + return + } if result != nil && result.ImageCount > 0 { reqLog.Warn("openai.forward_partial_error_with_image_result", zap.Int64("account_id", account.ID), @@ -3087,6 +3103,10 @@ func (h *OpenAIGatewayHandler) ensureForwardErrorResponse(c *gin.Context, stream if c == nil || c.Writer == nil { return false } + if c.Request != nil && errors.Is(c.Request.Context().Err(), context.Canceled) { + failoverClientGone(c) + return false + } // 先停 compact 心跳再读 Writer 状态，避免与心跳 goroutine 竞争。 compactKeepaliveCommitted := service.StopOpenAICompactSSEKeepaliveCommitted(c) if compactKeepaliveCommitted { diff --git a/backend/internal/service/openai_ws_forwarder_client_cancel_test.go b/backend/internal/service/openai_ws_forwarder_client_cancel_test.go new file mode 100644 index 0000000000..7045c78d16 --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_client_cancel_test.go @@ -0,0 +1,184 @@ +package service + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// TestForwardOpenAIWSV2_ClientCancellationDrainsWithoutSyntheticFailure +// reproduces the HTTP/SSE ingress bug: the client cancels after partial output, +// while the upstream WS still has a terminal event available for usage billing. +func TestForwardOpenAIWSV2_ClientCancellationDrainsWithoutSyntheticFailure(t *testing.T) { + gin.SetMode(gin.TestMode) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + writer := &cancelOnFirstWriteResponseWriter{cancel: cancel} + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(ctx) + c.Request.Header.Set("User-Agent", "unit-test-agent/1.0") + + captureConn := &openAIWSCancelSafeConn{openAIWSCaptureConn: &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.created","response":{"id":"resp_cancel_1","model":"gpt-5.5"}}`), + []byte(`{"type":"response.output_text.delta","delta":"partial"}`), + []byte(`{"type":"response.completed","response":{"id":"resp_cancel_1","model":"gpt-5.5","usage":{"input_tokens":3,"output_tokens":5}}}`), + }, + readDelays: []time.Duration{0, 0, 50 * time.Millisecond}, + }} + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true + cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(&openAIWSClientConnCancelDialer{conn: captureConn}) + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + account := &Account{ + ID: 9101, + Name: "openai-ws-client-cancel", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModeCtxPool, + }, + } + + result, err := svc.Forward(ctx, c, account, []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"input_text","text":"hello"}]}`)) + + require.NoError(t, err, "client cancellation must not surface as an upstream failure") + require.NotNil(t, result) + require.True(t, result.ClientDisconnect) + require.Equal(t, "resp_cancel_1", result.RequestID) + require.Equal(t, 3, result.Usage.InputTokens) + require.Equal(t, 5, result.Usage.OutputTokens) + require.NotContains(t, writer.body.String(), "response.failed") +} + +type cancelOnFirstWriteResponseWriter struct { + mu sync.Mutex + header http.Header + body bytes.Buffer + cancel context.CancelFunc + writes int + status int +} + +func (w *cancelOnFirstWriteResponseWriter) Header() http.Header { + if w.header == nil { + w.header = make(http.Header) + } + return w.header +} + +func (w *cancelOnFirstWriteResponseWriter) WriteHeader(status int) { + w.status = status +} + +func (w *cancelOnFirstWriteResponseWriter) Write(p []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + if w.writes == 0 && w.cancel != nil { + w.cancel() + } + w.writes++ + return w.body.Write(p) +} + +func (w *cancelOnFirstWriteResponseWriter) Flush() {} + +type openAIWSClientConnCancelDialer struct { + conn openAIWSClientConn +} + +func (d *openAIWSClientConnCancelDialer) Dial( + ctx context.Context, + wsURL string, + headers http.Header, + proxyURL string, +) (openAIWSClientConn, int, http.Header, error) { + return d.conn, 0, nil, nil +} + +// openAIWSCancelSafeConn leaves a delayed event queued when its read context is +// canceled, matching websocket frame semantics and allowing the drain read to +// consume that event with a detached context. +type openAIWSCancelSafeConn struct { + *openAIWSCaptureConn +} + +func (c *openAIWSCancelSafeConn) ReadMessage(ctx context.Context) ([]byte, error) { + if ctx == nil { + ctx = context.Background() + } + c.mu.Lock() + if c.closed { + c.mu.Unlock() + return nil, errOpenAIWSConnClosed + } + if len(c.events) == 0 { + c.mu.Unlock() + return nil, io.EOF + } + delay := time.Duration(0) + if len(c.readDelays) > 0 { + delay = c.readDelays[0] + } + event := c.events[0] + c.mu.Unlock() + if delay > 0 { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-timer.C: + } + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return nil, errOpenAIWSConnClosed + } + if len(c.events) == 0 { + return nil, io.EOF + } + c.events = c.events[1:] + if len(c.readDelays) > 0 { + c.readDelays = c.readDelays[1:] + } + return event, nil +} diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 81f8631e96..3507b5b26c 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -374,6 +374,59 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( firstEventType := "" lastEventType := "" upstreamTerminalEvent := "" + clientDisconnected := false + clientDisconnectDrainStartedAt := time.Time{} + readTimeout := s.openAIWSReadTimeout() + upstreamReadCtx := ctx + upstreamReadDetached := false + clientRequestCanceled := func() bool { + return ctx != nil && errors.Is(ctx.Err(), context.Canceled) + } + markClientDisconnected := func(cause string) { + if clientDisconnected { + return + } + clientDisconnected = true + clientDisconnectDrainStartedAt = time.Now() + if !upstreamReadDetached { + upstreamReadCtx = context.WithoutCancel(ctx) + upstreamReadDetached = true + } + logOpenAIWSModeInfo( + "client_disconnected account_id=%d conn_id=%s cause=%s events=%d token_events=%d", + account.ID, + connID, + cause, + eventCount, + tokenEventCount, + ) + } + markClientRequestCanceled := func() { + if clientRequestCanceled() { + markClientDisconnected("request_context_canceled") + } + } + resultWithUsage := func() *OpenAIForwardResult { + return &OpenAIForwardResult{ + RequestID: responseID, + ResponseID: responseID, + Usage: *usage, + Model: originalModel, + UpstreamModel: mappedModel, + UpstreamResponseModel: responseModelObserver.Model(), + UpstreamResponseModelConflict: responseModelObserver.Conflict(), + UpstreamResponseServiceTier: responseModelObserver.ServiceTier(), + ServiceTier: resolvedOpenAIUpstreamServiceTierFromObserver(responseModelObserver, extractOpenAIServiceTier(reqBody)), + ReasoningEffort: extractOpenAIReasoningEffort(reqBody, mappedModel, originalModel), + Stream: reqStream, + OpenAIWSMode: true, + UpstreamTerminalEvent: upstreamTerminalEvent, + ResponseHeaders: lease.HandshakeHeaders(), + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnected, + } + } var flusher http.Flusher if reqStream { @@ -392,7 +445,6 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( flusher = f } - clientDisconnected := false flushBatchSize := s.openAIWSEventFlushBatchSize() flushInterval := s.openAIWSEventFlushInterval() pendingFlushEvents := 0 @@ -425,7 +477,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( flushStreamWriter(forceFlush) return } - clientDisconnected = true + markClientDisconnected("downstream_write_error") logger.LegacyPrintf("service.openai_gateway", "[OpenAI WS Mode] client disconnected, continue draining upstream: account=%d", account.ID) } flushBufferedStreamEvents := func(reason string) { @@ -452,17 +504,32 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } } - readTimeout := s.openAIWSReadTimeout() + // Keep per-read timeouts unchanged for connected clients. Once a client + // disconnects, use the same timeout as a bounded total drain budget. var pendingJSONDocuments [][]byte +readLoop: for { + markClientRequestCanceled() var message []byte var readErr error + readUsedDetachedContext := upstreamReadDetached if len(pendingJSONDocuments) > 0 { message = pendingJSONDocuments[0] pendingJSONDocuments = pendingJSONDocuments[1:] } else { - message, readErr = lease.ReadMessageWithContextTimeout(ctx, readTimeout) + currentReadTimeout := readTimeout + if clientDisconnected && !clientDisconnectDrainStartedAt.IsZero() { + remaining := readTimeout - time.Since(clientDisconnectDrainStartedAt) + if remaining <= 0 { + lease.MarkBroken() + break readLoop + } + if remaining < currentReadTimeout { + currentReadTimeout = remaining + } + } + message, readErr = lease.ReadMessageWithContextTimeout(upstreamReadCtx, currentReadTimeout) if readErr == nil { if documents, repaired := splitOpenAIConcatenatedJSONDocuments(message); repaired { logOpenAIWSModeInfo( @@ -477,6 +544,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } } } + markClientRequestCanceled() if readErr == nil && !json.Valid(message) { eventType, _, _ := parseOpenAIWSEventEnvelope(message) if eventType == "" { @@ -515,12 +583,15 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen), truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen), ) + if clientDisconnected { + if !readUsedDetachedContext && errors.Is(readErr, context.Canceled) && clientRequestCanceled() { + continue + } + break + } if !wroteDownstream { return nil, wrapOpenAIWSFallback(classifyOpenAIWSReadFallbackReason(readErr), readErr) } - if clientDisconnected { - break - } setOpsUpstreamError(c, 0, sanitizeUpstreamErrorMessage(readErr.Error()), "") return nil, fmt.Errorf("openai ws read event: %w", readErr) } @@ -710,7 +781,13 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } } + if clientDisconnected && terminalEventCount == 0 { + return resultWithUsage(), fmt.Errorf("openai ws stream incomplete after client disconnect: %w", context.Canceled) + } if !reqStream { + if clientDisconnected { + return resultWithUsage(), nil + } if len(finalResponse) == 0 { logOpenAIWSModeInfo( "missing_final_response account_id=%d conn_id=%s events=%d token_events=%d terminal_events=%d wrote_downstream=%v", @@ -772,26 +849,10 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( clientDisconnected, ) - return &OpenAIForwardResult{ - RequestID: responseID, - Usage: *usage, - Model: originalModel, - UpstreamModel: mappedModel, - UpstreamResponseModel: responseModelObserver.Model(), - UpstreamResponseModelConflict: responseModelObserver.Conflict(), - UpstreamResponseServiceTier: responseModelObserver.ServiceTier(), - ImageCount: imageCounter.Count(), - ImageOutputSizes: imageCounter.Sizes(), - ServiceTier: resolvedOpenAIUpstreamServiceTierFromObserver(responseModelObserver, extractOpenAIServiceTier(reqBody)), - ReasoningEffort: extractOpenAIReasoningEffort(reqBody, mappedModel, originalModel), - RequestedReasoningEffort: CanonicalRequestedReasoningEffortFromReqBody(reqBody, originalModel, mappedModel), - Stream: reqStream, - OpenAIWSMode: true, - UpstreamTerminalEvent: upstreamTerminalEvent, - ResponseHeaders: lease.HandshakeHeaders(), - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, - }, nil + result := resultWithUsage() + result.ImageCount = imageCounter.Count() + result.ImageOutputSizes = imageCounter.Sizes() + return result, nil } // ProxyResponsesWebSocketFromClient 处理客户端入站 WebSocket（OpenAI Responses WS Mode）并转发到上游。