fix(knowledge): name embedding quota, key, and deadline failures in knowledge search (#8405)

* fix(knowledge): name embedding quota, key, and deadline failures in knowledge search

* fix(knowledge): keep keyless Ollama and mixed fallback quota attribution accurate

* test(embeddings): prove the quota pause through returned errors, not mock call counts

Author: waleedlatif1

apps/sim/lib/embeddings/client.test.ts

@@ -606,6 +606,46 @@ describe('knowledge embedding transport fallback', () => {
     expect(isBYOKEmbeddingCredentialRejection(workspaceError)).toBe(true)
   })
 
+  it('attributes quota exhaustion to the workspace key, including while its pause is open', async () => {
+    /**
+     * Only a credential's first request is refused; the provider would answer every later
+     * one. A second search can therefore fail only if the open pause refused it.
+     */
+    const refusedCredentials = new Set<string>()
+    fetchMock.mockImplementation(async (_url, init) => {
+      const credential = String((init as RequestInit).headers?.Authorization)
+      if (refusedCredentials.has(credential)) return jsonResponse(openAIBody([[1, 2]]))
+      refusedCredentials.add(credential)
+      return jsonResponse(
+        { error: { type: 'insufficient_quota', code: 'insufficient_quota' } },
+        429
+      )
+    })
+    const search = () =>
+      embedKnowledgeForDeployment(
+        ['hello'],
+        { ...options, taskType: 'query' as const, workspaceId: 'workspace-1' },
+        true
+      ).catch((error) => error)
+
+    mockGetBYOKKey.mockResolvedValue({ apiKey: 'workspace-openai-test', isBYOK: true })
+    const refused = await search()
+    const paused = await search()
+    expect(refused).toBeInstanceOf(EmbeddingQuotaExhaustedError)
+    expect(paused).toBeInstanceOf(EmbeddingQuotaExhaustedError)
+    expect(refused.isBYOK).toBe(true)
+    expect(paused.isBYOK).toBe(true)
+
+    mockGetBYOKKey.mockResolvedValue(null)
+    setEnv({ OPENAI_API_KEY: 'platform-openai-test' })
+    const platformRefused = await search()
+    const platformPaused = await search()
+    expect(platformRefused).toBeInstanceOf(EmbeddingQuotaExhaustedError)
+    expect(platformPaused).toBeInstanceOf(EmbeddingQuotaExhaustedError)
+    expect(platformRefused.isBYOK).toBe(false)
+    expect(platformPaused.isBYOK).toBe(false)
+  })
+
   it('ignores OpenRouter on hosted deployments', async () => {
     setEnv({ OPENAI_API_KEY: 'openai-test', OPENROUTER_API_KEY: 'or-test' })
     fetchMock.mockResolvedValue(jsonResponse(openAIBody([[1, 2]])))

apps/sim/lib/embeddings/client.ts

@@ -212,12 +212,22 @@ export const BYOK_EMBEDDING_CREDENTIAL_REJECTION_MESSAGE =
 export class EmbeddingQuotaExhaustedError extends EmbeddingAPIError {
   public readonly providerId: EmbeddingProviderKind
 
-  constructor(providerId: EmbeddingProviderKind, cause?: unknown) {
+  /**
+   * `isBYOK` must be passed when there is no provider response to read it from — an
+   * already-open quota pause or an admission refusal — so a workspace key's exhaustion
+   * is never reported as the platform's. Ollama takes no credential: its provider-level
+   * `isBYOK` only marks its tokens non-billable, so it never attributes to a customer key.
+   */
+  constructor(
+    providerId: EmbeddingProviderKind,
+    cause?: unknown,
+    isBYOK = cause instanceof EmbeddingAPIError && cause.isBYOK
+  ) {
     const status = cause instanceof EmbeddingAPIError ? cause.status : 429
     super(
       `The ${providerId} embedding credential has exhausted its available quota. Add credit or replace the credential before retrying.`,
       status,
-      cause instanceof EmbeddingAPIError && cause.isBYOK
+      isBYOK && providerId !== 'ollama'
     )
     this.name = 'EmbeddingQuotaExhaustedError'
     this.providerId = providerId
@@ -240,6 +250,18 @@ export function isEmbeddingQuotaExhaustion(error: unknown): boolean {
   return false
 }
 
+/**
+ * True when the operation failed on quota and a customer-managed credential is among
+ * the exhausted ones: adding credit to that key is what lets it run again, even when a
+ * platform fallback behind it is exhausted too.
+ */
+export function isBYOKEmbeddingQuotaExhaustion(error: unknown): boolean {
+  if (error instanceof AggregateError) {
+    return isEmbeddingQuotaExhaustion(error) && error.errors.some(isBYOKEmbeddingQuotaExhaustion)
+  }
+  return error instanceof EmbeddingAPIError && error.isBYOK && error.quotaExhausted === true
+}
+
 /**
  * True when a customer-managed embedding credential was rejected outright.
  * These failures require a key or permission change; retrying the same request
@@ -524,7 +546,7 @@ async function callEmbeddingAPI(
   return retryWithExponentialBackoff(
     async (operationSignal, deadlineAt) => {
       if (await isEmbeddingQuotaCircuitOpen(admissionIdentity)) {
-        throw new EmbeddingQuotaExhaustedError(providerId)
+        throw new EmbeddingQuotaExhaustedError(providerId, undefined, isBYOK)
       }
 
       try {
@@ -541,7 +563,7 @@ async function callEmbeddingAPI(
         })
       } catch (error) {
         if (error instanceof ProviderQuotaExhaustedError)
-          throw new EmbeddingQuotaExhaustedError(providerId, error)
+          throw new EmbeddingQuotaExhaustedError(providerId, error, isBYOK)
         throw error
       }
 
@@ -1243,7 +1265,7 @@ export async function assertKnowledgeEmbeddingCapacityForDeployment(
     const exhausted = await isEmbeddingQuotaCircuitOpen(embeddingAdmissionIdentity(provider))
     options.signal?.throwIfAborted()
     if (!exhausted) return
-    errors.push(new EmbeddingQuotaExhaustedError(provider.providerId))
+    errors.push(new EmbeddingQuotaExhaustedError(provider.providerId, undefined, provider.isBYOK))
   }
   if (errors.length === 1) throw errors[0]
   throw new AggregateError(

apps/sim/lib/embeddings/index.ts

@@ -20,6 +20,7 @@ export {
   embedOpenRouter,
   getEmbeddingAggregateItemLimit,
   isBYOKEmbeddingCredentialRejection,
+  isBYOKEmbeddingQuotaExhaustion,
   isEmbeddingQuotaExhaustion,
 } from '@/lib/embeddings/client'
 export { DEFAULT_OPENROUTER_EMBEDDING_MODEL } from '@/lib/embeddings/openrouter-models'

apps/sim/lib/internal/knowledge/execute-tool.ts

@@ -93,8 +93,9 @@ function projectError(
 ): Response {
   signal?.throwIfAborted()
   const projected = policy.project(error)
-  if (projected) return descriptorResponse(projected)
+  if (projected && projected.status < 500) return descriptorResponse(projected)
   logger.error(`[${requestId}] Knowledge tool execution failed`, { error })
+  if (projected) return descriptorResponse(projected)
   return descriptorResponse(
     policy.unhandled?.() ?? { status: 500, body: { error: 'Internal server error' } }
   )

apps/sim/lib/knowledge/api/route-policies.test.ts

@@ -7,7 +7,65 @@ import {
   WorkspaceApiKeyScopeAuthorizationError,
 } from '@/lib/core/application'
 import { OrchestrationError } from '@/lib/core/orchestration/types'
-import { v2KnowledgeErrorPolicies } from '@/lib/knowledge/api/route-policies'
+import { EmbeddingAPIError } from '@/lib/embeddings/api-error'
+import { EmbeddingQuotaExhaustedError } from '@/lib/embeddings/client'
+import {
+  internalKnowledgeErrorPolicies,
+  v2KnowledgeErrorPolicies,
+} from '@/lib/knowledge/api/route-policies'
+import { SearchDeadlineError } from '@/lib/knowledge/search/budget'
+
+function quotaError(isBYOK: boolean, providerId: 'openai' | 'ollama' = 'openai') {
+  return new EmbeddingQuotaExhaustedError(providerId, undefined, isBYOK)
+}
+
+describe('internal knowledge search error policy', () => {
+  it.each([
+    ['a workspace key out of quota', quotaError(true), 503, /this workspace's embedding API key/],
+    ['the platform key out of quota', quotaError(false), 503, /^Knowledge search is temporarily/],
+    [
+      'every fallback provider out of quota',
+      new AggregateError([quotaError(false), quotaError(false)]),
+      503,
+      /^Knowledge search is temporarily/,
+    ],
+    [
+      'a workspace key out of quota ahead of an exhausted platform fallback',
+      new AggregateError([quotaError(true), quotaError(false)]),
+      503,
+      /this workspace's embedding API key/,
+    ],
+    [
+      'a keyless Ollama server out of quota',
+      quotaError(true, 'ollama'),
+      503,
+      /^Knowledge search is temporarily/,
+    ],
+    [
+      'a rejected workspace key',
+      new EmbeddingAPIError('Embedding API failed: 401', 401, true),
+      502,
+      /was rejected/,
+    ],
+    ['the retrieval deadline', new SearchDeadlineError(), 504, /retrieval deadline/],
+  ])('names %s', (_case, error, status, message) => {
+    const response = internalKnowledgeErrorPolicies.search.project(error)
+    expect(response?.status).toBe(status)
+    expect((response?.body as { error: string }).error).toMatch(message)
+  })
+
+  it('leaves a platform key rejection and unknown failures to the generic server error', () => {
+    const policy = internalKnowledgeErrorPolicies.search
+    expect(
+      policy.project(new EmbeddingAPIError('Embedding API failed: 401', 401, false))
+    ).toBeNull()
+    expect(policy.project(new Error('connection reset'))).toBeNull()
+    expect(policy.unhandled?.()).toMatchObject({
+      status: 500,
+      body: { error: 'Failed to perform vector search' },
+    })
+  })
+})
 
 describe('v2 knowledge error policies', () => {
   it.each([

apps/sim/lib/knowledge/api/route-policies.ts

@@ -11,11 +11,17 @@ import {
 } from '@/lib/api/server/routes'
 import { isPayloadSizeLimitError } from '@/lib/core/utils/stream-limits'
 import { internalPersonalCredentialConnectionErrorPolicy } from '@/lib/credentials/api/route-policies'
+import {
+  isBYOKEmbeddingCredentialRejection,
+  isBYOKEmbeddingQuotaExhaustion,
+  isEmbeddingQuotaExhaustion,
+} from '@/lib/embeddings'
 import { KNOWLEDGE_DELEGATION_AUDIENCE } from '@/lib/knowledge/application/authorization'
 import { KnowledgeUsageLimitExceededError } from '@/lib/knowledge/application/billing'
 import { KnowledgeDocumentNotReadyError } from '@/lib/knowledge/application/chunk-errors'
 import { KnowledgeSearchProvenanceUnavailableError } from '@/lib/knowledge/application/search'
 import { KnowledgeDocumentUnsupportedMediaTypeError } from '@/lib/knowledge/application/upload-sessions'
+import { SearchDeadlineError } from '@/lib/knowledge/search/budget'
 import { SearchIndexDormantError } from '@/lib/sim-search/indexed/gate'
 import { v2Error } from '@/app/api/v2/lib/response'
 
@@ -40,6 +46,18 @@ const internalKnowledgeUploadErrorPolicy: InternalErrorPolicy = {
     internalErrorResponse(500, { error: 'Failed to process knowledge upload request' }),
 }
 
+const BYOK_EMBEDDING_QUOTA_SEARCH_MESSAGE =
+  "Knowledge search could not run: this workspace's embedding API key (Settings > Provider API keys) has no remaining quota. Add credit with the provider or replace the key."
+const BYOK_EMBEDDING_REJECTED_SEARCH_MESSAGE =
+  "Knowledge search could not run: this workspace's embedding API key (Settings > Provider API keys) was rejected. Update the key and try again."
+const PLATFORM_EMBEDDING_QUOTA_SEARCH_MESSAGE =
+  'Knowledge search is temporarily unavailable because the embedding provider has no remaining quota. Try again later.'
+
+/**
+ * Names the failures a caller can act on instead of collapsing them into the generic
+ * vector-search `500`. Every status stays `5xx`, so retry and alerting behavior keyed
+ * on server errors is unchanged; only the message and the specific code differ.
+ */
 const internalKnowledgeSearchErrorPolicy: InternalErrorPolicy = {
   project(error) {
     if (error instanceof KnowledgeUsageLimitExceededError) {
@@ -48,6 +66,18 @@ const internalKnowledgeSearchErrorPolicy: InternalErrorPolicy = {
     if (error instanceof KnowledgeSearchProvenanceUnavailableError) {
       return internalErrorResponse(422, { error: error.message })
     }
+    if (isBYOKEmbeddingQuotaExhaustion(error)) {
+      return internalErrorResponse(503, { error: BYOK_EMBEDDING_QUOTA_SEARCH_MESSAGE })
+    }
+    if (isEmbeddingQuotaExhaustion(error)) {
+      return internalErrorResponse(503, { error: PLATFORM_EMBEDDING_QUOTA_SEARCH_MESSAGE })
+    }
+    if (isBYOKEmbeddingCredentialRejection(error)) {
+      return internalErrorResponse(502, { error: BYOK_EMBEDDING_REJECTED_SEARCH_MESSAGE })
+    }
+    if (error instanceof SearchDeadlineError) {
+      return internalErrorResponse(504, { error: error.message })
+    }
     return internalOrchestrationErrorPolicy.project(error)
   },
   unhandled: () => internalErrorResponse(500, { error: 'Failed to perform vector search' }),