diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index 715371c29..f047d4d3b 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -1258,8 +1258,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http } if usage.CompletionTokens <= 0 { - str := responseText.String() - if len(str) > 0 { + if info.ReceivedResponseCount > 0 { usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) } else { usage = &dto.Usage{} diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 20ca53ed1..f5c1d769e 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -113,6 +113,7 @@ type RelayInfo struct { UserQuota int RelayFormat types.RelayFormat SendResponseCount int + ReceivedResponseCount int FinalPreConsumedQuota int // 最终预消耗的配额 // BillingSource indicates whether this request is billed from wallet quota or subscription. // "" or "wallet" => wallet; "subscription" => subscription diff --git a/relay/helper/stream_scanner.go b/relay/helper/stream_scanner.go index 13c32c675..4f3ab2363 100644 --- a/relay/helper/stream_scanner.go +++ b/relay/helper/stream_scanner.go @@ -90,10 +90,10 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon // 等待所有 goroutine 退出,最多等待5秒 done := make(chan struct{}) - go func() { + gopool.Go(func() { wg.Wait() close(done) - }() + }) select { case <-done: @@ -138,11 +138,11 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon case <-pingTicker.C: // 使用超时机制防止写操作阻塞 done := make(chan error, 1) - go func() { + gopool.Go(func() { writeMutex.Lock() defer writeMutex.Unlock() done <- PingData(c) - }() + }) select { case err := <-done: @@ -219,14 +219,14 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon data = strings.TrimSuffix(data, "\r") if !strings.HasPrefix(data, "[DONE]") { info.SetFirstResponseTime() - + info.ReceivedResponseCount++ // 使用超时机制防止写操作阻塞 done := make(chan bool, 1) - go func() { + gopool.Go(func() { writeMutex.Lock() defer writeMutex.Unlock() done <- dataHandler(data) - }() + }) select { case success := <-done: