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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JBacX6HGVhPMUySuMfANwn

Author: waleedlatif1

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 () => {

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

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' }))

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