improvement(knowledge): rank inside the permitted set for organization search (#7996)

* improvement(knowledge): rank inside the permitted set for organization search

* fix(knowledge): count a caller's reach inside the requested bases

Author: waleedlatif1

apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts

@@ -3,14 +3,18 @@ import type { Principal } from '@sim/auth/principal'
 import { db } from '@sim/db'
 import { document, embedding, knowledgeBase, organization, user, workspace } from '@sim/db/schema'
 import { generateId } from '@sim/utils/id'
-import { eq, inArray } from 'drizzle-orm'
+import { eq, inArray, sql } from 'drizzle-orm'
 import { afterAll, beforeAll, describe, expect, it } from 'vitest'
 import {
   createKnowledgeAclFixtureIds,
   seedKnowledgeAclFixture,
 } from '@/lib/knowledge/__integration__/seed-source-access-fixture'
 import { createKnowledgeAccessProvider } from '@/lib/knowledge/access/scope'
-import { retrieveKnowledgeSearch } from '@/lib/knowledge/search/queries'
+import {
+  resolvePermittedDocuments,
+  retrieveKnowledgeSearch,
+  VECTOR_PROBE_DOCUMENT_LIMIT,
+} from '@/lib/knowledge/search/queries'
 import { embeddingVectorValues } from '@/lib/knowledge/vector-columns'
 
 describe('API-key KB block fan-out', () => {
@@ -140,13 +144,45 @@ describe('API-key KB block fan-out', () => {
         expect(matching('AS visible')).toHaveLength(bases.length)
         expect(matching(') + 0 LIMIT')).toHaveLength(bases.length)
         expect(matching('scored_search_candidates')).toHaveLength(bases.length)
-        /** The probe enumerates visible documents; it never ranks them. */
+        /** The probe enumerates visible documents and reports saturation; it never ranks them. */
         expect(
-          statements.filter((query) => query.includes('AS id FROM') && !query.includes('ORDER BY'))
+          statements.filter(
+            (query) => query.includes('AS saturated') && !query.includes('ORDER BY')
+          )
         ).toHaveLength(bases.length)
       } finally {
         db.$client.options.debug = previousDebug
       }
     }
   )
+
+  it('bounds the permitted set by the requested bases, not by what the tokens reach elsewhere', async () => {
+    const crowded = generateId()
+    await db.insert(knowledgeBase).values({
+      id: crowded,
+      userId: ids.aliceId,
+      workspaceId: ids.workspaceId,
+      name: 'Crowded neighbour',
+    })
+    try {
+      /** Baseline tokens are shared by every tenant, so another base can hold more than the limit. */
+      await db.execute(sql`
+        INSERT INTO ${document} (id, knowledge_base_id, filename, file_url, file_size, mime_type,
+          processing_status, acl)
+        SELECT 'crowded-' || n, ${crowded}, 'crowded', 'https://fixture.invalid/crowded', 1,
+          'text/plain', 'completed', ARRAY['ws']::text[]
+        FROM generate_series(1, ${VECTOR_PROBE_DOCUMENT_LIMIT + 1}) AS n
+      `)
+      const permitted = await resolvePermittedDocuments({
+        knowledgeBaseIds: [bases[0].id],
+        access: { kind: 'user', userId: ids.bobId, tokens: ['pub', 'ws'] },
+      })
+      expect(permitted).toEqual({
+        kind: 'bounded',
+        documents: [{ id: bases[0].visible, connectorId: null }],
+      })
+    } finally {
+      await db.delete(knowledgeBase).where(eq(knowledgeBase.id, crowded))
+    }
+  })
 })

apps/sim/lib/knowledge/access/predicate.ts

@@ -193,6 +193,26 @@ export function knowledgeMetadataCandidateAccessCondition(
   return storedKnowledgeAccessCondition(scope, sql`true`)
 }
 
+/**
+ * A members-mode document is readable while one of the caller's active member identities on its
+ * connector still observes it, freshly. Correlated on the document so each check is a lookup on
+ * the observation primary key, which leads with `document_id`: phrased as a row-value `IN`
+ * inside the access predicate's `OR`, PostgreSQL instead hashes every observation in the table
+ * once per statement, a fixed cost paid by every query that carries the predicate.
+ */
+function memberObservationCondition(tokens: SQL, cutoff: SQL): SQL {
+  return sql`EXISTS (
+    SELECT 1 FROM ${knowledgeDocumentObservation}
+    JOIN ${knowledgeConnectorMember}
+      ON ${knowledgeConnectorMember.id} = ${knowledgeDocumentObservation.memberId}
+    WHERE ${knowledgeDocumentObservation.documentId} = ${document.id}
+      AND ${knowledgeConnectorMember.connectorId} = ${document.connectorId}
+      AND ${knowledgeConnectorMember.status} = 'active'
+      AND ${knowledgeConnectorMember.subjectToken} = ANY(${tokens})
+      AND GREATEST(${knowledgeDocumentObservation.lastSeenAt}, ${knowledgeConnectorMember.memberSyncedThrough}) > ${cutoff}
+  )`
+}
+
 function storedKnowledgeAccessCondition(
   scope: KnowledgeAccessScope | SystemAccessScope,
   liveSourceAccess: SQL
@@ -202,7 +222,7 @@ function storedKnowledgeAccessCondition(
   const tokens = textArrayLiteral(scope.tokens)
   const cutoff = sql`statement_timestamp() - (${SOURCE_ACL_MAX_AGE_MS} * interval '1 millisecond')`
   return sql`(
-    ${document.acl} && ${tokens}
+    ${aclOverlap(tokens)}
     AND NOT EXISTS (
       SELECT 1 FROM jsonb_array_elements(${document.aclRequirements}) AS required_clause(tokens)
       WHERE NOT (required_clause.tokens ?| ${tokens})
@@ -221,21 +241,33 @@ function storedKnowledgeAccessCondition(
             (${knowledgeConnector.accessMode} = 'workspace' AND ${document.acl} = ARRAY['ws']::text[])
             OR (${document.acl} <> ARRAY['ws']::text[] AND (
             (${knowledgeConnector.accessMode} = 'admin' AND ${document.aclVerifiedAt} > ${cutoff})
-            OR (${knowledgeConnector.accessMode} = 'members' AND (${document.id}, ${document.connectorId}) IN (
-              SELECT ${knowledgeDocumentObservation.documentId}, ${knowledgeConnectorMember.connectorId} FROM ${knowledgeDocumentObservation}
-              JOIN ${knowledgeConnectorMember}
-                ON ${knowledgeConnectorMember.id} = ${knowledgeDocumentObservation.memberId}
-              WHERE ${knowledgeConnectorMember.status} = 'active'
-                AND ${knowledgeConnectorMember.subjectToken} = ANY(${tokens})
-                AND GREATEST(${knowledgeDocumentObservation.lastSeenAt}, ${knowledgeConnectorMember.memberSyncedThrough}) > ${cutoff}
-            ))
+            OR (${knowledgeConnector.accessMode} = 'members' AND ${memberObservationCondition(tokens, cutoff)})
             ))
           )
       )
     )
   )`
 }
 
+/**
+ * The token half of the stored access predicate: the documents a caller's tokens reach before
+ * any source, freshness, or requirement check narrows them. It is a necessary condition of
+ * {@link knowledgeAccessCondition}, never a substitute for it.
+ *
+ * Paired with `deleted_at IS NULL` it matches `doc_acl_gin_idx` exactly, so a query can enumerate
+ * a member's reachable documents from that index alone. PostgreSQL cannot estimate array-overlap
+ * selectivity, so left to itself it intersects this highly selective bitmap with base-wide ones.
+ */
+export function knowledgeAclOverlapCondition(scope: KnowledgeAccessScope): SQL {
+  if (scope.tokens.length === 0) return sql`false`
+  return aclOverlap(textArrayLiteral(scope.tokens))
+}
+
+/** One spelling of the token overlap, so the probe's reach and the full predicate cannot drift. */
+function aclOverlap(tokens: SQL): SQL {
+  return sql`${document.acl} && ${tokens}`
+}
+
 /**
  * The pool uses fetch_types: false, so arrays must be constructed from scalar
  * parameters. A JSON scalar keeps large sets below PostgreSQL's bind limit.

apps/sim/lib/knowledge/search/diagnostics.ts

@@ -28,6 +28,7 @@ export type SearchStage =
   | 'access_scope'
   | 'defaults'
   | 'retrieval'
+  | 'permitted_documents'
   | 'result_provenance'
   | 'reranking'
   | 'usage_recording'
@@ -89,6 +90,14 @@ export interface SearchDiagnosticMetadata {
   vectorCandidateLimit?: number
   /** Visible documents the tractability probe enumerated, capped at its own document limit. */
   vectorProbeDocumentCount?: number
+  /**
+   * Whether a user-scoped search resolved its permitted documents before retrieval: `bounded`
+   * ranks inside that set, `unbounded` means it exceeded the probe's limit and both legs search
+   * the index with the access predicate applied per candidate.
+   */
+  permittedDocuments?: 'bounded' | 'unbounded'
+  /** Documents in a bounded permitted set. */
+  permittedDocumentCount?: number
   vectorCandidateCount?: number
   vectorCandidateDimensions?: number
   resultCount?: number