fix(knowledge): classify rejected BYOK embedding keys (#7226)

Author: waleedlatif1

apps/sim/background/knowledge-processing.test.ts

@@ -26,7 +26,7 @@ vi.mock('@/lib/knowledge/documents/service', () => ({
   processDocumentAsync: mockProcessDocumentAsync,
 }))
 
-import { EmbeddingQuotaExhaustedError } from '@/lib/embeddings/client'
+import { EmbeddingAPIError, EmbeddingQuotaExhaustedError } from '@/lib/embeddings/client'
 import { EMBEDDING_QUOTA_CIRCUIT_TTL_MS } from '@/lib/embeddings/quota-circuit'
 import {
   PermanentDocumentProcessingError,
@@ -331,6 +331,27 @@ describe('knowledge processing worker', () => {
     })
   })
 
+  it('returns an actionable outcome when customer-managed embedding credentials are rejected', async () => {
+    mockProcessDocumentAsync.mockRejectedValue(
+      new EmbeddingAPIError('Embedding API failed: 401', 401, true)
+    )
+
+    await expect(runDocumentProcessing(WORKSPACE_PAYLOAD)).resolves.toMatchObject({
+      success: false,
+      outcome: 'customer_configuration',
+      code: 'embedding_credentials_rejected',
+      error:
+        'The configured embedding API key was rejected. Update the key and retry this document.',
+    })
+  })
+
+  it('preserves task failure for rejected platform embedding credentials', async () => {
+    const platformError = new EmbeddingAPIError('Embedding API failed: 401', 401)
+    mockProcessDocumentAsync.mockRejectedValue(platformError)
+
+    await expect(runDocumentProcessing(WORKSPACE_PAYLOAD)).rejects.toBe(platformError)
+  })
+
   it('preserves normal retries for transient failures', async () => {
     const transientError = new Error('Database connection timed out')
     mockProcessDocumentAsync.mockRejectedValue(transientError)

apps/sim/background/knowledge-processing.ts

@@ -1,7 +1,12 @@
 import { createLogger } from '@sim/logger'
 import { task } from '@trigger.dev/sdk'
 import { env, envNumber } from '@/lib/core/config/env'
-import { EMBEDDING_QUOTA_EXHAUSTED_MESSAGE, isEmbeddingQuotaExhaustion } from '@/lib/embeddings'
+import {
+  BYOK_EMBEDDING_CREDENTIAL_REJECTION_MESSAGE,
+  EMBEDDING_QUOTA_EXHAUSTED_MESSAGE,
+  isBYOKEmbeddingCredentialRejection,
+  isEmbeddingQuotaExhaustion,
+} from '@/lib/embeddings'
 import {
   isPermanentDocumentProcessingError,
   isUsageLimitDocumentProcessingError,
@@ -112,6 +117,21 @@ export async function runDocumentProcessing(
         processingTime: Date.now() - startedAt,
       }
     }
+    if (isBYOKEmbeddingCredentialRejection(error)) {
+      logger.warn(`[${requestId}] Customer-managed embedding credentials were rejected`, {
+        filename: docData.filename,
+        status: error.status,
+      })
+      return {
+        success: false,
+        outcome: 'customer_configuration' as const,
+        code: 'embedding_credentials_rejected' as const,
+        documentId,
+        filename: docData.filename,
+        error: BYOK_EMBEDDING_CREDENTIAL_REJECTION_MESSAGE,
+        processingTime: Date.now() - startedAt,
+      }
+    }
     if (isPermanentDocumentProcessingError(error)) {
       logger.warn(`[${requestId}] Document cannot be processed without changing its content`, {
         code: error.code,

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

@@ -12,6 +12,7 @@ import {
   embed,
   embedKnowledgeForDeployment,
   embedOpenRouter,
+  isBYOKEmbeddingCredentialRejection,
   isEmbeddingQuotaExhaustion,
   isTransientEmbeddingError,
   MAX_EMBEDDING_SUCCESS_RESPONSE_BYTES,
@@ -341,6 +342,7 @@ describe('embed', () => {
     expect(error).toBeInstanceOf(Error)
     expect((error as Error).message).toMatch(/Embedding API failed: 401/)
     expect((error as Error).message).not.toContain(echoedSecret)
+    expect(isBYOKEmbeddingCredentialRejection(error)).toBe(true)
     // 401 is not retryable, so exactly one attempt is made.
     expect(fetchMock).toHaveBeenCalledTimes(1)
   })
@@ -992,6 +994,24 @@ describe('knowledge embedding transport fallback', () => {
     expect(result.isBYOK).toBe(true)
   })
 
+  it('distinguishes workspace credential rejection from a platform credential failure', async () => {
+    fetchMock.mockResolvedValue(jsonResponse({ error: 'invalid key' }, 401))
+
+    setEnv({ OPENAI_API_KEY: 'platform-openai-test' })
+    const platformError = await embedKnowledgeForDeployment(['hello'], options, true).catch(
+      (error) => error
+    )
+    expect(isBYOKEmbeddingCredentialRejection(platformError)).toBe(false)
+
+    mockGetBYOKKey.mockResolvedValue({ apiKey: 'workspace-openai-test', isBYOK: true })
+    const workspaceError = await embedKnowledgeForDeployment(
+      ['hello'],
+      { ...options, workspaceId: 'workspace-1' },
+      true
+    ).catch((error) => error)
+    expect(isBYOKEmbeddingCredentialRejection(workspaceError)).toBe(true)
+  })
+
   it('does not use OpenRouter for non-OpenAI knowledge models', async () => {
     setEnv({ GEMINI_API_KEY: 'gemini-test', OPENROUTER_API_KEY: 'or-test' })
     fetchMock.mockResolvedValue(
@@ -1349,4 +1369,11 @@ describe('knowledge embedding transport fallback', () => {
     expect(isTransientEmbeddingError(new EmbeddingAPIError('invalid key', 401))).toBe(false)
     expect(isTransientEmbeddingError(new DOMException('timed out', 'AbortError'))).toBe(true)
   })
+
+  it('does not misclassify quota-related BYOK rejections as authentication failures', () => {
+    const error = new EmbeddingAPIError('Embedding API failed: 403', 403, true)
+    error.quotaExhausted = true
+
+    expect(isBYOKEmbeddingCredentialRejection(error)).toBe(false)
+  })
 })

apps/sim/lib/embeddings/client.ts

@@ -135,6 +135,9 @@ const EMBEDDING_RETRY_BUDGET_MS = EMBEDDING_MAX_RETRIES * EMBEDDING_MAX_RETRY_DE
 export class EmbeddingAPIError extends Error {
   public status: number
 
+  /** True when the rejected request used a customer-managed credential. */
+  public readonly isBYOK: boolean
+
   /** Rejected for an exhausted balance rather than a recoverable rate limit. */
   public quotaExhausted?: boolean
 
@@ -144,10 +147,11 @@ export class EmbeddingAPIError extends Error {
    */
   public retryAfterMs?: number
 
-  constructor(message: string, status: number) {
+  constructor(message: string, status: number, isBYOK = false) {
     super(message)
     this.name = 'EmbeddingAPIError'
     this.status = status
+    this.isBYOK = isBYOK
   }
 }
 
@@ -170,6 +174,9 @@ export class EmbeddingOutputLimitError extends Error {
 export const EMBEDDING_QUOTA_EXHAUSTED_MESSAGE =
   'The embedding provider has exhausted its available quota. Add credit or replace the credential before retrying.'
 
+export const BYOK_EMBEDDING_CREDENTIAL_REJECTION_MESSAGE =
+  'The configured embedding API key was rejected. Update the key and retry this document.'
+
 /**
  * A provider credential has no remaining credit. This remains transient across
  * providers so a configured fallback can run, but it is terminal for the
@@ -182,7 +189,8 @@ export class EmbeddingQuotaExhaustedError extends EmbeddingAPIError {
     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
+      status,
+      cause instanceof EmbeddingAPIError && cause.isBYOK
     )
     this.name = 'EmbeddingQuotaExhaustedError'
     this.providerId = providerId
@@ -204,6 +212,21 @@ export function isEmbeddingQuotaExhaustion(error: unknown): boolean {
   return false
 }
 
+/**
+ * True when a customer-managed embedding credential was rejected outright.
+ * These failures require a key or permission change; retrying the same request
+ * cannot recover. Quota failures are classified separately even when a provider
+ * reports them with HTTP 403.
+ */
+export function isBYOKEmbeddingCredentialRejection(error: unknown): error is EmbeddingAPIError {
+  return (
+    error instanceof EmbeddingAPIError &&
+    error.isBYOK &&
+    !error.quotaExhausted &&
+    (error.status === 401 || error.status === 403)
+  )
+}
+
 /**
  * True when a rejection body reports an exhausted balance rather than a rate
  * limit. OpenAI returns 429 for both, but only a rate limit reopens: a spent
@@ -435,6 +458,7 @@ async function callEmbeddingAPI(
    */
   requestedDimensions: number | undefined,
   expectedDimensions: number | undefined,
+  isBYOK: boolean,
   signal?: AbortSignal
 ): Promise<{ embeddings: number[][]; totalTokens: number; dimensions: number }> {
   return retryWithExponentialBackoff(
@@ -470,7 +494,8 @@ async function callEmbeddingAPI(
         const classificationBody = await readEmbeddingErrorBody(response)
         const error = new EmbeddingAPIError(
           `Embedding API failed: ${response.status}`,
-          response.status
+          response.status,
+          isBYOK
         )
         error.quotaExhausted =
           isQuotaExhaustionBody(classificationBody) ||
@@ -639,12 +664,19 @@ async function embedWithProvider(
           provider.quotaCircuitIdentity,
           requestedDimensions,
           provider.dimensions,
+          provider.isBYOK,
           signal
         )
       } catch (error) {
         const message = `Failed to generate embeddings for batch ${i + 1}/${batches.length}:`
         if (isEmbeddingQuotaExhaustion(error)) {
           logger.warn(message, { providerId: provider.providerId, quotaExhausted: true })
+        } else if (isBYOKEmbeddingCredentialRejection(error)) {
+          logger.warn(message, {
+            providerId: provider.providerId,
+            outcome: 'customer_configuration',
+            status: error.status,
+          })
         } else {
           logger.error(message, error)
         }
@@ -818,6 +850,7 @@ export async function embedOpenRouter(
       quotaCircuitIdentity,
       options.dimensions,
       expectedDimensions,
+      true,
       options.signal
     )
 
@@ -1001,14 +1034,20 @@ export async function embedKnowledgeForDeployment(
             provider.providerId,
             provider.quotaCircuitIdentity,
             options.dimensions,
-            provider.dimensions
+            provider.dimensions,
+            provider.isBYOK
           )),
           provider,
         }))
       } catch (error) {
         const message = `Failed to generate embeddings for batch ${i + 1}/${batches.length}:`
         if (isEmbeddingQuotaExhaustion(error)) {
           logger.warn(message, { quotaExhausted: true })
+        } else if (isBYOKEmbeddingCredentialRejection(error)) {
+          logger.warn(message, {
+            outcome: 'customer_configuration',
+            status: error.status,
+          })
         } else {
           logger.error(message, error)
         }

apps/sim/lib/embeddings/index.ts

@@ -10,12 +10,14 @@ export {
   resolveDimensions,
 } from '@/lib/embeddings/catalog'
 export {
+  BYOK_EMBEDDING_CREDENTIAL_REJECTION_MESSAGE,
   EMBEDDING_QUOTA_EXHAUSTED_MESSAGE,
   EmbeddingOutputLimitError,
   embed,
   embedKnowledge,
   embedOpenRouter,
   getEmbeddingAggregateItemLimit,
+  isBYOKEmbeddingCredentialRejection,
   isEmbeddingQuotaExhaustion,
 } from '@/lib/embeddings/client'
 export { DEFAULT_OPENROUTER_EMBEDDING_MODEL } from '@/lib/embeddings/openrouter-models'

apps/sim/lib/knowledge/documents/document-processing-source.test.ts

@@ -73,8 +73,11 @@ import {
   markInsideTriggerRun,
   resetInsideTriggerRunForTests,
 } from '@/lib/core/config/trigger-runtime'
-import { EMBEDDING_QUOTA_EXHAUSTED_MESSAGE } from '@/lib/embeddings'
-import { EmbeddingQuotaExhaustedError } from '@/lib/embeddings/client'
+import {
+  BYOK_EMBEDDING_CREDENTIAL_REJECTION_MESSAGE,
+  EMBEDDING_QUOTA_EXHAUSTED_MESSAGE,
+} from '@/lib/embeddings'
+import { EmbeddingAPIError, EmbeddingQuotaExhaustedError } from '@/lib/embeddings/client'
 import {
   PermanentDocumentProcessingError,
   UsageLimitDocumentProcessingError,
@@ -742,6 +745,41 @@ describe('processDocumentAsync write guards', () => {
     expect(failure![0]).not.toHaveProperty('processingAttempts')
   })
 
+  it('dead-letters rejected customer-managed embedding credentials until the user retries', async () => {
+    dbChainMockFns.limit
+      .mockResolvedValueOnce([PERSISTED_CONTEXT])
+      .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW])
+      .mockResolvedValueOnce([{ id: 'document-1' }])
+    mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING])
+    mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue(
+      new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]])
+    )
+    mockProcessDocument.mockResolvedValue({
+      chunks: [{ text: 'Index me', metadata: { startIndex: 0, endIndex: 8 } }],
+      metadata: { chunkCount: 1, tokenCount: 2, characterCount: 8 },
+    })
+    mockGenerateEmbeddings.mockRejectedValue(
+      new EmbeddingAPIError('Embedding API failed: 401', 401, true)
+    )
+
+    await expect(
+      processDocumentAsync('knowledge-base-1', 'document-1', {
+        filename: 'report.docx',
+        fileUrl: 'https://example.com/report.docx',
+        fileSize: 1,
+        mimeType: 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+      })
+    ).rejects.toMatchObject({ status: 401, isBYOK: true })
+
+    const failure = dbChainMockFns.set.mock.calls.find(
+      (call) => (call[0] as Record<string, unknown> | undefined)?.processingStatus === 'failed'
+    )
+    expect(failure?.[0]).toMatchObject({
+      processingError: BYOK_EMBEDDING_CREDENTIAL_REJECTION_MESSAGE,
+      processingAttempts: MAX_PROCESSING_ATTEMPTS,
+    })
+  })
+
   it.each([
     { chargedAtDispatch: true, refundsAttempt: true },
     { chargedAtDispatch: false, refundsAttempt: false },