openai_quota_spark_window_test.go 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779
  1. package service
  2. import (
  3. "context"
  4. "crypto/ed25519"
  5. "crypto/rand"
  6. "crypto/x509"
  7. "encoding/base64"
  8. "encoding/json"
  9. "errors"
  10. "fmt"
  11. "net/http"
  12. "net/http/httptest"
  13. "net/url"
  14. "strings"
  15. "testing"
  16. "time"
  17. "github.com/imroc/req/v3"
  18. "github.com/stretchr/testify/require"
  19. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  20. )
  21. // ── stub helpers ─────────────────────────────────────────────────────────────
  22. // stubQuotaAccountRepo 是多账号 AccountRepository stub,仅实现 GetByID。
  23. type stubQuotaAccountRepo struct {
  24. AccountRepository
  25. accounts map[int64]*Account
  26. extraUpdates map[int64]map[string]any
  27. extraUpdateCalls int
  28. extraUpdateErr error
  29. }
  30. func (r *stubQuotaAccountRepo) GetByID(_ context.Context, id int64) (*Account, error) {
  31. acc, ok := r.accounts[id]
  32. if !ok {
  33. return nil, fmt.Errorf("account %d not found", id)
  34. }
  35. return acc, nil
  36. }
  37. func (r *stubQuotaAccountRepo) UpdateCredentials(_ context.Context, id int64, credentials map[string]any) error {
  38. acc, ok := r.accounts[id]
  39. if !ok {
  40. return fmt.Errorf("account %d not found", id)
  41. }
  42. acc.Credentials = credentials
  43. return nil
  44. }
  45. func (r *stubQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error {
  46. if r.extraUpdateErr != nil {
  47. return r.extraUpdateErr
  48. }
  49. r.extraUpdateCalls++
  50. if r.extraUpdates == nil {
  51. r.extraUpdates = make(map[int64]map[string]any)
  52. }
  53. r.extraUpdates[id] = updates
  54. return nil
  55. }
  56. // stubQuotaTokenCache 实现 OpenAITokenCache,返回预设静态 token。
  57. type stubQuotaTokenCache struct {
  58. tokens map[string]string
  59. }
  60. func (c *stubQuotaTokenCache) GetAccessToken(_ context.Context, key string) (string, error) {
  61. if t, ok := c.tokens[key]; ok {
  62. return t, nil
  63. }
  64. return "", errors.New("token not found")
  65. }
  66. func (c *stubQuotaTokenCache) SetAccessToken(_ context.Context, _ string, _ string, _ time.Duration) error {
  67. return nil
  68. }
  69. func (c *stubQuotaTokenCache) DeleteAccessToken(_ context.Context, _ string) error { return nil }
  70. func (c *stubQuotaTokenCache) AcquireRefreshLock(_ context.Context, _ string, _ time.Duration) (bool, error) {
  71. return true, nil
  72. }
  73. func (c *stubQuotaTokenCache) ReleaseRefreshLock(_ context.Context, _ string) error { return nil }
  74. // newQuotaRedirectingFactory 返回 PrivacyClientFactory,将请求重定向到 httptest.Server。
  75. func newQuotaRedirectingFactory(srv *httptest.Server) PrivacyClientFactory {
  76. targetURL, _ := url.Parse(srv.URL)
  77. return func(_ string) (*req.Client, error) {
  78. c := req.C().WrapRoundTripFunc(func(rt req.RoundTripper) req.RoundTripFunc {
  79. return func(r *req.Request) (*req.Response, error) {
  80. r.URL.Scheme = targetURL.Scheme
  81. r.URL.Host = targetURL.Host
  82. return rt.RoundTrip(r)
  83. }
  84. })
  85. return c, nil
  86. }
  87. }
  88. // ── Part A: buildCodexSparkWindowExtraUpdates ─────────────────────────────────
  89. // TestBuildCodexSparkWindowExtraUpdates_ContainsCodexKeys 验证:
  90. // - 产出包含 codex_5h_used_percent / codex_7d_used_percent
  91. // - 不含任何 codex_spark_ 前缀的 key(Method Z 前缀已禁止)
  92. // - 数值正确映射(primary 较短→5h,secondary 较长→7d)
  93. func TestBuildCodexSparkWindowExtraUpdates_ContainsCodexKeys(t *testing.T) {
  94. now := time.Now().UTC()
  95. usage := &OpenAIQuotaUsage{
  96. AdditionalRateLimits: []OpenAIAdditionalRateLimit{
  97. {
  98. MeteredFeature: "codex_bengalfox",
  99. RateLimit: &OpenAIRateLimit{
  100. PrimaryWindow: &OpenAIRateLimitWindow{
  101. UsedPercent: 0.42,
  102. LimitWindowSeconds: 18000, // 300 min = 5 h
  103. ResetAfterSeconds: 3600,
  104. },
  105. SecondaryWindow: &OpenAIRateLimitWindow{
  106. UsedPercent: 0.15,
  107. LimitWindowSeconds: 604800, // 7 d
  108. ResetAfterSeconds: 86400,
  109. },
  110. },
  111. },
  112. },
  113. }
  114. updates := buildCodexSparkWindowExtraUpdates(usage, now)
  115. require.NotNil(t, updates, "expected non-nil updates for valid codex_bengalfox entry")
  116. // 必须含有 codex_5h_* 和 codex_7d_* 键
  117. require.Contains(t, updates, "codex_5h_used_percent")
  118. require.Contains(t, updates, "codex_7d_used_percent")
  119. // 任何键不得含有 codex_spark_ 前缀(Method Z 已禁止)
  120. for k := range updates {
  121. require.False(t, strings.Contains(k, "codex_spark_"),
  122. "unexpected Method-Z prefix in key: %s", k)
  123. }
  124. // 数值验证(primary=5h, secondary=7d)
  125. require.InDelta(t, 0.42, updates["codex_5h_used_percent"], 1e-9)
  126. require.InDelta(t, 0.15, updates["codex_7d_used_percent"], 1e-9)
  127. }
  128. // TestBuildCodexSparkWindowExtraUpdates_NilUsage 验证 nil usage 返回 nil。
  129. func TestBuildCodexSparkWindowExtraUpdates_NilUsage(t *testing.T) {
  130. require.Nil(t, buildCodexSparkWindowExtraUpdates(nil, time.Now()))
  131. }
  132. // TestBuildCodexSparkWindowExtraUpdates_NoBengalfox 验证无 codex_bengalfox 条目时返回 nil。
  133. func TestBuildCodexSparkWindowExtraUpdates_NoBengalfox(t *testing.T) {
  134. usage := &OpenAIQuotaUsage{
  135. AdditionalRateLimits: []OpenAIAdditionalRateLimit{
  136. {MeteredFeature: "other_feature", RateLimit: &OpenAIRateLimit{}},
  137. },
  138. }
  139. require.Nil(t, buildCodexSparkWindowExtraUpdates(usage, time.Now()))
  140. }
  141. // ── Part C: ResetCredit 影子拒绝 ───────────────────────────────────────────
  142. // TestResetCreditShadowRejected 验证:
  143. // - ResetCredit(ctx, shadowID) 返回 ErrSparkShadowResetNotSupported
  144. // - 不触达上游(privacyClientFactory 为 nil,若调用则 panic)
  145. func TestResetCreditShadowRejected(t *testing.T) {
  146. pid := int64(100)
  147. shadow := &Account{
  148. ID: 200,
  149. ParentAccountID: &pid,
  150. Platform: PlatformOpenAI,
  151. Type: AccountTypeOAuth,
  152. QuotaDimension: QuotaDimensionSpark,
  153. }
  154. repo := &stubQuotaAccountRepo{
  155. accounts: map[int64]*Account{200: shadow},
  156. }
  157. // privacyClientFactory 故意为 nil —— 若流程误到上游则 prepareUpstreamCall 会先在
  158. // 配置检查处报错,但我们在此之前就应该拦截并返回 ErrSparkShadowResetNotSupported。
  159. svc := &OpenAIQuotaService{accountRepo: repo}
  160. _, err := svc.ResetCredit(context.Background(), 200)
  161. require.ErrorIs(t, err, ErrSparkShadowResetNotSupported,
  162. "shadow ResetCredit should return ErrSparkShadowResetNotSupported, got: %v", err)
  163. // 外审 F6:必须是结构化 409(而非裸 error→500)。
  164. require.Equal(t, http.StatusConflict, infraerrors.Code(err),
  165. "shadow ResetCredit 应映射为 409 Conflict 而非 500")
  166. }
  167. func TestResetCreditTargetedSendsStableCreditAndRedeemIDs(t *testing.T) {
  168. account := &Account{
  169. ID: 203, Platform: PlatformOpenAI, Type: AccountTypeOAuth,
  170. Credentials: map[string]any{"chatgpt_account_id": "account-targeted"},
  171. }
  172. repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
  173. tokenCache := &stubQuotaTokenCache{tokens: map[string]string{OpenAITokenCacheKey(account): "fake-token"}}
  174. tokenProvider := NewOpenAITokenProvider(repo, tokenCache, nil)
  175. var body map[string]string
  176. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  177. require.Equal(t, "/backend-api/wham/rate-limit-reset-credits/consume", r.URL.Path)
  178. require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
  179. w.Header().Set("content-type", "application/json")
  180. _, _ = w.Write([]byte(`{"code":"ok","windows_reset":1}`))
  181. }))
  182. defer srv.Close()
  183. svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
  184. result, err := svc.ResetCreditTargeted(context.Background(), account.ID, "credit-123", "redeem-456")
  185. require.NoError(t, err)
  186. require.Equal(t, "ok", result.Code)
  187. require.Equal(t, map[string]string{
  188. "credit_id": "credit-123", "redeem_request_id": "redeem-456",
  189. }, body)
  190. }
  191. func TestResetCreditAgentIdentityUsesAssertionAndRecoversInvalidTaskOnce(t *testing.T) {
  192. _, privateKey, err := ed25519.GenerateKey(rand.Reader)
  193. require.NoError(t, err)
  194. der, err := x509.MarshalPKCS8PrivateKey(privateKey)
  195. require.NoError(t, err)
  196. account := &Account{
  197. ID: 201,
  198. Platform: PlatformOpenAI,
  199. Type: AccountTypeOAuth,
  200. Credentials: map[string]any{
  201. "auth_mode": OpenAIAuthModeAgentIdentity,
  202. "agent_runtime_id": "runtime-reset-recovery",
  203. "agent_private_key": base64.StdEncoding.EncodeToString(der),
  204. "task_id": "task-reset-old",
  205. "chatgpt_account_id": "account-reset-recovery",
  206. },
  207. }
  208. repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
  209. resetCalls := 0
  210. registerCalls := 0
  211. var assertions []string
  212. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  213. w.Header().Set("content-type", "application/json")
  214. if strings.Contains(r.URL.Path, "/task/register") {
  215. registerCalls++
  216. _, _ = w.Write([]byte(`{"task_id":"task-reset-new"}`))
  217. return
  218. }
  219. resetCalls++
  220. assertions = append(assertions, r.Header.Get("authorization"))
  221. require.Equal(t, "account-reset-recovery", r.Header.Get("chatgpt-account-id"))
  222. if resetCalls == 1 {
  223. w.WriteHeader(http.StatusUnauthorized)
  224. _, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`))
  225. return
  226. }
  227. _, _ = w.Write([]byte(`{"code":"ok","windows_reset":2}`))
  228. }))
  229. defer srv.Close()
  230. oldBase := openAIAgentIdentityAuthAPIBaseURL
  231. openAIAgentIdentityAuthAPIBaseURL = srv.URL
  232. t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
  233. invalidator := &agentIdentityWSInvalidationRecorder{}
  234. svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
  235. svc.agentIdentityWS = invalidator
  236. result, err := svc.ResetCredit(context.Background(), account.ID)
  237. require.NoError(t, err)
  238. require.Equal(t, "ok", result.Code)
  239. require.Equal(t, 2, result.WindowsReset)
  240. require.Equal(t, 2, resetCalls)
  241. require.Equal(t, 1, registerCalls)
  242. require.Len(t, assertions, 2)
  243. require.True(t, strings.HasPrefix(assertions[0], "AgentAssertion "))
  244. require.True(t, strings.HasPrefix(assertions[1], "AgentAssertion "))
  245. require.NotEqual(t, assertions[0], assertions[1])
  246. require.Equal(t, "task-reset-new", account.GetCredential("task_id"))
  247. require.Equal(t, []int64{account.ID}, invalidator.accountIDs)
  248. }
  249. func TestResetCreditAgentIdentityReusesConcurrentlyRecoveredTask(t *testing.T) {
  250. _, privateKey, err := ed25519.GenerateKey(rand.Reader)
  251. require.NoError(t, err)
  252. der, err := x509.MarshalPKCS8PrivateKey(privateKey)
  253. require.NoError(t, err)
  254. account := &Account{
  255. ID: 202,
  256. Platform: PlatformOpenAI,
  257. Type: AccountTypeOAuth,
  258. Credentials: map[string]any{
  259. "auth_mode": OpenAIAuthModeAgentIdentity,
  260. "agent_runtime_id": "runtime-reset-concurrent",
  261. "agent_private_key": base64.StdEncoding.EncodeToString(der),
  262. "task_id": "task-reset-old",
  263. "chatgpt_account_id": "account-reset-concurrent",
  264. },
  265. }
  266. repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
  267. resetCalls := 0
  268. registerCalls := 0
  269. var assertions []string
  270. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  271. w.Header().Set("content-type", "application/json")
  272. if strings.Contains(r.URL.Path, "/task/register") {
  273. registerCalls++
  274. _, _ = w.Write([]byte(`{"task_id":"task-reset-unexpected"}`))
  275. return
  276. }
  277. resetCalls++
  278. assertions = append(assertions, r.Header.Get("authorization"))
  279. if resetCalls == 1 {
  280. credentials := shallowCopyMap(account.Credentials)
  281. credentials["task_id"] = "task-reset-concurrent"
  282. account.Credentials = credentials
  283. w.WriteHeader(http.StatusUnauthorized)
  284. _, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`))
  285. return
  286. }
  287. _, _ = w.Write([]byte(`{"code":"ok","windows_reset":1}`))
  288. }))
  289. defer srv.Close()
  290. oldBase := openAIAgentIdentityAuthAPIBaseURL
  291. openAIAgentIdentityAuthAPIBaseURL = srv.URL
  292. t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
  293. svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
  294. result, err := svc.ResetCredit(context.Background(), account.ID)
  295. require.NoError(t, err)
  296. require.Equal(t, "ok", result.Code)
  297. require.Equal(t, 2, resetCalls)
  298. require.Zero(t, registerCalls)
  299. require.Equal(t, "task-reset-old", decodeAgentAssertionTask(t, assertions[0]))
  300. require.Equal(t, "task-reset-concurrent", decodeAgentAssertionTask(t, assertions[1]))
  301. }
  302. // ── Part B: prepareUpstreamCall 影子 resolve ──────────────────────────────
  303. // TestPrepareUpstreamCallShadowResolve 验证影子账号(200)QueryUsage 时:
  304. // - 不因 chatgpt_account_id 为空而报错
  305. // - 使用母账号(100)的 chatgpt_account_id("org-parent123")
  306. //
  307. // 测试策略: 直接调用包内可见的 prepareUpstreamCall,注入 stubTokenCache(命中路径)
  308. // 和 stubQuotaAccountRepo(同时持有影子+母账号),绕开 /wham/usage HTTP 往返。
  309. // 这比 httptest 端到端 mock 更轻量且对实现细节的耦合更低。
  310. func TestPrepareUpstreamCallShadowResolve(t *testing.T) {
  311. ctx := context.Background()
  312. pid := int64(100)
  313. // 影子账号:无 chatgpt_account_id credentials
  314. shadow := &Account{
  315. ID: 200,
  316. ParentAccountID: &pid,
  317. Platform: PlatformOpenAI,
  318. Type: AccountTypeOAuth,
  319. Status: StatusActive,
  320. QuotaDimension: QuotaDimensionSpark,
  321. }
  322. // 母账号:有完整 credentials
  323. parent := &Account{
  324. ID: 100,
  325. Platform: PlatformOpenAI,
  326. Type: AccountTypeOAuth,
  327. Status: StatusActive,
  328. Credentials: map[string]any{
  329. "chatgpt_account_id": "org-parent123",
  330. },
  331. }
  332. repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{200: shadow, 100: parent}}
  333. // stubTokenCache 为母账号 cache key 提供 fake token(走缓存命中路径,无需真实刷新)
  334. tokenCache := &stubQuotaTokenCache{tokens: map[string]string{
  335. OpenAITokenCacheKey(parent): "fake-access-token",
  336. }}
  337. tokenProvider := NewOpenAITokenProvider(repo, tokenCache, nil)
  338. // privacyClientFactory 可以是任意合法工厂;prepareUpstreamCall 在返回前不调用它
  339. svc := NewOpenAIQuotaService(repo, nil, tokenProvider, func(_ string) (*req.Client, error) {
  340. return req.C(), nil
  341. })
  342. _, chatGPTAccountID, _, _, err := svc.prepareUpstreamCall(ctx, 200)
  343. require.NoError(t, err, "shadow resolve should succeed; got error: %v", err)
  344. require.Equal(t, "org-parent123", chatGPTAccountID,
  345. "prepareUpstreamCall should use parent's chatgpt_account_id after shadow resolve")
  346. }
  347. func TestQueryUsageAgentIdentityUsesAssertionWithoutOAuthToken(t *testing.T) {
  348. _, privateKey, err := ed25519.GenerateKey(rand.Reader)
  349. require.NoError(t, err)
  350. der, err := x509.MarshalPKCS8PrivateKey(privateKey)
  351. require.NoError(t, err)
  352. account := &Account{
  353. ID: 300,
  354. Platform: PlatformOpenAI,
  355. Type: AccountTypeOAuth,
  356. Credentials: map[string]any{
  357. "auth_mode": OpenAIAuthModeAgentIdentity,
  358. "agent_runtime_id": "runtime-quota",
  359. "agent_private_key": base64.StdEncoding.EncodeToString(der),
  360. "task_id": "task-quota",
  361. "chatgpt_account_id": "account-quota",
  362. "chatgpt_account_is_fedramp": true,
  363. },
  364. }
  365. repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
  366. var authorization string
  367. var accountHeader string
  368. var fedrampHeader string
  369. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  370. authorization = r.Header.Get("authorization")
  371. accountHeader = r.Header.Get("chatgpt-account-id")
  372. fedrampHeader = r.Header.Get("x-openai-fedramp")
  373. w.Header().Set("content-type", "application/json")
  374. _, _ = w.Write([]byte(`{"plan_type":"pro","rate_limit":{"allowed":true}}`))
  375. }))
  376. defer srv.Close()
  377. svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
  378. usage, err := svc.QueryUsage(context.Background(), account.ID)
  379. require.NoError(t, err)
  380. require.NotNil(t, usage)
  381. require.True(t, strings.HasPrefix(authorization, "AgentAssertion "))
  382. require.Equal(t, "account-quota", accountHeader)
  383. require.Equal(t, "true", fedrampHeader)
  384. }
  385. func TestQueryUsageAgentIdentityRecoversInvalidTaskOnce(t *testing.T) {
  386. _, privateKey, err := ed25519.GenerateKey(rand.Reader)
  387. require.NoError(t, err)
  388. der, err := x509.MarshalPKCS8PrivateKey(privateKey)
  389. require.NoError(t, err)
  390. account := &Account{
  391. ID: 301,
  392. Platform: PlatformOpenAI,
  393. Type: AccountTypeOAuth,
  394. Credentials: map[string]any{
  395. "auth_mode": OpenAIAuthModeAgentIdentity,
  396. "agent_runtime_id": "runtime-quota-recovery",
  397. "agent_private_key": base64.StdEncoding.EncodeToString(der),
  398. "task_id": "task-quota-old",
  399. "chatgpt_account_id": "account-quota-recovery",
  400. },
  401. }
  402. repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
  403. usageCalls := 0
  404. registerCalls := 0
  405. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  406. w.Header().Set("content-type", "application/json")
  407. if strings.Contains(r.URL.Path, "/task/register") {
  408. registerCalls++
  409. _, _ = w.Write([]byte(`{"task_id":"task-quota-new"}`))
  410. return
  411. }
  412. if strings.Contains(r.URL.Path, "rate-limit-reset-credits") {
  413. _, _ = w.Write([]byte(`{}`))
  414. return
  415. }
  416. usageCalls++
  417. if usageCalls == 1 {
  418. w.WriteHeader(http.StatusUnauthorized)
  419. _, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`))
  420. return
  421. }
  422. _, _ = w.Write([]byte(`{"plan_type":"pro","rate_limit":{"allowed":true}}`))
  423. }))
  424. defer srv.Close()
  425. oldBase := openAIAgentIdentityAuthAPIBaseURL
  426. openAIAgentIdentityAuthAPIBaseURL = srv.URL
  427. t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
  428. invalidator := &agentIdentityWSInvalidationRecorder{}
  429. svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
  430. svc.agentIdentityWS = invalidator
  431. usage, err := svc.QueryUsage(context.Background(), account.ID)
  432. require.NoError(t, err)
  433. require.NotNil(t, usage)
  434. require.Equal(t, 2, usageCalls)
  435. require.Equal(t, 1, registerCalls)
  436. require.Equal(t, "task-quota-new", account.GetCredential("task_id"))
  437. require.Equal(t, []int64{account.ID}, invalidator.accountIDs)
  438. }
  439. func TestParseOpenAIRateLimitResetCreditDetails_CompatibleContainers(t *testing.T) {
  440. tests := []struct {
  441. name string
  442. body string
  443. want []string
  444. }{
  445. {
  446. name: "credits",
  447. body: `{"credits":[{"id":"secret-id","expires_at":"2026-07-03T04:05:06Z"}]}`,
  448. want: []string{"2026-07-03T04:05:06Z"},
  449. },
  450. {
  451. name: "rate limit reset credits",
  452. body: `{"rate_limit_reset_credits":[{"expiresAt":"2026-07-04T04:05:06Z"}]}`,
  453. want: []string{"2026-07-04T04:05:06Z"},
  454. },
  455. {
  456. name: "items",
  457. body: `{"items":[{"expires_at":"2026-07-05T04:05:06Z"}]}`,
  458. want: []string{"2026-07-05T04:05:06Z"},
  459. },
  460. {
  461. name: "data",
  462. body: `{"data":[{"expires_at":"2026-07-06T04:05:06Z"}]}`,
  463. want: []string{"2026-07-06T04:05:06Z"},
  464. },
  465. {
  466. name: "array",
  467. body: `[{"expires_at":"2026-07-07T04:05:06Z"}]`,
  468. want: []string{"2026-07-07T04:05:06Z"},
  469. },
  470. }
  471. for _, tt := range tests {
  472. t.Run(tt.name, func(t *testing.T) {
  473. got, err := parseOpenAIRateLimitResetCreditDetails([]byte(tt.body))
  474. require.NoError(t, err)
  475. require.Len(t, got.Credits, len(tt.want))
  476. for i := range tt.want {
  477. require.Equal(t, tt.want[i], got.Credits[i].ExpiresAt)
  478. }
  479. encoded, err := json.Marshal(got.Credits)
  480. require.NoError(t, err)
  481. require.NotContains(t, string(encoded), "secret-id")
  482. })
  483. }
  484. }
  485. func TestQueryUsageIncludesResetCreditExpirations_EndToEnd(t *testing.T) {
  486. ctx := context.Background()
  487. account := &Account{
  488. ID: 100,
  489. Platform: PlatformOpenAI,
  490. Type: AccountTypeOAuth,
  491. Status: StatusActive,
  492. Credentials: map[string]any{
  493. "chatgpt_account_id": "org-parent123",
  494. },
  495. }
  496. repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: account}}
  497. tokenCache := &stubQuotaTokenCache{tokens: map[string]string{
  498. OpenAITokenCacheKey(account): "fake-token",
  499. }}
  500. tokenProvider := NewOpenAITokenProvider(repo, tokenCache, nil)
  501. var capturedBeta string
  502. var detailCalls int
  503. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  504. w.Header().Set("content-type", "application/json")
  505. switch r.URL.Path {
  506. case "/backend-api/wham/usage":
  507. _ = json.NewEncoder(w).Encode(OpenAIQuotaUsage{
  508. RateLimitResetCredits: &OpenAIRateLimitResetCredits{AvailableCount: 2},
  509. })
  510. case "/backend-api/wham/rate-limit-reset-credits":
  511. detailCalls++
  512. capturedBeta = r.Header.Get("OpenAI-Beta")
  513. require.Equal(t, "org-parent123", r.Header.Get("ChatGPT-Account-ID"))
  514. _, _ = w.Write([]byte(`{"credits":[{"id":"secret-credit-id","expires_at":"2026-07-03T04:05:06Z"},{"expiresAt":"2026-07-04T04:05:06Z"}]}`))
  515. default:
  516. http.NotFound(w, r)
  517. }
  518. }))
  519. defer srv.Close()
  520. svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
  521. usage, err := svc.QueryUsage(ctx, 100)
  522. require.NoError(t, err)
  523. require.NotNil(t, usage)
  524. require.NotNil(t, usage.RateLimitResetCredits)
  525. require.Equal(t, 2, usage.RateLimitResetCredits.AvailableCount)
  526. require.Equal(t, 1, detailCalls)
  527. require.Equal(t, openaiQuotaCodexBeta, capturedBeta)
  528. require.Equal(t, []OpenAIRateLimitResetCreditDetail{
  529. {ExpiresAt: "2026-07-03T04:05:06Z"},
  530. {ExpiresAt: "2026-07-04T04:05:06Z"},
  531. }, usage.RateLimitResetCredits.Credits)
  532. require.NoError(t, svc.CacheResetCreditsSnapshot(ctx, 100, usage.RateLimitResetCredits))
  533. require.Equal(t, &OpenAIRateLimitResetCredits{
  534. AvailableCount: 2,
  535. Credits: []OpenAIRateLimitResetCreditDetail{
  536. {ExpiresAt: "2026-07-03T04:05:06Z"},
  537. {ExpiresAt: "2026-07-04T04:05:06Z"},
  538. },
  539. }, repo.extraUpdates[100][openaiQuotaResetCreditsKey])
  540. encoded, err := json.Marshal(usage)
  541. require.NoError(t, err)
  542. require.NotContains(t, string(encoded), "secret-credit-id")
  543. }
  544. func TestQueryUsageResetCreditDetails401NonFatal(t *testing.T) {
  545. ctx := context.Background()
  546. account := &Account{
  547. ID: 100,
  548. Platform: PlatformOpenAI,
  549. Type: AccountTypeOAuth,
  550. Status: StatusActive,
  551. Credentials: map[string]any{
  552. "chatgpt_account_id": "org-parent123",
  553. },
  554. }
  555. repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: account}}
  556. tokenCache := &stubQuotaTokenCache{tokens: map[string]string{
  557. OpenAITokenCacheKey(account): "fake-token",
  558. }}
  559. tokenProvider := NewOpenAITokenProvider(repo, tokenCache, nil)
  560. var detailCalls int
  561. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  562. w.Header().Set("content-type", "application/json")
  563. switch r.URL.Path {
  564. case "/backend-api/wham/usage":
  565. _ = json.NewEncoder(w).Encode(OpenAIQuotaUsage{
  566. RateLimitResetCredits: &OpenAIRateLimitResetCredits{AvailableCount: 1},
  567. })
  568. case "/backend-api/wham/rate-limit-reset-credits":
  569. detailCalls++
  570. w.WriteHeader(http.StatusUnauthorized)
  571. _, _ = w.Write([]byte(`{"error":"unauthorized","id":"secret-error-id"}`))
  572. default:
  573. http.NotFound(w, r)
  574. }
  575. }))
  576. defer srv.Close()
  577. svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
  578. usage, err := svc.QueryUsage(ctx, 100)
  579. require.NoError(t, err)
  580. require.NotNil(t, usage)
  581. require.NotNil(t, usage.RateLimitResetCredits)
  582. require.Equal(t, 1, usage.RateLimitResetCredits.AvailableCount)
  583. require.Equal(t, 1, detailCalls)
  584. require.Empty(t, usage.RateLimitResetCredits.Credits)
  585. // A count without expiration details must not be persisted (the reader could
  586. // never age it out), and the previous snapshot must survive untouched.
  587. require.Error(t, svc.CacheResetCreditsSnapshot(ctx, 100, usage.RateLimitResetCredits))
  588. require.Empty(t, repo.extraUpdates)
  589. }
  590. func TestCacheResetCreditsSnapshot(t *testing.T) {
  591. ctx := context.Background()
  592. t.Run("zero count allows an empty expiration list", func(t *testing.T) {
  593. repo := &stubQuotaAccountRepo{}
  594. svc := &OpenAIQuotaService{accountRepo: repo}
  595. credits := &OpenAIRateLimitResetCredits{AvailableCount: 0}
  596. require.NoError(t, svc.CacheResetCreditsSnapshot(ctx, 100, credits))
  597. require.Equal(t, credits, repo.extraUpdates[100][openaiQuotaResetCreditsKey])
  598. })
  599. t.Run("missing expiration list preserves the cache", func(t *testing.T) {
  600. repo := &stubQuotaAccountRepo{}
  601. svc := &OpenAIQuotaService{accountRepo: repo}
  602. err := svc.CacheResetCreditsSnapshot(ctx, 100, &OpenAIRateLimitResetCredits{AvailableCount: 1})
  603. require.Error(t, err)
  604. require.Empty(t, repo.extraUpdates)
  605. })
  606. t.Run("empty expiration list with a positive count preserves the cache", func(t *testing.T) {
  607. repo := &stubQuotaAccountRepo{}
  608. svc := &OpenAIQuotaService{accountRepo: repo}
  609. err := svc.CacheResetCreditsSnapshot(ctx, 100, &OpenAIRateLimitResetCredits{
  610. AvailableCount: 2,
  611. Credits: []OpenAIRateLimitResetCreditDetail{},
  612. })
  613. require.Error(t, err)
  614. require.Empty(t, repo.extraUpdates)
  615. })
  616. t.Run("nil snapshot preserves the cache", func(t *testing.T) {
  617. repo := &stubQuotaAccountRepo{}
  618. svc := &OpenAIQuotaService{accountRepo: repo}
  619. require.Error(t, svc.CacheResetCreditsSnapshot(ctx, 100, nil))
  620. require.Empty(t, repo.extraUpdates)
  621. })
  622. t.Run("repository errors are returned", func(t *testing.T) {
  623. repo := &stubQuotaAccountRepo{extraUpdateErr: errors.New("database unavailable")}
  624. svc := &OpenAIQuotaService{accountRepo: repo}
  625. err := svc.CacheResetCreditsSnapshot(ctx, 100, &OpenAIRateLimitResetCredits{
  626. AvailableCount: 1,
  627. Credits: []OpenAIRateLimitResetCreditDetail{{ExpiresAt: "2026-07-03T04:05:06Z"}},
  628. })
  629. require.ErrorContains(t, err, "database unavailable")
  630. })
  631. }
  632. func TestCachePostResetSnapshot(t *testing.T) {
  633. repo := &stubQuotaAccountRepo{}
  634. svc := &OpenAIQuotaService{accountRepo: repo}
  635. credits := &OpenAIRateLimitResetCredits{AvailableCount: 0}
  636. usage := &OpenAIQuotaUsage{
  637. RateLimitResetCredits: credits,
  638. RateLimit: &OpenAIRateLimit{
  639. PrimaryWindow: &OpenAIRateLimitWindow{
  640. UsedPercent: 0, LimitWindowSeconds: 5 * 60 * 60, ResetAfterSeconds: 5 * 60 * 60,
  641. },
  642. SecondaryWindow: &OpenAIRateLimitWindow{
  643. UsedPercent: 0, LimitWindowSeconds: 7 * 24 * 60 * 60, ResetAfterSeconds: 7 * 24 * 60 * 60,
  644. },
  645. },
  646. }
  647. require.NoError(t, svc.CachePostResetSnapshot(context.Background(), 100, usage))
  648. require.Equal(t, 1, repo.extraUpdateCalls)
  649. require.Equal(t, credits, repo.extraUpdates[100][openaiQuotaResetCreditsKey])
  650. require.Equal(t, 0.0, repo.extraUpdates[100]["codex_5h_used_percent"])
  651. require.Equal(t, 0.0, repo.extraUpdates[100]["codex_7d_used_percent"])
  652. }
  653. // TestResetCreditGetByIDError_FailsClosed 验证守卫「失败关闭」语义:
  654. // 当守卫的 GetByID 发生瞬时错误时,ResetCredit 必须立即返回该错误,
  655. // 不得旁路进入 prepareUpstreamCall(否则影子账号会借 resolve 路径操作母账号)。
  656. //
  657. // 区分方法:privacyClientFactory/tokenProvider 留 nil;
  658. // - 旁路路径:prepareUpstreamCall 配置检查先命中,报 "not configured"
  659. // - 守卫正确关闭:报 "account not found"(来自守卫的 infraerrors)
  660. func TestResetCreditGetByIDError_FailsClosed(t *testing.T) {
  661. // 空 map:GetByID(200) 返回 "account 200 not found"
  662. repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{}}
  663. // tokenProvider / privacyClientFactory 故意为 nil:
  664. // 若代码泄漏到 prepareUpstreamCall,会因配置检查而报 "not configured" 而非 "account not found"。
  665. svc := &OpenAIQuotaService{accountRepo: repo}
  666. _, err := svc.ResetCredit(context.Background(), 200)
  667. require.Error(t, err, "GetByID error must propagate; got nil")
  668. require.NotContains(t, err.Error(), "not configured",
  669. "error reached prepareUpstreamCall config-check — guard did not fail-closed; got: %v", err)
  670. }
  671. // TestQueryUsageShadowResolve_EndToEnd 是端到端补充:通过 httptest 服务真实 /wham/usage
  672. // 路径,验证影子账号的 QueryUsage 能成功拿到服务器响应(header 由母账号注入)。
  673. func TestQueryUsageShadowResolve_EndToEnd(t *testing.T) {
  674. ctx := context.Background()
  675. pid := int64(100)
  676. shadow := &Account{
  677. ID: 200, ParentAccountID: &pid,
  678. Platform: PlatformOpenAI, Type: AccountTypeOAuth,
  679. Status: StatusActive, QuotaDimension: QuotaDimensionSpark,
  680. }
  681. parent := &Account{
  682. ID: 100, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive,
  683. Credentials: map[string]any{"chatgpt_account_id": "org-e2e-parent"},
  684. }
  685. repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{200: shadow, 100: parent}}
  686. tokenCache := &stubQuotaTokenCache{tokens: map[string]string{
  687. OpenAITokenCacheKey(parent): "fake-token-e2e",
  688. }}
  689. tokenProvider := NewOpenAITokenProvider(repo, tokenCache, nil)
  690. // httptest server 记录收到的 chatgpt-account-id header,返回空 usage JSON
  691. var capturedAccountID string
  692. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  693. capturedAccountID = r.Header.Get("chatgpt-account-id")
  694. w.Header().Set("content-type", "application/json")
  695. _ = json.NewEncoder(w).Encode(OpenAIQuotaUsage{})
  696. }))
  697. defer srv.Close()
  698. svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
  699. usage, err := svc.QueryUsage(ctx, 200)
  700. require.NoError(t, err)
  701. require.NotNil(t, usage)
  702. require.Equal(t, "org-e2e-parent", capturedAccountID,
  703. "upstream should receive parent's chatgpt-account-id; got: %s", capturedAccountID)
  704. }