fix(knowledge): stop one env knob from setting the embedding request fan-out

KB_CONFIG_CONCURRENCY_LIMIT was read in three places with three meanings: the
document-processing queue depth, the number of embedding requests issued
concurrently inside a single embed call, and (divided by five) the in-process
document concurrency. The first two multiply — every admitted task run reaches
the embed path and opens its own fan-out — so the default put roughly a thousand
requests in flight against one provider key. A rate limit is per key, so the
pipeline held itself at the limit, and no retry policy can absorb a load its own
concurrency is generating.

Each variable is now read by exactly one consumer, which also removes the drift
that hid this: the same variable was read with a different inline fallback in
each place, and since createEnv runs with skipValidation the declared defaults
never execute, so the fallbacks were the real ones and disagreed. The divisors
are gone and the previous effective values are the declared defaults, so only
the embedding fan-out changes: 50 to 8.

KB_CONFIG_BATCH_SIZE had the same conflation between chunks-per-embedding-request
and documents-per-batch, and is split the same way.

Rate-limit rejections also discarded what the provider said about when to come
back. The response headers were dropped when building EmbeddingAPIError, so the
retry loop's support for a server-stated wait was dead code on this path and
every attempt fired blind, exhausting the budget inside a window that had not
reopened. The headers now travel with the error the way fetchWithRetry already
does for connectors, and the wait is read from Retry-After or, failing that, the
reset header for whichever limit dimension is actually exhausted. Those carry a
Go duration rather than the epoch seconds the shared connector helper expects,
so the reading lives with the provider instead of changing retry behaviour for
every connector. The retry budget is sized against a rate-limit window rather
than a blip, since a 10s ceiling clamped every stated wait below the reopen time.

Author: waleedlatif1

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,

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,40 @@ 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)
+  })
+
   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)

apps/sim/lib/embeddings/client.ts

@@ -20,19 +20,34 @@ 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,
   EmbedOptions,
   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,21 @@ 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
+
 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'
@@ -188,10 +215,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,9 +252,20 @@ 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,
+      maxDelayMs: EMBEDDING_MAX_RETRY_DELAY_MS,
       retryCondition: isTransientEmbeddingError,
     }
   )

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<string, string>): { 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()
+  })
+})