diff --git a/controller/user.go b/controller/user.go index 9b8d931ec1f8..70a40a34318b 100644 --- a/controller/user.go +++ b/controller/user.go @@ -1498,7 +1498,11 @@ func UpdateUserSetting(c *gin.Context) { NotifyType: req.QuotaWarningType, QuotaWarningThreshold: req.QuotaWarningThreshold, UpstreamModelUpdateNotifyEnabled: upstreamModelUpdateNotifyEnabled, - AcceptUnsetRatioModel: req.AcceptUnsetModelRatioModel, + // F-27: only admins may enable accept-unset-ratio. Self-service by a + // regular user bypasses the "model price not configured" gate, letting + // them use deliberately unpriced models at the default ratio (37.5x), + // which is both an availability-control and pricing-control bypass. + AcceptUnsetRatioModel: user.Role >= common.RoleAdminUser && req.AcceptUnsetModelRatioModel, RecordIpLog: req.RecordIpLog, } diff --git a/go.mod b/go.mod index b0642f162db2..a060914995e3 100644 --- a/go.mod +++ b/go.mod @@ -90,7 +90,7 @@ require ( require ( github.com/DmitriyVTitov/size v1.5.0 // indirect github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6 // indirect - github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 // indirect github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 // indirect github.com/beorn7/perks v1.0.1 // indirect diff --git a/relay/channel/ali/rerank.go b/relay/channel/ali/rerank.go index ac2afbd3d3db..7403413f2342 100644 --- a/relay/channel/ali/rerank.go +++ b/relay/channel/ali/rerank.go @@ -59,6 +59,12 @@ func RerankHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayI CompletionTokens: 0, TotalTokens: aliResponse.Usage.TotalTokens, } + // F-53: fallback to the prompt estimate when the upstream omits usage so + // rerank requests are not billed as zero (F-26 residual for ali). + if usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 { + usage.PromptTokens = info.GetEstimatePromptTokens() + usage.TotalTokens = usage.PromptTokens + } rerankResponse := dto.RerankResponse{ Results: aliResponse.Output.Results, Usage: usage, diff --git a/relay/channel/aws/relay-aws.go b/relay/channel/aws/relay-aws.go index c4751b5af855..6e02d37e9299 100644 --- a/relay/channel/aws/relay-aws.go +++ b/relay/channel/aws/relay-aws.go @@ -346,6 +346,20 @@ func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) } // 构造OpenAI格式响应 + usage := dto.Usage{ + PromptTokens: novaResp.Usage.InputTokens, + CompletionTokens: novaResp.Usage.OutputTokens, + TotalTokens: novaResp.Usage.TotalTokens, + } + if usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 { + // F-60: fall back to the estimate when the upstream omits usage so + // Nova requests are not billed as zero. + var text string + if len(novaResp.Output.Message.Content) > 0 { + text = novaResp.Output.Message.Content[0].Text + } + usage = *service.ResponseText2Usage(c, text, info.UpstreamModelName, info.GetEstimatePromptTokens()) + } response := dto.OpenAITextResponse{ Id: helper.GetResponseID(c), Object: "chat.completion", @@ -359,13 +373,9 @@ func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) }, FinishReason: "stop", }}, - Usage: dto.Usage{ - PromptTokens: novaResp.Usage.InputTokens, - CompletionTokens: novaResp.Usage.OutputTokens, - TotalTokens: novaResp.Usage.TotalTokens, - }, + Usage: usage, } c.JSON(http.StatusOK, response) - return nil, &response.Usage + return nil, &usage } diff --git a/relay/channel/baidu/relay-baidu.go b/relay/channel/baidu/relay-baidu.go index ab74edc10c75..74cf2e65aef7 100644 --- a/relay/channel/baidu/relay-baidu.go +++ b/relay/channel/baidu/relay-baidu.go @@ -178,6 +178,15 @@ func baiduEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *ht return types.NewError(fmt.Errorf("%s", baiduResponse.ErrorMsg), types.ErrorCodeBadResponseBody), nil } fullTextResponse := embeddingResponseBaidu2OpenAI(&baiduResponse) + // F-26 family: Baidu embedding responses may omit usage; without a + // fallback the settle charges 0 and the pre-consume is refunded, making + // embeddings free. Fall back to the local request estimate. + if fullTextResponse.Usage.TotalTokens == 0 && fullTextResponse.Usage.PromptTokens == 0 && fullTextResponse.Usage.CompletionTokens == 0 { + fullTextResponse.Usage = dto.Usage{ + PromptTokens: info.GetEstimatePromptTokens(), + TotalTokens: info.GetEstimatePromptTokens(), + } + } jsonResponse, err := json.Marshal(fullTextResponse) if err != nil { return types.NewError(err, types.ErrorCodeBadResponseBody), nil diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index 2f424b32abdd..950c26489ede 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -238,6 +238,21 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud claudeInfo.Usage.ClaudeCacheCreation5mTokens = claudeResponse.Usage.GetCacheCreation5mTokens() claudeInfo.Usage.ClaudeCacheCreation1hTokens = claudeResponse.Usage.GetCacheCreation1hTokens() } + if claudeInfo.Usage.TotalTokens == 0 && + claudeInfo.Usage.PromptTokens == 0 && + claudeInfo.Usage.CompletionTokens == 0 { + // F-57: fall back to the estimate when the upstream omits usage so the + // response written to the client carries the same usage that + // settlement will bill (previously the fallback ran after the response + // was already serialized with zero usage). + var textBuilder strings.Builder + for _, block := range claudeResponse.Content { + if block.Text != nil && *block.Text != "" { + textBuilder.WriteString(*block.Text) + } + } + claudeInfo.Usage = service.ResponseText2Usage(c, textBuilder.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + } var responseData []byte switch info.RelayFormat { case types.RelayFormatOpenAI: diff --git a/relay/channel/cohere/relay-cohere.go b/relay/channel/cohere/relay-cohere.go index 30a3038c1363..14be1da8ea35 100644 --- a/relay/channel/cohere/relay-cohere.go +++ b/relay/channel/cohere/relay-cohere.go @@ -98,75 +98,83 @@ func cohereStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http } return 0, nil, nil }) - dataChan := make(chan string) - stopChan := make(chan bool) + // F-29: the reader goroutine must not block forever on an unbuffered + // channel when the client disconnects mid-stream (gin's c.Stream returns, + // leaving no receiver -> goroutine + upstream connection leak per request). + // A small buffer plus a done-select lets the reader exit promptly on + // client disconnect; stopChan send is non-blocking so EOF also cannot + // strand the goroutine. + dataChan := make(chan string, 64) go func() { for scanner.Scan() { data := scanner.Text() - dataChan <- data + select { + case dataChan <- data: + case <-c.Request.Context().Done(): + return + } } if err := scanner.Err(); err != nil { common.SysLog("error reading stream: " + err.Error()) } - stopChan <- true + close(dataChan) }() helper.SetEventStreamHeaders(c) isFirst := true c.Stream(func(w io.Writer) bool { - select { - case data := <-dataChan: - if isFirst { - isFirst = false - info.FirstResponseTime = time.Now() + data, ok := <-dataChan + if !ok { + c.Render(-1, common.CustomEvent{Data: "data: [DONE]"}) + return false + } + if isFirst { + isFirst = false + info.FirstResponseTime = time.Now() + } + data = strings.TrimSuffix(data, "\r") + var cohereResp CohereResponse + err := json.Unmarshal([]byte(data), &cohereResp) + if err != nil { + common.SysLog("error unmarshalling stream response: " + err.Error()) + return true + } + var openaiResp dto.ChatCompletionsStreamResponse + openaiResp.Id = responseId + openaiResp.Created = createdTime + openaiResp.Object = "chat.completion.chunk" + openaiResp.Model = info.UpstreamModelName + if cohereResp.IsFinished { + finishReason := stopReasonCohere2OpenAI(cohereResp.FinishReason) + openaiResp.Choices = []dto.ChatCompletionsStreamResponseChoice{ + { + Delta: dto.ChatCompletionsStreamResponseChoiceDelta{}, + Index: 0, + FinishReason: &finishReason, + }, } - data = strings.TrimSuffix(data, "\r") - var cohereResp CohereResponse - err := json.Unmarshal([]byte(data), &cohereResp) - if err != nil { - common.SysLog("error unmarshalling stream response: " + err.Error()) - return true + if cohereResp.Response != nil { + usage.PromptTokens = cohereResp.Response.Meta.BilledUnits.InputTokens + usage.CompletionTokens = cohereResp.Response.Meta.BilledUnits.OutputTokens } - var openaiResp dto.ChatCompletionsStreamResponse - openaiResp.Id = responseId - openaiResp.Created = createdTime - openaiResp.Object = "chat.completion.chunk" - openaiResp.Model = info.UpstreamModelName - if cohereResp.IsFinished { - finishReason := stopReasonCohere2OpenAI(cohereResp.FinishReason) - openaiResp.Choices = []dto.ChatCompletionsStreamResponseChoice{ - { - Delta: dto.ChatCompletionsStreamResponseChoiceDelta{}, - Index: 0, - FinishReason: &finishReason, - }, - } - if cohereResp.Response != nil { - usage.PromptTokens = cohereResp.Response.Meta.BilledUnits.InputTokens - usage.CompletionTokens = cohereResp.Response.Meta.BilledUnits.OutputTokens - } - } else { - openaiResp.Choices = []dto.ChatCompletionsStreamResponseChoice{ - { - Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ - Role: "assistant", - Content: &cohereResp.Text, - }, - Index: 0, + } else { + openaiResp.Choices = []dto.ChatCompletionsStreamResponseChoice{ + { + Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ + Role: "assistant", + Content: &cohereResp.Text, }, - } - responseText += cohereResp.Text + Index: 0, + }, } - jsonStr, err := json.Marshal(openaiResp) - if err != nil { - common.SysLog("error marshalling stream response: " + err.Error()) - return true - } - c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonStr)}) + responseText += cohereResp.Text + } + jsonStr, err := json.Marshal(openaiResp) + if err != nil { + common.SysLog("error marshalling stream response: " + err.Error()) return true - case <-stopChan: - c.Render(-1, common.CustomEvent{Data: "data: [DONE]"}) - return false } + c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonStr)}) + return true }) if usage.PromptTokens == 0 { usage = service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens()) @@ -190,6 +198,11 @@ func cohereHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo usage.PromptTokens = cohereResp.Meta.BilledUnits.InputTokens usage.CompletionTokens = cohereResp.Meta.BilledUnits.OutputTokens usage.TotalTokens = cohereResp.Meta.BilledUnits.InputTokens + cohereResp.Meta.BilledUnits.OutputTokens + if usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 { + // F-55: fall back to the estimate (prompt+completion+total) when the + // upstream omits usage, so generated output is not billed as zero. + usage = *service.ResponseText2Usage(c, cohereResp.Text, info.UpstreamModelName, info.GetEstimatePromptTokens()) + } var openaiResp dto.TextResponse openaiResp.Id = cohereResp.ResponseId diff --git a/relay/channel/dify/relay-dify.go b/relay/channel/dify/relay-dify.go index 2fcf2f5f2d58..07d8c67c5202 100644 --- a/relay/channel/dify/relay-dify.go +++ b/relay/channel/dify/relay-dify.go @@ -275,11 +275,17 @@ func difyHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respons if err != nil { return nil, types.NewError(err, types.ErrorCodeBadResponseBody) } + usage := difyResponse.MetaData.Usage + if usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 { + // F-58: fall back to the estimate when the upstream omits usage so + // non-stream Dify chat requests are not billed as zero. + usage = *service.ResponseText2Usage(c, difyResponse.Answer, info.UpstreamModelName, info.GetEstimatePromptTokens()) + } fullTextResponse := dto.OpenAITextResponse{ Id: difyResponse.ConversationId, Object: "chat.completion", Created: common.GetTimestamp(), - Usage: difyResponse.MetaData.Usage, + Usage: usage, } choice := dto.OpenAITextResponseChoice{ Index: 0, @@ -297,5 +303,5 @@ func difyHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respons c.Writer.Header().Set("Content-Type", "application/json") c.Writer.WriteHeader(resp.StatusCode) c.Writer.Write(jsonResponse) - return &difyResponse.MetaData.Usage, nil + return &usage, nil } diff --git a/relay/channel/mokaai/relay-mokaai.go b/relay/channel/mokaai/relay-mokaai.go index 71780216ce6b..d0711f8b3984 100644 --- a/relay/channel/mokaai/relay-mokaai.go +++ b/relay/channel/mokaai/relay-mokaai.go @@ -63,6 +63,14 @@ func mokaEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *htt if err != nil { return nil, types.NewError(err, types.ErrorCodeBadResponseBody) } + // F-54: fallback to the prompt estimate when the upstream omits usage so + // embedding requests are not billed as zero (F-26 residual for mokaai). + if baiduResponse.Usage.TotalTokens == 0 && + baiduResponse.Usage.PromptTokens == 0 && + baiduResponse.Usage.CompletionTokens == 0 { + baiduResponse.Usage.PromptTokens = info.GetEstimatePromptTokens() + baiduResponse.Usage.TotalTokens = baiduResponse.Usage.PromptTokens + } // if baiduResponse.ErrorMsg != "" { // return &dto.OpenAIErrorWithStatusCode{ // Error: dto.OpenAIError{ diff --git a/relay/channel/ollama/relay-ollama.go b/relay/channel/ollama/relay-ollama.go index e517a1e6aa03..08f90cdc3e92 100644 --- a/relay/channel/ollama/relay-ollama.go +++ b/relay/channel/ollama/relay-ollama.go @@ -325,6 +325,12 @@ func ollamaEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h data = append(data, dto.OpenAIEmbeddingResponseItem{Index: i, Object: "embedding", Embedding: emb}) } usage := &dto.Usage{PromptTokens: oResp.PromptEvalCount, CompletionTokens: 0, TotalTokens: oResp.PromptEvalCount} + // F-54: fallback to the prompt estimate when the upstream omits prompt + // eval count so embedding requests are not billed as zero. + if usage.TotalTokens == 0 && usage.PromptTokens == 0 { + usage.PromptTokens = info.GetEstimatePromptTokens() + usage.TotalTokens = usage.PromptTokens + } embResp := &dto.OpenAIEmbeddingResponse{Object: "list", Data: data, Model: info.UpstreamModelName, Usage: *usage} out, _ := common.Marshal(embResp) service.IOCopyBytesGracefully(c, resp, out) diff --git a/relay/channel/ollama/stream.go b/relay/channel/ollama/stream.go index a0d7839f9f6d..8edddf231580 100644 --- a/relay/channel/ollama/stream.go +++ b/relay/channel/ollama/stream.go @@ -104,6 +104,7 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http helper.SetEventStreamHeaders(c) scanner := helper.NewStreamScanner(resp.Body) usage := &dto.Usage{} + var responseText strings.Builder var model = info.UpstreamModelName var responseId = common.GetUUID() var created = time.Now().Unix() @@ -148,6 +149,7 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http }}, } if content != "" { + responseText.WriteString(content) delta.Choices[0].Delta.SetContentString(content) } if chunk.Message != nil && len(chunk.Message.Thinking) > 0 { @@ -190,6 +192,11 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http _ = helper.StringData(c, string(data)) } } + if usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 { + // F-62: apply the estimate before the usage frame is emitted so the + // client sees the same usage that settlement will bill. + usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + } // emit usage frame if final := helper.GenerateFinalUsageResponse(responseId, created, model, *usage); final != nil { if data, err := common.Marshal(final); err == nil { @@ -303,6 +310,11 @@ func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R } created := toUnix(lastChunk.CreatedAt) usage := &dto.Usage{PromptTokens: lastChunk.PromptEvalCount, CompletionTokens: lastChunk.EvalCount, TotalTokens: lastChunk.PromptEvalCount + lastChunk.EvalCount} + if usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 { + // F-62: fall back to the estimate when the upstream omits usage so + // non-stream ollama chat is not billed as zero. + usage = service.ResponseText2Usage(c, aggContent.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + } content := aggContent.String() finishReason := lastChunk.DoneReason if finishReason == "" { diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index ceca1af3b381..5c36fdfd3c41 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -48,6 +48,19 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http usage.PromptTokensDetails.CacheWriteTokens = responsesResponse.Usage.InputTokensDetails.CacheWriteTokens } } + if usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 { + // F-56: fall back to the estimate when the upstream omits usage so + // non-stream Responses requests are not billed as zero. + var outText strings.Builder + for _, out := range responsesResponse.Output { + for _, c := range out.Content { + if c.Text != "" { + outText.WriteString(c.Text) + } + } + } + usage = *service.ResponseText2Usage(c, outText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + } // Count actual tool invocations from Output (not tool declarations). for _, output := range responsesResponse.Output { switch output.Type { diff --git a/relay/channel/openai/relay_responses_compact.go b/relay/channel/openai/relay_responses_compact.go index ff30d36e417b..ea4a4c875ed4 100644 --- a/relay/channel/openai/relay_responses_compact.go +++ b/relay/channel/openai/relay_responses_compact.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/service" @@ -40,6 +41,13 @@ func OaiResponsesCompactionHandler(c *gin.Context, resp *http.Response) (*dto.Us usage.PromptTokensDetails.CacheWriteTokens = compactResp.Usage.InputTokensDetails.CacheWriteTokens } } + // F-26 family: when the upstream compaction response omits usage, fall + // back to the local request estimate so the pre-consume is not fully + // refunded (free compaction). + if usage.TotalTokens == 0 { + usage.PromptTokens = common.GetContextKeyInt(c, constant.ContextKeyEstimatedTokens) + usage.TotalTokens = usage.PromptTokens + } return &usage, nil } diff --git a/relay/channel/siliconflow/relay-siliconflow.go b/relay/channel/siliconflow/relay-siliconflow.go index 35079eace36e..335490a50aa7 100644 --- a/relay/channel/siliconflow/relay-siliconflow.go +++ b/relay/channel/siliconflow/relay-siliconflow.go @@ -29,6 +29,12 @@ func siliconflowRerankHandler(c *gin.Context, info *relaycommon.RelayInfo, resp CompletionTokens: siliconflowResp.Meta.Tokens.OutputTokens, TotalTokens: siliconflowResp.Meta.Tokens.InputTokens + siliconflowResp.Meta.Tokens.OutputTokens, } + // F-53: fallback to the prompt estimate when the upstream omits usage so + // rerank requests are not billed as zero (F-26 residual for siliconflow). + if usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 { + usage.PromptTokens = info.GetEstimatePromptTokens() + usage.TotalTokens = usage.PromptTokens + } rerankResp := &dto.RerankResponse{ Results: siliconflowResp.Results, Usage: *usage, diff --git a/relay/channel/xunfei/adaptor.go b/relay/channel/xunfei/adaptor.go index 2f8112f48b5d..603ba33f2270 100644 --- a/relay/channel/xunfei/adaptor.go +++ b/relay/channel/xunfei/adaptor.go @@ -88,9 +88,9 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom return nil, types.NewError(errors.New("request is nil"), types.ErrorCodeInvalidRequest) } if info.IsStream { - usage, err = xunfeiStreamHandler(c, *a.request, splits[0], splits[1], splits[2]) + usage, err = xunfeiStreamHandler(c, info, *a.request, splits[0], splits[1], splits[2]) } else { - usage, err = xunfeiHandler(c, *a.request, splits[0], splits[1], splits[2]) + usage, err = xunfeiHandler(c, info, *a.request, splits[0], splits[1], splits[2]) } return } diff --git a/relay/channel/xunfei/relay-xunfei.go b/relay/channel/xunfei/relay-xunfei.go index fc80ba77c256..c9f1ca91d522 100644 --- a/relay/channel/xunfei/relay-xunfei.go +++ b/relay/channel/xunfei/relay-xunfei.go @@ -1,6 +1,7 @@ package xunfei import ( + "context" "crypto/hmac" "crypto/sha256" "encoding/base64" @@ -13,9 +14,11 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" + relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relay/helper" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/types" + "github.com/QuantumNous/new-api/service" "github.com/samber/lo" "github.com/gin-gonic/gin" @@ -128,58 +131,61 @@ func buildXunfeiAuthUrl(hostUrl string, apiKey, apiSecret string) string { return callUrl } -func xunfeiStreamHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId string, apiSecret string, apiKey string) (*dto.Usage, *types.NewAPIError) { +func xunfeiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, textRequest dto.GeneralOpenAIRequest, appId string, apiSecret string, apiKey string) (*dto.Usage, *types.NewAPIError) { domain, authUrl := getXunfeiAuthUrl(c, apiKey, apiSecret, textRequest.Model) - dataChan, stopChan, err := xunfeiMakeRequest(textRequest, domain, authUrl, appId) + dataChan, err := xunfeiMakeRequest(textRequest, domain, authUrl, appId, c.Request.Context()) if err != nil { return nil, types.NewError(err, types.ErrorCodeDoRequestFailed) } helper.SetEventStreamHeaders(c) var usage dto.Usage + var responseText strings.Builder c.Stream(func(w io.Writer) bool { - select { - case xunfeiResponse := <-dataChan: - usage.PromptTokens += xunfeiResponse.Payload.Usage.Text.PromptTokens - usage.CompletionTokens += xunfeiResponse.Payload.Usage.Text.CompletionTokens - usage.TotalTokens += xunfeiResponse.Payload.Usage.Text.TotalTokens - response := streamResponseXunfei2OpenAI(&xunfeiResponse) - jsonResponse, err := json.Marshal(response) - if err != nil { - common.SysLog("error marshalling stream response: " + err.Error()) - return true - } - c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)}) - return true - case <-stopChan: + xunfeiResponse, ok := <-dataChan + if !ok { c.Render(-1, common.CustomEvent{Data: "data: [DONE]"}) return false } + usage.PromptTokens += xunfeiResponse.Payload.Usage.Text.PromptTokens + usage.CompletionTokens += xunfeiResponse.Payload.Usage.Text.CompletionTokens + usage.TotalTokens += xunfeiResponse.Payload.Usage.Text.TotalTokens + response := streamResponseXunfei2OpenAI(&xunfeiResponse) + if len(response.Choices) > 0 { + responseText.WriteString(response.Choices[0].Delta.GetContentString()) + } + jsonResponse, err := json.Marshal(response) + if err != nil { + common.SysLog("error marshalling stream response: " + err.Error()) + return true + } + c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)}) + return true }) + if usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 { + // F-59: fall back to the estimate when the upstream omits usage so + // streamed xunfei chat is not billed as zero. + usage = *service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + } return &usage, nil } -func xunfeiHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId string, apiSecret string, apiKey string) (*dto.Usage, *types.NewAPIError) { +func xunfeiHandler(c *gin.Context, info *relaycommon.RelayInfo, textRequest dto.GeneralOpenAIRequest, appId string, apiSecret string, apiKey string) (*dto.Usage, *types.NewAPIError) { domain, authUrl := getXunfeiAuthUrl(c, apiKey, apiSecret, textRequest.Model) - dataChan, stopChan, err := xunfeiMakeRequest(textRequest, domain, authUrl, appId) + dataChan, err := xunfeiMakeRequest(textRequest, domain, authUrl, appId, c.Request.Context()) if err != nil { return nil, types.NewError(err, types.ErrorCodeDoRequestFailed) } var usage dto.Usage var content string var xunfeiResponse XunfeiChatResponse - stop := false - for !stop { - select { - case xunfeiResponse = <-dataChan: - if len(xunfeiResponse.Payload.Choices.Text) == 0 { - continue - } - content += xunfeiResponse.Payload.Choices.Text[0].Content - usage.PromptTokens += xunfeiResponse.Payload.Usage.Text.PromptTokens - usage.CompletionTokens += xunfeiResponse.Payload.Usage.Text.CompletionTokens - usage.TotalTokens += xunfeiResponse.Payload.Usage.Text.TotalTokens - case stop = <-stopChan: + for xunfeiResponse = range dataChan { + if len(xunfeiResponse.Payload.Choices.Text) == 0 { + continue } + content += xunfeiResponse.Payload.Choices.Text[0].Content + usage.PromptTokens += xunfeiResponse.Payload.Usage.Text.PromptTokens + usage.CompletionTokens += xunfeiResponse.Payload.Usage.Text.CompletionTokens + usage.TotalTokens += xunfeiResponse.Payload.Usage.Text.TotalTokens } if len(xunfeiResponse.Payload.Choices.Text) == 0 { xunfeiResponse.Payload.Choices.Text = []XunfeiChatResponseTextItem{ @@ -190,6 +196,12 @@ func xunfeiHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId s } xunfeiResponse.Payload.Choices.Text[0].Content = content + if usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 { + // F-59: fall back to the estimate when the upstream omits usage so + // non-stream xunfei chat is not billed as zero. + usage = *service.ResponseText2Usage(c, content, info.UpstreamModelName, info.GetEstimatePromptTokens()) + } + response := responseXunfei2OpenAI(&xunfeiResponse) jsonResponse, err := json.Marshal(response) if err != nil { @@ -200,26 +212,28 @@ func xunfeiHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId s return &usage, nil } -func xunfeiMakeRequest(textRequest dto.GeneralOpenAIRequest, domain, authUrl, appId string) (chan XunfeiChatResponse, chan bool, error) { +func xunfeiMakeRequest(textRequest dto.GeneralOpenAIRequest, domain, authUrl, appId string, ctx context.Context) (chan XunfeiChatResponse, error) { d := websocket.Dialer{ HandshakeTimeout: 5 * time.Second, } conn, resp, err := d.Dial(authUrl, nil) if err != nil || resp.StatusCode != 101 { - return nil, nil, err + return nil, err } data := requestOpenAI2Xunfei(textRequest, appId, domain) err = conn.WriteJSON(data) if err != nil { - return nil, nil, err + return nil, err } - dataChan := make(chan XunfeiChatResponse) - stopChan := make(chan bool) + // F-29 family: buffered channels + done-select so a client disconnect + // (ctx cancel) cannot strand the reader goroutine on an unbuffered send. + dataChan := make(chan XunfeiChatResponse, 64) go func() { defer func() { conn.Close() + close(dataChan) }() for { _, msg, err := conn.ReadMessage() @@ -233,7 +247,11 @@ func xunfeiMakeRequest(textRequest dto.GeneralOpenAIRequest, domain, authUrl, ap common.SysLog("error unmarshalling stream response: " + err.Error()) break } - dataChan <- response + select { + case dataChan <- response: + case <-ctx.Done(): + return + } if response.Payload.Choices.Status == 2 { if err != nil { common.SysLog("error closing websocket connection: " + err.Error()) @@ -241,10 +259,9 @@ func xunfeiMakeRequest(textRequest dto.GeneralOpenAIRequest, domain, authUrl, ap break } } - stopChan <- true }() - return dataChan, stopChan, nil + return dataChan, nil } func apiVersion2domain(apiVersion string) string { diff --git a/relay/channel/zhipu/relay-zhipu.go b/relay/channel/zhipu/relay-zhipu.go index 0c280e2b705d..0c5a9f786f33 100644 --- a/relay/channel/zhipu/relay-zhipu.go +++ b/relay/channel/zhipu/relay-zhipu.go @@ -157,10 +157,14 @@ func streamMetaResponseZhipu2OpenAI(zhipuResponse *ZhipuStreamMetaResponse) (*dt func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { var usage *dto.Usage + var responseText strings.Builder scanner := helper.NewStreamScanner(resp.Body) scanner.Split(bufio.ScanLines) - dataChan := make(chan string) - metaChan := make(chan string) + // F-29 family: buffer the channels and let the reader exit on client + // disconnect; otherwise a mid-stream disconnect strands the goroutine + // forever on an unbuffered send (goroutine + connection leak per request). + dataChan := make(chan string, 64) + metaChan := make(chan string, 16) stopChan := make(chan bool) go func() { for scanner.Scan() { @@ -171,25 +175,47 @@ func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. continue } if line[:5] == "data:" { - dataChan <- line[5:] + select { + case dataChan <- line[5:]: + case <-c.Request.Context().Done(): + return + } if i != len(lines)-1 { - dataChan <- "\n" + select { + case dataChan <- "\n": + case <-c.Request.Context().Done(): + return + } } } else if line[:5] == "meta:" { - metaChan <- line[5:] + select { + case metaChan <- line[5:]: + case <-c.Request.Context().Done(): + return + } } } } if err := scanner.Err(); err != nil { common.SysLog("error reading stream: " + err.Error()) } - stopChan <- true + // Deliver the terminal signal reliably: a non-blocking send to an + // unbuffered channel drops the completion when the consumer is busy + // draining data/meta frames, leaving the stream hanging after EOF. + select { + case stopChan <- true: + case <-c.Request.Context().Done(): + return + } }() helper.SetEventStreamHeaders(c) c.Stream(func(w io.Writer) bool { select { case data := <-dataChan: response := streamResponseZhipu2OpenAI(data) + if len(response.Choices) > 0 { + responseText.WriteString(response.Choices[0].Delta.GetContentString()) + } jsonResponse, err := json.Marshal(response) if err != nil { common.SysLog("error marshalling stream response: " + err.Error()) @@ -219,6 +245,11 @@ func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. } }) service.CloseResponseBodyGracefully(resp) + if usage == nil || (usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0) { + // F-55: zhipu streams may omit the usage meta event entirely; fall + // back to the estimate so chat is not billed as zero. + usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + } return usage, nil } @@ -240,6 +271,17 @@ func zhipuHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respon }, resp.StatusCode) } fullTextResponse := responseZhipu2OpenAI(&zhipuResponse) + if fullTextResponse.Usage.TotalTokens == 0 && + fullTextResponse.Usage.PromptTokens == 0 && + fullTextResponse.Usage.CompletionTokens == 0 { + // F-55: fall back to the estimate (prompt+completion+total) when the + // upstream omits usage, so generated output is not billed as zero. + var textBuilder strings.Builder + for _, choice := range zhipuResponse.Data.Choices { + textBuilder.WriteString(strings.Trim(choice.Content, "\"")) + } + fullTextResponse.Usage = *service.ResponseText2Usage(c, textBuilder.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + } jsonResponse, err := json.Marshal(fullTextResponse) if err != nil { return nil, types.NewError(err, types.ErrorCodeBadResponseBody) diff --git a/relay/common_handler/rerank.go b/relay/common_handler/rerank.go index 1e7658da6968..06c8a22b18ac 100644 --- a/relay/common_handler/rerank.go +++ b/relay/common_handler/rerank.go @@ -65,6 +65,14 @@ func RerankHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } + // F-26: upstreams like Jina v1 / OpenAI-compatible rerank endpoints may + // omit usage entirely. Without a fallback the settle charges 0 and the + // pre-consume is fully refunded -> unlimited free rerank. Fall back to + // the local request estimate whenever upstream usage is missing. + if jinaResp.Usage.TotalTokens == 0 && jinaResp.Usage.PromptTokens == 0 && jinaResp.Usage.CompletionTokens == 0 { + jinaResp.Usage.PromptTokens = info.GetEstimatePromptTokens() + jinaResp.Usage.TotalTokens = jinaResp.Usage.PromptTokens + } jinaResp.Usage.PromptTokens = jinaResp.Usage.TotalTokens } diff --git a/relay/image_handler.go b/relay/image_handler.go index 690f229f3ff0..8c2a6b16643b 100644 --- a/relay/image_handler.go +++ b/relay/image_handler.go @@ -122,6 +122,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type imageN = *request.N } + usageZero := usage.(*dto.Usage).TotalTokens == 0 && usage.(*dto.Usage).PromptTokens == 0 if usage.(*dto.Usage).TotalTokens == 0 { usage.(*dto.Usage).TotalTokens = 1 } @@ -129,6 +130,15 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type usage.(*dto.Usage).PromptTokens = 1 } + // F-52: ensure per-image count ratio is set for UsePrice-priced image + // channels whose adaptor does not report it (e.g. Replicate), so n>1 + // requests are billed per generated image instead of once. Only apply the + // multiplier when the adaptor reported no usage: adaptors that already + // embed the image count in usage (e.g. ali) must not be multiplied again. + if usageZero && imageN > 0 && info.PriceData.UsePrice && !info.PriceData.HasOtherRatio("n") { + info.PriceData.AddOtherRatio("n", float64(imageN)) + } + quality := request.Quality if quality == "" { quality = "standard"