gateway_record_usage_test.go 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "strings"
  7. "testing"
  8. "time"
  9. "github.com/Wei-Shaw/sub2api/internal/config"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
  11. "github.com/stretchr/testify/require"
  12. )
  13. func newGatewayRecordUsageServiceForTest(usageRepo UsageLogRepository, userRepo UserRepository, subRepo UserSubscriptionRepository) *GatewayService {
  14. cfg := &config.Config{}
  15. cfg.Default.RateMultiplier = 1.1
  16. return NewGatewayService(
  17. nil,
  18. nil,
  19. usageRepo,
  20. nil,
  21. userRepo,
  22. subRepo,
  23. nil,
  24. nil,
  25. cfg,
  26. nil,
  27. nil,
  28. NewBillingService(cfg, nil),
  29. nil,
  30. &BillingCacheService{},
  31. nil,
  32. nil,
  33. &DeferredService{},
  34. nil,
  35. nil,
  36. nil,
  37. nil,
  38. nil,
  39. nil,
  40. nil,
  41. nil,
  42. nil,
  43. nil,
  44. nil, // userPlatformQuotaRepo
  45. )
  46. }
  47. func newGatewayRecordUsageServiceWithBillingRepoForTest(usageRepo UsageLogRepository, billingRepo UsageBillingRepository, userRepo UserRepository, subRepo UserSubscriptionRepository) *GatewayService {
  48. svc := newGatewayRecordUsageServiceForTest(usageRepo, userRepo, subRepo)
  49. svc.usageBillingRepo = billingRepo
  50. return svc
  51. }
  52. type openAIRecordUsageBestEffortLogRepoStub struct {
  53. UsageLogRepository
  54. bestEffortErr error
  55. createErr error
  56. bestEffortCalls int
  57. createCalls int
  58. lastLog *UsageLog
  59. lastCtxErr error
  60. }
  61. func (s *openAIRecordUsageBestEffortLogRepoStub) CreateBestEffort(ctx context.Context, log *UsageLog) error {
  62. s.bestEffortCalls++
  63. s.lastLog = log
  64. s.lastCtxErr = ctx.Err()
  65. return s.bestEffortErr
  66. }
  67. func (s *openAIRecordUsageBestEffortLogRepoStub) Create(ctx context.Context, log *UsageLog) (bool, error) {
  68. s.createCalls++
  69. s.lastLog = log
  70. s.lastCtxErr = ctx.Err()
  71. return false, s.createErr
  72. }
  73. func TestGatewayServiceRecordUsage_BillingUsesDetachedContext(t *testing.T) {
  74. usageRepo := &openAIRecordUsageLogRepoStub{inserted: false, err: context.DeadlineExceeded}
  75. userRepo := &openAIRecordUsageUserRepoStub{}
  76. subRepo := &openAIRecordUsageSubRepoStub{}
  77. quotaSvc := &openAIRecordUsageAPIKeyQuotaStub{}
  78. svc := newGatewayRecordUsageServiceForTest(usageRepo, userRepo, subRepo)
  79. reqCtx, cancel := context.WithCancel(context.Background())
  80. cancel()
  81. err := svc.RecordUsage(reqCtx, &RecordUsageInput{
  82. Result: &ForwardResult{
  83. RequestID: "gateway_detached_ctx",
  84. Usage: ClaudeUsage{
  85. InputTokens: 10,
  86. OutputTokens: 6,
  87. },
  88. Model: "claude-sonnet-4",
  89. Duration: time.Second,
  90. },
  91. APIKey: &APIKey{
  92. ID: 501,
  93. Quota: 100,
  94. },
  95. User: &User{ID: 601},
  96. Account: &Account{ID: 701},
  97. APIKeyService: quotaSvc,
  98. })
  99. require.NoError(t, err)
  100. require.Equal(t, 1, usageRepo.calls)
  101. require.Equal(t, 1, userRepo.deductCalls)
  102. require.NoError(t, userRepo.lastCtxErr)
  103. require.Equal(t, 1, quotaSvc.quotaCalls)
  104. require.NoError(t, quotaSvc.lastQuotaCtxErr)
  105. }
  106. func TestGatewayServiceRecordUsage_BillingFingerprintIncludesRequestPayloadHash(t *testing.T) {
  107. usageRepo := &openAIRecordUsageLogRepoStub{}
  108. billingRepo := &openAIRecordUsageBillingRepoStub{result: &UsageBillingApplyResult{Applied: true}}
  109. svc := newGatewayRecordUsageServiceWithBillingRepoForTest(usageRepo, billingRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  110. payloadHash := HashUsageRequestPayload([]byte(`{"messages":[{"role":"user","content":"hello"}]}`))
  111. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  112. Result: &ForwardResult{
  113. RequestID: "gateway_payload_hash",
  114. Usage: ClaudeUsage{
  115. InputTokens: 10,
  116. OutputTokens: 6,
  117. },
  118. Model: "claude-sonnet-4",
  119. Duration: time.Second,
  120. },
  121. APIKey: &APIKey{ID: 501, Quota: 100},
  122. User: &User{ID: 601},
  123. Account: &Account{ID: 701},
  124. RequestPayloadHash: payloadHash,
  125. })
  126. require.NoError(t, err)
  127. require.NotNil(t, billingRepo.lastCmd)
  128. require.Equal(t, payloadHash, billingRepo.lastCmd.RequestPayloadHash)
  129. }
  130. func TestGatewayServiceRecordUsage_BillingFingerprintFallsBackToContextRequestID(t *testing.T) {
  131. usageRepo := &openAIRecordUsageLogRepoStub{}
  132. billingRepo := &openAIRecordUsageBillingRepoStub{result: &UsageBillingApplyResult{Applied: true}}
  133. svc := newGatewayRecordUsageServiceWithBillingRepoForTest(usageRepo, billingRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  134. ctx := context.WithValue(context.Background(), ctxkey.RequestID, "req-local-123")
  135. err := svc.RecordUsage(ctx, &RecordUsageInput{
  136. Result: &ForwardResult{
  137. RequestID: "gateway_payload_fallback",
  138. Usage: ClaudeUsage{
  139. InputTokens: 10,
  140. OutputTokens: 6,
  141. },
  142. Model: "claude-sonnet-4",
  143. Duration: time.Second,
  144. },
  145. APIKey: &APIKey{ID: 501, Quota: 100},
  146. User: &User{ID: 601},
  147. Account: &Account{ID: 701},
  148. })
  149. require.NoError(t, err)
  150. require.NotNil(t, billingRepo.lastCmd)
  151. require.Equal(t, "local:req-local-123", billingRepo.lastCmd.RequestPayloadHash)
  152. }
  153. func TestGatewayServiceRecordUsage_PreservesRequestedAndUpstreamModels(t *testing.T) {
  154. usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
  155. svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  156. mappedModel := "claude-sonnet-4-20250514"
  157. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  158. Result: &ForwardResult{
  159. RequestID: "gateway_models_split",
  160. Usage: ClaudeUsage{InputTokens: 10, OutputTokens: 6},
  161. Model: "claude-sonnet-4",
  162. UpstreamModel: mappedModel,
  163. Duration: time.Second,
  164. },
  165. APIKey: &APIKey{ID: 501, Quota: 100},
  166. User: &User{ID: 601},
  167. Account: &Account{ID: 701},
  168. })
  169. require.NoError(t, err)
  170. require.NotNil(t, usageRepo.lastLog)
  171. require.Equal(t, "claude-sonnet-4", usageRepo.lastLog.Model)
  172. require.Equal(t, "claude-sonnet-4", usageRepo.lastLog.RequestedModel)
  173. require.NotNil(t, usageRepo.lastLog.UpstreamModel)
  174. require.Equal(t, mappedModel, *usageRepo.lastLog.UpstreamModel)
  175. }
  176. func TestGatewayServiceRecordUsage_GeminiFlashThinkingTierUsesCatalogPrice(t *testing.T) {
  177. for _, baseModel := range []string{"gemini-3.7-flash", "gemini-3.8-flash"} {
  178. t.Run(baseModel, func(t *testing.T) {
  179. usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
  180. userRepo := &openAIRecordUsageUserRepoStub{}
  181. svc := newGatewayRecordUsageServiceForTest(usageRepo, userRepo, &openAIRecordUsageSubRepoStub{})
  182. svc.billingService = NewBillingService(svc.cfg, &PricingService{pricingData: map[string]*LiteLLMModelPricing{
  183. baseModel: {InputCostPerToken: 0.75e-6, OutputCostPerToken: 3.75e-6, CacheReadInputTokenCost: 0.075e-6},
  184. }})
  185. svc.resolver = NewModelPricingResolver(nil, svc.billingService)
  186. group := &Group{ID: 27, Platform: PlatformGemini, RateMultiplier: 0.15}
  187. model := baseModel + "-medium"
  188. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  189. Result: &ForwardResult{
  190. RequestID: "gemini_thinking_tier",
  191. Model: model,
  192. UpstreamModel: model,
  193. Usage: ClaudeUsage{InputTokens: 8498, OutputTokens: 469, CacheReadInputTokens: 159248},
  194. Duration: time.Second,
  195. },
  196. APIKey: &APIKey{ID: 501, GroupID: &group.ID, Group: group},
  197. User: &User{ID: 601},
  198. Account: &Account{ID: 701, Platform: PlatformGemini, Type: AccountTypeAPIKey},
  199. })
  200. require.NoError(t, err)
  201. require.NotNil(t, usageRepo.lastLog)
  202. require.Equal(t, model, usageRepo.lastLog.Model)
  203. require.InDelta(t, 0.02007585, usageRepo.lastLog.TotalCost, 1e-12)
  204. require.InDelta(t, 0.0030113775, usageRepo.lastLog.ActualCost, 1e-12)
  205. require.InDelta(t, 0.0030113775, userRepo.lastAmount, 1e-12)
  206. })
  207. }
  208. }
  209. func TestGatewayServiceRecordUsage_PreservesChannelMappedUpstreamModel(t *testing.T) {
  210. usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
  211. svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  212. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  213. Result: &ForwardResult{
  214. RequestID: "gateway_channel_mapping_models",
  215. Usage: ClaudeUsage{InputTokens: 10, OutputTokens: 6},
  216. Model: "gpt-5.6-terra",
  217. UpstreamModel: "gpt-5.6-terra",
  218. Duration: time.Second,
  219. },
  220. APIKey: &APIKey{ID: 501, Quota: 100},
  221. User: &User{ID: 601},
  222. Account: &Account{ID: 701},
  223. ChannelUsageFields: ChannelUsageFields{
  224. OriginalModel: "gpt-5.6-sol",
  225. ChannelMappedModel: "gpt-5.6-terra",
  226. },
  227. })
  228. require.NoError(t, err)
  229. require.NotNil(t, usageRepo.lastLog)
  230. require.Equal(t, "gpt-5.6-sol", usageRepo.lastLog.RequestedModel)
  231. require.Equal(t, "gpt-5.6-terra", usageRepo.lastLog.Model)
  232. require.NotNil(t, usageRepo.lastLog.UpstreamModel)
  233. require.Equal(t, "gpt-5.6-terra", *usageRepo.lastLog.UpstreamModel)
  234. }
  235. func TestGatewayServiceRecordUsage_PreservesLoopedChannelAndAccountUpstreamModel(t *testing.T) {
  236. usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
  237. svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  238. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  239. Result: &ForwardResult{
  240. RequestID: "gateway_looped_mapping_models",
  241. Usage: ClaudeUsage{InputTokens: 10, OutputTokens: 6},
  242. Model: "gpt-5.6-terra",
  243. UpstreamModel: "gpt-5.6-sol",
  244. Duration: time.Second,
  245. },
  246. APIKey: &APIKey{ID: 501, Quota: 100},
  247. User: &User{ID: 601},
  248. Account: &Account{ID: 701},
  249. ChannelUsageFields: ChannelUsageFields{
  250. OriginalModel: "gpt-5.6-sol",
  251. ChannelMappedModel: "gpt-5.6-terra",
  252. },
  253. })
  254. require.NoError(t, err)
  255. require.NotNil(t, usageRepo.lastLog)
  256. require.Equal(t, "gpt-5.6-sol", usageRepo.lastLog.RequestedModel)
  257. require.Equal(t, "gpt-5.6-terra", usageRepo.lastLog.Model)
  258. require.NotNil(t, usageRepo.lastLog.UpstreamModel)
  259. require.Equal(t, "gpt-5.6-sol", *usageRepo.lastLog.UpstreamModel)
  260. }
  261. func TestGatewayServiceRecordUsage_EmptyImageSizeDefaultsBeforeBillingAndPersistence(t *testing.T) {
  262. imagePrice2K := 0.19
  263. groupID := int64(901)
  264. usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
  265. svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  266. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  267. Result: &ForwardResult{
  268. RequestID: "gateway_image_default_size",
  269. Model: "gemini-image",
  270. ImageCount: 1,
  271. ImageInputSize: "auto",
  272. Duration: time.Second,
  273. },
  274. APIKey: &APIKey{
  275. ID: 801,
  276. GroupID: i64p(groupID),
  277. Group: &Group{
  278. ID: groupID,
  279. RateMultiplier: 1.0,
  280. ImagePrice2K: &imagePrice2K,
  281. },
  282. },
  283. User: &User{ID: 601},
  284. Account: &Account{ID: 701},
  285. })
  286. require.NoError(t, err)
  287. require.NotNil(t, usageRepo.lastLog)
  288. require.Equal(t, 1, usageRepo.lastLog.ImageCount)
  289. require.NotNil(t, usageRepo.lastLog.ImageSize)
  290. require.Equal(t, ImageBillingSize2K, *usageRepo.lastLog.ImageSize)
  291. require.NotNil(t, usageRepo.lastLog.ImageInputSize)
  292. require.Equal(t, "auto", *usageRepo.lastLog.ImageInputSize)
  293. require.NotNil(t, usageRepo.lastLog.ImageSizeSource)
  294. require.Equal(t, ImageSizeSourceDefault, *usageRepo.lastLog.ImageSizeSource)
  295. require.InDelta(t, 0.19, usageRepo.lastLog.TotalCost, 1e-12)
  296. require.InDelta(t, 0.19, usageRepo.lastLog.ActualCost, 1e-12)
  297. }
  298. func TestGatewayServiceRecordUsage_PeakRateAffectsTokenModeImageOutputTokens(t *testing.T) {
  299. groupID := int64(902)
  300. usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
  301. userRepo := &openAIRecordUsageUserRepoStub{}
  302. svc := newGatewayRecordUsageServiceForTest(usageRepo, userRepo, &openAIRecordUsageSubRepoStub{})
  303. svc.resolver = newOpenAITokenImageChannelPricingResolverForTest(t, groupID, "gemini-image")
  304. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  305. Result: &ForwardResult{
  306. RequestID: "gateway_peak_image_tokens",
  307. Model: "gemini-image",
  308. ImageCount: 1,
  309. Usage: ClaudeUsage{
  310. InputTokens: 1000,
  311. OutputTokens: 600,
  312. ImageOutputTokens: 100,
  313. },
  314. Duration: time.Second,
  315. },
  316. APIKey: &APIKey{
  317. ID: 802,
  318. GroupID: i64p(groupID),
  319. Group: &Group{
  320. ID: groupID,
  321. RateMultiplier: 1.0,
  322. SubscriptionType: SubscriptionTypeSubscription,
  323. PeakRateEnabled: true,
  324. PeakStart: "00:00",
  325. PeakEnd: "23:59",
  326. PeakRateMultiplier: 3.0,
  327. },
  328. },
  329. User: &User{ID: 602},
  330. Account: &Account{ID: 702},
  331. })
  332. require.NoError(t, err)
  333. require.NotNil(t, usageRepo.lastLog)
  334. require.NotNil(t, usageRepo.lastLog.BillingMode)
  335. require.Equal(t, string(BillingModeToken), *usageRepo.lastLog.BillingMode)
  336. require.Equal(t, 3.0, usageRepo.lastLog.RateMultiplier)
  337. textInput := 1000 * 3e-6
  338. textOutput := 500 * 15e-6
  339. imageOutput := 100 * 15e-6
  340. expectedActual := (textInput + textOutput + imageOutput) * 3.0
  341. require.InDelta(t, textInput+textOutput+imageOutput, usageRepo.lastLog.TotalCost, 1e-12)
  342. require.InDelta(t, imageOutput, usageRepo.lastLog.ImageOutputCost, 1e-12)
  343. require.InDelta(t, expectedActual, usageRepo.lastLog.ActualCost, 1e-12)
  344. require.InDelta(t, expectedActual, userRepo.lastAmount, 1e-12)
  345. }
  346. func TestGatewayServiceRecordUsage_TimePricingUsesPricingAt(t *testing.T) {
  347. groupID := int64(904)
  348. requestStart := time.Date(2024, time.January, 2, 2, 0, 0, 0, time.UTC) // 上海 10:00
  349. usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
  350. userRepo := &openAIRecordUsageUserRepoStub{}
  351. svc := newGatewayRecordUsageServiceForTest(usageRepo, userRepo, &openAIRecordUsageSubRepoStub{})
  352. svc.resolver = newOpenAITokenImageChannelPricingResolverWithTimeForTest(t, groupID, "gpt-5.1", &ChannelTimePricing{
  353. Timezone: "Asia/Shanghai",
  354. Periods: []ChannelTimePricingPeriod{{StartTime: "09:00", EndTime: "12:00", Multiplier: 2}},
  355. })
  356. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  357. Result: &ForwardResult{
  358. RequestID: "gateway_time_pricing_request_start",
  359. Model: "gpt-5.1",
  360. Usage: ClaudeUsage{InputTokens: 1000, OutputTokens: 500},
  361. },
  362. APIKey: &APIKey{ID: 804, GroupID: i64p(groupID), Group: &Group{
  363. ID: groupID, RateMultiplier: 0.8, SubscriptionType: SubscriptionTypeSubscription,
  364. }},
  365. User: &User{ID: 604},
  366. Account: &Account{ID: 704},
  367. PricingAt: requestStart,
  368. })
  369. require.NoError(t, err)
  370. require.NotNil(t, usageRepo.lastLog)
  371. baseCost := 1000*3e-6 + 500*15e-6
  372. require.InDelta(t, baseCost*2, usageRepo.lastLog.TotalCost, 1e-12)
  373. require.InDelta(t, baseCost*2*0.8, usageRepo.lastLog.ActualCost, 1e-12)
  374. require.InDelta(t, 0.8, usageRepo.lastLog.RateMultiplier, 1e-12)
  375. }
  376. func TestGatewayServiceRecordUsage_UsesExplicitPricingAtForPeakRate(t *testing.T) {
  377. for _, platform := range []string{PlatformAnthropic, PlatformGemini, PlatformGrok, PlatformAntigravity} {
  378. t.Run(platform, func(t *testing.T) {
  379. groupID := int64(903)
  380. usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
  381. userRepo := &openAIRecordUsageUserRepoStub{}
  382. svc := newGatewayRecordUsageServiceForTest(usageRepo, userRepo, &openAIRecordUsageSubRepoStub{})
  383. svc.resolver = newOpenAITokenImageChannelPricingResolverForTest(t, groupID, "gemini-image")
  384. pricingAt := time.Date(2026, time.January, 1, 0, 30, 0, 0, time.UTC)
  385. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  386. Result: &ForwardResult{
  387. RequestID: "gateway_explicit_pricing_at_" + platform,
  388. Model: "gemini-image",
  389. ImageCount: 1,
  390. Usage: ClaudeUsage{
  391. InputTokens: 1000,
  392. OutputTokens: 600,
  393. ImageOutputTokens: 100,
  394. },
  395. },
  396. APIKey: &APIKey{
  397. ID: 803,
  398. GroupID: i64p(groupID),
  399. Group: &Group{
  400. ID: groupID,
  401. Platform: platform,
  402. RateMultiplier: 1.0,
  403. SubscriptionType: SubscriptionTypeSubscription,
  404. PeakRateEnabled: true,
  405. PeakStart: "00:00",
  406. PeakEnd: "01:00",
  407. PeakRateMultiplier: 3.0,
  408. },
  409. },
  410. User: &User{ID: 603},
  411. Account: &Account{ID: 703, Platform: platform},
  412. PricingAt: pricingAt,
  413. })
  414. require.NoError(t, err)
  415. require.NotNil(t, usageRepo.lastLog)
  416. require.Equal(t, 3.0, usageRepo.lastLog.RateMultiplier)
  417. })
  418. }
  419. }
  420. func TestGatewayServiceRecordUsage_DeepSeekAccountStatsUsesRequestPricingAtAndUpstreamModel(t *testing.T) {
  421. for _, model := range []struct {
  422. name string
  423. offPeakCost float64
  424. }{
  425. {"deepseek-v4-flash", 1000*2.2e-7 + 500*6.6e-7 + 1000*7e-9},
  426. {"deepseek-v4-pro", 1000*6.6e-7 + 500*1.98e-6 + 1000*2.2e-8},
  427. } {
  428. for _, slot := range []struct {
  429. name string
  430. pricingAt time.Time
  431. multiplier float64
  432. }{
  433. {"peak", time.Date(2026, time.August, 24, 2, 0, 0, 0, time.UTC), 2},
  434. {"off_peak", time.Date(2026, time.August, 24, 12, 0, 0, 0, time.UTC), 1},
  435. } {
  436. t.Run(model.name+"/"+slot.name, func(t *testing.T) {
  437. usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
  438. userRepo := &openAIRecordUsageUserRepoStub{}
  439. svc := newGatewayRecordUsageServiceForTest(usageRepo, userRepo, &openAIRecordUsageSubRepoStub{})
  440. groupID := int64(905)
  441. svc.channelService = newTestChannelServiceForStats(t, &Channel{ID: 1, Status: StatusActive}, groupID, PlatformDeepseek)
  442. svc.resolver = NewModelPricingResolver(svc.channelService, svc.billingService)
  443. alias := "customer-chat"
  444. inputPrice, outputPrice, cachePrice := 1e-6, 2e-6, 1e-7
  445. group := &Group{ID: groupID, Platform: PlatformDeepseek, RateMultiplier: 0.8,
  446. ModelPricing: []ChannelModelPricing{{
  447. Models: []string{alias}, BillingMode: BillingModeToken,
  448. InputPrice: &inputPrice, OutputPrice: &outputPrice, CacheReadPrice: &cachePrice,
  449. }},
  450. }
  451. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  452. Result: &ForwardResult{
  453. RequestID: "gateway_deepseek_account_stats_" + model.name + "_" + slot.name,
  454. Model: alias, UpstreamModel: model.name,
  455. Usage: ClaudeUsage{InputTokens: 1000, OutputTokens: 500, CacheReadInputTokens: 1000},
  456. },
  457. APIKey: &APIKey{ID: 805, GroupID: &groupID, Group: group},
  458. User: &User{ID: 605}, Account: &Account{ID: 705, Platform: PlatformDeepseek},
  459. PricingAt: slot.pricingAt,
  460. ChannelUsageFields: ChannelUsageFields{OriginalModel: alias, BillingModelSource: BillingModelSourceRequested},
  461. })
  462. require.NoError(t, err)
  463. require.NotNil(t, usageRepo.lastLog)
  464. log := usageRepo.lastLog
  465. require.Equal(t, alias, log.RequestedModel)
  466. require.NotNil(t, log.UpstreamModel)
  467. require.Equal(t, model.name, *log.UpstreamModel)
  468. require.WithinDuration(t, time.Now(), log.CreatedAt, time.Minute)
  469. require.False(t, log.CreatedAt.Equal(slot.pricingAt), "request pricing time must differ from record creation")
  470. customerTotal := 1000*inputPrice + 500*outputPrice + 1000*cachePrice
  471. require.InDelta(t, customerTotal, log.TotalCost, 1e-12)
  472. require.InDelta(t, customerTotal*0.8, log.ActualCost, 1e-12)
  473. require.Equal(t, 1, userRepo.deductCalls)
  474. require.InDelta(t, customerTotal*0.8, userRepo.lastAmount, 1e-12)
  475. require.NotNil(t, log.AccountStatsCost)
  476. require.InDelta(t, model.offPeakCost*slot.multiplier, *log.AccountStatsCost, 1e-12,
  477. "account cost must use the upstream model and historical PricingAt")
  478. })
  479. }
  480. }
  481. }
  482. func TestGatewayServiceRecordUsage_UsageLogWriteErrorDoesNotSkipBilling(t *testing.T) {
  483. usageRepo := &openAIRecordUsageLogRepoStub{inserted: false, err: MarkUsageLogCreateNotPersisted(context.Canceled)}
  484. userRepo := &openAIRecordUsageUserRepoStub{}
  485. subRepo := &openAIRecordUsageSubRepoStub{}
  486. quotaSvc := &openAIRecordUsageAPIKeyQuotaStub{}
  487. svc := newGatewayRecordUsageServiceForTest(usageRepo, userRepo, subRepo)
  488. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  489. Result: &ForwardResult{
  490. RequestID: "gateway_not_persisted",
  491. Usage: ClaudeUsage{
  492. InputTokens: 10,
  493. OutputTokens: 6,
  494. },
  495. Model: "claude-sonnet-4",
  496. Duration: time.Second,
  497. },
  498. APIKey: &APIKey{
  499. ID: 503,
  500. Quota: 100,
  501. },
  502. User: &User{ID: 603},
  503. Account: &Account{ID: 703},
  504. APIKeyService: quotaSvc,
  505. })
  506. require.NoError(t, err)
  507. require.Equal(t, 1, usageRepo.calls)
  508. require.Equal(t, 1, userRepo.deductCalls)
  509. require.Equal(t, 1, quotaSvc.quotaCalls)
  510. }
  511. func TestGatewayServiceRecordUsage_UsesFallbackRequestIDForUsageLog(t *testing.T) {
  512. usageRepo := &openAIRecordUsageLogRepoStub{}
  513. userRepo := &openAIRecordUsageUserRepoStub{}
  514. subRepo := &openAIRecordUsageSubRepoStub{}
  515. svc := newGatewayRecordUsageServiceForTest(usageRepo, userRepo, subRepo)
  516. ctx := context.WithValue(context.Background(), ctxkey.RequestID, "gateway-local-fallback")
  517. err := svc.RecordUsage(ctx, &RecordUsageInput{
  518. Result: &ForwardResult{
  519. RequestID: "",
  520. Usage: ClaudeUsage{
  521. InputTokens: 10,
  522. OutputTokens: 6,
  523. },
  524. Model: "claude-sonnet-4",
  525. Duration: time.Second,
  526. },
  527. APIKey: &APIKey{ID: 504},
  528. User: &User{ID: 604},
  529. Account: &Account{ID: 704},
  530. })
  531. require.NoError(t, err)
  532. require.NotNil(t, usageRepo.lastLog)
  533. require.Equal(t, "local:gateway-local-fallback", usageRepo.lastLog.RequestID)
  534. }
  535. func TestGatewayServiceRecordUsage_PrefersClientRequestIDOverUpstreamRequestID(t *testing.T) {
  536. usageRepo := &openAIRecordUsageLogRepoStub{}
  537. billingRepo := &openAIRecordUsageBillingRepoStub{result: &UsageBillingApplyResult{Applied: true}}
  538. svc := newGatewayRecordUsageServiceWithBillingRepoForTest(usageRepo, billingRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  539. ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-stable-123")
  540. ctx = context.WithValue(ctx, ctxkey.RequestID, "req-local-ignored")
  541. err := svc.RecordUsage(ctx, &RecordUsageInput{
  542. Result: &ForwardResult{
  543. RequestID: "upstream-volatile-456",
  544. Usage: ClaudeUsage{
  545. InputTokens: 10,
  546. OutputTokens: 6,
  547. },
  548. Model: "claude-sonnet-4",
  549. Duration: time.Second,
  550. },
  551. APIKey: &APIKey{ID: 506},
  552. User: &User{ID: 606},
  553. Account: &Account{ID: 706},
  554. })
  555. require.NoError(t, err)
  556. require.NotNil(t, billingRepo.lastCmd)
  557. require.Equal(t, "client:client-stable-123", billingRepo.lastCmd.RequestID)
  558. require.NotNil(t, usageRepo.lastLog)
  559. require.Equal(t, "client:client-stable-123", usageRepo.lastLog.RequestID)
  560. }
  561. func TestGatewayServiceRecordUsage_GeneratesRequestIDWhenAllSourcesMissing(t *testing.T) {
  562. usageRepo := &openAIRecordUsageLogRepoStub{}
  563. billingRepo := &openAIRecordUsageBillingRepoStub{result: &UsageBillingApplyResult{Applied: true}}
  564. svc := newGatewayRecordUsageServiceWithBillingRepoForTest(usageRepo, billingRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  565. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  566. Result: &ForwardResult{
  567. RequestID: "",
  568. Usage: ClaudeUsage{
  569. InputTokens: 10,
  570. OutputTokens: 6,
  571. },
  572. Model: "claude-sonnet-4",
  573. Duration: time.Second,
  574. },
  575. APIKey: &APIKey{ID: 507},
  576. User: &User{ID: 607},
  577. Account: &Account{ID: 707},
  578. })
  579. require.NoError(t, err)
  580. require.NotNil(t, billingRepo.lastCmd)
  581. require.True(t, strings.HasPrefix(billingRepo.lastCmd.RequestID, "generated:"))
  582. require.NotNil(t, usageRepo.lastLog)
  583. require.Equal(t, billingRepo.lastCmd.RequestID, usageRepo.lastLog.RequestID)
  584. }
  585. func TestGatewayServiceRecordUsage_DroppedUsageLogFallsBackToSyncCreate(t *testing.T) {
  586. // 计费成功后 best-effort 写入被丢弃(队列超时)时必须同步兜底,
  587. // 否则出现“已扣费但无 usage_log”的对账缺口(issue #3656)。
  588. usageRepo := &openAIRecordUsageBestEffortLogRepoStub{
  589. bestEffortErr: MarkUsageLogCreateDropped(errors.New("usage log best-effort queue full")),
  590. }
  591. billingRepo := &openAIRecordUsageBillingRepoStub{result: &UsageBillingApplyResult{Applied: true}}
  592. svc := newGatewayRecordUsageServiceWithBillingRepoForTest(usageRepo, billingRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  593. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  594. Result: &ForwardResult{
  595. RequestID: "gateway_drop_usage_log",
  596. Usage: ClaudeUsage{
  597. InputTokens: 10,
  598. OutputTokens: 6,
  599. },
  600. Model: "claude-sonnet-4",
  601. Duration: time.Second,
  602. },
  603. APIKey: &APIKey{ID: 508},
  604. User: &User{ID: 608},
  605. Account: &Account{ID: 708},
  606. })
  607. require.NoError(t, err)
  608. require.Equal(t, 1, usageRepo.bestEffortCalls)
  609. require.Equal(t, 1, usageRepo.createCalls)
  610. // 兜底调用使用的 ctx 必须仍然存活,不能带着已死的 ctx 走过场。
  611. require.NoError(t, usageRepo.lastCtxErr)
  612. }
  613. func TestGatewayServiceRecordUsage_BillingErrorWritesUnsettledUsageLog(t *testing.T) {
  614. usageRepo := &openAIRecordUsageLogRepoStub{}
  615. billingErr := errors.New("billing tx failed")
  616. billingRepo := &openAIRecordUsageBillingRepoStub{err: billingErr}
  617. userRepo := &openAIRecordUsageUserRepoStub{}
  618. subRepo := &openAIRecordUsageSubRepoStub{}
  619. svc := newGatewayRecordUsageServiceWithBillingRepoForTest(usageRepo, billingRepo, userRepo, subRepo)
  620. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  621. Result: &ForwardResult{
  622. RequestID: "gateway_billing_fail",
  623. Usage: ClaudeUsage{
  624. InputTokens: 10,
  625. OutputTokens: 6,
  626. },
  627. Model: "claude-sonnet-4",
  628. Duration: time.Second,
  629. },
  630. APIKey: &APIKey{ID: 505},
  631. User: &User{ID: 605},
  632. Account: &Account{ID: 705},
  633. })
  634. require.ErrorIs(t, err, billingErr)
  635. require.Equal(t, 1, billingRepo.calls)
  636. require.Equal(t, 1, usageRepo.calls)
  637. require.NotNil(t, usageRepo.lastLog)
  638. require.Equal(t, 10, usageRepo.lastLog.InputTokens)
  639. require.Equal(t, 6, usageRepo.lastLog.OutputTokens)
  640. require.Greater(t, usageRepo.lastLog.InputCost, 0.0)
  641. require.Greater(t, usageRepo.lastLog.OutputCost, 0.0)
  642. require.Greater(t, usageRepo.lastLog.TotalCost, 0.0)
  643. require.Zero(t, usageRepo.lastLog.ActualCost)
  644. }
  645. func TestGatewayServiceRecordUsage_ReasoningEffortPersisted(t *testing.T) {
  646. usageRepo := &openAIRecordUsageBestEffortLogRepoStub{}
  647. svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  648. effort := "max"
  649. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  650. Result: &ForwardResult{
  651. RequestID: "effort_test",
  652. Usage: ClaudeUsage{
  653. InputTokens: 10,
  654. OutputTokens: 5,
  655. },
  656. Model: "claude-opus-4-6",
  657. Duration: time.Second,
  658. ReasoningEffort: &effort,
  659. },
  660. APIKey: &APIKey{ID: 1},
  661. User: &User{ID: 1},
  662. Account: &Account{ID: 1},
  663. })
  664. require.NoError(t, err)
  665. require.NotNil(t, usageRepo.lastLog)
  666. require.NotNil(t, usageRepo.lastLog.ReasoningEffort)
  667. require.Equal(t, "max", *usageRepo.lastLog.ReasoningEffort)
  668. }
  669. func TestGatewayServiceRecordUsage_ReasoningEffortNil(t *testing.T) {
  670. usageRepo := &openAIRecordUsageBestEffortLogRepoStub{}
  671. svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  672. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  673. Result: &ForwardResult{
  674. RequestID: "no_effort_test",
  675. Usage: ClaudeUsage{
  676. InputTokens: 10,
  677. OutputTokens: 5,
  678. },
  679. Model: "claude-sonnet-4",
  680. Duration: time.Second,
  681. },
  682. APIKey: &APIKey{ID: 1},
  683. User: &User{ID: 1},
  684. Account: &Account{ID: 1},
  685. })
  686. require.NoError(t, err)
  687. require.NotNil(t, usageRepo.lastLog)
  688. require.Nil(t, usageRepo.lastLog.ReasoningEffort)
  689. }
  690. // newGatewayRecordUsageServiceWithResolverForTest mirrors production wiring for
  691. // token billing: a pricing resolver plus a grouped API key select the unified
  692. // billing path, which is the only one that honours the service tier.
  693. func newGatewayRecordUsageServiceWithResolverForTest(usageRepo UsageLogRepository) (*GatewayService, *APIKey) {
  694. svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
  695. svc.resolver = NewModelPricingResolver(nil, svc.billingService)
  696. groupID := int64(7)
  697. return svc, &APIKey{ID: 1, GroupID: &groupID, Group: &Group{ID: groupID, RateMultiplier: 1.0}}
  698. }
  699. func TestGatewayServiceRecordUsage_FastSpeedDowngradedByUpstreamResponse(t *testing.T) {
  700. usageRepo := &openAIRecordUsageBestEffortLogRepoStub{}
  701. svc, apiKey := newGatewayRecordUsageServiceWithResolverForTest(usageRepo)
  702. tier := "fast"
  703. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  704. Result: &ForwardResult{
  705. RequestID: "fast_downgraded_test",
  706. Usage: ClaudeUsage{InputTokens: 100, OutputTokens: 50},
  707. Model: "claude-opus-5",
  708. Duration: time.Second,
  709. ServiceTier: &tier,
  710. UpstreamResponseServiceTier: "standard",
  711. },
  712. APIKey: apiKey,
  713. User: &User{ID: 1},
  714. Account: &Account{ID: 1, Platform: PlatformAnthropic},
  715. })
  716. require.NoError(t, err)
  717. require.NotNil(t, usageRepo.lastLog)
  718. require.NotNil(t, usageRepo.lastLog.ServiceTier)
  719. require.Equal(t, "standard", *usageRepo.lastLog.ServiceTier)
  720. tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
  721. standardCost, err := svc.billingService.CalculateCost("claude-opus-5", tokens, 1.0)
  722. require.NoError(t, err)
  723. fastCost, err := svc.billingService.CalculateCostWithServiceTier("claude-opus-5", tokens, 1.0, "fast")
  724. require.NoError(t, err)
  725. require.Greater(t, fastCost.TotalCost, standardCost.TotalCost, "fast mode must carry a premium for the test to be meaningful")
  726. require.InDelta(t, standardCost.TotalCost, usageRepo.lastLog.TotalCost, 1e-10)
  727. }
  728. func TestGatewayServiceRecordUsage_FastSpeedHonouredKeepsPremium(t *testing.T) {
  729. usageRepo := &openAIRecordUsageBestEffortLogRepoStub{}
  730. svc, apiKey := newGatewayRecordUsageServiceWithResolverForTest(usageRepo)
  731. tier := "fast"
  732. err := svc.RecordUsage(context.Background(), &RecordUsageInput{
  733. Result: &ForwardResult{
  734. RequestID: "fast_honoured_test",
  735. Usage: ClaudeUsage{InputTokens: 100, OutputTokens: 50},
  736. Model: "claude-opus-5",
  737. Duration: time.Second,
  738. ServiceTier: &tier,
  739. UpstreamResponseServiceTier: "fast",
  740. },
  741. APIKey: apiKey,
  742. User: &User{ID: 1},
  743. Account: &Account{ID: 1, Platform: PlatformAnthropic},
  744. })
  745. require.NoError(t, err)
  746. require.NotNil(t, usageRepo.lastLog)
  747. require.Equal(t, "fast", *usageRepo.lastLog.ServiceTier)
  748. fastCost, err := svc.billingService.CalculateCostWithServiceTier("claude-opus-5", UsageTokens{InputTokens: 100, OutputTokens: 50}, 1.0, "fast")
  749. require.NoError(t, err)
  750. require.InDelta(t, fastCost.TotalCost, usageRepo.lastLog.TotalCost, 1e-10)
  751. }