From e74c690ac42f634797ca14a5847aa4874711b033 Mon Sep 17 00:00:00 2001 From: cyh Date: Thu, 17 Sep 2026 19:15:09 +0800 Subject: [PATCH] fix(openai): keep transport heartbeats outside semantic output --- .../service/openai_capacity_shed_test.go | 58 +++++- .../service/openai_gateway_passthrough.go | 9 +- .../service/openai_visible_ttft_test.go | 3 +- .../internal/service/openai_ws_http_bridge.go | 27 ++- .../service/openai_ws_http_bridge_test.go | 186 +++++++++++++++++- 5 files changed, 261 insertions(+), 22 deletions(-) diff --git a/backend/internal/service/openai_capacity_shed_test.go b/backend/internal/service/openai_capacity_shed_test.go index 537ec435c5..7d2f8c8dd2 100644 --- a/backend/internal/service/openai_capacity_shed_test.go +++ b/backend/internal/service/openai_capacity_shed_test.go @@ -125,6 +125,8 @@ func TestOpenAIStreamErrorFrameDoesNotStartClientOutput(t *testing.T) { {`{"type":"response.failed","response":{"error":{"code":"server_is_overloaded"}}}`, "response.failed", false}, {`{"type":"response.created","response":{"id":"resp_1"}}`, "response.created", false}, {`{"type":"response.in_progress","response":{"id":"resp_1"}}`, "response.in_progress", false}, + {`{"type":"keepalive"}`, "keepalive", false}, + {`{"type":"unknown_event"}`, "unknown_event", true}, {`{"type":"response.output_item.added","item":{"type":"reasoning","summary":[]}}`, "response.output_item.added", false}, {`{"type":"response.output_item.added","item":{"type":"reasoning","encrypted_content":"ciphertext"}}`, "response.output_item.added", true}, {`{"type":"response.reasoning_summary_part.added","part":{"type":"summary_text","text":""}}`, "response.reasoning_summary_part.added", false}, @@ -138,7 +140,7 @@ func TestOpenAIStreamErrorFrameDoesNotStartClientOutput(t *testing.T) { } } -func TestOpenAIStreamMetadataPreambleAndMessageOnlyOverloadFailOver(t *testing.T) { +func TestOpenAIStreamMetadataAndKeepaliveBeforeOverloadFailOver(t *testing.T) { gin.SetMode(gin.TestMode) largeMetadata := strings.Repeat("x", 16*1024) stream := strings.Join([]string{ @@ -151,6 +153,11 @@ func TestOpenAIStreamMetadataPreambleAndMessageOnlyOverloadFailOver(t *testing.T "event: response.reasoning_summary_part.added", `data: {"type":"response.reasoning_summary_part.added","part":{"type":"summary_text","text":""}}`, "", + "event: keepalive", + `data: {"type":"keepalive"}`, + "", + `data: {"type":"keepalive"}`, + "", "event: error", `data: {"type":"error","error":{"type":"service_unavailable_error","message":"Our servers are currently overloaded. Please try again later."}}`, "", @@ -203,6 +210,55 @@ func TestOpenAIStreamMetadataPreambleAndMessageOnlyOverloadFailOver(t *testing.T } } +func TestOpenAIStreamKeepaliveOnlyFailureDoesNotRecordFirstToken(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, passthrough := range []bool{false, true} { + name := "native" + if passthrough { + name = "passthrough" + } + t.Run(name, func(t *testing.T) { + svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{ + MaxLineSize: defaultMaxLineSize, + }}} + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/", nil) + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp_heartbeat"}}`, + "", + `data: {"type":"keepalive"}`, + "", + `data: {"type":"response.failed","response":{"id":"resp_heartbeat","status":"failed","error":{"code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}}}`, + "", + "", + }, "\n"))), + } + account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + var err error + if passthrough { + var result *openaiStreamingResultPassthrough + result, err = svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, account, time.Now(), "model", "model") + require.NotNil(t, result) + require.Nil(t, result.firstTokenMs) + } else { + var result *openaiStreamingResult + result, err = svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model") + require.NotNil(t, result) + require.Nil(t, result.firstTokenMs) + } + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.True(t, failoverErr.RetryableOnSameAccount) + require.True(t, failoverErr.RequestScopedTransient) + require.Empty(t, rec.Body.String()) + }) + } +} + // 回归用例（真实上游降载序列）：created → in_progress → error 帧 → response.failed。 // 期望仍然走 pre-output failover（同账号重试 + 请求级瞬时标记），且不向客户端写出任何字节。 func TestOpenAIStreamCapacityShedErrorFramePrecedingFailedStillFailsOver(t *testing.T) { diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index dd8f390079..25874b5a5a 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -1073,9 +1073,10 @@ func openAIStreamClientOutputStarted(c *gin.Context, localStarted bool) bool { return OpenAICompactKeepaliveAdjustedWrittenSize(c) >= 0 } -func openAIStreamEventIsPreamble(eventType string) bool { +// Lifecycle metadata and transport heartbeats are not model output. +func openAIStreamEventIsMetadata(eventType string) bool { switch strings.TrimSpace(eventType) { - case "response.created", "response.in_progress": + case "response.created", "response.in_progress", "keepalive": return true default: return false @@ -1180,7 +1181,7 @@ func openAIStreamDataStartsClientOutput(data, eventType string) bool { case "response.output_item.added", "response.content_part.added", "response.reasoning_summary_part.added": return openAIStreamAddedEventStartsClientOutput([]byte(trimmed), eventType) } - return !openAIStreamEventIsPreamble(eventType) + return !openAIStreamEventIsMetadata(eventType) } func openAIStreamItemHasVisibleOutput(item gjson.Result) bool { @@ -1258,7 +1259,7 @@ func openAIStreamDataStartsSemanticTTFT(data, eventType string) bool { payload := []byte(trimmed) return !openAIStreamFailedEventShouldFailover(payload, extractOpenAISSEErrorMessage(payload)) default: - return !openAIStreamEventIsPreamble(eventType) + return !openAIStreamEventIsMetadata(eventType) } } diff --git a/backend/internal/service/openai_visible_ttft_test.go b/backend/internal/service/openai_visible_ttft_test.go index 1dae291cfb..24a60e21e2 100644 --- a/backend/internal/service/openai_visible_ttft_test.go +++ b/backend/internal/service/openai_visible_ttft_test.go @@ -71,7 +71,7 @@ func TestOpenAIResponsesTTFTStartsAtCompletedImage(t *testing.T) { } } -func TestOpenAINativeMetadataDoesNotDisarmFirstOutputTimeout(t *testing.T) { +func TestOpenAINativeMetadataAndKeepaliveDoNotDisarmFirstOutputTimeout(t *testing.T) { gin.SetMode(gin.TestMode) svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{ MaxLineSize: defaultMaxLineSize, @@ -84,6 +84,7 @@ func TestOpenAINativeMetadataDoesNotDisarmFirstOutputTimeout(t *testing.T) { defer func() { _ = writer.Close() }() _, _ = io.WriteString(writer, "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_test\"}}\n\n") _, _ = io.WriteString(writer, "data: {\"type\":\"response.output_item.added\",\"item\":{\"id\":\"item_test\",\"type\":\"reasoning\",\"summary\":[]}}\n\n") + _, _ = io.WriteString(writer, "data: {\"type\":\"keepalive\"}\n\n") time.Sleep(1200 * time.Millisecond) }() diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index 780e58c86d..1c3b1e96e7 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -835,7 +835,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( s.handleGrokAccountUpstreamError(ctx, account, statusCode, resp.Header, upstreamMessage) } } - if !wroteDownstream && shouldFailover && (turn == 1 || statusCode == http.StatusTooManyRequests) { + // A disconnected client needs this attempt drained for usage, not replayed, + // even when only non-semantic heartbeats were delivered. + if !clientDisconnected && !wroteDownstream && shouldFailover && (turn == 1 || statusCode == http.StatusTooManyRequests) { if account.Platform == PlatformGrok { return nil, newOpenAIUpstreamFailoverError(statusCode, resp.Header, upstreamMessage, errMessage, false) } @@ -873,11 +875,12 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( } } if !clientDisconnected && !suppressClientMessage { + isKeepalive := eventType == "keepalive" stageBeforeSemanticOutput := turn == 1 && account.Platform == PlatformOpenAI && !wroteDownstream commitStagedMessages := !stageBeforeSemanticOutput || openAIStreamDataStartsClientOutput(string(clientMessage), eventType) || isOpenAIWSTerminalEvent(eventType) - if stageBeforeSemanticOutput && !commitStagedMessages { + if stageBeforeSemanticOutput && !commitStagedMessages && !isKeepalive { if pendingClientMessageBytes+int64(len(clientMessage)) > openAIFirstOutputStageMaxBytes { return nil, s.newOpenAIStreamFailoverError( c, @@ -892,9 +895,15 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( pendingClientMessages = append(pendingClientMessages, append([]byte(nil), clientMessage...)) pendingClientMessageBytes += int64(len(clientMessage)) } else { - messages := append(pendingClientMessages, clientMessage) - pendingClientMessages = nil - pendingClientMessageBytes = 0 + // Keep the client connection alive without committing this attempt + // or exposing its staged lifecycle metadata. + var messages [][]byte + if !isKeepalive { + messages = pendingClientMessages + pendingClientMessages = nil + pendingClientMessageBytes = 0 + } + messages = append(messages, clientMessage) for _, message := range messages { if err := writeClientMessage(message); err != nil { if isOpenAIWSClientDisconnectError(err) { @@ -915,7 +924,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( wroteDownstream, ) } - wroteDownstream = true + if !isKeepalive { + wroteDownstream = true + } } } } @@ -966,7 +977,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( } if err := scanner.Err(); err != nil { streamErr := fmt.Errorf("read upstream http bridge stream: %w", err) - if turn == 1 && !wroteDownstream { + if turn == 1 && !clientDisconnected && !wroteDownstream { return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, streamErr, true) } return resultWithUsage(), streamErr @@ -975,7 +986,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( if sawDone { terminalErr = errors.New("upstream http bridge stream sent [DONE] before terminal event") } - if turn == 1 && !wroteDownstream { + if turn == 1 && !clientDisconnected && !wroteDownstream { return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, terminalErr, true) } return resultWithUsage(), terminalErr diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index 5e23a51985..12f08c655a 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -9,6 +9,7 @@ import ( "net/http/httptest" "strings" "testing" + "testing/iotest" "time" "github.com/Wei-Shaw/sub2api/internal/config" @@ -1186,20 +1187,41 @@ func TestProxyOpenAIWSHTTPBridgeTurnBareErrorFollowedByCompletedUsesCompleted(t require.Equal(t, "response.completed", gjson.GetBytes(writes[1], "type").String()) } -func TestProxyOpenAIWSHTTPBridgeTurnStagesMetadataBeforeCapacityFailover(t *testing.T) { +func TestProxyOpenAIWSHTTPBridgeTurnStagesMetadataAndRelaysKeepaliveBeforeCapacityFailover(t *testing.T) { gin.SetMode(gin.TestMode) - body := strings.Join([]string{ + preamble := strings.Join([]string{ `data: {"type":"response.created","response":{"id":"resp_shed"}}`, "", `data: {"type":"response.in_progress","response":{"id":"resp_shed"}}`, "", - `data: {"type":"response.failed","response":{"id":"resp_shed","status":"failed","error":{"message":"Our servers are currently overloaded. Please try again later."}}}`, + `data: {"type":"keepalive"}`, + "", "", }, "\n") + reader, writer := io.Pipe() + defer func() { _ = reader.Close() }() + defer func() { _ = writer.Close() }() + keepaliveReceived := make(chan struct{}, 1) + writerDone := make(chan error, 1) + go func() { + _, writeErr := io.WriteString(writer, preamble) + if writeErr == nil { + // The upstream stays idle until the heartbeat reaches the client. + // Flushing it only with later output or EOF must not pass this test. + select { + case <-keepaliveReceived: + _, writeErr = io.WriteString(writer, "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_shed\",\"status\":\"failed\",\"error\":{\"message\":\"Our servers are currently overloaded. Please try again later.\"}}}\n\n") + case <-time.After(2 * time.Second): + writeErr = errors.New("keepalive was not forwarded while upstream was idle") + } + } + _ = writer.CloseWithError(writeErr) + writerDone <- writeErr + }() upstream := &httpUpstreamRecorder{resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{"X-Request-Id": []string{"rid-ws-bridge-capacity"}}, - Body: io.NopCloser(strings.NewReader(body)), + Body: reader, }} svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} account := &Account{ID: 12, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1} @@ -1214,16 +1236,21 @@ func TestProxyOpenAIWSHTTPBridgeTurnStagesMetadataBeforeCapacityFailover(t *test "gpt-5", "", "", "", "", 1, func(message []byte) error { writes = append(writes, append([]byte(nil), message...)) + if gjson.GetBytes(message, "type").String() == "keepalive" { + keepaliveReceived <- struct{}{} + } return nil }, ) + require.NoError(t, <-writerDone) require.Nil(t, result) var failoverErr *UpstreamFailoverError require.ErrorAs(t, err, &failoverErr) require.True(t, failoverErr.RetryableOnSameAccount) require.True(t, failoverErr.RequestScopedTransient) - require.Empty(t, writes) + require.Len(t, writes, 1) + require.JSONEq(t, `{"type":"keepalive"}`, string(writes[0])) } func TestProxyOpenAIWSHTTPBridgeTurnDoesNotReplayCapacityAfterSemanticOutput(t *testing.T) { @@ -1233,8 +1260,12 @@ func TestProxyOpenAIWSHTTPBridgeTurnDoesNotReplayCapacityAfterSemanticOutput(t * body := strings.Join([]string{ `data: {"type":"response.created","response":{"id":"resp_partial"}}`, "", + `data: {"type":"keepalive"}`, + "", `data: {"type":"response.output_text.delta","delta":"partial"}`, "", + `data: {"type":"keepalive"}`, + "", `data: {"type":"response.failed","response":{"id":"resp_partial","status":"failed","error":{"code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}}}`, "", }, "\n") @@ -1262,15 +1293,154 @@ func TestProxyOpenAIWSHTTPBridgeTurnDoesNotReplayCapacityAfterSemanticOutput(t * require.NotNil(t, result) require.NoError(t, err) - require.Len(t, writes, 3) + require.Len(t, writes, 5) + require.Equal(t, "keepalive", gjson.GetBytes(writes[0], "type").String()) + require.Equal(t, "response.created", gjson.GetBytes(writes[1], "type").String()) + require.Equal(t, "response.output_text.delta", gjson.GetBytes(writes[2], "type").String()) + require.Equal(t, "keepalive", gjson.GetBytes(writes[3], "type").String()) var failoverErr *UpstreamFailoverError require.False(t, errors.As(err, &failoverErr)) - require.Contains(t, string(writes[2]), `"code":"server_error"`) - require.NotContains(t, string(writes[2]), "server_is_overloaded") + require.Contains(t, string(writes[4]), `"code":"server_error"`) + require.NotContains(t, string(writes[4]), "server_is_overloaded") require.True(t, logSink.ContainsMessage("gateway.failover_suppressed_after_semantic_output")) require.True(t, logSink.ContainsFieldValue("path", "ws_http_bridge")) } +func TestProxyOpenAIWSHTTPBridgeTurnKeepaliveWriteFailure(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, tc := range []struct { + name string + writeErr error + wantDrain bool + }{ + {name: "write error", writeErr: errors.New("downstream write failed")}, + {name: "client disconnect", writeErr: io.EOF, wantDrain: true}, + } { + t.Run(tc.name, func(t *testing.T) { + body := strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp_disconnect"}}`, + "", + `data: {"type":"keepalive"}`, + "", + `data: {"type":"response.completed","response":{"id":"resp_disconnect","usage":{"input_tokens":3,"output_tokens":2}}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(body)), + }} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ID: 14, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1} + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`) + var writes [][]byte + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "sk-test", payload, len(payload), + "gpt-5", "", "", "", "", 1, + func(message []byte) error { + writes = append(writes, append([]byte(nil), message...)) + return tc.writeErr + }, + ) + + require.Len(t, writes, 1) + require.JSONEq(t, `{"type":"keepalive"}`, string(writes[0])) + if tc.wantDrain { + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 3, result.Usage.InputTokens) + require.Equal(t, 2, result.Usage.OutputTokens) + } else { + require.Nil(t, result) + require.ErrorIs(t, err, tc.writeErr) + var turnErr *openAIWSIngressTurnError + require.ErrorAs(t, err, &turnErr) + require.Equal(t, "write_client", turnErr.stage) + require.False(t, turnErr.wroteDownstream) + require.False(t, isOpenAIWSIngressTurnRetryable(err)) + } + }) + } +} + +func TestProxyOpenAIWSHTTPBridgeTurnDisconnectedBeforeOutputDoesNotFailOver(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, tc := range []struct { + name string + tail string + readErr error + wantTerminal bool + }{ + { + name: "overload followed by failed usage", + tail: "data: {\"type\":\"error\",\"error\":{\"code\":\"server_is_overloaded\",\"message\":\"overloaded\"}}\n\n" + + "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_disconnect\",\"status\":\"failed\",\"error\":{\"code\":\"server_is_overloaded\",\"message\":\"overloaded\"},\"usage\":{\"input_tokens\":3,\"output_tokens\":2}}}\n\n", + wantTerminal: true, + }, + {name: "upstream read failure", readErr: io.ErrUnexpectedEOF}, + {name: "upstream EOF without terminal"}, + {name: "done without terminal", tail: "data: [DONE]\n\n"}, + } { + t.Run(tc.name, func(t *testing.T) { + preamble := "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_disconnect\"}}\n\n" + + "data: {\"type\":\"keepalive\"}\n\n" + + "data: {\"type\":\"keepalive\"}\n\n" + var reader io.Reader = strings.NewReader(preamble + tc.tail) + if tc.readErr != nil { + reader = io.MultiReader(reader, iotest.ErrReader(tc.readErr)) + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{}, + Body: io.NopCloser(reader), + }} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ID: 15, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1} + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`) + var writes [][]byte + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "sk-test", payload, len(payload), + "gpt-5", "", "", "", "", 1, + func(message []byte) error { + writes = append(writes, append([]byte(nil), message...)) + if len(writes) == 1 { + return nil + } + return io.EOF + }, + ) + + require.NotNil(t, result) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr)) + require.Len(t, writes, 2) + for _, message := range writes { + require.JSONEq(t, `{"type":"keepalive"}`, string(message)) + } + require.Nil(t, result.FirstTokenMs) + if tc.wantTerminal { + require.NoError(t, err) + require.Equal(t, "response.failed", result.UpstreamTerminalEvent) + require.Equal(t, 3, result.Usage.InputTokens) + require.Equal(t, 2, result.Usage.OutputTokens) + } else { + require.Error(t, err) + if tc.readErr != nil { + require.ErrorIs(t, err, tc.readErr) + } + } + }) + } +} + func TestProxyOpenAIWSHTTPBridgeTurnRequiresTerminalEvent(t *testing.T) { gin.SetMode(gin.TestMode)