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..895ebf33289 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,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) diff --git a/apps/sim/lib/embeddings/client.test.ts b/apps/sim/lib/embeddings/client.test.ts index 5824e3988db..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, @@ -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 () => { 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( {