diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index c0041c0f1ce..9901b6889f9 100644 --- a/.github/workflows/test-build.yml +++ b/.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 diff --git a/apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts b/apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts new file mode 100644 index 00000000000..c5474d349a2 --- /dev/null +++ b/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(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(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[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([]) + }) +}) diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts index 0b73882e0db..223a7a5507f 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts @@ -33,6 +33,11 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => let sql: Sql const schemaName = `acl_write_${generateId().replaceAll('-', '')}` + const fannedOut = async () => + ( + await sql<{ document_id: string }[]>`SELECT document_id FROM fan_out ORDER BY document_id` + ).map((row) => row.document_id) + const projected = () => sql<{ id: string; acl: string[] | null }[]>` SELECT id, acl FROM embedding_search @@ -68,6 +73,15 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => )` } await installProjectionSourceAcl(sql) + /** + * Counts the documents whose write fired the projection fan-out: the trigger fires on any + * assignment of `acl`, changed or not, so this is the cost an unchanged write must not pay. + */ + await sql`CREATE TABLE fan_out (document_id text NOT NULL)` + await sql.unsafe(`CREATE FUNCTION count_fan_out() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN INSERT INTO fan_out VALUES (NEW.id); RETURN NEW; END; $$`) + await sql`CREATE TRIGGER count_fan_out AFTER UPDATE OF connector_id, acl ON document + FOR EACH ROW EXECUTE FUNCTION count_fan_out()` await sql`ALTER TABLE embedding_search DISABLE TRIGGER embedding_search_source_acl_set` await sql`ALTER TABLE embedding_keyword_tin DISABLE TRIGGER embedding_keyword_tin_source_acl_set` }, 60_000) @@ -79,7 +93,7 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => }) beforeEach(async () => { - await sql`TRUNCATE embedding_search, embedding_keyword_tin, document` + await sql`TRUNCATE embedding_search, embedding_keyword_tin, document, fan_out` await sql`INSERT INTO document (id, external_id, connector_id, acl, acl_verified_at) VALUES ('doc-same', 'file-same', 'admin', ARRAY[${ALICE}], now() - interval '1 day'), ('doc-moved', 'file-moved', 'admin', ARRAY[${ALICE}], now() - interval '1 day')` @@ -93,7 +107,7 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => } }) - it('refreshes the evidence of an unchanged ACL without rewriting any chunk projection row', async () => { + it('refreshes the evidence of an unchanged ACL without firing the projection fan-out', async () => { await expect(persist(new Map([['file-same', [ALICE]]]))).resolves.toEqual({ updated: 1, rejected: 0, @@ -103,6 +117,7 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => SELECT acl, acl_verified_at > now() AT TIME ZONE 'UTC' - interval '1 minute' AS fresh FROM document WHERE id = 'doc-same'` expect(stored).toEqual({ acl: [ALICE], fresh: true }) + expect(await fannedOut()).toEqual([]) expect((await projected()).filter((row) => row.id.includes('-same-'))).toEqual([ { id: 'kw-same-filled', acl: [ALICE] }, { id: 'kw-same-unfilled', acl: null }, @@ -111,7 +126,8 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => ]) }) - it('propagates a changed ACL to every chunk projection row, filled or not', async () => { + /** A chunk the backfill has not filled keeps a NULL ACL; the backfill copies the current one. */ + it('propagates a changed ACL to every filled chunk projection row', async () => { await expect(persist(new Map([['file-moved', [BOB]]]))).resolves.toEqual({ updated: 1, rejected: 0, @@ -121,9 +137,9 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => expect(stored.acl).toEqual([BOB]) expect((await projected()).filter((row) => row.id.includes('-moved-'))).toEqual([ { id: 'kw-moved-filled', acl: [BOB] }, - { id: 'kw-moved-unfilled', acl: [BOB] }, + { id: 'kw-moved-unfilled', acl: null }, { id: 'vec-moved-filled', acl: [BOB] }, - { id: 'vec-moved-unfilled', acl: [BOB] }, + { id: 'vec-moved-unfilled', acl: null }, ]) }) @@ -156,9 +172,10 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => { id: 'doc-moved', acl: [ALICE, BOB] }, { id: 'doc-same', acl: [ALICE] }, ]) + expect(await fannedOut()).toEqual(['doc-moved']) expect( - (await projected()).filter((row) => row.id.endsWith('-unfilled')).map((row) => row.acl) - ).toEqual([[ALICE, BOB], null, [ALICE, BOB], null]) + (await projected()).filter((row) => row.id.endsWith('-filled')).map((row) => row.acl) + ).toEqual([[ALICE, BOB], [ALICE], [ALICE, BOB], [ALICE]]) }) it('writes a changed ACL group larger than one change batch completely', async () => { const ids = Array.from( diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 129f0082ac2..a9e1cba7c41 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -2174,8 +2174,14 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise dateFilterCondition(params.filters) ? sql`EXISTS (SELECT 1 FROM ${document} WHERE ${and(sql`${document.id} = ranked_tin_chunks.document_id`, dateFilterCondition(params.filters))})` : undefined, + /** + * A filled row's mirrored source decides it; a row the fill has not reached carries no + * source, so only its document is asked, as the vector leg does. + */ excludedSources.length - ? sql`(ranked_tin_chunks.connector_id IS NULL OR NOT (ranked_tin_chunks.connector_id = ANY(${textArrayLiteral([...excludedSources])})))` + ? tinFilled + ? sql`(ranked_tin_chunks.connector_id IS NULL OR NOT (ranked_tin_chunks.connector_id = ANY(${textArrayLiteral([...excludedSources])})))` + : sql`((ranked_tin_chunks.acl IS NOT NULL AND (ranked_tin_chunks.connector_id IS NULL OR NOT (ranked_tin_chunks.connector_id = ANY(${textArrayLiteral([...excludedSources])})))) OR (ranked_tin_chunks.acl IS NULL AND NOT EXISTS (SELECT 1 FROM ${document} WHERE ${document.id} = ranked_tin_chunks.document_id AND ${document.connectorId} = ANY(${textArrayLiteral([...excludedSources])}))))` : undefined ) const documentConditions = (excludedSources: readonly string[]) => @@ -2188,6 +2194,23 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise ), excludeSearchSources(excludedSources) ) + /** + * A page read on the row takes each candidate's source from the row, which is what decides + * whether its live source proof is asked for. A row the fill has not reached carries no + * source, so until the fill is complete the page's own unfilled rows take it from their + * document: one primary-key read per row of the page, after its limit, never per ranked row. + */ + const onRowPage = (ranked: SQL) => + tinFilled + ? ranked + : sql` + SELECT paged.id, paged."documentId", + CASE WHEN paged.unfilled + THEN (SELECT ${document.connectorId} FROM ${document} WHERE ${document.id} = paged."documentId") + ELSE paged."connectorId" + END AS "connectorId", + paged.keyword_rank + FROM (${ranked}) AS paged` /** * One page from the top of Tin's ranking. The window of ranked chunks widens while too few of * them are readable to fill the page; if the widest window still cannot, the page is left to @@ -2239,13 +2262,17 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise * so a window of mostly unreadable chunks costs an array test per row, not a * document lookup. The full predicate follows at hydration. */ - sql` + onRowPage( + sql` SELECT ranked_tin_chunks.id, ranked_tin_chunks.document_id AS "documentId", - ranked_tin_chunks.connector_id AS "connectorId", ranked_tin_chunks.keyword_rank + ranked_tin_chunks.connector_id AS "connectorId", ranked_tin_chunks.keyword_rank${ + tinFilled ? sql`` : sql`, ranked_tin_chunks.acl IS NULL AS unfilled` + } FROM ranked_tin_chunks /* on-row visibility */ WHERE ranked_tin_chunks.enabled AND ${onRowKeywordVisibility(excludedSources)} ORDER BY ranked_tin_chunks.keyword_rank DESC, ranked_tin_chunks.id LIMIT ${pageLimit} OFFSET ${offset}` + ) : sql` SELECT ranked_tin_chunks.id, ${document.id} AS "documentId", ${document.connectorId} AS "connectorId",