fix(knowledge): wait for embedding admission instead of deferring documents for an hour

apps/sim/lib/core/rate-limiter/provider-admission.test.ts

@@ -46,6 +46,8 @@ describe('provider admission', () => {
         'provider:embedding:openai:hashed-credential:requests',
       ])
       expect(reservations[0].cost).toBe(50)
+      /** Enough burst for every concurrent document to start a batch; the rate still governs throughput. */
+      expect(reservations[1].config).toMatchObject({ maxTokens: 64, refillRate: 10 })
       expect(options.cooldownKeys).toHaveLength(2)
     }
   })
@@ -92,7 +94,7 @@ describe('provider admission', () => {
     })
     expect(consumeTokens).toHaveBeenCalledOnce()
     expect(consumeTokens.mock.calls[0][0]).toMatchObject([
-      { key: 'provider:ocr:openai:another-key:requests' },
+      { key: 'provider:ocr:openai:another-key:requests', config: { maxTokens: 2 } },
     ])
   })
   it('retains the cooldown when an admission storage call consumes the remaining deadline', async () => {

apps/sim/lib/core/rate-limiter/provider-admission.ts

@@ -18,6 +18,15 @@ interface ProviderAdmissionInput extends ProviderIdentity {
   maxWaitMs: number
 }
 
+/**
+ * Requests admitted in the same instant per embedding credential. The
+ * per-minute rate still governs sustained throughput; the burst only decides
+ * how many concurrent documents can start a batch together instead of losing a
+ * race for a handful of slots while the token budget sits unused.
+ */
+const EMBEDDING_REQUEST_BURST = 64
+const DEFAULT_REQUEST_BURST = 2
+
 /** A local admission wait expired; the document scheduler may retry the work later. */
 export class ProviderAdmissionTimeoutError extends Error {
   readonly retryable = false
@@ -68,15 +77,26 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P
     key: `${key}:requests`,
     cost: 1,
     config: {
-      maxTokens: Math.min(input.operation === 'embedding' ? 8 : 2, requestsPerMinute),
+      maxTokens: Math.min(
+        input.operation === 'embedding' ? EMBEDDING_REQUEST_BURST : DEFAULT_REQUEST_BURST,
+        requestsPerMinute
+      ),
       refillRate: requestsPerMinute / 60,
       refillIntervalMs: 1000,
     },
   })
 
+  /** When the bucket last said capacity returns, so a deadline hit after a sleep reports the wait still left. */
+  let capacityAvailableAt: number | undefined
   for (;;) {
     input.signal?.throwIfAborted()
-    if (Date.now() >= deadlineAt) throw new ProviderAdmissionTimeoutError()
+    if (Date.now() >= deadlineAt) {
+      const remainingMs =
+        capacityAvailableAt === undefined ? undefined : capacityAvailableAt - Date.now()
+      throw new ProviderAdmissionTimeoutError(
+        remainingMs !== undefined && remainingMs > 0 ? remainingMs : undefined
+      )
+    }
     if (await isProviderQuotaExhausted(input))
       throw new ProviderQuotaExhaustedError(input.providerId)
     let result: AtomicAdmissionResult
@@ -98,6 +118,7 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P
     }
     if (result.allowed) return
     const waitMs = Math.max(1, result.retryAfterMs)
+    if (Number.isFinite(waitMs)) capacityAvailableAt = Date.now() + waitMs
     if (!Number.isFinite(waitMs) || waitMs >= deadlineAt - Date.now()) {
       if (await isProviderQuotaExhausted(input))
         throw new ProviderQuotaExhaustedError(input.providerId)

apps/sim/lib/embeddings/client.test.ts

@@ -10,6 +10,7 @@ import {
   assertKnowledgeEmbeddingCapacityForDeployment,
   clampEmbeddingConcurrency,
   EMBEDDING_MAX_RETRIES,
+  EMBEDDING_RETRY_BUDGET_MS,
   EmbeddingAPIError,
   EmbeddingOutputLimitError,
   EmbeddingQuotaExhaustedError,
@@ -19,6 +20,7 @@ import {
   isBYOKEmbeddingCredentialRejection,
   isEmbeddingQuotaExhaustion,
   isTransientEmbeddingError,
+  KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS,
   MAX_EMBEDDING_SUCCESS_RESPONSE_BYTES,
 } from '@/lib/embeddings/client'
 
@@ -1801,12 +1803,20 @@ describe('durable embedding batches', () => {
     expect(fetchMock).not.toHaveBeenCalled()
   })
 
+  it('keeps the checkpointed admission wait inside the retry budget the processing deadline reserves', () => {
+    expect(KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS).toBeLessThan(EMBEDDING_RETRY_BUDGET_MS)
+  })
+
   it('limits checkpointed admission waits while retaining the interactive request budget', async () => {
     fetchMock.mockImplementation(() => Promise.resolve(jsonResponse(openAIBody([[1]], 7))))
     await embed(['text'], { apiKey: 'fixture-key', checkpoints: memoryCheckpoints() })
-    expect(mockAdmit).toHaveBeenLastCalledWith(expect.objectContaining({ maxWaitMs: 5000 }))
+    expect(mockAdmit).toHaveBeenLastCalledWith(
+      expect.objectContaining({ maxWaitMs: KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS })
+    )
     await embed(['text'], { apiKey: 'fixture-key' })
-    expect(mockAdmit.mock.lastCall?.[0].maxWaitMs).toBeGreaterThan(5000)
+    expect(mockAdmit.mock.lastCall?.[0].maxWaitMs).toBeGreaterThan(
+      KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS
+    )
   })
 
   it('drains admitted batches, resumes only missing requests and retains the complete token charge', async () => {

apps/sim/lib/embeddings/client.ts

@@ -147,7 +147,17 @@ export const EMBEDDING_MAX_RETRY_DELAY_MS = 30_000
  * is honored in full when it fits inside this deadline.
  */
 export const EMBEDDING_RETRY_BUDGET_MS = EMBEDDING_MAX_RETRIES * EMBEDDING_MAX_RETRY_DELAY_MS
-const KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS = 5000
+
+/**
+ * How long a checkpointed indexing batch waits for the shared admission bucket
+ * before the document yields its slot. Twenty concurrent documents fanning out
+ * eight batches each can queue for a couple of minutes behind the configured
+ * per-minute budget; yielding after a few seconds turned every such wait into a
+ * full re-dispatch with a minute-or-more delay. A minute of idle waiting is far
+ * cheaper than that round trip, and the per-request retry budget still bounds
+ * the whole attempt. Interactive callers keep the full request budget.
+ */
+export const KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS = 60_000
 
 export class EmbeddingAPIError extends Error {
   public status: number

apps/sim/lib/knowledge/documents/processing-provider-continuation.test.ts

@@ -15,6 +15,7 @@ import {
   MAX_PROCESSING_CONTINUATION_SLICES,
   MAX_PROVIDER_CONTINUATION_AGE_MS,
   MAX_PROVIDER_CONTINUATION_ATTEMPTS,
+  resolveAdmissionContinuationDelayMs,
   resolveProviderContinuationDelayMs,
   scheduleDocumentProcessingProviderContinuation,
 } from '@/lib/knowledge/documents/processing-provider-continuation'
@@ -135,6 +136,55 @@ describe('durable provider continuations', () => {
     })
   })
 
+  it('resumes soon after the local admission bucket turned a batch away, without spending a provider attempt', async () => {
+    const payload = {
+      ...PAYLOAD,
+      providerRetryCount: 2,
+      processingSliceCount: 7,
+      providerRetryStartedAt: new Date(NOW.getTime() - 3_600_000).toISOString(),
+    }
+    const continuation = await scheduleDocumentProcessingProviderContinuation(
+      payload,
+      new ProviderCapacityDeferredError('admission_timeout', { retryAfterMs: 3_000 }),
+      false
+    )
+    const delay = continuation.deferredUntil.getTime() - NOW.getTime()
+    expect(delay).toBeGreaterThanOrEqual(8_000)
+    expect(delay).toBeLessThanOrEqual(12_000)
+    expect(continuation.processingQueueToken).toBe('knowledge-slice-doc-1-pass-1-8')
+    expect(assertDocumentProcessingPayload(dispatch.mock.calls[0][0])).toMatchObject({
+      providerRetryCount: 2,
+      processingSliceCount: 8,
+      providerRetryStartedAt: payload.providerRetryStartedAt,
+    })
+  })
+
+  it('keeps the exponential ladder for provider-side throttling', async () => {
+    const continuation = await scheduleDocumentProcessingProviderContinuation(
+      { ...PAYLOAD, providerRetryCount: 3 },
+      new ProviderCapacityDeferredError('rate_limit', { retryAfterMs: 3_000 }),
+      false
+    )
+    const delay = continuation.deferredUntil.getTime() - NOW.getTime()
+    expect(delay).toBeGreaterThanOrEqual(8 * 60_000 * 0.8)
+    expect(continuation.processingQueueToken).toBe('knowledge-provider-doc-1-pass-1-4')
+  })
+
+  it('bounds admission resumes independently of provider retries', async () => {
+    await expect(
+      scheduleDocumentProcessingProviderContinuation(
+        {
+          ...PAYLOAD,
+          processingSliceCount: MAX_PROCESSING_CONTINUATION_SLICES,
+          providerRetryStartedAt: NOW.toISOString(),
+        },
+        new ProviderCapacityDeferredError('admission_timeout'),
+        false
+      )
+    ).rejects.toBeInstanceOf(ProviderCapacityContinuationExhaustedError)
+    expect(dispatch).not.toHaveBeenCalled()
+  })
+
   it('starts the same bounded recovery horizon when the first continuation is a processing slice', async () => {
     await scheduleDocumentProcessingProviderContinuation(
       PAYLOAD,
@@ -187,6 +237,23 @@ describe('durable provider continuations', () => {
     ).rejects.toBe(error)
   })
 
+  it('clamps and jitters the admission bucket wait', () => {
+    for (let i = 0; i < 20; i++) {
+      const stated = resolveAdmissionContinuationDelayMs(30_000)
+      expect(stated).toBeGreaterThanOrEqual(24_000)
+      expect(stated).toBeLessThanOrEqual(36_000)
+      const floored = resolveAdmissionContinuationDelayMs(800)
+      expect(floored).toBeGreaterThanOrEqual(8_000)
+      expect(floored).toBeLessThanOrEqual(12_000)
+      const capped = resolveAdmissionContinuationDelayMs(10 * 60_000)
+      expect(capped).toBeGreaterThanOrEqual(48_000)
+      expect(capped).toBeLessThanOrEqual(72_000)
+      const missing = resolveAdmissionContinuationDelayMs(undefined)
+      expect(missing).toBeGreaterThanOrEqual(12_000)
+      expect(missing).toBeLessThanOrEqual(18_000)
+    }
+  })
+
   it('bounds jittered polling without reducing provider minimums', () => {
     expect(resolveProviderContinuationDelayMs(1)).toBeGreaterThanOrEqual(48_000)
     expect(resolveProviderContinuationDelayMs(1)).toBeLessThanOrEqual(72_000)

apps/sim/lib/knowledge/documents/processing-provider-continuation.ts

@@ -14,6 +14,16 @@ export const MAX_PROVIDER_CONTINUATION_ATTEMPTS = 48
 export const MAX_PROCESSING_CONTINUATION_SLICES = 512
 export const MAX_PROVIDER_CONTINUATION_AGE_MS = 24 * 60 * 60 * 1000
 const MAX_PROVIDER_CONTINUATION_DELAY_MS = 60 * 60 * 1000
+/**
+ * Bounds for resuming after the deployment's own admission bucket ran out of
+ * wait budget. The bucket states when capacity returns, but that estimate does
+ * not know about the other documents waiting on it, so it is floored to keep
+ * re-dispatches apart and capped so a document never sits idle for long while
+ * the provider itself is healthy.
+ */
+const ADMISSION_RETRY_MIN_MS = 10_000
+const ADMISSION_RETRY_MAX_MS = 60_000
+const ADMISSION_RETRY_DEFAULT_MS = 15_000
 
 /** Server-stated waits are a lower bound, including when they exceed the ordinary polling cap. */
 export function resolveProviderContinuationDelayMs(attempt: number, retryAfterMs?: number): number {
@@ -31,6 +41,20 @@ export function resolveProviderContinuationDelayMs(attempt: number, retryAfterMs
   )
 }
 
+/**
+ * Delay after the local admission bucket declined a batch: the bucket's stated
+ * wait, clamped, with jitter so the documents it turned away do not return in
+ * lockstep. Provider-side throttling keeps {@link resolveProviderContinuationDelayMs}.
+ */
+export function resolveAdmissionContinuationDelayMs(retryAfterMs?: number): number {
+  const stated =
+    retryAfterMs !== undefined && Number.isFinite(retryAfterMs) && retryAfterMs > 0
+      ? retryAfterMs
+      : ADMISSION_RETRY_DEFAULT_MS
+  const clamped = Math.min(Math.max(stated, ADMISSION_RETRY_MIN_MS), ADMISSION_RETRY_MAX_MS)
+  return Math.round(backoffWithJitter(1, null, { baseMs: clamped, maxMs: clamped }))
+}
+
 /** Defers capacity pressure without spending another document dispatch or changing billing identity. */
 export async function scheduleDocumentProcessingProviderContinuation(
   payload: DocumentProcessingPayload,
@@ -39,9 +63,16 @@ export async function scheduleDocumentProcessingProviderContinuation(
   predecessorAdmissionCharged = false
 ): Promise<DocumentProcessingContinuation> {
   const now = Date.now()
+  /**
+   * A processing slice and a local admission timeout both mean the provider is
+   * fine and the document simply needs another turn: neither spends one of the
+   * bounded provider-failure attempts, and both count against the slice budget.
+   */
   const isProcessingSlice = error.reason === 'processing_budget'
-  const providerRetryCount = (payload.providerRetryCount ?? 0) + (isProcessingSlice ? 0 : 1)
-  const processingSliceCount = (payload.processingSliceCount ?? 0) + (isProcessingSlice ? 1 : 0)
+  const isAdmissionTimeout = error.reason === 'admission_timeout'
+  const isLocalYield = isProcessingSlice || isAdmissionTimeout
+  const providerRetryCount = (payload.providerRetryCount ?? 0) + (isLocalYield ? 0 : 1)
+  const processingSliceCount = (payload.processingSliceCount ?? 0) + (isLocalYield ? 1 : 0)
   const providerRetryStartedAt = payload.providerRetryStartedAt ?? new Date(now).toISOString()
   /** Tokenless legacy payloads retain a conservative handoff delay because their predecessor cannot be adopted safely. */
   const deferredUntil = new Date(
@@ -50,7 +81,9 @@ export async function scheduleDocumentProcessingProviderContinuation(
         ? payload.processingQueueToken
           ? 1000
           : 60_000
-        : resolveProviderContinuationDelayMs(providerRetryCount, error.retryAfterMs))
+        : isAdmissionTimeout
+          ? resolveAdmissionContinuationDelayMs(error.retryAfterMs)
+          : resolveProviderContinuationDelayMs(providerRetryCount, error.retryAfterMs))
   )
   if (
     providerRetryCount > MAX_PROVIDER_CONTINUATION_ATTEMPTS ||
@@ -62,8 +95,8 @@ export async function scheduleDocumentProcessingProviderContinuation(
   }
   const processingQueueToken = createDocumentProcessingContinuationToken(
     payload,
-    isProcessingSlice ? 'slice' : 'provider',
-    isProcessingSlice ? processingSliceCount : providerRetryCount
+    isLocalYield ? 'slice' : 'provider',
+    isLocalYield ? processingSliceCount : providerRetryCount
   )
   await dispatchDocumentProcessingContinuation(
     {