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

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

Every embedding batch on the indexing path waited at most five seconds for
the deployment's shared admission bucket. Twenty concurrent documents fanning
out eight batches each queue for minutes behind the configured per-minute
budget, so under load most batches timed out, the document stopped, and it
was re-dispatched with a delay that started at a minute and doubled to an
hour. The bucket's own estimate of when capacity returns was only a floor
under that ladder. During a bulk sync this produced thousands of hour-long
deferrals for a limiter we run ourselves, while the provider was healthy.

The knowledge path now waits up to a minute for admission, which is cheaper
than the re-dispatch it replaces and still bounded by the per-request retry
budget. The request bucket admits 64 concurrent starts instead of 8, so
documents that begin together no longer lose a race for slots while the
token budget sits unused. When an admission wait still expires, the document
resumes after the bucket's stated wait, clamped to 10 to 60 seconds with
jitter, and the yield counts against the processing-slice budget rather than
the provider-failure attempts, since the provider did nothing wrong. The
deadline path now carries the bucket's last stated wait so that estimate is
available to the scheduler.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* fix(knowledge): report only the admission wait still left at the deadline

The bucket's stated wait is stored as an absolute instant so a deadline hit
after a sleep carries the remainder, not the original duration. A test pins
the knowledge admission wait below the retry budget the processing deadline
reserves for each request.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: waleedlatif1

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