mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
1060 lines
44 KiB
Go
1060 lines
44 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"log/slog"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||
)
|
||
|
||
func (s *GatewayService) getUserGroupRateMultiplier(ctx context.Context, userID, groupID int64, groupDefaultMultiplier float64) float64 {
|
||
if s == nil {
|
||
return groupDefaultMultiplier
|
||
}
|
||
resolver := s.userGroupRateResolver
|
||
if resolver == nil {
|
||
resolver = newUserGroupRateResolver(
|
||
s.userGroupRateRepo,
|
||
s.userGroupRateCache,
|
||
resolveUserGroupRateCacheTTL(s.cfg),
|
||
&s.userGroupRateSF,
|
||
"service.gateway",
|
||
)
|
||
}
|
||
return resolver.Resolve(ctx, userID, groupID, groupDefaultMultiplier)
|
||
}
|
||
|
||
// ResolveUserGroupRateMultiplier resolves the same cached multiplier used by usage billing.
|
||
func (s *GatewayService) ResolveUserGroupRateMultiplier(ctx context.Context, userID, groupID int64, groupDefaultMultiplier float64) float64 {
|
||
return s.getUserGroupRateMultiplier(ctx, userID, groupID, groupDefaultMultiplier)
|
||
}
|
||
|
||
// RecordUsageInput 记录使用量的输入参数。
|
||
// 异步 worker 只接收计费所需快照，不能持有 ParsedRequest/RequestBodyRef 这类大请求体引用。
|
||
type RecordUsageInput struct {
|
||
Result *ForwardResult
|
||
APIKey *APIKey
|
||
User *User
|
||
Account *Account
|
||
Subscription *UserSubscription // 可选：订阅信息
|
||
InboundEndpoint string // 入站端点（客户端请求路径）
|
||
UpstreamEndpoint string // 上游端点（标准化后的上游路径）
|
||
UserAgent string // 请求的 User-Agent
|
||
IPAddress string // 请求的客户端 IP 地址
|
||
SessionID string // 客户端显式会话标识（session_id / X-Session-Id 等请求头），仅用于用量行会话关联
|
||
RequestPayloadHash string // 请求体语义哈希，用于降低 request_id 误复用时的静默误去重风险
|
||
ForceCacheBilling bool // 强制缓存计费：将 input_tokens 转为 cache_read 计费（用于粘性会话切换）
|
||
APIKeyService APIKeyQuotaUpdater // 可选：用于更新API Key配额
|
||
QuotaPlatform string // user×platform 配额计量平台：handler 在请求 ctx 内经 QuotaPlatform() 算定后传入（后扣运行在 worker 池 background ctx 上，取不到 ForcePlatform）
|
||
|
||
ChannelUsageFields // 渠道映射信息（由 handler 在 Forward 前解析）
|
||
}
|
||
|
||
// APIKeyQuotaUpdater defines the interface for updating API Key quota and rate limit usage
|
||
type APIKeyQuotaUpdater interface {
|
||
UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error
|
||
UpdateRateLimitUsage(ctx context.Context, apiKeyID int64, cost float64) error
|
||
}
|
||
|
||
type apiKeyAuthCacheInvalidator interface {
|
||
InvalidateAuthCacheByKey(ctx context.Context, key string)
|
||
}
|
||
|
||
type usageLogBestEffortWriter interface {
|
||
CreateBestEffort(ctx context.Context, log *UsageLog) error
|
||
}
|
||
|
||
// postUsageBillingParams 统一扣费所需的参数
|
||
type postUsageBillingParams struct {
|
||
Cost *CostBreakdown
|
||
User *User
|
||
APIKey *APIKey
|
||
Account *Account
|
||
Subscription *UserSubscription
|
||
RequestPayloadHash string
|
||
IsSubscriptionBill bool
|
||
AccountRateMultiplier float64
|
||
APIKeyService APIKeyQuotaUpdater
|
||
Platform string // 来自 APIKey 关联 Group 的平台标识
|
||
}
|
||
|
||
// PlatformFromAPIKey 从 APIKey 关联的 Group 推导 platform 名称。
|
||
// apiKey 为 nil 或 Group 信息缺失时返回空串（调用方据此 short-circuit quota 累加）。
|
||
// 导出供 handler 层调用。
|
||
func PlatformFromAPIKey(apiKey *APIKey) string {
|
||
if apiKey == nil || apiKey.Group == nil {
|
||
return ""
|
||
}
|
||
return apiKey.Group.Platform
|
||
}
|
||
|
||
// QuotaPlatform 返回 user×platform 配额计量使用的平台标识。
|
||
// 强制平台路由（如 /antigravity）优先按 ctx 中的 ForcePlatform 计量，否则回退到
|
||
// APIKey 关联 Group 的平台。
|
||
//
|
||
// 注意：必须用带 ForcePlatform 的请求 context 调用（如 handler 的 c.Request.Context()）。
|
||
// 后扣运行在 worker 池的 background ctx 上没有 ForcePlatform，因此后扣平台由 handler
|
||
// 预先算定、经 RecordUsageInput.QuotaPlatform 传入，不要在后扣链路用 worker ctx 调用本函数。
|
||
func QuotaPlatform(ctx context.Context, apiKey *APIKey) string {
|
||
if ctx != nil {
|
||
if fp, ok := ctx.Value(ctxkey.ForcePlatform).(string); ok && fp != "" {
|
||
return fp
|
||
}
|
||
}
|
||
if platform, ok := ResolvedTargetPlatformFromContext(ctx); ok {
|
||
return platform
|
||
}
|
||
platform := PlatformFromAPIKey(apiKey)
|
||
if platform == PlatformComposite {
|
||
return ""
|
||
}
|
||
return platform
|
||
}
|
||
|
||
func (p *postUsageBillingParams) shouldDeductAPIKeyQuota() bool {
|
||
return p.Cost.ActualCost > 0 && p.APIKey.Quota > 0 && p.APIKeyService != nil
|
||
}
|
||
|
||
func (p *postUsageBillingParams) shouldUpdateRateLimits() bool {
|
||
return p.Cost.ActualCost > 0 && p.APIKey.HasRateLimits() && p.APIKeyService != nil
|
||
}
|
||
|
||
func (p *postUsageBillingParams) shouldUpdateAccountQuota() bool {
|
||
return p.Cost.TotalCost > 0 && p.Account.IsAPIKeyOrBedrock() && p.Account.HasAnyQuotaLimit()
|
||
}
|
||
|
||
// postUsageBilling is the legacy fallback billing path used when the unified
|
||
// billing repo is unavailable (nil). Production uses applyUsageBilling → repo.Apply
|
||
// for atomic billing. This path only runs in tests or degraded mode.
|
||
func postUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *billingDeps) {
|
||
billingCtx, cancel := detachedBillingContext(ctx)
|
||
defer cancel()
|
||
|
||
cost := p.Cost
|
||
|
||
if p.IsSubscriptionBill {
|
||
// Subscription usage tracked by ActualCost so group rate multiplier
|
||
// consumes the quota at the expected speed.
|
||
if cost.ActualCost > 0 {
|
||
if err := deps.userSubRepo.IncrementUsage(billingCtx, p.Subscription.ID, cost.ActualCost); err != nil {
|
||
slog.Error("increment subscription usage failed", "subscription_id", p.Subscription.ID, "error", err)
|
||
}
|
||
}
|
||
} else {
|
||
if cost.ActualCost > 0 {
|
||
if err := deps.userRepo.DeductBalance(billingCtx, p.User.ID, cost.ActualCost); err != nil {
|
||
slog.Error("deduct balance failed", "user_id", p.User.ID, "error", err)
|
||
} else if deps.billingCacheService != nil {
|
||
if err := deps.billingCacheService.InvalidateUserBalance(billingCtx, p.User.ID); err != nil {
|
||
slog.Warn("invalidate balance cache after legacy deduction failed", "user_id", p.User.ID, "error", err)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
if p.shouldDeductAPIKeyQuota() {
|
||
if err := p.APIKeyService.UpdateQuotaUsed(billingCtx, p.APIKey.ID, cost.ActualCost); err != nil {
|
||
slog.Error("update api key quota failed", "api_key_id", p.APIKey.ID, "error", err)
|
||
}
|
||
}
|
||
|
||
if p.shouldUpdateRateLimits() {
|
||
if err := p.APIKeyService.UpdateRateLimitUsage(billingCtx, p.APIKey.ID, cost.ActualCost); err != nil {
|
||
slog.Error("update api key rate limit usage failed", "api_key_id", p.APIKey.ID, "error", err)
|
||
}
|
||
}
|
||
|
||
if p.shouldUpdateAccountQuota() {
|
||
accountCost := cost.TotalCost * p.AccountRateMultiplier
|
||
if err := deps.accountRepo.IncrementQuotaUsed(billingCtx, p.Account.ID, accountCost); err != nil {
|
||
slog.Error("increment account quota used failed", "account_id", p.Account.ID, "cost", accountCost, "error", err)
|
||
}
|
||
}
|
||
|
||
// Platform quota 累加（legacy 兜底路径）：仅对 standard（余额）模式生效；订阅模式豁免；仅对有 limit 的用户写
|
||
// - HasUserPlatformQuotaLimit 守卫:与正常路径对齐，无 limit 公司跳过
|
||
// - 新增 Redis 同步写:enforcement 走 Redis，legacy 路径也必须同步写，否则 preflight 看不到消费
|
||
// - flusher_enabled=false（降级）:保留原有同步直写 DB
|
||
// - flusher_enabled=true:跳过直写 DB，由 flusher 异步批量刷（markDirty 在 IncrementUserPlatformQuotaUsage 内部完成）
|
||
// - 失败仅记 ALERT log + counter，不阻断主扣费流程
|
||
if !p.IsSubscriptionBill && p.Platform != "" && cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil {
|
||
if deps.billingCacheService.HasUserPlatformQuotaLimit(billingCtx, p.User.ID, p.Platform) {
|
||
deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, cost.ActualCost)
|
||
if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled {
|
||
// 降级路径:flusher 未启用时保留原有同步直写 DB
|
||
if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(billingCtx, p.User.ID, p.Platform, cost.ActualCost, time.Now().UTC()); err != nil {
|
||
userPlatformQuotaDBIncrLegacyErrorTotal.Add(1)
|
||
logger.LegacyPrintf("service.gateway", "ALERT: legacy incr user platform quota DB failed user=%d platform=%s cost=%f: %v", p.User.ID, p.Platform, cost.ActualCost, err)
|
||
}
|
||
}
|
||
// flusher_enabled=true:不直写 DB，flusher 异步批量刷
|
||
}
|
||
}
|
||
|
||
// NOTE: finalizePostUsageBilling is NOT called here to avoid double-queuing
|
||
// cache updates. The legacy path does DB writes directly; the finalize path
|
||
// does cache queue + notifications. Notifications are dispatched separately
|
||
// by the caller after recording the usage log.
|
||
}
|
||
|
||
func resolveUsageBillingRequestID(ctx context.Context, upstreamRequestID string) string {
|
||
if ctx != nil {
|
||
if clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(clientRequestID) != "" {
|
||
return "client:" + strings.TrimSpace(clientRequestID)
|
||
}
|
||
if requestID, _ := ctx.Value(ctxkey.RequestID).(string); strings.TrimSpace(requestID) != "" {
|
||
return "local:" + strings.TrimSpace(requestID)
|
||
}
|
||
}
|
||
if requestID := strings.TrimSpace(upstreamRequestID); requestID != "" {
|
||
return requestID
|
||
}
|
||
return "generated:" + generateRequestID()
|
||
}
|
||
|
||
func resolveUsageBillingPayloadFingerprint(ctx context.Context, requestPayloadHash string) string {
|
||
if payloadHash := strings.TrimSpace(requestPayloadHash); payloadHash != "" {
|
||
return payloadHash
|
||
}
|
||
if ctx != nil {
|
||
if clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(clientRequestID) != "" {
|
||
return "client:" + strings.TrimSpace(clientRequestID)
|
||
}
|
||
if requestID, _ := ctx.Value(ctxkey.RequestID).(string); strings.TrimSpace(requestID) != "" {
|
||
return "local:" + strings.TrimSpace(requestID)
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
func buildUsageBillingCommand(requestID string, usageLog *UsageLog, p *postUsageBillingParams) *UsageBillingCommand {
|
||
if p == nil || p.Cost == nil || p.APIKey == nil || p.User == nil || p.Account == nil {
|
||
return nil
|
||
}
|
||
|
||
cmd := &UsageBillingCommand{
|
||
RequestID: requestID,
|
||
APIKeyID: p.APIKey.ID,
|
||
UserID: p.User.ID,
|
||
AccountID: p.Account.ID,
|
||
AccountType: p.Account.Type,
|
||
RequestPayloadHash: strings.TrimSpace(p.RequestPayloadHash),
|
||
}
|
||
if usageLog != nil {
|
||
cmd.Model = usageLog.Model
|
||
cmd.BillingType = usageLog.BillingType
|
||
cmd.InputTokens = usageLog.InputTokens
|
||
cmd.OutputTokens = usageLog.OutputTokens
|
||
cmd.CacheCreationTokens = usageLog.CacheCreationTokens
|
||
cmd.CacheReadTokens = usageLog.CacheReadTokens
|
||
cmd.ImageCount = usageLog.ImageCount
|
||
if usageLog.ServiceTier != nil {
|
||
cmd.ServiceTier = *usageLog.ServiceTier
|
||
}
|
||
if usageLog.ReasoningEffort != nil {
|
||
cmd.ReasoningEffort = *usageLog.ReasoningEffort
|
||
}
|
||
if usageLog.SubscriptionID != nil {
|
||
cmd.SubscriptionID = usageLog.SubscriptionID
|
||
}
|
||
}
|
||
|
||
// Record subscription / balance cost using ActualCost so the group (and any
|
||
// user-specific) rate multiplier consumes subscription quota at the expected
|
||
// speed. TotalCost remains the raw (pre-multiplier) value; downstream guards
|
||
// on "> 0" still correctly skip free subscriptions (RateMultiplier == 0).
|
||
if p.IsSubscriptionBill && p.Subscription != nil && p.Cost.TotalCost > 0 {
|
||
cmd.SubscriptionID = &p.Subscription.ID
|
||
cmd.SubscriptionCost = p.Cost.ActualCost
|
||
} else if p.Cost.ActualCost > 0 {
|
||
cmd.BalanceCost = p.Cost.ActualCost
|
||
}
|
||
|
||
if p.shouldDeductAPIKeyQuota() {
|
||
cmd.APIKeyQuotaCost = p.Cost.ActualCost
|
||
}
|
||
if p.shouldUpdateRateLimits() {
|
||
cmd.APIKeyRateLimitCost = p.Cost.ActualCost
|
||
}
|
||
if p.shouldUpdateAccountQuota() {
|
||
cmd.AccountQuotaCost = p.Cost.TotalCost * p.AccountRateMultiplier
|
||
}
|
||
|
||
cmd.Normalize()
|
||
return cmd
|
||
}
|
||
|
||
func applyUsageBilling(ctx context.Context, requestID string, usageLog *UsageLog, p *postUsageBillingParams, deps *billingDeps, repo UsageBillingRepository) (bool, error) {
|
||
if p == nil || deps == nil {
|
||
return false, nil
|
||
}
|
||
|
||
cmd := buildUsageBillingCommand(requestID, usageLog, p)
|
||
if cmd == nil || cmd.RequestID == "" || repo == nil {
|
||
postUsageBilling(ctx, p, deps)
|
||
return true, nil
|
||
}
|
||
|
||
billingCtx, cancel := detachedBillingContext(ctx)
|
||
defer cancel()
|
||
|
||
result, err := repo.Apply(billingCtx, cmd)
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
|
||
if result == nil || !result.Applied {
|
||
deps.deferredService.ScheduleLastUsedUpdate(p.Account.ID)
|
||
return false, nil
|
||
}
|
||
|
||
if result.APIKeyQuotaExhausted {
|
||
if invalidator, ok := p.APIKeyService.(apiKeyAuthCacheInvalidator); ok && p.APIKey != nil && p.APIKey.Key != "" {
|
||
invalidator.InvalidateAuthCacheByKey(billingCtx, p.APIKey.Key)
|
||
}
|
||
}
|
||
|
||
finalizePostUsageBilling(billingCtx, p, deps, result)
|
||
return true, nil
|
||
}
|
||
|
||
func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) {
|
||
if p == nil || p.Cost == nil || deps == nil {
|
||
return
|
||
}
|
||
|
||
if p.IsSubscriptionBill {
|
||
if p.Cost.ActualCost > 0 && p.User != nil && p.APIKey != nil && p.APIKey.GroupID != nil {
|
||
deps.billingCacheService.QueueUpdateSubscriptionUsage(p.User.ID, *p.APIKey.GroupID, p.Cost.ActualCost)
|
||
}
|
||
} else if p.Cost.ActualCost > 0 && p.User != nil {
|
||
syncBalanceCacheAfterDeduction(ctx, p, deps, result)
|
||
}
|
||
|
||
if p.Cost.ActualCost > 0 && p.APIKey != nil && p.APIKey.HasRateLimits() {
|
||
deps.billingCacheService.QueueUpdateAPIKeyRateLimitUsage(p.APIKey.ID, p.Cost.ActualCost)
|
||
}
|
||
|
||
deps.deferredService.ScheduleLastUsedUpdate(p.Account.ID)
|
||
|
||
// Platform quota 累加：仅在 standard（余额）模式生效；订阅模式豁免；仅对有 limit 的用户写
|
||
// Redis 同步写 + DB 异步持久化（flag=false 降级）或 flusher 异步刷（flag=true）:
|
||
// - HasUserPlatformQuotaLimit 守卫:无 limit 的公司跳过,避免无效写入 + 浪费 Redis 容量
|
||
// - Redis 同步:确保下次 preflight 立即看到最新 usage,把 TOCTOU 超支窗口
|
||
// 限制在并发 in-flight 请求数量内（旧实现的异步入队会让超支无限累积直到 worker 处理）
|
||
// - DB 异步(flusher_enabled=false):在独立 goroutine 中走 detached context,失败用 ALERT log 触发 oncall 对账
|
||
// - flusher_enabled=true:不直写 DB,由 flusher 异步批量刷（markDirty 已在 IncrementUserPlatformQuotaUsage 内部完成）
|
||
if !p.IsSubscriptionBill && p.Platform != "" && p.Cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil {
|
||
if deps.billingCacheService.HasUserPlatformQuotaLimit(ctx, p.User.ID, p.Platform) {
|
||
deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, p.Cost.ActualCost)
|
||
if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled {
|
||
// 降级路径:flusher 未启用时保留原有异步直写 DB
|
||
dbCtx, dbCancel := detachUpstreamContext(ctx)
|
||
userID, platform, cost := p.User.ID, p.Platform, p.Cost.ActualCost
|
||
go func() {
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
logger.LegacyPrintf("service.gateway", "ALERT: panic in user platform quota incr goroutine user=%d platform=%s: %v", userID, platform, r)
|
||
}
|
||
}()
|
||
defer dbCancel()
|
||
if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(dbCtx, userID, platform, cost, time.Now().UTC()); err != nil {
|
||
// 失败计数器:暴露给 GatewayUserPlatformQuotaIncrStats(),由 ops 面板做斜率告警。
|
||
userPlatformQuotaDBIncrErrorTotal.Add(1)
|
||
// ALERT 级别:DB 持久化失败意味着 Redis cache 失效后该笔 cost 永久丢失,
|
||
// 用户配额视图与实际消费会偏差,oncall 需要据此对账或人工补录。
|
||
logger.LegacyPrintf("service.gateway", "ALERT: incr user platform quota DB failed user=%d platform=%s cost=%f: %v", userID, platform, cost, err)
|
||
}
|
||
}()
|
||
}
|
||
// flusher_enabled=true:不直写 DB,flusher 异步批量刷
|
||
}
|
||
}
|
||
|
||
// Notification checks run async — all parameters are already captured,
|
||
// no dependency on the request context or upstream connection.
|
||
go notifyBalanceLow(p, deps, result)
|
||
go notifyAccountQuota(p, deps, result)
|
||
}
|
||
|
||
func syncBalanceCacheAfterDeduction(ctx context.Context, p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) {
|
||
if p == nil || p.Cost == nil || p.User == nil || deps == nil || deps.billingCacheService == nil {
|
||
return
|
||
}
|
||
if result != nil && result.NewBalance != nil && deps.billingCacheService.balanceBelowEligibilityThreshold(*result.NewBalance) {
|
||
if err := deps.billingCacheService.InvalidateUserBalance(ctx, p.User.ID); err != nil {
|
||
slog.Warn("invalidate balance cache after exhausted deduction failed",
|
||
"user_id", p.User.ID,
|
||
"new_balance", *result.NewBalance,
|
||
"balance_overdrafted", result.BalanceOverdrafted,
|
||
"error", err,
|
||
)
|
||
}
|
||
return
|
||
}
|
||
deps.billingCacheService.QueueDeductBalance(p.User.ID, p.Cost.ActualCost)
|
||
}
|
||
|
||
// notifyBalanceLow sends balance low notification after deduction.
|
||
// When result.NewBalance is available (from DB transaction RETURNING), it is used directly
|
||
// to reconstruct oldBalance, avoiding stale Redis reads and concurrent-deduction races.
|
||
func notifyBalanceLow(p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) {
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
slog.Error("panic in notifyBalanceLow", "recover", r)
|
||
}
|
||
}()
|
||
if p.IsSubscriptionBill || p.Cost.ActualCost <= 0 || p.User == nil || deps.balanceNotifyService == nil {
|
||
slog.Debug("notifyBalanceLow: skipped",
|
||
"is_subscription", p.IsSubscriptionBill,
|
||
"actual_cost", p.Cost.ActualCost,
|
||
"user_nil", p.User == nil,
|
||
"service_nil", deps.balanceNotifyService == nil,
|
||
)
|
||
return
|
||
}
|
||
|
||
oldBalance := resolveOldBalance(p, result)
|
||
slog.Debug("notifyBalanceLow: calling CheckBalanceAfterDeduction",
|
||
"user_id", p.User.ID,
|
||
"old_balance", oldBalance,
|
||
"cost", p.Cost.ActualCost,
|
||
"notify_enabled", p.User.BalanceNotifyEnabled,
|
||
"threshold", p.User.BalanceNotifyThreshold,
|
||
"result_has_new_balance", result != nil && result.NewBalance != nil,
|
||
)
|
||
deps.balanceNotifyService.CheckBalanceAfterDeduction(context.Background(), p.User, oldBalance, p.Cost.ActualCost)
|
||
}
|
||
|
||
// resolveOldBalance returns the pre-deduction balance.
|
||
// Prefers the DB transaction result (newBalance + cost) over snapshot.
|
||
func resolveOldBalance(p *postUsageBillingParams, result *UsageBillingApplyResult) float64 {
|
||
if result != nil && result.NewBalance != nil {
|
||
return *result.NewBalance + p.Cost.ActualCost
|
||
}
|
||
// Legacy fallback: snapshot balance from request context
|
||
return p.User.Balance
|
||
}
|
||
|
||
// notifyAccountQuota sends account quota threshold notification after increment.
|
||
// When result.QuotaState is available (from DB transaction RETURNING), it is passed directly
|
||
// to avoid a separate DB read that may see stale or concurrently-modified data.
|
||
func notifyAccountQuota(p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) {
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
slog.Error("panic in notifyAccountQuota", "recover", r)
|
||
}
|
||
}()
|
||
if p.Cost.TotalCost <= 0 || p.Account == nil || !p.Account.IsAPIKeyOrBedrock() || deps.balanceNotifyService == nil {
|
||
slog.Debug("notifyAccountQuota: skipped",
|
||
"total_cost", p.Cost.TotalCost,
|
||
"account_nil", p.Account == nil,
|
||
"is_apikey_or_bedrock", p.Account != nil && p.Account.IsAPIKeyOrBedrock(),
|
||
"service_nil", deps.balanceNotifyService == nil,
|
||
)
|
||
return
|
||
}
|
||
accountCost := p.Cost.TotalCost * p.AccountRateMultiplier
|
||
var quotaState *AccountQuotaState
|
||
if result != nil {
|
||
quotaState = result.QuotaState
|
||
}
|
||
slog.Debug("notifyAccountQuota: calling CheckAccountQuotaAfterIncrement",
|
||
"account_id", p.Account.ID,
|
||
"account_cost", accountCost,
|
||
"has_quota_state", quotaState != nil,
|
||
)
|
||
deps.balanceNotifyService.CheckAccountQuotaAfterIncrement(context.Background(), p.Account, accountCost, quotaState)
|
||
}
|
||
|
||
func detachedBillingContext(ctx context.Context) (context.Context, context.CancelFunc) {
|
||
base := context.Background()
|
||
if ctx != nil {
|
||
base = context.WithoutCancel(ctx)
|
||
}
|
||
return context.WithTimeout(base, postUsageBillingTimeout)
|
||
}
|
||
|
||
func detachStreamUpstreamContext(ctx context.Context, stream bool) (context.Context, context.CancelFunc) {
|
||
if ctx == nil {
|
||
return context.Background(), func() {}
|
||
}
|
||
if !stream {
|
||
return ctx, func() {}
|
||
}
|
||
return context.WithoutCancel(ctx), func() {}
|
||
}
|
||
|
||
func detachUpstreamContext(ctx context.Context) (context.Context, context.CancelFunc) {
|
||
if ctx == nil {
|
||
return context.Background(), func() {}
|
||
}
|
||
return context.WithoutCancel(ctx), func() {}
|
||
}
|
||
|
||
// billingDeps 扣费逻辑依赖的服务（由各 gateway service 提供）
|
||
type billingDeps struct {
|
||
accountRepo AccountRepository
|
||
userRepo UserRepository
|
||
userSubRepo UserSubscriptionRepository
|
||
billingCacheService *BillingCacheService
|
||
deferredService *DeferredService
|
||
balanceNotifyService *BalanceNotifyService
|
||
userPlatformQuotaRepo UserPlatformQuotaRepository
|
||
cfg *config.Config
|
||
}
|
||
|
||
func (s *GatewayService) billingDeps() *billingDeps {
|
||
return &billingDeps{
|
||
accountRepo: s.accountRepo,
|
||
userRepo: s.userRepo,
|
||
userSubRepo: s.userSubRepo,
|
||
billingCacheService: s.billingCacheService,
|
||
deferredService: s.deferredService,
|
||
balanceNotifyService: s.balanceNotifyService,
|
||
userPlatformQuotaRepo: s.userPlatformQuotaRepo,
|
||
cfg: s.cfg,
|
||
}
|
||
}
|
||
|
||
func writeUsageLogBestEffort(ctx context.Context, repo UsageLogRepository, usageLog *UsageLog, logKey string) {
|
||
if repo == nil || usageLog == nil {
|
||
return
|
||
}
|
||
usageCtx, cancel := detachedBillingContext(ctx)
|
||
defer cancel()
|
||
|
||
if writer, ok := repo.(usageLogBestEffortWriter); ok {
|
||
if err := writer.CreateBestEffort(usageCtx, usageLog); err != nil {
|
||
logger.LegacyPrintf(logKey, "Create usage log failed: %v", err)
|
||
// 计费已在此前完成，日志必须落库：dropped（批处理队列超时）同样走同步兜底，
|
||
// 否则会出现“已扣费但无 usage_log”的对账缺口（issue #3656）。
|
||
// 重复写入由 usage_logs 的 ON CONFLICT (request_id, api_key_id) DO NOTHING 防护。
|
||
fallbackCtx := usageCtx
|
||
if usageCtx.Err() != nil {
|
||
// usageCtx 已耗尽（best-effort 入队阻塞到期限）：换新的 detached 窗口，避免兜底必然失败。
|
||
var fallbackCancel context.CancelFunc
|
||
fallbackCtx, fallbackCancel = detachedBillingContext(context.Background())
|
||
defer fallbackCancel()
|
||
}
|
||
if _, syncErr := repo.Create(fallbackCtx, usageLog); syncErr != nil {
|
||
logger.LegacyPrintf(logKey, "Create usage log sync fallback failed: %v", syncErr)
|
||
}
|
||
}
|
||
return
|
||
}
|
||
|
||
if _, err := repo.Create(usageCtx, usageLog); err != nil {
|
||
logger.LegacyPrintf(logKey, "Create usage log failed: %v", err)
|
||
}
|
||
}
|
||
|
||
// recordUsageOpts 内部选项，参数化普通计费与长上下文计费的差异点。
|
||
type recordUsageOpts struct {
|
||
// 长上下文计费（仅 Gemini 路径需要）
|
||
LongContextThreshold int
|
||
LongContextMultiplier float64
|
||
}
|
||
|
||
// RecordUsage 记录使用量并扣费（或更新订阅用量）
|
||
func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInput) error {
|
||
return s.recordUsageCore(ctx, &recordUsageCoreInput{
|
||
Result: input.Result,
|
||
APIKey: input.APIKey,
|
||
User: input.User,
|
||
Account: input.Account,
|
||
Subscription: input.Subscription,
|
||
InboundEndpoint: input.InboundEndpoint,
|
||
UpstreamEndpoint: input.UpstreamEndpoint,
|
||
UserAgent: input.UserAgent,
|
||
IPAddress: input.IPAddress,
|
||
SessionID: input.SessionID,
|
||
RequestPayloadHash: input.RequestPayloadHash,
|
||
ForceCacheBilling: input.ForceCacheBilling,
|
||
APIKeyService: input.APIKeyService,
|
||
QuotaPlatform: input.QuotaPlatform,
|
||
ChannelUsageFields: input.ChannelUsageFields,
|
||
}, &recordUsageOpts{})
|
||
}
|
||
|
||
// RecordUsageLongContextInput 记录使用量的输入参数（支持长上下文双倍计费）
|
||
type RecordUsageLongContextInput struct {
|
||
Result *ForwardResult
|
||
APIKey *APIKey
|
||
User *User
|
||
Account *Account
|
||
Subscription *UserSubscription // 可选：订阅信息
|
||
InboundEndpoint string // 入站端点（客户端请求路径）
|
||
UpstreamEndpoint string // 上游端点（标准化后的上游路径）
|
||
UserAgent string // 请求的 User-Agent
|
||
IPAddress string // 请求的客户端 IP 地址
|
||
SessionID string // 客户端显式会话标识（session_id / X-Session-Id 等请求头），仅用于用量行会话关联
|
||
RequestPayloadHash string // 请求体语义哈希，用于降低 request_id 误复用时的静默误去重风险
|
||
LongContextThreshold int // 长上下文阈值（如 200000）
|
||
LongContextMultiplier float64 // 超出阈值部分的倍率（如 2.0）
|
||
ForceCacheBilling bool // 强制缓存计费：将 input_tokens 转为 cache_read 计费（用于粘性会话切换）
|
||
APIKeyService APIKeyQuotaUpdater // API Key 配额服务（可选）
|
||
QuotaPlatform string // user×platform 配额计量平台：handler 在请求 ctx 内经 QuotaPlatform() 算定后传入（后扣运行在 worker 池 background ctx 上，取不到 ForcePlatform）
|
||
|
||
ChannelUsageFields // 渠道映射信息（由 handler 在 Forward 前解析）
|
||
}
|
||
|
||
// RecordUsageWithLongContext 记录使用量并扣费，支持长上下文双倍计费（用于 Gemini）
|
||
func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input *RecordUsageLongContextInput) error {
|
||
return s.recordUsageCore(ctx, &recordUsageCoreInput{
|
||
Result: input.Result,
|
||
APIKey: input.APIKey,
|
||
User: input.User,
|
||
Account: input.Account,
|
||
Subscription: input.Subscription,
|
||
InboundEndpoint: input.InboundEndpoint,
|
||
UpstreamEndpoint: input.UpstreamEndpoint,
|
||
UserAgent: input.UserAgent,
|
||
IPAddress: input.IPAddress,
|
||
SessionID: input.SessionID,
|
||
RequestPayloadHash: input.RequestPayloadHash,
|
||
ForceCacheBilling: input.ForceCacheBilling,
|
||
APIKeyService: input.APIKeyService,
|
||
QuotaPlatform: input.QuotaPlatform,
|
||
ChannelUsageFields: input.ChannelUsageFields,
|
||
}, &recordUsageOpts{
|
||
LongContextThreshold: input.LongContextThreshold,
|
||
LongContextMultiplier: input.LongContextMultiplier,
|
||
})
|
||
}
|
||
|
||
// recordUsageCoreInput 是 recordUsageCore 的公共输入字段，从两种输入结构体中提取。
|
||
type recordUsageCoreInput struct {
|
||
Result *ForwardResult
|
||
APIKey *APIKey
|
||
User *User
|
||
Account *Account
|
||
Subscription *UserSubscription
|
||
InboundEndpoint string
|
||
UpstreamEndpoint string
|
||
UserAgent string
|
||
IPAddress string
|
||
SessionID string
|
||
RequestPayloadHash string
|
||
ForceCacheBilling bool
|
||
APIKeyService APIKeyQuotaUpdater
|
||
QuotaPlatform string
|
||
ChannelUsageFields
|
||
}
|
||
|
||
// recordUsageCore 是 RecordUsage 和 RecordUsageWithLongContext 的统一实现。
|
||
// LongContextThreshold > 0 时 Token 计费回退走 CalculateCostWithLongContext。
|
||
func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsageCoreInput, opts *recordUsageOpts) error {
|
||
result := input.Result
|
||
apiKey := input.APIKey
|
||
user := input.User
|
||
account := input.Account
|
||
subscription := input.Subscription
|
||
ApplyForwardImageBillingResolution(result)
|
||
|
||
// 强制缓存计费：将 input_tokens 转为 cache_read_input_tokens
|
||
// 用于粘性会话切换时的特殊计费处理
|
||
if input.ForceCacheBilling && result.Usage.InputTokens > 0 {
|
||
logger.LegacyPrintf("service.gateway", "force_cache_billing: %d input_tokens → cache_read_input_tokens (account=%d)",
|
||
result.Usage.InputTokens, account.ID)
|
||
result.Usage.CacheReadInputTokens += result.Usage.InputTokens
|
||
result.Usage.InputTokens = 0
|
||
}
|
||
|
||
// Cache TTL Override: 确保计费时 token 分类与账号设置一致。
|
||
// 账号级设置优先；全局 1h 请求注入开启时，默认把 usage 计费归回 5m。
|
||
cacheTTLOverridden := false
|
||
if overrideTarget, ok := s.resolveCacheTTLUsageOverrideTarget(ctx, account); ok {
|
||
applyCacheTTLOverride(&result.Usage, overrideTarget)
|
||
cacheTTLOverridden = (result.Usage.CacheCreation5mTokens + result.Usage.CacheCreation1hTokens) > 0
|
||
}
|
||
|
||
// 获取费率倍数（优先级：用户专属 > 分组默认 > 系统默认）
|
||
multiplier := 1.0
|
||
if s.cfg != nil {
|
||
multiplier = s.cfg.Default.RateMultiplier
|
||
}
|
||
if apiKey.GroupID != nil && apiKey.Group != nil {
|
||
groupDefault := apiKey.Group.RateMultiplier
|
||
multiplier = s.ResolveUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, groupDefault)
|
||
}
|
||
// token 倍率叠加高峰因子（token 计费含图片 token，图片按次倍率不受影响）。高峰因子按请求时刻现算，
|
||
// 不并入上面的 getUserGroupRateMultiplier，以免污染 user:group 倍率缓存。
|
||
multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, timezone.Now())
|
||
|
||
// 确定计费模型
|
||
concreteBillingModel := forwardResultBillingModel(result.Model, result.UpstreamModel)
|
||
billingModel := concreteBillingModel
|
||
if input.BillingModelSource == BillingModelSourceChannelMapped && input.ChannelMappedModel != "" {
|
||
billingModel = input.ChannelMappedModel
|
||
}
|
||
if input.BillingModelSource == BillingModelSourceRequested && input.OriginalModel != "" {
|
||
billingModel = input.OriginalModel
|
||
}
|
||
// composite 分组的公开别名（如 all/claude）会经 OriginalModel/ChannelMappedModel
|
||
// 进入上面的来源覆盖：任意别名查无价会静默落 $0，含家族词的别名则被价格表的
|
||
// 家族模糊匹配错计（如 Opus 流量按 Sonnet 兜底价）。除非管理员为别名显式配置了
|
||
// 渠道定价（OpenRouter 式自定价），composite 请求一律按实际转发的具体模型计费。
|
||
if apiKey.Group != nil && apiKey.Group.Platform == PlatformComposite {
|
||
billingModel = s.compositeBillableModel(ctx, apiKey, billingModel, concreteBillingModel)
|
||
}
|
||
// 通用兜底（与 OpenAI 路径的 usageBillingModelCandidates 语义对齐）：
|
||
// 选定模型查不到任何价格时回退到实际转发的具体模型。已定价流量不受影响。
|
||
billingModel = s.billableModelWithFallback(ctx, apiKey, billingModel, result.UpstreamModel, result.Model)
|
||
|
||
// 确定 RequestedModel（渠道映射前的原始模型）
|
||
requestedModel := result.Model
|
||
if input.OriginalModel != "" {
|
||
requestedModel = input.OriginalModel
|
||
}
|
||
|
||
// 计算费用
|
||
cost := s.calculateRecordUsageCost(ctx, result, apiKey, billingModel, multiplier, imageMultiplier, opts)
|
||
|
||
// 判断计费方式：订阅模式 vs 余额模式
|
||
isSubscriptionBilling := subscription != nil && apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
|
||
billingType := BillingTypeBalance
|
||
if isSubscriptionBilling {
|
||
billingType = BillingTypeSubscription
|
||
}
|
||
|
||
// 创建使用日志
|
||
accountRateMultiplier := account.BillingRateMultiplier()
|
||
usageLog := s.buildRecordUsageLog(ctx, input, result, apiKey, user, account, subscription,
|
||
requestedModel, multiplier, imageMultiplier, accountRateMultiplier, billingType, cacheTTLOverridden, cost, opts)
|
||
|
||
// 计算账号统计定价费用（使用最终上游模型匹配自定义规则）
|
||
if apiKey.GroupID != nil {
|
||
applyAccountStatsCost(ctx, usageLog, s.channelService, s.billingService,
|
||
account.ID, *apiKey.GroupID, result.UpstreamModel, result.Model,
|
||
// Anthropic's input_tokens excludes cache_read and cache_creation (billed separately);
|
||
// OpenAI gateway uses actualInputTokens which also excludes cache_read for the same reason.
|
||
UsageTokens{
|
||
InputTokens: result.Usage.InputTokens,
|
||
OutputTokens: result.Usage.OutputTokens,
|
||
CacheCreationTokens: result.Usage.CacheCreationInputTokens,
|
||
CacheReadTokens: result.Usage.CacheReadInputTokens,
|
||
ImageOutputTokens: result.Usage.ImageOutputTokens,
|
||
},
|
||
cost.TotalCost,
|
||
)
|
||
}
|
||
|
||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
||
writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway")
|
||
logger.LegacyPrintf("service.gateway", "[SIMPLE MODE] Usage recorded (not billed): user=%d, tokens=%d", usageLog.UserID, usageLog.TotalTokens())
|
||
s.deferredService.ScheduleLastUsedUpdate(account.ID)
|
||
return nil
|
||
}
|
||
|
||
// 配额平台由 handler 在请求 ctx 内经 QuotaPlatform() 算定并通过 input 传入；
|
||
// 后扣运行在 worker 池的 background ctx 上，无法再从 ctx 取 ForcePlatform。
|
||
// 缺省（未设置）时回退到分组平台，保持对其它调用方的兼容。
|
||
quotaPlatform := input.QuotaPlatform
|
||
if quotaPlatform == "" {
|
||
quotaPlatform = PlatformFromAPIKey(apiKey)
|
||
if quotaPlatform == PlatformComposite && account != nil {
|
||
quotaPlatform = account.Platform
|
||
}
|
||
}
|
||
requestID := usageLog.RequestID
|
||
_, billingErr := applyUsageBilling(ctx, requestID, usageLog, &postUsageBillingParams{
|
||
Cost: cost,
|
||
User: user,
|
||
APIKey: apiKey,
|
||
Account: account,
|
||
Subscription: subscription,
|
||
RequestPayloadHash: resolveUsageBillingPayloadFingerprint(ctx, input.RequestPayloadHash),
|
||
IsSubscriptionBill: isSubscriptionBilling,
|
||
AccountRateMultiplier: accountRateMultiplier,
|
||
APIKeyService: input.APIKeyService,
|
||
Platform: quotaPlatform,
|
||
}, s.billingDeps(), s.usageBillingRepo)
|
||
|
||
if billingErr != nil {
|
||
return billingErr
|
||
}
|
||
writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway")
|
||
|
||
return nil
|
||
}
|
||
|
||
// calculateRecordUsageCost 根据请求类型和选项计算费用。
|
||
func (s *GatewayService) calculateRecordUsageCost(
|
||
ctx context.Context,
|
||
result *ForwardResult,
|
||
apiKey *APIKey,
|
||
billingModel string,
|
||
multiplier float64,
|
||
imageMultiplier float64,
|
||
opts *recordUsageOpts,
|
||
) *CostBreakdown {
|
||
// 图片生成：渠道定价为 token 计费时走 token 路径，否则走图片计费
|
||
if result.ImageCount > 0 {
|
||
if resolved := s.resolveChannelPricing(ctx, billingModel, apiKey); resolved != nil && resolved.Mode == BillingModeToken {
|
||
return s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts)
|
||
}
|
||
return s.calculateImageCost(ctx, result, apiKey, billingModel, imageMultiplier)
|
||
}
|
||
|
||
// Token 计费
|
||
return s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts)
|
||
}
|
||
|
||
// compositeBillableModel 决定 composite 分组请求的计费模型：来源覆盖把计费模型
|
||
// 换成公开别名等非具体模型时，只有管理员为该名字显式配置了渠道定价才按其计费
|
||
// （OpenRouter 式自定价），否则回退到实际转发的具体模型，避免别名落入价格表的
|
||
// 家族模糊匹配（错价）或查无价（$0）。未发生来源覆盖时原样返回。
|
||
func (s *GatewayService) compositeBillableModel(ctx context.Context, apiKey *APIKey, billingModel, concreteBillingModel string) string {
|
||
if concreteBillingModel == "" || billingModel == concreteBillingModel {
|
||
return billingModel
|
||
}
|
||
if s.resolveChannelPricing(ctx, billingModel, apiKey) != nil {
|
||
return billingModel
|
||
}
|
||
logger.LegacyPrintf("service.gateway", "[Billing] composite billing model %q has no explicit channel pricing, billing by concrete model %q", billingModel, concreteBillingModel)
|
||
return concreteBillingModel
|
||
}
|
||
|
||
// billableModelWithFallback 在选定计费模型（可能是 composite 公开别名或未定价的映射名）
|
||
// 查不到任何价格（渠道价与全局价均无）时，按序回退到实际转发的具体模型，避免静默 $0 计费。
|
||
// 所有候选都无价时保持原值，走既有的 warn + 零成本路径。
|
||
func (s *GatewayService) billableModelWithFallback(ctx context.Context, apiKey *APIKey, billingModel string, fallbacks ...string) string {
|
||
if s.hasResolvableTokenPricing(ctx, billingModel, apiKey) {
|
||
return billingModel
|
||
}
|
||
for _, fallback := range fallbacks {
|
||
fallback = strings.TrimSpace(fallback)
|
||
if fallback == "" || fallback == billingModel {
|
||
continue
|
||
}
|
||
if s.hasResolvableTokenPricing(ctx, fallback, apiKey) {
|
||
logger.LegacyPrintf("service.gateway", "[Billing] billing model %q has no pricing, falling back to concrete model %q", billingModel, fallback)
|
||
return fallback
|
||
}
|
||
}
|
||
return billingModel
|
||
}
|
||
|
||
// hasResolvableTokenPricing 判断模型是否能在渠道定价或全局价格表中解析出 token 价格。
|
||
func (s *GatewayService) hasResolvableTokenPricing(ctx context.Context, model string, apiKey *APIKey) bool {
|
||
if strings.TrimSpace(model) == "" {
|
||
return false
|
||
}
|
||
if s.resolveChannelPricing(ctx, model, apiKey) != nil {
|
||
return true
|
||
}
|
||
if s.billingService == nil {
|
||
return false
|
||
}
|
||
_, err := s.billingService.GetModelPricing(model)
|
||
return err == nil
|
||
}
|
||
|
||
// resolveChannelPricing 检查指定模型是否存在渠道级别定价。
|
||
// 返回非 nil 的 ResolvedPricing 表示有渠道定价，nil 表示走默认定价路径。
|
||
func (s *GatewayService) resolveChannelPricing(ctx context.Context, billingModel string, apiKey *APIKey) *ResolvedPricing {
|
||
if s.resolver == nil || apiKey.Group == nil {
|
||
return nil
|
||
}
|
||
gid := apiKey.Group.ID
|
||
resolved := s.resolver.Resolve(ctx, PricingInput{Model: billingModel, GroupID: &gid})
|
||
if resolved.Source == PricingSourceChannel {
|
||
return resolved
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// calculateImageCost 计算图片生成费用：渠道级别定价优先，否则走按次计费。
|
||
func (s *GatewayService) calculateImageCost(
|
||
ctx context.Context,
|
||
result *ForwardResult,
|
||
apiKey *APIKey,
|
||
billingModel string,
|
||
multiplier float64,
|
||
) *CostBreakdown {
|
||
sizeTier := NormalizeImageBillingTierOrDefault(result.ImageSize)
|
||
groupConfig := imagePriceConfigFromAPIKey(apiKey)
|
||
if apiKeyHasConfiguredImagePrice(apiKey, sizeTier) {
|
||
return s.billingService.CalculateImageCost(billingModel, sizeTier, result.ImageCount, groupConfig, multiplier)
|
||
}
|
||
if resolved := s.resolveChannelPricing(ctx, billingModel, apiKey); resolved != nil {
|
||
tokens := UsageTokens{
|
||
InputTokens: result.Usage.InputTokens,
|
||
OutputTokens: result.Usage.OutputTokens,
|
||
ImageOutputTokens: result.Usage.ImageOutputTokens,
|
||
}
|
||
gid := apiKey.Group.ID
|
||
cost, err := s.billingService.CalculateCostUnified(CostInput{
|
||
Ctx: ctx,
|
||
Model: billingModel,
|
||
GroupID: &gid,
|
||
Tokens: tokens,
|
||
RequestCount: result.ImageCount,
|
||
SizeTier: sizeTier,
|
||
RateMultiplier: multiplier,
|
||
Resolver: s.resolver,
|
||
Resolved: resolved,
|
||
})
|
||
if err != nil {
|
||
logger.LegacyPrintf("service.gateway", "Calculate image token cost failed: %v", err)
|
||
return &CostBreakdown{ActualCost: 0}
|
||
}
|
||
return cost
|
||
}
|
||
|
||
return s.billingService.CalculateImageCost(billingModel, sizeTier, result.ImageCount, groupConfig, multiplier)
|
||
}
|
||
|
||
// calculateTokenCost 计算 Token 计费：根据 opts 决定走普通/长上下文/渠道统一计费。
|
||
func (s *GatewayService) calculateTokenCost(
|
||
ctx context.Context,
|
||
result *ForwardResult,
|
||
apiKey *APIKey,
|
||
billingModel string,
|
||
multiplier float64,
|
||
opts *recordUsageOpts,
|
||
) *CostBreakdown {
|
||
tokens := UsageTokens{
|
||
InputTokens: result.Usage.InputTokens,
|
||
OutputTokens: result.Usage.OutputTokens,
|
||
CacheCreationTokens: result.Usage.CacheCreationInputTokens,
|
||
CacheReadTokens: result.Usage.CacheReadInputTokens,
|
||
CacheCreation5mTokens: result.Usage.CacheCreation5mTokens,
|
||
CacheCreation1hTokens: result.Usage.CacheCreation1hTokens,
|
||
ImageOutputTokens: result.Usage.ImageOutputTokens,
|
||
}
|
||
|
||
var cost *CostBreakdown
|
||
var err error
|
||
|
||
// 优先尝试渠道定价 → CalculateCostUnified
|
||
if resolved := s.resolveChannelPricing(ctx, billingModel, apiKey); resolved != nil {
|
||
gid := apiKey.Group.ID
|
||
cost, err = s.billingService.CalculateCostUnified(CostInput{
|
||
Ctx: ctx,
|
||
Model: billingModel,
|
||
GroupID: &gid,
|
||
Tokens: tokens,
|
||
RequestCount: 1,
|
||
RateMultiplier: multiplier,
|
||
Resolver: s.resolver,
|
||
Resolved: resolved,
|
||
})
|
||
} else if opts.LongContextThreshold > 0 {
|
||
// 长上下文双倍计费（如 Gemini 200K 阈值）
|
||
cost, err = s.billingService.CalculateCostWithLongContext(billingModel, tokens, multiplier, opts.LongContextThreshold, opts.LongContextMultiplier)
|
||
} else {
|
||
cost, err = s.billingService.CalculateCost(billingModel, tokens, multiplier)
|
||
}
|
||
if err != nil {
|
||
logger.LegacyPrintf("service.gateway", "Calculate cost failed: %v", err)
|
||
return &CostBreakdown{ActualCost: 0}
|
||
}
|
||
return cost
|
||
}
|
||
|
||
// buildRecordUsageLog 构建使用日志并设置计费模式。
|
||
func (s *GatewayService) buildRecordUsageLog(
|
||
ctx context.Context,
|
||
input *recordUsageCoreInput,
|
||
result *ForwardResult,
|
||
apiKey *APIKey,
|
||
user *User,
|
||
account *Account,
|
||
subscription *UserSubscription,
|
||
requestedModel string,
|
||
multiplier float64,
|
||
imageMultiplier float64,
|
||
accountRateMultiplier float64,
|
||
billingType int8,
|
||
cacheTTLOverridden bool,
|
||
cost *CostBreakdown,
|
||
opts *recordUsageOpts,
|
||
) *UsageLog {
|
||
durationMs := int(result.Duration.Milliseconds())
|
||
requestID := resolveUsageBillingRequestID(ctx, result.RequestID)
|
||
usageLog := &UsageLog{
|
||
UserID: user.ID,
|
||
APIKeyID: apiKey.ID,
|
||
AccountID: account.ID,
|
||
RequestID: requestID,
|
||
Model: result.Model,
|
||
RequestedModel: requestedModel,
|
||
UpstreamModel: optionalNonEqualStringPtr(result.UpstreamModel, result.Model),
|
||
ReasoningEffort: result.ReasoningEffort,
|
||
InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint),
|
||
UpstreamEndpoint: optionalTrimmedStringPtr(input.UpstreamEndpoint),
|
||
InputTokens: result.Usage.InputTokens,
|
||
OutputTokens: result.Usage.OutputTokens,
|
||
CacheCreationTokens: result.Usage.CacheCreationInputTokens,
|
||
CacheReadTokens: result.Usage.CacheReadInputTokens,
|
||
CacheCreation5mTokens: result.Usage.CacheCreation5mTokens,
|
||
CacheCreation1hTokens: result.Usage.CacheCreation1hTokens,
|
||
ImageOutputTokens: result.Usage.ImageOutputTokens,
|
||
RateMultiplier: multiplier,
|
||
AccountRateMultiplier: &accountRateMultiplier,
|
||
BillingType: billingType,
|
||
BillingMode: resolveBillingMode(result, cost),
|
||
Stream: result.Stream,
|
||
DurationMs: &durationMs,
|
||
FirstTokenMs: result.FirstTokenMs,
|
||
ImageCount: result.ImageCount,
|
||
ImageSize: optionalTrimmedStringPtr(result.ImageSize),
|
||
ImageInputSize: optionalTrimmedStringPtr(result.ImageInputSize),
|
||
ImageOutputSize: optionalTrimmedStringPtr(result.ImageOutputSize),
|
||
ImageSizeSource: optionalTrimmedStringPtr(result.ImageSizeSource),
|
||
ImageSizeBreakdown: result.ImageSizeBreakdown,
|
||
CacheTTLOverridden: cacheTTLOverridden,
|
||
ChannelID: optionalInt64Ptr(input.ChannelID),
|
||
ModelMappingChain: optionalTrimmedStringPtr(input.ModelMappingChain),
|
||
UserAgent: optionalTrimmedStringPtr(input.UserAgent),
|
||
IPAddress: optionalTrimmedStringPtr(input.IPAddress),
|
||
SessionID: optionalTrimmedStringPtr(input.SessionID),
|
||
GroupID: apiKey.GroupID,
|
||
SubscriptionID: optionalSubscriptionID(subscription),
|
||
CreatedAt: time.Now(),
|
||
}
|
||
if result.ImageCount > 0 && (cost == nil || cost.BillingMode != string(BillingModeToken)) {
|
||
usageLog.RateMultiplier = imageMultiplier
|
||
}
|
||
if cost != nil {
|
||
usageLog.InputCost = cost.InputCost
|
||
usageLog.OutputCost = cost.OutputCost
|
||
usageLog.ImageOutputCost = cost.ImageOutputCost
|
||
usageLog.CacheCreationCost = cost.CacheCreationCost
|
||
usageLog.CacheReadCost = cost.CacheReadCost
|
||
usageLog.TotalCost = cost.TotalCost
|
||
usageLog.ActualCost = cost.ActualCost
|
||
usageLog.LongContextBillingApplied = cost.LongContextBillingApplied
|
||
}
|
||
|
||
return usageLog
|
||
}
|
||
|
||
// resolveBillingMode 根据计费结果和请求类型确定计费模式。
|
||
func resolveBillingMode(result *ForwardResult, cost *CostBreakdown) *string {
|
||
var mode string
|
||
switch {
|
||
case cost != nil && cost.BillingMode != "":
|
||
mode = cost.BillingMode
|
||
case result.ImageCount > 0:
|
||
mode = string(BillingModeImage)
|
||
default:
|
||
mode = string(BillingModeToken)
|
||
}
|
||
return &mode
|
||
}
|
||
|
||
func optionalSubscriptionID(subscription *UserSubscription) *int64 {
|
||
if subscription != nil {
|
||
return &subscription.ID
|
||
}
|
||
return nil
|
||
}
|