From 91a04d4de871440342833bfe2433224131224e55 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 21 Sep 2026 00:36:57 -0700 Subject: [PATCH 1/2] improvement(knowledge): count an unproven reach, read memos under the search deadline, and drop the dead search paths Correctness - A resolved scope's reach was assumed broad whenever the bases were small enough for the probe limit, without a probe having saturated; the shortcut now applies only to a saturated probe, and an unproven reach is counted - The reach denominator and the source-index catalog were read outside the search's deadline on a memo miss; both run under the leg's budget now - A page whose identity read recovered fewer rows than its slice ended the candidate stream early; short is judged by what the pool gave, not by what the read returned Round trips - The keyword leg re-read every base's kind per search to decide the Tin engine; the search carries whether its bases are search indexes, and the flag and index readiness are read together - The v1 route re-read each result's document for its name and URL, which the rows already carry from hydration; the read, its helper and the retrieval result's read-access field are gone - Input provenance is prepared alongside the other pre-model reads; result provenance is imported per document in parallel; projection fill and the source-index memo are read together - A reordered leg sorts once at the end instead of re-sorting the result map after every page Dead code - `getQueryStrategy` fields nothing read, a never-emitted authorization stage, a one-valued candidate-storage diagnostic, the unused live-source `current` getter, the unread usage-admission policy field, the detached enqueue path the script never took, and a tautological Tin guard --- .../app/api/knowledge/search/utils.test.ts | 10 - apps/sim/app/api/knowledge/utils.ts | 10 +- .../app/api/v1/knowledge/search/route.test.ts | 70 ++----- apps/sim/app/api/v1/knowledge/search/route.ts | 14 +- .../knowledge/search/route.provenance.test.ts | 5 - apps/sim/lib/knowledge/application/search.ts | 110 ++++++----- apps/sim/lib/knowledge/search/diagnostics.ts | 4 +- .../projection-source-acl-backfill.test.ts | 8 +- .../search/projection-source-acl-backfill.ts | 30 +-- apps/sim/lib/knowledge/search/queries.test.ts | 65 ++++++- apps/sim/lib/knowledge/search/queries.ts | 179 ++++++------------ .../knowledge/search/source-vector-indexes.ts | 8 +- .../search/tin-keyword-readiness.test.ts | 7 +- .../lib/knowledge/search/tin-keyword.test.ts | 46 ++--- apps/sim/lib/knowledge/search/tin-keyword.ts | 54 ++---- .../scripts/backfill-projection-source-acl.ts | 4 +- 16 files changed, 269 insertions(+), 355 deletions(-) diff --git a/apps/sim/app/api/knowledge/search/utils.test.ts b/apps/sim/app/api/knowledge/search/utils.test.ts index e0eac74f0a6..afc6312ccc3 100644 --- a/apps/sim/app/api/knowledge/search/utils.test.ts +++ b/apps/sim/app/api/knowledge/search/utils.test.ts @@ -877,14 +877,4 @@ describe('Knowledge Search Utils', () => { ) }) }) - - describe('getDocumentMetadataByIds', () => { - it('should handle empty input gracefully', async () => { - const { getDocumentMetadataByIds } = await import('@/lib/knowledge/search/queries') - - const result = await getDocumentMetadataByIds([]) - - expect(result).toEqual({}) - }) - }) }) diff --git a/apps/sim/app/api/knowledge/utils.ts b/apps/sim/app/api/knowledge/utils.ts index 3a6a4aa98a2..8e5ccfd7c57 100644 --- a/apps/sim/app/api/knowledge/utils.ts +++ b/apps/sim/app/api/knowledge/utils.ts @@ -8,6 +8,7 @@ interface KnowledgeBaseData { userId: string workspaceId?: string | null name: string + isSearchIndex: boolean description?: string | null tokenCount: number embeddingModel: string @@ -22,7 +23,13 @@ export interface KnowledgeBaseAccessResult { hasAccess: true knowledgeBase: Pick< KnowledgeBaseData, - 'id' | 'userId' | 'workspaceId' | 'name' | 'embeddingModel' | 'embeddingDimension' + | 'id' + | 'userId' + | 'workspaceId' + | 'name' + | 'isSearchIndex' + | 'embeddingModel' + | 'embeddingDimension' > } @@ -52,6 +59,7 @@ async function resolveKnowledgeBaseAccess( userId: knowledgeBase.userId, workspaceId: knowledgeBase.workspaceId, name: knowledgeBase.name, + isSearchIndex: knowledgeBase.isSearchIndex, embeddingModel: knowledgeBase.embeddingModel, embeddingDimension: knowledgeBase.embeddingDimension, }) diff --git a/apps/sim/app/api/v1/knowledge/search/route.test.ts b/apps/sim/app/api/v1/knowledge/search/route.test.ts index 018e2040a4c..a1c7b9aca7d 100644 --- a/apps/sim/app/api/v1/knowledge/search/route.test.ts +++ b/apps/sim/app/api/v1/knowledge/search/route.test.ts @@ -16,7 +16,6 @@ const { mockExecuteKnowledgeSearch, mockRetrievalStatus, mockGenerateSearchEmbedding, - mockGetDocumentMetadataByIds, mockGetDocumentTagDefinitions, mockAuthenticateRequest, mockValidateWorkspaceAccess, @@ -28,7 +27,6 @@ const { mockExecuteKnowledgeSearch: vi.fn(), mockRetrievalStatus: vi.fn(() => ({ status: 'complete', timedOutLegs: [] })), mockGenerateSearchEmbedding: vi.fn(), - mockGetDocumentMetadataByIds: vi.fn(), mockGetDocumentTagDefinitions: vi.fn(), mockAuthenticateRequest: vi.fn(), mockValidateWorkspaceAccess: vi.fn(), @@ -60,9 +58,7 @@ vi.mock('@/lib/knowledge/search/queries', () => ({ retrieveKnowledgeSearch: async (params: { access: unknown }) => ({ rows: await mockExecuteKnowledgeSearch(params), retrieval: mockRetrievalStatus(), - readAccess: params.access, }), - getDocumentMetadataByIds: mockGetDocumentMetadataByIds, })) vi.mock('@/app/api/knowledge/utils', () => knowledgeApiUtilsMock) @@ -132,7 +128,6 @@ describe('v1 knowledge search route — per-KB embedding model', () => { isBYOK: false, }) mockExecuteKnowledgeSearch.mockResolvedValue([]) - mockGetDocumentMetadataByIds.mockResolvedValue({}) mockGetDocumentTagDefinitions.mockResolvedValue([]) mockResolveBillingAttribution.mockImplementation( ({ actorUserId, workspaceId }: { actorUserId: string; workspaceId: string }) => @@ -164,7 +159,6 @@ describe('v1 knowledge search route — per-KB embedding model', () => { ) expect(mockExecuteKnowledgeSearch).toHaveBeenCalledOnce() expect(response.status).toBe(500) - expect(mockGetDocumentMetadataByIds).not.toHaveBeenCalled() }) it('retains the reader provider for ranked results and returned document metadata', async () => { @@ -193,17 +187,11 @@ describe('v1 knowledge search route — per-KB embedding model', () => { accessProvider: provider, }) ) - expect(mockGetDocumentMetadataByIds).toHaveBeenCalledWith([], access) }) - it.each([ - ['query', false], - ['query', true], - ['filters', false], - ['filters', true], - ] as const)( - 'omits newly denied content from %s results and counts when all denied is %s', - async (mode, allDenied) => { + it.each(['query', 'filters'] as const)( + 'renders the source card each %s result row carries', + async (mode) => { const access = { kind: 'user' as const, userId: 'user-1', tokens: ['reader-token'] } const provider = { get: vi.fn().mockResolvedValue(access), @@ -219,26 +207,17 @@ describe('v1 knowledge search route — per-KB embedding model', () => { { tagSlot: 'tag1', displayName: 'category', fieldType: 'text' }, ]) mockExecuteKnowledgeSearch.mockResolvedValue([ - { - documentId: 'revoked-document', - knowledgeBaseId: 'kb-1', - content: 'revoked page content', - tag1: 'revoked tag', - chunkIndex: 0, - distance: 0.1, - }, { documentId: 'allowed-document', knowledgeBaseId: 'kb-1', content: 'allowed page content', + filename: 'Allowed page', + sourceUrl: null, tag1: 'docs', chunkIndex: 0, distance: 0.2, }, ]) - mockGetDocumentMetadataByIds.mockResolvedValue( - allDenied ? {} : { 'allowed-document': { filename: 'Allowed page', sourceUrl: null } } - ) const response = await POST( createMockRequest('POST', { workspaceId: 'ws-1', @@ -250,24 +229,19 @@ describe('v1 knowledge search route — per-KB embedding model', () => { ) const body = await response.json() expect(response.status).toBe(200) - expect(mockGetDocumentMetadataByIds).toHaveBeenCalledWith( - ['revoked-document', 'allowed-document'], - access - ) - expect(body.data.results).toEqual( - allDenied - ? [] - : [ - expect.objectContaining({ - documentId: 'allowed-document', - documentName: 'Allowed page', - content: 'allowed page content', - metadata: { category: 'docs' }, - }), - ] + expect(mockExecuteKnowledgeSearch).toHaveBeenCalledWith( + expect.objectContaining({ access, accessProvider: provider }) ) - expect(body.data.totalResults).toBe(allDenied ? 0 : 1) - expect(JSON.stringify(body)).not.toContain('revoked') + expect(body.data.results).toEqual([ + expect.objectContaining({ + documentId: 'allowed-document', + documentName: 'Allowed page', + sourceUrl: null, + content: 'allowed page content', + metadata: { category: 'docs' }, + }), + ]) + expect(body.data.totalResults).toBe(1) } ) @@ -372,7 +346,7 @@ describe('v1 knowledge search route — per-KB embedding model', () => { expect(mockGenerateSearchEmbedding).not.toHaveBeenCalled() }) - it('surfaces sourceUrl from document metadata in search results', async () => { + it('surfaces the sourceUrl a result row carries', async () => { mockCheckKnowledgeBaseAccess.mockResolvedValueOnce({ hasAccess: true, knowledgeBase: baseKb('kb-confluence', 'text-embedding-3-small'), @@ -382,16 +356,12 @@ describe('v1 knowledge search route — per-KB embedding model', () => { documentId: 'doc-confluence', knowledgeBaseId: 'kb-confluence', content: 'page content', + filename: 'Runbook.md', + sourceUrl: 'https://example.atlassian.net/wiki/spaces/DOCS/pages/12345', chunkIndex: 0, distance: 0.1, }, ]) - mockGetDocumentMetadataByIds.mockResolvedValue({ - 'doc-confluence': { - filename: 'Runbook.md', - sourceUrl: 'https://example.atlassian.net/wiki/spaces/DOCS/pages/12345', - }, - }) const req = createMockRequest('POST', { workspaceId: 'ws-1', diff --git a/apps/sim/app/api/v1/knowledge/search/route.ts b/apps/sim/app/api/v1/knowledge/search/route.ts index 7a63690de18..2321fa5ed07 100644 --- a/apps/sim/app/api/v1/knowledge/search/route.ts +++ b/apps/sim/app/api/v1/knowledge/search/route.ts @@ -17,7 +17,6 @@ import { import { SearchDeadlineError } from '@/lib/knowledge/search/budget' import { resolveKnowledgeSearchDefaults } from '@/lib/knowledge/search/defaults' import { - getDocumentMetadataByIds, type KnowledgeRetrievalResult, retrieveKnowledgeSearch, type SearchResult, @@ -267,6 +266,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => { accessProvider, searchMode, boostRecency, + searchIndexOnly: accessibleKbs.every((kb) => kb.isSearchIndex), query, queryVector: { vector: JSON.stringify(queryEmbeddingResult.embedding), @@ -316,14 +316,11 @@ export const POST = withRouteHandler(async (request: NextRequest) => { /** v1 cannot express an incomplete search, so a leg that ran out of time fails the request. */ if (retrieved.retrieval.status === 'partial') throw new SearchDeadlineError() const results = retrieved.rows - const documentIds = results.map((r) => r.documentId) - const documentMetadataMap = await getDocumentMetadataByIds(documentIds, retrieved.readAccess) - const readableResults = results.filter((result) => documentMetadataMap[result.documentId]) return NextResponse.json({ success: true, data: { - results: readableResults.map((result) => { + results: results.map((result) => { const kbTagMap = tagDefinitionsMap[result.knowledgeBaseId] || {} const tags: Record = {} @@ -335,11 +332,10 @@ export const POST = withRouteHandler(async (request: NextRequest) => { } }) - const docMeta = documentMetadataMap[result.documentId] return { documentId: result.documentId, - documentName: docMeta?.filename || undefined, - sourceUrl: docMeta?.sourceUrl ?? null, + documentName: result.filename || undefined, + sourceUrl: result.sourceUrl, content: result.content, chunkIndex: result.chunkIndex, metadata: tags, @@ -349,7 +345,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => { query: query || '', knowledgeBaseIds: accessibleKbIds, topK, - totalResults: readableResults.length, + totalResults: results.length, }, }) } catch (error) { diff --git a/apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts b/apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts index 5dd6d6c8a22..224f8840330 100644 --- a/apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts +++ b/apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts @@ -13,7 +13,6 @@ const mocks = vi.hoisted(() => ({ checkActorUsage: vi.fn(), generateEmbedding: vi.fn(), executeSearch: vi.fn(), - getDocumentMetadata: vi.fn(), getTagDefinitions: vi.fn(), recordEmbeddingUsage: vi.fn(), })) @@ -66,7 +65,6 @@ vi.mock('@/lib/knowledge/search/queries', () => ({ rows: await mocks.executeSearch(...args), retrieval: { status: 'complete', timedOutLegs: [] }, }), - getDocumentMetadataByIds: mocks.getDocumentMetadata, })) vi.mock('@/lib/knowledge/tags/service', () => ({ @@ -240,9 +238,6 @@ beforeEach(() => { mocks.checkActorUsage.mockResolvedValue({ isExceeded: false }) mocks.generateEmbedding.mockResolvedValue({ embedding: [0.1], isBYOK: false }) mocks.executeSearch.mockResolvedValue([row]) - mocks.getDocumentMetadata.mockResolvedValue({ - 'document-1': { filename: 'synthetic.txt', sourceUrl: null }, - }) mocks.getTagDefinitions.mockResolvedValue([]) v2RouteMocks.preauthRate.mockResolvedValue(V2_PREAUTH_RATE_LIMIT_ALLOWED) v2RouteMocks.operationRate.mockResolvedValue(V2_OPERATION_RATE_LIMIT_ALLOWED) diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index 41fe47443e7..9f050870807 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -66,7 +66,6 @@ const logger = createLogger('KnowledgeSearchApplication') export const KNOWLEDGE_SEARCH_COST_POLICY = { maxKnowledgeBases: 20, maxTopK: 100, - usageAdmission: 'before_model_execution', } as const export class KnowledgeSearchProvenanceUnavailableError extends Error { @@ -383,38 +382,44 @@ export async function runKnowledgeSearch({ dimensions: toKbEmbeddingDimensions(selectedTarget.dimensions), } : undefined - const preparedRegistry = input.prepareModelInputProvenance - ? await measureSearchStage('input_provenance', () => - input.prepareModelInputProvenance!({ userId, workspaceId: context.workspaceId }) - ) - : undefined - const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry input.signal?.throwIfAborted() - const [access, searchDefaults, billingAttribution, tagDefinitions, rerankerCredential] = - await Promise.all([ - measureSearchStage('access_scope', () => context.access.get()), - measureSearchStage('defaults', () => - resolveKnowledgeSearchDefaults({ - workspaceId: context.workspaceId, - organizationId: context.organizationId, + const [ + access, + searchDefaults, + billingAttribution, + tagDefinitions, + rerankerCredential, + preparedRegistry, + ] = await Promise.all([ + measureSearchStage('access_scope', () => context.access.get()), + measureSearchStage('defaults', () => + resolveKnowledgeSearchDefaults({ + workspaceId: context.workspaceId, + organizationId: context.organizationId, - /** The signed-in person, if any; never the billing owner or a key's creator. */ - userId: resolvePrincipalSubjectUserId(principal) ?? undefined, - requestedMode: input.searchMode, - }) - ), - admit(), - /** The tag names the results are labelled with depend on the bases alone. */ - filters.length === 0 - ? measureSearchStage('tag_definitions', () => - getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds) - ) - : Promise.resolve(definitionsByKnowledgeBase), - /** A surface may ask to rerank; without a key for the workspace or the platform there is nothing to ask. */ - input.rerankerEnabled && hasQuery - ? hasRerankerCredential(context.workspaceId, input.rerankerApiKey) - : false, - ]) + /** The signed-in person, if any; never the billing owner or a key's creator. */ + userId: resolvePrincipalSubjectUserId(principal) ?? undefined, + requestedMode: input.searchMode, + }) + ), + admit(), + /** The tag names the results are labelled with depend on the bases alone. */ + filters.length === 0 + ? measureSearchStage('tag_definitions', () => + getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds) + ) + : Promise.resolve(definitionsByKnowledgeBase), + /** A surface may ask to rerank; without a key for the workspace or the platform there is nothing to ask. */ + input.rerankerEnabled && hasQuery + ? hasRerankerCredential(context.workspaceId, input.rerankerApiKey) + : false, + input.prepareModelInputProvenance + ? measureSearchStage('input_provenance', () => + input.prepareModelInputProvenance!({ userId, workspaceId: context.workspaceId }) + ) + : undefined, + ]) + const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry definitionsByKnowledgeBase = tagDefinitions input.signal?.throwIfAborted() /** Requested only once every prerequisite held: a search refused for any reason spends no model call. */ @@ -461,6 +466,7 @@ export async function runKnowledgeSearch({ } : undefined, structuredFilters: structuredFilters.length > 0 ? structuredFilters : undefined, + searchIndexOnly: context.knowledgeBases.every((knowledgeBase) => knowledgeBase.isSearchIndex), }) ) @@ -687,22 +693,32 @@ export async function runKnowledgeSearch({ } }) if (registry && provenanceSnapshot) { - for (const [documentId, document] of Object.entries(provenanceSnapshot.documentMetadata)) { - const renderedMetadata = results - .filter((result) => result.documentId === documentId) - .map((result) => ({ - documentName: result.documentName, - sourceUrl: result.sourceUrl, - metadata: result.metadata, - })) - if (renderedMetadata.length === 0) continue - if ( - !(await measureSearchStage('metadata_provenance', () => - importDurableSecretProvenance(registry, document.provenance, renderedMetadata) - )) - ) { - registry.markIncomplete('knowledge-result-provenance-unavailable') - } + const renderedByDocument = new Map< + string, + Array> + >() + for (const result of results) { + const rendered = renderedByDocument.get(result.documentId) ?? [] + rendered.push({ + documentName: result.documentName, + sourceUrl: result.sourceUrl, + metadata: result.metadata, + }) + renderedByDocument.set(result.documentId, rendered) + } + /** Each document's provenance stands alone, so they are imported together. */ + const imported = await measureSearchStage('metadata_provenance', () => + Promise.all( + Object.entries(provenanceSnapshot.documentMetadata).map(([documentId, document]) => { + const renderedMetadata = renderedByDocument.get(documentId) + return renderedMetadata + ? importDurableSecretProvenance(registry, document.provenance, renderedMetadata) + : true + }) + ) + ) + if (imported.includes(false)) { + registry.markIncomplete('knowledge-result-provenance-unavailable') } } annotateSearchDiagnostics({ resultCount: results.length }) diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index 57db3d6b7c9..fe0086eed4c 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -38,12 +38,10 @@ export type SearchStage = | 'usage_recording' | 'overage_billing' | 'tag_definitions' - | 'metadata.sql' | 'metadata_provenance' | 'activity_recording' | RetrievalLeg | `${RetrievalLeg}.candidates` - | `${RetrievalLeg}.authorization` | `${RetrievalLeg}.hydration` | `${RetrievalLeg}.connection_acquire` | `${RetrievalLeg}.sql` @@ -51,6 +49,7 @@ export type SearchStage = | 'vector.probe' | 'vector.page' | 'vector.projection_filled' + | 'vector.source_indexes' | 'keyword.projection_filled' | 'vector.exact_candidates' | 'vector.exact' @@ -86,7 +85,6 @@ export interface SearchDiagnosticMetadata { boostRecency?: boolean embeddingDimensions?: number vectorRanking?: 'exact' | 'exact-candidates' | 'projection-walk' | 'per-source' - vectorCandidateStorage?: 'stored-halfvec' /** * Whether the bounded traversal filled its candidate limit. `underfilled` means visibility * removed enough neighbours that the rerank pool is smaller than requested, which lowers recall diff --git a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts index 63d8d3faa5e..4a357e3e9a3 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts @@ -117,7 +117,7 @@ describe('enqueueProjectionSourceAclBackfill', () => { }) it('hands the backfill to the Trigger.dev worker when one is configured', async () => { - await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 }, true)).resolves.toEqual({ + await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 })).resolves.toEqual({ runId: 'run-1', }) expect(mockTasksTrigger).toHaveBeenCalledWith( @@ -127,10 +127,4 @@ describe('enqueueProjectionSourceAclBackfill', () => { ) expect(mockBackfill).not.toHaveBeenCalled() }) - - it('fills the projections detached in this process without one', async () => { - await expect(enqueueProjectionSourceAclBackfill({}, false)).resolves.toBeNull() - expect(mockTasksTrigger).not.toHaveBeenCalled() - await vi.waitFor(() => expect(mockBackfill).toHaveBeenCalledTimes(2)) - }) }) diff --git a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts index 43dabf7e409..6ac3947d5d3 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts @@ -7,9 +7,6 @@ import { import { createLogger } from '@sim/logger' import postgres from 'postgres' import { resolveTriggerRegion } from '@/lib/core/async-jobs/region' -import { env } from '@/lib/core/config/env' -import { isTriggerDevEnabled } from '@/lib/core/config/env-flags' -import { runDetached } from '@/lib/core/utils/background' import { prewarmSearchProjection } from '@/lib/knowledge/search/prewarm' const logger = createLogger('ProjectionSourceAclBackfill') @@ -86,23 +83,16 @@ export async function runProjectionSourceAclBackfill( } /** - * Starts the backfill the way the table backfill is started: on the deployment's Trigger.dev - * worker when one is configured, where bounded runs chain until both projections are filled, and - * detached in this process otherwise. Safe to call again at any time — a run only fills rows still - * unset. + * Starts the backfill on the deployment's Trigger.dev worker, where bounded runs chain until both + * projections are filled. Safe to call again at any time: a run only fills rows still unset. */ export async function enqueueProjectionSourceAclBackfill( - payload: ProjectionSourceAclBackfillPayload = {}, - useTrigger = Boolean(isTriggerDevEnabled && env.TRIGGER_SECRET_KEY) -): Promise<{ runId: string } | null> { - if (useTrigger) { - const { tasks } = await import('@trigger.dev/sdk') - const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, payload, { - region: await resolveTriggerRegion(), - }) - logger.info('Projection source and ACL backfill enqueued', { runId: handle.id }) - return { runId: handle.id } - } - runDetached(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, () => runProjectionSourceAclBackfill(payload)) - return null + payload: ProjectionSourceAclBackfillPayload = {} +): Promise<{ runId: string }> { + const { tasks } = await import('@trigger.dev/sdk') + const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, payload, { + region: await resolveTriggerRegion(), + }) + logger.info('Projection source and ACL backfill enqueued', { runId: handle.id }) + return { runId: handle.id } } diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 52266d4a23c..7ebc4a69797 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -587,6 +587,21 @@ describe('workspace-scoped vector retrieval', () => { expect(statements().filter((query) => isWalk(query.sql))).toHaveLength(1) }) + it('asks for the next page when an identity read recovers fewer rows than its slice', async () => { + const execute = dbChainMockFns.execute.getMockImplementation()! + let pages = 0 + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query) + /** The first slice's read finds one of its twenty rows still present; the pool has more. */ + if (isPageStatement(statement.sql)) return pages++ === 0 ? [ranked[0]] : ranked + return execute(query) + }) + queueTableRows(schemaMock.embedding, [ranked[0]]) + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far']) + expect(pages).toBe(2) + }) + it('sizes the pool to the pages asked for, doubling a pool the pages outran', () => { expect(vectorCandidatePoolLimit(20, undefined)).toBe(200) expect(vectorCandidatePoolLimit(150, undefined)).toBe(300) @@ -710,7 +725,6 @@ describe('workspace-scoped vector retrieval', () => { ).toEqual({ rows: [], retrieval: { status: 'partial', timedOutLegs: ['vector'] }, - readAccess: params.access, }) } ) @@ -819,7 +833,6 @@ describe('workspace-scoped vector retrieval', () => { expect(result).toEqual({ rows: [], retrieval: { status: 'partial', timedOutLegs: ['vector'] }, - readAccess: params.access, }) } for (const resume of release) resume() @@ -1537,6 +1550,7 @@ describe('permitted-document planner', () => { query: 'release', queryVector: params.queryVector!, permitted: unbounded, + searchIndexOnly: true, ...overrides, }) const tinStatements = () => @@ -1562,7 +1576,7 @@ describe('permitted-document planner', () => { const results = await keyword() expect(results.map((row) => row.id)).toEqual(['a']) expect(mockResolveTinKeywordQuery).toHaveBeenCalledWith( - ['org-index'], + true, 'release', 'english', params.budget @@ -1895,6 +1909,35 @@ describe('permitted-document planner', () => { expect(reachCounts()).toHaveLength(2) }) + it('counts a resolved reach against a small index instead of assuming it broad', async () => { + /** A bound inside the probe limit proves nothing without a saturated probe. */ + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query).sql + if (statement.includes('EXPLAIN')) + return [{ 'QUERY PLAN': [{ Plan: { 'Plan Rows': 1_000 } }] }] + if (statement.includes(') reached')) return [{ n: 100 }] + return [] + }) + const reachCounts = () => statements().filter((query) => query.sql.includes(') reached')) + const reach = await resolveReach( + ['org-index'], + scope('small-index'), + new SearchBudget('vector', performance.now() + 10_000), + { + connectors: { workspace: [], admin: [], members: [], liveProofRequired: [] }, + observers: { confirmed: [], observed: [] }, + memberSources: [], + connectorTypes: new Map(), + uploads: true, + } + ) + expect(reach).toEqual({ kind: 'unbounded', broad: false }) + expect(reachCounts()).toHaveLength(1) + /** The count is the search's own read: it runs inside the leg's deadline statement. */ + const countAt = statements().findIndex((query) => query.sql.includes(') reached')) + expect(statements()[countAt - 1].sql).toContain('statement_timeout') + }) + it('does not remember a reach whose count ran out of time', async () => { dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql @@ -2309,6 +2352,22 @@ describe('filters on a resolved scope', () => { expect(JSON.stringify(probes[0])).toContain('"type":"gte"') }) + it('tests a tag filter on every walked row of a planned scope', async () => { + traversedRows = [{ id: 'a' }] + rerankRows = [hit('a', 'src-a')] + queueTableRows(schemaMock.embedding, [{ ...rerankRows[0], tag1: 'release' }]) + await handleVectorOnlySearch({ + ...params, + permitted: { kind: 'unbounded', broad: true }, + accessPlan: plan(), + structuredFilters: [{ tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'release' }], + }) + const walks = statements().filter((query) => isWalk(query.sql)) + expect(walks).toHaveLength(1) + /** The filter is chunk-level and the row does not carry it, so the walk asks the document: no unfiltered row fills the pool. */ + expect(JSON.stringify(walks[0])).toContain('release') + }) + it('ranks a date-bounded set exactly even when a member source has its own index', async () => { indexedSourceRows = [{ name: 'idx', connectorId: 'member-src' }] exactRows = [{ id: 'a' }] diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 75cf97dc86f..3f0e4092cca 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -152,7 +152,7 @@ async function isProjectionFilled( return (await projectionFilled.fetch(projection, { context: { budget, stage } })) ?? false } -/** Forgets whether the projections were filled, after their rows changed. */ +/** Forgets whether the projections were filled; the memo is per process and otherwise expires on its own. */ export function forgetProjectionFilled(): void { projectionFilled.clear() } @@ -229,8 +229,8 @@ let hnswSettingsUnsupportedUntil = 0 */ async function withVectorScanSettings( run: (executor: SearchExecutor) => Promise, - budget?: SearchBudget, - stage: SearchStage = 'vector.candidate_search', + budget: SearchBudget | undefined, + stage: SearchStage, maxScanTuples: number = Number(CANDIDATE_HNSW_MAX_SCAN_TUPLES) ): Promise { const untuned = () => runSearchQuery(budget, stage, run) @@ -274,65 +274,6 @@ async function withVectorScanSettings( } } -export interface DocumentMetadata { - filename: string - sourceUrl: string | null - /** When the source last changed the document; null for uploads and sources that do not say. */ - sourceModifiedAt: Date | null - /** The connector the document was synced through; null for an upload. */ - connectorType: string | null -} - -/** - * Batch-fetch display metadata for documents referenced by search results, under the full read - * predicate and the scope the results were read under — with the live grants that scope resolved, - * so a gated source's result keeps its name and URL, and a revoked one loses them here too. - * Returns a map keyed by document id; missing ids indicate the document is no longer visible and - * should be skipped. - */ -export async function getDocumentMetadataByIds( - documentIds: string[], - access: KnowledgeAccessScope -): Promise> { - if (documentIds.length === 0) { - return {} - } - - const uniqueIds = [...new Set(documentIds)] - const documents = await measureSearchStage('metadata.sql', () => - db - .select({ - id: document.id, - filename: document.filename, - sourceUrl: document.sourceUrl, - sourceModifiedAt: document.sourceModifiedAt, - connectorType: knowledgeConnector.connectorType, - }) - .from(document) - .leftJoin(knowledgeConnector, eq(knowledgeConnector.id, document.connectorId)) - .where( - and( - inArray(document.id, uniqueIds), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt), - knowledgeAccessCondition(access) - ) - ) - ) - const map: Record = {} - documents.forEach((doc) => { - map[doc.id] = { - filename: doc.filename, - sourceUrl: doc.sourceUrl ?? null, - sourceModifiedAt: doc.sourceModifiedAt ?? null, - connectorType: doc.connectorType ?? null, - } - }) - - return map -} - export interface SearchResult { id: string content: string @@ -732,8 +673,6 @@ export interface LiveSourceAccess { gates: (connectorId: string) => boolean /** The caller's scope with its grants, and the gated sources those grants do not cover. */ resolve: () => Promise<{ access: KnowledgeAccessScope; denied: ReadonlySet }> - /** The scope content was read under: with its grants once they were resolved, else as given. */ - current: () => Promise } /** Binds a search's gated sources to one memoized resolution of the caller's grants. */ @@ -755,7 +694,6 @@ export function liveSourceAccessFor( } return { gates: (connectorId) => gated.has(connectorId), - current: async () => (pending ? (await pending).access : access), resolve: () => { pending ??= measureSearchStage('live_source_grants', async () => { const scopes = await mapWithConcurrency(pages, SOURCE_RANKING_CONCURRENCY, (page) => @@ -840,8 +778,9 @@ async function selectAuthorizedSearchResults(input: { ) if (!page.candidates.length) break scanned += page.candidates.length + /** Short means the ranking had fewer to give, not that a read recovered fewer than it asked. */ + lastPageShort = page.nextOffset - offset < pageSize offset = page.nextOffset - lastPageShort = page.candidates.length < pageSize pending = page.candidates.filter((candidate) => !considered.has(candidate.id)) if (!pending.length) { if (lastPageShort) break @@ -894,11 +833,6 @@ async function selectAuthorizedSearchResults(input: { if (row) results.set(row.id, row) if (!input.compareResults && results.size === input.topK) break } - if (input.compareResults) { - const ranked = [...results.values()].sort(input.compareResults).slice(0, input.topK) - results.clear() - for (const row of ranked) results.set(row.id, row) - } if (refill) { /** The rebuilt pages are a new stream of candidates, so the scan budget starts over. */ offset = 0 @@ -913,7 +847,9 @@ async function selectAuthorizedSearchResults(input: { if (!input.budget?.isTimeout(error)) throw error } input.signal?.throwIfAborted() - return [...results.values()] + const rows = [...results.values()] + /** A reordered leg keeps every page's rows until the end: a later page cannot displace what an earlier one ranked. */ + return input.compareResults ? rows.sort(input.compareResults).slice(0, input.topK) : rows } /** Keeps the candidates of sources the caller turned out not to hold out of a page. */ @@ -975,15 +911,9 @@ export function hybridCandidateCount(topK: number): number { } export function getQueryStrategy(kbCount: number, topK: number) { - const useParallel = kbCount > 4 || (kbCount > 2 && topK > 50) - const distanceThreshold = kbCount > 3 ? 0.8 : 1.0 - const parallelLimit = Math.ceil(topK / kbCount) + 5 - return { - useParallel, - distanceThreshold, - parallelLimit, - singleQueryOptimized: kbCount <= 2, + useParallel: kbCount > 4 || (kbCount > 2 && topK > 50), + distanceThreshold: kbCount > 3 ? 0.8 : 1.0, } } @@ -1251,19 +1181,21 @@ const saturatedReach = new LRUCache({ }) /** How many documents the bases hold: the denominator of a reach share, and it moves slowly. */ -const indexDocumentCounts = new LRUCache({ +const indexDocumentCounts = new LRUCache({ max: 1000, ttl: SATURATED_REACH_TTL_MS, /** * The planner's estimate of the bases' documents, from the statistics it already keeps: a share * threshold needs the order of magnitude, and counting every row to learn it costs more than the - * search it serves. + * search it serves. The read that misses is the search's own, under its deadline. */ - fetchMethod: async (key) => { - const [row] = await db.execute<{ 'QUERY PLAN': Array<{ Plan: { 'Plan Rows': number } }> }>(sql` + fetchMethod: async (key, _stale, { context: budget }) => { + const [row] = await runSearchQuery(budget, 'permitted_documents', (executor) => + executor.execute<{ 'QUERY PLAN': Array<{ Plan: { 'Plan Rows': number } }> }>(sql` EXPLAIN (FORMAT JSON) SELECT 1 FROM ${document} WHERE ${document.knowledgeBaseId} = ANY(${textArrayLiteral(key.split(','))}) AND ${document.deletedAt} IS NULL`) + ) /** An empty answer is not remembered; the bases may simply not have been analyzed yet. */ return Number(row?.['QUERY PLAN']?.[0]?.Plan?.['Plan Rows'] ?? 0) || undefined }, @@ -1307,20 +1239,25 @@ async function estimateFilteredDocuments( } /** - * Whether a saturated reach is broad: the caller reaches at least {@link BROAD_REACH_SHARE} of the - * bases' documents. Counted once against that bound and remembered with the saturation, so the - * first search after the window pays for it and the rest do not. + * Whether a reach is broad: the caller reaches at least {@link BROAD_REACH_SHARE} of the bases' + * documents. Counted once against that bound and remembered, so the first search after the + * window pays for it and the rest do not. A caller whose probe already saturated is known to + * reach past the probe's limit, so a bound inside that limit is met without counting. */ async function reachIsBroad( knowledgeBaseIds: string[], access: KnowledgeAccessScope, budget: SearchBudget | undefined, - plan: SearchAccessPlan | undefined + plan: SearchAccessPlan | undefined, + saturated: boolean ): Promise { if (access.kind !== 'user') return true - const total = (await indexDocumentCounts.fetch([...knowledgeBaseIds].sort().join(','))) ?? 0 + const total = + (await indexDocumentCounts.fetch([...knowledgeBaseIds].sort().join(','), { + context: budget, + })) ?? 0 const bound = Math.ceil(total * BROAD_REACH_SHARE) - if (bound <= VECTOR_PROBE_DOCUMENT_LIMIT) return true + if (saturated && bound <= VECTOR_PROBE_DOCUMENT_LIMIT) return true const [row] = await runSearchQuery(budget, 'permitted_documents', (executor) => executor.execute<{ n: number }>(sql` SELECT count(*) AS n FROM ( @@ -1385,7 +1322,7 @@ export async function resolveReach( const remembered = key ? saturatedReach.get(key) : undefined if (remembered) return { kind: 'unbounded', broad: remembered.broad } try { - const broad = await reachIsBroad(knowledgeBaseIds, access, budget, plan) + const broad = await reachIsBroad(knowledgeBaseIds, access, budget, plan, false) if (key) saturatedReach.set(key, { broad }) return { kind: 'unbounded', broad } } catch (error) { @@ -1417,10 +1354,11 @@ export async function resolvePermittedDocuments(params: { * A remembered reach says how much of the bases the caller reads, which a date filter does not * change; the filtered set still has to be enumerated, so under one the probe always runs. */ - const remembered = - key && !(params.accessPlan && (dateFilterCondition(params.filters) || params.filters?.source)) - ? saturatedReach.get(key) - : undefined + /** A plan under a date or source filter enumerates the filtered set directly; reach cannot stand in for it. */ + const filteredDirectly = Boolean( + params.accessPlan && (dateFilterCondition(params.filters) || params.filters?.source) + ) + const remembered = key && !filteredDirectly ? saturatedReach.get(key) : undefined if (remembered) { probe = { kind: 'saturated' } broad = remembered.broad @@ -1437,9 +1375,7 @@ export async function resolvePermittedDocuments(params: { params.access, params.budget, 'permitted_documents', - params.accessPlan && (dateFilterCondition(params.filters) || params.filters?.source) - ? 'direct' - : 'reach-first' + filteredDirectly ? 'direct' : 'reach-first' ) } catch (error) { if (!params.budget?.isTimeout(error)) throw error @@ -1451,7 +1387,8 @@ export async function resolvePermittedDocuments(params: { params.knowledgeBaseIds, params.access, params.budget, - params.accessPlan + params.accessPlan, + true ) } catch (error) { if (!params.budget?.isTimeout(error)) throw error @@ -1547,7 +1484,7 @@ async function selectSourceVectorCandidates(input: { }): Promise { const sources = planSourceVectorCandidates({ plan: input.plan, - indexedSources: await indexedVectorSources(), + indexedSources: await indexedVectorSources(input.budget), }) annotateSearchDiagnostics({ vectorRanking: 'per-source', @@ -1778,11 +1715,11 @@ async function selectVectorResults(params: SearchParams): Promise> = null if (params.permitted?.kind === 'unbounded' && tagFilterConditions.length === 0) { try { - tinQuery = await resolveTinKeywordQuery(knowledgeBaseIds, query, FTS_CONFIG, params.budget) + tinQuery = await resolveTinKeywordQuery( + params.searchIndexOnly === true, + query, + FTS_CONFIG, + params.budget + ) } catch (error) { /** A leg whose deadline passed before it ranked anything is short, not failed. */ if (!params.budget?.isTimeout(error)) throw error return [] } } - annotateSearchDiagnostics({ - ...(params.permitted?.kind === 'unbounded' - ? { keywordRanking: tinQuery ? 'tin' : 'gin' } - : {}), - }) + if (params.permitted?.kind === 'unbounded') + annotateSearchDiagnostics({ keywordRanking: tinQuery ? 'tin' : 'gin' }) const accessPlan = access.kind === 'user' ? params.accessPlan : undefined /** A filled projection decides readability on the ranked row alone; none of its rows needs the document. */ const tinFilled = @@ -2257,7 +2193,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise ? params.permitted.documents.map((entry) => entry.id) : undefined if (permittedIds?.length === 0) return { candidates: [], nextOffset: offset } - if (tinScope && (accessPlan || !permittedIds)) { + if (tinScope) { const tinPage = await selectTinPage(tinScope, limit, offset, excludedSources) if (tinPage) return tinPage annotateSearchDiagnostics({ keywordRanking: 'gin' }) @@ -2514,6 +2450,8 @@ export interface ExecuteKnowledgeSearchParams { queryVector?: KnowledgeQueryVector structuredFilters?: StructuredFilter[] filters?: WorkspaceSearchFilters + /** Every base is an organization search index; only those are projected for Tin ranking. */ + searchIndexOnly?: boolean } export interface RetrievalStatus { @@ -2524,11 +2462,9 @@ export interface RetrievalStatus { export interface KnowledgeRetrievalResult { rows: SearchResult[] retrieval: RetrievalStatus - /** The scope the returned content was read under; what may see these rows may see their metadata. */ - readAccess: KnowledgeAccessScope } -/** Legacy surfaces cannot silently present partial retrieval as complete. */ +/** Retrieval for a surface that cannot present a partial result as complete. */ export async function executeKnowledgeSearch( params: ExecuteKnowledgeSearchParams ): Promise { @@ -2563,14 +2499,12 @@ export async function retrieveKnowledgeSearch( } const finish = async (rows: SearchResult[]): Promise => { params.signal?.throwIfAborted() - const readAccess = (await liveSourceAccess?.current()) ?? access const timedOutLegs = Object.values(budgets) .filter((budget) => budget.timedOut) .map((budget) => budget.leg) return { rows: boostRecency ? applyRecencyBoost(rows) : rows, retrieval: { status: timedOutLegs.length ? 'partial' : 'complete', timedOutLegs }, - readAccess, } } /** @@ -2605,6 +2539,7 @@ export async function retrieveKnowledgeSearch( filters: params.filters, structuredFilters, liveSourceAccess, + searchIndexOnly: params.searchIndexOnly, } const hasQuery = Boolean(query?.trim()) const hasFilters = Boolean(structuredFilters?.length) diff --git a/apps/sim/lib/knowledge/search/source-vector-indexes.ts b/apps/sim/lib/knowledge/search/source-vector-indexes.ts index 3ac351199f4..9686c402ddd 100644 --- a/apps/sim/lib/knowledge/search/source-vector-indexes.ts +++ b/apps/sim/lib/knowledge/search/source-vector-indexes.ts @@ -4,6 +4,7 @@ import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' import { and, count, eq, isNull, sql } from 'drizzle-orm' import { LRUCache } from 'lru-cache' +import { runSearchQuery, type SearchBudget } from '@/lib/knowledge/search/budget' const logger = createLogger('SourceVectorIndexes') @@ -37,15 +38,18 @@ export function forgetIndexedVectorSources(): void { indexedSources.clear() } -export async function indexedVectorSources(): Promise> { +/** The sources with a graph of their own; a search that misses the memo reads under its own deadline. */ +export async function indexedVectorSources(budget?: SearchBudget): Promise> { const cached = indexedSources.get('sources') if (cached) return cached - const rows = await db.execute<{ connectorId: string | null }>(sql` + const rows = await runSearchQuery(budget, 'vector.source_indexes', (executor) => + executor.execute<{ connectorId: string | null }>(sql` SELECT substring(pg_get_expr(i.indpred, i.indrelid) from '''([0-9a-f-]+)''') AS "connectorId" FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid WHERE i.indrelid = 'embedding_search'::regclass AND c.relname LIKE 'embedding_search_src_%' AND i.indisvalid AND i.indisready`) + ) const sources = new Set( rows.map((row) => row.connectorId).filter((id): id is string => id !== null) ) diff --git a/apps/sim/lib/knowledge/search/tin-keyword-readiness.test.ts b/apps/sim/lib/knowledge/search/tin-keyword-readiness.test.ts index 1cbf5bbaa3c..f42b4bc2cfb 100644 --- a/apps/sim/lib/knowledge/search/tin-keyword-readiness.test.ts +++ b/apps/sim/lib/knowledge/search/tin-keyword-readiness.test.ts @@ -1,7 +1,7 @@ /** * @vitest-environment node */ -import { dbChainMockFns, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' +import { dbChainMockFns, resetDbChainMock } from '@sim/testing' import { expect, it, vi } from 'vitest' vi.mock('@/lib/core/config/feature-flags', () => ({ @@ -20,10 +20,9 @@ it('stays on the GIN projection while the Tin index is incomplete, and remembers if (text.includes('websearch_to_tsquery')) return [{ rendered: "'releas'" }] return [] }) - queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-index', isSearchIndex: true }]) - expect(await resolveTinKeywordQuery(['kb-index'], 'release', 'english', undefined)).toBeNull() + expect(await resolveTinKeywordQuery(true, 'release', 'english', undefined)).toBeNull() indexValid = true - expect(await resolveTinKeywordQuery(['kb-index'], 'release', 'english', undefined)).toBeNull() + expect(await resolveTinKeywordQuery(true, 'release', 'english', undefined)).toBeNull() const readinessReads = dbChainMockFns.execute.mock.calls.filter(([query]) => JSON.stringify(query).includes('indisvalid') ) diff --git a/apps/sim/lib/knowledge/search/tin-keyword.test.ts b/apps/sim/lib/knowledge/search/tin-keyword.test.ts index fabc8afc1ff..3e84779f978 100644 --- a/apps/sim/lib/knowledge/search/tin-keyword.test.ts +++ b/apps/sim/lib/knowledge/search/tin-keyword.test.ts @@ -1,7 +1,7 @@ /** * @vitest-environment node */ -import { dbChainMockFns, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' +import { dbChainMockFns, resetDbChainMock } from '@sim/testing' import { beforeEach, describe, expect, it, vi } from 'vitest' const { mockIsFeatureEnabled } = vi.hoisted(() => ({ @@ -15,10 +15,7 @@ vi.mock('@/lib/core/config/feature-flags', () => ({ import { SearchBudget, SearchDeadlineError } from '@/lib/knowledge/search/budget' import { resolveTinKeywordQuery } from '@/lib/knowledge/search/tin-keyword' -/** - * Base kinds are cached per process, so each case uses its own bases. Readiness is cached too; - * the incomplete-index case lives in its own file, where the cache starts empty. - */ +/** Readiness is cached per process; the incomplete-index case lives in its own file, where the cache starts empty. */ describe('resolveTinKeywordQuery', () => { let indexValid: boolean let rendered: string @@ -38,53 +35,40 @@ describe('resolveTinKeywordQuery', () => { it('is off while the rollout flag is off, without touching the database', async () => { mockIsFeatureEnabled.mockResolvedValue(false) + expect(await resolveTinKeywordQuery(true, 'release notes', 'english', undefined)).toBeNull() expect( - await resolveTinKeywordQuery(['kb-flag-off'], 'release notes', 'english', undefined) - ).toBeNull() - expect(dbChainMockFns.execute).not.toHaveBeenCalled() + dbChainMockFns.execute.mock.calls.some(([query]) => + JSON.stringify(query).includes('websearch_to_tsquery') + ) + ).toBe(false) }) - it('translates the analyzed query once every base is a search index', async () => { - queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-index', isSearchIndex: true }]) - expect(await resolveTinKeywordQuery(['kb-index'], 'release notes', 'english', undefined)).toBe( + it('translates the analyzed query when every base is a search index', async () => { + expect(await resolveTinKeywordQuery(true, 'release notes', 'english', undefined)).toBe( '("releas" AND "note")' ) }) - it('is off when any base is not a search index, since only those are projected', async () => { - queueTableRows(schemaMock.knowledgeBase, [ - { id: 'kb-index-2', isSearchIndex: true }, - { id: 'kb-plain', isSearchIndex: false }, - ]) - expect( - await resolveTinKeywordQuery( - ['kb-index-2', 'kb-plain'], - 'release notes', - 'english', - undefined - ) - ).toBeNull() + it('is off when a base is not a search index, without touching the database', async () => { + expect(await resolveTinKeywordQuery(false, 'release notes', 'english', undefined)).toBeNull() + expect(dbChainMockFns.execute).not.toHaveBeenCalled() }) it('is off for a query Tin cannot express', async () => { - queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-negation', isSearchIndex: true }]) rendered = "!'draft'" - expect(await resolveTinKeywordQuery(['kb-negation'], '-draft', 'english', undefined)).toBeNull() + expect(await resolveTinKeywordQuery(true, '-draft', 'english', undefined)).toBeNull() }) it('reads under the keyword budget, so an expired deadline ends the leg instead of querying', async () => { const expired = new SearchBudget('keyword', performance.now() - 1) await expect( - resolveTinKeywordQuery(['kb-expired'], 'release notes', 'english', expired) + resolveTinKeywordQuery(true, 'release notes', 'english', expired) ).rejects.toBeInstanceOf(SearchDeadlineError) - expect(dbChainMockFns.select).not.toHaveBeenCalled() expect(dbChainMockFns.execute).not.toHaveBeenCalled() }) it('falls back to GIN instead of failing the search when readiness cannot be read', async () => { mockIsFeatureEnabled.mockRejectedValue(new Error('config unavailable')) - expect( - await resolveTinKeywordQuery(['kb-error'], 'release notes', 'english', undefined) - ).toBeNull() + expect(await resolveTinKeywordQuery(true, 'release notes', 'english', undefined)).toBeNull() }) }) diff --git a/apps/sim/lib/knowledge/search/tin-keyword.ts b/apps/sim/lib/knowledge/search/tin-keyword.ts index fe1e95d06cf..3ee8f48d903 100644 --- a/apps/sim/lib/knowledge/search/tin-keyword.ts +++ b/apps/sim/lib/knowledge/search/tin-keyword.ts @@ -1,7 +1,7 @@ -import { EMBEDDING_KEYWORD_TIN_INDEX, knowledgeBase } from '@sim/db/schema' +import { EMBEDDING_KEYWORD_TIN_INDEX } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' -import { inArray, sql } from 'drizzle-orm' +import { sql } from 'drizzle-orm' import { LRUCache } from 'lru-cache' import { isFeatureEnabled } from '@/lib/core/config/feature-flags' import { runSearchQuery, type SearchBudget } from '@/lib/knowledge/search/budget' @@ -14,7 +14,6 @@ const logger = createLogger('TinKeywordSearch') * backfilled, and flipping the flag off takes effect within this window. */ const READINESS_TTL_MS = 60 * 1000 -const SEARCH_INDEX_TTL_MS = 10 * 60 * 1000 /** * Whether the Tin index exists and finished building, i.e. the projection is complete. The read @@ -34,37 +33,11 @@ const indexReadiness = new LRUCache<'index', boolean, SearchBudget | undefined>( }) /** - * Only organization search indexes are projected. `is_search_index` is only ever turned on, when a - * legacy base is adopted, so a stale answer just keeps that base on the GIN projection for one TTL. - */ -const searchIndexBases = new LRUCache({ - max: 10_000, - ttl: SEARCH_INDEX_TTL_MS, -}) - -async function allSearchIndexes( - knowledgeBaseIds: readonly string[], - budget: SearchBudget | undefined -): Promise { - const unknown = knowledgeBaseIds.filter((id) => searchIndexBases.get(id) === undefined) - if (unknown.length > 0) { - const rows = await runSearchQuery(budget, 'keyword.tin_readiness', (executor) => - executor - .select({ id: knowledgeBase.id, isSearchIndex: knowledgeBase.isSearchIndex }) - .from(knowledgeBase) - .where(inArray(knowledgeBase.id, unknown)) - ) - for (const row of rows) searchIndexBases.set(row.id, row.isSearchIndex) - } - return knowledgeBaseIds.every((id) => searchIndexBases.get(id) === true) -} - -/** - * The TINQL query that ranks `query` inside `knowledgeBaseIds`, or null when keyword search must - * keep the GIN projection: the rollout flag is off, the database has no complete Tin index, a - * base is not an organization search index (only those are projected), or the query uses a shape - * TINQL cannot express. The text is analyzed by the same `websearch_to_tsquery` the GIN path - * uses, so both engines match the same stemmed terms. + * The TINQL query that ranks `query` inside the search's bases, or null when keyword search must + * keep the GIN projection: a base is not an organization search index (only those are + * projected), the rollout flag is off, the database has no complete Tin index, or the query uses + * a shape TINQL cannot express. The text is analyzed by the same `websearch_to_tsquery` the GIN + * path uses, so both engines match the same stemmed terms. * * Every read runs under the keyword leg's `budget`, so deciding the engine cannot outlast the leg's * deadline. A read that fails for another reason, including a shared cache read cut short by @@ -72,16 +45,19 @@ async function allSearchIndexes( * cancellation propagates like any other keyword query's. */ export async function resolveTinKeywordQuery( - knowledgeBaseIds: readonly string[], + searchIndexOnly: boolean, query: string, ftsConfig: string, budget: SearchBudget | undefined ): Promise { - if (knowledgeBaseIds.length === 0) return null + if (!searchIndexOnly) return null try { - if (!(await isFeatureEnabled('knowledge-tin-keyword'))) return null - if (!(await indexReadiness.fetch('index', { context: budget }))) return null - if (!(await allSearchIndexes(knowledgeBaseIds, budget))) return null + /** The flag and the index are independent facts; the search waits for the slower one only. */ + const [enabled, ready] = await Promise.all([ + isFeatureEnabled('knowledge-tin-keyword'), + indexReadiness.fetch('index', { context: budget }), + ]) + if (!enabled || !ready) return null const [{ rendered }] = await runSearchQuery(budget, 'keyword.tin_query', (executor) => executor.execute<{ rendered: string }>( sql`SELECT websearch_to_tsquery(${ftsConfig}::regconfig, ${query})::text AS rendered` diff --git a/apps/sim/scripts/backfill-projection-source-acl.ts b/apps/sim/scripts/backfill-projection-source-acl.ts index def5181c49c..be982412df6 100644 --- a/apps/sim/scripts/backfill-projection-source-acl.ts +++ b/apps/sim/scripts/backfill-projection-source-acl.ts @@ -25,8 +25,8 @@ const logger = createLogger('BackfillProjectionSourceAcl') /** A script has no long-lived process to detach into, so without a worker it fills inline. */ async function main(): Promise { if (isTriggerDevEnabled && env.TRIGGER_SECRET_KEY) { - const handle = await enqueueProjectionSourceAclBackfill({}, true) - logger.info('Backfill enqueued on the Trigger.dev worker', handle ?? {}) + const handle = await enqueueProjectionSourceAclBackfill() + logger.info('Backfill enqueued on the Trigger.dev worker', handle) return } await runProjectionSourceAclBackfill({}) From 166a6d1c36d5f676d8ffafe07ccc523176891cb0 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 21 Sep 2026 00:51:07 -0700 Subject: [PATCH 2/2] fix(knowledge): remember a saturated reach only once its count succeeded, and read the Tin flag before the index --- apps/sim/lib/knowledge/search/queries.test.ts | 32 +++++++++++++++++++ apps/sim/lib/knowledge/search/queries.ts | 3 +- .../lib/knowledge/search/tin-keyword.test.ts | 6 +--- apps/sim/lib/knowledge/search/tin-keyword.ts | 9 ++---- 4 files changed, 38 insertions(+), 12 deletions(-) diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 7ebc4a69797..930021250c3 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -1909,6 +1909,38 @@ describe('permitted-document planner', () => { expect(reachCounts()).toHaveLength(2) }) + it('does not remember a saturated reach whose count ran out of time', async () => { + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query).sql + if (isProbeStatement(statement)) return [{ id: null, connectorId: null, saturated: true }] + if (statement.includes('EXPLAIN')) + return [{ 'QUERY PLAN': [{ Plan: { 'Plan Rows': 1_000_000 } }] }] + if (statement.includes(') reached')) + throw Object.assign(new Error('canceling statement due to statement timeout'), { + code: '57014', + }) + return [] + }) + const reachCounts = () => statements().filter((query) => query.sql.includes(') reached')) + const budget = () => new SearchBudget('vector', performance.now() + 10_000) + expect( + await resolvePermittedDocuments({ + knowledgeBaseIds: ['org-index'], + access: scope('timed-saturated'), + budget: budget(), + }) + ).toEqual({ kind: 'unbounded', broad: true }) + expect(reachCounts()).toHaveLength(1) + /** The next search probes and counts again rather than trusting a reach that was never measured. */ + await resolvePermittedDocuments({ + knowledgeBaseIds: ['org-index'], + access: scope('timed-saturated'), + budget: budget(), + }) + expect(probes()).toBe(2) + expect(reachCounts()).toHaveLength(2) + }) + it('counts a resolved reach against a small index instead of assuming it broad', async () => { /** A bound inside the probe limit proves nothing without a saturated probe. */ dbChainMockFns.execute.mockImplementation(async (query) => { diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 3f0e4092cca..0798f7bacf2 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -1390,10 +1390,11 @@ export async function resolvePermittedDocuments(params: { params.accessPlan, true ) + if (key) saturatedReach.set(key, { broad }) } catch (error) { if (!params.budget?.isTimeout(error)) throw error + /** A count that ran out of time decides this search only; the next one counts again. */ } - if (key) saturatedReach.set(key, { broad }) } } const permitted: PermittedDocuments = diff --git a/apps/sim/lib/knowledge/search/tin-keyword.test.ts b/apps/sim/lib/knowledge/search/tin-keyword.test.ts index 3e84779f978..ac82444c449 100644 --- a/apps/sim/lib/knowledge/search/tin-keyword.test.ts +++ b/apps/sim/lib/knowledge/search/tin-keyword.test.ts @@ -36,11 +36,7 @@ describe('resolveTinKeywordQuery', () => { it('is off while the rollout flag is off, without touching the database', async () => { mockIsFeatureEnabled.mockResolvedValue(false) expect(await resolveTinKeywordQuery(true, 'release notes', 'english', undefined)).toBeNull() - expect( - dbChainMockFns.execute.mock.calls.some(([query]) => - JSON.stringify(query).includes('websearch_to_tsquery') - ) - ).toBe(false) + expect(dbChainMockFns.execute).not.toHaveBeenCalled() }) it('translates the analyzed query when every base is a search index', async () => { diff --git a/apps/sim/lib/knowledge/search/tin-keyword.ts b/apps/sim/lib/knowledge/search/tin-keyword.ts index 3ee8f48d903..534cb42f094 100644 --- a/apps/sim/lib/knowledge/search/tin-keyword.ts +++ b/apps/sim/lib/knowledge/search/tin-keyword.ts @@ -52,12 +52,9 @@ export async function resolveTinKeywordQuery( ): Promise { if (!searchIndexOnly) return null try { - /** The flag and the index are independent facts; the search waits for the slower one only. */ - const [enabled, ready] = await Promise.all([ - isFeatureEnabled('knowledge-tin-keyword'), - indexReadiness.fetch('index', { context: budget }), - ]) - if (!enabled || !ready) return null + /** The flag is read from memory and decides whether the index is worth asking about at all. */ + if (!(await isFeatureEnabled('knowledge-tin-keyword'))) return null + if (!(await indexReadiness.fetch('index', { context: budget }))) return null const [{ rendered }] = await runSearchQuery(budget, 'keyword.tin_query', (executor) => executor.execute<{ rendered: string }>( sql`SELECT websearch_to_tsquery(${ftsConfig}::regconfig, ${query})::text AS rendered`