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

* 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.

* fix(knowledge): stop retrying an embedding wait we will not honor

Honoring the provider's stated wait introduced a case the retry budget could
not serve. When a provider states a reset longer than the ceiling, the loop
clamps every attempt to that ceiling, so the whole budget is spent inside a
window that has not reopened — and with five attempts at thirty seconds that
delayed the fallback provider by around two and a half minutes, where the
previous blind backoff reached it in about seven seconds.

A stated wait past the ceiling now refuses the retry outright. The error still
classifies as transient, and the fallback chain classifies separately through
shouldFallback, so the next provider is reached immediately instead of after the
budget burns down. Retrying was never going to succeed in that window, so
nothing is given up.

* fix(knowledge): measure a stated wait against the whole retry budget

Refusing to retry once the stated wait passed the per-attempt ceiling was too
blunt. Each wait is clamped individually but the attempts accumulate, so a
window a little longer than one clamped delay still reopens partway through the
budget: a 35s wait is reachable on the second attempt. Rejecting those stranded
a caller with no fallback provider, which would have recovered by waiting.

The comparison is now against the budget the attempts span in total. A window
inside it is retried and can recover; only one that outlasts every attempt is
unreachable, and that still fails fast so the fallback chain is reached at once
rather than after the budget burns down.

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,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)

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,16 +62,60 @@ 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'
     this.status = status
   }
 }
 
+/**
+ * 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),
     }
   )
 }