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
40 changes: 40 additions & 0 deletions apps/sim/lib/embeddings/client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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]])))
Expand Down
32 changes: 27 additions & 5 deletions apps/sim/lib/embeddings/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Comment thread
waleedlatif1 marked this conversation as resolved.
}

/**
* True when a customer-managed embedding credential was rejected outright.
* These failures require a key or permission change; retrying the same request
Expand Down Expand Up @@ -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 {
Expand All @@ -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
}

Expand Down Expand Up @@ -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(
Expand Down
1 change: 1 addition & 0 deletions apps/sim/lib/embeddings/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ export {
embedOpenRouter,
getEmbeddingAggregateItemLimit,
isBYOKEmbeddingCredentialRejection,
isBYOKEmbeddingQuotaExhaustion,
isEmbeddingQuotaExhaustion,
} from '@/lib/embeddings/client'
export { DEFAULT_OPENROUTER_EMBEDDING_MODEL } from '@/lib/embeddings/openrouter-models'
Expand Down
3 changes: 2 additions & 1 deletion apps/sim/lib/internal/knowledge/execute-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
waleedlatif1 marked this conversation as resolved.
return descriptorResponse(
policy.unhandled?.() ?? { status: 500, body: { error: 'Internal server error' } }
)
Expand Down
60 changes: 59 additions & 1 deletion apps/sim/lib/knowledge/api/route-policies.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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([
Expand Down
30 changes: 30 additions & 0 deletions apps/sim/lib/knowledge/api/route-policies.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'

Expand All @@ -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) {
Expand All @@ -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 })
Comment thread
waleedlatif1 marked this conversation as resolved.
}
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' }),
Expand Down
Loading