mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-09-14 00:01:53 +00:00
fix: avoid stale stream writes after client disconnect (#5710)
* fix: avoid stale stream writes after client disconnect * fix: wait for stream ping goroutines before returning * fix: log stream results after goroutine cleanup * fix: broadcast stream stop signals * fix: abort upstream on client disconnect and restore write error contracts Keep the goroutine-lifecycle fix (unconditional wg.Wait before returning the gin.Context, close resp.Body inside cleanup), but drop the drain-on-disconnect behavior: when the client goes away, cleanup now runs immediately so the upstream body is closed, the provider stops generating, and users are not billed for tokens produced after they disconnected. Also restore FlushWriter/StringData/PingData returning an error when the request context is done, so non-scanner relay loops (ollama, fake-stream, audio, image) keep their disconnect awareness instead of silently consuming the upstream to completion. ResponseChunkData now propagates write errors. Add a bounded per-write deadline (http.NewResponseController) before each locked stream write so a slow-but-connected client cannot block a write forever and hang the unconditional wg.Wait. --------- Co-authored-by: CaIon <i@caion.me>
This commit is contained in:
@@ -206,5 +206,5 @@ func sendResponsesStreamData(c *gin.Context, streamResponse dto.ResponsesStreamR
|
||||
if data == "" {
|
||||
return
|
||||
}
|
||||
helper.ResponseChunkData(c, streamResponse, data)
|
||||
_ = helper.ResponseChunkData(c, streamResponse, data)
|
||||
}
|
||||
|
||||
@@ -136,10 +136,10 @@ func writeOpenaiImageStreamChunk(c *gin.Context, data []byte) {
|
||||
}
|
||||
_ = common.Unmarshal(data, &payload)
|
||||
if eventName := strings.TrimSpace(payload.Type); eventName != "" {
|
||||
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", eventName)})
|
||||
_ = helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventName}, string(data))
|
||||
return
|
||||
}
|
||||
c.Render(-1, common.CustomEvent{Data: "data: " + string(data)})
|
||||
_ = helper.FlushWriter(c)
|
||||
_ = helper.StringData(c, string(data))
|
||||
}
|
||||
|
||||
// isOpenAIImageStreamErrorEvent detects upstream error chunks by JSON content
|
||||
@@ -269,19 +269,11 @@ func writeOpenaiImageStreamPayload(c *gin.Context, eventName string, payload any
|
||||
return err
|
||||
}
|
||||
if eventName != "" {
|
||||
if _, err := fmt.Fprintf(c.Writer, "event: %s\n", eventName); err != nil {
|
||||
return err
|
||||
}
|
||||
return helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventName}, string(data))
|
||||
}
|
||||
if _, err := fmt.Fprintf(c.Writer, "data: %s\n\n", data); err != nil {
|
||||
return err
|
||||
}
|
||||
return helper.FlushWriter(c)
|
||||
return helper.StringData(c, string(data))
|
||||
}
|
||||
|
||||
func writeOpenaiImageStreamDone(c *gin.Context) error {
|
||||
if _, err := fmt.Fprint(c.Writer, "data: [DONE]\n\n"); err != nil {
|
||||
return err
|
||||
}
|
||||
return helper.FlushWriter(c)
|
||||
return helper.StringData(c, "[DONE]")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user