Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 22 additions & 1 deletion apps/sim/background/knowledge-processing.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
22 changes: 21 additions & 1 deletion apps/sim/background/knowledge-processing.ts
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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,
Expand Down
27 changes: 27 additions & 0 deletions apps/sim/lib/embeddings/client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import {
embed,
embedKnowledgeForDeployment,
embedOpenRouter,
isBYOKEmbeddingCredentialRejection,
isEmbeddingQuotaExhaustion,
isTransientEmbeddingError,
MAX_EMBEDDING_SUCCESS_RESPONSE_BYTES,
Expand Down Expand Up @@ -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)
})
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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)
})
})
47 changes: 43 additions & 4 deletions apps/sim/lib/embeddings/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
}
}

Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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) ||
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -818,6 +850,7 @@ export async function embedOpenRouter(
quotaCircuitIdentity,
options.dimensions,
expectedDimensions,
true,
options.signal
)

Expand Down Expand Up @@ -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)
}
Expand Down
2 changes: 2 additions & 0 deletions apps/sim/lib/embeddings/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 },
Expand Down
Loading
Loading