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:
Seefs
2026-07-06 21:40:23 +08:00
committed by GitHub
co-authored by CaIon
parent fc26b88fd1
commit 153d7f01a2
6 changed files with 187 additions and 115 deletions
+21 -5
View File
@@ -25,7 +25,7 @@ func FlushWriter(c *gin.Context) (err error) {
return nil
}
if c.Request != nil && c.Request.Context().Err() != nil {
if requestContextDone(c) {
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
}
@@ -38,6 +38,10 @@ func FlushWriter(c *gin.Context) (err error) {
return nil
}
func requestContextDone(c *gin.Context) bool {
return c != nil && c.Request != nil && c.Request.Context().Err() != nil
}
func SetEventStreamHeaders(c *gin.Context) {
// 检查是否已经设置过头部
if _, exists := c.Get("event_stream_headers_set"); exists {
@@ -55,6 +59,10 @@ func SetEventStreamHeaders(c *gin.Context) {
}
func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error {
if requestContextDone(c) {
return nil
}
jsonData, err := common.Marshal(resp)
if err != nil {
common.SysError("error marshalling stream response: " + err.Error())
@@ -67,15 +75,23 @@ func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error {
}
func ClaudeChunkData(c *gin.Context, resp dto.ClaudeResponse, data string) {
if requestContextDone(c) {
return
}
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s\n", data)})
_ = FlushWriter(c)
}
func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) {
func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) error {
if requestContextDone(c) {
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
}
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)})
_ = FlushWriter(c)
return FlushWriter(c)
}
func StringData(c *gin.Context, str string) error {
@@ -83,7 +99,7 @@ func StringData(c *gin.Context, str string) error {
return errors.New("context or writer is nil")
}
if c.Request != nil && c.Request.Context().Err() != nil {
if requestContextDone(c) {
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
}
@@ -96,7 +112,7 @@ func PingData(c *gin.Context) error {
return errors.New("context or writer is nil")
}
if c.Request != nil && c.Request.Context().Err() != nil {
if requestContextDone(c) {
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
}