mirror of
https://github.com/simstudioai/sim.git
synced 2026-09-24 15:45:35 +08:00
feat(tools): queue hosted-key tool calls instead of failing with 429 (#4416)
* Add queueing for hosted keys * feat(rate-limiter): FIFO queue for hosted-key per-workspace fairness Replace the per-call distributed lock with a Redis-backed FIFO queue so callers within a workspace get strict ordering instead of racing the bucket. Adds heartbeat-based crash recovery and dead-head reaping in a single Lua script. Bumps Exa search hosted RPM from 5 to 60. * fix(rate-limiter): bound hosted-key queue wait to execution budget; fix heartbeat + telemetry Tie the per-workspace hosted-key queue wait to the surrounding execution budget instead of a flat 5-minute cap. acquireKey now accepts the execution AbortSignal (threaded from ExecutionContext): when present, the wait is bounded by the run's actual plan timeout / cancellation, with the enterprise async ceiling as a backstop; when absent it falls back to MAX_QUEUE_WAIT_MS. This lets long-running async (Trigger.dev) runs use their full budget while no longer letting a single queued call burn a short sync run's entire budget. Also addresses Greptile review: - P1: share one lastHeartbeatAt across all wait phases and cap every sleep to HEARTBEAT_REFRESH_INTERVAL_MS so a long low-RPM retryAfterMs can no longer let the head's heartbeat lapse mid-wait and break FIFO ordering. - P2: derive hostedKeyQueueWaited telemetry reason from the actual bottleneck (queue_position / dimension / actor_requests) instead of hardcoding it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * feat(rate-limiter): make hosted-key queue waits abort-interruptible Replace the plain capped sleeps in the queue-head and bucket-capacity wait loops with an interruptibleSleep that resolves early when the execution AbortSignal fires (timeout or cancellation), cleaning up its own timer and listener. Previously a cancelled/timed-out run could overshoot by up to the heartbeat cap (~10s) before the loop re-checked its budget; now it wakes within a tick. The cap remains for heartbeat renewal. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Theodore Li
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
7218185644
commit
4fa7e74c41
@@ -5,8 +5,13 @@ import type {
|
||||
TokenStatus,
|
||||
} from '@/lib/core/rate-limiter/storage'
|
||||
import { HostedKeyRateLimiter } from './hosted-key-rate-limiter'
|
||||
import { HEARTBEAT_REFRESH_INTERVAL_MS, type HostedKeyQueue } from './queue'
|
||||
import type { CustomRateLimit, PerRequestRateLimit } from './types'
|
||||
|
||||
/** Force the queue wait to give up on the first iteration by reporting a retry time
|
||||
* larger than the 5-minute MAX_QUEUE_WAIT_MS cap. */
|
||||
const RETRY_PAST_CAP_MS = 6 * 60 * 1000
|
||||
|
||||
interface MockAdapter {
|
||||
consumeTokens: Mock
|
||||
getTokenStatus: Mock
|
||||
@@ -19,10 +24,30 @@ const createMockAdapter = (): MockAdapter => ({
|
||||
resetBucket: vi.fn(),
|
||||
})
|
||||
|
||||
interface MockQueue {
|
||||
enqueue: Mock
|
||||
checkHead: Mock
|
||||
refreshHeartbeat: Mock
|
||||
dequeue: Mock
|
||||
}
|
||||
|
||||
/** Stub queue that defaults to "you're at the head, no waiting" — i.e. acts as if the
|
||||
* queue is empty or Redis is unavailable. Tests override per-call to simulate ordering. */
|
||||
const createMockQueue = (): MockQueue => {
|
||||
const queue: MockQueue = {
|
||||
enqueue: vi.fn().mockResolvedValue({ position: 0, enabled: true }),
|
||||
checkHead: vi.fn().mockResolvedValue('head'),
|
||||
refreshHeartbeat: vi.fn().mockResolvedValue(undefined),
|
||||
dequeue: vi.fn().mockResolvedValue(undefined),
|
||||
}
|
||||
return queue
|
||||
}
|
||||
|
||||
describe('HostedKeyRateLimiter', () => {
|
||||
const testProvider = 'exa'
|
||||
const envKeyPrefix = 'EXA_API_KEY'
|
||||
let mockAdapter: MockAdapter
|
||||
let mockQueue: MockQueue
|
||||
let rateLimiter: HostedKeyRateLimiter
|
||||
let originalEnv: NodeJS.ProcessEnv
|
||||
|
||||
@@ -34,7 +59,11 @@ describe('HostedKeyRateLimiter', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
mockAdapter = createMockAdapter()
|
||||
rateLimiter = new HostedKeyRateLimiter(mockAdapter as RateLimitStorageAdapter)
|
||||
mockQueue = createMockQueue()
|
||||
rateLimiter = new HostedKeyRateLimiter(
|
||||
mockAdapter as RateLimitStorageAdapter,
|
||||
mockQueue as unknown as HostedKeyQueue
|
||||
)
|
||||
|
||||
originalEnv = { ...process.env }
|
||||
process.env.EXA_API_KEY_COUNT = '3'
|
||||
@@ -72,11 +101,12 @@ describe('HostedKeyRateLimiter', () => {
|
||||
expect(result.error).toContain('No hosted keys configured')
|
||||
})
|
||||
|
||||
it('should rate limit billing actor when they exceed their limit', async () => {
|
||||
it('should rate limit billing actor when wait exceeds the queue cap', async () => {
|
||||
// resetAt past the 5-minute cap forces the wait loop to bail immediately.
|
||||
const rateLimitedResult: ConsumeResult = {
|
||||
allowed: false,
|
||||
tokensRemaining: 0,
|
||||
resetAt: new Date(Date.now() + 30000),
|
||||
resetAt: new Date(Date.now() + RETRY_PAST_CAP_MS),
|
||||
}
|
||||
mockAdapter.consumeTokens.mockResolvedValue(rateLimitedResult)
|
||||
|
||||
@@ -93,6 +123,33 @@ describe('HostedKeyRateLimiter', () => {
|
||||
expect(result.error).toContain('Rate limit exceeded')
|
||||
})
|
||||
|
||||
it('should wait for capacity then succeed when bucket refills within the cap', async () => {
|
||||
// First call: bucket empty, refills in 100ms (well under cap).
|
||||
// Second call: bucket has capacity, consumed.
|
||||
const blocked: ConsumeResult = {
|
||||
allowed: false,
|
||||
tokensRemaining: 0,
|
||||
resetAt: new Date(Date.now() + 100),
|
||||
}
|
||||
const allowed: ConsumeResult = {
|
||||
allowed: true,
|
||||
tokensRemaining: 9,
|
||||
resetAt: new Date(Date.now() + 60000),
|
||||
}
|
||||
mockAdapter.consumeTokens.mockResolvedValueOnce(blocked).mockResolvedValueOnce(allowed)
|
||||
|
||||
const result = await rateLimiter.acquireKey(
|
||||
testProvider,
|
||||
envKeyPrefix,
|
||||
perRequestRateLimit,
|
||||
'workspace-wait'
|
||||
)
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
expect(result.key).toBe('test-key-1')
|
||||
expect(mockAdapter.consumeTokens).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
it('should allow billing actor within their rate limit', async () => {
|
||||
const allowedResult: ConsumeResult = {
|
||||
allowed: true,
|
||||
@@ -184,6 +241,261 @@ describe('HostedKeyRateLimiter', () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe('FIFO queue ordering', () => {
|
||||
const allowed: ConsumeResult = {
|
||||
allowed: true,
|
||||
tokensRemaining: 9,
|
||||
resetAt: new Date(Date.now() + 60000),
|
||||
}
|
||||
|
||||
it('enqueues every call onto the per-workspace+provider queue', async () => {
|
||||
mockAdapter.consumeTokens.mockResolvedValue(allowed)
|
||||
|
||||
await rateLimiter.acquireKey(testProvider, envKeyPrefix, perRequestRateLimit, 'workspace-1')
|
||||
|
||||
expect(mockQueue.enqueue).toHaveBeenCalledWith(
|
||||
testProvider,
|
||||
'workspace-1',
|
||||
expect.any(String)
|
||||
)
|
||||
})
|
||||
|
||||
it('always dequeues at the end of a successful acquisition', async () => {
|
||||
mockAdapter.consumeTokens.mockResolvedValue(allowed)
|
||||
|
||||
await rateLimiter.acquireKey(testProvider, envKeyPrefix, perRequestRateLimit, 'workspace-1')
|
||||
|
||||
expect(mockQueue.dequeue).toHaveBeenCalledWith(
|
||||
testProvider,
|
||||
'workspace-1',
|
||||
expect.any(String)
|
||||
)
|
||||
})
|
||||
|
||||
it('always dequeues even when the call fails (no keys configured)', async () => {
|
||||
mockAdapter.consumeTokens.mockResolvedValue(allowed)
|
||||
process.env.EXA_API_KEY_COUNT = '0'
|
||||
|
||||
await rateLimiter.acquireKey(testProvider, envKeyPrefix, perRequestRateLimit, 'workspace-1')
|
||||
|
||||
expect(mockQueue.dequeue).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('waits at the head of the queue before consuming from the bucket', async () => {
|
||||
mockAdapter.consumeTokens.mockResolvedValue(allowed)
|
||||
// First two checkHead calls say we're waiting; third says we're up.
|
||||
mockQueue.checkHead
|
||||
.mockResolvedValueOnce('waiting')
|
||||
.mockResolvedValueOnce('waiting')
|
||||
.mockResolvedValueOnce('head')
|
||||
|
||||
const result = await rateLimiter.acquireKey(
|
||||
testProvider,
|
||||
envKeyPrefix,
|
||||
perRequestRateLimit,
|
||||
'workspace-1'
|
||||
)
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
expect(mockQueue.checkHead).toHaveBeenCalledTimes(3)
|
||||
// Bucket is only consumed once we reach the head.
|
||||
expect(mockAdapter.consumeTokens).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
it('refreshes the heartbeat while waiting at the head of the queue', async () => {
|
||||
mockAdapter.consumeTokens.mockResolvedValue(allowed)
|
||||
|
||||
// We need the wait loop to iterate long enough for HEARTBEAT_REFRESH_INTERVAL_MS
|
||||
// to elapse. Use fake timers so we don't actually sleep.
|
||||
vi.useFakeTimers()
|
||||
try {
|
||||
// Queue says we're waiting forever — except after some time we're at head.
|
||||
mockQueue.checkHead.mockImplementation(async () => {
|
||||
// Advance past the heartbeat interval each time we poll, then say we're up.
|
||||
vi.advanceTimersByTime(15_000)
|
||||
return mockQueue.checkHead.mock.calls.length >= 2 ? 'head' : 'waiting'
|
||||
})
|
||||
|
||||
const promise = rateLimiter.acquireKey(
|
||||
testProvider,
|
||||
envKeyPrefix,
|
||||
perRequestRateLimit,
|
||||
'workspace-1'
|
||||
)
|
||||
// Drain pending timers so the sleep() resolves.
|
||||
await vi.runAllTimersAsync()
|
||||
await promise
|
||||
|
||||
expect(mockQueue.refreshHeartbeat).toHaveBeenCalled()
|
||||
} finally {
|
||||
vi.useRealTimers()
|
||||
}
|
||||
})
|
||||
|
||||
it('returns 429 when the queue wait exceeds the cap', async () => {
|
||||
mockAdapter.consumeTokens.mockResolvedValue(allowed)
|
||||
mockQueue.checkHead.mockResolvedValue('waiting')
|
||||
|
||||
vi.useFakeTimers()
|
||||
try {
|
||||
const promise = rateLimiter.acquireKey(
|
||||
testProvider,
|
||||
envKeyPrefix,
|
||||
perRequestRateLimit,
|
||||
'workspace-1'
|
||||
)
|
||||
// Burn past the 5-minute cap.
|
||||
await vi.advanceTimersByTimeAsync(6 * 60 * 1000)
|
||||
const result = await promise
|
||||
|
||||
expect(result.success).toBe(false)
|
||||
expect(result.billingActorRateLimited).toBe(true)
|
||||
} finally {
|
||||
vi.useRealTimers()
|
||||
}
|
||||
})
|
||||
|
||||
it('treats "missing" status as proceed (queue evicted, fall through to bucket race)', async () => {
|
||||
mockAdapter.consumeTokens.mockResolvedValue(allowed)
|
||||
mockQueue.checkHead.mockResolvedValueOnce('missing')
|
||||
|
||||
const result = await rateLimiter.acquireKey(
|
||||
testProvider,
|
||||
envKeyPrefix,
|
||||
perRequestRateLimit,
|
||||
'workspace-1'
|
||||
)
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe('execution-budget-bounded waits', () => {
|
||||
it('bails immediately when the execution signal is already aborted', async () => {
|
||||
const blocked: ConsumeResult = {
|
||||
allowed: false,
|
||||
tokensRemaining: 0,
|
||||
resetAt: new Date(Date.now() + 100),
|
||||
}
|
||||
mockAdapter.consumeTokens.mockResolvedValue(blocked)
|
||||
|
||||
const result = await rateLimiter.acquireKey(
|
||||
testProvider,
|
||||
envKeyPrefix,
|
||||
perRequestRateLimit,
|
||||
'workspace-1',
|
||||
AbortSignal.abort()
|
||||
)
|
||||
|
||||
expect(result.success).toBe(false)
|
||||
expect(result.billingActorRateLimited).toBe(true)
|
||||
// Aborted budget => give up on the first bucket check rather than looping.
|
||||
expect(mockAdapter.consumeTokens).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
it('stops waiting promptly when the signal aborts mid-sleep', async () => {
|
||||
// Bucket reports a long refill, so the wait sleeps up to the heartbeat cap (10s).
|
||||
// Aborting mid-sleep must wake the wait within a tick, not after the full interval.
|
||||
const blocked: ConsumeResult = {
|
||||
allowed: false,
|
||||
tokensRemaining: 0,
|
||||
resetAt: new Date(Date.now() + 10_000),
|
||||
}
|
||||
mockAdapter.consumeTokens.mockResolvedValue(blocked)
|
||||
|
||||
const controller = new AbortController()
|
||||
const start = Date.now()
|
||||
const promise = rateLimiter.acquireKey(
|
||||
testProvider,
|
||||
envKeyPrefix,
|
||||
perRequestRateLimit,
|
||||
'workspace-1',
|
||||
controller.signal
|
||||
)
|
||||
// Let the first bucket check run and the sleep begin, then abort.
|
||||
await new Promise((resolve) => setTimeout(resolve, 20))
|
||||
controller.abort()
|
||||
const result = await promise
|
||||
|
||||
expect(result.success).toBe(false)
|
||||
expect(result.billingActorRateLimited).toBe(true)
|
||||
// Resolved well before the 10s capped sleep would otherwise have elapsed.
|
||||
expect(Date.now() - start).toBeLessThan(2000)
|
||||
})
|
||||
|
||||
it('keeps waiting past the no-signal fallback cap while the signal is live', async () => {
|
||||
// A live (non-aborted) signal means the run still has budget, so the wait must not
|
||||
// 429 at the 5-minute MAX_QUEUE_WAIT_MS fallback. The bucket frees up after ~7 min.
|
||||
const blocked: ConsumeResult = {
|
||||
allowed: false,
|
||||
tokensRemaining: 0,
|
||||
resetAt: new Date(Date.now() + 10_000),
|
||||
}
|
||||
const allowedResult: ConsumeResult = {
|
||||
allowed: true,
|
||||
tokensRemaining: 9,
|
||||
resetAt: new Date(Date.now() + 60_000),
|
||||
}
|
||||
mockAdapter.consumeTokens.mockResolvedValue(blocked)
|
||||
|
||||
vi.useFakeTimers()
|
||||
try {
|
||||
const promise = rateLimiter.acquireKey(
|
||||
testProvider,
|
||||
envKeyPrefix,
|
||||
perRequestRateLimit,
|
||||
'workspace-1',
|
||||
new AbortController().signal
|
||||
)
|
||||
// Burn well past the 5-minute fallback cap — without a signal this would have 429'd.
|
||||
await vi.advanceTimersByTimeAsync(7 * 60 * 1000)
|
||||
mockAdapter.consumeTokens.mockResolvedValue(allowedResult)
|
||||
await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_INTERVAL_MS)
|
||||
const result = await promise
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
expect(result.key).toBe('test-key-1')
|
||||
} finally {
|
||||
vi.useRealTimers()
|
||||
}
|
||||
})
|
||||
|
||||
it('refreshes the heartbeat during a long low-RPM bucket wait', async () => {
|
||||
// Provider with a long refill (retryAfterMs >> heartbeat TTL). The sleep must be
|
||||
// capped so the heartbeat is renewed and the head is not reaped mid-wait.
|
||||
const blocked: ConsumeResult = {
|
||||
allowed: false,
|
||||
tokensRemaining: 0,
|
||||
resetAt: new Date(Date.now() + 60_000),
|
||||
}
|
||||
const allowedResult: ConsumeResult = {
|
||||
allowed: true,
|
||||
tokensRemaining: 0,
|
||||
resetAt: new Date(Date.now() + 60_000),
|
||||
}
|
||||
mockAdapter.consumeTokens.mockResolvedValue(blocked)
|
||||
|
||||
vi.useFakeTimers()
|
||||
try {
|
||||
const promise = rateLimiter.acquireKey(
|
||||
testProvider,
|
||||
envKeyPrefix,
|
||||
perRequestRateLimit,
|
||||
'workspace-1',
|
||||
new AbortController().signal
|
||||
)
|
||||
await vi.advanceTimersByTimeAsync(3 * HEARTBEAT_REFRESH_INTERVAL_MS)
|
||||
mockAdapter.consumeTokens.mockResolvedValue(allowedResult)
|
||||
await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_INTERVAL_MS)
|
||||
await promise
|
||||
|
||||
expect(mockQueue.refreshHeartbeat).toHaveBeenCalled()
|
||||
} finally {
|
||||
vi.useRealTimers()
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
describe('acquireKey with custom rate limit', () => {
|
||||
const customRateLimit: CustomRateLimit = {
|
||||
mode: 'custom',
|
||||
@@ -197,11 +509,11 @@ describe('HostedKeyRateLimiter', () => {
|
||||
],
|
||||
}
|
||||
|
||||
it('should enforce requestsPerMinute for custom mode', async () => {
|
||||
it('should enforce requestsPerMinute for custom mode when wait exceeds the cap', async () => {
|
||||
const rateLimitedResult: ConsumeResult = {
|
||||
allowed: false,
|
||||
tokensRemaining: 0,
|
||||
resetAt: new Date(Date.now() + 30000),
|
||||
resetAt: new Date(Date.now() + RETRY_PAST_CAP_MS),
|
||||
}
|
||||
mockAdapter.consumeTokens.mockResolvedValue(rateLimitedResult)
|
||||
|
||||
@@ -246,7 +558,7 @@ describe('HostedKeyRateLimiter', () => {
|
||||
expect(mockAdapter.getTokenStatus).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
it('should block request when a dimension is depleted', async () => {
|
||||
it('should block request when a dimension wait exceeds the cap', async () => {
|
||||
const allowedConsume: ConsumeResult = {
|
||||
allowed: true,
|
||||
tokensRemaining: 4,
|
||||
@@ -258,7 +570,7 @@ describe('HostedKeyRateLimiter', () => {
|
||||
tokensAvailable: 0,
|
||||
maxTokens: 2000,
|
||||
lastRefillAt: new Date(),
|
||||
nextRefillAt: new Date(Date.now() + 45000),
|
||||
nextRefillAt: new Date(Date.now() + RETRY_PAST_CAP_MS),
|
||||
}
|
||||
mockAdapter.getTokenStatus.mockResolvedValue(depleted)
|
||||
|
||||
@@ -274,6 +586,39 @@ describe('HostedKeyRateLimiter', () => {
|
||||
expect(result.error).toContain('tokens')
|
||||
})
|
||||
|
||||
it('should wait for dimension capacity then succeed when budget refills', async () => {
|
||||
const allowedConsume: ConsumeResult = {
|
||||
allowed: true,
|
||||
tokensRemaining: 4,
|
||||
resetAt: new Date(Date.now() + 60000),
|
||||
}
|
||||
mockAdapter.consumeTokens.mockResolvedValue(allowedConsume)
|
||||
|
||||
const depleted: TokenStatus = {
|
||||
tokensAvailable: 0,
|
||||
maxTokens: 2000,
|
||||
lastRefillAt: new Date(),
|
||||
nextRefillAt: new Date(Date.now() + 100),
|
||||
}
|
||||
const refilled: TokenStatus = {
|
||||
tokensAvailable: 500,
|
||||
maxTokens: 2000,
|
||||
lastRefillAt: new Date(),
|
||||
nextRefillAt: new Date(Date.now() + 60000),
|
||||
}
|
||||
mockAdapter.getTokenStatus.mockResolvedValueOnce(depleted).mockResolvedValueOnce(refilled)
|
||||
|
||||
const result = await rateLimiter.acquireKey(
|
||||
testProvider,
|
||||
envKeyPrefix,
|
||||
customRateLimit,
|
||||
'workspace-dim-wait'
|
||||
)
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
expect(mockAdapter.getTokenStatus).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
it('should pre-check all dimensions and block on first depleted one', async () => {
|
||||
const multiDimensionConfig: CustomRateLimit = {
|
||||
mode: 'custom',
|
||||
@@ -309,7 +654,7 @@ describe('HostedKeyRateLimiter', () => {
|
||||
tokensAvailable: 0,
|
||||
maxTokens: 100,
|
||||
lastRefillAt: new Date(),
|
||||
nextRefillAt: new Date(Date.now() + 30000),
|
||||
nextRefillAt: new Date(Date.now() + RETRY_PAST_CAP_MS),
|
||||
}
|
||||
mockAdapter.getTokenStatus
|
||||
.mockResolvedValueOnce(tokensBudget)
|
||||
|
||||
@@ -1,9 +1,14 @@
|
||||
import { createLogger } from '@sim/logger'
|
||||
import { sleep } from '@sim/utils/helpers'
|
||||
import { generateShortId } from '@sim/utils/id'
|
||||
import { getMaxExecutionTimeout } from '@/lib/core/execution-limits'
|
||||
import {
|
||||
createStorageAdapter,
|
||||
type RateLimitStorageAdapter,
|
||||
type TokenBucketConfig,
|
||||
} from '@/lib/core/rate-limiter/storage'
|
||||
import { PlatformEvents } from '@/lib/core/telemetry'
|
||||
import { getHostedKeyQueue, HEARTBEAT_REFRESH_INTERVAL_MS, type HostedKeyQueue } from './queue'
|
||||
import {
|
||||
type AcquireKeyResult,
|
||||
type CustomRateLimit,
|
||||
@@ -16,6 +21,60 @@ import {
|
||||
|
||||
const logger = createLogger('HostedKeyRateLimiter')
|
||||
|
||||
/**
|
||||
* Fallback ceiling on how long a hosted-key acquisition waits for the per-workspace
|
||||
* bucket to refill when no execution `AbortSignal` is available to bound the wait (e.g.
|
||||
* a caller without a wired execution deadline, or Redis no-op mode). When a signal IS
|
||||
* provided, the wait is instead bounded by the surrounding execution budget — the signal
|
||||
* fires when the run hits its plan timeout or is cancelled — and this constant no longer
|
||||
* applies (see {@link ABSOLUTE_MAX_QUEUE_WAIT_MS} for the backstop in that case).
|
||||
*/
|
||||
const MAX_QUEUE_WAIT_MS = 5 * 60 * 1000
|
||||
|
||||
/**
|
||||
* Hard safety ceiling applied even when an execution `AbortSignal` is present, in case the
|
||||
* signal never fires. Matches the longest possible execution budget (enterprise async) so
|
||||
* it never truncates a legitimately long-running background run before its own deadline.
|
||||
*/
|
||||
const ABSOLUTE_MAX_QUEUE_WAIT_MS = getMaxExecutionTimeout()
|
||||
|
||||
/**
|
||||
* Floor on per-iteration sleep when the bucket reports `retryAfterMs <= 0`,
|
||||
* which can happen due to clock skew or sub-millisecond resets. Prevents a
|
||||
* tight retry loop hammering the storage adapter.
|
||||
*/
|
||||
const MIN_QUEUE_RETRY_DELAY_MS = 50
|
||||
|
||||
/**
|
||||
* Poll interval while waiting to reach the head of the FIFO queue. 200ms balances
|
||||
* acquisition latency (worst-case wait for advancement is one poll period) against
|
||||
* Redis load — at this cadence, N waiters generate N×5 EVAL/sec, which is fine for
|
||||
* the typical low-tens contention. Revisit if telemetry shows hot Redis under load.
|
||||
*/
|
||||
const QUEUE_HEAD_POLL_MS = 200
|
||||
|
||||
/**
|
||||
* Sleep for `ms`, resolving early if `signal` aborts. Cleans up its own timer and listener
|
||||
* so neither leaks. Callers don't need to distinguish an early (aborted) return from a normal
|
||||
* one — the surrounding wait loop re-checks its budget immediately after and bails when the
|
||||
* signal has fired. Falls back to a plain sleep when no signal is provided.
|
||||
*/
|
||||
function interruptibleSleep(ms: number, signal?: AbortSignal): Promise<void> {
|
||||
if (!signal) return sleep(ms)
|
||||
if (signal.aborted) return Promise.resolve()
|
||||
return new Promise<void>((resolve) => {
|
||||
const onAbort = () => {
|
||||
clearTimeout(timer)
|
||||
resolve()
|
||||
}
|
||||
const timer = setTimeout(() => {
|
||||
signal.removeEventListener('abort', onAbort)
|
||||
resolve()
|
||||
}, ms)
|
||||
signal.addEventListener('abort', onAbort, { once: true })
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves env var names for a numbered key prefix using a `{PREFIX}_COUNT` env var.
|
||||
* E.g. with `EXA_API_KEY_COUNT=5`, returns `['EXA_API_KEY_1', ..., 'EXA_API_KEY_5']`.
|
||||
@@ -41,6 +100,18 @@ interface AvailableKey {
|
||||
envVarName: string
|
||||
}
|
||||
|
||||
/**
|
||||
* Mutable heartbeat bookkeeping shared across every wait phase of a single `acquireKey`
|
||||
* call. Carrying one `lastHeartbeatAt` across the queue-head wait and the bucket waits
|
||||
* ensures the heartbeat-refresh cadence reflects the *actual* last write to Redis rather
|
||||
* than resetting per phase, which could otherwise let the ticket heartbeat lapse and get
|
||||
* reaped mid-wait (breaking FIFO ordering).
|
||||
*/
|
||||
interface WaitState {
|
||||
/** Epoch ms of the last heartbeat write to Redis. */
|
||||
lastHeartbeatAt: number
|
||||
}
|
||||
|
||||
/**
|
||||
* HostedKeyRateLimiter provides:
|
||||
* 1. Per-billing-actor rate limiting (enforced - blocks actors who exceed their limit)
|
||||
@@ -52,11 +123,13 @@ interface AvailableKey {
|
||||
*/
|
||||
export class HostedKeyRateLimiter {
|
||||
private storage: RateLimitStorageAdapter
|
||||
private queue: HostedKeyQueue
|
||||
/** Round-robin counter per provider for even key distribution */
|
||||
private roundRobinCounters = new Map<string, number>()
|
||||
|
||||
constructor(storage?: RateLimitStorageAdapter) {
|
||||
constructor(storage?: RateLimitStorageAdapter, queue?: HostedKeyQueue) {
|
||||
this.storage = storage ?? createStorageAdapter()
|
||||
this.queue = queue ?? getHostedKeyQueue()
|
||||
}
|
||||
|
||||
private buildActorStorageKey(provider: string, billingActorId: string): string {
|
||||
@@ -179,71 +252,341 @@ export class HostedKeyRateLimiter {
|
||||
* Acquire an available key via round-robin selection.
|
||||
*
|
||||
* For both modes:
|
||||
* 1. Per-billing-actor request rate limiting (enforced): blocks actors who exceed their request limit
|
||||
* 1. Per-billing-actor request rate limiting (enforced): the call enqueues itself
|
||||
* onto a per-workspace+provider FIFO queue. Only the head of the queue attempts
|
||||
* to consume from the token bucket, guaranteeing strict ordering across callers
|
||||
* within a workspace. Different workspaces have independent queues and don't
|
||||
* block each other.
|
||||
* 2. Round-robin key selection: cycles through available keys for even distribution
|
||||
*
|
||||
* For `custom` mode additionally:
|
||||
* 3. Pre-checks dimension budgets: blocks if any dimension is already depleted
|
||||
* 3. Pre-checks dimension budgets: head waits on dimension refill the same way it
|
||||
* waits on actor request capacity.
|
||||
*
|
||||
* The wait is bounded by the surrounding execution budget: when `signal` is provided it
|
||||
* fires at the run's plan timeout (or on cancellation), so a queued call uses its full
|
||||
* available budget rather than a flat cap. When no signal is available the wait falls
|
||||
* back to `MAX_QUEUE_WAIT_MS`. On exhaustion the call returns today's 429 result. The
|
||||
* ticket is removed from the queue on exit regardless of success or failure.
|
||||
*
|
||||
* @param envKeyPrefix - Env var prefix (e.g. 'EXA_API_KEY'). Keys resolved via `{prefix}_COUNT`.
|
||||
* @param billingActorId - The billing actor (typically workspace ID) to rate limit against
|
||||
* @param signal - Optional execution `AbortSignal`; bounds the queue wait to the run's budget.
|
||||
*/
|
||||
async acquireKey(
|
||||
provider: string,
|
||||
envKeyPrefix: string,
|
||||
config: HostedKeyRateLimitConfig,
|
||||
billingActorId: string
|
||||
billingActorId: string,
|
||||
signal?: AbortSignal
|
||||
): Promise<AcquireKeyResult> {
|
||||
if (config.requestsPerMinute) {
|
||||
const rateLimitResult = await this.checkActorRateLimit(provider, billingActorId, config)
|
||||
if (rateLimitResult) {
|
||||
const ticketId = generateShortId()
|
||||
const startedAt = Date.now()
|
||||
const waitState: WaitState = { lastHeartbeatAt: startedAt }
|
||||
const enqueueResult = await this.queue.enqueue(provider, billingActorId, ticketId)
|
||||
|
||||
try {
|
||||
// Wait for our turn at the head of the queue (no-op when Redis unavailable).
|
||||
const headStatus = await this.waitForQueueHead(
|
||||
provider,
|
||||
billingActorId,
|
||||
ticketId,
|
||||
startedAt,
|
||||
waitState,
|
||||
signal
|
||||
)
|
||||
if (headStatus.timedOut) {
|
||||
PlatformEvents.hostedKeyQueueWaitExceeded({
|
||||
provider,
|
||||
workspaceId: billingActorId,
|
||||
waitedMs: Date.now() - startedAt,
|
||||
reason: 'queue_position',
|
||||
})
|
||||
return {
|
||||
success: false,
|
||||
billingActorRateLimited: true,
|
||||
retryAfterMs: rateLimitResult.retryAfterMs,
|
||||
error: `Rate limit exceeded. Please wait ${Math.ceil(rateLimitResult.retryAfterMs / 1000)} seconds. If you're getting throttled frequently, consider adding your own API key under Settings > BYOK to avoid shared rate limits.`,
|
||||
retryAfterMs: MAX_QUEUE_WAIT_MS,
|
||||
error: `Rate limit exceeded — request waited too long in the queue. If you're getting throttled frequently, consider adding your own API key under Settings > BYOK to avoid shared rate limits.`,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (config.mode === 'custom' && config.dimensions.length > 0) {
|
||||
const dimensionResult = await this.preCheckDimensions(provider, billingActorId, config)
|
||||
if (dimensionResult) {
|
||||
let dimensionWaited = false
|
||||
|
||||
if (config.requestsPerMinute) {
|
||||
const rateLimitResult = await this.waitForActorCapacity(
|
||||
provider,
|
||||
billingActorId,
|
||||
ticketId,
|
||||
config,
|
||||
startedAt,
|
||||
waitState,
|
||||
signal
|
||||
)
|
||||
if (rateLimitResult.rateLimited) {
|
||||
return {
|
||||
success: false,
|
||||
billingActorRateLimited: true,
|
||||
retryAfterMs: rateLimitResult.retryAfterMs,
|
||||
error: `Rate limit exceeded. Please wait ${Math.ceil(rateLimitResult.retryAfterMs / 1000)} seconds. If you're getting throttled frequently, consider adding your own API key under Settings > BYOK to avoid shared rate limits.`,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (config.mode === 'custom' && config.dimensions.length > 0) {
|
||||
const dimensionResult = await this.waitForDimensionCapacity(
|
||||
provider,
|
||||
billingActorId,
|
||||
ticketId,
|
||||
config,
|
||||
startedAt,
|
||||
waitState,
|
||||
signal
|
||||
)
|
||||
if (dimensionResult.rateLimited) {
|
||||
return {
|
||||
success: false,
|
||||
billingActorRateLimited: true,
|
||||
retryAfterMs: dimensionResult.retryAfterMs,
|
||||
error: `Rate limit exceeded for ${dimensionResult.dimension}. Please wait ${Math.ceil(dimensionResult.retryAfterMs / 1000)} seconds. If you're getting throttled frequently, consider adding your own API key under Settings > BYOK to avoid shared rate limits.`,
|
||||
}
|
||||
}
|
||||
dimensionWaited = dimensionResult.waited
|
||||
}
|
||||
|
||||
const totalWaitedMs = Date.now() - startedAt
|
||||
if (enqueueResult.enabled && (enqueueResult.position > 0 || totalWaitedMs > 100)) {
|
||||
// Attribute the wait to its dominant cause: queue depth takes precedence (it's
|
||||
// reported alongside queuePosition), otherwise the bucket phase that actually slept.
|
||||
const reason: 'queue_position' | 'actor_requests' | 'dimension' =
|
||||
enqueueResult.position > 0
|
||||
? 'queue_position'
|
||||
: dimensionWaited
|
||||
? 'dimension'
|
||||
: 'actor_requests'
|
||||
PlatformEvents.hostedKeyQueueWaited({
|
||||
provider,
|
||||
workspaceId: billingActorId,
|
||||
waitedMs: totalWaitedMs,
|
||||
attempts: 1,
|
||||
reason,
|
||||
queuePosition: enqueueResult.position,
|
||||
})
|
||||
}
|
||||
|
||||
const envKeys = resolveEnvKeys(envKeyPrefix)
|
||||
const availableKeys = this.getAvailableKeys(envKeys)
|
||||
|
||||
if (availableKeys.length === 0) {
|
||||
logger.warn(`No hosted keys configured for provider ${provider}`)
|
||||
return {
|
||||
success: false,
|
||||
billingActorRateLimited: true,
|
||||
retryAfterMs: dimensionResult.retryAfterMs,
|
||||
error: `Rate limit exceeded for ${dimensionResult.dimension}. Please wait ${Math.ceil(dimensionResult.retryAfterMs / 1000)} seconds. If you're getting throttled frequently, consider adding your own API key under Settings > BYOK to avoid shared rate limits.`,
|
||||
error: `No hosted keys configured for ${provider}`,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const envKeys = resolveEnvKeys(envKeyPrefix)
|
||||
const availableKeys = this.getAvailableKeys(envKeys)
|
||||
const counter = this.roundRobinCounters.get(provider) ?? 0
|
||||
const selected = availableKeys[counter % availableKeys.length]
|
||||
this.roundRobinCounters.set(provider, counter + 1)
|
||||
|
||||
logger.debug(`Selected hosted key for ${provider}`, {
|
||||
provider,
|
||||
keyIndex: selected.keyIndex,
|
||||
envVarName: selected.envVarName,
|
||||
})
|
||||
|
||||
if (availableKeys.length === 0) {
|
||||
logger.warn(`No hosted keys configured for provider ${provider}`)
|
||||
return {
|
||||
success: false,
|
||||
error: `No hosted keys configured for ${provider}`,
|
||||
success: true,
|
||||
key: selected.key,
|
||||
keyIndex: selected.keyIndex,
|
||||
envVarName: selected.envVarName,
|
||||
}
|
||||
} finally {
|
||||
// Always remove our ticket so the next caller can advance, regardless of whether
|
||||
// we succeeded, hit the cap, or threw. Best-effort; safe to call multiple times.
|
||||
await this.queue.dequeue(provider, billingActorId, ticketId)
|
||||
}
|
||||
}
|
||||
|
||||
const counter = this.roundRobinCounters.get(provider) ?? 0
|
||||
const selected = availableKeys[counter % availableKeys.length]
|
||||
this.roundRobinCounters.set(provider, counter + 1)
|
||||
/**
|
||||
* Remaining time budget for waiting, in milliseconds. When an execution `AbortSignal` is
|
||||
* present it governs the wait: the budget is exhausted the moment the signal aborts (the
|
||||
* run hit its plan timeout or was cancelled), with {@link ABSOLUTE_MAX_QUEUE_WAIT_MS} as a
|
||||
* backstop should the signal never fire. Without a signal we fall back to the flat
|
||||
* {@link MAX_QUEUE_WAIT_MS} ceiling.
|
||||
*/
|
||||
private remainingWaitBudgetMs(startedAt: number, signal?: AbortSignal): number {
|
||||
if (signal?.aborted) return 0
|
||||
const ceiling = signal ? ABSOLUTE_MAX_QUEUE_WAIT_MS : MAX_QUEUE_WAIT_MS
|
||||
return ceiling - (Date.now() - startedAt)
|
||||
}
|
||||
|
||||
logger.debug(`Selected hosted key for ${provider}`, {
|
||||
provider,
|
||||
keyIndex: selected.keyIndex,
|
||||
envVarName: selected.envVarName,
|
||||
})
|
||||
/** Refresh the ticket heartbeat if the refresh interval has elapsed since the last write. */
|
||||
private async maybeRefreshHeartbeat(
|
||||
provider: string,
|
||||
billingActorId: string,
|
||||
ticketId: string,
|
||||
waitState: WaitState
|
||||
): Promise<void> {
|
||||
if (Date.now() - waitState.lastHeartbeatAt >= HEARTBEAT_REFRESH_INTERVAL_MS) {
|
||||
await this.queue.refreshHeartbeat(provider, billingActorId, ticketId)
|
||||
waitState.lastHeartbeatAt = Date.now()
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
success: true,
|
||||
key: selected.key,
|
||||
keyIndex: selected.keyIndex,
|
||||
envVarName: selected.envVarName,
|
||||
/**
|
||||
* Sleep before the next bucket re-check, refreshing the heartbeat first if due. The sleep
|
||||
* is capped at {@link HEARTBEAT_REFRESH_INTERVAL_MS} so that no single wait can outlive the
|
||||
* heartbeat TTL — even when the bucket reports a long `retryAfterMs` (e.g. low-RPM
|
||||
* providers). Without this cap a multi-second sleep could let the heartbeat lapse, the
|
||||
* head get reaped as dead, and a second caller advance and race us for the bucket. The
|
||||
* sleep also resolves early if `signal` aborts, so a cancelled/timed-out run stops waiting
|
||||
* promptly rather than overshooting by up to the cap.
|
||||
*/
|
||||
private async heartbeatAwareSleep(
|
||||
provider: string,
|
||||
billingActorId: string,
|
||||
ticketId: string,
|
||||
desiredMs: number,
|
||||
waitState: WaitState,
|
||||
signal?: AbortSignal
|
||||
): Promise<void> {
|
||||
await this.maybeRefreshHeartbeat(provider, billingActorId, ticketId, waitState)
|
||||
const sleepMs = Math.min(
|
||||
Math.max(MIN_QUEUE_RETRY_DELAY_MS, desiredMs),
|
||||
HEARTBEAT_REFRESH_INTERVAL_MS
|
||||
)
|
||||
await interruptibleSleep(sleepMs, signal)
|
||||
}
|
||||
|
||||
/**
|
||||
* Block until our ticket reaches the head of the queue. Refreshes the heartbeat on a
|
||||
* regular cadence so we don't get reaped as dead. Returns `timedOut: true` once the wait
|
||||
* budget is exhausted before reaching the head (see {@link remainingWaitBudgetMs}).
|
||||
*
|
||||
* No-op when Redis is unavailable (queue.enqueue returns enabled=false and checkHead
|
||||
* always returns 'head').
|
||||
*/
|
||||
private async waitForQueueHead(
|
||||
provider: string,
|
||||
billingActorId: string,
|
||||
ticketId: string,
|
||||
startedAt: number,
|
||||
waitState: WaitState,
|
||||
signal?: AbortSignal
|
||||
): Promise<{ timedOut: boolean }> {
|
||||
while (true) {
|
||||
const status = await this.queue.checkHead(provider, billingActorId, ticketId)
|
||||
if (status === 'head') return { timedOut: false }
|
||||
|
||||
// 'missing' shouldn't normally happen — the queue list TTL (10min) outlives a typical
|
||||
// wait — but if it does (e.g. Redis flushed mid-wait), treat as "you're up" so the
|
||||
// caller proceeds to the bucket race rather than hanging forever.
|
||||
if (status === 'missing') return { timedOut: false }
|
||||
|
||||
if (this.remainingWaitBudgetMs(startedAt, signal) <= 0) {
|
||||
return { timedOut: true }
|
||||
}
|
||||
|
||||
await this.maybeRefreshHeartbeat(provider, billingActorId, ticketId, waitState)
|
||||
await interruptibleSleep(QUEUE_HEAD_POLL_MS, signal)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Wait for actor request-rate capacity. Called once we're at the head of the FIFO
|
||||
* queue, so other callers can't race us for the next token — they're blocked behind us
|
||||
* at queue level. Re-checks the bucket until the wait budget is exhausted (accounting for
|
||||
* time already spent waiting in the queue). `waited` reports whether the loop ever slept.
|
||||
*/
|
||||
private async waitForActorCapacity(
|
||||
provider: string,
|
||||
billingActorId: string,
|
||||
ticketId: string,
|
||||
config: HostedKeyRateLimitConfig,
|
||||
startedAt: number,
|
||||
waitState: WaitState,
|
||||
signal?: AbortSignal
|
||||
): Promise<
|
||||
{ rateLimited: false; waited: boolean } | { rateLimited: true; retryAfterMs: number }
|
||||
> {
|
||||
let waited = false
|
||||
|
||||
while (true) {
|
||||
const result = await this.checkActorRateLimit(provider, billingActorId, config)
|
||||
if (!result) return { rateLimited: false, waited }
|
||||
|
||||
const remaining = this.remainingWaitBudgetMs(startedAt, signal)
|
||||
if (remaining <= 0 || result.retryAfterMs > remaining) {
|
||||
PlatformEvents.hostedKeyQueueWaitExceeded({
|
||||
provider,
|
||||
workspaceId: billingActorId,
|
||||
waitedMs: Date.now() - startedAt,
|
||||
reason: 'actor_requests',
|
||||
})
|
||||
return { rateLimited: true, retryAfterMs: result.retryAfterMs }
|
||||
}
|
||||
|
||||
waited = true
|
||||
await this.heartbeatAwareSleep(
|
||||
provider,
|
||||
billingActorId,
|
||||
ticketId,
|
||||
result.retryAfterMs,
|
||||
waitState,
|
||||
signal
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Wait for custom-mode dimension capacity. `preCheckDimensions` is read-only — it does
|
||||
* not consume — so re-running it after a sleep is safe and does not double-charge.
|
||||
* Post-execution `reportUsage` performs the actual consumption. `waited` reports whether
|
||||
* the loop ever slept.
|
||||
*/
|
||||
private async waitForDimensionCapacity(
|
||||
provider: string,
|
||||
billingActorId: string,
|
||||
ticketId: string,
|
||||
config: CustomRateLimit,
|
||||
startedAt: number,
|
||||
waitState: WaitState,
|
||||
signal?: AbortSignal
|
||||
): Promise<
|
||||
| { rateLimited: false; waited: boolean }
|
||||
| { rateLimited: true; retryAfterMs: number; dimension: string }
|
||||
> {
|
||||
let waited = false
|
||||
|
||||
while (true) {
|
||||
const result = await this.preCheckDimensions(provider, billingActorId, config)
|
||||
if (!result) return { rateLimited: false, waited }
|
||||
|
||||
const remaining = this.remainingWaitBudgetMs(startedAt, signal)
|
||||
if (remaining <= 0 || result.retryAfterMs > remaining) {
|
||||
PlatformEvents.hostedKeyQueueWaitExceeded({
|
||||
provider,
|
||||
workspaceId: billingActorId,
|
||||
waitedMs: Date.now() - startedAt,
|
||||
reason: 'dimension',
|
||||
dimension: result.dimension,
|
||||
})
|
||||
return {
|
||||
rateLimited: true,
|
||||
retryAfterMs: result.retryAfterMs,
|
||||
dimension: result.dimension,
|
||||
}
|
||||
}
|
||||
|
||||
waited = true
|
||||
await this.heartbeatAwareSleep(
|
||||
provider,
|
||||
billingActorId,
|
||||
ticketId,
|
||||
result.retryAfterMs,
|
||||
waitState,
|
||||
signal
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
import { redisConfigMock, redisConfigMockFns } from '@sim/testing'
|
||||
import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
|
||||
import { HostedKeyQueue } from './queue'
|
||||
|
||||
vi.mock('@/lib/core/config/redis', () => redisConfigMock)
|
||||
|
||||
interface MockPipeline {
|
||||
rpush: Mock
|
||||
expire: Mock
|
||||
set: Mock
|
||||
lrem: Mock
|
||||
del: Mock
|
||||
exec: Mock
|
||||
}
|
||||
|
||||
interface MockRedis {
|
||||
multi: Mock
|
||||
set: Mock
|
||||
eval: Mock
|
||||
pipeline: MockPipeline
|
||||
}
|
||||
|
||||
function createFakeRedis(): MockRedis {
|
||||
const pipeline: MockPipeline = {
|
||||
rpush: vi.fn(),
|
||||
expire: vi.fn(),
|
||||
set: vi.fn(),
|
||||
lrem: vi.fn(),
|
||||
del: vi.fn(),
|
||||
exec: vi.fn(),
|
||||
}
|
||||
// Pipeline methods return the pipeline for chaining.
|
||||
pipeline.rpush.mockReturnValue(pipeline)
|
||||
pipeline.expire.mockReturnValue(pipeline)
|
||||
pipeline.set.mockReturnValue(pipeline)
|
||||
pipeline.lrem.mockReturnValue(pipeline)
|
||||
pipeline.del.mockReturnValue(pipeline)
|
||||
|
||||
return {
|
||||
multi: vi.fn(() => pipeline),
|
||||
set: vi.fn(),
|
||||
eval: vi.fn(),
|
||||
pipeline,
|
||||
}
|
||||
}
|
||||
|
||||
const provider = 'exa'
|
||||
const workspaceId = 'workspace-1'
|
||||
const ticketId = 'ticket-1'
|
||||
|
||||
describe('HostedKeyQueue', () => {
|
||||
let queue: HostedKeyQueue
|
||||
let mockRedis: MockRedis
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
mockRedis = createFakeRedis()
|
||||
redisConfigMockFns.mockGetRedisClient.mockReturnValue(mockRedis)
|
||||
queue = new HostedKeyQueue()
|
||||
})
|
||||
|
||||
describe('enqueue', () => {
|
||||
it('returns position 0 when first in line', async () => {
|
||||
// RPUSH returns new list length; first push -> 1.
|
||||
mockRedis.pipeline.exec.mockResolvedValueOnce([
|
||||
[null, 1],
|
||||
[null, 1],
|
||||
[null, 'OK'],
|
||||
])
|
||||
|
||||
const result = await queue.enqueue(provider, workspaceId, ticketId)
|
||||
|
||||
expect(result).toEqual({ position: 0, enabled: true })
|
||||
expect(mockRedis.pipeline.rpush).toHaveBeenCalledWith(
|
||||
'hosted-queue:exa:workspace-1',
|
||||
ticketId
|
||||
)
|
||||
expect(mockRedis.pipeline.set).toHaveBeenCalledWith(
|
||||
'hosted-queue-tkt:exa:workspace-1:ticket-1',
|
||||
'1',
|
||||
'EX',
|
||||
expect.any(Number)
|
||||
)
|
||||
})
|
||||
|
||||
it('returns higher position when others are ahead', async () => {
|
||||
// Length 5 after push -> position 4.
|
||||
mockRedis.pipeline.exec.mockResolvedValueOnce([
|
||||
[null, 5],
|
||||
[null, 1],
|
||||
[null, 'OK'],
|
||||
])
|
||||
|
||||
const result = await queue.enqueue(provider, workspaceId, ticketId)
|
||||
|
||||
expect(result.position).toBe(4)
|
||||
})
|
||||
|
||||
it('falls back to enabled=false when Redis is unavailable', async () => {
|
||||
redisConfigMockFns.mockGetRedisClient.mockReturnValueOnce(null)
|
||||
|
||||
const result = await queue.enqueue(provider, workspaceId, ticketId)
|
||||
|
||||
expect(result).toEqual({ position: 0, enabled: false })
|
||||
})
|
||||
|
||||
it('falls back to enabled=false on Redis error', async () => {
|
||||
mockRedis.pipeline.exec.mockRejectedValueOnce(new Error('connection lost'))
|
||||
|
||||
const result = await queue.enqueue(provider, workspaceId, ticketId)
|
||||
|
||||
expect(result.enabled).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe('checkHead', () => {
|
||||
it('returns "head" when our ticket is at the head', async () => {
|
||||
mockRedis.eval.mockResolvedValueOnce('head')
|
||||
|
||||
const status = await queue.checkHead(provider, workspaceId, ticketId)
|
||||
|
||||
expect(status).toBe('head')
|
||||
})
|
||||
|
||||
it('returns "waiting" when someone else is the head', async () => {
|
||||
mockRedis.eval.mockResolvedValueOnce('waiting')
|
||||
|
||||
const status = await queue.checkHead(provider, workspaceId, ticketId)
|
||||
|
||||
expect(status).toBe('waiting')
|
||||
})
|
||||
|
||||
it('returns "missing" when our ticket is not in the queue', async () => {
|
||||
mockRedis.eval.mockResolvedValueOnce('missing')
|
||||
|
||||
const status = await queue.checkHead(provider, workspaceId, ticketId)
|
||||
|
||||
expect(status).toBe('missing')
|
||||
})
|
||||
|
||||
it('passes queue list key, heartbeat prefix, and ticketId to the Lua script', async () => {
|
||||
mockRedis.eval.mockResolvedValueOnce('head')
|
||||
|
||||
await queue.checkHead(provider, workspaceId, ticketId)
|
||||
|
||||
expect(mockRedis.eval).toHaveBeenCalledWith(
|
||||
expect.stringContaining('lindex'),
|
||||
1,
|
||||
'hosted-queue:exa:workspace-1',
|
||||
'hosted-queue-tkt:exa:workspace-1:',
|
||||
ticketId
|
||||
)
|
||||
})
|
||||
|
||||
it('fails open to "head" on Redis error so callers do not hang', async () => {
|
||||
mockRedis.eval.mockRejectedValueOnce(new Error('boom'))
|
||||
|
||||
const status = await queue.checkHead(provider, workspaceId, ticketId)
|
||||
|
||||
expect(status).toBe('head')
|
||||
})
|
||||
|
||||
it('returns "head" no-op when Redis is unavailable', async () => {
|
||||
redisConfigMockFns.mockGetRedisClient.mockReturnValueOnce(null)
|
||||
|
||||
const status = await queue.checkHead(provider, workspaceId, ticketId)
|
||||
|
||||
expect(status).toBe('head')
|
||||
})
|
||||
})
|
||||
|
||||
describe('refreshHeartbeat', () => {
|
||||
it('writes the heartbeat key with TTL', async () => {
|
||||
mockRedis.set.mockResolvedValueOnce('OK')
|
||||
|
||||
await queue.refreshHeartbeat(provider, workspaceId, ticketId)
|
||||
|
||||
expect(mockRedis.set).toHaveBeenCalledWith(
|
||||
'hosted-queue-tkt:exa:workspace-1:ticket-1',
|
||||
'1',
|
||||
'EX',
|
||||
expect.any(Number)
|
||||
)
|
||||
})
|
||||
|
||||
it('is a no-op when Redis is unavailable', async () => {
|
||||
redisConfigMockFns.mockGetRedisClient.mockReturnValueOnce(null)
|
||||
|
||||
await expect(queue.refreshHeartbeat(provider, workspaceId, ticketId)).resolves.toBeUndefined()
|
||||
expect(mockRedis.set).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe('dequeue', () => {
|
||||
it('removes the ticket from the list and deletes the heartbeat', async () => {
|
||||
mockRedis.pipeline.exec.mockResolvedValueOnce([
|
||||
[null, 1],
|
||||
[null, 1],
|
||||
])
|
||||
|
||||
await queue.dequeue(provider, workspaceId, ticketId)
|
||||
|
||||
expect(mockRedis.pipeline.lrem).toHaveBeenCalledWith(
|
||||
'hosted-queue:exa:workspace-1',
|
||||
1,
|
||||
ticketId
|
||||
)
|
||||
expect(mockRedis.pipeline.del).toHaveBeenCalledWith(
|
||||
'hosted-queue-tkt:exa:workspace-1:ticket-1'
|
||||
)
|
||||
})
|
||||
|
||||
it('is a no-op when Redis is unavailable', async () => {
|
||||
redisConfigMockFns.mockGetRedisClient.mockReturnValueOnce(null)
|
||||
|
||||
await expect(queue.dequeue(provider, workspaceId, ticketId)).resolves.toBeUndefined()
|
||||
expect(mockRedis.multi).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('swallows errors so callers do not throw on cleanup', async () => {
|
||||
mockRedis.pipeline.exec.mockRejectedValueOnce(new Error('connection lost'))
|
||||
|
||||
await expect(queue.dequeue(provider, workspaceId, ticketId)).resolves.toBeUndefined()
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,203 @@
|
||||
import { createLogger } from '@sim/logger'
|
||||
import { toError } from '@sim/utils/errors'
|
||||
import { getRedisClient } from '@/lib/core/config/redis'
|
||||
|
||||
const logger = createLogger('HostedKeyQueue')
|
||||
|
||||
/**
|
||||
* Per-ticket heartbeat TTL. Refreshed by the head while it's actively waiting
|
||||
* on the bucket. If the holder crashes, the heartbeat key expires, and the next
|
||||
* caller sees the head as dead and removes it (lazy cleanup).
|
||||
*/
|
||||
const TICKET_HEARTBEAT_TTL_SECONDS = 30
|
||||
|
||||
/** How often the head should refresh its heartbeat while waiting. */
|
||||
export const HEARTBEAT_REFRESH_INTERVAL_MS = 10_000
|
||||
|
||||
/**
|
||||
* TTL on the queue list itself. Set on every enqueue. Prevents abandoned queues
|
||||
* (whole workspace went silent) from sticking around forever in Redis.
|
||||
*/
|
||||
const QUEUE_LIST_TTL_SECONDS = 600
|
||||
|
||||
const queueListKey = (provider: string, billingActorId: string): string =>
|
||||
`hosted-queue:${provider}:${billingActorId}`
|
||||
|
||||
const heartbeatKey = (provider: string, billingActorId: string, ticketId: string): string =>
|
||||
`hosted-queue-tkt:${provider}:${billingActorId}:${ticketId}`
|
||||
|
||||
/**
|
||||
* Atomically reap any dead head, then return our ticket's status. Combines what
|
||||
* would otherwise be 3 round-trips (reap, LINDEX, LPOS) into one EVAL — meaningful
|
||||
* because callers poll this every ~200ms while waiting in the queue.
|
||||
*
|
||||
* `KEYS[1]` = queue list key. `ARGV[1]` = heartbeat key prefix. `ARGV[2]` = our ticketId.
|
||||
*
|
||||
* Reaping is bounded: at most one dead head is removed per call. If multiple dead
|
||||
* tickets pile up at the head, subsequent polls will clean them one by one. This
|
||||
* keeps the script O(1) rather than O(N) and is sufficient because queue depth
|
||||
* is bounded by concurrent callers per workspace (typically tens).
|
||||
*
|
||||
* Returns one of: "head", "waiting", "missing".
|
||||
*/
|
||||
const CHECK_HEAD_SCRIPT = `
|
||||
local head = redis.call("lindex", KEYS[1], 0)
|
||||
if head and redis.call("exists", ARGV[1] .. head) == 0 then
|
||||
redis.call("lrem", KEYS[1], 1, head)
|
||||
head = redis.call("lindex", KEYS[1], 0)
|
||||
end
|
||||
if not head then
|
||||
return "missing"
|
||||
end
|
||||
if head == ARGV[2] then
|
||||
return "head"
|
||||
end
|
||||
if redis.call("lpos", KEYS[1], ARGV[2]) == false then
|
||||
return "missing"
|
||||
end
|
||||
return "waiting"
|
||||
`
|
||||
|
||||
export interface EnqueueResult {
|
||||
/** Position at the moment of enqueue (0 = head, you go next). */
|
||||
position: number
|
||||
/** Whether Redis was available — false means we're in no-op mode. */
|
||||
enabled: boolean
|
||||
}
|
||||
|
||||
/**
|
||||
* Per-workspace+provider FIFO queue for hosted-key acquisitions.
|
||||
*
|
||||
* Callers `enqueue` to claim a position, then `waitForHead` until they're at
|
||||
* the head, then attempt to consume from the token bucket. On success or cap
|
||||
* exceeded, they `dequeue` to make room for the next caller.
|
||||
*
|
||||
* No-op when Redis is unavailable: every method returns "you're the head /
|
||||
* empty / etc." so the rate limiter falls back to plain bucket racing.
|
||||
*/
|
||||
export class HostedKeyQueue {
|
||||
/**
|
||||
* Push a ticket onto the tail of the queue and write a heartbeat. Returns the
|
||||
* position at enqueue time (0 = head, ready to proceed).
|
||||
*/
|
||||
async enqueue(
|
||||
provider: string,
|
||||
billingActorId: string,
|
||||
ticketId: string
|
||||
): Promise<EnqueueResult> {
|
||||
const redis = getRedisClient()
|
||||
if (!redis) {
|
||||
return { position: 0, enabled: false }
|
||||
}
|
||||
|
||||
const listKey = queueListKey(provider, billingActorId)
|
||||
const hbKey = heartbeatKey(provider, billingActorId, ticketId)
|
||||
|
||||
try {
|
||||
const pipeline = redis.multi()
|
||||
pipeline.rpush(listKey, ticketId)
|
||||
pipeline.expire(listKey, QUEUE_LIST_TTL_SECONDS)
|
||||
pipeline.set(hbKey, '1', 'EX', TICKET_HEARTBEAT_TTL_SECONDS)
|
||||
const results = await pipeline.exec()
|
||||
// results[0] is the rpush response: [err, length]
|
||||
const length = results?.[0] && typeof results[0][1] === 'number' ? results[0][1] : 1
|
||||
// Position is length - 1 (just-pushed at the tail).
|
||||
return { position: length - 1, enabled: true }
|
||||
} catch (error) {
|
||||
logger.warn(`Queue enqueue failed for ${listKey}`, { error: toError(error).message })
|
||||
return { position: 0, enabled: false }
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check whether `ticketId` is currently at the head of the queue. If the head
|
||||
* is a different ticket but its heartbeat has expired (caller crashed), reap
|
||||
* it and re-check on the next poll.
|
||||
*
|
||||
* Returns:
|
||||
* - "head": you're at the head, proceed to consume from the bucket
|
||||
* - "waiting": someone else is the head and they're alive
|
||||
* - "missing": your ticket isn't in the queue at all (e.g. queue list TTL
|
||||
* expired); caller should re-enqueue or treat as enabled=false
|
||||
*/
|
||||
async checkHead(
|
||||
provider: string,
|
||||
billingActorId: string,
|
||||
ticketId: string
|
||||
): Promise<'head' | 'waiting' | 'missing'> {
|
||||
const redis = getRedisClient()
|
||||
if (!redis) {
|
||||
return 'head'
|
||||
}
|
||||
|
||||
const listKey = queueListKey(provider, billingActorId)
|
||||
const hbPrefix = `hosted-queue-tkt:${provider}:${billingActorId}:`
|
||||
|
||||
try {
|
||||
const result = (await redis.eval(CHECK_HEAD_SCRIPT, 1, listKey, hbPrefix, ticketId)) as
|
||||
| 'head'
|
||||
| 'waiting'
|
||||
| 'missing'
|
||||
return result
|
||||
} catch (error) {
|
||||
logger.warn(`Queue checkHead failed for ${listKey}`, { error: toError(error).message })
|
||||
// Fail-open: treat as head so the caller proceeds rather than hanging.
|
||||
return 'head'
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Refresh the ticket's heartbeat. Called periodically by the head while it's
|
||||
* waiting on the bucket so it doesn't get reaped as dead.
|
||||
*/
|
||||
async refreshHeartbeat(
|
||||
provider: string,
|
||||
billingActorId: string,
|
||||
ticketId: string
|
||||
): Promise<void> {
|
||||
const redis = getRedisClient()
|
||||
if (!redis) return
|
||||
|
||||
const hbKey = heartbeatKey(provider, billingActorId, ticketId)
|
||||
try {
|
||||
await redis.set(hbKey, '1', 'EX', TICKET_HEARTBEAT_TTL_SECONDS)
|
||||
} catch (error) {
|
||||
logger.warn(`Queue heartbeat refresh failed for ${hbKey}`, {
|
||||
error: toError(error).message,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Remove a ticket from the queue and its heartbeat key. Best-effort; safe to
|
||||
* call multiple times. LREM count=1 removes at most one matching entry.
|
||||
*/
|
||||
async dequeue(provider: string, billingActorId: string, ticketId: string): Promise<void> {
|
||||
const redis = getRedisClient()
|
||||
if (!redis) return
|
||||
|
||||
const listKey = queueListKey(provider, billingActorId)
|
||||
const hbKey = heartbeatKey(provider, billingActorId, ticketId)
|
||||
try {
|
||||
const pipeline = redis.multi()
|
||||
pipeline.lrem(listKey, 1, ticketId)
|
||||
pipeline.del(hbKey)
|
||||
await pipeline.exec()
|
||||
} catch (error) {
|
||||
logger.warn(`Queue dequeue failed for ${listKey}`, { error: toError(error).message })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let cachedQueue: HostedKeyQueue | null = null
|
||||
|
||||
export function getHostedKeyQueue(): HostedKeyQueue {
|
||||
if (!cachedQueue) {
|
||||
cachedQueue = new HostedKeyQueue()
|
||||
}
|
||||
return cachedQueue
|
||||
}
|
||||
|
||||
export function resetHostedKeyQueue(): void {
|
||||
cachedQueue = null
|
||||
}
|
||||
@@ -1002,6 +1002,50 @@ export const PlatformEvents = {
|
||||
})
|
||||
},
|
||||
|
||||
/**
|
||||
* Track a successful hosted-key acquisition that had to wait — either for a slot at
|
||||
* the head of the FIFO queue, or for the actor/dimension bucket to refill once at the
|
||||
* head. `queuePosition` is the position at the moment of enqueue (0 = ready to proceed).
|
||||
*/
|
||||
hostedKeyQueueWaited: (attrs: {
|
||||
provider: string
|
||||
workspaceId: string
|
||||
waitedMs: number
|
||||
attempts: number
|
||||
reason: 'actor_requests' | 'dimension' | 'queue_position'
|
||||
dimension?: string
|
||||
queuePosition?: number
|
||||
}) => {
|
||||
trackPlatformEvent('platform.hosted_key.queue_waited', {
|
||||
'provider.id': attrs.provider,
|
||||
'workspace.id': attrs.workspaceId,
|
||||
'queue.waited_ms': attrs.waitedMs,
|
||||
'queue.attempts': attrs.attempts,
|
||||
'queue.reason': attrs.reason,
|
||||
...(attrs.dimension && { 'queue.dimension': attrs.dimension }),
|
||||
...(attrs.queuePosition != null && { 'queue.position': attrs.queuePosition }),
|
||||
})
|
||||
},
|
||||
|
||||
/**
|
||||
* Track a hosted-key acquisition that exceeded the queue wait cap and fell back to a 429.
|
||||
*/
|
||||
hostedKeyQueueWaitExceeded: (attrs: {
|
||||
provider: string
|
||||
workspaceId: string
|
||||
waitedMs: number
|
||||
reason: 'actor_requests' | 'dimension' | 'queue_position'
|
||||
dimension?: string
|
||||
}) => {
|
||||
trackPlatformEvent('platform.hosted_key.queue_wait_exceeded', {
|
||||
'provider.id': attrs.provider,
|
||||
'workspace.id': attrs.workspaceId,
|
||||
'queue.waited_ms': attrs.waitedMs,
|
||||
'queue.reason': attrs.reason,
|
||||
...(attrs.dimension && { 'queue.dimension': attrs.dimension }),
|
||||
})
|
||||
},
|
||||
|
||||
/**
|
||||
* Track chat deployed (workflow deployed as chat interface)
|
||||
*/
|
||||
|
||||
+86
-2
@@ -272,7 +272,8 @@ async function injectHostedKeyIfNeeded(
|
||||
provider,
|
||||
envKeyPrefix,
|
||||
rateLimit,
|
||||
billingActorId
|
||||
billingActorId,
|
||||
executionContext?.abortSignal
|
||||
)
|
||||
|
||||
if (!acquireResult.success && acquireResult.billingActorRateLimited) {
|
||||
@@ -318,6 +319,49 @@ async function injectHostedKeyIfNeeded(
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Re-acquire a hosted key after upstream-429 retries have been exhausted. Calls
|
||||
* `acquireKey` (which now blocks on the per-workspace bucket) and re-injects the
|
||||
* fresh key into `params`. Returns false if no key could be obtained — caller
|
||||
* should re-throw the original upstream 429.
|
||||
*
|
||||
* Does not consult BYOK. We only enter this path from inside the hosted-key
|
||||
* branch of `executeTool`, so BYOK has already been ruled out for this call.
|
||||
*/
|
||||
async function reacquireHostedKey(
|
||||
tool: ToolConfig,
|
||||
params: Record<string, unknown>,
|
||||
executionContext: ExecutionContext | undefined,
|
||||
requestId: string
|
||||
): Promise<boolean> {
|
||||
if (!tool.hosting) return false
|
||||
const { envKeyPrefix, apiKeyParam, byokProviderId, rateLimit } = tool.hosting
|
||||
const { workspaceId } = resolveToolScope(params, executionContext)
|
||||
if (!workspaceId) return false
|
||||
|
||||
const provider = byokProviderId || tool.id
|
||||
const acquireResult = await getHostedKeyRateLimiter().acquireKey(
|
||||
provider,
|
||||
envKeyPrefix,
|
||||
rateLimit,
|
||||
workspaceId,
|
||||
executionContext?.abortSignal
|
||||
)
|
||||
|
||||
if (!acquireResult.success || !acquireResult.key) {
|
||||
logger.warn(
|
||||
`[${requestId}] Re-acquire of hosted key for ${tool.id} failed: ${acquireResult.error ?? 'unknown'}`
|
||||
)
|
||||
return false
|
||||
}
|
||||
|
||||
params[apiKeyParam] = acquireResult.key
|
||||
logger.info(
|
||||
`[${requestId}] Re-acquired hosted key for ${tool.id} (${acquireResult.envVarName}) after upstream throttling`
|
||||
)
|
||||
return true
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if an error is a rate limit (throttling) or quota exhaustion error.
|
||||
* Some providers (e.g. Perplexity) return 401/403 with "insufficient_quota"
|
||||
@@ -344,11 +388,23 @@ interface RetryContext {
|
||||
toolId: string
|
||||
envVarName: string
|
||||
executionContext?: ExecutionContext
|
||||
/**
|
||||
* Optional callback invoked after the local exponential backoff has been exhausted by
|
||||
* upstream 429s. Should re-enter the per-workspace hosted-key queue (which now blocks
|
||||
* on the bucket) and return a fresh execution thunk bound to the newly acquired key.
|
||||
* If the callback returns null, we give up and re-throw the last error.
|
||||
*/
|
||||
reacquireAfterRetriesExhausted?: () => Promise<(() => Promise<unknown>) | null>
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a function with exponential backoff retry for rate limiting errors.
|
||||
* Only used for hosted key requests. Tracks rate limit events via telemetry.
|
||||
*
|
||||
* On terminal upstream 429, optionally re-enters the hosted-key queue (which waits for
|
||||
* the per-workspace bucket to refill) and retries once with a freshly acquired key.
|
||||
* This handles the case where the upstream provider's limit is tighter than ours — we
|
||||
* re-queue the call instead of surfacing the error.
|
||||
*/
|
||||
async function executeWithRetry<T>(
|
||||
fn: () => Promise<T>,
|
||||
@@ -356,7 +412,8 @@ async function executeWithRetry<T>(
|
||||
maxRetries = 3,
|
||||
baseDelayMs = 1000
|
||||
): Promise<T> {
|
||||
const { requestId, toolId, envVarName, executionContext } = context
|
||||
const { requestId, toolId, envVarName, executionContext, reacquireAfterRetriesExhausted } =
|
||||
context
|
||||
let lastError: unknown
|
||||
|
||||
for (let attempt = 0; attempt <= maxRetries; attempt++) {
|
||||
@@ -367,6 +424,23 @@ async function executeWithRetry<T>(
|
||||
|
||||
if (!isRateLimitError(error) || attempt === maxRetries) {
|
||||
if (isRateLimitError(error) && attempt === maxRetries) {
|
||||
if (reacquireAfterRetriesExhausted) {
|
||||
try {
|
||||
const requeued = await reacquireAfterRetriesExhausted()
|
||||
if (requeued) {
|
||||
logger.warn(
|
||||
`[${requestId}] Upstream retries exhausted for ${toolId} (${envVarName}); re-queued and retrying once with fresh key`
|
||||
)
|
||||
return (await requeued()) as T
|
||||
}
|
||||
} catch (requeueError) {
|
||||
logger.error(
|
||||
`[${requestId}] Re-queue after exhausted upstream retries failed for ${toolId}`,
|
||||
{ error: toError(requeueError).message }
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
PlatformEvents.hostedKeyUserThrottled({
|
||||
toolId,
|
||||
reason: 'upstream_retries_exhausted',
|
||||
@@ -1099,6 +1173,16 @@ export async function executeTool(
|
||||
toolId,
|
||||
envVarName: hostedKeyInfo.envVarName!,
|
||||
executionContext,
|
||||
reacquireAfterRetriesExhausted: async () => {
|
||||
const reacquired = await reacquireHostedKey(
|
||||
tool,
|
||||
contextParams,
|
||||
executionContext,
|
||||
requestId
|
||||
)
|
||||
if (!reacquired) return null
|
||||
return () => executeToolRequest(toolId, tool, contextParams)
|
||||
},
|
||||
})
|
||||
: await executeToolRequest(toolId, tool, contextParams, signal)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user