From dba94ddfb6c26be5e9c6af7466921988b54db9f0 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 9 Sep 2026 23:53:58 -0700 Subject: [PATCH 1/2] 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 --- .../rate-limiter/provider-admission.test.ts | 4 +- .../core/rate-limiter/provider-admission.ts | 19 +++++- apps/sim/lib/embeddings/client.test.ts | 9 ++- apps/sim/lib/embeddings/client.ts | 12 +++- .../processing-provider-continuation.test.ts | 67 +++++++++++++++++++ .../processing-provider-continuation.ts | 43 ++++++++++-- 6 files changed, 143 insertions(+), 11 deletions(-) diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts index 1c7a0d01b2a..0307ed6f683 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts +++ b/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 () => { diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.ts b/apps/sim/lib/core/rate-limiter/provider-admission.ts index 4cb66586952..fefcc24eb32 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.ts +++ b/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,20 @@ 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, }, }) + /** The bucket's last stated wait, so a deadline hit between polls still reports when capacity returns. */ + let lastRetryAfterMs: number | undefined for (;;) { input.signal?.throwIfAborted() - if (Date.now() >= deadlineAt) throw new ProviderAdmissionTimeoutError() + if (Date.now() >= deadlineAt) throw new ProviderAdmissionTimeoutError(lastRetryAfterMs) if (await isProviderQuotaExhausted(input)) throw new ProviderQuotaExhaustedError(input.providerId) let result: AtomicAdmissionResult @@ -98,6 +112,7 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P } if (result.allowed) return const waitMs = Math.max(1, result.retryAfterMs) + if (Number.isFinite(waitMs)) lastRetryAfterMs = waitMs if (!Number.isFinite(waitMs) || waitMs >= deadlineAt - Date.now()) { if (await isProviderQuotaExhausted(input)) throw new ProviderQuotaExhaustedError(input.providerId) diff --git a/apps/sim/lib/embeddings/client.test.ts b/apps/sim/lib/embeddings/client.test.ts index 5824e3988db..52b80a5b53c 100644 --- a/apps/sim/lib/embeddings/client.test.ts +++ b/apps/sim/lib/embeddings/client.test.ts @@ -19,6 +19,7 @@ import { isBYOKEmbeddingCredentialRejection, isEmbeddingQuotaExhaustion, isTransientEmbeddingError, + KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS, MAX_EMBEDDING_SUCCESS_RESPONSE_BYTES, } from '@/lib/embeddings/client' @@ -1804,9 +1805,13 @@ describe('durable embedding batches', () => { 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 () => { diff --git a/apps/sim/lib/embeddings/client.ts b/apps/sim/lib/embeddings/client.ts index 14406a4e50c..0c6728d35e7 100644 --- a/apps/sim/lib/embeddings/client.ts +++ b/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 diff --git a/apps/sim/lib/knowledge/documents/processing-provider-continuation.test.ts b/apps/sim/lib/knowledge/documents/processing-provider-continuation.test.ts index 10a275a344d..b8c13fa16e5 100644 --- a/apps/sim/lib/knowledge/documents/processing-provider-continuation.test.ts +++ b/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) diff --git a/apps/sim/lib/knowledge/documents/processing-provider-continuation.ts b/apps/sim/lib/knowledge/documents/processing-provider-continuation.ts index 266194bb324..3dcccb4a060 100644 --- a/apps/sim/lib/knowledge/documents/processing-provider-continuation.ts +++ b/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 { 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( { From 83a1c7d328bb9c99a67d5a037be1318099b47927 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 10 Sep 2026 00:05:54 -0700 Subject: [PATCH 2/2] 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 --- .../lib/core/rate-limiter/provider-admission.ts | 14 ++++++++++---- apps/sim/lib/embeddings/client.test.ts | 5 +++++ 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.ts b/apps/sim/lib/core/rate-limiter/provider-admission.ts index fefcc24eb32..895ebf33289 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.ts @@ -86,11 +86,17 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P }, }) - /** The bucket's last stated wait, so a deadline hit between polls still reports when capacity returns. */ - let lastRetryAfterMs: number | undefined + /** 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(lastRetryAfterMs) + 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 @@ -112,7 +118,7 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P } if (result.allowed) return const waitMs = Math.max(1, result.retryAfterMs) - if (Number.isFinite(waitMs)) lastRetryAfterMs = waitMs + 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) diff --git a/apps/sim/lib/embeddings/client.test.ts b/apps/sim/lib/embeddings/client.test.ts index 52b80a5b53c..70ff9e97646 100644 --- a/apps/sim/lib/embeddings/client.test.ts +++ b/apps/sim/lib/embeddings/client.test.ts @@ -10,6 +10,7 @@ import { assertKnowledgeEmbeddingCapacityForDeployment, clampEmbeddingConcurrency, EMBEDDING_MAX_RETRIES, + EMBEDDING_RETRY_BUDGET_MS, EmbeddingAPIError, EmbeddingOutputLimitError, EmbeddingQuotaExhaustedError, @@ -1802,6 +1803,10 @@ 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() })