xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-11-05 22:03:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

修复流模式下客户端中断导致的计费异常问题

Browse Source
This commit is contained in:
CaIon
2023-11-26 01:19:20 +08:00
parent a15f2d546e
commit aa29868def
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
controller/relay-openai.go
View File

@@ -9,6 +9,7 @@ import (
"net/http" "net/http"
"one-api/common" "one-api/common"
"strings" "strings"
"sync"
) )
func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*OpenAIErrorWithStatusCode, string) { 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) dataChan := make(chan string)
stopChan := make(chan bool) stopChan := make(chan bool)
var wg sync.WaitGroup
go func() { go func() {
wg.Add(1)
var streamItems []string var streamItems []string
for scanner.Scan() { for scanner.Scan() {
data := scanner.Text() data := scanner.Text()
@@ -71,6 +74,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
} }
} }
} }
wg.Done()
stopChan <- true stopChan <- true
}() }()
setEventStreamHeaders(c) setEventStreamHeaders(c)
@@ -92,6 +96,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
if err != nil { if err != nil {
return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), "" return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), ""
} }
wg.Wait()
return nil, responseTextBuilder.String() return nil, responseTextBuilder.String()
} }