From 08d3b4db0b7f97bd1eae0ed892602b013e456c95 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 13 Oct 2026 11:29:34 -0700 Subject: [PATCH 1/2] improvement(knowledge): finish Confluence passes past unreadable attachments, rank unauthenticated keyword search in one statement - Confluence: a 403/404 on one parent's attachment listing is skipped with a warning instead of recorded as a listing failure, so the pass can finish and advance lastSyncAt. The page still marks the listing reconciliation-unsafe, so attachments it did not list are kept, never deleted. Missing scopes and server failures stay listing failures. - Connector counts: the member engine keeps the recorded document count on an unfinished slice once a pass has finished, the same gate the content engine uses. The count query and gate now live in sync-primitives. - Keyword search: callers with no person behind them rank through the same materialized candidate statement signed-in readers use, instead of ORDER BY ts_rank_cd over a per-chunk join. Same ACL predicate and global ordering, ties now broken by chunk id; still unbudgeted. --- .../connectors/confluence/attachments.test.ts | 84 +++++++++++++------ apps/sim/connectors/confluence/attachments.ts | 36 ++++++-- ...ber-sync-engine-content-credential.test.ts | 45 +++++++++- .../connectors/member-sync-engine.ts | 18 ++-- .../lib/knowledge/connectors/sync-engine.ts | 50 +++-------- .../knowledge/connectors/sync-primitives.ts | 32 +++++++ apps/sim/lib/knowledge/search/queries.ts | 61 +++----------- 7 files changed, 192 insertions(+), 134 deletions(-) diff --git a/apps/sim/connectors/confluence/attachments.test.ts b/apps/sim/connectors/confluence/attachments.test.ts index bdca87f1567..35353e909d7 100644 --- a/apps/sim/connectors/confluence/attachments.test.ts +++ b/apps/sim/connectors/confluence/attachments.test.ts @@ -215,41 +215,73 @@ describe('Confluence attachment listing', () => { expect(last.hasMore).toBe(false) }) - it.each([403, 401])( - 'preserves parent pages and cumulative partial progress for attachment access failure %s', + it('preserves parent pages and cumulative partial progress when the attachment scope is missing', async () => { + fetchMock.mockImplementation(async () => + Response.json({ code: 401, message: 'Unauthorized; scope does not match' }, { status: 401 }) + ) + const listParents = vi.fn( + async (cursor: string | undefined): Promise => ({ + documents: [parent(cursor ? 'p2' : 'p1')], + hasMore: !cursor, + nextCursor: cursor ? undefined : 'next', + }) + ) + const first = await listConfluenceAttachments({ ...INPUT, listParents }) + const context: Record = {} + const last = await listConfluenceAttachments({ + ...INPUT, + listParents, + cursor: first.nextCursor, + syncContext: context, + }) + expect(first.documents.map((doc) => doc.externalId)).toEqual(['p1']) + expect(last.documents.map((doc) => doc.externalId)).toEqual(['p2']) + expect(last.listingFailures).toEqual({ + count: 2, + samples: ['p1', 'p2'].map((scope) => ({ + scope, + operation: 'confluence.attachments.list', + status: 401, + reasons: ['attachment_scope_mismatch'], + })), + }) + expect(last.reconciliationSafe).toBe(false) + expect(context.reconciliationUnsafe).toBe(true) + }) + + /** + * One page whose attachments the credential cannot open must not leave every pass incomplete, + * which would keep the connector from ever finishing a sync. Its attachments were not listed, + * so the page still withholds deletion of anything it did not list. + */ + it.each([403, 404])( + 'skips a parent whose attachments answer %s without failing the listing or deleting what it did not list', async (status) => { - fetchMock.mockImplementation(async () => - status === 401 - ? Response.json({ code: 401, message: 'Unauthorized; scope does not match' }, { status }) - : new Response(null, { status }) - ) + fetchMock.mockImplementation(async (input) => { + const parentId = new URL(String(input)).pathname.split('/').at(-2) ?? '' + if (parentId === 'p1') return new Response(null, { status }) + return Response.json({ results: [file({ id: `${parentId}-file`, pageId: parentId })] }) + }) const listParents = vi.fn( - async (cursor: string | undefined): Promise => ({ - documents: [parent(cursor ? 'p2' : 'p1')], - hasMore: !cursor, - nextCursor: cursor ? undefined : 'next', + async (): Promise => ({ + documents: [parent('p1'), parent('p2')], + hasMore: false, }) ) - const first = await listConfluenceAttachments({ ...INPUT, listParents }) const context: Record = {} - const last = await listConfluenceAttachments({ + const result = await listConfluenceAttachments({ ...INPUT, listParents, - cursor: first.nextCursor, syncContext: context, }) - expect(first.documents.map((doc) => doc.externalId)).toEqual(['p1']) - expect(last.documents.map((doc) => doc.externalId)).toEqual(['p2']) - expect(last.listingFailures).toEqual({ - count: 2, - samples: ['p1', 'p2'].map((scope) => ({ - scope, - operation: 'confluence.attachments.list', - status, - reasons: [status === 401 ? 'attachment_scope_mismatch' : 'attachment_access_denied'], - })), - }) - expect(last.reconciliationSafe).toBe(false) + expect(result.documents.map((doc) => doc.externalId)).toEqual([ + 'p1', + 'p2', + 'attachment:page:p2:p2-file', + ]) + expect(result.hasMore).toBe(false) + expect(result.listingFailures).toBeUndefined() + expect(result.reconciliationSafe).toBe(false) expect(context.reconciliationUnsafe).toBe(true) } ) diff --git a/apps/sim/connectors/confluence/attachments.ts b/apps/sim/connectors/confluence/attachments.ts index 16e17f66dc6..4c873ed08d1 100644 --- a/apps/sim/connectors/confluence/attachments.ts +++ b/apps/sim/connectors/confluence/attachments.ts @@ -378,6 +378,7 @@ export async function listConfluenceAttachments( ): Promise { const state = decodeCursor(input.cursor, input.syncContext?.totalDocsFetched) const documents: ExternalDocument[] = [] + let skippedUnreadableParent = false if (state.parents.length === 0 && !state.parentsDone) { const parentContext = { ...input.syncContext, totalDocsFetched: state.parentsListed } const page = await input.listParents(state.parentCursor, parentContext) @@ -413,15 +414,30 @@ export async function listConfluenceAttachments( }) } if (state.serverFailureStreak >= MAX_CONSECUTIVE_SERVER_FAILURES) throw listing.error - state.failures ??= { count: 0, samples: [] } - state.failures.count += 1 - if (state.failures.samples.length < MAX_LISTING_FAILURE_SAMPLES) { - state.failures.samples.push({ - scope: parent.id, - operation: 'confluence.attachments.list', + if (listing.reason === 'attachment_access_denied') { + /** + * The parent was listed a moment ago, so a 403 or 404 on its attachments is that one + * item's restriction, not a scope the source failed to list: skipping it leaves the pass + * complete. Its unlisted attachments are not evidence of deletion, so the page still + * marks the listing unsafe and reconciliation keeps them. + */ + logger.warn('Skipping a Confluence parent whose attachments are not readable', { + parentId: parent.id, + parentType: parent.type, status: listing.status, - reasons: [listing.reason], }) + skippedUnreadableParent = true + } else { + state.failures ??= { count: 0, samples: [] } + state.failures.count += 1 + if (state.failures.samples.length < MAX_LISTING_FAILURE_SAMPLES) { + state.failures.samples.push({ + scope: parent.id, + operation: 'confluence.attachments.list', + status: listing.status, + reasons: [listing.reason], + }) + } } state.parents.shift() state.attachmentCursor = undefined @@ -453,7 +469,8 @@ export async function listConfluenceAttachments( state.attachmentCursor = nextCursor if (!nextCursor) state.parents.shift() } - if (state.failures && input.syncContext) input.syncContext.reconciliationUnsafe = true + const reconciliationUnsafe = Boolean(state.failures) || skippedUnreadableParent + if (reconciliationUnsafe && input.syncContext) input.syncContext.reconciliationUnsafe = true const hasMore = state.parents.length > 0 || !state.parentsDone const nextCursor = hasMore ? CURSOR_PREFIX + JSON.stringify(state) : undefined if (nextCursor && Buffer.byteLength(nextCursor, 'utf8') > MAX_CURSOR_BYTES) { @@ -463,7 +480,8 @@ export async function listConfluenceAttachments( documents, hasMore, nextCursor, - ...(state.failures ? { listingFailures: state.failures, reconciliationSafe: false } : {}), + ...(state.failures ? { listingFailures: state.failures } : {}), + ...(reconciliationUnsafe ? { reconciliationSafe: false } : {}), } } diff --git a/apps/sim/lib/knowledge/connectors/member-sync-engine-content-credential.test.ts b/apps/sim/lib/knowledge/connectors/member-sync-engine-content-credential.test.ts index 9984b583bbb..1821476c930 100644 --- a/apps/sim/lib/knowledge/connectors/member-sync-engine-content-credential.test.ts +++ b/apps/sim/lib/knowledge/connectors/member-sync-engine-content-credential.test.ts @@ -185,6 +185,8 @@ function arrange( existingDocument?: boolean openMemberFeed?: boolean contentFresh?: boolean + /** A content pass that finished before this run, past the content interval. */ + finishedPass?: { lastSyncAt: Date; lastSyncDocCount: number } forceContentRefresh?: boolean noDueMembers?: boolean syncIntervalMinutes?: number @@ -211,7 +213,8 @@ function arrange( memberSyncStatus: 'idle', memberSyncConsecutiveFailures: 0, syncIntervalMinutes: options.syncIntervalMinutes ?? 1440, - lastSyncAt: options.contentFresh ? new Date() : null, + lastSyncAt: options.finishedPass?.lastSyncAt ?? (options.contentFresh ? new Date() : null), + lastSyncDocCount: options.finishedPass?.lastSyncDocCount ?? null, directoryCheckpoint: options.directoryCheckpoint ?? null, sourceConfig: { folderId: 'shared-folder', adminEmail: 'admin@example.com' }, accessRewritePending: false, @@ -280,7 +283,8 @@ function arrange( }, ]) } - queueTableRows(schemaMock.document, [{ count: 1 }]) + if (!options.finishedPass || !options.contentIncomplete) + queueTableRows(schemaMock.document, [{ count: 1 }]) } if (options.members && !options.noDueMembers) queueTableRows(schemaMock.document, [ @@ -553,6 +557,43 @@ describe('member engine with a dedicated content credential', () => { expect(bounded).toEqual([]) }) + /** + * A finished pass recorded the count, and an unfinished slice of the next pass keeps it rather + * than scanning the whole connector again; the count is replaced when that pass finishes. + */ + it('keeps the recorded document count on an unfinished slice once a pass has finished', async () => { + const run = arrange({ + contentIncomplete: true, + finishedPass: { lastSyncAt: new Date(Date.now() - 2 * 86_400_000), lastSyncDocCount: 7 }, + }) + const now = Date.now() + const clock = vi.spyOn(Date, 'now') + mocks.list.mockResolvedValue({ + documents: Array.from({ length: 26 }, (_, i) => ({ + ...serviceDocument, + externalId: `content-${i}`, + })), + hasMore: false, + }) + mocks.get.mockImplementation(async (_token, _config, externalId: string) => { + clock.mockReturnValue(now + 35 * 60_000) + return { ...serviceDocument, externalId, content: 'Service content', contentDeferred: false } + }) + try { + const result = await run() + expect(result.error).toBeUndefined() + expect(result.docsAdded).toBeGreaterThan(0) + expect(dbChainMockFns.values).toHaveBeenCalledWith( + expect.objectContaining({ status: 'partial', listedCount: null }) + ) + expect( + dbChainMockFns.set.mock.calls.some(([value]) => value && 'lastSyncDocCount' in value) + ).toBe(false) + } finally { + clock.mockRestore() + } + }) + it('reserves time for member permissions when a slow dedicated content page has more batches', async () => { const run = arrange({ members: true, contentIncomplete: true }) const now = Date.now() diff --git a/apps/sim/lib/knowledge/connectors/member-sync-engine.ts b/apps/sim/lib/knowledge/connectors/member-sync-engine.ts index ce6b23c5846..2ed416cb4bf 100644 --- a/apps/sim/lib/knowledge/connectors/member-sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/member-sync-engine.ts @@ -121,7 +121,9 @@ import { ConnectorSyncCapacityError, classifyListing, classifySuspectListing, + countLiveConnectorDocuments, createSyncRunState, + keepsRecordedDocCount, loadPageCorpus, processDocOps, RETRY_WINDOW_DAYS, @@ -1711,17 +1713,9 @@ async function syncDedicatedMemberContent(input: { run.result.listingIncomplete = true const contentNotice = pass.holdNotice /** Counted before the lease transaction: a scan of the whole connector never holds its row. */ - const [{ count }] = await db - .select({ count: sql`count(*)::int` }) - .from(document) - .where( - and( - eq(document.connectorId, run.connectorId), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt) - ) - ) + const count = keepsRecordedDocCount(pass.complete, connector) + ? null + : await countLiveConnectorDocuments(run.connectorId) await withMemberLease(run, async (tx) => { const now = new Date() await tx.insert(knowledgeConnectorSyncLog).values({ @@ -1746,7 +1740,7 @@ async function syncDedicatedMemberContent(input: { ? { lastSyncAt: new Date(pass.checkpoint.startedAt) } : {}), ...(pass.complete ? { listingCheckpoint: null } : {}), - lastSyncDocCount: count, + ...(count === null ? {} : { lastSyncDocCount: count }), lastSyncError: contentNotice, updatedAt: now, }) diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.ts b/apps/sim/lib/knowledge/connectors/sync-engine.ts index 87f7a61e317..e568485741b 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.ts @@ -89,6 +89,8 @@ import { ConnectorDeletedException, ConnectorSyncCapacityError, checkSyncTargetPresence, + countLiveConnectorDocuments, + keepsRecordedDocCount, RETRY_WINDOW_DAYS, shouldRunIncrementalSync, sweepStuckDocuments, @@ -368,22 +370,6 @@ export function isContentPassIncomplete( ) } -/** Live documents the connector owns: what the connector list shows as its document count. */ -async function countLiveConnectorDocuments(connectorId: string): Promise { - const [row] = await db - .select({ count: sql`count(*)::int` }) - .from(document) - .where( - and( - eq(document.connectorId, connectorId), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt) - ) - ) - return row?.count ?? 0 -} - /** * Atomically publishes the completed log and connector terminal state. * @@ -435,29 +421,19 @@ export async function completeSuccessfulSync( * owned count the next pass's mass-deletion guard starts from, so a tally of this run's * adds and deletes would drift; when it cannot be read the previous count stands, which the * guard already treats as a floor. - * - * Once a pass has finished, a slice that leaves the next pass unfinished keeps the count that - * finished pass recorded, the same pass `lastSyncAt` still names, so a large connector is not - * counted again on every slice of a pass whose count is replaced when it finishes. Unfinished - * slices never advance `lastSyncAt`, so until a pass finishes every slice is counted and a - * first pass shows its progress. */ - const keepRecordedCount = - contentPass !== undefined && - !contentPass.complete && - lockedConnector?.lastSyncAt != null && - lockedConnector.lastSyncDocCount !== null - const actualDocCount = keepRecordedCount - ? null - : await countLiveConnectorDocuments(connectorId).catch((error: unknown) => { - const diagnostic = getConnectorFailureDiagnostic(error) - if (diagnostic?.category !== 'database') throw error - logger.warn('Could not count connector documents; keeping the previous count', { - connectorId, - diagnostic, + const actualDocCount = + contentPass !== undefined && keepsRecordedDocCount(contentPass.complete, lockedConnector) + ? null + : await countLiveConnectorDocuments(connectorId).catch((error: unknown) => { + const diagnostic = getConnectorFailureDiagnostic(error) + if (diagnostic?.category !== 'database') throw error + logger.warn('Could not count connector documents; keeping the previous count', { + connectorId, + diagnostic, + }) + return null }) - return null - }) try { return await db.transaction(async (tx) => { const [lockedKnowledgeBase] = await tx diff --git a/apps/sim/lib/knowledge/connectors/sync-primitives.ts b/apps/sim/lib/knowledge/connectors/sync-primitives.ts index 958190fe37c..ead9db8a351 100644 --- a/apps/sim/lib/knowledge/connectors/sync-primitives.ts +++ b/apps/sim/lib/knowledge/connectors/sync-primitives.ts @@ -653,6 +653,38 @@ export function resolvePreviousOwnedCount( return Math.max(lastSyncDocCount ?? 0, ownedDocCount) } +/** Live documents the connector owns: what the connector list shows as its document count. */ +export async function countLiveConnectorDocuments(connectorId: string): Promise { + const [row] = await db + .select({ count: sql`count(*)::int` }) + .from(document) + .where( + and( + eq(document.connectorId, connectorId), + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt) + ) + ) + return row?.count ?? 0 +} + +/** + * Whether a slice keeps the document count the connector holds instead of counting again. + * + * Once a pass has finished, a slice that leaves the next pass unfinished keeps the count that + * finished pass recorded, the same pass `lastSyncAt` still names, so a large connector is not + * counted again on every slice of a pass whose count is replaced when it finishes. Unfinished + * slices never advance `lastSyncAt`, so until a pass finishes every slice is counted and a first + * pass shows its progress. + */ +export function keepsRecordedDocCount( + passComplete: boolean, + connector: { lastSyncAt: Date | null; lastSyncDocCount: number | null } | null +): boolean { + return !passComplete && connector?.lastSyncAt != null && connector.lastSyncDocCount != null +} + /** * Fraction of a connector's owned documents that a single reconciliation pass * may remove before the pass is held. diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 0e55d49ad21..a2b4a4d23b3 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -67,8 +67,9 @@ interface DocumentReadAccess { /** The candidate predicate every leg ranks under. */ rankCondition: SQL /** - * A signed-in reader: their tag and keyword legs run under the search deadline and rank in one - * statement. A caller with no person behind it runs those legs unbudgeted, as it always has. + * A signed-in reader: their tag and keyword legs run under the search deadline, and their tag leg + * ranks in one statement. A caller with no person behind it runs those legs unbudgeted, as it + * always has. */ signedIn: boolean /** Present only when a searched base holds a source whose reader must be proven live. */ @@ -385,11 +386,10 @@ export async function handleVectorSearch( * query makes Postgres detoast the chunk's vector and compute a distance for *every* full-text * match before the `LIMIT` applies — work that scales with how common the query term is rather * than with `topK`. Ranking therefore touches no vectors, and only the rows that survive the limit - * are hydrated. A signed-in reader ranks in one statement under the search deadline (see - * {@link rankSignedInKeywordCandidates}); a caller with no person behind it ranks unbudgeted, many - * bases in parallel, each up to the leg's whole `topK` so the merged ranking is the global one. A - * search holding a source whose reader must be proven live pages its ranking instead, so the - * proof is asked for only when one of its candidates is read. + * are hydrated. Every caller ranks every base in one statement (see + * {@link rankKeywordCandidates}): a signed-in reader under the search deadline, a caller with no + * person behind it unbudgeted. A search holding a source whose reader must be proven live pages + * its ranking instead, so the proof is asked for only when one of its candidates is read. */ async function executeKeywordSearch( params: KeywordSearchParams, @@ -418,7 +418,7 @@ async function executeKeywordSearch( 'keyword', budget ) - const signedInRanking = { + const ranking = { knowledgeBaseIds, tsQuery, tagFilterConditions, @@ -439,8 +439,8 @@ async function executeKeywordSearch( budget, topK, selectPage: async (limit, offset, excludedSources) => { - const candidates = await rankSignedInKeywordCandidates({ - ...signedInRanking, + const candidates = await rankKeywordCandidates({ + ...ranking, exclusion: excludeSearchSources(excludedSources), limit, offset, @@ -451,43 +451,8 @@ async function executeKeywordSearch( }) } - const rankExpr = sql`ts_rank_cd(${embedding.contentTsv}, ${tsQuery})` - const visibility = getVisibilityConditions(params.filters, read.rankCondition) - /** Ranking pass: ids and relevance only, so no vector is read. */ - const rankRows = (kbScope: SQL) => - db - .select({ id: embedding.id, keywordRank: rankExpr.as('keyword_rank') }) - .from(embedding) - .innerJoin(document, eq(embedding.documentId, document.id)) - .where( - and( - kbScope, - ...visibility, - sql`${embedding.contentTsv} @@ ${tsQuery}`, - ...tagFilterConditions - ) - ) - .orderBy(sql`${rankExpr} DESC`) - .limit(topK) - return shortOnDeadline(budget, async () => { - let topIds: string[] - if (read.signedIn) { - const ranked = await rankSignedInKeywordCandidates({ ...signedInRanking, limit: topK }) - topIds = ranked.map((row) => row.id) - } else if (getQueryStrategy(knowledgeBaseIds.length, topK).useParallel) { - const perBase = await Promise.all( - knowledgeBaseIds.map((kbId) => rankRows(eq(embedding.knowledgeBaseId, kbId))) - ) - topIds = perBase - .flat() - .sort((a, b) => b.keywordRank - a.keywordRank) - .slice(0, topK) - .map((row) => row.id) - } else { - const ranked = await rankRows(inArray(embedding.knowledgeBaseId, knowledgeBaseIds)) - topIds = ranked.map((row) => row.id) - } + const topIds = (await rankKeywordCandidates({ ...ranking, limit: topK })).map((row) => row.id) if (topIds.length === 0) { return [] } @@ -498,8 +463,8 @@ async function executeKeywordSearch( }) } -/** A signed-in reader's keyword ranking over the source chunks, in one statement; see {@link keywordCandidateRankingQuery}. */ -function rankSignedInKeywordCandidates(input: { +/** A keyword ranking over the source chunks, in one statement; see {@link keywordCandidateRankingQuery}. */ +function rankKeywordCandidates(input: { knowledgeBaseIds: string[] tsQuery: SQL tagFilterConditions: SQL[] From 8bdc00f1b23448b98b1a090c3ea7dbff28115a37 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 13 Oct 2026 11:39:59 -0700 Subject: [PATCH 2/2] test(knowledge): compare single-statement keyword ranking against the two-step reference --- .../keyword-ranking.integration.ts | 415 ++++++++++++++++++ 1 file changed, 415 insertions(+) create mode 100644 apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts diff --git a/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts b/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts new file mode 100644 index 00000000000..f6dc67b9131 --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts @@ -0,0 +1,415 @@ +/** + * The keyword leg of a search with no signed-in reader behind it, against real PostgreSQL: ranking + * every base in one statement answers what ranking each base (or all of them) in the old two-step + * pass did, tie groups compared as sets, and admits nothing the reader may not see. + * + * Every chunk's vector points away from the query vector, so the vector leg of a hybrid search + * finds nothing under its distance threshold and the fused list is the keyword ranking alone. + */ +import type { Principal } from '@sim/auth/principal' +import { db } from '@sim/db' +import { document, embedding, knowledgeBase, organization, user, workspace } from '@sim/db/schema' +import { deleteWorkspaceFixture } from '@sim/db/testing/workspace-fixtures' +import { generateId } from '@sim/utils/id' +import { and, desc, eq, inArray, type SQL, sql } from 'drizzle-orm' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { + createKnowledgeAclFixtureIds, + seedKnowledgeAclFixture, +} from '@/lib/knowledge/__integration__/seed-source-access-fixture' +import { confluencePageAcl } from '@/lib/knowledge/access/confluence-permissions' +import { knowledgeAccessCondition } from '@/lib/knowledge/access/predicate' +import { resolveKnowledgeAccessScope } from '@/lib/knowledge/access/scope' +import type { KnowledgeAccessScope } from '@/lib/knowledge/access/types' +import { FTS_CONFIG, getVisibilityConditions } from '@/lib/knowledge/search/candidates' +import { retrieveKnowledgeSearch } from '@/lib/knowledge/search/queries' +import { embeddingVectorValues } from '@/lib/knowledge/vector-columns' + +const DIMENSIONS = 1536 +const QUERY = 'orion' +const queryVector = { + vector: JSON.stringify([1, ...Array(DIMENSIONS - 1).fill(0)]), + dimensions: DIMENSIONS, + model: 'text-embedding-3-small', +} as const +const AWAY_FROM_QUERY = [-1, ...Array(DIMENSIONS - 1).fill(0)] + +/** More bases than the old pass ranked in one statement, so it took its per-base branch. */ +const EXTRA_BASE_COUNT = 5 +const DOCUMENTS_PER_BASE = 3 +const CHUNKS_PER_DOCUMENT = 2 +/** Hidden documents outrank every visible one, so admitting any of them would surface it first. */ +const HIDDEN_OCCURRENCES = 9 + +type HiddenKind = + | 'excluded' + | 'deleted' + | 'archived' + | 'disabled-document' + | 'disabled-chunk' + | 'pending' + | 'missing-group' + | 'stale-source' + | 'unsearched-base' + | 'restricted-source' + | 'user-only' + +/** Hidden from every reader. */ +const HIDDEN_FROM_ALL: HiddenKind[] = [ + 'excluded', + 'deleted', + 'archived', + 'disabled-document', + 'disabled-chunk', + 'pending', + 'missing-group', + 'stale-source', + 'unsearched-base', +] +/** Readable to the signed-in owner of the source groups, never to a workspace key. */ +const HIDDEN_FROM_WORKSPACE_KEY: HiddenKind[] = [ + ...HIDDEN_FROM_ALL, + 'restricted-source', + 'user-only', +] + +interface RankedKeywordRow { + id: string + keywordRank: number +} + +/** + * The unauthenticated keyword ranking as it stood before it moved into one statement: the oracle + * each case is compared against. + */ +async function rankBeforeSingleStatement(input: { + knowledgeBaseIds: string[] + topK: number + query: string + access: KnowledgeAccessScope +}): Promise { + const { knowledgeBaseIds, topK, query } = input + const tsQuery = sql`websearch_to_tsquery(${FTS_CONFIG}, ${query})` + const tagFilterConditions: SQL[] = [] + const rankExpr = sql`ts_rank_cd(${embedding.contentTsv}, ${tsQuery})` + const visibility = getVisibilityConditions(undefined, knowledgeAccessCondition(input.access)) + const rankRows = (kbScope: SQL) => + db + .select({ id: embedding.id, keywordRank: rankExpr.as('keyword_rank') }) + .from(embedding) + .innerJoin(document, eq(embedding.documentId, document.id)) + .where( + and( + kbScope, + ...visibility, + sql`${embedding.contentTsv} @@ ${tsQuery}`, + ...tagFilterConditions + ) + ) + .orderBy(sql`${rankExpr} DESC`) + .limit(topK) + + const kbCount = knowledgeBaseIds.length + const useParallel = kbCount > 4 || (kbCount > 2 && topK > 50) + if (useParallel) { + const perBase = await Promise.all( + knowledgeBaseIds.map((kbId) => rankRows(eq(embedding.knowledgeBaseId, kbId))) + ) + return perBase + .flat() + .sort((a, b) => b.keywordRank - a.keywordRank) + .slice(0, topK) + } + return rankRows(inArray(embedding.knowledgeBaseId, knowledgeBaseIds)) +} + +/** Candidates the hybrid keyword leg ranks before the fused list is trimmed to `topK`. */ +function hybridLegTopK(topK: number) { + return Math.min(Math.max(topK * 3, 50), 200) +} + +/** + * Asserts `actual` is the reference ranking's head: distinct scores in the same order, each tie + * group holding the same members (only the group the cut falls in may hold a subset). + */ +function expectSameRanking(actual: string[], reference: RankedKeywordRow[], topK: number) { + expect(new Set(actual).size).toBe(actual.length) + expect(actual).toHaveLength(Math.min(topK, reference.length)) + let cursor = 0 + let groupStart = 0 + while (cursor < actual.length && groupStart < reference.length) { + const score = reference[groupStart].keywordRank + let groupEnd = groupStart + while (groupEnd < reference.length && reference[groupEnd].keywordRank === score) groupEnd++ + const group = new Set(reference.slice(groupStart, groupEnd).map((row) => row.id)) + const taken = actual.slice(cursor, cursor + group.size) + for (const id of taken) { + expect(group.has(id), `${id} belongs to the tie group scored ${score}`).toBe(true) + } + cursor += taken.length + groupStart = groupEnd + } + expect(cursor).toBe(actual.length) +} + +describe('unauthenticated keyword ranking in one statement', () => { + const ids = createKnowledgeAclFixtureIds() + const extraBaseIds = Array.from({ length: EXTRA_BASE_COUNT }, () => generateId()) + const unsearchedBaseId = generateId() + const allBaseIds = [ids.knowledgeBaseId, ...extraBaseIds] + const hiddenChunkIds = new Map() + const workspaceKey: Principal = { + kind: 'workspace_api_key', + workspaceId: ids.workspaceId, + keyId: 'fixture-key', + } + const alice: Principal = { kind: 'session', userId: ids.aliceId, sessionId: 'fixture-alice' } + const restrictedAcl = confluencePageAcl({ + providerId: 'google-drive', + tenantId: 'fixture-tenant', + spacePrincipals: [{ kind: 'group', id: 'space' }], + restrictionChain: [[{ kind: 'group', id: 'parent' }], [{ kind: 'group', id: 'page' }]], + }) + + function chunkValues(input: { + documentId: string + knowledgeBaseId: string + chunkIndex: number + occurrences: number + enabled?: boolean + }) { + const id = generateId() + const content = `${Array(input.occurrences).fill(QUERY).join(' ')} survey notes` + return { + id, + documentId: input.documentId, + knowledgeBaseId: input.knowledgeBaseId, + chunkIndex: input.chunkIndex, + chunkHash: id, + content, + contentLength: content.length, + tokenCount: input.occurrences + 2, + startOffset: 0, + endOffset: content.length, + enabled: input.enabled ?? true, + ...embeddingVectorValues(DIMENSIONS, AWAY_FROM_QUERY), + } + } + + function documentValues(documentId: string, knowledgeBaseId: string) { + return { + id: documentId, + knowledgeBaseId, + filename: `${documentId}.txt`, + fileUrl: `https://fixture.invalid/${documentId}`, + fileSize: 12, + mimeType: 'text/plain', + processingStatus: 'completed', + acl: ['ws'], + } + } + + beforeAll(async () => { + await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' }) + await db.insert(knowledgeBase).values( + [...extraBaseIds, unsearchedBaseId].map((id, index) => ({ + id, + userId: ids.aliceId, + workspaceId: ids.workspaceId, + name: `Keyword ranking fixture ${index}`, + })) + ) + + const documents: ReturnType[] = [] + const chunks: ReturnType[] = [] + /** Occurrence counts repeat across bases and documents, so equal scores span bases. */ + allBaseIds.forEach((knowledgeBaseId, baseIndex) => { + for (let documentIndex = 0; documentIndex < DOCUMENTS_PER_BASE; documentIndex++) { + const documentId = generateId() + documents.push(documentValues(documentId, knowledgeBaseId)) + for (let chunkIndex = 0; chunkIndex < CHUNKS_PER_DOCUMENT; chunkIndex++) { + chunks.push( + chunkValues({ + documentId, + knowledgeBaseId, + chunkIndex, + occurrences: 1 + ((baseIndex + documentIndex + chunkIndex) % 4), + }) + ) + } + chunks.push( + chunkValues({ + documentId, + knowledgeBaseId, + chunkIndex: CHUNKS_PER_DOCUMENT, + occurrences: 0, + }) + ) + } + }) + + const hidden: Array<{ + kind: HiddenKind + document: Partial + chunkEnabled?: boolean + }> = [ + { kind: 'excluded', document: { userExcluded: true } }, + { kind: 'deleted', document: { deletedAt: new Date() } }, + { kind: 'archived', document: { archivedAt: new Date() } }, + { kind: 'disabled-document', document: { enabled: false } }, + { kind: 'disabled-chunk', document: {}, chunkEnabled: false }, + { kind: 'pending', document: { processingStatus: 'pending' } }, + { + kind: 'missing-group', + document: { + connectorId: ids.connectorId, + acl: ['g:google-drive:fixture-tenant:missing'], + aclVerifiedAt: new Date(), + }, + }, + { + kind: 'stale-source', + document: { + connectorId: ids.connectorId, + acl: [...restrictedAcl.acl], + aclRequirements: restrictedAcl.requirements.map((clause) => [...clause]), + aclVerifiedAt: new Date(0), + }, + }, + { kind: 'unsearched-base', document: {} }, + { + kind: 'restricted-source', + document: { + connectorId: ids.connectorId, + acl: [...restrictedAcl.acl], + aclRequirements: restrictedAcl.requirements.map((clause) => [...clause]), + aclVerifiedAt: new Date(), + }, + }, + { kind: 'user-only', document: { acl: [`u:${ids.aliceId}@fixture.test`] } }, + ] + const hiddenDocuments: (typeof document.$inferInsert)[] = [] + for (const entry of hidden) { + const documentId = generateId() + const knowledgeBaseId = + entry.kind === 'unsearched-base' ? unsearchedBaseId : ids.knowledgeBaseId + hiddenDocuments.push({ ...documentValues(documentId, knowledgeBaseId), ...entry.document }) + const chunk = chunkValues({ + documentId, + knowledgeBaseId, + chunkIndex: 0, + occurrences: HIDDEN_OCCURRENCES, + enabled: entry.chunkEnabled, + }) + hiddenChunkIds.set(entry.kind, chunk.id) + chunks.push(chunk) + } + + await db.insert(document).values([...documents, ...hiddenDocuments]) + await db.insert(embedding).values(chunks) + }) + + afterAll(async () => { + await deleteWorkspaceFixture(db, eq(workspace.id, ids.workspaceId)) + await db.delete(organization).where(eq(organization.id, ids.organizationId)) + await db.delete(user).where(inArray(user.id, [ids.aliceId, ids.bobId])) + await db.$client.end() + }) + + async function compare(input: { + principal: Principal + knowledgeBaseIds: string[] + topK: number + hiddenKinds: HiddenKind[] + }) { + const access = await resolveKnowledgeAccessScope(input.principal, { + workspaceId: ids.workspaceId, + }) + const { rows, retrieval } = await retrieveKnowledgeSearch({ + knowledgeBaseIds: input.knowledgeBaseIds, + topK: input.topK, + access, + searchMode: 'hybrid', + query: QUERY, + queryVector, + }) + expect(retrieval).toEqual({ status: 'complete', timedOutLegs: [] }) + const reference = await rankBeforeSingleStatement({ + knowledgeBaseIds: input.knowledgeBaseIds, + topK: hybridLegTopK(input.topK), + query: QUERY, + access, + }) + const actual = rows.map((row) => row.id) + expectSameRanking(actual, reference, input.topK) + + const hiddenIds = input.hiddenKinds.map((kind) => hiddenChunkIds.get(kind)) + expect(hiddenIds.every((id) => typeof id === 'string')).toBe(true) + for (const id of hiddenIds) { + expect(actual).not.toContain(id) + expect(reference.map((row) => row.id)).not.toContain(id) + } + return { actual, reference } + } + + it('seeds scores tied across bases at several distinct levels', async () => { + const scores = await db + .select({ + rank: sql`ts_rank_cd(${embedding.contentTsv}, websearch_to_tsquery(${FTS_CONFIG}, ${QUERY}))`, + bases: sql`count(DISTINCT ${embedding.knowledgeBaseId})::int`, + }) + .from(embedding) + .where( + and( + inArray(embedding.knowledgeBaseId, allBaseIds), + sql`${embedding.contentTsv} @@ websearch_to_tsquery(${FTS_CONFIG}, ${QUERY})` + ) + ) + .groupBy(sql`1`) + .orderBy(desc(sql`1`)) + expect(scores.length).toBeGreaterThan(2) + expect(scores.filter((score) => score.bases > 1).length).toBeGreaterThan(1) + }) + + describe.each([ + { label: 'every base (old per-base branch)', baseCount: allBaseIds.length }, + { label: 'two bases (old single-statement branch)', baseCount: 2 }, + ])('$label', ({ baseCount }) => { + const knowledgeBaseIds = allBaseIds.slice(0, baseCount) + + it.each([5, 12, 40])( + 'matches the two-step ranking for a workspace key at topK %i', + async (topK) => { + await compare({ + principal: workspaceKey, + knowledgeBaseIds, + topK, + hiddenKinds: HIDDEN_FROM_WORKSPACE_KEY, + }) + } + ) + + it.each([5, 12, 40])( + 'matches the two-step ranking for a session reader without an access provider at topK %i', + async (topK) => { + await compare({ + principal: alice, + knowledgeBaseIds, + topK, + hiddenKinds: HIDDEN_FROM_ALL, + }) + } + ) + }) + + it('ranks a source document only for the reader its groups admit', async () => { + const restrictedId = hiddenChunkIds.get('restricted-source') + const { actual } = await compare({ + principal: alice, + knowledgeBaseIds: allBaseIds, + topK: 5, + hiddenKinds: HIDDEN_FROM_ALL, + }) + expect(actual).toContain(restrictedId) + }) +})