fix(knowledge): apply the source and date filters inside a resolved scope's ranking (#8059)

* fix(knowledge): apply the source and date filters inside a resolved scope's ranking

- a source filter confines the access plan itself, so every on-row predicate, the reach, and the sources the legs walk or rank are that kind of source alone; `upload` keeps only source-less documents
- a date filter enumerates the documents it admits off a new `(knowledge_base_id, source_modified_at)` index and ranks them exactly while the planner estimates the window within the probe's limit; a wider window is walked with the date tested through the document, on the row and on the ranked keyword row
- a date-bounded set is ranked exactly even when a member source has its own index, since a walk cannot see the date
- the reach is keyed and counted by the plan's sources
- a keyword leg whose deadline passes before its ranking is resolved is short rather than failed

* fix(knowledge): enumerate a small filtered source so both legs rank inside it

* chore(db): mark the document date lookup index concurrent in the schema

* fix(knowledge): estimate a filter's size under a short deadline of its own

Author: waleedlatif1

apps/sim/lib/knowledge/access/connector-eligibility.ts

@@ -22,14 +22,15 @@ import { searchIntegrationAccessCondition } from '@/lib/knowledge/search/integra
  */
 async function resolveConnectorEligibility(
   knowledgeBaseIds: readonly string[]
-): Promise<KnowledgeConnectorEligibility> {
+): Promise<{ eligibility: KnowledgeConnectorEligibility; types: Map<string, string> }> {
   const eligibility: {
     workspace: string[]
     admin: string[]
     members: string[]
     liveProofRequired: string[]
   } = { workspace: [], admin: [], members: [], liveProofRequired: [] }
-  if (knowledgeBaseIds.length === 0) return eligibility
+  const types = new Map<string, string>()
+  if (knowledgeBaseIds.length === 0) return { eligibility, types }
   const rows = await db
     .select({
       id: knowledgeConnector.id,
@@ -53,12 +54,13 @@ async function resolveConnectorEligibility(
     else if (row.accessMode === 'admin') eligibility.admin.push(row.id)
     else if (row.accessMode === 'members') eligibility.members.push(row.id)
     else continue
+    types.set(row.id, row.connectorType)
     const live =
       (row.connectorType === 'github' && row.githubRepository) ||
       (row.connectorType === 'confluence' && row.accessMode === 'admin')
     if (live) eligibility.liveProofRequired.push(row.id)
   }
-  return eligibility
+  return { eligibility, types }
 }
 
 /**
@@ -116,7 +118,8 @@ export async function resolveSearchAccessPlan(
   knowledgeBaseIds: readonly string[],
   access: KnowledgeAccessScope
 ): Promise<SearchAccessPlan> {
-  const connectors = await resolveConnectorEligibility(knowledgeBaseIds)
+  const { eligibility: connectors, types: connectorTypes } =
+    await resolveConnectorEligibility(knowledgeBaseIds)
   const { observers, memberSources } = await resolveMemberObservers(access, connectors.members)
-  return { connectors, observers, memberSources }
+  return { connectors, observers, memberSources, connectorTypes, uploads: true }
 }

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

@@ -30,6 +30,7 @@ const {
   knowledgeAccessCondition,
   knowledgeCandidateAccessConditionForConnectors,
   projectionCandidateAccessCondition,
+  restrictSearchAccessPlan,
   knowledgeMetadataCandidateAccessCondition,
 } = await import('@/lib/knowledge/access/predicate')
 const { confluencePageAcl } = await import('@/lib/knowledge/access/confluence-permissions')
@@ -557,7 +558,17 @@ describe.runIf(Boolean(databaseUrl))('knowledge ACLs in PostgreSQL', () => {
     /** What `resolveSearchAccessPlan` resolves for this caller: their member identity, confirmed. */
     const observers = { confirmed: [{ id: 'm-alice', connectorId: 'members' }], observed: [] }
     const perRow = knowledgeMetadataCandidateAccessCondition(scope)
-    const plan = { connectors: eligibility, observers, memberSources: ['members'] }
+    const plan = {
+      connectors: eligibility,
+      observers,
+      memberSources: ['members'],
+      connectorTypes: new Map([
+        ['ws-mode', 'slack'],
+        ['admin', 'google_drive'],
+        ['members', 'slack'],
+      ]),
+      uploads: true,
+    }
     const perQuery = knowledgeCandidateAccessConditionForConnectors(scope, plan)
     for (const id of [...cases.map(([documentId]) => documentId), 'upload-doc']) {
       expect([id, await admits(perQuery, id)]).toEqual([id, await admits(perRow, id)])
@@ -655,6 +666,44 @@ describe.runIf(Boolean(databaseUrl))('knowledge ACLs in PostgreSQL', () => {
         'admin-current'
       )
     ).toBe(false)
+    /**
+     * A plan confined to one kind of source admits that kind alone, on the document and on the
+     * row, and a plan confined to uploads admits only documents without a source.
+     */
+    const drive = restrictSearchAccessPlan(plan, 'google_drive')
+    const driveOnRow = new PgDialect().sqlToQuery(
+      projectionCandidateAccessCondition(schema.embeddingSearch, scope, drive)
+    )
+    const driveOnRowAdmits = async (id: string) => {
+      const rows = await connection.unsafe(
+        `SELECT 1 FROM embedding_search WHERE ${driveOnRow.sql} AND document_id = $${driveOnRow.params.length + 1}`,
+        [...(driveOnRow.params as string[]), id]
+      )
+      return rows.length > 0
+    }
+    const drivePerQuery = knowledgeCandidateAccessConditionForConnectors(scope, drive)
+    for (const [id, expected] of [
+      ['admin-current', true],
+      ['workspace-doc', false],
+      ['members-current', false],
+      ['upload-doc', false],
+    ] as const) {
+      expect([id, await admits(drivePerQuery, id)]).toEqual([id, expected])
+      expect([id, await driveOnRowAdmits(id)]).toEqual([id, expected])
+    }
+    /** Uploads carry the workspace ACL, so a caller with that token reads them and nothing sourced. */
+    await connection.unsafe("INSERT INTO document(id, acl) VALUES ('upload-mine', ARRAY['ws'])")
+    const wsScope = { ...scope, tokens: [...scope.tokens, 'ws'] }
+    const uploadsPerQuery = knowledgeCandidateAccessConditionForConnectors(
+      wsScope,
+      restrictSearchAccessPlan(plan, 'upload')
+    )
+    expect(await admits(knowledgeMetadataCandidateAccessCondition(wsScope), 'upload-mine')).toBe(
+      true
+    )
+    expect(await admits(uploadsPerQuery, 'upload-mine')).toBe(true)
+    expect(await admits(uploadsPerQuery, 'workspace-doc')).toBe(false)
+    expect(await admits(uploadsPerQuery, 'admin-current')).toBe(false)
     /** A connector left out of the resolution is refused, however current its documents are. */
     expect(
       await admits(

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

@@ -18,9 +18,14 @@ process.env.DATABASE_URL ??= 'postgresql://user:pass@localhost:5432/test'
 
 const { PgDialect } = await import('drizzle-orm/pg-core')
 const { embeddingSearch } = await import('@sim/db/schema')
-const { knowledgeAccessCondition, projectionCandidateAccessCondition } = await import(
-  '@/lib/knowledge/access/predicate'
-)
+const { sql: rawSql } = await import('drizzle-orm')
+const sqlColumn = (name: string) => rawSql.raw(`"row"."${name}"`)
+const {
+  knowledgeAccessCondition,
+  knowledgeCandidateAccessConditionForConnectors,
+  projectionCandidateAccessCondition,
+  restrictSearchAccessPlan,
+} = await import('@/lib/knowledge/access/predicate')
 const { SYSTEM_ACCESS_SCOPE } = await import('@/lib/knowledge/access/types')
 
 function render(condition: ReturnType<typeof knowledgeAccessCondition>) {
@@ -32,6 +37,8 @@ describe('projectionCandidateAccessCondition', () => {
     connectors: { workspace: ['ws-src'], admin: [], members: [], liveProofRequired: [] },
     observers: { confirmed: [], observed: [] },
     memberSources: [],
+    connectorTypes: new Map(),
+    uploads: true,
   }
 
   it('decides a filled row on its mirrored columns and an unfilled row on its document', () => {
@@ -46,8 +53,9 @@ describe('projectionCandidateAccessCondition', () => {
       '("embedding_search"."acl" IS NULL AND EXISTS (\n    SELECT 1 FROM "document"\n    WHERE "document"."id" = "embedding_search"."document_id"\n      AND ('
     )
     expect(sql).toContain('"document"."acl" && ARRAY[$1, $2]::text[]')
-    expect(sql).toMatch(
-      /OR \("embedding_search"\."acl" && ARRAY\[\$\d+, \$\d+\]::text\[\]\n {4}AND \("embedding_search"\."connector_id" IS NULL OR "embedding_search"\."connector_id" = ANY\(ARRAY\[\$\d+\]::text\[\]\)\)\)\)$/
+    expect(sql).toContain('OR ("embedding_search"."acl" && ARRAY[')
+    expect(sql).toContain(
+      'AND ("embedding_search"."connector_id" IS NULL OR "embedding_search"."connector_id" = ANY(ARRAY['
     )
     expect(params.slice(0, 2)).toEqual(['ws', 'u:alice'])
     expect(params.slice(-3)).toEqual(['ws', 'u:alice', 'ws-src'])
@@ -128,3 +136,78 @@ describe('knowledgeAccessCondition', () => {
     expect(sql).not.toContain('"document"."acl"')
   })
 })
+
+describe('restrictSearchAccessPlan', () => {
+  const plan = {
+    connectors: {
+      workspace: ['slack-ws'],
+      admin: ['drive-admin', 'confluence-admin'],
+      members: ['slack-members'],
+      liveProofRequired: ['confluence-admin'],
+    },
+    observers: {
+      confirmed: [{ id: 'm-1', connectorId: 'slack-members' }],
+      observed: [{ id: 'm-2', connectorId: 'drive-admin' }],
+    },
+    memberSources: ['slack-members'],
+    connectorTypes: new Map([
+      ['slack-ws', 'slack'],
+      ['slack-members', 'slack'],
+      ['drive-admin', 'google_drive'],
+      ['confluence-admin', 'confluence'],
+    ]),
+    uploads: true,
+  }
+  const reader = { kind: 'user' as const, userId: 'u', tokens: ['u:reader@example.com'] }
+
+  it('keeps only the connectors of that kind, in every list, and leaves uploads out', () => {
+    const slack = restrictSearchAccessPlan(plan, 'slack')
+    expect(slack.connectors).toEqual({
+      workspace: ['slack-ws'],
+      admin: [],
+      members: ['slack-members'],
+      liveProofRequired: [],
+    })
+    expect(slack.observers.confirmed).toEqual([{ id: 'm-1', connectorId: 'slack-members' }])
+    expect(slack.observers.observed).toEqual([])
+    expect(slack.memberSources).toEqual(['slack-members'])
+    expect(slack.uploads).toBe(false)
+  })
+
+  it('keeps no connector for uploads, which have none', () => {
+    const uploads = restrictSearchAccessPlan(plan, 'upload')
+    expect(uploads.connectors).toEqual({
+      workspace: [],
+      admin: [],
+      members: [],
+      liveProofRequired: [],
+    })
+    expect(uploads.memberSources).toEqual([])
+    expect(uploads.uploads).toBe(true)
+  })
+
+  it('drops source-less rows from both predicates once uploads are out of scope', () => {
+    const rowSql = (restricted: typeof plan) =>
+      render(
+        projectionCandidateAccessCondition(
+          {
+            connectorId: sqlColumn('connector_id'),
+            acl: sqlColumn('acl'),
+            documentId: sqlColumn('document_id'),
+          },
+          reader,
+          restricted
+        )
+      ).sql
+    expect(rowSql(plan)).toContain('IS NULL OR')
+    expect(rowSql(restrictSearchAccessPlan(plan, 'slack'))).not.toContain(
+      '"row"."connector_id" IS NULL'
+    )
+    const documentSql = (restricted: typeof plan) =>
+      render(knowledgeCandidateAccessConditionForConnectors(reader, restricted)).sql
+    expect(documentSql(plan)).toContain('"document"."connector_id" IS NULL OR')
+    expect(documentSql(restrictSearchAccessPlan(plan, 'slack'))).not.toContain(
+      '"document"."connector_id" IS NULL'
+    )
+  })
+})

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

@@ -232,6 +232,35 @@ export interface SearchAccessPlan {
   observers: KnowledgeMemberObservers
   /** Connectors the caller is an active member of, whose documents they read broadly. */
   memberSources: readonly string[]
+  /** Each eligible connector's type, so a search may be confined to one kind of source. */
+  connectorTypes: ReadonlyMap<string, string>
+  /** Whether documents without a source — uploads — are in scope. */
+  uploads: boolean
+}
+
+/**
+ * The plan confined to one kind of source: the connectors of that type keep their eligibility and
+ * the rest lose it, so every predicate built from the plan — on the row and on the document — and
+ * every source the legs walk or rank are that kind alone. `upload` keeps only source-less documents.
+ */
+export function restrictSearchAccessPlan(plan: SearchAccessPlan, source: string): SearchAccessPlan {
+  const keep = (id: string) => source !== 'upload' && plan.connectorTypes.get(id) === source
+  const kept = (ids: readonly string[]) => ids.filter(keep)
+  return {
+    connectors: {
+      workspace: kept(plan.connectors.workspace),
+      admin: kept(plan.connectors.admin),
+      members: kept(plan.connectors.members),
+      liveProofRequired: kept(plan.connectors.liveProofRequired),
+    },
+    observers: {
+      confirmed: plan.observers.confirmed.filter((observer) => keep(observer.connectorId)),
+      observed: plan.observers.observed.filter((observer) => keep(observer.connectorId)),
+    },
+    memberSources: kept(plan.memberSources),
+    connectorTypes: plan.connectorTypes,
+    uploads: source === 'upload',
+  }
 }
 
 export interface KnowledgeConnectorEligibility {
@@ -290,12 +319,14 @@ export function knowledgeCandidateAccessConditionForConnectors(
       ))
     )`
   }
+  const workspaceOwned = plan.uploads
+    ? sql`(${document.connectorId} IS NULL OR ${inConnectors(eligibility.workspace)})`
+    : inConnectors(eligibility.workspace)
   return sql`(
     ${aclOverlap(tokens)}
     AND ${aclRequirementsSatisfied(tokens)}
     AND (
-      ((${document.connectorId} IS NULL OR ${inConnectors(eligibility.workspace)})
-        AND ${document.acl} = ARRAY['ws']::text[])
+      (${workspaceOwned} AND ${document.acl} = ARRAY['ws']::text[])
       OR ${mirrored(eligibility.admin, sql`${document.aclVerifiedAt} > ${cutoff}`)}
       OR ${mirrored(eligibility.members, resolvedObservationCondition(plan.observers, cutoff))}
     )
@@ -340,13 +371,15 @@ export function projectionCandidateAccessCondition(
     ...plan.connectors.admin,
     ...plan.connectors.members,
   ]
+  const owned = plan.uploads
+    ? sql`(${projection.connectorId} IS NULL OR ${inSources(mirrored)})`
+    : inSources(mirrored)
   const unfilled = sql`(${projection.acl} IS NULL AND EXISTS (
     SELECT 1 FROM ${document}
     WHERE ${document.id} = ${projection.documentId}
       AND ${knowledgeCandidateAccessConditionForConnectors(scope, plan)}
   ))`
-  return sql`(${unfilled} OR (${projection.acl} && ${tokens}
-    AND (${projection.connectorId} IS NULL OR ${inSources(mirrored)})))`
+  return sql`(${unfilled} OR (${projection.acl} && ${tokens} AND ${owned}))`
 }
 
 /**