Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
 
 
 

391 řádky
14 KiB

  1. package service
  2. import (
  3. "fmt"
  4. "net/http"
  5. "strings"
  6. "sync"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/constant"
  9. "github.com/QuantumNous/new-api/logger"
  10. "github.com/QuantumNous/new-api/model"
  11. relaycommon "github.com/QuantumNous/new-api/relay/common"
  12. "github.com/QuantumNous/new-api/types"
  13. "github.com/bytedance/gopkg/util/gopool"
  14. "github.com/gin-gonic/gin"
  15. )
  16. // ---------------------------------------------------------------------------
  17. // BillingSession — 统一计费会话
  18. // ---------------------------------------------------------------------------
  19. // BillingSession 封装单次请求的预扣费/结算/退款生命周期。
  20. // 实现 relaycommon.BillingSettler 接口。
  21. type BillingSession struct {
  22. relayInfo *relaycommon.RelayInfo
  23. funding FundingSource
  24. preConsumedQuota int // 实际预扣额度(信任用户可能为 0)
  25. tokenConsumed int // 令牌额度实际扣减量
  26. fundingSettled bool // funding.Settle 已成功,资金来源已提交
  27. settled bool // Settle 全部完成(资金 + 令牌)
  28. refunded bool // Refund 已调用
  29. mu sync.Mutex
  30. }
  31. // Settle 根据实际消耗额度进行结算。
  32. // 资金来源和令牌额度分两步提交:若资金来源已提交但令牌调整失败,
  33. // 会标记 fundingSettled 防止 Refund 对已提交的资金来源执行退款。
  34. func (s *BillingSession) Settle(actualQuota int) error {
  35. s.mu.Lock()
  36. defer s.mu.Unlock()
  37. if s.settled {
  38. return nil
  39. }
  40. delta := actualQuota - s.preConsumedQuota
  41. if delta == 0 {
  42. s.settled = true
  43. return nil
  44. }
  45. // 1) 调整资金来源(仅在尚未提交时执行,防止重复调用)
  46. if !s.fundingSettled {
  47. if err := s.funding.Settle(delta); err != nil {
  48. return err
  49. }
  50. s.fundingSettled = true
  51. }
  52. // 2) 调整令牌额度
  53. var tokenErr error
  54. if !s.relayInfo.IsPlayground {
  55. if delta > 0 {
  56. tokenErr = model.DecreaseTokenQuota(s.relayInfo.TokenId, s.relayInfo.TokenKey, delta)
  57. } else {
  58. tokenErr = model.IncreaseTokenQuota(s.relayInfo.TokenId, s.relayInfo.TokenKey, -delta)
  59. }
  60. if tokenErr != nil {
  61. // 资金来源已提交,令牌调整失败只能记录日志;标记 settled 防止 Refund 误退资金
  62. common.SysLog(fmt.Sprintf("error adjusting token quota after funding settled (userId=%d, tokenId=%d, delta=%d): %s",
  63. s.relayInfo.UserId, s.relayInfo.TokenId, delta, tokenErr.Error()))
  64. }
  65. }
  66. // 3) 更新 relayInfo 上的订阅 PostDelta(用于日志)
  67. if s.funding.Source() == BillingSourceSubscription {
  68. s.relayInfo.SubscriptionPostDelta += int64(delta)
  69. }
  70. s.settled = true
  71. return tokenErr
  72. }
  73. // Refund 退还所有预扣费,幂等安全,异步执行。
  74. func (s *BillingSession) Refund(c *gin.Context) {
  75. s.mu.Lock()
  76. if s.settled || s.refunded || !s.needsRefundLocked() {
  77. s.mu.Unlock()
  78. return
  79. }
  80. s.refunded = true
  81. s.mu.Unlock()
  82. logger.LogInfo(c, fmt.Sprintf("用户 %d 请求失败, 返还预扣费(token_quota=%s, funding=%s)",
  83. s.relayInfo.UserId,
  84. logger.FormatQuota(s.tokenConsumed),
  85. s.funding.Source(),
  86. ))
  87. // 复制需要的值到闭包中
  88. tokenId := s.relayInfo.TokenId
  89. tokenKey := s.relayInfo.TokenKey
  90. isPlayground := s.relayInfo.IsPlayground
  91. tokenConsumed := s.tokenConsumed
  92. funding := s.funding
  93. gopool.Go(func() {
  94. // 1) 退还资金来源
  95. if err := funding.Refund(); err != nil {
  96. common.SysLog("error refunding billing source: " + err.Error())
  97. }
  98. // 2) 退还令牌额度
  99. if tokenConsumed > 0 && !isPlayground {
  100. if err := model.IncreaseTokenQuota(tokenId, tokenKey, tokenConsumed); err != nil {
  101. common.SysLog("error refunding token quota: " + err.Error())
  102. }
  103. }
  104. })
  105. }
  106. // NeedsRefund 返回是否存在需要退还的预扣状态。
  107. func (s *BillingSession) NeedsRefund() bool {
  108. s.mu.Lock()
  109. defer s.mu.Unlock()
  110. return s.needsRefundLocked()
  111. }
  112. func (s *BillingSession) needsRefundLocked() bool {
  113. if s.settled || s.refunded || s.fundingSettled {
  114. // fundingSettled 时资金来源已提交结算,不能再退预扣费
  115. return false
  116. }
  117. if s.tokenConsumed > 0 {
  118. return true
  119. }
  120. // 订阅可能在 tokenConsumed=0 时仍预扣了额度
  121. if sub, ok := s.funding.(*SubscriptionFunding); ok && sub.preConsumed > 0 {
  122. return true
  123. }
  124. // 同步用户可能在 tokenConsumed=0 时仍预扣了额度
  125. if synced, ok := s.funding.(*SyncedUserFunding); ok && synced.consumed > 0 {
  126. return true
  127. }
  128. return false
  129. }
  130. // GetPreConsumedQuota 返回实际预扣的额度。
  131. func (s *BillingSession) GetPreConsumedQuota() int {
  132. return s.preConsumedQuota
  133. }
  134. // ---------------------------------------------------------------------------
  135. // PreConsume — 统一预扣费入口(含信任额度旁路)
  136. // ---------------------------------------------------------------------------
  137. // preConsume 执行预扣费:信任检查 -> 令牌预扣 -> 资金来源预扣。
  138. // 任一步骤失败时原子回滚已完成的步骤。
  139. func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.NewAPIError {
  140. effectiveQuota := quota
  141. // ---- 信任额度旁路 ----
  142. if s.shouldTrust(c) {
  143. effectiveQuota = 0
  144. logger.LogInfo(c, fmt.Sprintf("用户 %d 额度充足, 信任且不需要预扣费 (funding=%s)", s.relayInfo.UserId, s.funding.Source()))
  145. } else if effectiveQuota > 0 {
  146. logger.LogInfo(c, fmt.Sprintf("用户 %d 需要预扣费 %s (funding=%s)", s.relayInfo.UserId, logger.FormatQuota(effectiveQuota), s.funding.Source()))
  147. }
  148. // ---- 1) 预扣令牌额度 ----
  149. if effectiveQuota > 0 {
  150. if err := PreConsumeTokenQuota(s.relayInfo, effectiveQuota); err != nil {
  151. return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
  152. }
  153. s.tokenConsumed = effectiveQuota
  154. }
  155. // ---- 2) 预扣资金来源 ----
  156. if err := s.funding.PreConsume(effectiveQuota); err != nil {
  157. // 预扣费失败,回滚令牌额度
  158. if s.tokenConsumed > 0 && !s.relayInfo.IsPlayground {
  159. if rollbackErr := model.IncreaseTokenQuota(s.relayInfo.TokenId, s.relayInfo.TokenKey, s.tokenConsumed); rollbackErr != nil {
  160. common.SysLog(fmt.Sprintf("error rolling back token quota (userId=%d, tokenId=%d, amount=%d, fundingErr=%s): %s",
  161. s.relayInfo.UserId, s.relayInfo.TokenId, s.tokenConsumed, err.Error(), rollbackErr.Error()))
  162. }
  163. s.tokenConsumed = 0
  164. }
  165. // TODO: model 层应定义哨兵错误(如 ErrNoActiveSubscription),用 errors.Is 替代字符串匹配
  166. errMsg := err.Error()
  167. if strings.Contains(errMsg, "no active subscription") || strings.Contains(errMsg, "subscription quota insufficient") {
  168. return types.NewErrorWithStatusCode(fmt.Errorf("订阅额度不足或未配置订阅: %s", errMsg), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
  169. }
  170. return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
  171. }
  172. s.preConsumedQuota = effectiveQuota
  173. // ---- 同步 RelayInfo 兼容字段 ----
  174. s.syncRelayInfo()
  175. return nil
  176. }
  177. // shouldTrust 统一信任额度检查,适用于钱包和订阅。
  178. func (s *BillingSession) shouldTrust(c *gin.Context) bool {
  179. // 异步任务(ForcePreConsume=true)必须预扣全额,不允许信任旁路
  180. if s.relayInfo.ForcePreConsume {
  181. return false
  182. }
  183. trustQuota := common.GetTrustQuota()
  184. if trustQuota <= 0 {
  185. return false
  186. }
  187. // 检查令牌是否充足
  188. tokenTrusted := s.relayInfo.TokenUnlimited
  189. if !tokenTrusted {
  190. tokenQuota := c.GetInt("token_quota")
  191. tokenTrusted = tokenQuota > trustQuota
  192. }
  193. if !tokenTrusted {
  194. return false
  195. }
  196. switch s.funding.Source() {
  197. case BillingSourceWallet:
  198. return s.relayInfo.UserQuota > trustQuota
  199. case BillingSourceSubscription:
  200. // 订阅不能启用信任旁路。原因:
  201. // 1. PreConsumeUserSubscription 要求 amount>0 来创建预扣记录并锁定订阅
  202. // 2. SubscriptionFunding.PreConsume 忽略参数,始终用 s.amount 预扣
  203. // 3. 若信任旁路将 effectiveQuota 设为 0,会导致 preConsumedQuota 与实际订阅预扣不一致
  204. return false
  205. default:
  206. return false
  207. }
  208. }
  209. // syncRelayInfo 将 BillingSession 的状态同步到 RelayInfo 的兼容字段上。
  210. func (s *BillingSession) syncRelayInfo() {
  211. info := s.relayInfo
  212. info.FinalPreConsumedQuota = s.preConsumedQuota
  213. info.BillingSource = s.funding.Source()
  214. if sub, ok := s.funding.(*SubscriptionFunding); ok {
  215. info.SubscriptionId = sub.subscriptionId
  216. info.SubscriptionPreConsumed = sub.preConsumed
  217. info.SubscriptionPostDelta = 0
  218. info.SubscriptionAmountTotal = sub.AmountTotal
  219. info.SubscriptionAmountUsedAfterPreConsume = sub.AmountUsedAfter
  220. info.SubscriptionPlanId = sub.PlanId
  221. info.SubscriptionPlanTitle = sub.PlanTitle
  222. } else {
  223. info.SubscriptionId = 0
  224. info.SubscriptionPreConsumed = 0
  225. }
  226. }
  227. // ---------------------------------------------------------------------------
  228. // NewBillingSession 工厂 — 根据计费偏好创建会话并处理回退
  229. // ---------------------------------------------------------------------------
  230. // NewBillingSession 根据用户计费偏好创建 BillingSession,处理 subscription_first / wallet_first 的回退。
  231. func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preConsumedQuota int) (*BillingSession, *types.NewAPIError) {
  232. if relayInfo == nil {
  233. return nil, types.NewError(fmt.Errorf("relayInfo is nil"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
  234. }
  235. // 检查是否为同步用户(从上下文中获取,避免热路径数据库查询)
  236. userSource := common.GetContextKeyString(c, constant.ContextKeyUserSource)
  237. common.SysLog(fmt.Sprintf("[RegionSync] NewBillingSession: userId=%d, userSource=%s, preConsumedQuota=%d, requestId=%s", relayInfo.UserId, userSource, preConsumedQuota, relayInfo.RequestId))
  238. if userSource == common.UserSourceSynced {
  239. user, userErr := model.GetUserById(relayInfo.UserId, false)
  240. if userErr == nil {
  241. common.SysLog(fmt.Sprintf("[RegionSync] synced user detected: userId=%d, remoteUserId=%d, syncedQuota=%d, requestId=%s", user.Id, user.RemoteUserId, user.SyncedQuota, relayInfo.RequestId))
  242. return newSyncedUserBillingSession(c, relayInfo, user, preConsumedQuota)
  243. }
  244. common.SysError(fmt.Sprintf("[RegionSync] GetUserById failed for synced user: userId=%d, err=%v", relayInfo.UserId, userErr))
  245. }
  246. pref := common.NormalizeBillingPreference(relayInfo.UserSetting.BillingPreference)
  247. // 钱包路径需要先检查用户额度
  248. tryWallet := func() (*BillingSession, *types.NewAPIError) {
  249. userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
  250. if err != nil {
  251. return nil, types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
  252. }
  253. if userQuota <= 0 {
  254. return nil, types.NewErrorWithStatusCode(
  255. fmt.Errorf("用户额度不足, 剩余额度: %s", logger.FormatQuota(userQuota)),
  256. types.ErrorCodeInsufficientUserQuota, http.StatusForbidden,
  257. types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
  258. }
  259. if userQuota-preConsumedQuota < 0 {
  260. return nil, types.NewErrorWithStatusCode(
  261. fmt.Errorf("预扣费额度失败, 用户剩余额度: %s, 需要预扣费额度: %s", logger.FormatQuota(userQuota), logger.FormatQuota(preConsumedQuota)),
  262. types.ErrorCodeInsufficientUserQuota, http.StatusForbidden,
  263. types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
  264. }
  265. relayInfo.UserQuota = userQuota
  266. session := &BillingSession{
  267. relayInfo: relayInfo,
  268. funding: &WalletFunding{userId: relayInfo.UserId},
  269. }
  270. if apiErr := session.preConsume(c, preConsumedQuota); apiErr != nil {
  271. return nil, apiErr
  272. }
  273. return session, nil
  274. }
  275. trySubscription := func() (*BillingSession, *types.NewAPIError) {
  276. subConsume := int64(preConsumedQuota)
  277. if subConsume <= 0 {
  278. subConsume = 1
  279. }
  280. session := &BillingSession{
  281. relayInfo: relayInfo,
  282. funding: &SubscriptionFunding{
  283. requestId: relayInfo.RequestId,
  284. userId: relayInfo.UserId,
  285. modelName: relayInfo.OriginModelName,
  286. amount: subConsume,
  287. },
  288. }
  289. // 必须传 subConsume 而非 preConsumedQuota,保证 SubscriptionFunding.amount、
  290. // preConsume 参数和 FinalPreConsumedQuota 三者一致,避免订阅多扣费。
  291. if apiErr := session.preConsume(c, int(subConsume)); apiErr != nil {
  292. return nil, apiErr
  293. }
  294. return session, nil
  295. }
  296. switch pref {
  297. case "subscription_only":
  298. return trySubscription()
  299. case "wallet_only":
  300. return tryWallet()
  301. case "wallet_first":
  302. session, err := tryWallet()
  303. if err != nil {
  304. if err.GetErrorCode() == types.ErrorCodeInsufficientUserQuota {
  305. return trySubscription()
  306. }
  307. return nil, err
  308. }
  309. return session, nil
  310. case "subscription_first":
  311. fallthrough
  312. default:
  313. hasSub, subCheckErr := model.HasActiveUserSubscription(relayInfo.UserId)
  314. if subCheckErr != nil {
  315. return nil, types.NewError(subCheckErr, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
  316. }
  317. if !hasSub {
  318. return tryWallet()
  319. }
  320. session, apiErr := trySubscription()
  321. if apiErr != nil {
  322. if apiErr.GetErrorCode() == types.ErrorCodeInsufficientUserQuota {
  323. return tryWallet()
  324. }
  325. return nil, apiErr
  326. }
  327. return session, nil
  328. }
  329. }
  330. // newSyncedUserBillingSession 为同步用户创建 BillingSession
  331. // 同步用户的余额在 master 节点,本地使用 SyncedUserFunding 管理预扣费和结算
  332. func newSyncedUserBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, user *model.User, preConsumedQuota int) (*BillingSession, *types.NewAPIError) {
  333. syncedQuota := user.SyncedQuota
  334. if syncedQuota <= 0 {
  335. syncedQuota = user.Quota
  336. }
  337. session := &BillingSession{
  338. relayInfo: relayInfo,
  339. funding: NewSyncedUserFunding(
  340. user.Id,
  341. user.RemoteUserId,
  342. relayInfo.RequestId,
  343. syncedQuota,
  344. ),
  345. }
  346. relayInfo.UserQuota = syncedQuota
  347. if apiErr := session.preConsume(c, preConsumedQuota); apiErr != nil {
  348. return nil, apiErr
  349. }
  350. return session, nil
  351. }