From 1e536ee7d9e12afdc74cab8e1d677d66e35c4ad1 Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Sun, 7 Jul 2024 01:01:55 +0800
Subject: [PATCH] fix: streaming timeout

---
 relay/channel/openai/relay-openai.go | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go
index 5c2acf4..a26d93b 100644
--- a/relay/channel/openai/relay-openai.go
+++ b/relay/channel/openai/relay-openai.go
@@ -133,13 +133,19 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 	}()
 	service.SetEventStreamHeaders(c)
 	isFirst := true
+	ticker := time.NewTicker(time.Duration(constant.StreamingTimeout))
+	defer ticker.Stop()
 	c.Stream(func(w io.Writer) bool {
 		select {
+		case <-ticker.C:
+			common.LogError(c, "reading data from upstream timeout")
+			return false
 		case data := <-dataChan:
 			if isFirst {
 				isFirst = false
 				info.FirstResponseTime = time.Now()
 			}
+			ticker.Reset(time.Duration(constant.StreamingTimeout))
 			if strings.HasPrefix(data, "data: [DONE]") {
 				data = data[:12]
 			}