Merge pull request #3534 from mxyhi/fix/openai-quota-platform-post-billing

fix(openai): preserve quota platform in usage billing
This commit is contained in:
Wesley Liddick
2026-06-29 09:21:44 +08:00
committed by GitHub
6 changed files with 86 additions and 1 deletions
@@ -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 {
@@ -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(
@@ -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 {
@@ -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(
@@ -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
}
@@ -6096,6 +6096,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
@@ -6370,6 +6371,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,
@@ -6381,7 +6389,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
}()