From 82c9d397f7ad77f2c6aa92b36f7326fb1fbf059a Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 10 Sep 2026 02:26:01 -0700 Subject: [PATCH 1/4] fix(search): reserve interactive embeddings from their own admission lane A search embeds its query through the same per-credential admission bucket as bulk indexing, so during a crawl one short interactive request competed with hundreds of batches for the bucket's refill. Non-checkpointed callers now reserve from a separate lane under the same identity; the provider cooldown and quota gates stay shared, so a provider 429 or an exhausted balance still pauses every caller. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01JBacX6HGVhPMUySuMfANwn --- .../core/rate-limiter/provider-admission.test.ts | 13 +++++++++++++ .../lib/core/rate-limiter/provider-admission.ts | 14 ++++++++++++-- apps/sim/lib/embeddings/client.test.ts | 5 +++-- apps/sim/lib/embeddings/client.ts | 13 +++++++++++-- 4 files changed, 39 insertions(+), 6 deletions(-) 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..06437be37ed 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts @@ -86,6 +86,19 @@ describe('provider admission', () => { ) }) + it('reserves an interactive lane from its own buckets while sharing the provider gates', async () => { + await waitForProviderAdmission({ ...INPUT, lane: 'interactive' }) + const [reservations, options] = consumeTokens.mock.calls[0] + expect(reservations.map((item: { key: string }) => item.key)).toEqual([ + 'provider:embedding:openai:hashed-credential:interactive:tokens', + 'provider:embedding:openai:hashed-credential:interactive:requests', + ]) + expect(options.cooldownKeys).toEqual([ + 'provider:embedding:openai:hashed-credential:cooldown', + 'provider:embedding:openai:hashed-credential:quota', + ]) + }) + 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..98c76048951 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.ts @@ -12,10 +12,19 @@ export interface ProviderIdentity { operation: 'embedding' | 'ocr' | 'rerank' } +/** + * A lane reserves from its own request and token buckets under the same + * provider identity, so a person waiting on one embedding never queues behind + * a bulk crawl's batches. Cooldown and quota gates stay per identity: a + * provider pause or an exhausted balance still stops every lane. + */ +export type ProviderAdmissionLane = 'interactive' + interface ProviderAdmissionInput extends ProviderIdentity { inputTokens?: number signal?: AbortSignal maxWaitMs: number + lane?: ProviderAdmissionLane } /** @@ -49,6 +58,7 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P input.signal?.throwIfAborted() const deadlineAt = Date.now() + input.maxWaitMs const key = providerKey(input) + const bucketKey = input.lane ? `${key}:${input.lane}` : key const requestsPerMinute = input.operation === 'embedding' ? envNumber(env.KB_CONFIG_EMBEDDING_REQUESTS_PER_MINUTE, 600, { min: 1 }) @@ -64,7 +74,7 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P throw new Error('Embedding request exceeds the configured per-credential token budget') } reservations.push({ - key: `${key}:tokens`, + key: `${bucketKey}:tokens`, cost: input.inputTokens, config: { maxTokens: tokensPerMinute, @@ -74,7 +84,7 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P }) } reservations.push({ - key: `${key}:requests`, + key: `${bucketKey}:requests`, cost: 1, config: { maxTokens: Math.min( diff --git a/apps/sim/lib/embeddings/client.test.ts b/apps/sim/lib/embeddings/client.test.ts index 70ff9e97646..51b1375806a 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 on their own 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, lane: undefined }) ) await embed(['text'], { apiKey: 'fixture-key' }) + expect(mockAdmit).toHaveBeenLastCalledWith(expect.objectContaining({ lane: 'interactive' })) 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..071df4d76b1 100644 --- a/apps/sim/lib/embeddings/client.ts +++ b/apps/sim/lib/embeddings/client.ts @@ -11,6 +11,7 @@ import { } from '@/lib/core/config/env-capabilities' import { isHosted } from '@/lib/core/config/env-flags' import { + type ProviderAdmissionLane, ProviderQuotaExhaustedError, recordProviderCooldown, waitForProviderAdmission, @@ -544,7 +545,8 @@ async function callEmbeddingAPI( expectedDimensions: number | undefined, isBYOK: boolean, signal?: AbortSignal, - admissionWaitMs = EMBEDDING_RETRY_BUDGET_MS + admissionWaitMs = EMBEDDING_RETRY_BUDGET_MS, + lane?: ProviderAdmissionLane ): Promise<{ embeddings: number[][]; totalTokens: number; dimensions: number }> { const admissionIdentity = embeddingAdmissionIdentity({ providerId, quotaCircuitIdentity, isBYOK }) return retryWithExponentialBackoff( @@ -563,6 +565,7 @@ async function callEmbeddingAPI( ), signal: operationSignal, maxWaitMs: Math.min(admissionWaitMs, Math.max(0, deadlineAt - Date.now())), + lane, }) } catch (error) { if (error instanceof ProviderQuotaExhaustedError) @@ -848,7 +851,13 @@ async function callCheckpointedEmbeddingBatch( provider.dimensions, provider.isBYOK, signal, - checkpoints ? KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS : undefined + checkpoints ? KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS : undefined, + /** + * Checkpoints mark the bulk indexing path. Everything else has a person + * waiting on it, so it reserves from the interactive lane and never + * queues behind a crawl's batches. + */ + checkpoints ? undefined : 'interactive' ) if (identity) await checkpoints!.save(identity, result, signal) return result From fe50edb003a6fa05d6a55d4aeb0d6ae6df84f4ed Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 10 Sep 2026 02:39:19 -0700 Subject: [PATCH 2/4] fix(search): cap the bulk embedding lane below the credential budget Every caller reserves from the credential's aggregate buckets, and the bulk lane additionally reserves from a bucket capped at 90% of that budget. The aggregate can no longer exceed the configured budget, and interactive callers always find headroom instead of a queue behind a crawl's batches. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01JBacX6HGVhPMUySuMfANwn --- .../rate-limiter/provider-admission.test.ts | 31 +++++++-- .../core/rate-limiter/provider-admission.ts | 65 +++++++++++-------- apps/sim/lib/embeddings/client.test.ts | 4 +- apps/sim/lib/embeddings/client.ts | 8 +-- 4 files changed, 68 insertions(+), 40 deletions(-) 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 06437be37ed..a141856b512 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts @@ -86,17 +86,36 @@ describe('provider admission', () => { ) }) - it('reserves an interactive lane from its own buckets while sharing the provider gates', async () => { + it('caps the bulk lane below the aggregate budget so interactive callers keep headroom', async () => { + await waitForProviderAdmission({ ...INPUT, lane: 'bulk' }) await waitForProviderAdmission({ ...INPUT, lane: 'interactive' }) - const [reservations, options] = consumeTokens.mock.calls[0] - expect(reservations.map((item: { key: string }) => item.key)).toEqual([ - 'provider:embedding:openai:hashed-credential:interactive:tokens', - 'provider:embedding:openai:hashed-credential:interactive:requests', + const [bulkReservations, bulkOptions] = consumeTokens.mock.calls[0] + expect(bulkReservations).toMatchObject([ + { + key: 'provider:embedding:openai:hashed-credential:tokens', + config: { maxTokens: 600_000, refillRate: 10_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([ + expect(bulkOptions.cooldownKeys).toEqual([ 'provider:embedding:openai:hashed-credential:cooldown', 'provider:embedding:openai:hashed-credential:quota', ]) + const [interactiveReservations, interactiveOptions] = consumeTokens.mock.calls[1] + expect(interactiveReservations.map((item: { key: string }) => item.key)).toEqual([ + 'provider:embedding:openai:hashed-credential:tokens', + 'provider:embedding:openai:hashed-credential:requests', + ]) + expect(interactiveOptions.cooldownKeys).toEqual(bulkOptions.cooldownKeys) }) it('isolates another credential and does not impose token costs on OCR', async () => { diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.ts b/apps/sim/lib/core/rate-limiter/provider-admission.ts index 98c76048951..0aee45a1014 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.ts @@ -13,12 +13,16 @@ export interface ProviderIdentity { } /** - * A lane reserves from its own request and token buckets under the same - * provider identity, so a person waiting on one embedding never queues behind - * a bulk crawl's batches. Cooldown and quota gates stay per identity: a - * provider pause or an exhausted balance still stops every lane. + * Every caller reserves from the credential's aggregate buckets, so the + * configured budget is never exceeded. The bulk lane also reserves from a + * bucket capped at {@link BULK_LANE_SHARE} of that budget, which leaves an + * interactive caller headroom instead of a queue behind a crawl's batches. + * Cooldown and quota gates stay per identity: a provider pause or an exhausted + * balance still stops every lane. */ -export type ProviderAdmissionLane = 'interactive' +export type ProviderAdmissionLane = 'bulk' | 'interactive' + +const BULK_LANE_SHARE = 0.9 interface ProviderAdmissionInput extends ProviderIdentity { inputTokens?: number @@ -58,43 +62,48 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P input.signal?.throwIfAborted() const deadlineAt = Date.now() + input.maxWaitMs const key = providerKey(input) - const bucketKey = input.lane ? `${key}:${input.lane}` : key const requestsPerMinute = input.operation === 'embedding' ? envNumber(env.KB_CONFIG_EMBEDDING_REQUESTS_PER_MINUTE, 600, { min: 1 }) : 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 + if (tokenBudget && tokenBudget.cost > tokenBudget.perMinute) { + 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) { + const maxTokens = Math.floor(tokenBudget.perMinute * share) + reservations.push({ + key: `${bucketKey}:tokens`, + cost: tokenBudget.cost, + config: { maxTokens, refillRate: maxTokens / 60, refillIntervalMs: 1000 }, + }) } reservations.push({ - key: `${bucketKey}:tokens`, - cost: input.inputTokens, + key: `${bucketKey}:requests`, + cost: 1, config: { - maxTokens: tokensPerMinute, - refillRate: tokensPerMinute / 60, + maxTokens: Math.floor(requestBurst * share), + refillRate: (requestsPerMinute * share) / 60, refillIntervalMs: 1000, }, }) } - reservations.push({ - key: `${bucketKey}: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.lane === '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 51b1375806a..ac514d4bc39 100644 --- a/apps/sim/lib/embeddings/client.test.ts +++ b/apps/sim/lib/embeddings/client.test.ts @@ -1807,11 +1807,11 @@ describe('durable embedding batches', () => { expect(KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS).toBeLessThan(EMBEDDING_RETRY_BUDGET_MS) }) - it('limits checkpointed admission waits and keeps interactive callers on their own lane', 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, lane: undefined }) + expect.objectContaining({ maxWaitMs: KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS, lane: 'bulk' }) ) await embed(['text'], { apiKey: 'fixture-key' }) expect(mockAdmit).toHaveBeenLastCalledWith(expect.objectContaining({ lane: 'interactive' })) diff --git a/apps/sim/lib/embeddings/client.ts b/apps/sim/lib/embeddings/client.ts index 071df4d76b1..e8901ec1283 100644 --- a/apps/sim/lib/embeddings/client.ts +++ b/apps/sim/lib/embeddings/client.ts @@ -853,11 +853,11 @@ async function callCheckpointedEmbeddingBatch( signal, checkpoints ? KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS : undefined, /** - * Checkpoints mark the bulk indexing path. Everything else has a person - * waiting on it, so it reserves from the interactive lane and never - * queues behind a crawl's batches. + * Checkpoints mark the bulk indexing path, which may never take the whole + * credential budget. Everything else has a person waiting on it and uses + * the headroom the bulk lane leaves. */ - checkpoints ? undefined : 'interactive' + checkpoints ? 'bulk' : 'interactive' ) if (identity) await checkpoints!.save(identity, result, signal) return result From 2370b873af3b3672990f745475ba63d5bcfc4f0e Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 10 Sep 2026 02:45:01 -0700 Subject: [PATCH 3/4] fix(search): keep a bulk lane bucket large enough for one valid reservation At the smallest supported budgets the 90% share floored below a single request or token cost and locked the bulk lane. The bucket now holds at least the reservation it is asked for, while the refill rate keeps the share. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01JBacX6HGVhPMUySuMfANwn --- .../rate-limiter/provider-admission.test.ts | 28 ++++++++++++++++++- .../core/rate-limiter/provider-admission.ts | 10 +++++-- 2 files changed, 34 insertions(+), 4 deletions(-) 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 a141856b512..e1077134a38 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts @@ -3,10 +3,16 @@ */ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -const { consumeTokens, getCooldownUntil, setCooldownUntil } = vi.hoisted(() => ({ +const { consumeTokens, getCooldownUntil, setCooldownUntil, mockEnv } = vi.hoisted(() => ({ consumeTokens: vi.fn(), getCooldownUntil: vi.fn(), setCooldownUntil: vi.fn(), + mockEnv: {} as Record, +})) +vi.mock('@/lib/core/config/env', () => ({ + env: mockEnv, + envNumber: (value: string | undefined, fallback: number) => + value === undefined ? fallback : Number(value), })) vi.mock('@/lib/core/rate-limiter/storage/factory', () => ({ createStorageAdapter: () => ({ @@ -118,6 +124,26 @@ describe('provider admission', () => { expect(interactiveOptions.cooldownKeys).toEqual(bulkOptions.cooldownKeys) }) + it('never shrinks a bulk bucket below one valid reservation', async () => { + mockEnv.KB_CONFIG_EMBEDDING_REQUESTS_PER_MINUTE = '1' + mockEnv.KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE = '100' + try { + await waitForProviderAdmission({ ...INPUT, inputTokens: 95, lane: 'bulk' }) + } finally { + mockEnv.KB_CONFIG_EMBEDDING_REQUESTS_PER_MINUTE = undefined + mockEnv.KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE = undefined + } + expect(consumeTokens.mock.calls[0][0]).toMatchObject([ + { key: 'provider:embedding:openai:hashed-credential:tokens', config: { maxTokens: 100 } }, + { key: 'provider:embedding:openai:hashed-credential:requests', config: { maxTokens: 1 } }, + { key: 'provider:embedding:openai:hashed-credential:bulk:tokens', config: { maxTokens: 95 } }, + { + 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 0aee45a1014..131c9ef1150 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.ts @@ -83,20 +83,24 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P requestsPerMinute ) const reservations: TokenBucketReservation[] = [] + /** A lane bucket always holds at least one valid reservation, so a tiny budget cannot lock the lane. */ const reserveBuckets = (bucketKey: string, share: number) => { if (tokenBudget) { - const maxTokens = Math.floor(tokenBudget.perMinute * share) reservations.push({ key: `${bucketKey}:tokens`, cost: tokenBudget.cost, - config: { maxTokens, refillRate: maxTokens / 60, refillIntervalMs: 1000 }, + config: { + maxTokens: Math.max(tokenBudget.cost, Math.floor(tokenBudget.perMinute * share)), + refillRate: (tokenBudget.perMinute * share) / 60, + refillIntervalMs: 1000, + }, }) } reservations.push({ key: `${bucketKey}:requests`, cost: 1, config: { - maxTokens: Math.floor(requestBurst * share), + maxTokens: Math.max(1, Math.floor(requestBurst * share)), refillRate: (requestsPerMinute * share) / 60, refillIntervalMs: 1000, }, From 8ef98ea55894117d4a51336748e0ef270b9bd792 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 10 Sep 2026 02:49:41 -0700 Subject: [PATCH 4/4] refactor(search): derive the bulk admission lane from one flag A single bulk flag now selects the shorter admission wait and the capped lane, the lane cap is a pure function of configuration, and a bulk batch larger than that cap is rejected up front instead of resizing a shared bucket. Tests use the shared env mock. Co-Authored-By: Claude Fable 5.1 --- .../rate-limiter/provider-admission.test.ts | 61 ++++++++----------- .../core/rate-limiter/provider-admission.ts | 24 ++++---- apps/sim/lib/embeddings/client.test.ts | 4 +- apps/sim/lib/embeddings/client.ts | 17 ++---- 4 files changed, 43 insertions(+), 63 deletions(-) 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 e1077134a38..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,18 +1,13 @@ /** * @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, mockEnv } = vi.hoisted(() => ({ +const { consumeTokens, getCooldownUntil, setCooldownUntil } = vi.hoisted(() => ({ consumeTokens: vi.fn(), getCooldownUntil: vi.fn(), setCooldownUntil: vi.fn(), - mockEnv: {} as Record, -})) -vi.mock('@/lib/core/config/env', () => ({ - env: mockEnv, - envNumber: (value: string | undefined, fallback: number) => - value === undefined ? fallback : Number(value), })) vi.mock('@/lib/core/rate-limiter/storage/factory', () => ({ createStorageAdapter: () => ({ @@ -41,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)]) @@ -92,15 +90,11 @@ describe('provider admission', () => { ) }) - it('caps the bulk lane below the aggregate budget so interactive callers keep headroom', async () => { - await waitForProviderAdmission({ ...INPUT, lane: 'bulk' }) - await waitForProviderAdmission({ ...INPUT, lane: 'interactive' }) - const [bulkReservations, bulkOptions] = consumeTokens.mock.calls[0] - expect(bulkReservations).toMatchObject([ - { - key: 'provider:embedding:openai:hashed-credential:tokens', - config: { maxTokens: 600_000, refillRate: 10_000 }, - }, + 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', @@ -112,31 +106,24 @@ describe('provider admission', () => { config: { maxTokens: 57, refillRate: 9 }, }, ]) - expect(bulkOptions.cooldownKeys).toEqual([ + expect(options.cooldownKeys).toEqual([ 'provider:embedding:openai:hashed-credential:cooldown', 'provider:embedding:openai:hashed-credential:quota', ]) - const [interactiveReservations, interactiveOptions] = consumeTokens.mock.calls[1] - expect(interactiveReservations.map((item: { key: string }) => item.key)).toEqual([ - 'provider:embedding:openai:hashed-credential:tokens', - 'provider:embedding:openai:hashed-credential:requests', - ]) - expect(interactiveOptions.cooldownKeys).toEqual(bulkOptions.cooldownKeys) }) - it('never shrinks a bulk bucket below one valid reservation', async () => { - mockEnv.KB_CONFIG_EMBEDDING_REQUESTS_PER_MINUTE = '1' - mockEnv.KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE = '100' - try { - await waitForProviderAdmission({ ...INPUT, inputTokens: 95, lane: 'bulk' }) - } finally { - mockEnv.KB_CONFIG_EMBEDDING_REQUESTS_PER_MINUTE = undefined - mockEnv.KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE = undefined - } - expect(consumeTokens.mock.calls[0][0]).toMatchObject([ - { key: 'provider:embedding:openai:hashed-credential:tokens', config: { maxTokens: 100 } }, - { key: 'provider:embedding:openai:hashed-credential:requests', config: { maxTokens: 1 } }, - { key: 'provider:embedding:openai:hashed-credential:bulk:tokens', config: { maxTokens: 95 } }, + 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 }, diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.ts b/apps/sim/lib/core/rate-limiter/provider-admission.ts index 131c9ef1150..b9b191c2f81 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.ts @@ -13,22 +13,19 @@ export interface ProviderIdentity { } /** - * Every caller reserves from the credential's aggregate buckets, so the - * configured budget is never exceeded. The bulk lane also reserves from a - * bucket capped at {@link BULK_LANE_SHARE} of that budget, which leaves an - * interactive caller headroom instead of a queue behind a crawl's batches. - * Cooldown and quota gates stay per identity: a provider pause or an exhausted - * balance still stops every lane. + * 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. */ -export type ProviderAdmissionLane = 'bulk' | 'interactive' - const BULK_LANE_SHARE = 0.9 interface ProviderAdmissionInput extends ProviderIdentity { inputTokens?: number signal?: AbortSignal maxWaitMs: number - lane?: ProviderAdmissionLane + /** Bulk work is capped at {@link BULK_LANE_SHARE}; cooldown and quota gates still stop every caller. */ + bulk?: boolean } /** @@ -75,7 +72,8 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P perMinute: envNumber(env.KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE, 600_000, { min: 1 }), } : undefined - if (tokenBudget && tokenBudget.cost > tokenBudget.perMinute) { + 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( @@ -83,14 +81,13 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P requestsPerMinute ) const reservations: TokenBucketReservation[] = [] - /** A lane bucket always holds at least one valid reservation, so a tiny budget cannot lock the lane. */ const reserveBuckets = (bucketKey: string, share: number) => { if (tokenBudget) { reservations.push({ key: `${bucketKey}:tokens`, cost: tokenBudget.cost, config: { - maxTokens: Math.max(tokenBudget.cost, Math.floor(tokenBudget.perMinute * share)), + maxTokens: Math.floor(tokenBudget.perMinute * share), refillRate: (tokenBudget.perMinute * share) / 60, refillIntervalMs: 1000, }, @@ -100,6 +97,7 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P key: `${bucketKey}:requests`, cost: 1, config: { + /** 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, @@ -107,7 +105,7 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P }) } reserveBuckets(key, 1) - if (input.lane === 'bulk') reserveBuckets(`${key}:bulk`, BULK_LANE_SHARE) + 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 ac514d4bc39..5efc3104276 100644 --- a/apps/sim/lib/embeddings/client.test.ts +++ b/apps/sim/lib/embeddings/client.test.ts @@ -1811,10 +1811,10 @@ describe('durable embedding batches', () => { 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, lane: 'bulk' }) + expect.objectContaining({ maxWaitMs: KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS, bulk: true }) ) await embed(['text'], { apiKey: 'fixture-key' }) - expect(mockAdmit).toHaveBeenLastCalledWith(expect.objectContaining({ lane: 'interactive' })) + 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 e8901ec1283..846ba759f0d 100644 --- a/apps/sim/lib/embeddings/client.ts +++ b/apps/sim/lib/embeddings/client.ts @@ -11,7 +11,6 @@ import { } from '@/lib/core/config/env-capabilities' import { isHosted } from '@/lib/core/config/env-flags' import { - type ProviderAdmissionLane, ProviderQuotaExhaustedError, recordProviderCooldown, waitForProviderAdmission, @@ -545,9 +544,10 @@ async function callEmbeddingAPI( expectedDimensions: number | undefined, isBYOK: boolean, signal?: AbortSignal, - admissionWaitMs = EMBEDDING_RETRY_BUDGET_MS, - lane?: ProviderAdmissionLane + /** 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) => { @@ -565,7 +565,7 @@ async function callEmbeddingAPI( ), signal: operationSignal, maxWaitMs: Math.min(admissionWaitMs, Math.max(0, deadlineAt - Date.now())), - lane, + bulk, }) } catch (error) { if (error instanceof ProviderQuotaExhaustedError) @@ -798,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, @@ -851,13 +852,7 @@ async function callCheckpointedEmbeddingBatch( provider.dimensions, provider.isBYOK, signal, - checkpoints ? KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS : undefined, - /** - * Checkpoints mark the bulk indexing path, which may never take the whole - * credential budget. Everything else has a person waiting on it and uses - * the headroom the bulk lane leaves. - */ - checkpoints ? 'bulk' : 'interactive' + checkpoints !== undefined ) if (identity) await checkpoints!.save(identity, result, signal) return result