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

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