diff --git a/apps/sim/lib/core/config/env.ts b/apps/sim/lib/core/config/env.ts index 315a9bb4349..8d9d784bd32 100644 --- a/apps/sim/lib/core/config/env.ts +++ b/apps/sim/lib/core/config/env.ts @@ -390,8 +390,11 @@ export const env = createEnv({ KB_CONFIG_RETRY_FACTOR: z.number().optional().default(2), // Retry backoff factor KB_CONFIG_MIN_TIMEOUT: z.number().optional().default(1000), // Min timeout in ms KB_CONFIG_MAX_TIMEOUT: z.number().optional().default(10000), // Max timeout in ms - KB_CONFIG_CONCURRENCY_LIMIT: z.number().optional().default(50), // Concurrent embedding API calls + KB_CONFIG_CONCURRENCY_LIMIT: z.number().optional().default(20), // Concurrent document-processing task runs (Trigger.dev queue depth) + KB_CONFIG_EMBEDDING_CONCURRENCY: z.number().optional().default(8), // Concurrent embedding API requests within one embed call + KB_CONFIG_DOCUMENT_CONCURRENCY: z.number().optional().default(4), // Concurrent documents in the in-process (non-Trigger) path KB_CONFIG_BATCH_SIZE: z.number().optional().default(2000), // Chunks to process per embedding batch + KB_CONFIG_DOCUMENT_BATCH_SIZE: z.number().optional().default(10), // Documents per batch in the in-process (non-Trigger) path KB_CONFIG_DELAY_BETWEEN_BATCHES: z.number().optional().default(0), // Delay between batches in ms (0 for max speed) KB_CONFIG_DELAY_BETWEEN_DOCUMENTS: z.number().optional().default(50), // Delay between documents in ms KB_CONFIG_CHUNK_CONCURRENCY: z.number().optional().default(10), // Concurrent PDF chunk OCR processing @@ -727,6 +730,13 @@ export { getEnv } * `z.number()` arrive as raw strings when sourced from `process.env` or Helm. * Use this helper anywhere a numeric env override is consumed to normalize the * type at the boundary instead of relying on JS implicit coercion. + * + * Skipping validation also means the schema never runs, so a `.default(...)` in + * the declaration above never executes: **the fallback passed here is the real + * default**, and the declared one is documentation. Keep the two in agreement — + * a variable read in more than one place with a different fallback each time has + * no single default at all, which is how one knob came to set both the + * document-processing queue depth and the embedding request fan-out. */ export function envNumber( value: number | string | undefined | null, diff --git a/apps/sim/lib/embeddings/client.test.ts b/apps/sim/lib/embeddings/client.test.ts index d2931d9fe26..0f2d4e66eaa 100644 --- a/apps/sim/lib/embeddings/client.test.ts +++ b/apps/sim/lib/embeddings/client.test.ts @@ -4,6 +4,7 @@ import { resetEnvMock, setEnv } from '@sim/testing' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { + EMBEDDING_MAX_RETRIES, EmbeddingAPIError, embed, embedKnowledgeForDeployment, @@ -28,11 +29,13 @@ vi.mock('@/lib/api-key/byok', () => ({ const originalFetch = global.fetch -function jsonResponse(body: unknown, status = 200): Response { +function jsonResponse(body: unknown, status = 200, responseHeaders?: HeadersInit): Response { return { ok: status >= 200 && status < 300, status, statusText: String(status), + // A real Response always carries these; the failure path reads them for rate-limit signals. + headers: new Headers(responseHeaders), json: async () => body, text: async () => JSON.stringify(body), } as Response @@ -677,11 +680,12 @@ describe('knowledge embedding transport fallback', () => { await vi.runAllTimersAsync() const result = await pending - expect(fetchMock).toHaveBeenCalledTimes(5) - expect(fetchMock.mock.calls.slice(0, 4).every(([url]) => url.includes('api.openai.com'))).toBe( - true - ) - expect(fetchMock.mock.calls[4][0]).toBe('https://openrouter.ai/api/v1/embeddings') + const attempts = EMBEDDING_MAX_RETRIES + 1 + expect(fetchMock).toHaveBeenCalledTimes(attempts + 1) + expect( + fetchMock.mock.calls.slice(0, attempts).every(([url]) => url.includes('api.openai.com')) + ).toBe(true) + expect(fetchMock.mock.calls[attempts][0]).toBe('https://openrouter.ai/api/v1/embeddings') expect(projectInputs).toHaveBeenCalledOnce() expect(result.embeddings).toEqual([[7, 8]]) }) @@ -713,13 +717,111 @@ describe('knowledge embedding transport fallback', () => { .filter(([url]) => url === 'https://openrouter.ai/api/v1/embeddings') .flatMap(([, init]) => JSON.parse((init as RequestInit).body as string).input as string[]) expect(openRouterInputs).toEqual([secondInput]) - expect(fetchMock).toHaveBeenCalledTimes(6) + // The succeeding batch, every attempt on the failing one, then its fallback. + expect(fetchMock).toHaveBeenCalledTimes(1 + (EMBEDDING_MAX_RETRIES + 1) + 1) expect(result.embeddings).toEqual([[1], [2]]) expect(result.totalTokens).toBe(6) expect(result.billableTokens).toBe(3) expect(result.isBYOK).toBe(false) }) + /** + * The retry loop replaces its own backoff with a provider-stated wait, but only + * if the wait reaches it. Nothing downstream of the transport could see the + * response headers, so a rate-limited embedding request retried blind. + */ + it('carries the provider-stated retry wait onto the thrown error', async () => { + vi.useFakeTimers() + setEnv({ OPENAI_API_KEY: 'openai-test' }) + fetchMock.mockResolvedValue({ + ok: false, + status: 429, + statusText: '429', + headers: new Headers({ 'retry-after': '42' }), + json: async () => ({ error: 'rate limited' }), + text: async () => 'rate limited', + } as Response) + + const pending = embed(['hello'], { ...options, apiKey: 'openai-test' }).catch((e) => e) + await vi.runAllTimersAsync() + const error = await pending + + expect(error).toBeInstanceOf(EmbeddingAPIError) + expect(error.status).toBe(429) + expect(error.retryAfterMs).toBe(42_000) + }) + + /** + * A stated wait past the ceiling cannot be honoured, so retrying only clamps + * every attempt below the reopen time and spends the budget for nothing. The + * error still classifies as transient, so the fallback chain takes over at + * once rather than after the retries burn down. + */ + it('does not retry a wait it cannot honour, falling back immediately', async () => { + vi.useFakeTimers() + setEnv({ OPENAI_API_KEY: 'openai-test', OPENROUTER_API_KEY: 'or-test' }) + const fetchMock = vi.fn().mockImplementation(async (url: string) => + url === 'https://api.openai.com/v1/embeddings' + ? ({ + ok: false, + status: 429, + statusText: '429', + // Six minutes, far past EMBEDDING_MAX_RETRY_DELAY_MS. + headers: new Headers({ + 'x-ratelimit-remaining-tokens': '0', + 'x-ratelimit-reset-tokens': '6m0s', + }), + json: async () => ({ error: 'rate limited' }), + text: async () => 'rate limited', + } as Response) + : jsonResponse(openAIBody([[9, 9]], 2)) + ) + vi.stubGlobal('fetch', fetchMock) + + const pending = embedKnowledgeForDeployment(['hello'], options, false) + await vi.runAllTimersAsync() + const result = await pending + + const openAICalls = fetchMock.mock.calls.filter(([url]) => url.includes('api.openai.com')) + expect(openAICalls).toHaveLength(1) + expect(result.embeddings).toEqual([[9, 9]]) + }) + + /** + * A window shorter than the whole budget is still reachable: the individual + * waits are clamped below it but they accumulate, so a later attempt lands + * after it reopens. Refusing these would strand a single-provider caller that + * had only to wait a little longer than one clamped delay. + */ + it('keeps retrying a wait the budget can outlast, and recovers', async () => { + vi.useFakeTimers() + setEnv({ OPENAI_API_KEY: 'openai-test' }) + let call = 0 + const fetchMock = vi.fn().mockImplementation(async () => { + call++ + if (call === 1) { + return { + ok: false, + status: 429, + statusText: '429', + // Above the per-attempt ceiling, well inside the total budget. + headers: new Headers({ 'retry-after': '35' }), + json: async () => ({ error: 'rate limited' }), + text: async () => 'rate limited', + } as Response + } + return jsonResponse(openAIBody([[4, 4]], 2)) + }) + vi.stubGlobal('fetch', fetchMock) + + const pending = embed(['hello'], { ...options, apiKey: 'openai-test' }) + await vi.runAllTimersAsync() + const result = await pending + + expect(fetchMock).toHaveBeenCalledTimes(2) + expect(result.embeddings).toEqual([[4, 4]]) + }) + it('classifies only transient embedding failures for failover', () => { expect(isTransientEmbeddingError(new EmbeddingAPIError('unavailable', 503))).toBe(true) expect(isTransientEmbeddingError(new EmbeddingAPIError('rate limited', 429))).toBe(true) diff --git a/apps/sim/lib/embeddings/client.ts b/apps/sim/lib/embeddings/client.ts index 912574c77cc..3383de1eaf4 100644 --- a/apps/sim/lib/embeddings/client.ts +++ b/apps/sim/lib/embeddings/client.ts @@ -20,6 +20,7 @@ import { import { resolveProviderKey } from '@/lib/embeddings/keys' import { DEFAULT_OPENROUTER_EMBEDDING_MODEL } from '@/lib/embeddings/openrouter-models' import { getAdapterFactory } from '@/lib/embeddings/providers' +import { resolveEmbeddingRetryDelayMs } from '@/lib/embeddings/rate-limit' import type { EmbeddingProviderAdapter, EmbeddingTaskType, @@ -27,12 +28,26 @@ import type { EmbedResult, OpenRouterEmbedOptions, } from '@/lib/embeddings/types' -import { isRetryableError, retryWithExponentialBackoff } from '@/lib/knowledge/documents/utils' +import { + attachRetryHeaders, + isRetryableError, + retryWithExponentialBackoff, +} from '@/lib/knowledge/documents/utils' import { batchByTokenLimit, estimateTokenCount, truncateToTokenLimit } from '@/lib/tokenization' const logger = createLogger('EmbeddingClient') -const MAX_CONCURRENT_BATCHES = envNumber(env.KB_CONFIG_CONCURRENCY_LIMIT, 50) +/** + * Embedding requests issued concurrently within a single embed call. + * + * A provider's rate limit is per API key, so this multiplies with however many + * documents are being processed at once: the document-processing queue admits + * {@link env.KB_CONFIG_CONCURRENCY_LIMIT} task runs, each reaching here. It was + * previously read from that same variable, so one knob set both factors and the + * product reached four figures of in-flight requests against one key — enough to + * hold a provider at its limit indefinitely, which no retry policy can absorb. + */ +const MAX_CONCURRENT_BATCHES = envNumber(env.KB_CONFIG_EMBEDDING_CONCURRENCY, 8) const EMBEDDING_REQUEST_TIMEOUT_MS = 60_000 /** @@ -47,9 +62,29 @@ const EMBEDDING_REQUEST_TIMEOUT_MS = 60_000 */ const BATCH_TOKEN_TARGET = 8192 +/** Retries after the initial attempt, per embedding request. */ +export const EMBEDDING_MAX_RETRIES = 5 + +/** Ceiling on a single wait between embedding attempts, including a provider-stated one. */ +export const EMBEDDING_MAX_RETRY_DELAY_MS = 30_000 + +/** + * Longest a request can stay in the retry loop: every attempt waits at most the + * ceiling, so the budget is what the attempts span in total. A provider window + * that reopens inside this is still reachable even though each individual wait + * is clamped below it. + */ +const EMBEDDING_RETRY_BUDGET_MS = EMBEDDING_MAX_RETRIES * EMBEDDING_MAX_RETRY_DELAY_MS + export class EmbeddingAPIError extends Error { public status: number + /** + * Wait the provider asked for, read from the rejected response. Consumed by + * {@link retryWithExponentialBackoff}, which prefers it over its own backoff. + */ + public retryAfterMs?: number + constructor(message: string, status: number) { super(message) this.name = 'EmbeddingAPIError' @@ -57,6 +92,30 @@ export class EmbeddingAPIError extends Error { } } +/** + * True when the provider's stated wait outlasts the entire retry budget. + * + * The comparison is against {@link EMBEDDING_RETRY_BUDGET_MS} rather than the + * per-attempt ceiling, because a window longer than one wait is still reachable: + * the attempts are clamped individually but accumulate, so a 35s window reopens + * before the second one lands. Only a window outlasting every attempt is + * genuinely unreachable, and retrying into it spends the budget waiting on + * something that cannot happen. + * + * The error stays transient — it just is not worth retrying here — so refusing + * the retry surfaces it immediately and the fallback chain, which classifies + * separately via `shouldFallback`, reaches the next provider at once. Where no + * fallback is configured the request fails either way; this only decides whether + * it fails now or after the budget burns down for nothing. + */ +function statedWaitOutlastsBudget(error: unknown): boolean { + return ( + error instanceof EmbeddingAPIError && + error.retryAfterMs !== undefined && + error.retryAfterMs > EMBEDDING_RETRY_BUDGET_MS + ) +} + export function isTransientEmbeddingError(error: unknown): boolean { if (error instanceof EmbeddingAPIError) { return error.status === 429 || error.status >= 500 @@ -188,10 +247,27 @@ async function callEmbeddingAPI( if (!response.ok) { const errorText = await response.text() - throw new EmbeddingAPIError( + const error = new EmbeddingAPIError( `Embedding API failed: ${response.status} ${response.statusText} - ${errorText}`, response.status ) + + /** + * Carry the provider's own answer to "when may I retry" onto the error, + * the way `fetchWithRetry` does for connectors. Without it the retry + * loop had nothing but blind exponential backoff and would exhaust every + * attempt inside a rate-limit window that had not yet reopened. + * + * The headers travel non-enumerably so the retry condition can re-read + * them without the bag reaching a log line. + */ + attachRetryHeaders(error, response.headers) + const waitMs = resolveEmbeddingRetryDelayMs(response.headers) + if (waitMs !== null) { + error.retryAfterMs = waitMs + } + + throw error } const json = await response.json() @@ -208,10 +284,22 @@ async function callEmbeddingAPI( return { embeddings, totalTokens } }, { - maxRetries: 3, + /** + * Sized against a rate-limit window rather than a transient blip. The + * provider states its reset in tens of seconds, and the loop clamps that + * stated wait to `maxDelayMs` — at the previous 10s ceiling every attempt + * fired before the window reopened, so the budget was spent without one + * retry landing in the reopened window. + * + * Bounded so a fully saturated provider cannot outlive the task: five + * attempts at the ceiling is well inside `KB_CONFIG_MAX_DURATION`, and + * batches wait concurrently rather than one after another. + */ + maxRetries: EMBEDDING_MAX_RETRIES, initialDelayMs: 1000, - maxDelayMs: 10000, - retryCondition: isTransientEmbeddingError, + maxDelayMs: EMBEDDING_MAX_RETRY_DELAY_MS, + retryCondition: (error) => + isTransientEmbeddingError(error) && !statedWaitOutlastsBudget(error), } ) } diff --git a/apps/sim/lib/embeddings/rate-limit.test.ts b/apps/sim/lib/embeddings/rate-limit.test.ts new file mode 100644 index 00000000000..a2d73c1714a --- /dev/null +++ b/apps/sim/lib/embeddings/rate-limit.test.ts @@ -0,0 +1,101 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it } from 'vitest' +import { parseGoDurationMs, resolveEmbeddingRetryDelayMs } from '@/lib/embeddings/rate-limit' + +function headers(values: Record): { get(name: string): string | null } { + return { get: (name: string) => values[name] ?? null } +} + +describe('parseGoDurationMs', () => { + it('reads the shapes OpenAI documents for its reset headers', () => { + expect(parseGoDurationMs('1s')).toBe(1000) + expect(parseGoDurationMs('6m0s')).toBe(360_000) + expect(parseGoDurationMs('23h47m36.648s')).toBeCloseTo(85_656_648, 0) + }) + + /** + * `m` and `ms` share a prefix, so a parser that matched the shorter unit first + * would read a twelve-millisecond wait as a twelve-minute one — a 60,000x + * overstatement that would stall ingestion on a healthy provider. + */ + it('does not confuse milliseconds with minutes', () => { + expect(parseGoDurationMs('12ms')).toBe(12) + expect(parseGoDurationMs('12m')).toBe(720_000) + }) + + it('refuses a value the format does not fully describe', () => { + expect(parseGoDurationMs('')).toBeNull() + expect(parseGoDurationMs(' ')).toBeNull() + expect(parseGoDurationMs('soon')).toBeNull() + expect(parseGoDurationMs('60')).toBeNull() + expect(parseGoDurationMs('1s later')).toBeNull() + expect(parseGoDurationMs('1y')).toBeNull() + }) +}) + +describe('resolveEmbeddingRetryDelayMs', () => { + it('prefers Retry-After, which names the wait directly', () => { + expect( + resolveEmbeddingRetryDelayMs( + headers({ 'retry-after': '20', 'x-ratelimit-reset-tokens': '6m0s' }) + ) + ).toBe(20_000) + }) + + /** + * The cap in `parseRetryAfter` defaults to 30s. The retry loop owns the clamp + * against its own ceiling, so a longer stated wait must arrive intact rather + * than being silently truncated on the way in. + */ + it('passes a long Retry-After through uncapped', () => { + expect(resolveEmbeddingRetryDelayMs(headers({ 'retry-after': '120' }))).toBe(120_000) + }) + + it('falls back to the reset of the dimension that is actually exhausted', () => { + expect( + resolveEmbeddingRetryDelayMs( + headers({ + 'x-ratelimit-remaining-tokens': '0', + 'x-ratelimit-reset-tokens': '45s', + 'x-ratelimit-remaining-requests': '4999', + 'x-ratelimit-reset-requests': '6m0s', + }) + ) + ).toBe(45_000) + }) + + it('waits out the longer window when both dimensions are exhausted', () => { + expect( + resolveEmbeddingRetryDelayMs( + headers({ + 'x-ratelimit-remaining-tokens': '0', + 'x-ratelimit-reset-tokens': '45s', + 'x-ratelimit-remaining-requests': '0', + 'x-ratelimit-reset-requests': '6m0s', + }) + ) + ).toBe(360_000) + }) + + /** + * Providers stamp these headers on every response. A reset read while quota + * remains says when the window rolls over, not when this request may be + * retried — using it would pin an unrelated failure to a flat wait. + */ + it('says nothing when no dimension is exhausted', () => { + expect( + resolveEmbeddingRetryDelayMs( + headers({ + 'x-ratelimit-remaining-tokens': '150000', + 'x-ratelimit-reset-tokens': '6m0s', + }) + ) + ).toBeNull() + }) + + it('says nothing when the response carries no rate-limit headers', () => { + expect(resolveEmbeddingRetryDelayMs(headers({}))).toBeNull() + }) +}) diff --git a/apps/sim/lib/embeddings/rate-limit.ts b/apps/sim/lib/embeddings/rate-limit.ts new file mode 100644 index 00000000000..da128b7a094 --- /dev/null +++ b/apps/sim/lib/embeddings/rate-limit.ts @@ -0,0 +1,92 @@ +import { parseRetryAfter } from '@sim/utils/retry' + +/** + * Reads the wait an embedding provider states when it rejects a request for + * rate limiting, in milliseconds, or `null` when the response says nothing + * usable and the caller should fall back to its own backoff. + * + * Kept out of the shared connector helper in `@/lib/knowledge/documents/utils` + * deliberately. That module reads `x-ratelimit-reset` as UTC epoch seconds, + * which is what GitHub and X document; the OpenAI family spells the same idea as + * `x-ratelimit-reset-tokens` carrying a Go duration. Teaching the shared reader + * both spellings would change the retry behaviour of every connector that goes + * through `fetchWithRetry`, so the provider-specific reading stays here. + */ +interface HeaderReader { + get(name: string): string | null +} + +/** Milliseconds per Go duration unit. Longer spellings first — `ms` must win over `m`. */ +const GO_DURATION_UNITS = [ + ['ns', 1e-6], + ['us', 1e-3], + ['µs', 1e-3], + ['ms', 1], + ['h', 3_600_000], + ['m', 60_000], + ['s', 1000], +] as const + +const GO_DURATION_PART = /(\d+(?:\.\d+)?)(ns|us|µs|ms|h|m|s)/gy + +/** + * Parses a Go duration such as `6m0s`, `12ms`, or `23h47m36.648s` — the format + * OpenAI uses for its rate-limit reset headers — into milliseconds. + * + * Returns `null` for anything the format does not fully describe, so a value in + * some other shape falls back to backoff rather than being half-read. The sticky + * match and the end-of-input check are what enforce that: `m` and `ms` share a + * prefix, so a parser scanning loosely would read `12ms` as twelve minutes. + */ +export function parseGoDurationMs(value: string): number | null { + const trimmed = value.trim() + if (trimmed.length === 0) return null + + GO_DURATION_PART.lastIndex = 0 + let total = 0 + let matched = false + + while (GO_DURATION_PART.lastIndex < trimmed.length) { + const part = GO_DURATION_PART.exec(trimmed) + if (!part) return null + const unit = GO_DURATION_UNITS.find(([name]) => name === part[2]) + if (!unit) return null + total += Number(part[1]) * unit[1] + matched = true + } + + return matched ? total : null +} + +/** Paired remaining/reset headers, per rate-limited dimension. */ +const OPENAI_LIMIT_DIMENSIONS = [ + { remaining: 'x-ratelimit-remaining-tokens', reset: 'x-ratelimit-reset-tokens' }, + { remaining: 'x-ratelimit-remaining-requests', reset: 'x-ratelimit-reset-requests' }, +] as const + +/** + * Resolves the stated wait for a rejected embedding request. + * + * `Retry-After` wins when present, since the provider is naming the wait + * directly. Otherwise the reset header is read for whichever dimension is + * actually exhausted — a tokens-per-minute rejection and a requests-per-minute + * rejection reopen at different times, and the reset for a dimension with quota + * left says nothing about when this request may be retried. When both are + * exhausted the longer wait is the one that governs. + */ +export function resolveEmbeddingRetryDelayMs(headers: HeaderReader): number | null { + const retryAfterMs = parseRetryAfter(headers.get('retry-after'), Number.POSITIVE_INFINITY) + if (retryAfterMs !== null && retryAfterMs > 0) return retryAfterMs + + let longest: number | null = null + for (const dimension of OPENAI_LIMIT_DIMENSIONS) { + if (headers.get(dimension.remaining) !== '0') continue + const reset = headers.get(dimension.reset) + if (!reset) continue + const waitMs = parseGoDurationMs(reset) + if (waitMs === null || waitMs <= 0) continue + longest = longest === null ? waitMs : Math.max(longest, waitMs) + } + + return longest +} diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 5a209ddc72a..2c5904773a4 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -297,10 +297,19 @@ function withTimeout( ]) } +/** + * Limits for the in-process document path. + * + * Both values used to be derived from variables owned by other subsystems — + * documents-at-once from the task-queue depth, documents-per-batch from the + * chunks-per-embedding-request size — so tuning either of those silently moved + * this one too, by a factor set by the divisor rather than by intent. The + * divisors are gone and the previous effective values (4 and 10) are now the + * declared defaults. + */ const PROCESSING_CONFIG = { - maxConcurrentDocuments: - Math.max(1, Math.floor(envNumber(env.KB_CONFIG_CONCURRENCY_LIMIT, 20) / 5)) || 4, - batchSize: Math.max(1, Math.floor(envNumber(env.KB_CONFIG_BATCH_SIZE, 20) / 2)) || 10, + maxConcurrentDocuments: envNumber(env.KB_CONFIG_DOCUMENT_CONCURRENCY, 4, { min: 1 }), + batchSize: envNumber(env.KB_CONFIG_DOCUMENT_BATCH_SIZE, 10, { min: 1 }), delayBetweenBatches: envNumber(env.KB_CONFIG_DELAY_BETWEEN_BATCHES, 100) * 2, delayBetweenDocuments: envNumber(env.KB_CONFIG_DELAY_BETWEEN_DOCUMENTS, 50) * 2, }