fix(knowledge): increase hosted rerank capacity (#7912)

Author: TheodoreSpeaks

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

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()

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[] = []

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<string, Date>(),
 }))
+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<string, unknown>)[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 () => {