rateLimitMocking.ts 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. /**
  2. * Facade for rate limit header processing
  3. * This isolates mock logic from production code
  4. */
  5. import { APIError } from '@anthropic-ai/sdk'
  6. import {
  7. applyMockHeaders,
  8. checkMockFastModeRateLimit,
  9. getMockHeaderless429Message,
  10. getMockHeaders,
  11. isMockFastModeRateLimitScenario,
  12. shouldProcessMockLimits,
  13. } from './mockRateLimits.js'
  14. /**
  15. * Process headers, applying mocks if /mock-limits command is active
  16. */
  17. export function processRateLimitHeaders(
  18. headers: globalThis.Headers,
  19. ): globalThis.Headers {
  20. // Only apply mocks for Ant employees using /mock-limits command
  21. if (shouldProcessMockLimits()) {
  22. return applyMockHeaders(headers)
  23. }
  24. return headers
  25. }
  26. /**
  27. * Check if we should process rate limits (either real subscriber or /mock-limits command)
  28. */
  29. export function shouldProcessRateLimits(isSubscriber: boolean): boolean {
  30. return isSubscriber || shouldProcessMockLimits()
  31. }
  32. /**
  33. * Check if mock rate limits should throw a 429 error
  34. * Returns the error to throw, or null if no error should be thrown
  35. * @param currentModel The model being used for the current request
  36. * @param isFastModeActive Whether fast mode is currently active (for fast-mode-only mocks)
  37. */
  38. export function checkMockRateLimitError(
  39. currentModel: string,
  40. isFastModeActive?: boolean,
  41. ): APIError | null {
  42. if (!shouldProcessMockLimits()) {
  43. return null
  44. }
  45. const headerlessMessage = getMockHeaderless429Message()
  46. if (headerlessMessage) {
  47. return new APIError(
  48. 429,
  49. { error: { type: 'rate_limit_error', message: headerlessMessage } },
  50. headerlessMessage,
  51. // eslint-disable-next-line eslint-plugin-n/no-unsupported-features/node-builtins
  52. new globalThis.Headers(),
  53. )
  54. }
  55. const mockHeaders = getMockHeaders()
  56. if (!mockHeaders) {
  57. return null
  58. }
  59. // Check if we should throw a 429 error
  60. // Only throw if:
  61. // 1. Status is rejected AND
  62. // 2. Either no overage headers OR overage is also rejected
  63. // 3. For Opus-specific limits, only throw if actually using an Opus model
  64. const status = mockHeaders['anthropic-ratelimit-unified-status']
  65. const overageStatus =
  66. mockHeaders['anthropic-ratelimit-unified-overage-status']
  67. const rateLimitType =
  68. mockHeaders['anthropic-ratelimit-unified-representative-claim']
  69. // Check if this is an Opus-specific rate limit
  70. const isOpusLimit = rateLimitType === 'seven_day_opus'
  71. // Check if current model is an Opus model (handles all variants including aliases)
  72. const isUsingOpus = currentModel.includes('opus')
  73. // For Opus limits, only throw 429 if actually using Opus
  74. // This simulates the real API behavior where fallback to Sonnet succeeds
  75. if (isOpusLimit && !isUsingOpus) {
  76. return null
  77. }
  78. // Check for mock fast mode rate limits (handles expiry, countdown, etc.)
  79. if (isMockFastModeRateLimitScenario()) {
  80. const fastModeHeaders = checkMockFastModeRateLimit(isFastModeActive)
  81. if (fastModeHeaders === null) {
  82. return null
  83. }
  84. // Create a mock 429 error with the fast mode headers
  85. const error = new APIError(
  86. 429,
  87. { error: { type: 'rate_limit_error', message: '超出频率限制' } },
  88. '超出频率限制',
  89. // eslint-disable-next-line eslint-plugin-n/no-unsupported-features/node-builtins
  90. new globalThis.Headers(
  91. Object.entries(fastModeHeaders).filter(([_, v]) => v !== undefined) as [
  92. string,
  93. string,
  94. ][],
  95. ),
  96. )
  97. return error
  98. }
  99. const shouldThrow429 =
  100. status === 'rejected' && (!overageStatus || overageStatus === 'rejected')
  101. if (shouldThrow429) {
  102. // Create a mock 429 error with the appropriate headers
  103. const error = new APIError(
  104. 429,
  105. { error: { type: 'rate_limit_error', message: '超出频率限制' } },
  106. '超出频率限制',
  107. // eslint-disable-next-line eslint-plugin-n/no-unsupported-features/node-builtins
  108. new globalThis.Headers(
  109. Object.entries(mockHeaders).filter(([_, v]) => v !== undefined) as [
  110. string,
  111. string,
  112. ][],
  113. ),
  114. )
  115. return error
  116. }
  117. return null
  118. }
  119. /**
  120. * Check if this is a mock 429 error that shouldn't be retried
  121. */
  122. export function isMockRateLimitError(error: APIError): boolean {
  123. return shouldProcessMockLimits() && error.status === 429
  124. }
  125. /**
  126. * Check if /mock-limits command is currently active (for UI purposes)
  127. */
  128. export { shouldProcessMockLimits }