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