diff --git a/apps/sim/lib/copilot/request/tools/executor.test.ts b/apps/sim/lib/copilot/request/tools/executor.test.ts index c30ca83db03..8b5415a44b8 100644 --- a/apps/sim/lib/copilot/request/tools/executor.test.ts +++ b/apps/sim/lib/copilot/request/tools/executor.test.ts @@ -146,6 +146,42 @@ function buildPendingToolCall(): ToolCallState { } } +describe('tool result size diagnostics', () => { + beforeEach(() => { + vi.clearAllMocks() + completeAsyncToolCall.mockResolvedValue(null) + markAsyncToolRunning.mockResolvedValue(null) + upsertAsyncToolCall.mockResolvedValue(null) + }) + + it.each(['é🔎', { content: 'é🔎' }])( + 'records UTF-8 bytes after result projection for %j', + async (output) => { + executeTool.mockResolvedValueOnce({ success: true, output }) + const toolCall = buildPendingToolCall() + const context = buildStreamingContext(toolCall) + const endSpan = vi.spyOn(context.trace, 'endSpan') + + const completion = await executeToolAndReport(toolCall.id, context, { + userId: 'user-1', + workflowId: 'workflow-1', + resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry(), + }) + + expect(completion.status).toBe(MothershipStreamV1ToolOutcome.success) + const serialized = + typeof completion.data === 'string' ? completion.data : JSON.stringify(completion.data) + expect(endSpan).toHaveBeenCalledWith( + expect.objectContaining({ + kind: 'tool.execute', + attributes: expect.objectContaining({ outputBytes: Buffer.byteLength(serialized) }), + }), + 'ok' + ) + } + ) +}) + describe('toolWatchdogTimeoutMs', () => { it('gives request-scoped MCP tools the long-running watchdog', () => { expect(toolWatchdogTimeoutMs('mcp-363de040-web_search_exa')).toBe(TOOL_WATCHDOG_LONG_RUNNING_MS) diff --git a/apps/sim/lib/copilot/request/tools/executor.ts b/apps/sim/lib/copilot/request/tools/executor.ts index d9c6b523e9d..e7c03b320e0 100644 --- a/apps/sim/lib/copilot/request/tools/executor.ts +++ b/apps/sim/lib/copilot/request/tools/executor.ts @@ -125,11 +125,11 @@ function summarizeToolResultForSpan(result: { const output = (result as { output: unknown }).output if (typeof output === 'string') { summary.outputKind = 'string' - summary.outputBytes = output.length + summary.outputBytes = Buffer.byteLength(output) } else if (output && typeof output === 'object') { summary.outputKind = Array.isArray(output) ? 'array' : 'object' try { - summary.outputBytes = JSON.stringify(output).length + summary.outputBytes = Buffer.byteLength(JSON.stringify(output)) } catch { summary.outputBytes = 0 } @@ -143,7 +143,7 @@ function summarizeToolResultForSpan(result: { } } else if (output !== undefined && output !== null) { summary.outputKind = typeof output - summary.outputBytes = String(output).length + summary.outputBytes = Buffer.byteLength(String(output)) } return summary } diff --git a/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.test.ts b/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.test.ts index 91747f3629a..1e669a7b45b 100644 --- a/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.test.ts +++ b/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.test.ts @@ -1,7 +1,15 @@ /** @vitest-environment node */ import { beforeEach, describe, expect, it, vi } from 'vitest' -const mocks = vi.hoisted(() => ({ search: vi.fn(), read: vi.fn(), authorizeChat: vi.fn() })) +const mocks = vi.hoisted(() => ({ + search: vi.fn(), + read: vi.fn(), + authorizeChat: vi.fn(), + info: vi.fn(), +})) +vi.mock('@sim/logger', () => ({ + createLogger: () => ({ info: mocks.info, error: vi.fn(), warn: vi.fn() }), +})) vi.mock('@/lib/copilot/chat/organization-chats', () => ({ authorizeOrganizationChatDelegation: { execute: mocks.authorizeChat }, })) @@ -179,6 +187,48 @@ describe('Assistant retrieval tools', () => { }) ) }) + it.each([0, 20, 50])( + 'measures UTF-8 bytes for %i passages without logging their content', + async (count) => { + const content = 'Confidential passage é🔎'.repeat(100) + mocks.search.mockResolvedValueOnce({ + knowledgeBases: [{ id: 'index', name: 'Enterprise Search' }], + results: Array.from({ length: count }, (_, index) => ({ + knowledgeBaseId: 'index', + documentId: `doc-${index % 4}`, + documentName: 'Private title', + sourceUrl: null, + sourceModifiedAt: null, + metadata: {}, + content, + chunkIndex: index, + similarity: 1, + })), + }) + + const output = await searchWorkspaceServerTool.execute( + { query: 'Private query', ...(count === 50 ? { topK: 50 } : {}) }, + context + ) + + expect(output.success).toBe(true) + expect(mocks.info).toHaveBeenCalledWith( + 'Knowledge search completed', + expect.objectContaining({ + toolCallId: 'call', + toolResultBytes: Buffer.byteLength(JSON.stringify(output)), + passageBytes: count * Buffer.byteLength(content), + maxPassageBytes: count ? Buffer.byteLength(content) : 0, + uniqueDocumentCount: Math.min(count, 4), + }) + ) + const logged = JSON.stringify(mocks.info.mock.calls) + expect(logged).not.toContain('Confidential passage') + expect(logged).not.toContain('Private title') + expect(logged).not.toContain('Private query') + } + ) + it('returns stable citation IDs with internal links for uploaded documents', async () => { const result = await searchWorkspaceServerTool.execute({ query: 'orion' }, context) expect(result).toMatchObject({ diff --git a/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.ts b/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.ts index 4fad2478865..01f6cb833f5 100644 --- a/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.ts +++ b/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.ts @@ -16,6 +16,12 @@ import { } from '@/lib/knowledge/application/workspace-search' import { sourceAuthor } from '@/lib/knowledge/search/author' import { createKnowledgeDocumentCitation } from '@/lib/knowledge/search/citation' +import { + annotateSearchDiagnostics, + measureSearchStage, + recordSearchStageDuration, + withSearchDiagnostics, +} from '@/lib/knowledge/search/diagnostics' import { intersectWorkspaceSearchFilters } from '@/lib/knowledge/search/filters' import { connectorDisplayName } from '@/lib/sim-search/connectors' import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection' @@ -37,77 +43,99 @@ const CITATION_INSTRUCTION = export const searchWorkspaceServerTool: BaseServerTool = { name: 'search_workspace', async execute(raw, context?: ServerToolContext) { - try { - const scope = requireCopilotKnowledgeScope(context) - const { query, topK, ...requestedFilters } = searchInputSchema.parse(raw) - const registry = context?.resolvedSecretTraceRegistry - if (!registry) throw new Error('Knowledge result provenance is unavailable') - const projected = projectResolvedSecretModelContent(query, registry) - if (!projected.safe || typeof projected.value !== 'string') { - return { - success: false, - message: 'Search query contains protected content. Rephrase the query.', - } - } - const input = { - query: projected.value, - topK, - filters: intersectWorkspaceSearchFilters(requestedFilters, context?.assistantSearch), + return withSearchDiagnostics( + { surface: context?.searchSurface ?? 'copilot', - resultSecretRegistry: registry, - signal: context?.abortSignal, - } as const - const result = - scope.kind === 'organization' - ? await executeCopilotOrganizationKnowledgeUseCase(context, searchOrganizationKnowledge, { - ...input, - organizationId: scope.organizationId, - }) - : await executeCopilotKnowledgeUseCase(context, searchWorkspaceKnowledge, { - ...input, - workspaceId: scope.workspaceId, + toolCallId: context?.toolCallId, + executionId: context?.executionId, + }, + async () => { + try { + const inputStarted = performance.now() + const scope = requireCopilotKnowledgeScope(context) + const { query, topK, ...requestedFilters } = searchInputSchema.parse(raw) + const registry = context?.resolvedSecretTraceRegistry + if (!registry) throw new Error('Knowledge result provenance is unavailable') + const projected = projectResolvedSecretModelContent(query, registry) + if (!projected.safe || typeof projected.value !== 'string') { + return { + success: false, + message: 'Search query contains protected content. Rephrase the query.', + } + } + const input = { + query: projected.value, + topK, + filters: intersectWorkspaceSearchFilters(requestedFilters, context?.assistantSearch), + surface: context?.searchSurface ?? 'copilot', + resultSecretRegistry: registry, + signal: context?.abortSignal, + } as const + recordSearchStageDuration('tool_input', performance.now() - inputStarted) + const result = await measureSearchStage('tool_application', () => + scope.kind === 'organization' + ? executeCopilotOrganizationKnowledgeUseCase(context, searchOrganizationKnowledge, { + ...input, + organizationId: scope.organizationId, + }) + : executeCopilotKnowledgeUseCase(context, searchWorkspaceKnowledge, { + ...input, + workspaceId: scope.workspaceId, + }) + ) + return await measureSearchStage('tool_presentation', () => { + const names = new Map(result.knowledgeBases.map((base) => [base.id, base.name])) + const output = { + success: true, + message: `Found ${result.results.length} passages. ${CITATION_INSTRUCTION}`, + data: { + query, + results: result.results.map((item) => ({ + documentId: item.documentId, + knowledgeBaseId: item.knowledgeBaseId, + knowledgeBaseName: names.get(item.knowledgeBaseId) ?? '', + siteName: item.connectorType + ? connectorDisplayName(item.connectorType) + : names.get(item.knowledgeBaseId), + documentName: item.documentName, + sourceUrl: item.sourceUrl, + connectorType: item.connectorType, + sourceModifiedAt: item.sourceModifiedAt?.toISOString() ?? null, + author: sourceAuthor(item.metadata), + content: item.content, + chunkIndex: item.chunkIndex, + similarity: item.similarity, + ...createKnowledgeDocumentCitation({ + scope, + knowledgeBaseId: item.knowledgeBaseId, + documentId: item.documentId, + sourceUrl: item.sourceUrl, + baseUrl: getBaseUrl(), + }), + })), + }, + } + const passageBytes = output.data.results.map((item) => Buffer.byteLength(item.content)) + annotateSearchDiagnostics({ + toolResultBytes: Buffer.byteLength(JSON.stringify(output)), + passageBytes: passageBytes.reduce((total, bytes) => total + bytes, 0), + maxPassageBytes: Math.max(0, ...passageBytes), + uniqueDocumentCount: new Set(output.data.results.map((item) => item.documentId)).size, }) - const names = new Map(result.knowledgeBases.map((base) => [base.id, base.name])) - return { - success: true, - message: `Found ${result.results.length} passages. ${CITATION_INSTRUCTION}`, - data: { - query, - results: result.results.map((item) => ({ - documentId: item.documentId, - knowledgeBaseId: item.knowledgeBaseId, - knowledgeBaseName: names.get(item.knowledgeBaseId) ?? '', - siteName: item.connectorType - ? connectorDisplayName(item.connectorType) - : names.get(item.knowledgeBaseId), - documentName: item.documentName, - sourceUrl: item.sourceUrl, - connectorType: item.connectorType, - sourceModifiedAt: item.sourceModifiedAt?.toISOString() ?? null, - author: sourceAuthor(item.metadata), - content: item.content, - chunkIndex: item.chunkIndex, - similarity: item.similarity, - ...createKnowledgeDocumentCitation({ - scope, - knowledgeBaseId: item.knowledgeBaseId, - documentId: item.documentId, - sourceUrl: item.sourceUrl, - baseUrl: getBaseUrl(), - }), - })), - }, - } - } catch (error) { - logger.error('Workspace search failed', { error }) - return { - success: false, - message: - error instanceof z.ZodError - ? 'Invalid search arguments' - : messageForCopilotKnowledgeError(error), + return output + }) + } catch (error) { + logger.error('Workspace search failed', { error }) + return { + success: false, + message: + error instanceof z.ZodError + ? 'Invalid search arguments' + : messageForCopilotKnowledgeError(error), + } + } } - } + ) }, } diff --git a/apps/sim/lib/knowledge/__integration__/application-acl.integration.ts b/apps/sim/lib/knowledge/__integration__/application-acl.integration.ts index 48d30b587e3..209f2142ebb 100644 --- a/apps/sim/lib/knowledge/__integration__/application-acl.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/application-acl.integration.ts @@ -62,7 +62,10 @@ import { seedKnowledgeMemberFixture, } from '@/lib/knowledge/__integration__/seed-source-access-fixture' import { confluencePageAcl } from '@/lib/knowledge/access/confluence-permissions' -import { knowledgeAccessCondition } from '@/lib/knowledge/access/predicate' +import { + knowledgeAccessCondition, + knowledgeMetadataCandidateAccessCondition, +} from '@/lib/knowledge/access/predicate' import { createKnowledgeAccessProvider } from '@/lib/knowledge/access/scope' import { listKnowledgeChunks } from '@/lib/knowledge/application/chunks' import { readKnowledgeDocument } from '@/lib/knowledge/application/documents' @@ -299,9 +302,16 @@ describe('indexed source content through real application access', () => { return result.results.map((row) => row.documentId) } - it.each(['workspace', 'admin', 'members'] as const)( - 'allows a remaining workspace ACL only in workspace mode, not during a %s transition', - async (accessMode) => { + it.each([ + ['workspace', true], + ['workspace', false], + ['admin', true], + ['admin', false], + ['members', true], + ['members', false], + ] as const)( + 'requires settled workspace mode for a remaining workspace ACL (%s, rewrite pending=%s)', + async (accessMode, accessRewritePending) => { const [savedConnector] = await db .select({ accessMode: knowledgeConnector.accessMode, @@ -324,18 +334,25 @@ describe('indexed source content through real application access', () => { .where(eq(document.id, documentId)) await db .update(knowledgeConnector) - .set({ accessMode, accessRewritePending: true }) + .set({ accessMode, accessRewritePending }) .where(eq(knowledgeConnector.id, connectorId)) - const visible = await db - .select({ id: document.id }) - .from(document) - .where( - and( - eq(document.id, documentId), - knowledgeAccessCondition({ kind: 'workspace', tokens: ['pub', 'ws'] }) + for (const accessCondition of [ + knowledgeMetadataCandidateAccessCondition, + knowledgeAccessCondition, + ]) { + const visible = await db + .select({ id: document.id }) + .from(document) + .where( + and( + eq(document.id, documentId), + accessCondition({ kind: 'workspace', tokens: ['pub', 'ws'] }) + ) ) + expect(visible.map((row) => row.id)).toEqual( + accessMode === 'workspace' && !accessRewritePending ? [documentId] : [] ) - expect(visible.map((row) => row.id)).toEqual(accessMode === 'workspace' ? [documentId] : []) + } } finally { await db .update(knowledgeConnector) diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts new file mode 100644 index 00000000000..5646bd707d1 --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -0,0 +1,558 @@ +/** Real Assistant tool, application authorization, PostgreSQL/pgvector, and result processing. */ +import { readFileSync, statSync, writeFileSync } from 'node:fs' +import { db } from '@sim/db' +import { + copilotChats, + credential, + credentialGroup, + document, + knowledgeBase, + knowledgeConnector, + member, + organization, + user, + workspace, +} from '@sim/db/schema' +import { createLogger, Logger } from '@sim/logger' +import { generateId } from '@sim/utils/id' +import { and, eq, inArray, sql } from 'drizzle-orm' +import { afterAll, beforeAll, describe, expect, it, type MockInstance, vi } from 'vitest' +import { z } from 'zod' +import { searchWorkspaceServerTool } from '@/lib/copilot/tools/server/knowledge/workspace-search' +import { seedSearchReaderFixture } from '@/lib/knowledge/__integration__/seed-search-reader-fixture' +import { + createKnowledgeAclFixtureIds, + seedKnowledgeAclFixture, +} from '@/lib/knowledge/__integration__/seed-source-access-fixture' +import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-search' +import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters' +import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' + +/** Initialize controlled provider configuration before the real application modules load. */ +vi.hoisted(() => { + if (process.env.KNOWLEDGE_SEARCH_PERFORMANCE_TEST === 'true') { + Object.assign(process.env, { + OPENAI_API_KEY: 'isolated-embedding-http-fixture', + CONFLUENCE_CLIENT_ID: 'isolated-confluence-fixture-client', + CONFLUENCE_CLIENT_SECRET: 'isolated-confluence-fixture-secret', + }) + } +}) + +const enabled = process.env.KNOWLEDGE_SEARCH_PERFORMANCE_TEST === 'true' +const chunkCount = Number(process.env.KNOWLEDGE_SEARCH_PERFORMANCE_CHUNKS ?? 20_000) +const dimensions = 1536 +const chunksPerDocument = 4 +const batchSize = 1000 +const logger = createLogger('SearchLatencyIntegration') +const fixtureSchema = z.object({ + aliceId: z.uuid(), + bobId: z.uuid(), + workspaceId: z.uuid(), + organizationId: z.uuid(), + knowledgeBaseId: z.uuid(), + connectorId: z.uuid(), + lockId: z.uuid(), + groups: z.array(z.string()).length(3), + groupIds: z.array(z.uuid()).length(3), +}) +const reuseFile = enabled ? process.env.KNOWLEDGE_SEARCH_PERFORMANCE_REUSE_REPORT_FILE : undefined +function readFixtureReport(file: string) { + if (statSync(file).size > 8 * 1024 * 1024) throw new Error('Fixture report exceeds 8 MiB') + return z + .object({ fixture: fixtureSchema, unrelatedFixture: fixtureSchema }) + .parse(JSON.parse(readFileSync(file, 'utf8'))) +} +const reused = reuseFile ? readFixtureReport(reuseFile) : undefined +const ids = reused?.fixture ?? createKnowledgeAclFixtureIds() +const unrelated = reused?.unrelatedFixture ?? createKnowledgeAclFixtureIds() +const organizationChatId = generateId() +const queryVector = Array.from({ length: dimensions }, (_, index) => (index === 0 ? 1 : 0)) +const captured: CapturedQuery[] = [] +const report: Record = { + fixture: ids, + unrelatedFixture: unrelated, + method: { + chunkCount, + dimensions, + chunksPerDocument, + sql: 'Captured from the real Assistant tool; no hand-written search query', + providers: + 'Embedding and source-permission HTTP responses are controlled; internal search and authorization code is real', + vectors: + 'Normalized topic clusters with deterministic dense noise; not semantic-quality evaluation', + cache: 'First and repeated samples; no claim of a cold operating-system cache', + }, +} +let capture = false +let embeddingCalls = 0 +let readerCalls = 0 +let readerRevoked = false +const previousDebug = db.$client.options.debug +let diagnosticLog: MockInstance | undefined + +interface CapturedQuery { + query: string + parameters: NonNullable[1]> +} + +interface ExplainNode { + 'Node Type': string + 'Actual Rows': number + 'Index Name'?: string + Plans?: ExplainNode[] +} + +const explainNodeSchema: z.ZodType = z.lazy(() => + z + .object({ + 'Node Type': z.string(), + 'Actual Rows': z.number(), + 'Index Name': z.string().optional(), + Plans: z.array(explainNodeSchema).optional(), + }) + .passthrough() +) +const explainSchema = z.array(z.object({ Plan: explainNodeSchema }).passthrough()).length(1) + +function usesVectorIndex(node: ExplainNode): boolean { + return ( + node['Index Name'] === 'embedding_vector_hnsw_idx' || + (node.Plans?.some(usesVectorIndex) ?? false) + ) +} + +function saveReport() { + const file = process.env.KNOWLEDGE_SEARCH_PERFORMANCE_REPORT_FILE + if (file) writeFileSync(file, JSON.stringify(report, null, 2), { mode: 0o600 }) +} + +const diagnosticSchema = z + .object({ + surface: z.enum(['dashboard', 'copilot']), + outcome: z.literal('success'), + elapsedMs: z.number(), + toolResultBytes: z.number().int().nonnegative().optional(), + passageBytes: z.number().int().nonnegative().optional(), + maxPassageBytes: z.number().int().nonnegative().optional(), + uniqueDocumentCount: z.number().int().nonnegative().optional(), + stages: z.record( + z.string(), + z.object({ + count: z.number(), + totalMs: z.number(), + maxMs: z.number(), + errors: z.number(), + }) + ), + }) + .passthrough() + +const resultSchema = z.object({ + success: z.literal(true), + data: z.object({ + results: z.array( + z.object({ documentId: z.string(), content: z.string(), knowledgeBaseId: z.string() }) + ), + }), +}) + +async function search( + userId = ids.aliceId, + query = 'Orion deployment', + filters: WorkspaceSearchFilters = {}, + organizationScope = false +) { + return resultSchema.parse( + await searchWorkspaceServerTool.execute( + { query, topK: 15, ...filters }, + { + userId, + ...(organizationScope + ? { organizationId: ids.organizationId, chatId: organizationChatId } + : { workspaceId: ids.workspaceId }), + requestMode: 'assistant', + toolCallId: generateId(), + copilotToolExecution: true, + resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry([], { + userId, + ...(organizationScope ? {} : { workspaceId: ids.workspaceId }), + }), + } + ) + ) +} + +async function sample(label: string, run: () => ReturnType) { + captured.length = 0 + diagnosticLog?.mockClear() + const start = performance.now() + capture = true + let result: Awaited> + try { + result = await run() + } finally { + capture = false + } + const milliseconds = performance.now() - start + const completed = + diagnosticLog?.mock.calls.filter(([message]) => message === 'Knowledge search completed') ?? [] + expect(completed).toHaveLength(1) + const diagnostics = diagnosticSchema.parse(completed[0][1]) + expect(diagnostics.stages.embedding.count).toBe(1) + expect(diagnostics.stages.retrieval.count).toBe(1) + if (diagnostics.surface === 'copilot') { + const passageBytes = result.data.results.map((row) => Buffer.byteLength(row.content)) + expect(diagnostics.passageBytes).toBe(passageBytes.reduce((total, bytes) => total + bytes, 0)) + expect(diagnostics.maxPassageBytes).toBe(Math.max(0, ...passageBytes)) + expect(diagnostics.uniqueDocumentCount).toBe( + new Set(result.data.results.map((row) => row.documentId)).size + ) + expect(diagnostics.toolResultBytes).toBeGreaterThan(diagnostics.passageBytes!) + } + expect(captured.length).toBeLessThan(300) + const searches = captured.filter( + (item) => + item.query.includes('from "embedding"') && + (item.query.includes('order by') || item.query.includes('limit')) + ) + const plans = [] + for (const query of searches) { + const plan = await db.$client.begin(async (tx) => { + await tx.unsafe("SET LOCAL hnsw.iterative_scan = 'relaxed_order'") + await tx.unsafe('SET LOCAL hnsw.max_scan_tuples = 20000') + return tx.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${query.query}`, query.parameters) + }) + plans.push({ + kind: query.query.includes('keyword_rank') + ? 'keyword' + : query.query.includes('order by') + ? 'vector' + : 'probe', + query: query.query, + parameters: query.parameters, + plan: explainSchema.parse(plan[0]['QUERY PLAN']), + }) + } + report[label] = { + milliseconds, + diagnostics, + queryCount: captured.length, + resultCount: result.data.results.length, + plans, + } + saveReport() + logger.info(label, { + milliseconds, + queryCount: captured.length, + resultCount: result.data.results.length, + }) + return { result, plans, diagnostics } +} + +describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpus', () => { + beforeAll(async () => { + if ( + !Number.isInteger(chunkCount) || + chunkCount < 10_000 || + chunkCount > 200_000 || + chunkCount % batchSize !== 0 + ) + throw new Error( + 'KNOWLEDGE_SEARCH_PERFORMANCE_CHUNKS must be a multiple of 1000 from 10000 to 200000' + ) + vi.stubGlobal('fetch', async (input: string | URL | Request, init?: RequestInit) => { + const url = input instanceof Request ? input.url : String(input) + if ( + url === 'https://api.atlassian.com/ex/confluence/search-fixture/wiki/rest/api/user/current' + ) { + expect(new Headers(init?.headers).get('Authorization')).toBe('Bearer fixture-search-reader') + readerCalls++ + return readerRevoked + ? new Response(null, { status: 403 }) + : Response.json({ type: 'known', accountId: ids.aliceId }) + } + if (url !== 'https://api.openai.com/v1/embeddings') + throw new Error(`Unexpected outbound request in search fixture: ${new URL(url).origin}`) + const body = z + .object({ input: z.array(z.string()).length(1), encoding_format: z.literal('base64') }) + .parse(JSON.parse(String(init?.body))) + embeddingCalls += body.input.length + const bytes = Buffer.alloc(dimensions * 4) + queryVector.forEach((value, index) => bytes.writeFloatLE(value, index * 4)) + return Response.json({ + data: [{ embedding: bytes.toString('base64') }], + usage: { total_tokens: 4 }, + }) + }) + if (reused) { + const owners = await db + .select({ id: workspace.id, ownerId: workspace.ownerId }) + .from(workspace) + .where(inArray(workspace.id, [ids.workspaceId, unrelated.workspaceId])) + expect(owners).toEqual( + expect.arrayContaining([ + { id: ids.workspaceId, ownerId: ids.aliceId }, + { id: unrelated.workspaceId, ownerId: unrelated.aliceId }, + ]) + ) + for (const fixture of [ids, unrelated]) { + const [size] = await db.execute<{ count: number }>( + sql`SELECT count(*)::int AS count FROM embedding WHERE knowledge_base_id = ${fixture.knowledgeBaseId}` + ) + expect(size.count).toBe(fixture === ids ? chunkCount : chunkCount / 2) + } + await db + .update(knowledgeBase) + .set({ workspaceId: ids.workspaceId, organizationId: null }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + await db + .update(knowledgeConnector) + .set({ connectorType: 'google_drive', credentialId: null, sourceConfig: {} }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + await db.delete(credential).where(eq(credential.workspaceId, ids.workspaceId)) + await db.delete(credentialGroup).where(eq(credentialGroup.workspaceId, ids.workspaceId)) + await db + .delete(copilotChats) + .where( + and( + eq(copilotChats.organizationId, ids.organizationId), + eq(copilotChats.userId, ids.aliceId) + ) + ) + await db + .delete(member) + .where(and(eq(member.organizationId, ids.organizationId), eq(member.userId, ids.aliceId))) + await db.execute( + sql`UPDATE document SET acl = ARRAY[${`u:${ids.aliceId}@fixture.test`}], user_excluded = false, acl_verified_at = statement_timestamp() WHERE knowledge_base_id = ${ids.knowledgeBaseId}` + ) + } else { + await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' }) + await seedKnowledgeAclFixture(unrelated, { connectorType: 'google_drive' }) + await db + .update(knowledgeBase) + .set({ isSearchIndex: true }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + const indexes = await db.execute<{ + indexname: string + indexdef: string + }>(sql`SELECT indexname, indexdef FROM pg_indexes + WHERE tablename = 'embedding' AND indexdef LIKE '% USING hnsw %'`) + for (const index of indexes) + await db.execute(sql`DROP INDEX ${sql.identifier(index.indexname)}`) + for (const fixture of [ids, unrelated]) { + const count = fixture === ids ? chunkCount : chunkCount / 2 + for (let first = 0; first < count / chunksPerDocument; first += batchSize) { + const last = Math.min(first + batchSize, count / chunksPerDocument) - 1 + await db.execute(sql`INSERT INTO document + (id, knowledge_base_id, connector_id, external_id, filename, file_url, file_size, mime_type, processing_status, acl, acl_verified_at) + SELECT ${fixture.workspaceId} || '-doc-' || n, ${fixture.knowledgeBaseId}, ${fixture.connectorId}, n::text, + 'Deployment guide ' || n, 'https://fixture.invalid/document/' || n, 12000, 'text/plain', 'completed', + ARRAY[${`u:${fixture.aliceId}@fixture.test`}]::text[], statement_timestamp() + FROM generate_series(${first}::int, ${last}::int) n`) + } + for (let first = 0; first < count; first += batchSize) { + const last = Math.min(first + batchSize, count) - 1 + await db.transaction(async (tx) => { + await tx.execute(sql`SET LOCAL jit = off`) + await tx.execute(sql`INSERT INTO embedding + (id, knowledge_base_id, document_id, chunk_index, chunk_hash, content, content_length, token_count, start_offset, end_offset, embedding) + SELECT ${fixture.workspaceId} || '-chunk-' || n, ${fixture.knowledgeBaseId}, ${fixture.workspaceId} || '-doc-' || (n / ${chunksPerDocument}), + n % ${chunksPerDocument}, 'hash-' || n, + CASE WHEN n % 8 = 0 THEN 'Orion deployment reference. ' ELSE 'Engineering operations reference. ' END || + (SELECT string_agg(md5(n::text || ':' || paragraph::text), ' ') FROM generate_series(1, 90) paragraph), + 3000, 750, 0, 3000, + l2_normalize(ARRAY(SELECT (CASE WHEN coordinate = n % 32 + 1 THEN 1 ELSE 0 END + + 0.025 * sin(n::double precision * coordinate * 12.9898 + coordinate * 78.233))::real + FROM generate_series(1, ${dimensions}) coordinate)::vector(1536)) + FROM generate_series(${first}::int, ${last}::int) n`) + }) + } + logger.info('Synthetic corpus loaded', { chunks: count }) + } + for (const index of indexes) await db.execute(sql.raw(index.indexdef)) + } + await db.execute(sql`ANALYZE document`) + await db.execute(sql`ANALYZE embedding`) + report.server = ( + await db.execute(sql`SELECT version(), current_setting('work_mem') AS work_mem, + (SELECT extversion FROM pg_extension WHERE extname = 'vector') AS pgvector`) + )[0] + db.$client.options.debug = (_connection, query, parameters) => { + if (capture && captured.length < 300) captured.push({ query, parameters: [...parameters] }) + } + diagnosticLog = vi.spyOn(Logger.prototype, 'info') + }, 60 * 60_000) + + afterAll(async () => { + diagnosticLog?.mockRestore() + db.$client.options.debug = previousDebug + vi.unstubAllGlobals() + saveReport() + for (const fixture of process.env.KNOWLEDGE_SEARCH_PERFORMANCE_KEEP_DATABASE === 'true' + ? [] + : [ids, unrelated]) { + await db.delete(workspace).where(eq(workspace.id, fixture.workspaceId)) + await db.delete(organization).where(eq(organization.id, fixture.organizationId)) + await db.delete(user).where(eq(user.id, fixture.aliceId)) + await db.delete(user).where(eq(user.id, fixture.bobId)) + } + await db.$client.end() + }, 120_000) + + it('records first and repeated application searches with the actual SQL plans', async () => { + for (let iteration = 0; iteration < 2; iteration++) { + const { result, plans } = await sample(`broad.${iteration}`, () => search()) + expect(result.data.results).toHaveLength(15) + expect(result.data.results.every((row) => row.knowledgeBaseId === ids.knowledgeBaseId)).toBe( + true + ) + expect(plans.length).toBeGreaterThanOrEqual(2) + const vectorPlans = plans.filter((plan) => plan.kind === 'vector') + expect(vectorPlans).toHaveLength(1) + expect(usesVectorIndex(vectorPlans[0].plan[0].Plan)).toBe(true) + } + expect(embeddingCalls).toBe(2) + }, 180_000) + + it('compares the Search tab and Assistant with the same person, query and index', async () => { + const dashboard = await sample('dashboard', async () => { + const result = await searchScopedKnowledge.execute({ + principal: { kind: 'session', userId: ids.aliceId, sessionId: 'fixture-dashboard' }, + input: { + workspaceId: ids.workspaceId, + query: 'Orion deployment', + topK: 15, + surface: 'dashboard', + }, + }) + return resultSchema.parse({ success: true, data: result }) + }) + const assistant = await sample('assistant.comparison', () => search()) + expect(dashboard.diagnostics.surface).toBe('dashboard') + expect(assistant.diagnostics.surface).toBe('copilot') + expect(dashboard.diagnostics.stages.result_provenance).toBeUndefined() + expect(assistant.diagnostics.stages.result_provenance.count).toBe(1) + expect(dashboard.result.data.results).toHaveLength(15) + expect(assistant.result.data.results).toHaveLength(15) + const dashboardVector = dashboard.plans.filter((plan) => plan.kind === 'vector') + const assistantVector = assistant.plans.filter((plan) => plan.kind === 'vector') + expect(dashboardVector).toHaveLength(1) + expect(assistantVector).toHaveLength(1) + expect(dashboardVector[0].query).toBe(assistantVector[0].query) + expect(dashboardVector[0].parameters).toEqual(assistantVector[0].parameters) + expect(usesVectorIndex(dashboardVector[0].plan[0].Plan)).toBe(true) + }, 180_000) + + it('keeps inaccessible content out of an otherwise identical search', async () => { + const { result } = await sample('denied', () => search(ids.bobId)) + expect(result.data.results).toEqual([]) + }, 180_000) + + it('ranks a small permission scope by its bounded IDs without a corpus-wide vector probe', async () => { + const documentIds = [0, 8, 16].map((index) => `${ids.workspaceId}-doc-${index}`) + await db + .update(document) + .set({ acl: [`u:${ids.aliceId}@fixture.test`, `u:${ids.bobId}@fixture.test`] }) + .where(inArray(document.id, documentIds)) + try { + const { result, plans } = await sample('small-scope', () => search(ids.bobId)) + expect(result.data.results.length).toBeGreaterThan(0) + expect(result.data.results.every((row) => documentIds.includes(row.documentId))).toBe(true) + const probe = plans.filter((plan) => plan.kind === 'probe') + expect(probe).toHaveLength(1) + expect(probe[0].query).not.toContain('<=>') + expect(probe[0].plan[0].Plan['Actual Rows']).toBe(12) + const vector = plans.filter((plan) => plan.kind === 'vector') + expect(vector).toHaveLength(1) + expect(vector[0].query).toContain('"embedding"."id" in') + } finally { + await db + .update(document) + .set({ acl: [`u:${ids.aliceId}@fixture.test`] }) + .where(inArray(document.id, documentIds)) + } + }, 180_000) + + it('applies selective document scope and exclusion before ranking', async () => { + const documentIds = [0, 8, 16, 24, 32].map((index) => `${ids.workspaceId}-doc-${index}`) + await db.update(document).set({ userExcluded: true }).where(eq(document.id, documentIds[0])) + try { + const { result } = await sample('selective', () => + search(ids.aliceId, 'Orion deployment', { documentIds }) + ) + expect(result.data.results.length).toBeGreaterThan(0) + expect(new Set(result.data.results.map((row) => row.documentId))).toEqual( + new Set(documentIds.slice(1)) + ) + expect( + result.data.results.every((row) => documentIds.slice(1).includes(row.documentId)) + ).toBe(true) + } finally { + await db.update(document).set({ userExcluded: false }).where(eq(document.id, documentIds[0])) + } + }, 180_000) + + it('runs two independent Assistant searches concurrently', async () => { + const start = performance.now() + const results = await Promise.all([search(), search(ids.aliceId, 'Engineering operations')]) + report.concurrent = { + milliseconds: performance.now() - start, + resultCounts: results.map((result) => result.data.results.length), + } + saveReport() + for (const result of results) expect(result.data.results).toHaveLength(15) + }, 180_000) + + it('checks live reader access on every search, including after revocation', async () => { + await seedSearchReaderFixture(ids) + const allowed = await sample('live.allowed', () => search()) + expect(allowed.result.data.results).toHaveLength(15) + expect(readerCalls).toBeGreaterThan(0) + readerRevoked = true + const before = readerCalls + const denied = await sample('live.revoked', () => search()) + for (const probe of denied.plans.filter((plan) => plan.kind === 'probe')) { + expect(probe.query).not.toContain('<=>') + } + expect(denied.result.data.results).toEqual([]) + expect(readerCalls).toBeGreaterThan(before) + readerRevoked = false + const restored = await sample('live.restored', () => search()) + expect(restored.result.data.results).toHaveLength(15) + }, 180_000) + + it('uses the same indexed retrieval through a persisted private organization Assistant chat', async () => { + await db.insert(member).values({ + id: generateId(), + organizationId: ids.organizationId, + userId: ids.aliceId, + role: 'owner', + }) + await db.insert(copilotChats).values({ + id: organizationChatId, + organizationId: ids.organizationId, + userId: ids.aliceId, + type: 'mothership', + }) + await db + .update(knowledgeBase) + .set({ workspaceId: null, organizationId: ids.organizationId }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + await db + .update(knowledgeConnector) + .set({ connectorType: 'google_drive', credentialId: null, sourceConfig: {} }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + await db.execute( + sql`UPDATE document SET acl = ARRAY[${`u:${ids.aliceId}@fixture.test`}] WHERE knowledge_base_id = ${ids.knowledgeBaseId}` + ) + await db.execute(sql`ANALYZE document`) + const { result } = await sample('organization', () => + search(ids.aliceId, 'Orion deployment', {}, true) + ) + expect(result.data.results).toHaveLength(15) + expect(result.data.results.every((row) => row.knowledgeBaseId === ids.knowledgeBaseId)).toBe( + true + ) + }, 180_000) +}) diff --git a/apps/sim/lib/knowledge/__integration__/seed-search-reader-fixture.ts b/apps/sim/lib/knowledge/__integration__/seed-search-reader-fixture.ts new file mode 100644 index 00000000000..8896912a40c --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/seed-search-reader-fixture.ts @@ -0,0 +1,107 @@ +import { createHash } from 'node:crypto' +import { db } from '@sim/db' +import { + credential, + credentialGroup, + credentialGroupEnrollment, + knowledgeConnector, +} from '@sim/db/schema' +import { generateId } from '@sim/utils/id' +import { eq, sql } from 'drizzle-orm' +import { encryptSecret } from '@/lib/core/security/encryption' +import { getCredentialGroupProviderAdapterByProviderId } from '@/lib/credential-groups/provider-registry' +import { encryptManagedOAuthTokenSet } from '@/lib/credentials/managed-oauth' +import type { createKnowledgeAclFixtureIds } from '@/lib/knowledge/__integration__/seed-source-access-fixture' + +/** Converts the existing synthetic corpus to a central crawl with a real managed reader. */ +export async function seedSearchReaderFixture( + ids: ReturnType +) { + const policy = await getCredentialGroupProviderAdapterByProviderId('confluence').getPolicy( + undefined, + { + workspaceId: ids.workspaceId, + } + ) + const groupId = generateId() + const optionId = generateId() + const crawlerId = generateId() + const enrollmentId = generateId() + await db.insert(credentialGroup).values({ + id: groupId, + workspaceId: ids.workspaceId, + publicId: generateId(), + name: 'Search benchmark readers', + options: [ + { + id: optionId, + provider: 'confluence', + label: 'Confluence fixture', + required: false, + status: 'active', + authorizationAppId: policy.authorizationAppId, + requiredScopes: policy.requiredScopes, + scopeVersion: policy.scopeVersion, + }, + ], + }) + await db.insert(credential).values({ + id: crawlerId, + workspaceId: ids.workspaceId, + type: 'service_account', + providerId: 'atlassian-service-account', + displayName: 'Synthetic crawler', + createdBy: ids.aliceId, + encryptedServiceAccountKey: ( + await encryptSecret( + JSON.stringify({ + type: 'atlassian_service_account', + cloudId: 'search-fixture', + domain: 'search-fixture.atlassian.net', + apiToken: 'fixture-crawler-never-used-for-reading', + }) + ) + ).encrypted, + }) + await db.insert(credentialGroupEnrollment).values({ + id: enrollmentId, + credentialGroupId: groupId, + userId: ids.aliceId, + email: `${ids.aliceId}@fixture.test`, + status: 'completed', + invitationTokenHash: createHash('sha256').update(generateId()).digest('hex'), + invitationExpiresAt: new Date(Date.now() + 3600000), + invitedAt: new Date(), + }) + await db.insert(credential).values({ + id: generateId(), + workspaceId: ids.workspaceId, + type: 'managed_oauth', + displayName: 'Synthetic search reader', + createdBy: ids.aliceId, + providerId: 'confluence', + providerSubjectId: ids.aliceId, + authorizationAppId: policy.authorizationAppId, + credentialGroupEnrollmentId: enrollmentId, + credentialGroupOptionId: optionId, + managedOauthScopeVersion: policy.scopeVersion, + managedOauthStatus: 'active', + grantedScopes: policy.requiredScopes, + grantedAt: new Date(), + encryptedOauthTokenSet: await encryptManagedOAuthTokenSet({ + accessToken: 'fixture-search-reader', + }), + accessTokenExpiresAt: new Date(Date.now() + 3600000), + }) + await db + .update(knowledgeConnector) + .set({ + connectorType: 'confluence', + credentialId: crawlerId, + sourceConfig: { domain: 'search-fixture.atlassian.net' }, + }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + await db.execute(sql`UPDATE document SET acl = ARRAY[${`s:confluence:-:${ids.aliceId}`}], + acl_verified_at = statement_timestamp() WHERE knowledge_base_id = ${ids.knowledgeBaseId}`) + await db.execute(sql`ANALYZE document`) +} diff --git a/apps/sim/lib/knowledge/application/search-diagnostics.ts b/apps/sim/lib/knowledge/application/search-diagnostics.ts new file mode 100644 index 00000000000..da404865994 --- /dev/null +++ b/apps/sim/lib/knowledge/application/search-diagnostics.ts @@ -0,0 +1,45 @@ +import type { AuthorizingUseCase } from '@/lib/core/application' +import type { knowledgeOperations } from '@/lib/knowledge/application/operations' +import type { SearchKnowledgeInput } from '@/lib/knowledge/application/search' +import { + measureSearchStage, + type SearchStage, + withSearchDiagnostics, +} from '@/lib/knowledge/search/diagnostics' + +/** Preserve the operation and authorization contract while timing the whole use case. */ +export function instrumentSearchUseCase< + I extends Omit & { + workspaceId?: string | null + organizationId?: string | null + }, + R, +>( + stage: Extract, + useCase: AuthorizingUseCase +): AuthorizingUseCase { + return { + ...useCase, + execute: (args) => + withSearchDiagnostics( + { + surface: args.input.surface ?? 'other', + principalKind: args.principal.kind, + scopeKind: args.input.organizationId + ? 'organization' + : args.input.workspaceId + ? 'workspace' + : undefined, + topK: args.input.topK, + documentFilterCount: args.input.filters?.documentIds?.length ?? 0, + hasSourceFilter: Boolean(args.input.filters?.source), + hasDateFilter: Boolean(args.input.filters?.modifiedAfter), + tagFilterCount: args.input.tagFilters?.length ?? 0, + hasProvenance: Boolean( + args.input.resultSecretRegistry || args.input.prepareModelInputProvenance + ), + }, + () => measureSearchStage(stage, () => useCase.execute(args)) + ), + } +} diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index 9568e817c9d..1bcd42670b8 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -33,6 +33,7 @@ import { resolveKnowledgeWorkspaceContext, } from '@/lib/knowledge/application/contexts' import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { instrumentSearchUseCase } from '@/lib/knowledge/application/search-diagnostics' import { ALL_TAG_SLOTS } from '@/lib/knowledge/constants' import { getEmbeddingModelInfo, toKbEmbeddingDimensions } from '@/lib/knowledge/embedding-models' import { generateSearchEmbedding, type KbEmbeddingTarget } from '@/lib/knowledge/embeddings' @@ -41,6 +42,7 @@ import { rerank } from '@/lib/knowledge/reranker' import type { RerankerStatus } from '@/lib/knowledge/reranker-models' import { recordOrganizationSearchActivity } from '@/lib/knowledge/search/activity' import { resolveKnowledgeSearchDefaults } from '@/lib/knowledge/search/defaults' +import { annotateSearchDiagnostics, measureSearchStage } from '@/lib/knowledge/search/diagnostics' import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters' import { executeKnowledgeSearch, @@ -252,13 +254,20 @@ async function resolveKnowledgeSearchContext( } } -export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ +const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ operation: knowledgeOperations.search, resolveContext: ({ principal, input }: { principal: Principal; input: SearchKnowledgeInput }) => - resolveKnowledgeSearchContext(input, principal), + measureSearchStage('knowledge_context', () => resolveKnowledgeSearchContext(input, principal)), async execute({ principal, input, context }) { + annotateSearchDiagnostics({ + scopeKind: context.organizationId ? 'organization' : 'workspace', + knowledgeBaseCount: context.knowledgeBases.length, + }) input.signal?.throwIfAborted() - if (context.organizationId) await requireOrganizationSearchAvailable(context.organizationId) + if (context.organizationId) + await measureSearchStage('availability', () => + requireOrganizationSearchAvailable(context.organizationId!) + ) const requestId = generateRequestId() const hasQuery = Boolean(input.query?.trim()) const filters = input.tagFilters ?? [] @@ -276,11 +285,17 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ ) const billingAttribution = hasQuery ? input.resolveBillingAttribution && context.workspaceId - ? await input.resolveBillingAttribution(context.workspaceId) - : await resolveKnowledgeBillingAttribution(principal, context) + ? await measureSearchStage('billing_attribution', () => + input.resolveBillingAttribution!(context.workspaceId!) + ) + : await measureSearchStage('billing_attribution', () => + resolveKnowledgeBillingAttribution(principal, context) + ) : undefined if (shouldMeter && billingAttribution) { - const usage = await checkAttributedUsageLimits(billingAttribution) + const usage = await measureSearchStage('usage_admission', () => + checkAttributedUsageLimits(billingAttribution) + ) if (usage.isExceeded) { throw new KnowledgeUsageLimitExceededError( usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' @@ -292,7 +307,9 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ let structuredFilters: StructuredFilter[] = [] let definitionsByKnowledgeBase = new Map() if (filters.length > 0) { - const built = await resolveKnowledgeTagFilters(filters, knowledgeBaseIds) + const built = await measureSearchStage('tag_filters', () => + resolveKnowledgeTagFilters(filters, knowledgeBaseIds) + ) structuredFilters = built.structuredFilters definitionsByKnowledgeBase = built.definitionsByKnowledgeBase } @@ -333,32 +350,44 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ } : undefined const preparedRegistry = input.prepareModelInputProvenance - ? await input.prepareModelInputProvenance({ userId, workspaceId: context.workspaceId }) + ? await measureSearchStage('input_provenance', () => + input.prepareModelInputProvenance!({ userId, workspaceId: context.workspaceId }) + ) : undefined const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry input.signal?.throwIfAborted() const [queryEmbedding, access, searchDefaults] = await Promise.all([ hasQuery - ? runWithKnowledgeModelInputProvenance(resultSecretRegistry, () => - generateSearchEmbedding( - input.query!, - embeddingTarget!, - context.workspaceId, - input.signal + ? measureSearchStage('embedding', () => + runWithKnowledgeModelInputProvenance(resultSecretRegistry, () => + generateSearchEmbedding( + input.query!, + embeddingTarget!, + context.workspaceId, + input.signal + ) ) ) : Promise.resolve(null), - context.access.get(), - resolveKnowledgeSearchDefaults({ - workspaceId: context.workspaceId, - organizationId: context.organizationId, + 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, - }), + /** The signed-in person, if any; never the billing owner or a key's creator. */ + userId: resolvePrincipalSubjectUserId(principal) ?? undefined, + requestedMode: input.searchMode, + }) + ), ]) input.signal?.throwIfAborted() + annotateSearchDiagnostics({ + accessScopeKind: access.kind, + searchMode: searchDefaults.searchMode, + boostRecency: searchDefaults.boostRecency, + embeddingDimensions: embeddingTarget?.dimensions, + }) const useReranker = Boolean(input.rerankerEnabled && hasQuery) const candidateTopK = useReranker ? input.rerankerInputCount !== undefined @@ -368,24 +397,26 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ ) : Math.min(KNOWLEDGE_SEARCH_COST_POLICY.maxTopK, input.topK * 4) : input.topK - let rows = await executeKnowledgeSearch({ - knowledgeBaseIds, - topK: candidateTopK, - filters: input.filters, - access, - accessProvider: context.access, - signal: input.signal, - searchMode: searchDefaults.searchMode, - boostRecency: searchDefaults.boostRecency, - query: input.query, - queryVector: hasQuery - ? { - vector: JSON.stringify(queryEmbedding?.embedding ?? null), - dimensions: embeddingTarget!.dimensions, - } - : undefined, - structuredFilters: structuredFilters.length > 0 ? structuredFilters : undefined, - }) + let rows = await measureSearchStage('retrieval', () => + executeKnowledgeSearch({ + knowledgeBaseIds, + topK: candidateTopK, + filters: input.filters, + access, + accessProvider: context.access, + signal: input.signal, + searchMode: searchDefaults.searchMode, + boostRecency: searchDefaults.boostRecency, + query: input.query, + queryVector: hasQuery + ? { + vector: JSON.stringify(queryEmbedding?.embedding ?? null), + dimensions: embeddingTarget!.dimensions, + } + : undefined, + structuredFilters: structuredFilters.length > 0 ? structuredFilters : undefined, + }) + ) input.signal?.throwIfAborted() /** Public callers have no input envelope, but persisted reranker inputs still need provenance. */ @@ -404,10 +435,12 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ ReturnType > | null = null if (registry) { - provenanceSnapshot = await importKnowledgeSearchResultSecretProvenance({ - registry, - results: rows, - }) + provenanceSnapshot = await measureSearchStage('result_provenance', () => + importKnowledgeSearchResultSecretProvenance({ + registry, + results: rows, + }) + ) if (!provenanceSnapshot.imported) { registry.markIncomplete('knowledge-result-provenance-unavailable') if (useReranker) { @@ -456,18 +489,20 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ if (useReranker && input.rerankerModel && rows.length > 0) { const candidateCount = rows.length try { - const reranked = await runWithKnowledgeModelInputProvenance(registry, () => - rerank( - input.query!, - rows.map((row) => ({ id: row.id, text: row.content })), - { - model: input.rerankerModel!, - topN: input.topK, - workspaceId: context.workspaceId, + const reranked = await measureSearchStage('reranking', () => + runWithKnowledgeModelInputProvenance(registry, () => + rerank( + input.query!, + rows.map((row) => ({ id: row.id, text: row.content })), + { + model: input.rerankerModel!, + topN: input.topK, + workspaceId: context.workspaceId, - apiKey: input.rerankerApiKey, - signal: input.signal, - } + apiKey: input.rerankerApiKey, + signal: input.signal, + } + ) ) ) rerankerBilled = true @@ -537,29 +572,34 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ } if (shouldMeter && billingAttribution && baseCost && baseCost.total > 0) { try { - await recordUsage({ - userId, - ...(context.workspaceId ? { workspaceId: context.workspaceId } : {}), - ...toBillingContext(billingAttribution), - entries: [ - { - category: 'model', - source: 'knowledge-base', - description: embeddingModel, - cost: baseCost.total, - sourceReference: `kb-search:${requestId}`, - }, - ], - }) - await checkAndBillPayerOverageThreshold(billingAttribution.billingEntity) + await measureSearchStage('usage_recording', () => + recordUsage({ + userId, + ...(context.workspaceId ? { workspaceId: context.workspaceId } : {}), + ...toBillingContext(billingAttribution), + entries: [ + { + category: 'model', + source: 'knowledge-base', + description: embeddingModel, + cost: baseCost.total, + sourceReference: `kb-search:${requestId}`, + }, + ], + }) + ) + await measureSearchStage('overage_billing', () => + checkAndBillPayerOverageThreshold(billingAttribution.billingEntity) + ) } catch (error) { logger.error('Failed to record Knowledge search usage', { error }) } } if (filters.length === 0) { - definitionsByKnowledgeBase = - await getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds) + definitionsByKnowledgeBase = await measureSearchStage('tag_definitions', () => + getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds) + ) } const tagMaps = new Map( [...definitionsByKnowledgeBase].map(([knowledgeBaseId, definitions]) => [ @@ -572,11 +612,13 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ * a model may see, but the source card's modified time and connector type * are only carried here, under the same access predicate as the search. */ - const basicDocumentMetadata = await getDocumentMetadataByIds( - rows.map((row) => row.documentId), - access, - context.access, - input.signal + const basicDocumentMetadata = await measureSearchStage('metadata', () => + getDocumentMetadataByIds( + rows.map((row) => row.documentId), + access, + context.access, + input.signal + ) ) const results = rows .filter((row) => basicDocumentMetadata[row.documentId]) @@ -625,12 +667,14 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ if (renderedMetadata.length === 0) continue if (document.provenance.status === 'unknown' && !knowledgeEnforced) unrecordedCount += 1 if ( - !(await importDurableSecretProvenance( - registry, - document.provenance, - renderedMetadata, - 'knowledge', - { reportUnrecorded: false } + !(await measureSearchStage('metadata_provenance', () => + importDurableSecretProvenance( + registry, + document.provenance, + renderedMetadata, + 'knowledge', + { reportUnrecorded: false } + ) )) ) { registry.markIncomplete('knowledge-result-provenance-unavailable') @@ -652,6 +696,7 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ }) } } + annotateSearchDiagnostics({ resultCount: results.length }) const cost = baseCost ? { input: baseCost.input, @@ -691,12 +736,14 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ afterSuccess: async ({ principal, context, input, result }) => { const actorUserId = resolvePrincipalSubjectUserId(principal) if (context.organizationId && actorUserId) { - await recordOrganizationSearchActivity({ - organizationId: context.organizationId, - userId: actorUserId, - surface: input.surface ?? 'other', - results: result.results, - }) + await measureSearchStage('activity_recording', () => + recordOrganizationSearchActivity({ + organizationId: context.organizationId, + userId: actorUserId, + surface: input.surface ?? 'other', + results: result.results, + }) + ) } PlatformEvents.knowledgeBaseSearched({ knowledgeBaseId: result.knowledgeBaseId, @@ -720,3 +767,8 @@ export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ }) }, }) + +export const searchKnowledge = instrumentSearchUseCase( + 'knowledge_application', + searchKnowledgeUseCase +) diff --git a/apps/sim/lib/knowledge/application/workspace-search.ts b/apps/sim/lib/knowledge/application/workspace-search.ts index 4c44f62cf0e..983c7819011 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.ts @@ -9,7 +9,9 @@ import { } from '@/lib/knowledge/application/contexts' import { knowledgeOperations } from '@/lib/knowledge/application/operations' import { type SearchKnowledgeInput, searchKnowledge } from '@/lib/knowledge/application/search' +import { instrumentSearchUseCase } from '@/lib/knowledge/application/search-diagnostics' import { recordOrganizationSearchActivity } from '@/lib/knowledge/search/activity' +import { measureSearchStage } from '@/lib/knowledge/search/diagnostics' import { findSearchIndex, findWorkspaceSearchIndex } from '@/lib/knowledge/search/search-index' export type SearchWorkspaceKnowledgeInput = Omit< @@ -20,13 +22,15 @@ export type SearchWorkspaceKnowledgeInput = Omit< } /** Search and Assistant share the workspace's canonical Enterprise Search index. */ -export const searchWorkspaceKnowledge = defineAuthorizedKnowledgeUseCase({ +const searchWorkspaceKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ operation: knowledgeOperations.search, resolveContext: ({ input }: { input: SearchWorkspaceKnowledgeInput }) => - resolveKnowledgeWorkspaceContext(input), + measureSearchStage('scope_resolution', () => resolveKnowledgeWorkspaceContext(input)), async execute({ principal, input, context }) { input.signal?.throwIfAborted() - const index = await findWorkspaceSearchIndex(context.workspaceId) + const index = await measureSearchStage('index_resolution', () => + findWorkspaceSearchIndex(context.workspaceId) + ) if (!index) return { results: [], query: input.query ?? '', knowledgeBases: [] } return searchKnowledge.execute({ principal, @@ -35,22 +39,29 @@ export const searchWorkspaceKnowledge = defineAuthorizedKnowledgeUseCase({ }, }) +export const searchWorkspaceKnowledge = instrumentSearchUseCase( + 'workspace_application', + searchWorkspaceKnowledgeUseCase +) + export type SearchOrganizationKnowledgeInput = Omit< SearchWorkspaceKnowledgeInput, 'workspaceId' > & { organizationId: string } /** Organization Search and Assistant resolve the same index and provider ACLs. */ -export const searchOrganizationKnowledge = defineAuthorizedKnowledgeUseCase({ +const searchOrganizationKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ operation: knowledgeOperations.search, resolveContext: ({ input }: { input: SearchOrganizationKnowledgeInput }) => - resolveKnowledgeOrganizationContext(input), + measureSearchStage('scope_resolution', () => resolveKnowledgeOrganizationContext(input)), async execute({ principal, input, context }) { input.signal?.throwIfAborted() - const index = await findSearchIndex({ - kind: 'organization', - organizationId: context.organizationId, - }) + const index = await measureSearchStage('index_resolution', () => + findSearchIndex({ + kind: 'organization', + organizationId: context.organizationId, + }) + ) if (!index) { if (context.organizationId) { await requireOrganizationSearchAvailable(context.organizationId) @@ -66,10 +77,18 @@ export const searchOrganizationKnowledge = defineAuthorizedKnowledgeUseCase({ } return { results: [], query: input.query ?? '', knowledgeBases: [] } } - return searchKnowledge.execute({ principal, input: { ...input, knowledgeBaseIds: [index.id] } }) + return searchKnowledge.execute({ + principal, + input: { ...input, knowledgeBaseIds: [index.id] }, + }) }, }) +export const searchOrganizationKnowledge = instrumentSearchUseCase( + 'organization_application', + searchOrganizationKnowledgeUseCase +) + export type SearchScopedKnowledgeInput = Omit< SearchKnowledgeInput, 'knowledgeBaseIds' | 'workspaceId' | 'organizationId' @@ -77,13 +96,15 @@ export type SearchScopedKnowledgeInput = Omit< ResourceOwner /** The routed owner selects the index; current membership and provider ACLs select its documents. */ -export const searchScopedKnowledge = defineAuthorizedKnowledgeUseCase({ +const searchScopedKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ operation: knowledgeOperations.search, resolveContext: ({ input }: { input: SearchScopedKnowledgeInput }) => - resolveKnowledgeOwnerContext(input), + measureSearchStage('scope_resolution', () => resolveKnowledgeOwnerContext(input)), async execute({ principal, input, context }) { input.signal?.throwIfAborted() - const index = await findSearchIndex(resourceScopeFromOwner(context)) + const index = await measureSearchStage('index_resolution', () => + findSearchIndex(resourceScopeFromOwner(context)) + ) if (!index) { if (context.organizationId) { await requireOrganizationSearchAvailable(context.organizationId) @@ -110,3 +131,8 @@ export const searchScopedKnowledge = defineAuthorizedKnowledgeUseCase({ }) }, }) + +export const searchScopedKnowledge = instrumentSearchUseCase( + 'scoped_application', + searchScopedKnowledgeUseCase +) diff --git a/apps/sim/lib/knowledge/search/diagnostics.test.ts b/apps/sim/lib/knowledge/search/diagnostics.test.ts new file mode 100644 index 00000000000..21cf17c8e07 --- /dev/null +++ b/apps/sim/lib/knowledge/search/diagnostics.test.ts @@ -0,0 +1,126 @@ +/** @vitest-environment node */ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const logs = vi.hoisted(() => ({ info: vi.fn() })) +vi.mock('@sim/logger', () => ({ createLogger: () => logs })) + +import { + annotateSearchDiagnostics, + measureSearchStage, + withSearchDiagnostics, +} from '@/lib/knowledge/search/diagnostics' + +beforeEach(() => { + vi.clearAllMocks() + vi.useFakeTimers({ toFake: ['performance', 'setInterval', 'clearInterval'] }) +}) +afterEach(() => vi.useRealTimers()) + +function deferred() { + let resolve!: () => void + const promise = new Promise((done) => { + resolve = done + }) + return { promise, resolve } +} + +describe('search pipeline diagnostics', () => { + it('reports a stalled stage and clears its heartbeat on completion', async () => { + const gate = deferred() + const result = withSearchDiagnostics({ surface: 'copilot', toolCallId: 'fixture-tool' }, () => + measureSearchStage('embedding', () => gate.promise) + ) + await vi.advanceTimersByTimeAsync(5000) + expect(logs.info).toHaveBeenCalledWith( + 'Knowledge search still running', + expect.objectContaining({ + toolCallId: 'fixture-tool', + elapsedMs: 5000, + activeStages: [{ stage: 'embedding', elapsedMs: 5000 }], + }) + ) + gate.resolve() + await result + expect(logs.info).toHaveBeenLastCalledWith( + 'Knowledge search completed', + expect.objectContaining({ + outcome: 'success', + activeStages: [], + stages: { embedding: { count: 1, totalMs: 5000, maxMs: 5000, errors: 0 } }, + }) + ) + await vi.advanceTimersByTimeAsync(10_000) + expect(logs.info).toHaveBeenCalledTimes(2) + expect(vi.getTimerCount()).toBe(0) + }) + + it('keeps concurrent searches separate and emits only one summary for nested use cases', async () => { + const first = deferred() + const second = deferred() + const a = withSearchDiagnostics({ surface: 'copilot', toolCallId: 'first' }, () => + withSearchDiagnostics({ topK: 15 }, async () => { + await measureSearchStage('result_provenance', () => first.promise) + annotateSearchDiagnostics({ resultCount: 15 }) + }) + ) + const b = withSearchDiagnostics({ surface: 'dashboard' }, () => + measureSearchStage('retrieval', () => second.promise) + ) + await vi.advanceTimersByTimeAsync(25) + second.resolve() + await b + await vi.advanceTimersByTimeAsync(50) + first.resolve() + await a + expect(logs.info).toHaveBeenCalledTimes(2) + const dashboard = logs.info.mock.calls[0][1] + const assistant = logs.info.mock.calls[1][1] + expect(dashboard).toMatchObject({ surface: 'dashboard', elapsedMs: 25 }) + expect(dashboard).not.toHaveProperty('toolCallId') + expect(dashboard.stages).not.toHaveProperty('result_provenance') + expect(assistant).toMatchObject({ + toolCallId: 'first', + topK: 15, + resultCount: 15, + elapsedMs: 75, + }) + expect(assistant.searchId).not.toBe(dashboard.searchId) + expect(assistant.stages).not.toHaveProperty('retrieval') + }) + + it('preserves failures, records repeated stage errors, and does not log error content', async () => { + const error = new Error('private provider response') + await expect( + withSearchDiagnostics({ surface: 'copilot' }, async () => { + await measureSearchStage('vector.authorization', () => 'allowed') + return measureSearchStage('vector.authorization', () => { + throw error + }) + }) + ).rejects.toBe(error) + expect(logs.info).toHaveBeenLastCalledWith( + 'Knowledge search completed', + expect.objectContaining({ + outcome: 'error', + stages: { + 'vector.authorization': { count: 2, totalMs: 0, maxMs: 0, errors: 1 }, + }, + }) + ) + expect(JSON.stringify(logs.info.mock.calls)).not.toContain(error.message) + expect(vi.getTimerCount()).toBe(0) + }) + + it('reports tool failures returned as values without changing the result', async () => { + const result = { success: false, message: 'private error' } + expect(await withSearchDiagnostics({}, async () => result)).toBe(result) + expect(logs.info.mock.calls[0][1].outcome).toBe('error') + expect(JSON.stringify(logs.info.mock.calls)).not.toContain(result.message) + }) + + it('does not create diagnostics outside a search invocation', async () => { + expect(await measureSearchStage('embedding', () => 42)).toBe(42) + expect(logs.info).not.toHaveBeenCalled() + expect(vi.getTimerCount()).toBe(0) + }) +}) diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts new file mode 100644 index 00000000000..09672e26a76 --- /dev/null +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -0,0 +1,173 @@ +import { AsyncLocalStorage } from 'node:async_hooks' +import { createLogger } from '@sim/logger' +import { generateId } from '@sim/utils/id' + +const logger = createLogger('KnowledgeSearchDiagnostics', { logLevel: 'INFO' }) +const PROGRESS_INTERVAL_MS = 5000 + +type RetrievalLeg = 'vector' | 'keyword' | 'tags' +export type SearchStage = + | 'tool_input' + | 'tool_application' + | 'tool_presentation' + | 'workspace_application' + | 'organization_application' + | 'scoped_application' + | 'knowledge_application' + | 'scope_resolution' + | 'knowledge_context' + | 'index_resolution' + | 'availability' + | 'billing_attribution' + | 'usage_admission' + | 'tag_filters' + | 'input_provenance' + | 'embedding' + | 'access_scope' + | 'defaults' + | 'retrieval' + | 'result_provenance' + | 'reranking' + | 'usage_recording' + | 'overage_billing' + | 'tag_definitions' + | 'metadata' + | 'metadata.authorization' + | 'metadata.sql' + | 'metadata_provenance' + | 'activity_recording' + | RetrievalLeg + | `${RetrievalLeg}.candidates` + | `${RetrievalLeg}.authorization` + | `${RetrievalLeg}.hydration` + | 'vector.connection_acquire' + | 'vector.settings' + | 'vector.probe' + | 'vector.ann' + | 'vector.exact' + +/** Fixed, content-free fields. Never pass queries, filters, document identities, SQL, or errors. */ +export interface SearchDiagnosticMetadata { + surface?: 'dashboard' | 'mcp' | 'copilot' | 'workflow' | 'api' | 'slack' | 'other' + toolCallId?: string + executionId?: string + scopeKind?: 'workspace' | 'organization' + accessScopeKind?: 'workspace' | 'user' + principalKind?: string + topK?: number + knowledgeBaseCount?: number + documentFilterCount?: number + hasSourceFilter?: boolean + hasDateFilter?: boolean + tagFilterCount?: number + hasProvenance?: boolean + searchMode?: 'hybrid' | 'vector' + boostRecency?: boolean + embeddingDimensions?: number + resultCount?: number + /** Tool output before the executor's final egress projection; counts only, never content. */ + toolResultBytes?: number + passageBytes?: number + maxPassageBytes?: number + uniqueDocumentCount?: number +} + +interface StageTiming { + count: number + totalMs: number + maxMs: number + errors: number +} + +interface SearchTrace { + searchId: string + startedAt: number + metadata: SearchDiagnosticMetadata + stages: Partial> + active: Map +} + +const traces = new AsyncLocalStorage() +const roundMs = (value: number) => Math.round(value * 100) / 100 + +export function annotateSearchDiagnostics(metadata: SearchDiagnosticMetadata): void { + const trace = traces.getStore() + if (trace) Object.assign(trace.metadata, metadata) +} + +export function recordSearchStageDuration(stage: SearchStage, milliseconds: number): void { + const trace = traces.getStore() + if (!trace) return + const timing = (trace.stages[stage] ??= { count: 0, totalMs: 0, maxMs: 0, errors: 0 }) + timing.count++ + timing.totalMs = roundMs(timing.totalMs + milliseconds) + timing.maxMs = roundMs(Math.max(timing.maxMs, milliseconds)) +} + +/** Timings include waiting on the dependency; nested and parallel stages must not be summed. */ +export async function measureSearchStage( + stage: SearchStage, + run: () => T | PromiseLike +): Promise { + const trace = traces.getStore() + if (!trace) return run() + const span = Symbol(stage) + const startedAt = performance.now() + trace.active.set(span, { stage, startedAt }) + try { + return await run() + } catch (error) { + const timing = (trace.stages[stage] ??= { count: 0, totalMs: 0, maxMs: 0, errors: 0 }) + timing.errors++ + throw error + } finally { + trace.active.delete(span) + recordSearchStageDuration(stage, performance.now() - startedAt) + } +} + +/** One correlated summary per invocation, plus active stages every five seconds while stalled. */ +export async function withSearchDiagnostics( + metadata: SearchDiagnosticMetadata, + run: () => Promise +): Promise { + if (traces.getStore()) { + annotateSearchDiagnostics(metadata) + return run() + } + const trace: SearchTrace = { + searchId: generateId(), + startedAt: performance.now(), + metadata: { ...metadata }, + stages: {}, + active: new Map(), + } + return traces.run(trace, async () => { + const snapshot = () => ({ + searchId: trace.searchId, + ...trace.metadata, + elapsedMs: roundMs(performance.now() - trace.startedAt), + stages: structuredClone(trace.stages), + activeStages: [...trace.active.values()].map(({ stage, startedAt }) => ({ + stage, + elapsedMs: roundMs(performance.now() - startedAt), + })), + }) + const timer = setInterval(() => { + logger.info('Knowledge search still running', snapshot()) + }, PROGRESS_INTERVAL_MS) + timer.unref() + let outcome = 'error' + try { + const result = await run() + outcome = + result && typeof result === 'object' && 'success' in result && result.success === false + ? 'error' + : 'success' + return result + } finally { + clearInterval(timer) + logger.info('Knowledge search completed', { ...snapshot(), outcome }) + } + }) +} diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index bb56e09e2d7..e9beb8eccc7 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -523,6 +523,159 @@ describe('live repository authorization follows ranked candidates', () => { afterEach(() => vi.useRealTimers()) + it('bounds broad vector ranking before metadata and reorders relaxed candidates before trimming', async () => { + queueTableRows( + schemaMock.embedding, + Array.from({ length: 200 }, (_, index) => candidate(`probe-${index}`, 'allowed-source')) + ) + queueTableRows(schemaMock.embedding, [ + { ...candidate('far', 'allowed-source'), distance: 0.3 }, + { ...candidate('near', 'allowed-source'), distance: 0.1 }, + ...Array.from({ length: 18 }, (_, index) => candidate(`other-${index}`, 'allowed-source')), + ]) + queueTableRows(schemaMock.embedding, [ + { id: 'far', content: 'Far authorized passage', distance: 0.3 }, + { id: 'near', content: 'Near authorized passage', distance: 0.1 }, + ]) + const rows = await handleVectorOnlySearch({ + ...params, + structuredFilters: undefined, + topK: 1, + }) + expect(rows.map((row) => row.id)).toEqual(['near']) + expect(dbChainMockFns.orderBy.mock.calls[0]).toHaveLength(1) + expect(Object.keys(dbChainMockFns.select.mock.calls[1][0])).toEqual(['id', 'distance']) + expect(dbChainMockFns.limit.mock.invocationCallOrder[1]).toBeLessThan( + dbChainMockFns.select.mock.invocationCallOrder[2] + ) + expect(JSON.stringify(dbChainMockFns.where.mock.calls[1][0])).toContain('OFFSET 0') + expect(getForConnectors).toHaveBeenCalledExactlyOnceWith(['allowed-source'], undefined) + expect(JSON.stringify(dbChainMockFns.where.mock.calls[2][0])).toContain('github_read_grant') + }) + + it('finishes empty scopes after the bounded probe without scanning HNSW or calling providers', async () => { + queueTableRows(schemaMock.embedding, []) + expect(await handleVectorOnlySearch({ ...params, structuredFilters: undefined })).toEqual([]) + expect(dbChainMockFns.select).toHaveBeenCalledOnce() + expect(dbChainMockFns.limit).toHaveBeenCalledExactlyOnceWith(200) + expect(dbChainMockFns.orderBy).not.toHaveBeenCalled() + expect(getForConnectors).not.toHaveBeenCalled() + }) + + it('reads vectors only for the bounded IDs when a broad scope has few candidates', async () => { + queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')]) + queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')]) + queueTableRows(schemaMock.embedding, [ + { id: 'selected', content: 'Verified small scope', distance: 0.1 }, + ]) + expect(await handleVectorOnlySearch({ ...params, structuredFilters: undefined })).toEqual([ + { id: 'selected', content: 'Verified small scope', distance: 0.1 }, + ]) + expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id']) + expect(JSON.stringify(dbChainMockFns.where.mock.calls[0][0])).not.toContain('<=>') + expect( + hasMockCondition( + dbChainMockFns.where.mock.calls[1][0], + (node) => + node.type === 'inArray' && + node.column === schemaMock.embedding.id && + Array.isArray(node.values) && + node.values.length === 1 && + node.values[0] === 'selected' + ) + ).toBe(true) + expect(getForConnectors).toHaveBeenCalledExactlyOnceWith(['allowed-source'], undefined) + }) + + it('falls back to exact ranking when the approximate page cannot fill its limit', async () => { + queueTableRows( + schemaMock.embedding, + Array.from({ length: 200 }, (_, index) => candidate(`probe-${index}`, 'allowed-source')) + ) + queueTableRows(schemaMock.embedding, [candidate('partial', 'allowed-source')]) + queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')]) + queueTableRows(schemaMock.embedding, [ + { id: 'selected', content: 'Verified fallback', distance: 0.1 }, + ]) + expect(await handleVectorOnlySearch({ ...params, structuredFilters: undefined })).toEqual([ + { id: 'selected', content: 'Verified fallback', distance: 0.1 }, + ]) + expect(render(dbChainMockFns.orderBy.mock.calls.at(-1)![0]).sql).toContain('+ 0') + expect(JSON.stringify(dbChainMockFns.where.mock.calls.at(-1)![0])).toContain( + 'github_read_grant' + ) + }) + + it('restarts exact ranking at zero and advances past already considered ANN candidates', async () => { + const probe = Array.from({ length: 200 }, (_, index) => + candidate(`probe-${index}`, 'allowed-source') + ) + const approximate = Array.from({ length: 20 }, (_, index) => + candidate(`approximate-${index}`, 'allowed-source') + ) + queueTableRows(schemaMock.embedding, probe) + queueTableRows(schemaMock.embedding, approximate) + queueTableRows(schemaMock.embedding, []) + queueTableRows(schemaMock.embedding, probe) + queueTableRows(schemaMock.embedding, []) + queueTableRows(schemaMock.embedding, approximate) + queueTableRows(schemaMock.embedding, probe) + queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')]) + queueTableRows(schemaMock.embedding, [ + { id: 'selected', content: 'Reachable after the exact restart', distance: 0.1 }, + ]) + + const rows = await handleVectorOnlySearch({ ...params, structuredFilters: undefined }) + + expect(rows.map((row) => row.id)).toEqual(['selected']) + expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [20], [0], [20]]) + expect(getForConnectors).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(3) + }) + + it('keeps the nearest exact results when an earlier ANN page hydrated only a farther result', async () => { + const probe = Array.from({ length: 200 }, (_, index) => + candidate(`probe-${index}`, 'allowed-source') + ) + queueTableRows(schemaMock.embedding, probe) + queueTableRows(schemaMock.embedding, [ + { ...candidate('far', 'allowed-source'), distance: 0.7 }, + ...Array.from({ length: 19 }, (_, index) => candidate(`hidden-${index}`, 'allowed-source')), + ]) + queueTableRows(schemaMock.embedding, [{ id: 'far', content: 'Far result', distance: 0.7 }]) + queueTableRows(schemaMock.embedding, probe) + queueTableRows(schemaMock.embedding, []) + queueTableRows(schemaMock.embedding, [ + candidate('near', 'allowed-source'), + candidate('nearer', 'allowed-source'), + candidate('far', 'allowed-source'), + ]) + queueTableRows(schemaMock.embedding, [ + { id: 'near', content: 'Near result', distance: 0.2 }, + { id: 'nearer', content: 'Nearest result', distance: 0.1 }, + ]) + + const rows = await handleVectorOnlySearch({ + ...params, + topK: 2, + structuredFilters: undefined, + }) + + expect(rows.map((row) => row.id)).toEqual(['nearer', 'near']) + expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [20], [0]]) + expect( + hasMockCondition( + dbChainMockFns.where.mock.calls.at(-1)![0], + (node) => + node.type === 'inArray' && + node.column === schemaMock.embedding.id && + Array.isArray(node.values) && + node.values.length === 2 && + !node.values.includes('far') + ) + ).toBe(true) + }) + it.each(['vector', 'tag-vector', 'tags', 'keyword'] as const)( '%s ranks identifiers before verification and loads content under the full predicate', async (mode) => { diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index fc7b69969ab..17becb53e74 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -9,6 +9,7 @@ import { } from '@/lib/knowledge/access/predicate' import type { KnowledgeAccessProvider, KnowledgeAccessScope } from '@/lib/knowledge/access/types' import type { KbEmbeddingDimensions } from '@/lib/knowledge/embedding-models' +import { measureSearchStage, recordSearchStageDuration } from '@/lib/knowledge/search/diagnostics' import { workspaceSearchFilterConditions } from '@/lib/knowledge/search/filter-conditions' import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters' import { applyRecencyBoost, RRF_K } from '@/lib/knowledge/search/recency' @@ -44,12 +45,16 @@ async function withVectorScanSettings( run: (executor: SearchExecutor) => Promise ): Promise { if (Date.now() < hnswSettingsUnsupportedUntil) return run(db) + const acquireStarted = performance.now() let applyingSettings = false try { return await db.transaction(async (tx) => { + recordSearchStageDuration('vector.connection_acquire', performance.now() - acquireStarted) applyingSettings = true - await tx.execute( - sql`SELECT set_config('hnsw.iterative_scan', 'relaxed_order', true), set_config('hnsw.max_scan_tuples', ${HNSW_MAX_SCAN_TUPLES}, true)` + await measureSearchStage('vector.settings', () => + tx.execute( + sql`SELECT set_config('hnsw.iterative_scan', 'relaxed_order', true), set_config('hnsw.max_scan_tuples', ${HNSW_MAX_SCAN_TUPLES}, true)` + ) ) applyingSettings = false return run(tx) @@ -92,28 +97,31 @@ export async function getDocumentMetadataByIds( const uniqueIds = [...new Set(documentIds)] const authorizedAccess = accessProvider - ? await accessProvider.getForDocuments(uniqueIds, signal) + ? await measureSearchStage('metadata.authorization', () => + accessProvider.getForDocuments(uniqueIds, signal) + ) : access - const documents = await 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(authorizedAccess) + 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(authorizedAccess) + ) ) - ) - + ) const map: Record = {} documents.forEach((doc) => { map[doc.id] = { @@ -424,6 +432,12 @@ function getVisibilityConditions( ] } +/** Each ranking strategy owns its cursor; ANN offsets cannot paginate exact ordering. */ +interface SearchReadCandidatePage { + candidates: SearchReadCandidate[] + nextOffset: number +} + interface SearchReadCandidate { id: string documentId: string @@ -457,6 +471,7 @@ const LIVE_SEARCH_BUDGET_MS = 8000 * consume every result slot. The existing vector tuple budget also bounds candidate work. */ async function selectAuthorizedSearchResults(input: { + leg: 'vector' | 'keyword' | 'tags' accessProvider: KnowledgeAccessProvider filters?: WorkspaceSearchFilters signal?: AbortSignal @@ -465,13 +480,15 @@ async function selectAuthorizedSearchResults(input: { limit: number, offset: number, excludedSources: readonly string[] - ) => Promise + ) => Promise + compareResults?: (a: SearchResult, b: SearchResult) => number hydrate: (ids: string[], access: KnowledgeAccessScope) => Promise }): Promise { const deadline = Date.now() + LIVE_SEARCH_BUDGET_MS const pageSize = Math.min(LIVE_SEARCH_PAGE_SIZE, Math.max(input.topK, 20)) const results = new Map() const excludedSources = new Set() + const considered = new Set() let scanned = 0 let offset = 0 while ( @@ -480,9 +497,18 @@ async function selectAuthorizedSearchResults(input: { Date.now() < deadline ) { input.signal?.throwIfAborted() - const candidates = await input.selectPage(pageSize, offset, [...excludedSources]) - if (!candidates.length) break - scanned += candidates.length + const page = await measureSearchStage(`${input.leg}.candidates`, () => + input.selectPage(pageSize, offset, [...excludedSources]) + ) + if (!page.candidates.length) break + scanned += page.candidates.length + offset = page.nextOffset + const candidates = page.candidates.filter((candidate) => !considered.has(candidate.id)) + for (const candidate of candidates) considered.add(candidate.id) + if (!candidates.length) { + if (page.candidates.length < pageSize) break + continue + } /** Candidate and hydration queries enforce this source filter; connector types are immutable. */ const connectorIds = input.filters?.source && @@ -496,7 +522,9 @@ async function selectAuthorizedSearchResults(input: { ) ), ] - const access = await input.accessProvider.getForConnectors(connectorIds, input.signal) + const access = await measureSearchStage(`${input.leg}.authorization`, () => + input.accessProvider.getForConnectors(connectorIds, input.signal) + ) input.signal?.throwIfAborted() const grantedSources = new Set( access.kind === 'user' @@ -515,21 +543,25 @@ async function selectAuthorizedSearchResults(input: { ) excludedSources.add(candidate.connectorId) } - const hydrated = await input.hydrate( - candidates.map((candidate) => candidate.id), - access + const hydrated = await measureSearchStage(`${input.leg}.hydration`, () => + input.hydrate( + candidates.map((candidate) => candidate.id), + access + ) ) const byId = new Map(hydrated.map((row) => [row.id, row])) for (const candidate of candidates) { const row = byId.get(candidate.id) if (row) results.set(row.id, row) - if (results.size === input.topK) break + if (!input.compareResults && results.size === input.topK) break } - if (excludedSources.size > excludedBefore) offset = 0 - else { - offset += candidates.length - if (candidates.length < pageSize) 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 (excludedSources.size > excludedBefore) offset = 0 + else if (page.candidates.length < pageSize) break } input.signal?.throwIfAborted() return [...results.values()] @@ -593,12 +625,13 @@ export async function handleTagOnlySearch(params: SearchParams): Promise - db + selectPage: async (limit, offset, excludedSources) => { + const candidates = await db .select(SEARCH_READ_CANDIDATE_FIELDS) .from(embedding) .innerJoin(document, eq(embedding.documentId, document.id)) @@ -615,7 +648,9 @@ export async function handleTagOnlySearch(params: SearchParams): Promise hydrateSearchCandidates( ids, @@ -712,43 +747,115 @@ export async function handleVectorOnlySearch(params: SearchParams): Promise a.distance - b.distance) } -/** The vector transaction ends after ranking, before any provider authorization request starts. */ -function selectLiveVectorResults( +/** + * Keep broad candidate visibility correlated with the vector scan. Flattening + * the document join can make PostgreSQL prefer sorting every vector (whose + * TOAST reads it undercosts) before checking access. OFFSET 0 keeps that + * visibility check inside the scan, before LIMIT. Only bounded identities join + * back for source metadata; content still requires live authorization below. + * Small scopes and underfilled approximate pages use exact ranking, so selective + * permissions do not force a fruitless index walk or lose reachable matches. + */ +async function selectLiveVectorResults( params: SearchParams, accessProvider: KnowledgeAccessProvider, distance: SQL, filters: (SQL | undefined)[] ): Promise { const conditions = [inArray(embedding.knowledgeBaseId, params.knowledgeBaseIds), ...filters] - return selectAuthorizedSearchResults({ + let useExactRanking = false + const rows = await selectAuthorizedSearchResults({ + leg: 'vector', accessProvider, filters: params.filters, signal: params.signal, topK: params.topK, + compareResults: (a, b) => a.distance - b.distance, selectPage: (limit, offset, excludedSources) => - withVectorScanSettings((executor) => - executor - .select({ ...SEARCH_READ_CANDIDATE_FIELDS, distance: distance.as('distance') }) + withVectorScanSettings(async (executor) => { + const visibility = [ + ...getVisibilityConditions( + params.access, + params.filters, + knowledgeMetadataCandidateAccessCondition(params.access) + ), + excludeSearchSources(excludedSources), + ] + /** Adding zero prevents an underfilled HNSW scan from being chosen again for fallback. */ + const exactPage = async (candidateIds?: string[]) => { + const exactOffset = useExactRanking ? offset : 0 + useExactRanking = true + const candidates = await measureSearchStage('vector.exact', () => + executor + .select({ ...SEARCH_READ_CANDIDATE_FIELDS, distance: distance.as('distance') }) + .from(embedding) + .innerJoin(document, eq(embedding.documentId, document.id)) + .where( + and( + ...conditions, + ...visibility, + candidateIds ? inArray(embedding.id, candidateIds) : undefined + ) + ) + .orderBy(sql`(${distance}) + 0`, embedding.id) + .limit(limit) + .offset(exactOffset) + ) + return { candidates, nextOffset: exactOffset + candidates.length } + } + if (params.filters?.documentIds?.length || params.structuredFilters?.length) { + return exactPage() + } + /** Probe visibility without vector reads; revoked scopes must not detoast the corpus. */ + const probe = await measureSearchStage('vector.probe', () => + executor + .select({ id: embedding.id }) + .from(embedding) + .innerJoin(document, eq(embedding.documentId, document.id)) + .where(and(inArray(embedding.knowledgeBaseId, params.knowledgeBaseIds), ...visibility)) + .limit(LIVE_SEARCH_PAGE_SIZE) + ) + if (probe.length === 0) return { candidates: [], nextOffset: offset } + if (probe.length < LIVE_SEARCH_PAGE_SIZE) { + return exactPage(probe.map((candidate) => candidate.id)) + } + if (useExactRanking) return exactPage() + const ranked = executor + .select({ id: embedding.id, distance: distance.as('distance') }) .from(embedding) - .innerJoin(document, eq(embedding.documentId, document.id)) .where( and( ...conditions, - ...getVisibilityConditions( - params.access, - params.filters, - knowledgeMetadataCandidateAccessCondition(params.access) - ), - excludeSearchSources(excludedSources) + sql`EXISTS ( + SELECT 1 FROM ${document} + WHERE ${and(eq(document.id, embedding.documentId), ...visibility)} + OFFSET 0 + )` ) ) - .orderBy(distance, embedding.id) + .orderBy(distance) .limit(limit) .offset(offset) - ), + .as('ranked_search_candidates') + const page = await measureSearchStage('vector.ann', () => + executor + .select({ ...SEARCH_READ_CANDIDATE_FIELDS, distance: ranked.distance }) + .from(ranked) + .innerJoin(embedding, eq(embedding.id, ranked.id)) + .innerJoin(document, eq(document.id, embedding.documentId)) + .orderBy(ranked.distance, ranked.id) + ) + if (page.length < limit) return exactPage() + return { + candidates: page.sort((a, b) => a.distance - b.distance), + nextOffset: offset + page.length, + } + }), hydrate: (ids, authorized) => hydrateSearchCandidates(ids, authorized, distance.as('distance'), params.filters, conditions), }) + /** Relaxed HNSW scans can return adjacent pages out of distance order. */ + return rows.sort((a, b) => a.distance - b.distance) } /** @@ -835,12 +942,13 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise ...tagFilterConditions, ] return selectAuthorizedSearchResults({ + leg: 'keyword', accessProvider: params.accessProvider, filters: params.filters, signal: params.signal, topK, - selectPage: (limit, offset, excludedSources) => - db + selectPage: async (limit, offset, excludedSources) => { + const candidates = await db .select({ ...SEARCH_READ_CANDIDATE_FIELDS, keywordRank: rankExpr.as('keyword_rank') }) .from(embedding) .innerJoin(document, eq(embedding.documentId, document.id)) @@ -857,7 +965,9 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise ) .orderBy(sql`${rankExpr} DESC`, embedding.id) .limit(limit) - .offset(offset), + .offset(offset) + return { candidates, nextOffset: offset + candidates.length } + }, hydrate: (ids, authorized) => hydrateSearchCandidates( ids, @@ -1087,15 +1197,17 @@ export async function executeKnowledgeSearch( if (!hasFilters) { throw new Error('A search query or tag filters are required') } - return await handleTagOnlySearch({ - knowledgeBaseIds, - topK, - structuredFilters, - access, - accessProvider: params.accessProvider, - signal: params.signal, - filters: params.filters, - }) + return await measureSearchStage('tags', () => + handleTagOnlySearch({ + knowledgeBaseIds, + topK, + structuredFilters, + access, + accessProvider: params.accessProvider, + signal: params.signal, + filters: params.filters, + }) + ) } if (!queryVector) { @@ -1110,29 +1222,30 @@ export async function executeKnowledgeSearch( */ const legTopK = searchMode === 'hybrid' ? hybridCandidateCount(topK) : topK - const vectorSearch = hasFilters - ? handleTagAndVectorSearch({ - knowledgeBaseIds, - topK: legTopK, - structuredFilters, - queryVector, - distanceThreshold, - access, - accessProvider: params.accessProvider, - signal: params.signal, - filters: params.filters, - }) - : handleVectorOnlySearch({ - knowledgeBaseIds, - topK: legTopK, - queryVector, - distanceThreshold, - access, - accessProvider: params.accessProvider, - signal: params.signal, - filters: params.filters, - }) - + const vectorSearch = measureSearchStage('vector', () => + hasFilters + ? handleTagAndVectorSearch({ + knowledgeBaseIds, + topK: legTopK, + structuredFilters, + queryVector, + distanceThreshold, + access, + accessProvider: params.accessProvider, + signal: params.signal, + filters: params.filters, + }) + : handleVectorOnlySearch({ + knowledgeBaseIds, + topK: legTopK, + queryVector, + distanceThreshold, + access, + accessProvider: params.accessProvider, + signal: params.signal, + filters: params.filters, + }) + ) if (searchMode === 'vector') { const results = await vectorSearch return boostRecency ? applyRecencyBoost(results) : results @@ -1142,17 +1255,19 @@ export async function executeKnowledgeSearch( * The lexical leg is best-effort: a failure there falls back to vector-only * results rather than failing the whole search. */ - const keywordSearch = executeKeywordSearch({ - knowledgeBaseIds, - topK: legTopK, - query: query!, - queryVector, - structuredFilters, - access, - accessProvider: params.accessProvider, - signal: params.signal, - filters: params.filters, - }).catch((error) => { + const keywordSearch = measureSearchStage('keyword', () => + executeKeywordSearch({ + knowledgeBaseIds, + topK: legTopK, + query: query!, + queryVector, + structuredFilters, + access, + accessProvider: params.accessProvider, + signal: params.signal, + filters: params.filters, + }) + ).catch((error) => { logger.warn('Keyword search leg failed; falling back to vector-only results', { error: getErrorMessage(error, 'Unknown error'), }) diff --git a/package.json b/package.json index 768902f6a06..ca71919251a 100644 --- a/package.json +++ b/package.json @@ -16,6 +16,7 @@ "dev:full:capped": "bunx concurrently -n \"App,Realtime\" -c \"cyan,magenta\" \"cd apps/sim && bun run dev:capped\" \"cd apps/realtime && bun run dev\"", "test": "bun run test:setup && bun run test:scripts && turbo run test", "test:setup": "bun run --cwd packages/sim-setup test", + "test:search-performance": "KNOWLEDGE_SEARCH_PERFORMANCE_TEST=true bun --no-env-file scripts/test-knowledge-acls.ts search-latency", "format": "turbo run format", "format:check": "turbo run format:check", "lint": "turbo run lint", diff --git a/packages/testing/src/mocks/database.mock.test.ts b/packages/testing/src/mocks/database.mock.test.ts index 5d60587a385..5a93aee7071 100644 --- a/packages/testing/src/mocks/database.mock.test.ts +++ b/packages/testing/src/mocks/database.mock.test.ts @@ -97,6 +97,7 @@ describe('database mock', () => { .where({}) .orderBy(workflowTable.id) .limit(1) + .offset(1) .as('ranked') expect(ranked.id).toEqual({ diff --git a/packages/testing/src/mocks/database.mock.ts b/packages/testing/src/mocks/database.mock.ts index a7e8151360e..b12125af63c 100644 --- a/packages/testing/src/mocks/database.mock.ts +++ b/packages/testing/src/mocks/database.mock.ts @@ -293,7 +293,7 @@ const lazyRowsThenable = (getRows: RowsSupplier): any => ({ // `.limit(1).for('update')` row-lock form. const limitBuilder = (getRows: RowsSupplier, fields: SelectedFields = {}) => { const thenable = lazyRowsThenable(getRows) - thenable.offset = spyOrDefault(offset, () => lazyRowsThenable(getRows)) + thenable.offset = spyOrDefault(offset, () => limitBuilder(getRows, fields)) thenable.for = spyOrDefault(forClause, () => limitBuilder(getRows, fields)) thenable.as = spyOrDefault(asAlias, (alias: string) => subqueryFields(fields, alias)) return thenable diff --git a/scripts/test-knowledge-acls.ts b/scripts/test-knowledge-acls.ts index 60b5c390d30..e13e665ea36 100644 --- a/scripts/test-knowledge-acls.ts +++ b/scripts/test-knowledge-acls.ts @@ -22,7 +22,10 @@ const testFilters = process.argv.slice(2) if (testFilters.some((filter) => filter.startsWith('-')) || (scale && testFilters.length)) { throw new Error('Pass only filename filters, and do not combine them with the scale suite') } -const keepScaleDatabase = scale && process.env.KNOWLEDGE_SCALE_KEEP_DATABASE === 'true' +const keepDatabase = + (scale && process.env.KNOWLEDGE_SCALE_KEEP_DATABASE === 'true') || + (process.env.KNOWLEDGE_SEARCH_PERFORMANCE_TEST === 'true' && + process.env.KNOWLEDGE_SEARCH_PERFORMANCE_KEEP_DATABASE === 'true') const scaleReportFile = process.env.KNOWLEDGE_SCALE_REPORT_FILE ?? path.join(tmpdir(), `${container}.json`) @@ -96,8 +99,8 @@ try { if (!/^127\.0\.0\.1:\d+$/.test(endpoint)) throw new Error('Unexpected disposable Postgres endpoint') const databaseUrl = `postgresql://postgres@${endpoint}/${database}` - if (keepScaleDatabase) - logger.info('Retaining disposable scale database for follow-up measurements', { + if (keepDatabase) + logger.info('Retaining disposable benchmark database for follow-up measurements', { container, databaseUrl, }) @@ -192,6 +195,6 @@ try { try { if (redisStarted) run('docker', ['stop', redisContainer], { capture: true }) } finally { - if (started && !keepScaleDatabase) run('docker', ['stop', container], { capture: true }) + if (started && !keepDatabase) run('docker', ['stop', container], { capture: true }) } }