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

controller/relay-openai.go

@@ -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()
 }
 }