From aa29868defb8231c9613f9ad80cc13902b4272bc Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Sun, 26 Nov 2023 01:19:20 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E6=B5=81=E6=A8=A1=E5=BC=8F?=
 =?UTF-8?q?=E4=B8=8B=E5=AE=A2=E6=88=B7=E7=AB=AF=E4=B8=AD=E6=96=AD=E5=AF=BC?=
 =?UTF-8?q?=E8=87=B4=E7=9A=84=E8=AE=A1=E8=B4=B9=E5=BC=82=E5=B8=B8=E9=97=AE?=
 =?UTF-8?q?=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 controller/relay-openai.go | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/controller/relay-openai.go b/controller/relay-openai.go
index 33c5b8f4..19b09c96 100644
--- a/controller/relay-openai.go
+++ b/controller/relay-openai.go
@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"one-api/common"
 	"strings"
+	"sync"
 )
 
 func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*OpenAIErrorWithStatusCode, string) {
@@ -28,7 +29,9 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 	})
 	dataChan := make(chan string)
 	stopChan := make(chan bool)
+	var wg sync.WaitGroup
 	go func() {
+		wg.Add(1)
 		var streamItems []string
 		for scanner.Scan() {
 			data := scanner.Text()
@@ -71,6 +74,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 				}
 			}
 		}
+		wg.Done()
 		stopChan <- true
 	}()
 	setEventStreamHeaders(c)
@@ -92,6 +96,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 	if err != nil {
 		return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), ""
 	}
+	wg.Wait()
 	return nil, responseTextBuilder.String()
 }