From a0db81094e0a70745fb5068eb04153d970a06422 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 28 Aug 2026 13:21:36 -0700 Subject: [PATCH] fix(knowledge): classify rejected BYOK embedding keys --- .../background/knowledge-processing.test.ts | 23 ++++++++- apps/sim/background/knowledge-processing.ts | 22 ++++++++- apps/sim/lib/embeddings/client.test.ts | 27 +++++++++++ apps/sim/lib/embeddings/client.ts | 47 +++++++++++++++++-- apps/sim/lib/embeddings/index.ts | 2 + .../document-processing-source.test.ts | 42 ++++++++++++++++- apps/sim/lib/knowledge/documents/service.ts | 31 ++++++++++-- 7 files changed, 181 insertions(+), 13 deletions(-) diff --git a/apps/sim/background/knowledge-processing.test.ts b/apps/sim/background/knowledge-processing.test.ts index 741beb58390..7e4423df541 100644 --- a/apps/sim/background/knowledge-processing.test.ts +++ b/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) diff --git a/apps/sim/background/knowledge-processing.ts b/apps/sim/background/knowledge-processing.ts index 8ff7e1aedf2..41efe6c6674 100644 --- a/apps/sim/background/knowledge-processing.ts +++ b/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, diff --git a/apps/sim/lib/embeddings/client.test.ts b/apps/sim/lib/embeddings/client.test.ts index 37286786e26..cd19f32f156 100644 --- a/apps/sim/lib/embeddings/client.test.ts +++ b/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) + }) }) diff --git a/apps/sim/lib/embeddings/client.ts b/apps/sim/lib/embeddings/client.ts index 6e1682ba3b2..c9e0a623345 100644 --- a/apps/sim/lib/embeddings/client.ts +++ b/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,7 +1034,8 @@ export async function embedKnowledgeForDeployment( provider.providerId, provider.quotaCircuitIdentity, options.dimensions, - provider.dimensions + provider.dimensions, + provider.isBYOK )), provider, })) @@ -1009,6 +1043,11 @@ export async function embedKnowledgeForDeployment( 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) } diff --git a/apps/sim/lib/embeddings/index.ts b/apps/sim/lib/embeddings/index.ts index 3e022df6bb2..9bd367954fa 100644 --- a/apps/sim/lib/embeddings/index.ts +++ b/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' diff --git a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts b/apps/sim/lib/knowledge/documents/document-processing-source.test.ts index 23cbed90464..96f2aee25ac 100644 --- a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts +++ b/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 | 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 }, diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 05b0f5af21a..164d71b04a2 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -64,8 +64,10 @@ import { isInsideTriggerRun } from '@/lib/core/config/trigger-runtime' import { OrchestrationError } from '@/lib/core/orchestration/types' import { mapWithConcurrency } from '@/lib/core/utils/concurrency' import { + BYOK_EMBEDDING_CREDENTIAL_REJECTION_MESSAGE, EMBEDDING_QUOTA_EXHAUSTED_MESSAGE, getEmbeddingAggregateItemLimit, + isBYOKEmbeddingCredentialRejection, isEmbeddingQuotaExhaustion, } from '@/lib/embeddings' import { @@ -1271,6 +1273,13 @@ async function dispatchInProcess( }) return true } + if (isBYOKEmbeddingCredentialRejection(error)) { + logger.warn(`[${requestId}] Customer-managed embedding credentials were rejected`, { + documentId: p.documentId, + status: error.status, + }) + return true + } const message = processingClaimed ? 'In-process document processing failed' : 'In-process document dispatch failed before claiming the document' @@ -1848,6 +1857,7 @@ export async function processDocumentAsync( } catch (error) { const processingTime = Date.now() - startTime const embeddingQuotaExhausted = isEmbeddingQuotaExhaustion(error) + const byokCredentialRejected = isBYOKEmbeddingCredentialRejection(error) const usageLimitExceeded = isUsageLimitDocumentProcessingError(error) const permanentError = toPermanentDocumentProcessingError(error, processingFilename) let recordedError = permanentError ?? error @@ -1862,22 +1872,32 @@ export async function processDocumentAsync( } } const quotaContinuationFailed = quotaContinuationAttempted && !quotaDeferredUntil - const errorMessage = embeddingQuotaExhausted - ? quotaContinuationFailed - ? getErrorMessage(recordedError, 'Embedding quota continuation dispatch failed') - : EMBEDDING_QUOTA_EXHAUSTED_MESSAGE - : getErrorMessage(recordedError, 'Unknown error') + const errorMessage = byokCredentialRejected + ? BYOK_EMBEDDING_CREDENTIAL_REJECTION_MESSAGE + : embeddingQuotaExhausted + ? quotaContinuationFailed + ? getErrorMessage(recordedError, 'Embedding quota continuation dispatch failed') + : EMBEDDING_QUOTA_EXHAUSTED_MESSAGE + : getErrorMessage(recordedError, 'Unknown error') const logContext = { errorType: toError(recordedError).name, knowledgeBaseId, mimeType: docData.mimeType, fileSize: docData.fileSize, + ...(byokCredentialRejected + ? { + code: 'embedding_credentials_rejected', + outcome: 'customer_configuration', + status: error.status, + } + : {}), } const logMessage = quotaDeferredUntil ? `[${documentId}] Deferred document processing after ${processingTime}ms:` : `[${documentId}] Failed to process document after ${processingTime}ms:` if ( (embeddingQuotaExhausted && !quotaContinuationFailed) || + byokCredentialRejected || usageLimitExceeded || permanentError ) { @@ -1898,6 +1918,7 @@ export async function processDocumentAsync( processingDeferredUntil: quotaDeferredUntil, processingCompletedAt: quotaDeferredUntil ? null : new Date(), ...(permanentError || + byokCredentialRejected || (embeddingQuotaExhausted && attemptContext?.quotaContinuationExhausted) ? { processingAttempts: MAX_PROCESSING_ATTEMPTS } : (embeddingQuotaExhausted || usageLimitExceeded) && attemptContext?.chargedAtDispatch