Skip to content

Commit b58fa3d

Browse files
committed
fix(helper): improve error handling in FlushWriter and related functions
1 parent 1c167c1 commit b58fa3d

File tree

2 files changed

+42
-15
lines changed

2 files changed

+42
-15
lines changed

relay/channel/gemini/relay-gemini-native.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,10 +94,10 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, info *relaycommon.RelayIn
9494
helper.SetEventStreamHeaders(c)
9595

9696
return geminiStreamHandler(c, info, resp, func(data string, geminiResponse *dto.GeminiChatResponse) bool {
97-
// 直接发送 GeminiChatResponse 响应
9897
err := helper.StringData(c, data)
9998
if err != nil {
100-
logger.LogError(c, err.Error())
99+
logger.LogError(c, "failed to write stream data: "+err.Error())
100+
return false
101101
}
102102
info.SendResponseCount++
103103
return true

relay/helper/common.go

Lines changed: 40 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,28 @@ import (
1414
"github.com/gorilla/websocket"
1515
)
1616

17-
func FlushWriter(c *gin.Context) error {
18-
if c.Writer == nil {
17+
func FlushWriter(c *gin.Context) (err error) {
18+
defer func() {
19+
if r := recover(); r != nil {
20+
err = fmt.Errorf("flush panic recovered: %v", r)
21+
}
22+
}()
23+
24+
if c == nil || c.Writer == nil {
1925
return nil
2026
}
21-
if flusher, ok := c.Writer.(http.Flusher); ok {
22-
flusher.Flush()
23-
return nil
27+
28+
if c.Request != nil && c.Request.Context().Err() != nil {
29+
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
2430
}
25-
return errors.New("streaming error: flusher not found")
31+
32+
flusher, ok := c.Writer.(http.Flusher)
33+
if !ok {
34+
return errors.New("streaming error: flusher not found")
35+
}
36+
37+
flusher.Flush()
38+
return nil
2639
}
2740

2841
func SetEventStreamHeaders(c *gin.Context) {
@@ -66,17 +79,31 @@ func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data st
6679
}
6780

6881
func StringData(c *gin.Context, str string) error {
69-
//str = strings.TrimPrefix(str, "data: ")
70-
//str = strings.TrimSuffix(str, "\r")
82+
if c == nil || c.Writer == nil {
83+
return errors.New("context or writer is nil")
84+
}
85+
86+
if c.Request != nil && c.Request.Context().Err() != nil {
87+
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
88+
}
89+
7190
c.Render(-1, common.CustomEvent{Data: "data: " + str})
72-
_ = FlushWriter(c)
73-
return nil
91+
return FlushWriter(c)
7492
}
7593

7694
func PingData(c *gin.Context) error {
77-
c.Writer.Write([]byte(": PING\n\n"))
78-
_ = FlushWriter(c)
79-
return nil
95+
if c == nil || c.Writer == nil {
96+
return errors.New("context or writer is nil")
97+
}
98+
99+
if c.Request != nil && c.Request.Context().Err() != nil {
100+
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
101+
}
102+
103+
if _, err := c.Writer.Write([]byte(": PING\n\n")); err != nil {
104+
return fmt.Errorf("write ping data failed: %w", err)
105+
}
106+
return FlushWriter(c)
80107
}
81108

82109
func ObjectData(c *gin.Context, object interface{}) error {

0 commit comments

Comments
 (0)