openai_ws_ratelimit_signal_test.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580
  1. package service
  2. import (
  3. "context"
  4. "io"
  5. "net/http"
  6. "net/http/httptest"
  7. "strconv"
  8. "strings"
  9. "testing"
  10. "time"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  12. coderws "github.com/coder/websocket"
  13. "github.com/gin-gonic/gin"
  14. "github.com/gorilla/websocket"
  15. "github.com/stretchr/testify/require"
  16. )
  17. type openAIWSRateLimitSignalRepo struct {
  18. stubOpenAIAccountRepo
  19. rateLimitCalls []time.Time
  20. updateExtra []map[string]any
  21. }
  22. type openAICodexSnapshotAsyncRepo struct {
  23. stubOpenAIAccountRepo
  24. updateExtraCh chan map[string]any
  25. rateLimitCh chan time.Time
  26. }
  27. type openAICodexExtraListRepo struct {
  28. stubOpenAIAccountRepo
  29. rateLimitCh chan time.Time
  30. }
  31. func (r *openAIWSRateLimitSignalRepo) SetRateLimited(_ context.Context, _ int64, resetAt time.Time) error {
  32. r.rateLimitCalls = append(r.rateLimitCalls, resetAt)
  33. return nil
  34. }
  35. func (r *openAIWSRateLimitSignalRepo) UpdateExtra(_ context.Context, _ int64, updates map[string]any) error {
  36. copied := make(map[string]any, len(updates))
  37. for k, v := range updates {
  38. copied[k] = v
  39. }
  40. r.updateExtra = append(r.updateExtra, copied)
  41. return nil
  42. }
  43. func (r *openAICodexSnapshotAsyncRepo) SetRateLimited(_ context.Context, _ int64, resetAt time.Time) error {
  44. if r.rateLimitCh != nil {
  45. r.rateLimitCh <- resetAt
  46. }
  47. return nil
  48. }
  49. func (r *openAICodexSnapshotAsyncRepo) UpdateExtra(_ context.Context, _ int64, updates map[string]any) error {
  50. if r.updateExtraCh != nil {
  51. copied := make(map[string]any, len(updates))
  52. for k, v := range updates {
  53. copied[k] = v
  54. }
  55. r.updateExtraCh <- copied
  56. }
  57. return nil
  58. }
  59. func (r *openAICodexExtraListRepo) SetRateLimited(_ context.Context, _ int64, resetAt time.Time) error {
  60. if r.rateLimitCh != nil {
  61. r.rateLimitCh <- resetAt
  62. }
  63. return nil
  64. }
  65. func (r *openAICodexExtraListRepo) ListWithFilters(_ context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) {
  66. _ = platform
  67. _ = accountType
  68. _ = status
  69. _ = search
  70. _ = groupID
  71. _ = privacyMode
  72. return r.accounts, &pagination.PaginationResult{Total: int64(len(r.accounts)), Page: params.Page, PageSize: params.PageSize}, nil
  73. }
  74. func TestOpenAIGatewayService_Forward_WSv2ErrorEventUsageLimitPersistsRateLimit(t *testing.T) {
  75. gin.SetMode(gin.TestMode)
  76. resetAt := time.Now().Add(2 * time.Hour).Unix()
  77. upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
  78. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  79. conn, err := upgrader.Upgrade(w, r, nil)
  80. if err != nil {
  81. t.Errorf("upgrade websocket failed: %v", err)
  82. return
  83. }
  84. defer func() { _ = conn.Close() }()
  85. var req map[string]any
  86. if err := conn.ReadJSON(&req); err != nil {
  87. t.Errorf("read ws request failed: %v", err)
  88. return
  89. }
  90. _ = conn.WriteJSON(map[string]any{
  91. "type": "error",
  92. "error": map[string]any{
  93. "code": "rate_limit_exceeded",
  94. "type": "usage_limit_reached",
  95. "message": "The usage limit has been reached",
  96. "resets_at": resetAt,
  97. },
  98. })
  99. }))
  100. defer wsServer.Close()
  101. rec := httptest.NewRecorder()
  102. c, _ := gin.CreateTestContext(rec)
  103. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  104. c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
  105. upstream := &httpUpstreamRecorder{
  106. resp: &http.Response{
  107. StatusCode: http.StatusOK,
  108. Header: http.Header{"Content-Type": []string{"application/json"}},
  109. Body: io.NopCloser(strings.NewReader(`{"id":"resp_http_should_not_run"}`)),
  110. },
  111. }
  112. cfg := newOpenAIWSV2TestConfig()
  113. cfg.Security.URLAllowlist.Enabled = false
  114. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  115. account := Account{
  116. ID: 501,
  117. Name: "openai-ws-rate-limit-event",
  118. Platform: PlatformOpenAI,
  119. Type: AccountTypeAPIKey,
  120. Status: StatusActive,
  121. Schedulable: true,
  122. Concurrency: 1,
  123. Credentials: map[string]any{
  124. "api_key": "sk-test",
  125. "base_url": wsServer.URL,
  126. },
  127. Extra: map[string]any{
  128. "responses_websockets_v2_enabled": true,
  129. },
  130. }
  131. repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}}
  132. rateSvc := &RateLimitService{accountRepo: repo}
  133. svc := &OpenAIGatewayService{
  134. accountRepo: repo,
  135. rateLimitService: rateSvc,
  136. httpUpstream: upstream,
  137. cache: &stubGatewayCache{},
  138. cfg: cfg,
  139. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  140. toolCorrector: NewCodexToolCorrector(),
  141. }
  142. body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
  143. result, err := svc.Forward(context.Background(), c, &account, body)
  144. require.Error(t, err)
  145. require.Nil(t, result)
  146. require.Equal(t, http.StatusTooManyRequests, rec.Code)
  147. require.Nil(t, upstream.lastReq, "WS 限流 error event 不应回退到同账号 HTTP")
  148. require.Len(t, repo.rateLimitCalls, 1)
  149. require.WithinDuration(t, time.Unix(resetAt, 0), repo.rateLimitCalls[0], 2*time.Second)
  150. }
  151. func TestOpenAIGatewayService_Forward_WSv2Handshake429PersistsRateLimit(t *testing.T) {
  152. gin.SetMode(gin.TestMode)
  153. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  154. w.Header().Set("x-codex-primary-used-percent", "100")
  155. w.Header().Set("x-codex-primary-reset-after-seconds", "7200")
  156. w.Header().Set("x-codex-primary-window-minutes", "10080")
  157. w.Header().Set("x-codex-secondary-used-percent", "3")
  158. w.Header().Set("x-codex-secondary-reset-after-seconds", "1800")
  159. w.Header().Set("x-codex-secondary-window-minutes", "300")
  160. w.WriteHeader(http.StatusTooManyRequests)
  161. _, _ = w.Write([]byte(`{"error":{"type":"rate_limit_exceeded","message":"rate limited"}}`))
  162. }))
  163. defer server.Close()
  164. rec := httptest.NewRecorder()
  165. c, _ := gin.CreateTestContext(rec)
  166. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  167. c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
  168. upstream := &httpUpstreamRecorder{
  169. resp: &http.Response{
  170. StatusCode: http.StatusOK,
  171. Header: http.Header{"Content-Type": []string{"application/json"}},
  172. Body: io.NopCloser(strings.NewReader(`{"id":"resp_http_should_not_run"}`)),
  173. },
  174. }
  175. cfg := newOpenAIWSV2TestConfig()
  176. cfg.Security.URLAllowlist.Enabled = false
  177. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  178. account := Account{
  179. ID: 502,
  180. Name: "openai-ws-rate-limit-handshake",
  181. Platform: PlatformOpenAI,
  182. Type: AccountTypeAPIKey,
  183. Status: StatusActive,
  184. Schedulable: true,
  185. Concurrency: 1,
  186. Credentials: map[string]any{
  187. "api_key": "sk-test",
  188. "base_url": server.URL,
  189. },
  190. Extra: map[string]any{
  191. "responses_websockets_v2_enabled": true,
  192. },
  193. }
  194. repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}}
  195. rateSvc := &RateLimitService{accountRepo: repo}
  196. svc := &OpenAIGatewayService{
  197. accountRepo: repo,
  198. rateLimitService: rateSvc,
  199. httpUpstream: upstream,
  200. cache: &stubGatewayCache{},
  201. cfg: cfg,
  202. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  203. toolCorrector: NewCodexToolCorrector(),
  204. }
  205. body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
  206. result, err := svc.Forward(context.Background(), c, &account, body)
  207. require.Error(t, err)
  208. require.Nil(t, result)
  209. require.Equal(t, http.StatusTooManyRequests, rec.Code)
  210. require.Nil(t, upstream.lastReq, "WS 握手 429 不应回退到同账号 HTTP")
  211. require.Len(t, repo.rateLimitCalls, 1)
  212. require.NotEmpty(t, repo.updateExtra, "握手 429 的 x-codex 头应立即落库")
  213. require.Contains(t, repo.updateExtra[0], "codex_usage_updated_at")
  214. }
  215. func TestOpenAIGatewayService_Forward_WSv2Handshake502RecordsModelTransient(t *testing.T) {
  216. gin.SetMode(gin.TestMode)
  217. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  218. w.Header().Set("x-request-id", "req-ws-502")
  219. w.WriteHeader(http.StatusBadGateway)
  220. _, _ = w.Write([]byte(`{"error":{"type":"server_error","message":"bad gateway"}}`))
  221. }))
  222. defer server.Close()
  223. cfg := newOpenAIWSV2TestConfig()
  224. cfg.Security.URLAllowlist.Enabled = false
  225. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  226. account := Account{
  227. ID: 504,
  228. Platform: PlatformOpenAI,
  229. Type: AccountTypeAPIKey,
  230. Status: StatusActive,
  231. Schedulable: true,
  232. Concurrency: 1,
  233. Credentials: map[string]any{"api_key": "sk-test", "base_url": server.URL},
  234. Extra: map[string]any{"responses_websockets_v2_enabled": true},
  235. }
  236. svc := &OpenAIGatewayService{
  237. cfg: cfg,
  238. rateLimitService: NewRateLimitService(transientCooldownAccountRepo{}, nil, cfg, nil, nil),
  239. httpUpstream: &httpUpstreamRecorder{},
  240. cache: &stubGatewayCache{},
  241. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  242. toolCorrector: NewCodexToolCorrector(),
  243. }
  244. body := []byte(`{"model":"gpt-5.5","stream":false,"input":"hello"}`)
  245. for range 2 {
  246. rec := httptest.NewRecorder()
  247. c, _ := gin.CreateTestContext(rec)
  248. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  249. c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
  250. result, err := svc.Forward(context.Background(), c, &account, body)
  251. require.Error(t, err)
  252. require.Nil(t, result)
  253. }
  254. require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(&account, "gpt-5.5"))
  255. }
  256. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_ErrorEventUsageLimitPersistsRateLimit(t *testing.T) {
  257. gin.SetMode(gin.TestMode)
  258. cfg := newOpenAIWSV2TestConfig()
  259. cfg.Security.URLAllowlist.Enabled = false
  260. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  261. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  262. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  263. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  264. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  265. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  266. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  267. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  268. resetAt := time.Now().Add(90 * time.Minute).Unix()
  269. captureConn := &openAIWSCaptureConn{
  270. events: [][]byte{
  271. []byte(`{"type":"error","error":{"code":"rate_limit_exceeded","type":"usage_limit_reached","message":"The usage limit has been reached","resets_at":PLACEHOLDER}}`),
  272. },
  273. }
  274. captureConn.events[0] = []byte(strings.ReplaceAll(string(captureConn.events[0]), "PLACEHOLDER", strconv.FormatInt(resetAt, 10)))
  275. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  276. pool := newOpenAIWSConnPool(cfg)
  277. pool.setClientDialerForTest(captureDialer)
  278. account := Account{
  279. ID: 503,
  280. Name: "openai-ingress-rate-limit",
  281. Platform: PlatformOpenAI,
  282. Type: AccountTypeAPIKey,
  283. Status: StatusActive,
  284. Schedulable: true,
  285. Concurrency: 1,
  286. Credentials: map[string]any{
  287. "api_key": "sk-test",
  288. },
  289. Extra: map[string]any{
  290. "responses_websockets_v2_enabled": true,
  291. },
  292. }
  293. repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}}
  294. rateSvc := &RateLimitService{accountRepo: repo}
  295. svc := &OpenAIGatewayService{
  296. accountRepo: repo,
  297. rateLimitService: rateSvc,
  298. httpUpstream: &httpUpstreamRecorder{},
  299. cache: &stubGatewayCache{},
  300. cfg: cfg,
  301. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  302. toolCorrector: NewCodexToolCorrector(),
  303. openaiWSPool: pool,
  304. }
  305. serverErrCh := make(chan error, 1)
  306. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  307. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  308. if err != nil {
  309. serverErrCh <- err
  310. return
  311. }
  312. defer func() { _ = conn.CloseNow() }()
  313. rec := httptest.NewRecorder()
  314. ginCtx, _ := gin.CreateTestContext(rec)
  315. req := r.Clone(r.Context())
  316. req.Header = req.Header.Clone()
  317. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  318. ginCtx.Request = req
  319. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  320. msgType, firstMessage, readErr := conn.Read(readCtx)
  321. cancel()
  322. if readErr != nil {
  323. serverErrCh <- readErr
  324. return
  325. }
  326. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  327. serverErrCh <- io.ErrUnexpectedEOF
  328. return
  329. }
  330. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, &account, "sk-test", firstMessage, nil)
  331. }))
  332. defer wsServer.Close()
  333. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  334. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  335. cancelDial()
  336. require.NoError(t, err)
  337. defer func() { _ = clientConn.CloseNow() }()
  338. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  339. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`))
  340. cancelWrite()
  341. require.NoError(t, err)
  342. select {
  343. case serverErr := <-serverErrCh:
  344. require.Error(t, serverErr)
  345. var failoverErr *UpstreamFailoverError
  346. require.ErrorAs(t, serverErr, &failoverErr)
  347. require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
  348. require.Len(t, repo.rateLimitCalls, 1)
  349. require.WithinDuration(t, time.Unix(resetAt, 0), repo.rateLimitCalls[0], 2*time.Second)
  350. case <-time.After(5 * time.Second):
  351. t.Fatal("等待 ingress websocket 结束超时")
  352. }
  353. }
  354. func TestOpenAIGatewayService_UpdateCodexUsageSnapshot_ExhaustedSnapshotDoesNotSetRateLimit(t *testing.T) {
  355. repo := &openAICodexSnapshotAsyncRepo{
  356. updateExtraCh: make(chan map[string]any, 1),
  357. rateLimitCh: make(chan time.Time, 1),
  358. }
  359. svc := &OpenAIGatewayService{accountRepo: repo}
  360. snapshot := &OpenAICodexUsageSnapshot{
  361. PrimaryUsedPercent: ptrFloat64WS(100),
  362. PrimaryResetAfterSeconds: ptrIntWS(3600),
  363. PrimaryWindowMinutes: ptrIntWS(10080),
  364. SecondaryUsedPercent: ptrFloat64WS(12),
  365. SecondaryResetAfterSeconds: ptrIntWS(1200),
  366. SecondaryWindowMinutes: ptrIntWS(300),
  367. }
  368. svc.updateCodexUsageSnapshot(context.Background(), 601, snapshot)
  369. select {
  370. case updates := <-repo.updateExtraCh:
  371. require.Equal(t, 100.0, updates["codex_7d_used_percent"])
  372. case <-time.After(2 * time.Second):
  373. t.Fatal("等待 codex 快照落库超时")
  374. }
  375. select {
  376. case resetAt := <-repo.rateLimitCh:
  377. t.Fatalf("不应因仅写入快照而生成运行时限流时间: %v", resetAt)
  378. case <-time.After(2 * time.Second):
  379. }
  380. }
  381. func TestOpenAIGatewayService_UpdateCodexUsageSnapshot_NonExhaustedSnapshotDoesNotSetRateLimit(t *testing.T) {
  382. repo := &openAICodexSnapshotAsyncRepo{
  383. updateExtraCh: make(chan map[string]any, 1),
  384. rateLimitCh: make(chan time.Time, 1),
  385. }
  386. svc := &OpenAIGatewayService{accountRepo: repo}
  387. snapshot := &OpenAICodexUsageSnapshot{
  388. PrimaryUsedPercent: ptrFloat64WS(94),
  389. PrimaryResetAfterSeconds: ptrIntWS(3600),
  390. PrimaryWindowMinutes: ptrIntWS(10080),
  391. SecondaryUsedPercent: ptrFloat64WS(22),
  392. SecondaryResetAfterSeconds: ptrIntWS(1200),
  393. SecondaryWindowMinutes: ptrIntWS(300),
  394. }
  395. svc.updateCodexUsageSnapshot(context.Background(), 602, snapshot)
  396. select {
  397. case <-repo.updateExtraCh:
  398. case <-time.After(2 * time.Second):
  399. t.Fatal("等待 codex 快照落库超时")
  400. }
  401. select {
  402. case resetAt := <-repo.rateLimitCh:
  403. t.Fatalf("不应写入运行时限流时间: %v", resetAt)
  404. case <-time.After(200 * time.Millisecond):
  405. }
  406. }
  407. func TestOpenAIGatewayService_UpdateCodexUsageSnapshot_ThrottlesExtraWrites(t *testing.T) {
  408. repo := &openAICodexSnapshotAsyncRepo{
  409. updateExtraCh: make(chan map[string]any, 2),
  410. }
  411. svc := &OpenAIGatewayService{
  412. accountRepo: repo,
  413. codexSnapshotThrottle: newAccountWriteThrottle(time.Hour),
  414. }
  415. snapshot := &OpenAICodexUsageSnapshot{
  416. PrimaryUsedPercent: ptrFloat64WS(94),
  417. PrimaryResetAfterSeconds: ptrIntWS(3600),
  418. PrimaryWindowMinutes: ptrIntWS(10080),
  419. SecondaryUsedPercent: ptrFloat64WS(22),
  420. SecondaryResetAfterSeconds: ptrIntWS(1200),
  421. SecondaryWindowMinutes: ptrIntWS(300),
  422. }
  423. svc.updateCodexUsageSnapshot(context.Background(), 777, snapshot)
  424. svc.updateCodexUsageSnapshot(context.Background(), 777, snapshot)
  425. select {
  426. case <-repo.updateExtraCh:
  427. case <-time.After(2 * time.Second):
  428. t.Fatal("等待第一次 codex 快照落库超时")
  429. }
  430. select {
  431. case updates := <-repo.updateExtraCh:
  432. t.Fatalf("unexpected second codex snapshot write: %v", updates)
  433. case <-time.After(200 * time.Millisecond):
  434. }
  435. }
  436. func ptrFloat64WS(v float64) *float64 { return &v }
  437. func ptrIntWS(v int) *int { return &v }
  438. func TestOpenAIGatewayService_GetSchedulableAccount_ExhaustedCodexExtraDoesNotSetRateLimit(t *testing.T) {
  439. resetAt := time.Now().Add(6 * 24 * time.Hour)
  440. account := Account{
  441. ID: 701,
  442. Platform: PlatformOpenAI,
  443. Type: AccountTypeOAuth,
  444. Status: StatusActive,
  445. Schedulable: true,
  446. Concurrency: 1,
  447. Extra: map[string]any{
  448. "codex_7d_used_percent": 100.0,
  449. "codex_7d_reset_at": resetAt.UTC().Format(time.RFC3339),
  450. },
  451. }
  452. repo := &openAICodexExtraListRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}, rateLimitCh: make(chan time.Time, 1)}
  453. svc := &OpenAIGatewayService{accountRepo: repo}
  454. fresh, err := svc.getSchedulableAccount(context.Background(), account.ID)
  455. require.NoError(t, err)
  456. require.NotNil(t, fresh)
  457. require.Nil(t, fresh.RateLimitResetAt)
  458. select {
  459. case persisted := <-repo.rateLimitCh:
  460. t.Fatalf("不应将已耗尽的 codex extra 提升为运行时限流状态: %v", persisted)
  461. case <-time.After(2 * time.Second):
  462. }
  463. }
  464. func TestAdminService_ListAccounts_ExhaustedCodexExtraDoesNotSetRateLimit(t *testing.T) {
  465. resetAt := time.Now().Add(4 * 24 * time.Hour)
  466. repo := &openAICodexExtraListRepo{
  467. stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{{
  468. ID: 702,
  469. Platform: PlatformOpenAI,
  470. Type: AccountTypeOAuth,
  471. Status: StatusActive,
  472. Schedulable: true,
  473. Concurrency: 1,
  474. Extra: map[string]any{
  475. "codex_7d_used_percent": 100.0,
  476. "codex_7d_reset_at": resetAt.UTC().Format(time.RFC3339),
  477. },
  478. }}},
  479. rateLimitCh: make(chan time.Time, 1),
  480. }
  481. svc := &adminServiceImpl{accountRepo: repo}
  482. accounts, total, err := svc.ListAccounts(context.Background(), 1, 20, PlatformOpenAI, AccountTypeOAuth, "", "", 0, "", "", "")
  483. require.NoError(t, err)
  484. require.Equal(t, int64(1), total)
  485. require.Len(t, accounts, 1)
  486. require.Nil(t, accounts[0].RateLimitResetAt)
  487. select {
  488. case persisted := <-repo.rateLimitCh:
  489. t.Fatalf("不应在账号列表查询时将 codex extra 持久化为运行时限流状态: %v", persisted)
  490. case <-time.After(2 * time.Second):
  491. }
  492. }
  493. func TestOpenAIWSErrorHTTPStatusFromRaw_UsageLimitReachedIs429(t *testing.T) {
  494. require.Equal(t, http.StatusTooManyRequests, openAIWSErrorHTTPStatusFromRaw("", "usage_limit_reached"))
  495. require.Equal(t, http.StatusTooManyRequests, openAIWSErrorHTTPStatusFromRaw("rate_limit_exceeded", ""))
  496. }
  497. func TestOpenAIWSRateLimitFailoverError_OAuthHonorsRetryAfterCooldown(t *testing.T) {
  498. svc := &OpenAIGatewayService{}
  499. headers := http.Header{"Retry-After": []string{"30"}}
  500. body := []byte(`{"error":{"type":"rate_limit_error","message":"limited"}}`)
  501. oauthErr := svc.newOpenAIWSRateLimitFailoverError(&Account{
  502. ID: 904,
  503. Platform: PlatformOpenAI,
  504. Type: AccountTypeOAuth,
  505. }, headers, body, "limited")
  506. require.False(t, oauthErr.RetryableOnSameAccount)
  507. require.True(t, oauthErr.SameAccountRetryDeadline.IsZero())
  508. require.Zero(t, oauthErr.SameAccountRetryDelay)
  509. require.Equal(t, body, oauthErr.ResponseBody)
  510. require.Equal(t, "30", oauthErr.ResponseHeaders.Get("Retry-After"))
  511. apiKeyErr := svc.newOpenAIWSRateLimitFailoverError(&Account{
  512. ID: 905,
  513. Platform: PlatformOpenAI,
  514. Type: AccountTypeAPIKey,
  515. }, headers, body, "limited")
  516. require.False(t, apiKeyErr.RetryableOnSameAccount)
  517. require.True(t, apiKeyErr.SameAccountRetryDeadline.IsZero())
  518. require.Zero(t, apiKeyErr.SameAccountRetryDelay)
  519. }