Improve organization search ranking and bound rerank admission

Author: waleedlatif1

apps/sim/lib/core/rate-limiter/provider-admission.test.ts

@@ -35,7 +35,66 @@ describe('provider admission', () => {
     consumeTokens.mockResolvedValue({ allowed: true, tokensRemaining: 1, resetAt: new Date() })
   })
 
-  afterEach(() => vi.useRealTimers())
+  afterEach(() => {
+    vi.useRealTimers()
+    vi.unstubAllEnvs()
+  })
+
+  it.each([
+    { requestsPerMinute: 1, capacity: 1 },
+    { requestsPerMinute: 60, capacity: 2 },
+    { requestsPerMinute: 120, capacity: 2 },
+    { requestsPerMinute: 121, capacity: 3 },
+    { requestsPerMinute: 600, capacity: 10 },
+  ])(
+    'admits a complete rerank refill at $requestsPerMinute requests per minute',
+    async ({ requestsPerMinute, capacity }) => {
+      vi.stubEnv('KB_CONFIG_RERANK_REQUESTS_PER_MINUTE', String(requestsPerMinute))
+
+      await waitForProviderAdmission({ ...INPUT, operation: 'rerank', providerId: 'cohere' })
+
+      expect(consumeTokens).toHaveBeenCalledExactlyOnceWith(
+        [
+          {
+            key: 'provider:rerank:cohere:hashed-credential:requests',
+            cost: 1,
+            config: {
+              maxTokens: capacity,
+              refillRate: requestsPerMinute / 60,
+              refillIntervalMs: 1000,
+            },
+          },
+        ],
+        expect.objectContaining({
+          cooldownKeys: [
+            'provider:rerank:cohere:hashed-credential:cooldown',
+            'provider:rerank:cohere:hashed-credential:quota',
+          ],
+        })
+      )
+    }
+  )
+
+  it.each([
+    { operation: 'embedding' as const, capacity: 8, refillRate: 10 },
+    { operation: 'ocr' as const, capacity: 2, refillRate: 1 },
+    { operation: 'rerank' as const, capacity: 2, refillRate: 1 },
+  ])('preserves default $operation request limits', async ({ operation, capacity, refillRate }) => {
+    vi.stubEnv('KB_CONFIG_EMBEDDING_REQUESTS_PER_MINUTE', undefined)
+    vi.stubEnv('KB_CONFIG_OCR_REQUESTS_PER_MINUTE', undefined)
+    vi.stubEnv('KB_CONFIG_RERANK_REQUESTS_PER_MINUTE', undefined)
+
+    await waitForProviderAdmission({ ...INPUT, operation })
+
+    expect(consumeTokens.mock.calls[0][0]).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          key: `provider:${operation}:openai:hashed-credential:requests`,
+          config: { maxTokens: capacity, refillRate, refillIntervalMs: 1000 },
+        }),
+      ])
+    )
+  })
 
   it('shares both credential dimensions in one reservation across concurrent callers', async () => {
     await Promise.all([waitForProviderAdmission(INPUT), waitForProviderAdmission(INPUT)])

apps/sim/lib/core/rate-limiter/provider-admission.ts

@@ -46,6 +46,13 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P
       : 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 })
+  /** Rerank capacity must retain a complete refill from the backends' whole-second ticks. */
+  const requestBurst =
+    input.operation === 'embedding'
+      ? 8
+      : input.operation === 'rerank'
+        ? Math.max(2, Math.ceil(requestsPerMinute / 60))
+        : 2
   const reservations: TokenBucketReservation[] = []
   if (input.operation === 'embedding' && input.inputTokens) {
     const tokensPerMinute = envNumber(env.KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE, 600_000, {
@@ -68,7 +75,7 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P
     key: `${key}:requests`,
     cost: 1,
     config: {
-      maxTokens: Math.min(input.operation === 'embedding' ? 8 : 2, requestsPerMinute),
+      maxTokens: Math.min(requestBurst, requestsPerMinute),
       refillRate: requestsPerMinute / 60,
       refillIntervalMs: 1000,
     },

apps/sim/lib/knowledge/application/contexts.test.ts

@@ -6,6 +6,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const mocks = vi.hoisted(() => ({
   getKnowledgeBase: vi.fn(),
+  getKnowledgeBaseWithCounts: vi.fn(),
   getDocument: vi.fn(),
   getDocumentById: vi.fn(),
   getTag: vi.fn(),
@@ -22,7 +23,8 @@ vi.mock('@/lib/knowledge/access/scope', () => ({
 }))
 
 vi.mock('@/lib/knowledge/service', () => ({
-  getKnowledgeBaseById: mocks.getKnowledgeBase,
+  getActiveKnowledgeBaseReference: mocks.getKnowledgeBase,
+  getKnowledgeBaseById: mocks.getKnowledgeBaseWithCounts,
 }))
 
 vi.mock('@/lib/knowledge/documents/service', () => ({
@@ -66,10 +68,76 @@ describe('knowledge application contexts', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     mocks.getKnowledgeBase.mockResolvedValue(knowledgeBase)
+    mocks.getKnowledgeBaseWithCounts.mockResolvedValue({
+      ...knowledgeBase,
+      docCount: 3,
+      tokenCount: 1536,
+    })
     mocks.loadWorkspace.mockResolvedValue(workspace)
     mocks.loadWorkspaceIncludingArchived.mockResolvedValue(workspace)
   })
 
+  it('resolves child-resource context without loading display counts', async () => {
+    const context = await resolveActiveKnowledgeResourceContext(
+      { knowledgeBaseId: 'knowledge-1', assertedWorkspaceId: 'workspace-1' },
+      principal
+    )
+
+    expect(context.knowledgeBase).toEqual(knowledgeBase)
+    expect(mocks.getKnowledgeBaseWithCounts).not.toHaveBeenCalled()
+    expect(mocks.createAccessProvider).toHaveBeenCalledWith(principal, {
+      workspaceId: 'workspace-1',
+    })
+  })
+
+  it('retains display counts in workspace knowledge-base detail context', async () => {
+    const context = await resolveActiveKnowledgeBaseContext(
+      { knowledgeBaseId: 'knowledge-1' },
+      principal
+    )
+
+    expect(context.knowledgeBase).toMatchObject({ docCount: 3, tokenCount: 1536 })
+    expect(mocks.getKnowledgeBase).not.toHaveBeenCalled()
+  })
+
+  it('conceals a missing or archived knowledge-base reference before loading its owner', async () => {
+    mocks.getKnowledgeBase.mockResolvedValueOnce(null)
+
+    await expect(
+      resolveActiveKnowledgeResourceContext({ knowledgeBaseId: 'archived' }, principal)
+    ).rejects.toMatchObject({ code: 'not_found' })
+    expect(mocks.loadWorkspace).not.toHaveBeenCalled()
+    expect(mocks.createAccessProvider).not.toHaveBeenCalled()
+  })
+
+  it('conceals an organization reference outside the asserted organization', async () => {
+    mocks.getKnowledgeBase.mockResolvedValueOnce({
+      id: 'org-index',
+      workspaceId: null,
+      organizationId: 'org-canonical',
+    })
+
+    await expect(
+      resolveActiveKnowledgeResourceContext(
+        { knowledgeBaseId: 'org-index', assertedOrganizationId: 'org-other' },
+        principal
+      )
+    ).rejects.toMatchObject({ code: 'not_found' })
+    expect(mocks.createAccessProvider).not.toHaveBeenCalled()
+  })
+
+  it('conceals a reference with conflicting owners', async () => {
+    mocks.getKnowledgeBase.mockResolvedValueOnce({
+      ...knowledgeBase,
+      organizationId: 'org-canonical',
+    })
+
+    await expect(
+      resolveActiveKnowledgeResourceContext({ knowledgeBaseId: 'knowledge-1' }, principal)
+    ).rejects.toMatchObject({ code: 'not_found' })
+    expect(mocks.createAccessProvider).not.toHaveBeenCalled()
+  })
+
   it('uses the canonical active-workspace loader', async () => {
     await expect(resolveKnowledgeWorkspaceContext({ workspaceId: 'workspace-1' })).resolves.toBe(
       workspace

apps/sim/lib/knowledge/application/contexts.ts

@@ -21,7 +21,11 @@ import {
   getRestorableKnowledgeBase,
   type RestorableKnowledgeBase,
 } from '@/lib/knowledge/orchestration/restore'
-import { getKnowledgeBaseById } from '@/lib/knowledge/service'
+import {
+  type ActiveKnowledgeBaseReference,
+  getActiveKnowledgeBaseReference,
+  getKnowledgeBaseById,
+} from '@/lib/knowledge/service'
 import { getTagDefinitionById } from '@/lib/knowledge/tags/service'
 import type { DocumentTagDefinition } from '@/lib/knowledge/tags/types'
 import type { KnowledgeBaseWithCounts } from '@/lib/knowledge/types'
@@ -71,7 +75,7 @@ export interface ActiveKnowledgeBaseContext
 export type ActiveKnowledgeResourceBaseContext = KnowledgeResourceContext &
   KnowledgeAccessBearingContext & {
     knowledgeBaseId: string
-    knowledgeBase: KnowledgeBaseWithCounts
+    knowledgeBase: ActiveKnowledgeBaseReference
   }
 
 export type ActiveKnowledgeDocumentContext = ActiveKnowledgeResourceBaseContext & {
@@ -227,7 +231,7 @@ export async function resolveActiveKnowledgeResourceContext(
   },
   principal: Principal
 ): Promise<ActiveKnowledgeResourceBaseContext> {
-  const knowledgeBase = await getKnowledgeBaseById(input.knowledgeBaseId)
+  const knowledgeBase = await getActiveKnowledgeBaseReference(input.knowledgeBaseId)
   if (
     !knowledgeBase ||
     (input.assertedOrganizationId !== undefined &&