fix(knowledge): read an unfilled keyword candidate's source from its document (#8176)

* fix(knowledge): read an unfilled keyword candidate's source from its document

The Tin keyword page took each candidate's source from its projection row. A
row the source and ACL fill has not reached carries a NULL source, so a chunk
from a source that needs a live reader proof was never recognized as one: the
proof was not resolved, hydration ran without the caller's grants, and the
document was hidden from a member who can read it.

Until the keyword projection is filled, a page's unfilled rows now take their
source from the document, one primary-key read per row of the page after its
limit. The excluded-sources filter asks the document the same way, as the
vector leg already does, so a denied source's unfilled rows stop taking slots
on a rebuilt page. The filled path's statement is unchanged.

* test(knowledge): assert the ACL write's fan-out, not unfilled rows the trigger now skips

The projection trigger no longer writes an ACL onto chunks the backfill has not
filled, so the ACL write tests can no longer see an unchanged write through
those rows. They now count the document writes that fire the fan-out trigger,
and expect a changed ACL on filled chunks only.

* fix(knowledge): read the document only for unfilled rows in the keyword source exclusion

Author: waleedlatif1

.github/workflows/test-build.yml

@@ -276,6 +276,7 @@ jobs:
           lib/knowledge/__integration__/member-scope-renewal.integration.ts
           lib/knowledge/__integration__/slack-empty-threads.integration.ts
           lib/knowledge/__integration__/kb-block-search.integration.ts
+          lib/knowledge/__integration__/unfilled-projection-source.integration.ts
           lib/core/outbox/service.integration.ts
           lib/knowledge/__integration__/connector-upload.integration.ts
           lib/uploads/contexts/organization-logo/application.integration.ts

apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts

@@ -0,0 +1,380 @@
+/**
+ * A search candidate's source decides whether the caller's live source proof is resolved before
+ * its content is read. These fixtures put a GitHub installation source's chunks on projection rows
+ * the source and ACL fill has not reached (`acl` and `connector_id` NULL), and check that such a
+ * chunk still reaches a member who holds the installation grant, stays hidden from one who does
+ * not, and is left out of a page once its source is known to be denied.
+ */
+import { createHash } from 'node:crypto'
+import { db } from '@sim/db'
+import {
+  credential,
+  credentialGroup,
+  credentialGroupEnrollment,
+  document,
+  embedding,
+  embeddingKeywordTin,
+  embeddingSearch,
+  knowledgeConnector,
+  knowledgeConnectorMember,
+  knowledgeDocumentObservation,
+  organization,
+  user,
+  workspace,
+} from '@sim/db/schema'
+import { generateId } from '@sim/utils/id'
+import { isRecordLike } from '@sim/utils/object'
+import { eq, inArray, sql } from 'drizzle-orm'
+import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
+
+vi.mock('@/lib/knowledge/search/tin-keyword', () => ({
+  resolveTinKeywordQuery: async () => 'fixture',
+}))
+
+import {
+  createKnowledgeAclFixtureIds,
+  seedKnowledgeAclFixture,
+} from '@/lib/knowledge/__integration__/seed-source-access-fixture'
+import type { SearchAccessPlan } from '@/lib/knowledge/access/predicate'
+import type {
+  GitHubInstallationReadGrant,
+  KnowledgeAccessProvider,
+  UserAccessScope,
+} from '@/lib/knowledge/access/types'
+import {
+  executeKeywordSearch,
+  forgetProjectionFilled,
+  handleVectorOnlySearch,
+  liveSourceAccessFor,
+} from '@/lib/knowledge/search/queries'
+import { GITHUB_INSTALLATION_PROVIDER_ID } from '@/lib/oauth/github-installation-types'
+
+const ids = createKnowledgeAclFixtureIds()
+const connectorId = generateId()
+const contentCredentialId = generateId()
+const groupId = generateId()
+const optionId = generateId()
+const documentId = generateId()
+const embeddingId = generateId()
+const repositoryId = '4242'
+const members = {
+  alice: { id: generateId(), subject: 'alice-gh', credentialId: generateId() },
+  bob: { id: generateId(), subject: 'bob-gh', credentialId: generateId() },
+}
+const userIdOf = (who: 'alice' | 'bob') => (who === 'alice' ? ids.aliceId : ids.bobId)
+const subjectToken = (subject: string) => `s:github-repositories:-:${subject}`
+const queryVector = {
+  vector: JSON.stringify([1, ...Array<number>(1535).fill(0)]),
+  dimensions: 1536 as const,
+  model: 'text-embedding-3-small',
+}
+
+/** The shims stand in for the Tin extension, which the test database does not carry. */
+let createdTinShims = false
+
+const scopeFor = (who: 'alice' | 'bob'): UserAccessScope => ({
+  kind: 'user',
+  userId: userIdOf(who),
+  tokens: [
+    'pub',
+    subjectToken(members[who].subject),
+    `u:${userIdOf(who)}@fixture.test`,
+    'ws',
+  ].sort(),
+})
+
+const planFor = (who: 'alice' | 'bob'): SearchAccessPlan => ({
+  connectors: {
+    workspace: [],
+    admin: [],
+    members: [connectorId],
+    liveProofRequired: [connectorId],
+  },
+  observers: { confirmed: [{ id: members[who].id, connectorId }], observed: [] },
+  memberSources: [connectorId],
+  connectorTypes: new Map([[connectorId, 'github']]),
+  uploads: false,
+})
+
+/** Alice's reader credential backs a real installation grant; Bob holds none. */
+const aliceGrant: GitHubInstallationReadGrant = {
+  connectorId,
+  contentCredentialId,
+  readerCredentialId: members.alice.credentialId,
+  readerSubjectToken: subjectToken(members.alice.subject),
+  repositoryId,
+}
+
+function searchInputs(who: 'alice' | 'bob') {
+  const access = scopeFor(who)
+  const accessPlan = planFor(who)
+  const granted = who === 'alice' ? { ...access, githubInstallationGrants: [aliceGrant] } : access
+  const accessProvider: KnowledgeAccessProvider = {
+    get: async () => access,
+    getForConnectors: async () => granted,
+    getForDocuments: async () => granted,
+    liveSourceConnectorCondition: async () => null,
+  }
+  return {
+    knowledgeBaseIds: [ids.knowledgeBaseId],
+    topK: 5,
+    access,
+    accessProvider,
+    accessPlan,
+    liveSourceAccess: liveSourceAccessFor(access, accessPlan, accessProvider),
+    queryVector,
+  }
+}
+
+const keywordIds = async (who: 'alice' | 'bob') =>
+  (
+    await executeKeywordSearch({
+      ...searchInputs(who),
+      query: 'fixture',
+      permitted: { kind: 'unbounded', broad: false },
+      searchIndexOnly: true,
+    })
+  ).map((row) => row.id)
+
+const vectorIds = async (who: 'alice' | 'bob') =>
+  (
+    await handleVectorOnlySearch({
+      ...searchInputs(who),
+      distanceThreshold: 2,
+      permitted: { kind: 'unbounded', broad: true },
+    })
+  ).map((row) => row.id)
+
+async function setProjection(state: 'filled' | 'unfilled') {
+  for (const table of [embeddingSearch, embeddingKeywordTin]) {
+    await db
+      .update(table)
+      .set(
+        state === 'filled'
+          ? {
+              connectorId,
+              acl: [subjectToken(members.alice.subject), subjectToken(members.bob.subject)].sort(),
+            }
+          : { connectorId: null, acl: null }
+      )
+      .where(eq(table.id, embeddingId))
+  }
+  forgetProjectionFilled()
+}
+
+beforeAll(async () => {
+  await seedKnowledgeAclFixture(ids)
+  const now = new Date()
+  await db
+    .update(user)
+    .set({ emailVerified: true })
+    .where(inArray(user.id, [ids.aliceId, ids.bobId]))
+  await db.insert(credential).values({
+    id: contentCredentialId,
+    workspaceId: ids.workspaceId,
+    type: 'service_account',
+    displayName: 'Fixture GitHub installation',
+    createdBy: ids.aliceId,
+    providerId: GITHUB_INSTALLATION_PROVIDER_ID,
+  })
+  await db.insert(credentialGroup).values({
+    id: groupId,
+    workspaceId: ids.workspaceId,
+    publicId: generateId(),
+    name: 'GitHub readers',
+    options: [
+      {
+        id: optionId,
+        provider: 'github-repositories',
+        label: 'GitHub fixture',
+        authorizationAppId: 'fixture-app',
+        requiredScopes: ['repo'],
+        scopeVersion: 1,
+        required: false,
+        status: 'active',
+      },
+    ],
+  } as typeof credentialGroup.$inferInsert)
+  for (const who of ['alice', 'bob'] as const) {
+    const [enrollment] = await db
+      .insert(credentialGroupEnrollment)
+      .values({
+        id: generateId(),
+        credentialGroupId: groupId,
+        userId: userIdOf(who),
+        email: `${userIdOf(who)}@fixture.test`,
+        status: 'completed',
+        invitationTokenHash: createHash('sha256').update(generateId()).digest('hex'),
+        invitationExpiresAt: new Date(Date.now() + 60 * 60 * 1000),
+        invitedAt: now,
+      })
+      .returning({ id: credentialGroupEnrollment.id })
+    await db.insert(credential).values({
+      id: members[who].credentialId,
+      workspaceId: ids.workspaceId,
+      type: 'managed_oauth',
+      displayName: 'Fixture GitHub reader',
+      providerId: 'github-repositories',
+      authorizationAppId: 'fixture-app',
+      credentialGroupEnrollmentId: enrollment!.id,
+      credentialGroupOptionId: optionId,
+      managedOauthScopeVersion: 1,
+      providerSubjectId: members[who].subject,
+      providerTenantId: '',
+      managedOauthStatus: 'active',
+      grantedScopes: ['repo'],
+      encryptedOauthTokenSet: 'fixture-not-an-oauth-token',
+      grantedAt: now,
+      createdBy: userIdOf(who),
+    })
+  }
+  await db.insert(knowledgeConnector).values({
+    id: connectorId,
+    knowledgeBaseId: ids.knowledgeBaseId,
+    connectorType: 'github',
+    sourceConfig: { githubRepositoryId: repositoryId },
+    accessMode: 'members',
+    status: 'active',
+    credentialId: contentCredentialId,
+    credentialGroupId: groupId,
+    credentialGroupOptionId: optionId,
+  })
+  await db.insert(knowledgeConnectorMember).values(
+    (['alice', 'bob'] as const).map((who) => ({
+      id: members[who].id,
+      workspaceId: ids.workspaceId,
+      connectorId,
+      credentialId: members[who].credentialId,
+      subjectToken: subjectToken(members[who].subject),
+      status: 'active',
+      memberSyncedThrough: now,
+    }))
+  )
+  await db.insert(document).values({
+    id: documentId,
+    connectorId,
+    knowledgeBaseId: ids.knowledgeBaseId,
+    externalId: 'fixture-file',
+    filename: 'readme.md',
+    fileUrl: 'https://fixture.test/readme',
+    fileSize: 12,
+    mimeType: 'text/plain',
+    processingStatus: 'completed',
+    acl: [subjectToken(members.alice.subject), subjectToken(members.bob.subject)].sort(),
+  })
+  await db.insert(knowledgeDocumentObservation).values(
+    (['alice', 'bob'] as const).map((who) => ({
+      documentId,
+      memberId: members[who].id,
+      lastSeenAt: now,
+      runId: generateId(),
+    }))
+  )
+  await db.insert(embedding).values({
+    id: embeddingId,
+    documentId,
+    knowledgeBaseId: ids.knowledgeBaseId,
+    chunkIndex: 0,
+    chunkHash: 'fixture-hash',
+    content: 'fixture readme',
+    contentLength: 14,
+    tokenCount: 2,
+    startOffset: 0,
+    endOffset: 14,
+    embeddingModel: 'text-embedding-3-small',
+    embedding: [1, ...Array<number>(1535).fill(0)],
+  })
+  await db.insert(embeddingKeywordTin).values({
+    id: embeddingId,
+    knowledgeBaseId: ids.knowledgeBaseId,
+    documentId,
+    enabled: true,
+    content: 'fixture readme',
+  })
+  const [tin] = await db.execute<{ present: boolean }>(
+    sql`SELECT to_regnamespace('tin') IS NOT NULL AS present`
+  )
+  if (!tin?.present) {
+    createdTinShims = true
+    await db.execute(
+      sql.raw(`CREATE SCHEMA tin;
+      CREATE FUNCTION tin.full_score(tid) RETURNS double precision LANGUAGE sql IMMUTABLE AS 'SELECT 1.0::float8';
+      CREATE FUNCTION knowledge_tin_base_token(text) RETURNS text LANGUAGE sql IMMUTABLE AS $$SELECT 'kb'$$;
+      CREATE FUNCTION tin_fixture_match(text, text) RETURNS boolean LANGUAGE sql IMMUTABLE AS 'SELECT true';
+      CREATE OPERATOR ==> (LEFTARG = text, RIGHTARG = text, FUNCTION = tin_fixture_match);`)
+    )
+  }
+})
+
+afterAll(async () => {
+  if (createdTinShims) {
+    await db.execute(
+      sql.raw(`DROP OPERATOR IF EXISTS ==> (text, text);
+      DROP FUNCTION IF EXISTS tin_fixture_match(text, text);
+      DROP FUNCTION IF EXISTS knowledge_tin_base_token(text);
+      DROP SCHEMA IF EXISTS tin CASCADE;`)
+    )
+  }
+  await db.delete(embeddingKeywordTin).where(eq(embeddingKeywordTin.id, embeddingId))
+  await db.delete(workspace).where(eq(workspace.id, ids.workspaceId))
+  await db.delete(credentialGroup).where(eq(credentialGroup.id, groupId))
+  await db.delete(organization).where(eq(organization.id, ids.organizationId))
+  await db.delete(user).where(inArray(user.id, [ids.aliceId, ids.bobId]))
+  forgetProjectionFilled()
+})
+
+describe('a chunk whose projection row the fill has not reached', () => {
+  beforeEach(() => setProjection('unfilled'))
+
+  it('reaches the member holding the installation grant through the keyword ranking', async () => {
+    expect(await keywordIds('alice')).toEqual([embeddingId])
+  })
+
+  it('reaches the member holding the installation grant through the vector ranking', async () => {
+    expect(await vectorIds('alice')).toEqual([embeddingId])
+  })
+
+  it('stays hidden from a member without the grant', async () => {
+    expect(await keywordIds('bob')).toEqual([])
+    expect(await vectorIds('bob')).toEqual([])
+  })
+
+  describe('once its source is known to be denied', () => {
+    /** Each Tin ranking statement's page of candidates, in the order the search read them. */
+    const pages: Array<{ candidates: unknown[] }> = []
+    beforeEach(() => {
+      pages.length = 0
+      const execute = db.execute.bind(db)
+      vi.spyOn(db, 'execute').mockImplementation((async (query: Parameters<typeof execute>[0]) => {
+        const rows = await execute(query)
+        const [row] = Array.from(rows)
+        if (isRecordLike(row) && 'ranked' in row && Array.isArray(row.candidates))
+          pages.push({ candidates: row.candidates })
+        return rows
+      }) as typeof db.execute)
+    })
+    afterEach(() => vi.restoreAllMocks())
+
+    it('carries the source read from its document and is left out of the rebuilt keyword page', async () => {
+      expect(await keywordIds('bob')).toEqual([])
+      expect(pages.length).toBeGreaterThanOrEqual(2)
+      expect(pages[0]!.candidates).toEqual([{ id: embeddingId, documentId, connectorId }])
+      expect(pages.at(-1)!.candidates).toEqual([])
+    })
+  })
+})
+
+describe('a chunk whose projection row is filled', () => {
+  beforeEach(() => setProjection('filled'))
+
+  it('ranks on the row as before and is read only by the member holding the grant', async () => {
+    const [{ unfilled }] = await db.execute<{ unfilled: boolean }>(
+      sql`SELECT EXISTS (SELECT 1 FROM ${embeddingKeywordTin} WHERE ${embeddingKeywordTin.acl} IS NULL) AS unfilled`
+    )
+    expect(unfilled).toBe(false)
+    expect(await keywordIds('alice')).toEqual([embeddingId])
+    expect(await keywordIds('bob')).toEqual([])
+    expect(await vectorIds('alice')).toEqual([embeddingId])
+    expect(await vectorIds('bob')).toEqual([])
+  })
+})