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 <noreply@anthropic.com>

Author: waleedlatif1

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<string, string | undefined>,
-}))
-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 },

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,22 +72,22 @@ 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(
     input.operation === 'embedding' ? EMBEDDING_REQUEST_BURST : DEFAULT_REQUEST_BURST,
     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,14 +97,15 @@ 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,
       },
     })
   }
   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

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
     )

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<T, R>(
   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