From 82553c4dcabd80a0e9ab60851e9a039e7dd6f619 Mon Sep 17 00:00:00 2001 From: mxyhi Date: Sun, 28 Jun 2026 11:13:09 +0800 Subject: [PATCH] fix(openai): preserve quota platform in usage billing --- .../handler/openai_chat_completions.go | 2 + backend/internal/handler/openai_embeddings.go | 2 + .../handler/openai_gateway_handler.go | 6 ++ backend/internal/handler/openai_images.go | 2 + .../openai_quota_platform_contract_test.go | 65 +++++++++++++++++++ .../service/openai_gateway_service.go | 10 ++- 6 files changed, 86 insertions(+), 1 deletion(-) create mode 100644 backend/internal/handler/openai_quota_platform_contract_test.go diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index a91ebf8bd9..ca43a2cff3 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -298,6 +298,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { clientIP := ip.GetClientIP(c) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account) + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) cyberBlocked := service.GetOpsCyberPolicy(c) != nil h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { @@ -312,6 +313,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { UserAgent: userAgent, IPAddress: clientIP, APIKeyService: h.apiKeyService, + QuotaPlatform: quotaPlatform, ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { diff --git a/backend/internal/handler/openai_embeddings.go b/backend/internal/handler/openai_embeddings.go index e538deacb0..a80c7f7d96 100644 --- a/backend/internal/handler/openai_embeddings.go +++ b/backend/internal/handler/openai_embeddings.go @@ -219,6 +219,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { clientIP := ip.GetClientIP(c) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ @@ -232,6 +233,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { UserAgent: userAgent, IPAddress: clientIP, APIKeyService: h.apiKeyService, + QuotaPlatform: quotaPlatform, ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), }); err != nil { logger.L().With( diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index cc62121d94..1a7cf357dc 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -513,6 +513,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { requestPayloadHash := service.HashUsageRequestPayload(body) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account) + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 cyberBlocked := service.GetOpsCyberPolicy(c) != nil @@ -529,6 +530,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, + QuotaPlatform: quotaPlatform, ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { @@ -926,6 +928,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { requestPayloadHash := service.HashUsageRequestPayload(body) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account) + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) cyberBlocked := service.GetOpsCyberPolicy(c) != nil h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { @@ -941,6 +944,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, + QuotaPlatform: quotaPlatform, ChannelUsageFields: channelMappingMsg.ToUsageFields(reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { @@ -1523,6 +1527,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account) + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) cyberBlocked := service.GetOpsCyberPolicy(c) != nil h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) { if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{ @@ -1537,6 +1542,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, + QuotaPlatform: quotaPlatform, ChannelUsageFields: channelMappingWS.ToUsageFields(reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 6ce053a3e0..31d245833d 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -341,6 +341,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { } inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) upstreamModel := "" if result != nil { @@ -359,6 +360,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, + QuotaPlatform: quotaPlatform, ChannelUsageFields: channelMapping.ToUsageFields(requestModel, upstreamModel), }); err != nil { logger.L().With( diff --git a/backend/internal/handler/openai_quota_platform_contract_test.go b/backend/internal/handler/openai_quota_platform_contract_test.go new file mode 100644 index 0000000000..12ceb3ca63 --- /dev/null +++ b/backend/internal/handler/openai_quota_platform_contract_test.go @@ -0,0 +1,65 @@ +package handler + +import ( + "go/ast" + "go/parser" + "go/token" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestOpenAIRecordUsageInputsCarryQuotaPlatform(t *testing.T) { + files := []string{ + "openai_gateway_handler.go", + "openai_chat_completions.go", + "openai_embeddings.go", + "openai_images.go", + } + + for _, name := range files { + t.Run(name, func(t *testing.T) { + fset := token.NewFileSet() + file, err := parser.ParseFile(fset, filepath.Join(".", name), nil, 0) + require.NoError(t, err) + + var missing []token.Position + ast.Inspect(file, func(node ast.Node) bool { + literal, ok := node.(*ast.CompositeLit) + if !ok || !isOpenAIRecordUsageInputLiteral(literal.Type) { + return true + } + if !compositeLiteralHasKey(literal, "QuotaPlatform") { + missing = append(missing, fset.Position(literal.Lbrace)) + } + return true + }) + + require.Empty(t, missing, "OpenAI usage post-billing must receive request-time QuotaPlatform") + }) + } +} + +func isOpenAIRecordUsageInputLiteral(expr ast.Expr) bool { + selector, ok := expr.(*ast.SelectorExpr) + if !ok { + return false + } + pkg, ok := selector.X.(*ast.Ident) + return ok && pkg.Name == "service" && selector.Sel.Name == "OpenAIRecordUsageInput" +} + +func compositeLiteralHasKey(literal *ast.CompositeLit, key string) bool { + for _, elt := range literal.Elts { + pair, ok := elt.(*ast.KeyValueExpr) + if !ok { + continue + } + ident, ok := pair.Key.(*ast.Ident) + if ok && ident.Name == key { + return true + } + } + return false +} diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 1def695c91..e345e23eef 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -6092,6 +6092,7 @@ type OpenAIRecordUsageInput struct { IPAddress string // 请求的客户端 IP 地址 RequestPayloadHash string APIKeyService APIKeyQuotaUpdater + QuotaPlatform string // user×platform quota platform resolved by the handler before async billing. // CyberBlocked 为 true 时把该用量行标记为 cyber（request_type=cyber），计费逻辑不变。 CyberBlocked bool ChannelUsageFields @@ -6366,6 +6367,13 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec return nil } + // Async usage billing runs outside the original request context, so it + // cannot recover ForcePlatform there. Fall back for internal/test callers. + quotaPlatform := input.QuotaPlatform + if quotaPlatform == "" { + quotaPlatform = PlatformFromAPIKey(apiKey) + } + billingErr := func() error { _, err := applyUsageBilling(ctx, requestID, usageLog, &postUsageBillingParams{ Cost: cost, @@ -6377,7 +6385,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec IsSubscriptionBill: isSubscriptionBilling, AccountRateMultiplier: accountRateMultiplier, APIKeyService: input.APIKeyService, - Platform: PlatformFromAPIKey(apiKey), + Platform: quotaPlatform, }, s.billingDeps(), s.usageBillingRepo) return err }()