diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts index 0307ed6f683..76d71ce4a81 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts @@ -1,6 +1,7 @@ /** * @vitest-environment node */ +import { resetEnvMock, setEnv } from '@sim/testing/mocks/env.mock' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' const { consumeTokens, getCooldownUntil, setCooldownUntil } = vi.hoisted(() => ({ @@ -35,7 +36,10 @@ describe('provider admission', () => { consumeTokens.mockResolvedValue({ allowed: true, tokensRemaining: 1, resetAt: new Date() }) }) - afterEach(() => vi.useRealTimers()) + afterEach(() => { + vi.useRealTimers() + resetEnvMock() + }) it('shares both credential dimensions in one reservation across concurrent callers', async () => { await Promise.all([waitForProviderAdmission(INPUT), waitForProviderAdmission(INPUT)]) @@ -86,6 +90,47 @@ describe('provider admission', () => { ) }) + it('caps bulk work below the aggregate budget so interactive callers keep headroom', async () => { + await waitForProviderAdmission({ ...INPUT, bulk: true }) + const [reservations, options] = consumeTokens.mock.calls[0] + expect(reservations).toMatchObject([ + { key: 'provider:embedding:openai:hashed-credential:tokens', config: { maxTokens: 600_000 } }, + { key: 'provider:embedding:openai:hashed-credential:requests', config: { maxTokens: 64 } }, + { + key: 'provider:embedding:openai:hashed-credential:bulk:tokens', + cost: 50, + config: { maxTokens: 540_000, refillRate: 9_000 }, + }, + { + key: 'provider:embedding:openai:hashed-credential:bulk:requests', + config: { maxTokens: 57, refillRate: 9 }, + }, + ]) + expect(options.cooldownKeys).toEqual([ + 'provider:embedding:openai:hashed-credential:cooldown', + 'provider:embedding:openai:hashed-credential:quota', + ]) + }) + + it('rejects a bulk batch the lane can never hold and keeps one request slot at a minimal burst', async () => { + setEnv({ + KB_CONFIG_EMBEDDING_REQUESTS_PER_MINUTE: '1', + KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE: '100', + }) + await expect( + waitForProviderAdmission({ ...INPUT, inputTokens: 95, bulk: true }) + ).rejects.toThrow('exceeds the configured per-credential token budget') + await waitForProviderAdmission({ ...INPUT, inputTokens: 95 }) + await waitForProviderAdmission({ ...INPUT, inputTokens: 90, bulk: true }) + expect(consumeTokens.mock.calls[1][0].slice(2)).toMatchObject([ + { key: 'provider:embedding:openai:hashed-credential:bulk:tokens', config: { maxTokens: 90 } }, + { + key: 'provider:embedding:openai:hashed-credential:bulk:requests', + config: { maxTokens: 1 }, + }, + ]) + }) + it('isolates another credential and does not impose token costs on OCR', async () => { await waitForProviderAdmission({ ...INPUT, diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.ts b/apps/sim/lib/core/rate-limiter/provider-admission.ts index 895ebf33289..b9b191c2f81 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.ts @@ -12,10 +12,20 @@ export interface ProviderIdentity { operation: 'embedding' | 'ocr' | 'rerank' } +/** + * Share of a credential's budget the bulk lane may use. Every caller reserves + * from the aggregate buckets, so the budget is never exceeded; bulk callers + * also reserve from buckets capped at this share, which leaves an interactive + * caller headroom instead of a queue behind a crawl's batches. + */ +const BULK_LANE_SHARE = 0.9 + interface ProviderAdmissionInput extends ProviderIdentity { inputTokens?: number signal?: AbortSignal maxWaitMs: number + /** Bulk work is capped at {@link BULK_LANE_SHARE}; cooldown and quota gates still stop every caller. */ + bulk?: boolean } /** @@ -55,36 +65,47 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P : input.operation === 'ocr' ? envNumber(env.KB_CONFIG_OCR_REQUESTS_PER_MINUTE, 60, { min: 1 }) : envNumber(env.KB_CONFIG_RERANK_REQUESTS_PER_MINUTE, 60, { min: 1 }) + const tokenBudget = + input.operation === 'embedding' && input.inputTokens + ? { + cost: input.inputTokens, + perMinute: envNumber(env.KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE, 600_000, { min: 1 }), + } + : undefined + const laneShare = input.bulk ? BULK_LANE_SHARE : 1 + if (tokenBudget && tokenBudget.cost > Math.floor(tokenBudget.perMinute * laneShare)) { + throw new Error('Embedding request exceeds the configured per-credential token budget') + } + const requestBurst = Math.min( + input.operation === 'embedding' ? EMBEDDING_REQUEST_BURST : DEFAULT_REQUEST_BURST, + requestsPerMinute + ) const reservations: TokenBucketReservation[] = [] - if (input.operation === 'embedding' && input.inputTokens) { - const tokensPerMinute = envNumber(env.KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE, 600_000, { - min: 1, - }) - if (input.inputTokens > tokensPerMinute) { - throw new Error('Embedding request exceeds the configured per-credential token budget') + const reserveBuckets = (bucketKey: string, share: number) => { + if (tokenBudget) { + reservations.push({ + key: `${bucketKey}:tokens`, + cost: tokenBudget.cost, + config: { + maxTokens: Math.floor(tokenBudget.perMinute * share), + refillRate: (tokenBudget.perMinute * share) / 60, + refillIntervalMs: 1000, + }, + }) } reservations.push({ - key: `${key}:tokens`, - cost: input.inputTokens, + key: `${bucketKey}:requests`, + cost: 1, config: { - maxTokens: tokensPerMinute, - refillRate: tokensPerMinute / 60, + /** A burst of one leaves no share to carve out, so the lane then matches the aggregate. */ + maxTokens: Math.max(1, Math.floor(requestBurst * share)), + refillRate: (requestsPerMinute * share) / 60, refillIntervalMs: 1000, }, }) } - reservations.push({ - key: `${key}:requests`, - cost: 1, - config: { - maxTokens: Math.min( - input.operation === 'embedding' ? EMBEDDING_REQUEST_BURST : DEFAULT_REQUEST_BURST, - requestsPerMinute - ), - refillRate: requestsPerMinute / 60, - refillIntervalMs: 1000, - }, - }) + reserveBuckets(key, 1) + if (input.bulk) reserveBuckets(`${key}:bulk`, BULK_LANE_SHARE) /** When the bucket last said capacity returns, so a deadline hit after a sleep reports the wait still left. */ let capacityAvailableAt: number | undefined diff --git a/apps/sim/lib/embeddings/client.test.ts b/apps/sim/lib/embeddings/client.test.ts index 70ff9e97646..5efc3104276 100644 --- a/apps/sim/lib/embeddings/client.test.ts +++ b/apps/sim/lib/embeddings/client.test.ts @@ -1807,13 +1807,14 @@ describe('durable embedding batches', () => { expect(KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS).toBeLessThan(EMBEDDING_RETRY_BUDGET_MS) }) - it('limits checkpointed admission waits while retaining the interactive request budget', async () => { + it('limits checkpointed admission waits and keeps interactive callers off the bulk lane', async () => { fetchMock.mockImplementation(() => Promise.resolve(jsonResponse(openAIBody([[1]], 7)))) await embed(['text'], { apiKey: 'fixture-key', checkpoints: memoryCheckpoints() }) expect(mockAdmit).toHaveBeenLastCalledWith( - expect.objectContaining({ maxWaitMs: KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS }) + expect.objectContaining({ maxWaitMs: KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS, bulk: true }) ) await embed(['text'], { apiKey: 'fixture-key' }) + expect(mockAdmit).toHaveBeenLastCalledWith(expect.objectContaining({ bulk: false })) expect(mockAdmit.mock.lastCall?.[0].maxWaitMs).toBeGreaterThan( KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS ) diff --git a/apps/sim/lib/embeddings/client.ts b/apps/sim/lib/embeddings/client.ts index 0c6728d35e7..846ba759f0d 100644 --- a/apps/sim/lib/embeddings/client.ts +++ b/apps/sim/lib/embeddings/client.ts @@ -544,8 +544,10 @@ async function callEmbeddingAPI( expectedDimensions: number | undefined, isBYOK: boolean, signal?: AbortSignal, - admissionWaitMs = EMBEDDING_RETRY_BUDGET_MS + /** Bulk indexing waits briefly and is capped below the credential budget; everything else has a person waiting on it. */ + bulk = false ): Promise<{ embeddings: number[][]; totalTokens: number; dimensions: number }> { + const admissionWaitMs = bulk ? KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS : EMBEDDING_RETRY_BUDGET_MS const admissionIdentity = embeddingAdmissionIdentity({ providerId, quotaCircuitIdentity, isBYOK }) return retryWithExponentialBackoff( async (operationSignal, deadlineAt) => { @@ -563,6 +565,7 @@ async function callEmbeddingAPI( ), signal: operationSignal, maxWaitMs: Math.min(admissionWaitMs, Math.max(0, deadlineAt - Date.now())), + bulk, }) } catch (error) { if (error instanceof ProviderQuotaExhaustedError) @@ -795,6 +798,7 @@ async function mapEmbeddingBatches( return results.map((result) => result!.value) } +/** Checkpoints mark the bulk indexing path; every other caller is interactive. */ async function callCheckpointedEmbeddingBatch( batch: string[], batchIndex: number, @@ -848,7 +852,7 @@ async function callCheckpointedEmbeddingBatch( provider.dimensions, provider.isBYOK, signal, - checkpoints ? KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS : undefined + checkpoints !== undefined ) if (identity) await checkpoints!.save(identity, result, signal) return result