From d719c6e2fd2fb96b050ddf62531a4459d45988d1 Mon Sep 17 00:00:00 2001 From: Theodore Li Date: Wed, 16 Sep 2026 19:23:13 -0700 Subject: [PATCH] fix(knowledge): increase hosted rerank capacity --- apps/sim/lib/core/config/env.ts | 5 +- .../rate-limiter/provider-admission.test.ts | 132 ++++++++++++++++++ .../core/rate-limiter/provider-admission.ts | 26 +++- apps/sim/lib/knowledge/reranker.test.ts | 73 +++++++++- apps/sim/lib/knowledge/reranker.ts | 16 ++- 5 files changed, 240 insertions(+), 12 deletions(-) diff --git a/apps/sim/lib/core/config/env.ts b/apps/sim/lib/core/config/env.ts index 404f367860d..2fcfd17badb 100644 --- a/apps/sim/lib/core/config/env.ts +++ b/apps/sim/lib/core/config/env.ts @@ -471,7 +471,10 @@ export const env = createEnv({ KB_CONFIG_MISTRAL_OCR_MAX_CONCURRENT: z.number().int().positive().max(64).optional().default(2), /** JSON map from API-key SHA-256 fingerprints to organization IDs; keys in one org share capacity. */ MISTRAL_OCR_QUOTA_GROUPS: z.string().optional(), - KB_CONFIG_RERANK_REQUESTS_PER_MINUTE: z.number().positive().optional().default(60), + /** Explicit override for all rerank credentials; otherwise defaults to 60, or 600 for hosted Cohere. */ + KB_CONFIG_RERANK_REQUESTS_PER_MINUTE: z.number().positive().optional(), + /** Overrides the shared rerank setting only for Sim-hosted Cohere credentials. */ + KB_CONFIG_HOSTED_RERANK_REQUESTS_PER_MINUTE: z.number().positive().optional(), KB_CONFIG_DOCUMENT_CONCURRENCY: z.number().optional().default(4), // Concurrent documents in the in-process (non-Trigger) path KB_CONFIG_BATCH_SIZE: z.number().optional().default(2000), // Chunks to process per embedding batch KB_CONFIG_DOCUMENT_BATCH_SIZE: z.number().optional().default(10), // Documents per batch in the in-process (non-Trigger) path diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts index 76d71ce4a81..6e3c80bcacf 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts @@ -1,6 +1,7 @@ /** * @vitest-environment node */ +import { dbChainMockFns, resetDbChainMock } from '@sim/testing' import { resetEnvMock, setEnv } from '@sim/testing/mocks/env.mock' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' @@ -18,6 +19,7 @@ vi.mock('@/lib/core/rate-limiter/storage/factory', () => ({ })) import { waitForProviderAdmission } from '@/lib/core/rate-limiter/provider-admission' +import { DbTokenBucket } from '@/lib/core/rate-limiter/storage/db-token-bucket' import { retryWithExponentialBackoff } from '@/lib/knowledge/documents/utils' const INPUT = { @@ -32,6 +34,11 @@ describe('provider admission', () => { beforeEach(() => { vi.useFakeTimers() vi.clearAllMocks() + resetDbChainMock() + setEnv({ + KB_CONFIG_RERANK_REQUESTS_PER_MINUTE: undefined, + KB_CONFIG_HOSTED_RERANK_REQUESTS_PER_MINUTE: undefined, + }) getCooldownUntil.mockResolvedValue(null) consumeTokens.mockResolvedValue({ allowed: true, tokensRemaining: 1, resetAt: new Date() }) }) @@ -66,6 +73,131 @@ describe('provider admission', () => { expect(consumeTokens).toHaveBeenCalledTimes(2) }) + it.each([ + { isHostedCredential: true, maxTokens: 16, refillRate: 10 }, + { isHostedCredential: false, maxTokens: 2, refillRate: 1 }, + { isHostedCredential: undefined, maxTokens: 2, refillRate: 1 }, + ])('selects the rerank budget for hosted=$isHostedCredential', async (fixture) => { + await waitForProviderAdmission({ + ...INPUT, + operation: 'rerank', + providerId: 'cohere', + isHostedCredential: fixture.isHostedCredential, + }) + expect(consumeTokens.mock.calls[0][0]).toEqual([ + { + key: 'provider:rerank:cohere:hashed-credential:requests', + cost: 1, + config: { + maxTokens: fixture.maxTokens, + refillRate: fixture.refillRate, + refillIntervalMs: 1000, + }, + }, + ]) + }) + + it('preserves the shared override unless a hosted-specific override is set', async () => { + setEnv({ KB_CONFIG_RERANK_REQUESTS_PER_MINUTE: '120' }) + const input = { ...INPUT, operation: 'rerank' as const, providerId: 'cohere' } + await waitForProviderAdmission({ ...input, isHostedCredential: true }) + await waitForProviderAdmission(input) + setEnv({ KB_CONFIG_HOSTED_RERANK_REQUESTS_PER_MINUTE: '300' }) + await waitForProviderAdmission({ ...input, isHostedCredential: true }) + await waitForProviderAdmission(input) + expect(consumeTokens.mock.calls.map(([reservations]) => reservations[0].config)).toEqual([ + { maxTokens: 16, refillRate: 2, refillIntervalMs: 1000 }, + { maxTokens: 2, refillRate: 2, refillIntervalMs: 1000 }, + { maxTokens: 16, refillRate: 5, refillIntervalMs: 1000 }, + { maxTokens: 2, refillRate: 2, refillIntervalMs: 1000 }, + ]) + }) + + it('caps the hosted burst when the configured minute budget is smaller', async () => { + setEnv({ KB_CONFIG_HOSTED_RERANK_REQUESTS_PER_MINUTE: '1' }) + await waitForProviderAdmission({ + ...INPUT, + operation: 'rerank', + providerId: 'cohere', + isHostedCredential: true, + }) + expect(consumeTokens.mock.calls[0][0][0].config).toMatchObject({ + maxTokens: 1, + refillRate: 1 / 60, + }) + }) + + it.each(['0', '-1', '', 'invalid', 'Infinity'])( + 'rejects an invalid hosted rerank override (%s) before spending capacity', + async (value) => { + setEnv({ KB_CONFIG_HOSTED_RERANK_REQUESTS_PER_MINUTE: value }) + await expect( + waitForProviderAdmission({ + ...INPUT, + operation: 'rerank', + providerId: 'cohere', + isHostedCredential: true, + }) + ).rejects.toThrow('Hosted rerank requests per minute must be finite and at least 1') + expect(consumeTokens).not.toHaveBeenCalled() + } + ) + + it.each([ + { operation: 'embedding', providerId: 'openai', maxTokens: 64, refillRate: 10 }, + { operation: 'ocr', providerId: 'mistral', maxTokens: 2, refillRate: 1 }, + { operation: 'rerank', providerId: 'another-provider', maxTokens: 2, refillRate: 1 }, + ] as const)('preserves the $operation budget for $providerId', async (fixture) => { + setEnv({ KB_CONFIG_HOSTED_RERANK_REQUESTS_PER_MINUTE: '300' }) + await waitForProviderAdmission({ + ...INPUT, + operation: fixture.operation, + providerId: fixture.providerId, + isHostedCredential: true, + }) + const reservations = consumeTokens.mock.calls[0][0] + expect(reservations.at(-1).config).toMatchObject({ + maxTokens: fixture.maxTokens, + refillRate: fixture.refillRate, + }) + }) + + it('sustains 600 hosted reranks per minute through the real bucket refill calculation', async () => { + const input = { + ...INPUT, + operation: 'rerank' as const, + providerId: 'cohere', + isHostedCredential: true, + maxWaitMs: 1, + } + let stored: { key: string; tokens: string; lastRefillAt: Date } | undefined + dbChainMockFns.values.mockImplementation((rows) => { + stored ??= rows.find((row: { key: string }) => row.key.endsWith(':requests')) + return { onConflictDoNothing: vi.fn().mockResolvedValue(undefined) } + }) + dbChainMockFns.limit.mockImplementation(async () => [stored]) + dbChainMockFns.set.mockImplementation((values) => { + Object.assign(stored!, values) + return { where: vi.fn().mockResolvedValue(undefined) } + }) + const bucket = new DbTokenBucket() + consumeTokens.mockImplementation((reservations, options) => + bucket.consumeTokensAtomically(reservations, options) + ) + + for (let request = 0; request < 16; request++) await waitForProviderAdmission(input) + await expect(waitForProviderAdmission(input)).rejects.toMatchObject({ retryAfterMs: 1000 }) + for (let second = 0; second < 60; second++) { + await vi.advanceTimersByTimeAsync(1000) + for (let request = 0; request < 10; request++) await waitForProviderAdmission(input) + await expect(waitForProviderAdmission(input)).rejects.toMatchObject({ retryAfterMs: 1000 }) + } + expect(stored?.tokens).toBe('0') + expect(new Set(consumeTokens.mock.calls.map(([reservations]) => reservations[0].key))).toEqual( + new Set(['provider:rerank:cohere:hashed-credential:requests']) + ) + }) + it('stops waiting immediately when the caller aborts', async () => { consumeTokens.mockResolvedValue({ allowed: false, retryAfterMs: 5000 }) const controller = new AbortController() diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.ts b/apps/sim/lib/core/rate-limiter/provider-admission.ts index b9b191c2f81..b516db90a16 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.ts @@ -21,6 +21,8 @@ export interface ProviderIdentity { const BULK_LANE_SHARE = 0.9 interface ProviderAdmissionInput extends ProviderIdentity { + /** True only for platform-owned credentials resolved on hosted Sim. Does not change bucket identity. */ + isHostedCredential?: boolean inputTokens?: number signal?: AbortSignal maxWaitMs: number @@ -35,8 +37,20 @@ interface ProviderAdmissionInput extends ProviderIdentity { * race for a handful of slots while the token budget sits unused. */ const EMBEDDING_REQUEST_BURST = 64 +const HOSTED_RERANK_REQUEST_BURST = 16 const DEFAULT_REQUEST_BURST = 2 +function hostedRerankRequestsPerMinute(): number { + const configured = + env.KB_CONFIG_HOSTED_RERANK_REQUESTS_PER_MINUTE ?? env.KB_CONFIG_RERANK_REQUESTS_PER_MINUTE + if (configured === undefined) return 600 + const requestsPerMinute = Number(configured) + if (!Number.isFinite(requestsPerMinute) || requestsPerMinute < 1) { + throw new Error('Hosted rerank requests per minute must be finite and at least 1') + } + return requestsPerMinute +} + /** A local admission wait expired; the document scheduler may retry the work later. */ export class ProviderAdmissionTimeoutError extends Error { readonly retryable = false @@ -59,12 +73,16 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P input.signal?.throwIfAborted() const deadlineAt = Date.now() + input.maxWaitMs const key = providerKey(input) + const isHostedRerank = + input.operation === 'rerank' && input.providerId === 'cohere' && input.isHostedCredential const requestsPerMinute = input.operation === 'embedding' ? envNumber(env.KB_CONFIG_EMBEDDING_REQUESTS_PER_MINUTE, 600, { min: 1 }) : input.operation === 'ocr' ? envNumber(env.KB_CONFIG_OCR_REQUESTS_PER_MINUTE, 60, { min: 1 }) - : envNumber(env.KB_CONFIG_RERANK_REQUESTS_PER_MINUTE, 60, { min: 1 }) + : isHostedRerank + ? hostedRerankRequestsPerMinute() + : envNumber(env.KB_CONFIG_RERANK_REQUESTS_PER_MINUTE, 60, { min: 1 }) const tokenBudget = input.operation === 'embedding' && input.inputTokens ? { @@ -77,7 +95,11 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P throw new Error('Embedding request exceeds the configured per-credential token budget') } const requestBurst = Math.min( - input.operation === 'embedding' ? EMBEDDING_REQUEST_BURST : DEFAULT_REQUEST_BURST, + input.operation === 'embedding' + ? EMBEDDING_REQUEST_BURST + : isHostedRerank + ? HOSTED_RERANK_REQUEST_BURST + : DEFAULT_REQUEST_BURST, requestsPerMinute ) const reservations: TokenBucketReservation[] = [] diff --git a/apps/sim/lib/knowledge/reranker.test.ts b/apps/sim/lib/knowledge/reranker.test.ts index 75ea03141a2..420bb229833 100644 --- a/apps/sim/lib/knowledge/reranker.test.ts +++ b/apps/sim/lib/knowledge/reranker.test.ts @@ -2,6 +2,8 @@ * @vitest-environment node */ import { setupGlobalFetchMock } from '@sim/testing/mocks' +import { setEnv } from '@sim/testing/mocks/env.mock' +import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing/mocks/env-flags.mock' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import type { AtomicAdmissionOptions, @@ -13,6 +15,8 @@ const admission = vi.hoisted(() => ({ setCooldown: vi.fn(), cooldowns: new Map(), })) +const { getBYOKKey } = vi.hoisted(() => ({ getBYOKKey: vi.fn() })) +vi.mock('@/lib/api-key/byok', () => ({ getBYOKKey })) vi.mock('@/lib/core/rate-limiter/storage/factory', () => ({ createStorageAdapter: () => ({ consumeTokensAtomically: admission.consume, @@ -31,6 +35,15 @@ const envSnapshot = { ...env } describe('Knowledge reranker model boundary', () => { beforeEach(() => { vi.clearAllMocks() + setEnvFlags({ isHosted: true }) + getBYOKKey.mockResolvedValue(null) + setEnv({ + KB_CONFIG_RERANK_REQUESTS_PER_MINUTE: undefined, + KB_CONFIG_HOSTED_RERANK_REQUESTS_PER_MINUTE: undefined, + COHERE_API_KEY_1: undefined, + COHERE_API_KEY_2: undefined, + COHERE_API_KEY_3: undefined, + }) admission.cooldowns.clear() admission.consume.mockImplementation( async (_reservations: readonly TokenBucketReservation[], options: AtomicAdmissionOptions) => { @@ -55,11 +68,59 @@ describe('Knowledge reranker model boundary', () => { afterEach(() => { vi.useRealTimers() + resetEnvFlagsMock() vi.unstubAllGlobals() for (const key of Object.keys(env)) delete (env as Record)[key] Object.assign(env, envSnapshot) }) + it.each([ + { hosted: true, source: 'env', expectedKey: 'cohere-key', burst: 16, refill: 10 }, + { hosted: true, source: 'rotation', expectedKey: 'rotating-key', burst: 16, refill: 10 }, + { hosted: true, source: 'workspace', expectedKey: 'byok-key', burst: 2, refill: 1 }, + { hosted: true, source: 'organization', expectedKey: 'byok-key', burst: 2, refill: 1 }, + { hosted: false, source: 'user', expectedKey: 'user-key', burst: 2, refill: 1 }, + { hosted: false, source: 'env', expectedKey: 'cohere-key', burst: 2, refill: 1 }, + { hosted: false, source: 'rotation', expectedKey: 'rotating-key', burst: 2, refill: 1 }, + { hosted: false, source: 'workspace', expectedKey: 'byok-key', burst: 2, refill: 1 }, + { hosted: false, source: 'organization', expectedKey: 'byok-key', burst: 2, refill: 1 }, + ])('uses the $source credential budget on hosted=$hosted', async (fixture) => { + setEnvFlags({ isHosted: fixture.hosted }) + const isBYOK = fixture.source === 'workspace' || fixture.source === 'organization' + if (isBYOK) { + getBYOKKey.mockResolvedValue({ apiKey: 'byok-key', scope: fixture.source, isBYOK: true }) + } + if (fixture.source === 'rotation') { + setEnv({ COHERE_API_KEY: undefined, COHERE_API_KEY_1: 'rotating-key' }) + } + const result = await rerank('query', [{ id: 'one', text: 'content' }], { + model: 'rerank-v4.0-fast', + workspaceId: 'fixture-workspace', + apiKey: fixture.hosted || fixture.source === 'user' ? 'user-key' : undefined, + }) + expect(result.isBYOK).toBe(isBYOK) + expect(fetch).toHaveBeenCalledWith( + 'https://api.cohere.com/v2/rerank', + expect.objectContaining({ + headers: expect.objectContaining({ Authorization: `Bearer ${fixture.expectedKey}` }), + }) + ) + expect(admission.consume.mock.calls[0][0]).toMatchObject([ + { config: { maxTokens: fixture.burst, refillRate: fixture.refill, refillIntervalMs: 1000 } }, + ]) + if (fixture.source === 'user') expect(getBYOKKey).not.toHaveBeenCalled() + else expect(getBYOKKey).toHaveBeenCalledWith('fixture-workspace', 'cohere') + }) + + it('fails before admission when no credential is configured', async () => { + setEnv({ COHERE_API_KEY: undefined }) + await expect( + rerank('query', [{ id: 'one', text: 'content' }], { model: 'rerank-v4.0-fast' }) + ).rejects.toThrow('No Cohere API key configured') + expect(admission.consume).not.toHaveBeenCalled() + expect(fetch).not.toHaveBeenCalled() + }) + it('projects query and documents at egress while returning the original item', async () => { const registry = new ResolvedSecretTraceRegistry([ { name: 'TOKEN', plaintext: 'secret-value', encryptedValue: 'encrypted-token' }, @@ -132,10 +193,16 @@ describe('Knowledge reranker model boundary', () => { vi.mocked(fetch).mockResolvedValueOnce( new Response('{}', { status: 429, headers: { 'Retry-After': '2' } }) ) - const first = rerank('first', [{ id: 'one', text: 'content' }], { model: 'rerank-v4.0-fast' }) + const first = rerank('first', [{ id: 'one', text: 'content' }], { + model: 'rerank-v4.0-fast', + workspaceId: 'fixture-workspace-one', + }) await vi.advanceTimersByTimeAsync(0) expect(admission.setCooldown).toHaveBeenCalledOnce() - const second = rerank('second', [{ id: 'two', text: 'content' }], { model: 'rerank-v4.0-fast' }) + const second = rerank('second', [{ id: 'two', text: 'content' }], { + model: 'rerank-v4.0-fast', + workspaceId: 'fixture-workspace-two', + }) await vi.advanceTimersByTimeAsync(1999) expect(fetch).toHaveBeenCalledTimes(1) await vi.advanceTimersByTimeAsync(1) @@ -147,7 +214,7 @@ describe('Knowledge reranker model boundary', () => { expect(new Set(reservations.map((item) => item.key)).size).toBe(1) expect(reservations[0].key).toMatch(/^provider:rerank:cohere:[a-f0-9]{64}:requests$/) expect(reservations[0].key).not.toContain('cohere-key') - expect(reservations[0].config.refillRate).toBe(1) + expect(reservations[0].config).toMatchObject({ maxTokens: 16, refillRate: 10 }) }) it('bounds repeated 429s to four attempts with no timer left behind', async () => { diff --git a/apps/sim/lib/knowledge/reranker.ts b/apps/sim/lib/knowledge/reranker.ts index 56440caa2dc..085f8da7d47 100644 --- a/apps/sim/lib/knowledge/reranker.ts +++ b/apps/sim/lib/knowledge/reranker.ts @@ -68,7 +68,7 @@ class RerankAPIError extends Error { async function resolveCohereKey( workspaceId?: string | null, userApiKey?: string -): Promise<{ apiKey: string; isBYOK: boolean }> { +): Promise<{ apiKey: string; isBYOK: boolean; isHostedCredential: boolean }> { /** * Mirrors the agent block hosted-key pattern (`injectHostedKeyIfNeeded`): * on self-hosted the user-supplied key from the block field flows through @@ -76,20 +76,20 @@ async function resolveCohereKey( * platform env, so any user-supplied value is ignored. */ if (!isHosted && userApiKey) { - return { apiKey: userApiKey, isBYOK: false } + return { apiKey: userApiKey, isBYOK: false, isHostedCredential: false } } if (workspaceId) { const byokResult = await getBYOKKey(workspaceId, 'cohere') if (byokResult) { logger.info('Using BYOK key for Cohere reranker', { scope: byokResult.scope }) - return { apiKey: byokResult.apiKey, isBYOK: true } + return { apiKey: byokResult.apiKey, isBYOK: true, isHostedCredential: false } } } if (env.COHERE_API_KEY) { - return { apiKey: env.COHERE_API_KEY, isBYOK: false } + return { apiKey: env.COHERE_API_KEY, isBYOK: false, isHostedCredential: isHosted } } try { - return { apiKey: getRotatingApiKey('cohere'), isBYOK: false } + return { apiKey: getRotatingApiKey('cohere'), isBYOK: false, isHostedCredential: isHosted } } catch { throw new Error( 'No Cohere API key configured. Set COHERE_API_KEY_1/2/3 (rotation) or COHERE_API_KEY.' @@ -135,7 +135,10 @@ export async function rerank( throw new Error(`Unsupported reranker model: ${options.model}`) } - const { apiKey, isBYOK } = await resolveCohereKey(options.workspaceId, options.apiKey) + const { apiKey, isBYOK, isHostedCredential } = await resolveCohereKey( + options.workspaceId, + options.apiKey + ) const cappedItems = items.length > MAX_DOCUMENTS_PER_RERANK ? items.slice(0, MAX_DOCUMENTS_PER_RERANK) : items if (items.length > MAX_DOCUMENTS_PER_RERANK) { @@ -155,6 +158,7 @@ export async function rerank( async (signal) => { await waitForProviderAdmission({ ...identity, + isHostedCredential, signal, maxWaitMs: Math.max(0, deadlineAt - Date.now()), })