fix(search): reserve interactive embeddings from their own admission lane

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,

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

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
     )

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<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,
@@ -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