From 5d0e2b71a651d2ee96ba0642bee5bde6e641c00e Mon Sep 17 00:00:00 2001 From: jedi95 Date: Fri, 28 Aug 2026 04:14:42 -0600 Subject: [PATCH] fix(relay): emit terminal error event when upstream stream fails mid-flight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Upstream TCP resets and scanner errors mid-SSE were previously swallowed: StreamScannerHandler recorded the failure in StreamStatus, but the stream handlers emitted a normal completion envelope (final usage chunk + [DONE] / response.completed) unconditionally. Strict OpenAI-compatible clients then saw content chunks followed by a success-shaped envelope with no semantic terminal event (finish_reason / response.completed) — indistinguishable from a truncated successful stream and unsafe to retry. Add EmitRelayFailureTerminal(): when StreamStatus reports an abnormal end (scanner_error, timeout, client_gone, panic, ping_fail), emit an in-band error chunk (empty choices, structured error object with the end reason) after the forwarded content, then terminate with [DONE] without the synthetic success envelope. Wired into: - OaiStreamHandler (chat completions): guard runs after the last forwarded chunk; on failure, skip the synthetic usage/final-response path entirely. - ResponsesStreamHandler (/v1/responses): same guard before return. Normal streams are byte-identical to before (guard is a no-op on done/eof/handler_stop). Fault injection covered by unit tests in relay/helper/stream_terminal_guard_test.go. Refs: QuantumNous/new-api#7059, #6547, #6594, #6649 --- relay/channel/openai/relay-openai.go | 11 ++ relay/channel/openai/relay_responses.go | 5 + relay/helper/relay_terminal_guard.go | 73 +++++++++++++ relay/helper/stream_terminal_guard_test.go | 121 +++++++++++++++++++++ 4 files changed, 210 insertions(+) create mode 100644 relay/helper/relay_terminal_guard.go create mode 100644 relay/helper/stream_terminal_guard_test.go diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 9a0619eb27f5..8c25502882f4 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -178,6 +178,17 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re } } + // Terminal-state guard (#7059): if the upstream relay ended abnormally + // (connection reset, scanner error, timeout), emit an in-band relay + // failure chunk after the forwarded content so strict clients can + // distinguish a truncated stream from a successful one, then terminate + // with [DONE] without the normal completion envelope (a synthetic + // final usage chunk would make the failure look like a success). + if helper.EmitRelayFailureTerminal(c, info) { + helper.Done(c) + return usage, nil + } + if !containStreamUsage { usage = service.ResponseText2Usage(c, responseTextBuilder.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) usage.CompletionTokens += toolCount * 7 diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index ceca1af3b381..b7338e8e2f28 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -174,5 +174,10 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + // Terminal-state guard (#7059): forward a terminal failure event when + // the upstream relay ended abnormally instead of letting the stream + // close with only a success-looking envelope (or bare EOF). + helper.EmitRelayFailureTerminal(c, info) + return usage, nil } diff --git a/relay/helper/relay_terminal_guard.go b/relay/helper/relay_terminal_guard.go new file mode 100644 index 000000000000..bc6069ec3b97 --- /dev/null +++ b/relay/helper/relay_terminal_guard.go @@ -0,0 +1,73 @@ +package helper + +import ( + "fmt" + + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/gin-gonic/gin" +) + +// EmitRelayFailureTerminal emits an in-band terminal error event on the +// downstream SSE stream when the upstream relay ended abnormally +// (scanner_error, timeout, client_gone, panic, ...), before the normal +// [DONE] terminator. +// +// Background: StreamScannerHandler records the accurate relay end-state in +// info.StreamStatus, but stream handlers historically emitted a normal +// completion envelope unconditionally (final usage chunk + data: [DONE] / +// response.completed) even when the upstream connection had died mid-stream +// (see QuantumNous/new-api#7059). Strict downstream clients (e.g. OpenAI SDK +// consumers that validate finish_reason / terminal events) then observe a +// stream whose envelope claims success while the semantic terminal event +// (finish_reason, response.completed) is missing — indistinguishable from a +// truncated successful stream. +// +// The error chunk uses the OpenAI chat-completions stream chunk shape with +// empty choices and error fields in the model_extra position +// (error_type/error_message), the same convention some providers use for +// in-stream validation errors and which strict OpenAI-compatible clients +// already special-case (choices empty ⇒ not a content chunk; error fields +// ⇒ retryable relay failure). The stream is still terminated with [DONE] +// so every client — strict or lenient — sees a syntactically complete SSE +// stream; the semantic distinction lives in the error chunk. +// +// It returns true if a failure event was emitted (caller should skip any +// further normal-completion chunks it would have appended). +func EmitRelayFailureTerminal(c *gin.Context, info *relaycommon.RelayInfo) bool { + if info == nil || info.StreamStatus == nil || info.StreamStatus.IsNormalEnd() { + return false + } + + errText := "" + if info.StreamStatus.EndError != nil { + errText = info.StreamStatus.EndError.Error() + } + // Truncate the upstream error string: it can embed IPs/hosts; keep enough + // to diagnose, not enough to leak infra details wholesale. + if len(errText) > 256 { + errText = errText[:256] + } + + reason := string(info.StreamStatus.EndReason) + if reason == "" { + reason = "unknown" + } + + // Chat-completions-shaped error chunk. Empty choices signals "not a + // content chunk"; the structured error object matches the in-stream + // provider-error convention that OpenAI-compatible SDKs already parse. + chunk := map[string]any{ + "id": fmt.Sprintf("relay-error-%s", reason), + "object": "chat.completion.chunk", + "choices": []any{}, + "error": map[string]any{ + "message": fmt.Sprintf("stream relay failure: %s: %s", reason, errText), + "type": "relay_stream_error", + "code": "upstream_stream_failure_" + reason, + "param": nil, + }, + } + _ = ObjectData(c, chunk) + _ = FlushWriter(c) + return true +} diff --git a/relay/helper/stream_terminal_guard_test.go b/relay/helper/stream_terminal_guard_test.go new file mode 100644 index 000000000000..8b303658a762 --- /dev/null +++ b/relay/helper/stream_terminal_guard_test.go @@ -0,0 +1,121 @@ +package helper + +import ( + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/relay/common" + "github.com/gin-gonic/gin" +) + +// buildStreamCtx builds a gin test context with an SSE-capable response +// recorder plus a RelayInfo carrying a fresh StreamStatus. +func buildStreamCtx(t *testing.T) (*gin.Context, *common.RelayInfo, *httptest.ResponseRecorder) { + t.Helper() + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil) + info := relayInfoForTest() + return c, info, w +} + +// upstreamResetBody returns a ReadCloser that delivers some SSE content then +// simulates a mid-stream connection reset (read error, not clean EOF). +type resetReader struct { + data []byte + pos int +} + +func (r *resetReader) Read(p []byte) (int, error) { + if r.pos < len(r.data) { + n := copy(p, r.data[r.pos:]) + r.pos += n + return n, nil + } + return 0, errors.New("read tcp: connection reset by peer") +} + +func (r *resetReader) Close() error { return nil } + +func TestStreamScannerHandlerUpstreamResetMarksScannerError(t *testing.T) { + c, info, _ := buildStreamCtx(t) + body := "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n" + + "data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n" + resp := &http.Response{Body: io.NopCloser(&resetReader{data: []byte(body)})} + + StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {}) + + if info.StreamStatus.EndReason != common.StreamEndReasonScannerErr { + t.Fatalf("expected end_reason=scanner_error, got %q", info.StreamStatus.EndReason) + } + if !strings.Contains(info.StreamStatus.Summary(), "reason=scanner_error") { + t.Fatalf("unexpected summary: %s", info.StreamStatus.Summary()) + } + t.Logf("StreamStatus after upstream reset: %s", info.StreamStatus.Summary()) +} + +func TestEmitRelayFailureTerminalNormalEndIsNoop(t *testing.T) { + c, info, _ := buildStreamCtx(t) + info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil) + + if EmitRelayFailureTerminal(c, info) { + t.Fatal("guard must not fire for a normal end") + } +} + +func TestEmitRelayFailureTerminalEmitsErrorChunk(t *testing.T) { + c, info, w := buildStreamCtx(t) + info.StreamStatus.SetEndReason(common.StreamEndReasonScannerErr, errors.New("read tcp: connection reset by peer")) + + if !EmitRelayFailureTerminal(c, info) { + t.Fatal("guard must fire for scanner_error") + } + out := w.Body.String() + if !strings.Contains(out, "relay_stream_error") { + t.Fatalf("error chunk missing relay_stream_error type, got: %s", out) + } + if !strings.Contains(out, "scanner_error") { + t.Fatalf("error chunk missing end reason, got: %s", out) + } +} + +func TestFullFaultPathResetThenGuard(t *testing.T) { + c, info, w := buildStreamCtx(t) + body := "data: {\"choices\":[{\"delta\":{\"content\":\"partial answer\"}}]}\n\n" + resp := &http.Response{Body: io.NopCloser(&resetReader{data: []byte(body)})} + + StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) { + // mimic OaiStreamHandler's dataHandler: forward each data chunk + _ = StringData(c, "data: "+data+"\n\n") + _ = FlushWriter(c) + }) + if fired := EmitRelayFailureTerminal(c, info); !fired { + t.Fatal("guard did not fire after upstream reset") + } + out := w.Body.String() + for _, want := range []string{ + "partial answer", // content was forwarded + "relay_stream_error", // failure is marked + "scanner_error", // reason is actionable + } { + if !strings.Contains(out, want) { + t.Fatalf("stream output missing %q, got: %s", want, out) + } + } + if strings.Contains(out, "\"finish_reason\":\"stop\"") { + t.Fatal("failure stream must not carry a success finish_reason") + } + t.Logf("combined stream output: %s", out) +} + +// relayInfoForTest builds the minimal RelayInfo shape StreamScannerHandler needs. +func relayInfoForTest() *common.RelayInfo { + info := &common.RelayInfo{} + info.StreamStatus = common.NewStreamStatus() + return info +}