Fix timeout 500

bi1101 · May 2, 2024 · a5bfcf4 · a5bfcf4
1 parent f6b2ca6
commit a5bfcf4
Showing 1 changed file with 21 additions and 7 deletions.
diff --git a/internal/chatgpt/request.go b/internal/chatgpt/request.go
@@ -545,7 +545,7 @@ func Handler(c *gin.Context, response *http.Response, secret *tokens.Secret, pro
 	var isWSS = false
 	var convId string
 	var respId string
-	var wssUrl string
+	//var wssUrl string
 	var connInfo *connInfo
 	var wsSeq int
 	var isWSInterrupt bool = false
@@ -560,14 +560,22 @@ func Handler(c *gin.Context, response *http.Response, secret *tokens.Secret, pro
 		}
 		var wssResponse chatgpt_types.ChatGPTWSSResponse
 		json.NewDecoder(response.Body).Decode(&wssResponse)
-		wssUrl = wssResponse.WssUrl
+		//wssUrl = wssResponse.WssUrl
 		respId = wssResponse.ResponseId
 		convId = wssResponse.ConversationId
 	}
+	const readTimeout = 2 * time.Second
+
 	for {
 		var line string
 		var err error
 		if isWSS {
+			if connInfo.conn != nil {
+				connInfo.conn.SetReadDeadline(time.Now().Add(readTimeout))
+			} else {
+				c.JSON(500, gin.H{"error": "WebSocket connection is closed or not initialized"})
+				return "", nil
+			}
 			var messageType int
 			var message []byte
 			if isWSInterrupt {
@@ -585,11 +593,17 @@ func Handler(c *gin.Context, response *http.Response, secret *tokens.Secret, pro
 		reader:
 			messageType, message, err = connInfo.conn.ReadMessage()
 			if err != nil {
-				connInfo.ticker.Stop()
-				connInfo.conn.Close()
-				connInfo.conn = nil
-				err := createWSConn(wssUrl, connInfo, proxy, 0)
-				if err != nil {
+				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+					println("Read timeout:", err.Error())
+					connInfo.ticker.Stop()
+					connInfo.conn.Close()
+					connInfo.conn = nil
+					c.JSON(500, gin.H{"error": "WebSocket read took too long"})
+				} else {
+					println(err.Error())
+					connInfo.ticker.Stop()
+					connInfo.conn.Close()
+					connInfo.conn = nil
 					c.JSON(500, gin.H{"error": err.Error()})
 					return "", nil
 				}