From 507de5a563abd7261e35493aa62fe78a79170f64 Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Fri, 11 Sep 2026 15:47:27 -0700 Subject: [PATCH 1/2] fix(search): bound retrieval and progressively read document evidence --- .../knowledge/search/route.provenance.test.ts | 5 +- .../lib/api/contracts/knowledge/documents.ts | 19 +- .../lib/copilot/generated/tool-catalog-v1.ts | 25 +- .../lib/copilot/generated/tool-schemas-v1.ts | 25 +- .../copilot/generated/trace-attributes-v1.ts | 2 + .../server/knowledge/workspace-search.test.ts | 18 +- .../server/knowledge/workspace-search.ts | 182 +++--- .../gitlab-live.integration.ts | 3 +- .../search-latency.integration.ts | 286 ++++++++- .../application/read-search-document.test.ts | 86 ++- .../application/read-search-document.ts | 108 +++- .../lib/knowledge/application/search.test.ts | 5 +- apps/sim/lib/knowledge/application/search.ts | 22 +- .../knowledge/application/workspace-search.ts | 22 +- apps/sim/lib/knowledge/chunks/service.ts | 10 +- apps/sim/lib/knowledge/chunks/types.ts | 2 + apps/sim/lib/knowledge/search/budget.test.ts | 55 ++ apps/sim/lib/knowledge/search/budget.ts | 115 ++++ apps/sim/lib/knowledge/search/diagnostics.ts | 9 +- apps/sim/lib/knowledge/search/queries.test.ts | 2 +- apps/sim/lib/knowledge/search/queries.ts | 562 +++++++++--------- apps/sim/lib/knowledge/search/snippet.test.ts | 29 + apps/sim/lib/knowledge/search/snippet.ts | 39 ++ 23 files changed, 1222 insertions(+), 409 deletions(-) create mode 100644 apps/sim/lib/knowledge/search/budget.test.ts create mode 100644 apps/sim/lib/knowledge/search/budget.ts diff --git a/apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts b/apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts index 3971a4c4ca5..7e847cb2797 100644 --- a/apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts +++ b/apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts @@ -59,7 +59,10 @@ vi.mock('@/lib/knowledge/embeddings', () => ({ vi.mock('@/lib/knowledge/search/queries', () => ({ generateSearchEmbedding: mocks.generateEmbedding, - executeKnowledgeSearch: mocks.executeSearch, + retrieveKnowledgeSearch: async (...args: unknown[]) => ({ + rows: await mocks.executeSearch(...args), + retrieval: { status: 'complete', timedOutLegs: [] }, + }), getDocumentMetadataByIds: mocks.getDocumentMetadata, })) diff --git a/apps/sim/lib/api/contracts/knowledge/documents.ts b/apps/sim/lib/api/contracts/knowledge/documents.ts index a5219138b05..4b586ff6765 100644 --- a/apps/sim/lib/api/contracts/knowledge/documents.ts +++ b/apps/sim/lib/api/contracts/knowledge/documents.ts @@ -429,8 +429,23 @@ export const readSearchDocumentResultSchema = z.object({ knowledgeBaseId: z.string().min(1), documentName: z.string().nullable(), sourceUrl: z.string().nullable(), - chunks: z.array(z.object({ content: z.string(), chunkIndex: z.number().int().min(0) })).max(50), + chunks: z + .array( + z.object({ + content: z.string().max(8000), + chunkIndex: z.number().int().min(0), + startOffset: z.number().int().min(0), + endOffset: z.number().int().min(0), + totalCharacters: z.number().int().min(0), + }) + ) + .max(8), hasMore: z.boolean(), - nextOffset: z.number().int().min(0).nullable(), + next: z + .object({ + startChunkIndex: z.number().int().min(0), + startOffset: z.number().int().min(0), + }) + .nullable(), }) export type ReadSearchDocumentResult = z.output diff --git a/apps/sim/lib/copilot/generated/tool-catalog-v1.ts b/apps/sim/lib/copilot/generated/tool-catalog-v1.ts index b65bebaf7d1..22fb32d4a10 100644 --- a/apps/sim/lib/copilot/generated/tool-catalog-v1.ts +++ b/apps/sim/lib/copilot/generated/tool-catalog-v1.ts @@ -4750,16 +4750,24 @@ export const ReadDocument: ToolCatalogEntry = { type: 'string', }, limit: { - default: 20, - description: 'Maximum number of chunks to read.', - maximum: 50, + default: 3, + description: + 'Maximum chunks to read; the server may return fewer to fit its text budget. Follow next when more context is needed.', + maximum: 8, minimum: 1, type: 'integer', }, - offset: { - default: 0, - description: 'Number of chunks to skip.', - maximum: 5000, + startChunkIndex: { + description: + "Inclusive chunk index from search or a previous read's next object. Gaps from disabled chunks are skipped.", + maximum: 2147483647, + minimum: 0, + type: 'integer', + }, + startOffset: { + description: + 'UTF-16 character offset within startChunkIndex. Omit to read the chunk from its start, or copy next.startOffset to continue a partial chunk.', + maximum: 2147483647, minimum: 0, type: 'integer', }, @@ -5616,7 +5624,8 @@ export const SearchWorkspace: ToolCatalogEntry = { }, topK: { default: 20, - description: 'Maximum number of matching chunks to return.', + description: + 'Maximum number of matching passage previews to return. Retrieval ranking is independent of preview length.', maximum: 50, minimum: 1, type: 'integer', diff --git a/apps/sim/lib/copilot/generated/tool-schemas-v1.ts b/apps/sim/lib/copilot/generated/tool-schemas-v1.ts index 388d341ae1f..72e603a3613 100644 --- a/apps/sim/lib/copilot/generated/tool-schemas-v1.ts +++ b/apps/sim/lib/copilot/generated/tool-schemas-v1.ts @@ -4694,16 +4694,24 @@ export const TOOL_RUNTIME_SCHEMAS: Record = { type: 'string', }, limit: { - default: 20, - description: 'Maximum number of chunks to read.', - maximum: 50, + default: 3, + description: + 'Maximum chunks to read; the server may return fewer to fit its text budget. Follow next when more context is needed.', + maximum: 8, minimum: 1, type: 'integer', }, - offset: { - default: 0, - description: 'Number of chunks to skip.', - maximum: 5000, + startChunkIndex: { + description: + "Inclusive chunk index from search or a previous read's next object. Gaps from disabled chunks are skipped.", + maximum: 2147483647, + minimum: 0, + type: 'integer', + }, + startOffset: { + description: + 'UTF-16 character offset within startChunkIndex. Omit to read the chunk from its start, or copy next.startOffset to continue a partial chunk.', + maximum: 2147483647, minimum: 0, type: 'integer', }, @@ -5518,7 +5526,8 @@ export const TOOL_RUNTIME_SCHEMAS: Record = { }, topK: { default: 20, - description: 'Maximum number of matching chunks to return.', + description: + 'Maximum number of matching passage previews to return. Retrieval ranking is independent of preview length.', maximum: 50, minimum: 1, type: 'integer', diff --git a/apps/sim/lib/copilot/generated/trace-attributes-v1.ts b/apps/sim/lib/copilot/generated/trace-attributes-v1.ts index fdc1a5da7ca..df4a808f809 100644 --- a/apps/sim/lib/copilot/generated/trace-attributes-v1.ts +++ b/apps/sim/lib/copilot/generated/trace-attributes-v1.ts @@ -478,6 +478,7 @@ export const TraceAttr = { LlmStreamChunks: 'llm.stream.chunks', LlmStreamFirstChunkBytes: 'llm.stream.first_chunk_bytes', LlmStreamFirstChunkMs: 'llm.stream.first_chunk_ms', + LlmStreamFirstTokenMs: 'llm.stream.first_token_ms', LlmStreamOpenMs: 'llm.stream.open_ms', LlmStreamTotalMs: 'llm.stream.total_ms', LockAcquired: 'lock.acquired', @@ -1159,6 +1160,7 @@ export const TraceAttrValues: readonly TraceAttrValue[] = [ 'llm.stream.chunks', 'llm.stream.first_chunk_bytes', 'llm.stream.first_chunk_ms', + 'llm.stream.first_token_ms', 'llm.stream.open_ms', 'llm.stream.total_ms', 'lock.acquired', 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 1e669a7b45b..7edc4ba294f 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 @@ -61,6 +61,7 @@ describe('Assistant retrieval tools', () => { beforeEach(() => { vi.clearAllMocks() mocks.search.mockResolvedValue({ + retrieval: { status: 'complete', timedOutLegs: [] }, knowledgeBases: [{ id: 'index', name: 'Enterprise Search' }], results: [ { @@ -83,7 +84,7 @@ describe('Assistant retrieval tools', () => { sourceUrl: 'https://source.test/doc', chunks: [{ content: 'body', chunkIndex: 0 }], hasMore: false, - nextOffset: null, + next: null, }) }) it('pins organization and private chat while reusing the canonical search index and citations', async () => { @@ -192,6 +193,7 @@ describe('Assistant retrieval tools', () => { async (count) => { const content = 'Confidential passage Γ©πŸ”Ž'.repeat(100) mocks.search.mockResolvedValueOnce({ + retrieval: { status: 'complete', timedOutLegs: [] }, knowledgeBases: [{ id: 'index', name: 'Enterprise Search' }], results: Array.from({ length: count }, (_, index) => ({ knowledgeBaseId: 'index', @@ -217,8 +219,9 @@ describe('Assistant retrieval tools', () => { expect.objectContaining({ toolCallId: 'call', toolResultBytes: Buffer.byteLength(JSON.stringify(output)), - passageBytes: count * Buffer.byteLength(content), - maxPassageBytes: count ? Buffer.byteLength(content) : 0, + passageBytes: count * Buffer.byteLength(content.slice(0, 1200)), + originalPassageBytes: count * Buffer.byteLength(content), + maxPassageBytes: count ? Buffer.byteLength(content.slice(0, 1200)) : 0, uniqueDocumentCount: Math.min(count, 4), }) ) @@ -245,6 +248,7 @@ describe('Assistant retrieval tools', () => { }) it('projects the provider name for connected-source citations instead of the index name', async () => { mocks.search.mockResolvedValueOnce({ + retrieval: { status: 'complete', timedOutLegs: [] }, knowledgeBases: [{ id: 'index', name: 'Sim Search' }], results: [ { @@ -292,20 +296,20 @@ describe('Assistant retrieval tools', () => { }) it('reads a selected document through the shared use case and rejects unbounded pages', async () => { expect( - await readDocumentServerTool.execute({ documentId: 'doc', offset: 20 }, context) + await readDocumentServerTool.execute({ documentId: 'doc', startChunkIndex: 20 }, context) ).toMatchObject({ success: true }) expect(mocks.read).toHaveBeenCalledWith( expect.objectContaining({ input: expect.objectContaining({ assertedWorkspaceId: 'workspace', filters: context.assistantSearch, - offset: 20, - limit: 20, + startChunkIndex: 20, + limit: 3, }), }) ) expect( - await readDocumentServerTool.execute({ documentId: 'doc', limit: 10000 }, context) + await readDocumentServerTool.execute({ documentId: 'doc', limit: 9 }, context) ).toMatchObject({ success: false }) expect(mocks.read).toHaveBeenCalledOnce() }) 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 01f6cb833f5..87cdebc8836 100644 --- a/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.ts +++ b/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.ts @@ -15,6 +15,7 @@ import { searchWorkspaceKnowledge, } from '@/lib/knowledge/application/workspace-search' import { sourceAuthor } from '@/lib/knowledge/search/author' +import { SearchDeadlineError } from '@/lib/knowledge/search/budget' import { createKnowledgeDocumentCitation } from '@/lib/knowledge/search/citation' import { annotateSearchDiagnostics, @@ -23,6 +24,7 @@ import { withSearchDiagnostics, } from '@/lib/knowledge/search/diagnostics' import { intersectWorkspaceSearchFilters } from '@/lib/knowledge/search/filters' +import { matchPassage } from '@/lib/knowledge/search/snippet' import { connectorDisplayName } from '@/lib/sim-search/connectors' import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection' @@ -33,8 +35,9 @@ const searchInputSchema = workspaceSearchFiltersSchema.extend({ }) const readInputSchema = z.object({ documentId: z.string().min(1).max(200), - offset: z.number().int().min(0).max(5000).default(0), - limit: z.number().int().min(1).max(50).default(20), + limit: z.number().int().min(1).max(8).default(3), + startChunkIndex: z.number().int().min(0).max(2147483647).optional(), + startOffset: z.number().int().min(0).max(2147483647).optional(), }) const CITATION_INSTRUCTION = @@ -46,6 +49,7 @@ export const searchWorkspaceServerTool: BaseServerTool = { return withSearchDiagnostics( { surface: context?.searchSurface ?? 'copilot', + operation: 'search_workspace', toolCallId: context?.toolCallId, executionId: context?.executionId, }, @@ -63,9 +67,11 @@ export const searchWorkspaceServerTool: BaseServerTool = { message: 'Search query contains protected content. Rephrase the query.', } } + const safeQuery = projected.value const input = { - query: projected.value, + query: safeQuery, topK, + allowPartialResults: true, filters: intersectWorkspaceSearchFilters(requestedFilters, context?.assistantSearch), surface: context?.searchSurface ?? 'copilot', resultSecretRegistry: registry, @@ -87,38 +93,48 @@ export const searchWorkspaceServerTool: BaseServerTool = { const names = new Map(result.knowledgeBases.map((base) => [base.id, base.name])) const output = { success: true, - message: `Found ${result.results.length} passages. ${CITATION_INSTRUCTION}`, + message: `${result.retrieval.status === 'partial' ? 'Partial search: a retrieval branch reached its deadline. These results cannot establish absence or completeness. ' : ''}Found ${result.results.length} passage previews. Read a document at its chunkIndex for more context. ${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, + retrieval: result.retrieval, + results: result.results.map((item) => { + const content = projectResolvedSecretModelContent(item.content, registry) + if (!content.safe || typeof content.value !== 'string') + throw new Error('Knowledge result provenance is unavailable') + return { 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, - baseUrl: getBaseUrl(), - }), - })), + connectorType: item.connectorType, + sourceModifiedAt: item.sourceModifiedAt?.toISOString() ?? null, + author: sourceAuthor(item.metadata), + ...matchPassage(content.value, safeQuery, 1200), + 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), + originalPassageBytes: result.results.reduce( + (total, item) => total + Buffer.byteLength(item.content), + 0 + ), maxPassageBytes: Math.max(0, ...passageBytes), uniqueDocumentCount: new Set(output.data.results.map((item) => item.documentId)).size, }) @@ -128,10 +144,13 @@ export const searchWorkspaceServerTool: BaseServerTool = { logger.error('Workspace search failed', { error }) return { success: false, + retryable: error instanceof SearchDeadlineError, message: - error instanceof z.ZodError - ? 'Invalid search arguments' - : messageForCopilotKnowledgeError(error), + error instanceof SearchDeadlineError + ? error.message + : error instanceof z.ZodError + ? 'Invalid search arguments' + : messageForCopilotKnowledgeError(error), } } } @@ -142,50 +161,69 @@ export const searchWorkspaceServerTool: BaseServerTool = { export const readDocumentServerTool: BaseServerTool = { name: 'read_document', async execute(raw, context?: ServerToolContext) { - try { - const scope = requireCopilotKnowledgeScope(context) - const input = readInputSchema.parse(raw) - const registry = context?.resolvedSecretTraceRegistry - if (!registry) throw new Error('Knowledge result provenance is unavailable') - const readInput = { - ...input, - ...(scope.kind === 'organization' - ? { assertedOrganizationId: scope.organizationId } - : { assertedWorkspaceId: scope.workspaceId }), - filters: intersectWorkspaceSearchFilters( - { documentIds: [input.documentId] }, - context?.assistantSearch - ), - resultSecretRegistry: registry, - signal: context?.abortSignal, - } - const result = - scope.kind === 'organization' - ? await executeCopilotOrganizationKnowledgeUseCase(context, readSearchDocument, readInput) - : await executeCopilotKnowledgeUseCase(context, readSearchDocument, readInput) - return { - success: true, - message: CITATION_INSTRUCTION, - data: { - ...result, - ...createKnowledgeDocumentCitation({ - scope, - knowledgeBaseId: result.knowledgeBaseId, - documentId: result.documentId, - sourceUrl: result.sourceUrl, - baseUrl: getBaseUrl(), - }), - }, - } - } catch (error) { - logger.error('Document read failed', { error }) - return { - success: false, - message: - error instanceof z.ZodError - ? 'Invalid document arguments' - : messageForCopilotKnowledgeError(error), + return withSearchDiagnostics( + { + surface: context?.searchSurface ?? 'copilot', + toolCallId: context?.toolCallId, + executionId: context?.executionId, + operation: 'read_document', + }, + async () => { + try { + const scope = requireCopilotKnowledgeScope(context) + const input = readInputSchema.parse(raw) + const registry = context?.resolvedSecretTraceRegistry + if (!registry) throw new Error('Knowledge result provenance is unavailable') + const readInput = { + ...input, + ...(scope.kind === 'organization' + ? { assertedOrganizationId: scope.organizationId } + : { assertedWorkspaceId: scope.workspaceId }), + filters: intersectWorkspaceSearchFilters( + { documentIds: [input.documentId] }, + context?.assistantSearch + ), + resultSecretRegistry: registry, + signal: context?.abortSignal, + } + const result = await measureSearchStage('document_read', () => + scope.kind === 'organization' + ? executeCopilotOrganizationKnowledgeUseCase(context, readSearchDocument, readInput) + : executeCopilotKnowledgeUseCase(context, readSearchDocument, readInput) + ) + const output = { + success: true, + message: CITATION_INSTRUCTION, + data: { + ...result, + ...createKnowledgeDocumentCitation({ + scope, + knowledgeBaseId: result.knowledgeBaseId, + documentId: result.documentId, + sourceUrl: result.sourceUrl, + baseUrl: getBaseUrl(), + }), + }, + } + annotateSearchDiagnostics({ + toolResultBytes: Buffer.byteLength(JSON.stringify(output)), + passageBytes: result.chunks.reduce( + (total, chunk) => total + Buffer.byteLength(chunk.content), + 0 + ), + }) + return output + } catch (error) { + logger.error('Document read failed', { error }) + return { + success: false, + message: + error instanceof z.ZodError + ? 'Invalid document arguments' + : messageForCopilotKnowledgeError(error), + } + } } - } + ) }, } diff --git a/apps/sim/lib/knowledge/__integration__/gitlab-live.integration.ts b/apps/sim/lib/knowledge/__integration__/gitlab-live.integration.ts index 0ffcd293d9f..e515978e548 100644 --- a/apps/sim/lib/knowledge/__integration__/gitlab-live.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/gitlab-live.integration.ts @@ -829,8 +829,7 @@ describe.skipIf(!fixtureFile)('live self-hosted GitLab ingestion and permission input: { documentId: ordinary.id, assertedOrganizationId: ids.organizationId, - offset: 0, - limit: 10, + limit: 3, resultSecretRegistry: new ResolvedSecretTraceRegistry(), }, }) diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index 5646bd707d1..b1792cee86d 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -6,6 +6,7 @@ import { credential, credentialGroup, document, + embedding, knowledgeBase, knowledgeConnector, member, @@ -18,13 +19,27 @@ 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 type { + MothershipStreamV1CheckpointPausePayload, + MothershipStreamV1ToolCallDescriptor, +} from '@/lib/copilot/generated/mothership-stream-v1' +import { isContractStreamEventEnvelope } from '@/lib/copilot/request/session/contract' +import { + readDocumentServerTool, + 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 { + SearchBudget, + SearchDeadlineError, + type SearchExecutor, +} from '@/lib/knowledge/search/budget' +import type { SearchStage } from '@/lib/knowledge/search/diagnostics' import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' @@ -39,6 +54,7 @@ vi.hoisted(() => { } }) +const externalFetch = globalThis.fetch const enabled = process.env.KNOWLEDGE_SEARCH_PERFORMANCE_TEST === 'true' const chunkCount = Number(process.env.KNOWLEDGE_SEARCH_PERFORMANCE_CHUNKS ?? 20_000) const dimensions = 1536 @@ -400,7 +416,106 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu await db.$client.end() }, 120_000) + it('cancels slow SQL on the server and restores pooled connection settings', async () => { + const budget = new SearchBudget('keyword', performance.now() + 200) + const started = performance.now() + await expect( + budget.query('keyword.sql', (tx) => tx.execute(sql`SELECT pg_sleep(5)`)) + ).rejects.toBeInstanceOf(SearchDeadlineError) + expect(performance.now() - started).toBeLessThan(2000) + const [active] = await db.execute<{ count: number }>( + sql`SELECT count(*)::int AS count FROM pg_stat_activity WHERE pid <> pg_backend_pid() AND state = 'active' AND query = 'SELECT pg_sleep(5)'` + ) + expect(active.count).toBe(0) + const [settings] = await db.execute<{ timeout: string }>( + sql`SELECT current_setting('statement_timeout') AS timeout` + ) + expect(settings.timeout).toBe('0') + }) + + it('expires waiting for a saturated pool without executing abandoned work', async () => { + let release!: () => void + const released = new Promise((resolve) => { + release = resolve + }) + let markSaturated!: () => void + const saturated = new Promise((resolve) => { + markSaturated = resolve + }) + let acquired = 0 + const holders = Array.from({ length: db.$client.options.max }, () => + db.transaction(async () => { + if (++acquired === db.$client.options.max) markSaturated() + await released + }) + ) + const run = vi.fn((tx: SearchExecutor) => tx.execute(sql`SELECT 1`)) + const started = performance.now() + try { + await saturated + const budget = new SearchBudget('keyword', performance.now() + 200) + await expect(budget.query('keyword.sql', run)).rejects.toBeInstanceOf(SearchDeadlineError) + expect(performance.now() - started).toBeLessThan(2000) + } finally { + release() + await Promise.all(holders) + } + await db.execute(sql`SELECT 1`) + expect(run).not.toHaveBeenCalled() + }) + + it.each(['keyword', 'both'] as const)( + 'handles %s SQL branches exceeding the deadline explicitly', + async (delayedLegs) => { + const query = SearchBudget.prototype.query + const delayed = vi.spyOn(SearchBudget.prototype, 'query').mockImplementation(function ( + this: SearchBudget, + stage: SearchStage, + run: (executor: SearchExecutor) => PromiseLike + ): Promise { + return query.call(this, stage, async (tx) => { + if (delayedLegs === 'both' || this.leg === 'keyword') + await tx.execute(sql`SELECT pg_sleep(9)`) + return run(tx) + }) as Promise + }) + try { + const result = await searchWorkspaceServerTool.execute( + { query: 'Orion deployment', topK: 15 }, + { + userId: ids.aliceId, + workspaceId: ids.workspaceId, + toolCallId: generateId(), + copilotToolExecution: true, + resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry([], { + userId: ids.aliceId, + workspaceId: ids.workspaceId, + }), + } + ) + if (delayedLegs === 'both') { + expect(result).toMatchObject({ success: false, retryable: true }) + return + } + expect(result).toMatchObject({ + success: true, + data: { retrieval: { status: 'partial', timedOutLegs: ['keyword'] } }, + }) + const parsed = resultSchema.parse(result) + expect(parsed.data.results.length).toBeGreaterThan(0) + expect( + parsed.data.results.every((row) => row.knowledgeBaseId === ids.knowledgeBaseId) + ).toBe(true) + report['deadline.partial'] = { resultCount: parsed.data.results.length } + } finally { + delayed.mockRestore() + } + }, + 30_000 + ) + it('records first and repeated application searches with the actual SQL plans', async () => { + const before = embeddingCalls for (let iteration = 0; iteration < 2; iteration++) { const { result, plans } = await sample(`broad.${iteration}`, () => search()) expect(result.data.results).toHaveLength(15) @@ -412,7 +527,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu expect(vectorPlans).toHaveLength(1) expect(usesVectorIndex(vectorPlans[0].plan[0].Plan)).toBe(true) } - expect(embeddingCalls).toBe(2) + expect(embeddingCalls - before).toBe(2) }, 180_000) it('compares the Search tab and Assistant with the same person, query and index', async () => { @@ -555,4 +670,171 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu true ) }, 180_000) + /** Opt in with local Sim and Go URLs; uses the real configured provider, billing adapter, and async resume protocol. */ + it.skipIf(!process.env.KNOWLEDGE_SEARCH_ASSISTANT_URL)( + 'answers through local Go Assistant with progressive reads and citations', + async () => { + const assistantUrl = new URL(process.env.KNOWLEDGE_SEARCH_ASSISTANT_URL!) + const simUrl = new URL(process.env.KNOWLEDGE_SEARCH_SIM_URL!) + for (const url of [assistantUrl, simUrl]) { + if (!['127.0.0.1', 'localhost'].includes(url.hostname)) + throw new Error( + 'Assistant integration requires local servers using the disposable test databases' + ) + } + const apiKey = process.env.KNOWLEDGE_SEARCH_ASSISTANT_API_KEY! + const internalKey = process.env.KNOWLEDGE_SEARCH_SIM_INTERNAL_KEY! + expect(apiKey).toBeTruthy() + expect(internalKey).toBeTruthy() + const chunkId = `${ids.workspaceId}-chunk-0` + const [original] = await db + .select({ + content: embedding.content, + contentLength: embedding.contentLength, + tokenCount: embedding.tokenCount, + }) + .from(embedding) + .where(eq(embedding.id, chunkId)) + .limit(1) + const longContent = + 'Orion deployment guide. The activation phrase and final checksum appear at the end.\n' + + 'Review the deployment stages in order. Preserve the rollback procedure.\n'.repeat(320) + + '\nActivation phrase: SILVER COMET\nFinal checksum: K7M2-84\n' + const registry = new ResolvedSecretTraceRegistry([], { userId: ids.aliceId }) + const calls: Array<{ + name: string + arguments: unknown + milliseconds: number + bytes: number + }> = [] + let answer = '' + const started = performance.now() + try { + await db + .update(embedding) + .set({ + content: longContent, + contentLength: longContent.length, + tokenCount: Math.ceil(longContent.length / 4), + }) + .where(eq(embedding.id, chunkId)) + const admission = await externalFetch(new URL('/api/copilot/api-keys/validate', simUrl), { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-api-key': internalKey, + 'x-sim-billing-protocol': 'legacy-v0', + }, + body: JSON.stringify({ + userId: ids.aliceId, + organizationId: ids.organizationId, + chatId: organizationChatId, + }), + }) + expect(admission.status, await admission.text()).toBe(200) + let path = '/api/mothership' + let body: Record = { + message: + 'Search for the Orion deployment guide that mentions an activation phrase and final checksum. Read enough of that document to report both values and cite it.', + version: '3.0.0', + mode: 'assistant', + userId: ids.aliceId, + organizationId: ids.organizationId, + chatId: organizationChatId, + } + for (let round = 0; round < 8; round++) { + const response = await externalFetch(new URL(path, assistantUrl), { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-api-key': apiKey, + 'x-sim-billing-protocol': 'legacy-v0', + }, + body: JSON.stringify(body), + signal: AbortSignal.timeout(120000), + }) + const wire = await response.text() + expect(response.status, wire.slice(0, 2000)).toBe(200) + expect(Buffer.byteLength(wire)).toBeLessThan(2 * 1024 * 1024) + const pending = new Map() + let checkpoint: MothershipStreamV1CheckpointPausePayload | undefined + let streamId = '' + for (const line of wire.split('\n')) { + if (!line.startsWith('data:')) continue + const raw = line.slice(5).trim() + if (!raw || raw === '[DONE]') continue + const event: unknown = JSON.parse(raw) + if (!isContractStreamEventEnvelope(event)) + throw new Error('Assistant returned an invalid generated stream envelope') + streamId = event.stream.streamId + if (event.type === 'error') throw new Error(JSON.stringify(event.payload)) + if (event.type === 'text') answer += event.payload.text + if ( + event.type === 'tool' && + event.payload.phase === 'call' && + !event.payload.partial && + event.payload.arguments + ) + pending.set(event.payload.toolCallId, event.payload) + if (event.type === 'run' && event.payload.kind === 'checkpoint_pause') + checkpoint = event.payload + } + if (!checkpoint) break + const results = await Promise.all( + checkpoint.pendingToolCallIds.map(async (callId) => { + const call = pending.get(callId) + if (!call) throw new Error('Checkpoint referenced an absent tool call') + const tool = + call.toolName === 'search_workspace' + ? searchWorkspaceServerTool + : call.toolName === 'read_document' + ? readDocumentServerTool + : undefined + if (!tool) throw new Error(`Unexpected Assistant tool: ${call.toolName}`) + const toolStarted = performance.now() + const result = await tool.execute(call.arguments, { + userId: ids.aliceId, + organizationId: ids.organizationId, + chatId: organizationChatId, + toolCallId: callId, + copilotToolExecution: true, + requestMode: 'assistant', + resolvedSecretTraceRegistry: registry, + }) + calls.push({ + name: call.toolName, + arguments: call.arguments, + milliseconds: performance.now() - toolStarted, + bytes: Buffer.byteLength(JSON.stringify(result)), + }) + const { success } = z.object({ success: z.boolean() }).parse(result) + return { callId, name: call.toolName, success, data: result } + }) + ) + path = '/api/tools/resume' + body = { + checkpointId: checkpoint.checkpointId, + streamId, + userId: ids.aliceId, + organizationId: ids.organizationId, + chatId: organizationChatId, + results, + } + report['assistant.live.progress'] = { rounds: round + 1, calls } + saveReport() + } + report['assistant.live'] = { milliseconds: performance.now() - started, calls, answer } + saveReport() + expect(answer).toContain('SILVER COMET') + expect(answer).toContain('K7M2-84') + expect(answer).toContain('') + expect(calls.some((call) => call.name === 'read_document')).toBe(true) + expect(calls.some((call) => call.name === 'search_workspace')).toBe(true) + expect(calls.every((call) => call.bytes < 40000)).toBe(true) + } finally { + await db.update(embedding).set(original).where(eq(embedding.id, chunkId)) + } + }, + 10 * 60_000 + ) }) diff --git a/apps/sim/lib/knowledge/application/read-search-document.test.ts b/apps/sim/lib/knowledge/application/read-search-document.test.ts index ddd93ad5b0a..8a7212c8820 100644 --- a/apps/sim/lib/knowledge/application/read-search-document.test.ts +++ b/apps/sim/lib/knowledge/application/read-search-document.test.ts @@ -56,8 +56,7 @@ const context = { const input = { documentId: 'doc', assertedWorkspaceId: 'workspace', - offset: 0, - limit: 20, + limit: 3, filters: { source: 'slack', documentIds: ['doc'] }, resultSecretRegistry: new ResolvedSecretTraceRegistry([], { userId: 'reader', @@ -89,7 +88,7 @@ describe('Assistant document read', () => { await expect(readSearchDocument.execute({ principal, input })).resolves.toMatchObject({ documentId: 'doc', chunks: [{ content: 'body', chunkIndex: 0 }], - nextOffset: null, + next: null, }) expect(mocks.context).toHaveBeenCalledWith({ ...input, knowledgeBaseId: 'index' }, principal) expect(mocks.chunks).toHaveBeenCalledWith( @@ -97,8 +96,7 @@ describe('Assistant document read', () => { expect.objectContaining({ documentFilters: input.filters, enabled: 'true', - offset: 0, - limit: 20, + limit: 3, }), expect.any(String), access @@ -194,3 +192,81 @@ describe('organization Search document reads', () => { expect(mocks.chunks).not.toHaveBeenCalled() }) }) + +describe('precise bounded passage expansion', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.permission.mockResolvedValue('read') + mocks.context.mockResolvedValue(context) + mocks.provenance.mockResolvedValue({ imported: true, documentMetadata: {} }) + }) + + it('projects a secret spanning the page boundary before slicing it', async () => { + const secret = 'private-token-that-crosses-the-boundary' + const registry = new ResolvedSecretTraceRegistry([ + { name: 'TOKEN', plaintext: secret, encryptedValue: 'ciphertext' }, + ]) + mocks.provenance.mockImplementationOnce(async () => { + registry.recordResolved('TOKEN', secret) + return { imported: true, documentMetadata: {} } + }) + mocks.chunks.mockResolvedValue({ + chunks: [{ id: 'secret-chunk', chunkIndex: 0, content: `${'x'.repeat(7990) + secret}tail` }], + pagination: { total: 1, hasMore: false }, + }) + const result = await readSearchDocument.execute({ + principal, + input: { ...input, resultSecretRegistry: registry }, + }) + expect(result.chunks[0].content).toContain('{{TOKEN}}') + expect(JSON.stringify(result)).not.toContain('private-token') + }) + + it('continues a long chunk before advancing across disabled chunk gaps', async () => { + const content = 'Evidence πŸ”Ž\n'.repeat(1100) + mocks.chunks.mockResolvedValue({ + chunks: [ + { id: 'c7', chunkIndex: 7, content }, + { id: 'c11', chunkIndex: 11, content: 'next enabled passage' }, + ], + pagination: { total: 2, hasMore: false }, + }) + const first = await readSearchDocument.execute({ + principal, + input: { ...input, startChunkIndex: 7 }, + }) + expect(first.chunks).toHaveLength(1) + expect(first.chunks[0].content.length).toBeLessThanOrEqual(8000) + expect(first.next).toEqual({ startChunkIndex: 7, startOffset: first.chunks[0].endOffset }) + const second = await readSearchDocument.execute({ + principal, + input: { ...input, ...first.next! }, + }) + expect(first.chunks[0].content + second.chunks[0].content).toBe(content) + expect(second.chunks[1].chunkIndex).toBe(11) + expect(second.next).toBeNull() + expect(mocks.chunks).toHaveBeenCalledWith( + 'doc', + expect.objectContaining({ startChunkIndex: 7, requireEnabledDocument: true }), + expect.any(String), + access + ) + }) + + it('rejects positions without an anchor and stale within-chunk continuation', async () => { + await expect( + readSearchDocument.execute({ principal, input: { ...input, startOffset: 2 } }) + ).rejects.toThrow('startOffset requires startChunkIndex') + expect(mocks.chunks).not.toHaveBeenCalled() + mocks.chunks.mockResolvedValue({ + chunks: [{ id: 'c8', chunkIndex: 8, content: 'replacement' }], + pagination: { total: 1, hasMore: false }, + }) + await expect( + readSearchDocument.execute({ + principal, + input: { ...input, startChunkIndex: 7, startOffset: 100 }, + }) + ).rejects.toThrow('no longer available') + }) +}) diff --git a/apps/sim/lib/knowledge/application/read-search-document.ts b/apps/sim/lib/knowledge/application/read-search-document.ts index ac48bcdbbe3..9d5dfa8116d 100644 --- a/apps/sim/lib/knowledge/application/read-search-document.ts +++ b/apps/sim/lib/knowledge/application/read-search-document.ts @@ -8,9 +8,12 @@ import { KnowledgeDocumentNotReadyError } from '@/lib/knowledge/application/chun import { resolveCanonicalActiveKnowledgeDocumentContext } from '@/lib/knowledge/application/contexts' import { knowledgeOperations } from '@/lib/knowledge/application/operations' import { queryChunks } from '@/lib/knowledge/chunks/service' +import { measureSearchStage } from '@/lib/knowledge/search/diagnostics' import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters' import { findSearchIndex } from '@/lib/knowledge/search/search-index' +import { passageWindow } from '@/lib/knowledge/search/snippet' import { importKnowledgeSearchResultSecretProvenance } from '@/lib/knowledge/secret-provenance' +import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection' import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' export interface ReadSearchDocumentInput { @@ -18,12 +21,17 @@ export interface ReadSearchDocumentInput { assertedWorkspaceId?: string assertedOrganizationId?: string filters?: WorkspaceSearchFilters - offset: number limit: number + startChunkIndex?: number + startOffset?: number resultSecretRegistry: ResolvedSecretTraceRegistry signal?: AbortSignal } +/** Bounds model text to at most 24KB of UTF-8, with continuation even inside a large chunk. */ +const READ_PAGE_CHARACTERS = 8000 +const READ_PAGE_CHUNKS = 8 + /** Reads enabled indexed passages with the same document scope and ACLs as search. */ export const readSearchDocument = defineAuthorizedKnowledgeUseCase({ operation: knowledgeOperations.readDocument, @@ -50,16 +58,22 @@ export const readSearchDocument = defineAuthorizedKnowledgeUseCase({ async execute({ input, context }): Promise { input.signal?.throwIfAborted() if ( - !Number.isInteger(input.offset) || - input.offset < 0 || - input.offset > 5000 || !Number.isInteger(input.limit) || input.limit < 1 || - input.limit > 50 + input.limit > READ_PAGE_CHUNKS || + (input.startChunkIndex !== undefined && + (!Number.isSafeInteger(input.startChunkIndex) || + input.startChunkIndex < 0 || + input.startChunkIndex > 2147483647)) || + (input.startOffset !== undefined && + (!Number.isSafeInteger(input.startOffset) || + input.startOffset < 0 || + input.startOffset > 2147483647 || + input.startChunkIndex === undefined)) ) { throw new OrchestrationError( 'validation', - 'Document reads require offset 0–5000 and limit 1–50' + 'Document reads require limit 1–8 and nonnegative chunk positions; startOffset requires startChunkIndex' ) } if (context.document.processingStatus !== 'completed') { @@ -68,24 +82,30 @@ export const readSearchDocument = defineAuthorizedKnowledgeUseCase({ if (!context.knowledgeBase.isSearchIndex) throw new OrchestrationError('not_found', 'Document not found') if (!context.document.enabled) throw new OrchestrationError('not_found', 'Document not found') - const page = await queryChunks( - context.documentId, - { - offset: input.offset, - limit: input.limit, - enabled: 'true', - sortBy: 'chunkIndex', - sortOrder: 'asc', - documentFilters: input.filters, - }, - generateRequestId(), - await context.access.get() + const access = await measureSearchStage('access_scope', () => context.access.get()) + const page = await measureSearchStage('document_read.sql', () => + queryChunks( + context.documentId, + { + limit: input.limit, + startChunkIndex: input.startChunkIndex, + requireEnabledDocument: true, + enabled: 'true', + sortBy: 'chunkIndex', + sortOrder: 'asc', + documentFilters: input.filters, + }, + generateRequestId(), + access + ) ) if (page.pagination.total === 0) throw new OrchestrationError('not_found', 'Document not found') - const provenance = await importKnowledgeSearchResultSecretProvenance({ - registry: input.resultSecretRegistry, - results: page.chunks.map((chunk) => ({ ...chunk, documentId: context.documentId })), - }) + const provenance = await measureSearchStage('result_provenance', () => + importKnowledgeSearchResultSecretProvenance({ + registry: input.resultSecretRegistry, + results: page.chunks.map((chunk) => ({ ...chunk, documentId: context.documentId })), + }) + ) if (!provenance.imported) throw new Error('Knowledge result provenance is unavailable') const metadata = provenance.documentMetadata[context.documentId] if ( @@ -99,15 +119,53 @@ export const readSearchDocument = defineAuthorizedKnowledgeUseCase({ ) { throw new Error('Knowledge document provenance is unavailable') } + /** Project complete strings before slicing, so a window cannot expose part of a secret. */ + const projectedChunks = page.chunks.map((chunk) => { + const projected = projectResolvedSecretModelContent(chunk.content, input.resultSecretRegistry) + if (!projected.safe || typeof projected.value !== 'string') + throw new Error('Knowledge result provenance is unavailable') + return { ...chunk, content: projected.value } + }) + if ( + input.startOffset && + (projectedChunks[0]?.chunkIndex !== input.startChunkIndex || + input.startOffset >= projectedChunks[0].content.length) + ) { + throw new OrchestrationError( + 'validation', + 'The passage position is no longer available; search again or read from the chunk start' + ) + } + let remaining = READ_PAGE_CHARACTERS + const chunks: ReadSearchDocumentResult['chunks'] = [] + let next: ReadSearchDocumentResult['next'] = null + for (const chunk of projectedChunks) { + if (remaining < 2) { + next = { startChunkIndex: chunk.chunkIndex, startOffset: 0 } + break + } + const start = chunk.chunkIndex === input.startChunkIndex ? (input.startOffset ?? 0) : 0 + const excerpt = passageWindow(chunk.content, start, remaining) + chunks.push({ chunkIndex: chunk.chunkIndex, ...excerpt }) + remaining -= excerpt.content.length + if (excerpt.endOffset < chunk.content.length) { + next = { startChunkIndex: chunk.chunkIndex, startOffset: excerpt.endOffset } + break + } + } + const last = chunks.at(-1) + if (!next && page.pagination.hasMore && last) { + next = { startChunkIndex: last.chunkIndex + 1, startOffset: 0 } + } input.signal?.throwIfAborted() return { documentId: context.documentId, knowledgeBaseId: context.knowledgeBaseId, documentName: metadata?.filename ?? null, sourceUrl: metadata?.sourceUrl ?? null, - chunks: page.chunks.map(({ content, chunkIndex }) => ({ content, chunkIndex })), - hasMore: page.pagination.hasMore, - nextOffset: page.pagination.hasMore ? input.offset + page.chunks.length : null, + chunks, + hasMore: next !== null, + next, } }, }) diff --git a/apps/sim/lib/knowledge/application/search.test.ts b/apps/sim/lib/knowledge/application/search.test.ts index 49aa6d88584..9adf6797f47 100644 --- a/apps/sim/lib/knowledge/application/search.test.ts +++ b/apps/sim/lib/knowledge/application/search.test.ts @@ -89,7 +89,10 @@ vi.mock('@/lib/knowledge/embeddings', () => ({ vi.mock('@/lib/knowledge/search/queries', () => ({ generateSearchEmbedding: mocks.generateEmbedding, - executeKnowledgeSearch: mocks.executeSearch, + retrieveKnowledgeSearch: async (...args: unknown[]) => ({ + rows: await mocks.executeSearch(...args), + retrieval: { status: 'complete', timedOutLegs: [] }, + }), getDocumentMetadataByIds: mocks.getDocumentMetadata, })) diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index 1bcd42670b8..b67f2fef3fc 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -41,12 +41,14 @@ import { runWithKnowledgeModelInputProvenance } from '@/lib/knowledge/model-inpu import { rerank } from '@/lib/knowledge/reranker' import type { RerankerStatus } from '@/lib/knowledge/reranker-models' import { recordOrganizationSearchActivity } from '@/lib/knowledge/search/activity' +import { SearchDeadlineError } from '@/lib/knowledge/search/budget' 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, getDocumentMetadataByIds, + type RetrievalStatus, + retrieveKnowledgeSearch, type SearchResult, } from '@/lib/knowledge/search/queries' import { importKnowledgeSearchResultSecretProvenance } from '@/lib/knowledge/secret-provenance' @@ -88,6 +90,8 @@ export class KnowledgeSearchProvenanceUnavailableError extends Error { export type KnowledgeSearchTagFilter = KnowledgeTagNameFilter export interface SearchKnowledgeInput { + /** Only surfaces displaying retrieval status may accept incomplete evidence. */ + allowPartialResults?: boolean /** Optional assertion from a trusted adapter or public contract. */ workspaceId?: string organizationId?: string @@ -154,6 +158,7 @@ interface KnowledgeSearchCost { } export interface SearchKnowledgeResult { + retrieval: RetrievalStatus results: KnowledgeSearchItem[] query: string knowledgeBaseIds: string[] @@ -397,8 +402,8 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ ) : Math.min(KNOWLEDGE_SEARCH_COST_POLICY.maxTopK, input.topK * 4) : input.topK - let rows = await measureSearchStage('retrieval', () => - executeKnowledgeSearch({ + const retrieved = await measureSearchStage('retrieval', () => + retrieveKnowledgeSearch({ knowledgeBaseIds, topK: candidateTopK, filters: input.filters, @@ -418,6 +423,13 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ }) ) + if (retrieved.retrieval.status === 'partial' && !input.allowPartialResults) + throw new SearchDeadlineError() + annotateSearchDiagnostics({ + retrievalStatus: retrieved.retrieval.status, + timedOutLegs: retrieved.retrieval.timedOutLegs, + }) + let rows = retrieved.rows input.signal?.throwIfAborted() /** Public callers have no input envelope, but persisted reranker inputs still need provenance. */ const registrySubjectUserId = resolvePrincipalSubjectUserId(principal) @@ -697,6 +709,9 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ } } annotateSearchDiagnostics({ resultCount: results.length }) + if (retrieved.retrieval.status === 'partial' && results.length === 0) { + throw new SearchDeadlineError() + } const cost = baseCost ? { input: baseCost.input, @@ -715,6 +730,7 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ } : undefined return { + retrieval: retrieved.retrieval, results, query: input.query ?? '', knowledgeBaseIds, diff --git a/apps/sim/lib/knowledge/application/workspace-search.ts b/apps/sim/lib/knowledge/application/workspace-search.ts index 983c7819011..81924891d5a 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.ts @@ -31,7 +31,13 @@ const searchWorkspaceKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ const index = await measureSearchStage('index_resolution', () => findWorkspaceSearchIndex(context.workspaceId) ) - if (!index) return { results: [], query: input.query ?? '', knowledgeBases: [] } + if (!index) + return { + results: [], + query: input.query ?? '', + knowledgeBases: [], + retrieval: { status: 'complete' as const, timedOutLegs: [] }, + } return searchKnowledge.execute({ principal, input: { ...input, workspaceId: context.workspaceId, knowledgeBaseIds: [index.id] }, @@ -75,7 +81,12 @@ const searchOrganizationKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ results: [], }) } - return { results: [], query: input.query ?? '', knowledgeBases: [] } + return { + results: [], + query: input.query ?? '', + knowledgeBases: [], + retrieval: { status: 'complete' as const, timedOutLegs: [] }, + } } return searchKnowledge.execute({ principal, @@ -118,7 +129,12 @@ const searchScopedKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ results: [], }) } - return { results: [], query: input.query ?? '', knowledgeBases: [] } + return { + results: [], + query: input.query ?? '', + knowledgeBases: [], + retrieval: { status: 'complete' as const, timedOutLegs: [] }, + } } return searchKnowledge.execute({ principal, diff --git a/apps/sim/lib/knowledge/chunks/service.ts b/apps/sim/lib/knowledge/chunks/service.ts index 02a018367c9..f471c084f81 100644 --- a/apps/sim/lib/knowledge/chunks/service.ts +++ b/apps/sim/lib/knowledge/chunks/service.ts @@ -3,7 +3,7 @@ import { document, embedding, knowledgeBase } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { sha256Hex } from '@sim/security/hash' import { generateId } from '@sim/utils/id' -import { and, eq, inArray, isNull, sql } from 'drizzle-orm' +import { and, eq, gte, inArray, isNull, sql } from 'drizzle-orm' import { type KeysetKey, keysetColumns, @@ -123,7 +123,13 @@ export async function queryChunks( * keyset resume narrows the *page*, and folding it into the count would turn * a total into a remainder that shrinks with every page. */ - const pageConditions = [...conditions, resumeKeyset(keys, cursorKeys, sortOrder)] + const pageConditions = [ + ...conditions, + resumeKeyset(keys, cursorKeys, sortOrder), + filters.startChunkIndex === undefined + ? undefined + : gte(embedding.chunkIndex, filters.startChunkIndex), + ] const rows = await db .select({ diff --git a/apps/sim/lib/knowledge/chunks/types.ts b/apps/sim/lib/knowledge/chunks/types.ts index 66c20011a18..4225a2e28e6 100644 --- a/apps/sim/lib/knowledge/chunks/types.ts +++ b/apps/sim/lib/knowledge/chunks/types.ts @@ -13,6 +13,8 @@ export interface ChunkFilters { enabled?: 'true' | 'false' | 'all' limit?: number offset?: number + /** Inclusive indexed position; unlike an ordinal offset it survives disabled chunk gaps. */ + startChunkIndex?: number sortBy?: ChunkSortBy sortOrder?: 'asc' | 'desc' /** Keyset position from a previous page. Never combined with `offset`. */ diff --git a/apps/sim/lib/knowledge/search/budget.test.ts b/apps/sim/lib/knowledge/search/budget.test.ts new file mode 100644 index 00000000000..d5e04601737 --- /dev/null +++ b/apps/sim/lib/knowledge/search/budget.test.ts @@ -0,0 +1,55 @@ +/** @vitest-environment node */ +import { db } from '@sim/db' +import { afterEach, describe, expect, it, vi } from 'vitest' +import { SearchBudget, SearchDeadlineError } from '@/lib/knowledge/search/budget' + +afterEach(() => vi.restoreAllMocks()) + +describe('search SQL deadline', () => { + it('does not start a fallback with a fresh budget', async () => { + let now = 0 + vi.spyOn(performance, 'now').mockImplementation(() => now) + const budget = new SearchBudget('vector', 100) + const run = vi.fn(async () => { + now = 101 + return ['candidate'] + }) + await budget.query('vector.ann', run) + await expect(budget.query('vector.exact', run)).rejects.toBeInstanceOf(SearchDeadlineError) + expect(run).toHaveBeenCalledTimes(1) + }) + + it('returns at the acquisition deadline and skips work when the queued transaction later starts', async () => { + vi.useFakeTimers() + try { + let begin: (() => void) | undefined + let finished: Promise | undefined + vi.spyOn(db, 'transaction').mockImplementation((callback) => { + finished = new Promise((resolve) => { + begin = resolve + }).then(() => callback(db as never)) + return finished as ReturnType + }) + const run = vi.fn(async () => ['private result']) + const budget = new SearchBudget('keyword', performance.now() + 50) + const pending = budget.query('keyword.sql', run) + const assertion = expect(pending).rejects.toBeInstanceOf(SearchDeadlineError) + await vi.advanceTimersByTimeAsync(60) + await assertion + begin!() + await expect(finished).rejects.toBeInstanceOf(SearchDeadlineError) + expect(run).not.toHaveBeenCalled() + } finally { + vi.useRealTimers() + } + }) + + it('preserves cancellation and unexpected errors instead of labeling them incomplete evidence', async () => { + const controller = new AbortController() + const budget = new SearchBudget('keyword', performance.now() - 1, controller.signal) + const denied = new Error('access revoked') + expect(budget.isTimeout(denied)).toBe(false) + controller.abort(denied) + expect(() => budget.isTimeout(new SearchDeadlineError())).toThrow(denied) + }) +}) diff --git a/apps/sim/lib/knowledge/search/budget.ts b/apps/sim/lib/knowledge/search/budget.ts new file mode 100644 index 00000000000..7f8a0eaeecf --- /dev/null +++ b/apps/sim/lib/knowledge/search/budget.ts @@ -0,0 +1,115 @@ +import { db } from '@sim/db' +import { getPostgresErrorCode } from '@sim/utils/errors' +import { sql } from 'drizzle-orm' +import type { DbTransaction } from '@/lib/db/types' +import { + measureSearchStage, + recordSearchStageDuration, + type SearchStage, +} from '@/lib/knowledge/search/diagnostics' + +export const SEARCH_RETRIEVAL_BUDGET_MS = 8000 +export type RetrievalLeg = 'vector' | 'keyword' | 'tags' +export type SearchExecutor = Pick + +export class SearchDeadlineError extends Error { + constructor() { + super('Search reached its retrieval deadline. Retry with a narrower query or source filter.') + this.name = 'SearchDeadlineError' + } +} + +/** One leg's state, using the hybrid request's shared deadline across every refill and fallback. */ +export class SearchBudget { + timedOut = false + + constructor( + readonly leg: RetrievalLeg, + readonly deadline: number, + readonly signal?: AbortSignal + ) {} + + remaining(): number { + this.signal?.throwIfAborted() + const remaining = Math.ceil(this.deadline - performance.now()) + if (remaining <= 0) { + this.timedOut = true + throw new SearchDeadlineError() + } + return remaining + } + + isTimeout(error: unknown): boolean { + this.signal?.throwIfAborted() + if (error instanceof SearchDeadlineError || getPostgresErrorCode(error) === '57014') { + this.timedOut = true + return true + } + return false + } + + /** + * Only acquisition is raced. An expired queued transaction rolls back before executing work. + * Once acquired, PostgreSQL cancels the statement and we await rollback/release before returning. + * This avoids postgres.js cancellation packets racing a subsequent query on a pooled connection. + */ + async query( + stage: SearchStage, + run: (executor: SearchExecutor) => PromiseLike + ): Promise { + const started = performance.now() + const remaining = this.remaining() + let acquired = false + let expired = false + let timer: ReturnType | undefined + let rejectAcquisition: (error: unknown) => void = () => {} + const onAbort = () => { + if (!acquired) { + expired = true + rejectAcquisition(this.signal?.reason) + } + } + const acquisition = new Promise((_, reject) => { + rejectAcquisition = reject + timer = setTimeout(() => { + expired = true + this.timedOut = true + reject(new SearchDeadlineError()) + }, remaining) + this.signal?.addEventListener('abort', onAbort, { once: true }) + }) + const work = db.transaction(async (tx) => { + acquired = true + clearTimeout(timer) + this.signal?.removeEventListener('abort', onAbort) + if (expired) throw new SearchDeadlineError() + recordSearchStageDuration(`${this.leg}.connection_acquire`, performance.now() - started) + const timeout = String(this.remaining()) + await tx.execute(sql`SELECT set_config('statement_timeout', ${timeout}, true)`) + this.remaining() + return measureSearchStage(stage, () => run(tx)) + }) + try { + const result = await Promise.race([work, acquisition]) + this.signal?.throwIfAborted() + return result + } catch (error) { + if (this.isTimeout(error)) throw new SearchDeadlineError() + throw error + } finally { + clearTimeout(timer) + this.signal?.removeEventListener('abort', onAbort) + if (!acquired) + recordSearchStageDuration(`${this.leg}.connection_acquire`, performance.now() - started) + } + } +} + +/** Execute SQL with stage diagnostics and, for live search, the shared retrieval deadline. */ +export function runSearchQuery( + budget: SearchBudget | undefined, + stage: SearchStage, + run: (executor: SearchExecutor) => PromiseLike +): Promise { + return budget ? budget.query(stage, run) : measureSearchStage(stage, () => run(db)) +} diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index 09672e26a76..bc8e5c4e921 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -7,6 +7,8 @@ const PROGRESS_INTERVAL_MS = 5000 type RetrievalLeg = 'vector' | 'keyword' | 'tags' export type SearchStage = + | 'document_read' + | 'document_read.sql' | 'tool_input' | 'tool_application' | 'tool_presentation' @@ -40,7 +42,8 @@ export type SearchStage = | `${RetrievalLeg}.candidates` | `${RetrievalLeg}.authorization` | `${RetrievalLeg}.hydration` - | 'vector.connection_acquire' + | `${RetrievalLeg}.connection_acquire` + | `${RetrievalLeg}.sql` | 'vector.settings' | 'vector.probe' | 'vector.ann' @@ -48,6 +51,7 @@ export type SearchStage = /** Fixed, content-free fields. Never pass queries, filters, document identities, SQL, or errors. */ export interface SearchDiagnosticMetadata { + operation?: 'search_workspace' | 'read_document' surface?: 'dashboard' | 'mcp' | 'copilot' | 'workflow' | 'api' | 'slack' | 'other' toolCallId?: string executionId?: string @@ -68,6 +72,9 @@ export interface SearchDiagnosticMetadata { /** Tool output before the executor's final egress projection; counts only, never content. */ toolResultBytes?: number passageBytes?: number + originalPassageBytes?: number + retrievalStatus?: 'complete' | 'partial' + timedOutLegs?: RetrievalLeg[] maxPassageBytes?: number uniqueDocumentCount?: number } diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index e9beb8eccc7..a1b9728d06e 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -630,7 +630,7 @@ describe('live repository authorization follows ranked candidates', () => { 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) + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2) }) it('keeps the nearest exact results when an earlier ANN page hydrated only a farther result', async () => { diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 17becb53e74..91ff5d12957 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -9,6 +9,14 @@ import { } from '@/lib/knowledge/access/predicate' import type { KnowledgeAccessProvider, KnowledgeAccessScope } from '@/lib/knowledge/access/types' import type { KbEmbeddingDimensions } from '@/lib/knowledge/embedding-models' +import { + type RetrievalLeg, + runSearchQuery, + SEARCH_RETRIEVAL_BUDGET_MS, + SearchBudget, + SearchDeadlineError, + type SearchExecutor, +} from '@/lib/knowledge/search/budget' import { measureSearchStage, recordSearchStageDuration } from '@/lib/knowledge/search/diagnostics' import { workspaceSearchFilterConditions } from '@/lib/knowledge/search/filter-conditions' import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters' @@ -33,8 +41,6 @@ const HNSW_SETTINGS_UNSUPPORTED_RETRY_MS = 10 * 60 * 1000 let hnswSettingsUnsupportedUntil = 0 -type SearchExecutor = Pick - /** * Shared HNSW indexes can be selective on KB scope, access, or tags, even for * small workspace searches. Iterative scans keep looking within a bounded @@ -42,14 +48,17 @@ type SearchExecutor = Pick * older extensions retry without tuning until the compatibility cooldown expires. */ async function withVectorScanSettings( - run: (executor: SearchExecutor) => Promise + run: (executor: SearchExecutor) => Promise, + budget?: SearchBudget ): Promise { - if (Date.now() < hnswSettingsUnsupportedUntil) return run(db) + const untuned = () => runSearchQuery(budget, 'vector.ann', run) + if (Date.now() < hnswSettingsUnsupportedUntil) return untuned() const acquireStarted = performance.now() let applyingSettings = false try { - return await db.transaction(async (tx) => { - recordSearchStageDuration('vector.connection_acquire', performance.now() - acquireStarted) + const tuned = async (tx: SearchExecutor) => { + if (!budget) + recordSearchStageDuration('vector.connection_acquire', performance.now() - acquireStarted) applyingSettings = true await measureSearchStage('vector.settings', () => tx.execute( @@ -57,15 +66,20 @@ async function withVectorScanSettings( ) ) applyingSettings = false + if (budget) + await tx.execute( + sql`SELECT set_config('statement_timeout', ${String(budget.remaining())}, true)` + ) return run(tx) - }) + } + return await (budget ? budget.query('vector.ann', tuned) : db.transaction(tuned)) } catch (error) { if (!applyingSettings || getPostgresErrorCode(error) !== UNDEFINED_OBJECT_SQLSTATE) throw error hnswSettingsUnsupportedUntil = Date.now() + HNSW_SETTINGS_UNSUPPORTED_RETRY_MS logger.warn('pgvector iterative scan is unavailable; vector legs run without it', { error: getErrorMessage(error), }) - return run(db) + return untuned() } } @@ -186,6 +200,7 @@ export interface SearchParams { access: KnowledgeAccessScope accessProvider?: KnowledgeAccessProvider signal?: AbortSignal + budget?: SearchBudget structuredFilters?: StructuredFilter[] filters?: WorkspaceSearchFilters queryVector?: KnowledgeQueryVector @@ -475,6 +490,7 @@ async function selectAuthorizedSearchResults(input: { accessProvider: KnowledgeAccessProvider filters?: WorkspaceSearchFilters signal?: AbortSignal + budget?: SearchBudget topK: number selectPage: ( limit: number, @@ -491,77 +507,82 @@ async function selectAuthorizedSearchResults(input: { const considered = new Set() let scanned = 0 let offset = 0 - while ( - results.size < input.topK && - scanned < Number(HNSW_MAX_SCAN_TUPLES) && - Date.now() < deadline - ) { - input.signal?.throwIfAborted() - 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 && - input.filters.source !== 'github' && - input.filters.source !== 'confluence' - ? [] - : [ - ...new Set( - candidates.flatMap((candidate) => - candidate.connectorId ? [candidate.connectorId] : [] - ) - ), - ] - const access = await measureSearchStage(`${input.leg}.authorization`, () => - input.accessProvider.getForConnectors(connectorIds, input.signal) - ) - input.signal?.throwIfAborted() - const grantedSources = new Set( - access.kind === 'user' - ? [ - ...(access.githubInstallationGrants?.map((grant) => grant.connectorId) ?? []), - ...(access.confluenceSiteGrants?.map((grant) => grant.connectorId) ?? []), - ] - : [] - ) - const excludedBefore = excludedSources.size - for (const candidate of candidates) { - if ( - candidate.liveAuthorizationSource && - candidate.connectorId && - !grantedSources.has(candidate.connectorId) + try { + while ( + results.size < input.topK && + scanned < Number(HNSW_MAX_SCAN_TUPLES) && + (input.budget !== undefined || Date.now() < deadline) + ) { + input.signal?.throwIfAborted() + input.budget?.remaining() + const page = await measureSearchStage(`${input.leg}.candidates`, () => + input.selectPage(pageSize, offset, [...excludedSources]) ) - excludedSources.add(candidate.connectorId) - } - const hydrated = await measureSearchStage(`${input.leg}.hydration`, () => - input.hydrate( - candidates.map((candidate) => candidate.id), - access + 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 && + input.filters.source !== 'github' && + input.filters.source !== 'confluence' + ? [] + : [ + ...new Set( + candidates.flatMap((candidate) => + candidate.connectorId ? [candidate.connectorId] : [] + ) + ), + ] + const access = await measureSearchStage(`${input.leg}.authorization`, () => + input.accessProvider.getForConnectors(connectorIds, input.signal) ) - ) - 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 (!input.compareResults && results.size === input.topK) break - } - if (input.compareResults) { - const ranked = [...results.values()].sort(input.compareResults).slice(0, input.topK) - results.clear() - for (const row of ranked) results.set(row.id, row) + input.signal?.throwIfAborted() + const grantedSources = new Set( + access.kind === 'user' + ? [ + ...(access.githubInstallationGrants?.map((grant) => grant.connectorId) ?? []), + ...(access.confluenceSiteGrants?.map((grant) => grant.connectorId) ?? []), + ] + : [] + ) + const excludedBefore = excludedSources.size + for (const candidate of candidates) { + if ( + candidate.liveAuthorizationSource && + candidate.connectorId && + !grantedSources.has(candidate.connectorId) + ) + excludedSources.add(candidate.connectorId) + } + 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 (!input.compareResults && results.size === input.topK) break + } + if (input.compareResults) { + const ranked = [...results.values()].sort(input.compareResults).slice(0, input.topK) + results.clear() + for (const row of ranked) results.set(row.id, row) + } + if (excludedSources.size > excludedBefore) offset = 0 + else if (page.candidates.length < pageSize) break } - if (excludedSources.size > excludedBefore) offset = 0 - else if (page.candidates.length < pageSize) break + } catch (error) { + if (!input.budget?.isTimeout(error)) throw error } input.signal?.throwIfAborted() return [...results.values()] @@ -578,15 +599,19 @@ function hydrateSearchCandidates( access: KnowledgeAccessScope, distance: SQL | SQL.Aliased, filters: WorkspaceSearchFilters | undefined, - conditions: (SQL | undefined)[] + conditions: (SQL | undefined)[], + leg: RetrievalLeg, + budget?: SearchBudget ) { - return db - .select(getSearchResultFields(distance)) - .from(embedding) - .innerJoin(document, eq(embedding.documentId, document.id)) - .where( - and(inArray(embedding.id, ids), ...getVisibilityConditions(access, filters), ...conditions) - ) + return runSearchQuery(budget, `${leg}.sql`, (executor) => + executor + .select(getSearchResultFields(distance)) + .from(embedding) + .innerJoin(document, eq(embedding.documentId, document.id)) + .where( + and(inArray(embedding.id, ids), ...getVisibilityConditions(access, filters), ...conditions) + ) + ) } /** Candidates each hybrid leg retrieves before the fused list is trimmed to `topK`. */ @@ -629,26 +654,29 @@ export async function handleTagOnlySearch(params: SearchParams): Promise { - const candidates = await db - .select(SEARCH_READ_CANDIDATE_FIELDS) - .from(embedding) - .innerJoin(document, eq(embedding.documentId, document.id)) - .where( - and( - ...conditions, - ...getVisibilityConditions( - access, - params.filters, - knowledgeMetadataCandidateAccessCondition(access) - ), - excludeSearchSources(excludedSources) + const candidates = await runSearchQuery(params.budget, 'tags.sql', (executor) => + executor + .select(SEARCH_READ_CANDIDATE_FIELDS) + .from(embedding) + .innerJoin(document, eq(embedding.documentId, document.id)) + .where( + and( + ...conditions, + ...getVisibilityConditions( + access, + params.filters, + knowledgeMetadataCandidateAccessCondition(access) + ), + excludeSearchSources(excludedSources) + ) ) - ) - .orderBy(embedding.id) - .limit(limit) - .offset(offset) + .orderBy(embedding.id) + .limit(limit) + .offset(offset) + ) return { candidates, nextOffset: offset + candidates.length } }, hydrate: (ids, authorized) => @@ -657,7 +685,9 @@ export async function handleTagOnlySearch(params: SearchParams): Promise`0`.as('distance'), params.filters, - conditions + conditions, + 'tags', + params.budget ), }) } @@ -769,57 +799,58 @@ async function selectLiveVectorResults( accessProvider, filters: params.filters, signal: params.signal, + budget: params.budget, topK: params.topK, compareResults: (a, b) => a.distance - b.distance, - selectPage: (limit, offset, excludedSources) => - 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', () => + selectPage: async (limit, offset, excludedSources) => { + 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 runSearchQuery(params.budget, 'vector.exact', (executor) => executor - .select({ id: embedding.id }) + .select({ ...SEARCH_READ_CANDIDATE_FIELDS, distance: distance.as('distance') }) .from(embedding) .innerJoin(document, eq(embedding.documentId, document.id)) - .where(and(inArray(embedding.knowledgeBaseId, params.knowledgeBaseIds), ...visibility)) - .limit(LIVE_SEARCH_PAGE_SIZE) + .where( + and( + ...conditions, + ...visibility, + candidateIds ? inArray(embedding.id, candidateIds) : undefined + ) + ) + .orderBy(sql`(${distance}) + 0`, embedding.id) + .limit(limit) + .offset(exactOffset) ) - 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() + 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 runSearchQuery(params.budget, 'vector.probe', (executor) => + 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 ann = async (executor: SearchExecutor) => { const ranked = executor .select({ id: embedding.id, distance: distance.as('distance') }) .from(embedding) @@ -837,22 +868,30 @@ async function selectLiveVectorResults( .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, - } - }), + return 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) + } + const page = await withVectorScanSettings(ann, params.budget) + 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), + hydrateSearchCandidates( + ids, + authorized, + distance.as('distance'), + params.filters, + conditions, + 'vector', + params.budget + ), }) /** Relaxed HNSW scans can return adjacent pages out of distance order. */ return rows.sort((a, b) => a.distance - b.distance) @@ -893,6 +932,7 @@ export interface KeywordSearchParams { access: KnowledgeAccessScope accessProvider?: KnowledgeAccessProvider signal?: AbortSignal + budget?: SearchBudget query: string /** Query embedding, so keyword-only hits still carry a real cosine distance. */ queryVector: KnowledgeQueryVector @@ -946,26 +986,29 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise accessProvider: params.accessProvider, filters: params.filters, signal: params.signal, + budget: params.budget, topK, 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)) - .where( - and( - ...conditions, - ...getVisibilityConditions( - access, - params.filters, - knowledgeMetadataCandidateAccessCondition(access) - ), - excludeSearchSources(excludedSources) + const candidates = await runSearchQuery(params.budget, 'keyword.sql', (executor) => + executor + .select({ ...SEARCH_READ_CANDIDATE_FIELDS, keywordRank: rankExpr.as('keyword_rank') }) + .from(embedding) + .innerJoin(document, eq(embedding.documentId, document.id)) + .where( + and( + ...conditions, + ...getVisibilityConditions( + access, + params.filters, + knowledgeMetadataCandidateAccessCondition(access) + ), + excludeSearchSources(excludedSources) + ) ) - ) - .orderBy(sql`${rankExpr} DESC`, embedding.id) - .limit(limit) - .offset(offset) + .orderBy(sql`${rankExpr} DESC`, embedding.id) + .limit(limit) + .offset(offset) + ) return { candidates, nextOffset: offset + candidates.length } }, hydrate: (ids, authorized) => @@ -974,7 +1017,9 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise authorized, embeddingDistance(queryVector.dimensions, queryVector.vector).as('distance'), params.filters, - conditions + conditions, + 'keyword', + params.budget ), }) } @@ -1171,14 +1216,29 @@ export interface ExecuteKnowledgeSearchParams { filters?: WorkspaceSearchFilters } -/** - * Single retrieval entry point shared by the internal and v1 search routes. - * Callers remain responsible for auth, embedding generation, billing, and for - * rejecting requests that carry neither a query nor tag filters. - */ +export interface RetrievalStatus { + status: 'complete' | 'partial' + timedOutLegs: Array<'vector' | 'keyword' | 'tags'> +} + +export interface KnowledgeRetrievalResult { + rows: SearchResult[] + retrieval: RetrievalStatus +} + +/** Legacy surfaces cannot silently present partial retrieval as complete. */ export async function executeKnowledgeSearch( params: ExecuteKnowledgeSearchParams ): Promise { + const result = await retrieveKnowledgeSearch(params) + if (result.retrieval.status === 'partial') throw new SearchDeadlineError() + return result.rows +} + +/** Shared hybrid retrieval, with explicit completeness for surfaces that can represent it. */ +export async function retrieveKnowledgeSearch( + params: ExecuteKnowledgeSearchParams +): Promise { const { knowledgeBaseIds, topK, @@ -1189,99 +1249,69 @@ export async function executeKnowledgeSearch( access, boostRecency = false, } = params - + params.signal?.throwIfAborted() + const deadline = performance.now() + SEARCH_RETRIEVAL_BUDGET_MS + const budgets = { + vector: new SearchBudget('vector', deadline, params.signal), + keyword: new SearchBudget('keyword', deadline, params.signal), + tags: new SearchBudget('tags', deadline, params.signal), + } + const finish = (rows: SearchResult[]): KnowledgeRetrievalResult => { + params.signal?.throwIfAborted() + const timedOutLegs = Object.values(budgets) + .filter((budget) => budget.timedOut) + .map((budget) => budget.leg) + if (timedOutLegs.length && !rows.length) throw new SearchDeadlineError() + return { + rows: boostRecency ? applyRecencyBoost(rows) : rows, + retrieval: { status: timedOutLegs.length ? 'partial' : 'complete', timedOutLegs }, + } + } + const common = { + knowledgeBaseIds, + access, + accessProvider: params.accessProvider, + signal: params.signal, + filters: params.filters, + structuredFilters, + } const hasQuery = Boolean(query?.trim()) - const hasFilters = Boolean(structuredFilters && structuredFilters.length > 0) - + const hasFilters = Boolean(structuredFilters?.length) if (!hasQuery) { - if (!hasFilters) { - throw new Error('A search query or tag filters are required') - } - return await measureSearchStage('tags', () => - handleTagOnlySearch({ - knowledgeBaseIds, - topK, - structuredFilters, - access, - accessProvider: params.accessProvider, - signal: params.signal, - filters: params.filters, - }) + if (!hasFilters) throw new Error('A search query or tag filters are required') + return finish( + await measureSearchStage('tags', () => + handleTagOnlySearch({ ...common, topK, budget: budgets.tags }) + ) ) } - - if (!queryVector) { - throw new Error('Query vector is required when searching with a query') - } - + if (!queryVector) throw new Error('Query vector is required when searching with a query') const { distanceThreshold } = getQueryStrategy(knowledgeBaseIds.length, topK) - /** - * Hybrid fuses two rankings, so each leg retrieves more than the caller - * asked for: a chunk that both legs rank just below `topK` is a strong - * signal the fused list must be able to surface. - */ const legTopK = searchMode === 'hybrid' ? hybridCandidateCount(topK) : topK - + const vectorParams = { + ...common, + topK: legTopK, + queryVector, + distanceThreshold, + budget: budgets.vector, + } 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, - }) + hasFilters ? handleTagAndVectorSearch(vectorParams) : handleVectorOnlySearch(vectorParams) ) - if (searchMode === 'vector') { - const results = await vectorSearch - return boostRecency ? applyRecencyBoost(results) : results - } - - /** - * The lexical leg is best-effort: a failure there falls back to vector-only - * results rather than failing the whole search. - */ + if (searchMode === 'vector') return finish(await vectorSearch) const keywordSearch = measureSearchStage('keyword', () => executeKeywordSearch({ - knowledgeBaseIds, + ...common, 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'), + budget: budgets.keyword, }) - return [] as SearchResult[] - }) - - const [vectorResults, keywordResults] = await Promise.all([vectorSearch, keywordSearch]) - - /** - * Lexical leg first: on a total tie it wins, which is the behavior this mode - * exists for β€” an exact-token chunk the vector leg ranked below its distance - * threshold is precisely what a caller opted into hybrid to recover, and at - * `topK: 1` something has to win. - */ - const fused = fuseByReciprocalRank([keywordResults, vectorResults], topK) - return boostRecency ? applyRecencyBoost(fused) : fused + ) + const legs = await Promise.allSettled([vectorSearch, keywordSearch]) + /** Wait for both legs to release SQL resources; only deadline failures permit partial success. */ + for (const leg of legs) if (leg.status === 'rejected') throw leg.reason + const vectorResults = legs[0].status === 'fulfilled' ? legs[0].value : [] + const keywordResults = legs[1].status === 'fulfilled' ? legs[1].value : [] + return finish(fuseByReciprocalRank([keywordResults, vectorResults], topK)) } diff --git a/apps/sim/lib/knowledge/search/snippet.test.ts b/apps/sim/lib/knowledge/search/snippet.test.ts index 7dbec47ab5a..684e444d4d4 100644 --- a/apps/sim/lib/knowledge/search/snippet.test.ts +++ b/apps/sim/lib/knowledge/search/snippet.test.ts @@ -4,7 +4,9 @@ import { describe, expect, it } from 'vitest' import { findTermMatches, + matchPassage, matchSnippet, + passageWindow, queryTerms, SNIPPET_LENGTH, stripLeadingHeaders, @@ -138,3 +140,30 @@ describe('matchSnippet', () => { expect(snippet.endsWith('…')).toBe(true) }) }) + +describe('verbatim model passages', () => { + it('preserves formatting and exposes the location of evidence past the opening', () => { + const content = `Subject: Release\n\n${'Unrelated text. '.repeat(150)}\n\tdeploy_token = "orion"\n${'Trailing text. '.repeat(100)}` + const preview = matchPassage(content, 'orion', 1200) + expect(preview.content).toContain('\n\tdeploy_token = "orion"\n') + expect(preview.content).toBe(content.slice(preview.startOffset, preview.endOffset)) + expect(preview.content.length).toBeLessThanOrEqual(1200) + expect(preview.startOffset).toBeGreaterThan(0) + expect(preview.totalCharacters).toBe(content.length) + }) + + it('reassembles long Unicode chunks without splitting characters or losing text', () => { + const content = 'δΈ­Γ©πŸ”Ž\n code();\n'.repeat(1900) + let position = 0 + let rebuilt = '' + while (position < content.length) { + const page = passageWindow(content, position, 8000) + expect(page.content.isWellFormed()).toBe(true) + expect(Buffer.byteLength(page.content)).toBeLessThanOrEqual(24000) + expect(page.endOffset).toBeGreaterThan(position) + rebuilt += page.content + position = page.endOffset + } + expect(rebuilt).toBe(content) + }) +}) diff --git a/apps/sim/lib/knowledge/search/snippet.ts b/apps/sim/lib/knowledge/search/snippet.ts index 648e7c4060a..dbc866c77fb 100644 --- a/apps/sim/lib/knowledge/search/snippet.ts +++ b/apps/sim/lib/knowledge/search/snippet.ts @@ -97,6 +97,45 @@ function alignToCodePoint(text: string, index: number): number { return unit >= 0xdc00 && unit <= 0xdfff ? index - 1 : index } +export interface PassageExcerpt { + content: string + /** Positions are UTF-16 code units in the input chunk. */ + startOffset: number + endOffset: number + totalCharacters: number +} + +/** A bounded, verbatim window; offsets make omitted text explicitly recoverable. */ +export function passageWindow( + content: string, + startOffset: number, + maxCharacters: number +): PassageExcerpt { + const start = alignToCodePoint(content, Math.min(startOffset, content.length)) + const end = alignToCodePoint(content, Math.min(start + maxCharacters, content.length)) + return { + content: content.slice(start, end), + startOffset: start, + endOffset: end, + totalCharacters: content.length, + } +} + +/** Search evidence preserves source formatting and anchors on the same matches as Search's UI. */ +export function matchPassage( + content: string, + query: string, + maxCharacters: number +): PassageExcerpt { + if (content.length <= maxCharacters) return passageWindow(content, 0, maxCharacters) + const anchor = findTermMatches(content, queryTerms(query)).reduce( + (best, match) => (!best || match.length > best.length ? match : best), + undefined + ) + const start = anchor ? Math.max(0, anchor.index - Math.floor(maxCharacters / 3)) : 0 + return passageWindow(content, start, maxCharacters) +} + /** * The passage of a document a search result shows: a window around the longest * matching query term, keeping the earliest occurrence on ties. This favors a From 39fe2b85e1d80313ba4a14f6cbe27cbcfc8391e8 Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Fri, 11 Sep 2026 16:03:42 -0700 Subject: [PATCH 2/2] fix(search): protect query output and validate document continuation --- .../[knowledgeBaseId]/[documentId]/page.tsx | 29 +++++++++++-------- .../[documentId]/search-params.ts | 25 ++++++++++++++++ .../server/knowledge/workspace-search.test.ts | 20 +++++++++++-- .../server/knowledge/workspace-search.ts | 2 +- .../application/read-search-document.test.ts | 11 +++++++ .../application/read-search-document.ts | 10 ++++--- .../workspace-search.activity.test.ts | 1 + 7 files changed, 78 insertions(+), 20 deletions(-) create mode 100644 apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/search-params.ts diff --git a/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/page.tsx b/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/page.tsx index ebabd37f6fc..6cd806b52bc 100644 --- a/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/page.tsx +++ b/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/page.tsx @@ -1,16 +1,21 @@ import { ChipLink } from '@sim/emcn' import { notFound, redirect } from 'next/navigation' +import type { SearchParams } from 'nuqs/server' import { readSearchDocumentResultSchema } from '@/lib/api/contracts/knowledge/documents' import { getSession } from '@/lib/auth' import { OrchestrationError } from '@/lib/core/orchestration/types' import { readSearchDocument } from '@/lib/knowledge/application/read-search-document' import { buildAuthCrossLink } from '@/app/(auth)/auth-redirect' +import { + loadDocumentReadParams, + serializeDocumentReadParams, +} from '@/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/search-params' import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' interface OrganizationDocumentPageProps { params: Promise<{ organizationId: string; knowledgeBaseId: string; documentId: string }> - searchParams: Promise<{ offset?: string }> + searchParams: Promise } export default async function OrganizationDocumentPage({ @@ -18,15 +23,15 @@ export default async function OrganizationDocumentPage({ searchParams, }: OrganizationDocumentPageProps) { const { organizationId, knowledgeBaseId, documentId } = await params - const { offset: rawOffset } = await searchParams - const offset = rawOffset === undefined ? 0 : Number(rawOffset) - if (!Number.isInteger(offset) || offset < 0 || offset > 5000) notFound() + const position = await loadDocumentReadParams(searchParams, { strict: true }).catch(() => + notFound() + ) const href = `/o/${encodeURIComponent(organizationId)}/knowledge/${encodeURIComponent(knowledgeBaseId)}/${encodeURIComponent(documentId)}` const session = await getSession() if (!session?.user) { redirect( buildAuthCrossLink('/login', { - callbackUrl: offset ? `${href}?offset=${offset}` : href, + callbackUrl: serializeDocumentReadParams(href, position), isInviteFlow: false, }) ) @@ -39,15 +44,15 @@ export default async function OrganizationDocumentPage({ input: { documentId, assertedOrganizationId: organizationId, - offset, - limit: 20, + ...position, + limit: 3, resultSecretRegistry: registry, }, }) } catch (error) { if ( error instanceof OrchestrationError && - (error.code === 'not_found' || error.code === 'forbidden') + (error.code === 'not_found' || error.code === 'forbidden' || error.code === 'validation') ) notFound() throw error @@ -71,11 +76,11 @@ export default async function OrganizationDocumentPage({

))} diff --git a/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/search-params.ts b/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/search-params.ts new file mode 100644 index 00000000000..42f9dd224e8 --- /dev/null +++ b/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/search-params.ts @@ -0,0 +1,25 @@ +import { createLoader, createParser, createSerializer } from 'nuqs/server' + +const parseAsDocumentPosition = createParser({ + parse: (value) => { + if (!/^\d+$/.test(value)) return null + const position = Number(value) + return Number.isSafeInteger(position) && position <= 2147483647 ? position : null + }, + serialize: String, +}).withDefault(0) + +export const documentReadParams = { + startChunkIndex: parseAsDocumentPosition, + startOffset: parseAsDocumentPosition, +} + +const documentReadUrlKeys = { + urlKeys: { + startChunkIndex: 'start-chunk-index', + startOffset: 'start-offset', + }, +} as const + +export const loadDocumentReadParams = createLoader(documentReadParams, documentReadUrlKeys) +export const serializeDocumentReadParams = createSerializer(documentReadParams, documentReadUrlKeys) 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 7edc4ba294f..b83a791fc06 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 @@ -35,9 +35,6 @@ vi.mock('@/lib/knowledge/application/read-search-document', () => ({ execute: mocks.read, }, })) -vi.mock('@/executor/utils/resolved-secret-content-projection', () => ({ - projectResolvedSecretModelContent: (value: unknown) => ({ safe: true, value }), -})) import { readDocumentServerTool, @@ -188,6 +185,23 @@ describe('Assistant retrieval tools', () => { }) ) }) + it('returns only the projected query to the model', async () => { + const secret = 'private-resolved-query-token' + const registry = new ResolvedSecretTraceRegistry([ + { name: 'TOKEN', plaintext: secret, encryptedValue: 'ciphertext' }, + ]) + registry.recordResolved('TOKEN', secret) + const result = await searchWorkspaceServerTool.execute( + { query: `Find ${secret}` }, + { ...context, resolvedSecretTraceRegistry: registry } + ) + expect(result).toMatchObject({ success: true, data: { query: 'Find {{TOKEN}}' } }) + expect(mocks.search).toHaveBeenCalledWith( + expect.objectContaining({ input: expect.objectContaining({ query: 'Find {{TOKEN}}' }) }) + ) + expect(JSON.stringify(result)).not.toContain(secret) + }) + it.each([0, 20, 50])( 'measures UTF-8 bytes for %i passages without logging their content', async (count) => { 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 87cdebc8836..8b014c4d17c 100644 --- a/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.ts +++ b/apps/sim/lib/copilot/tools/server/knowledge/workspace-search.ts @@ -95,7 +95,7 @@ export const searchWorkspaceServerTool: BaseServerTool = { success: true, message: `${result.retrieval.status === 'partial' ? 'Partial search: a retrieval branch reached its deadline. These results cannot establish absence or completeness. ' : ''}Found ${result.results.length} passage previews. Read a document at its chunkIndex for more context. ${CITATION_INSTRUCTION}`, data: { - query, + query: safeQuery, retrieval: result.retrieval, results: result.results.map((item) => { const content = projectResolvedSecretModelContent(item.content, registry) diff --git a/apps/sim/lib/knowledge/application/read-search-document.test.ts b/apps/sim/lib/knowledge/application/read-search-document.test.ts index 8a7212c8820..ad3b3831013 100644 --- a/apps/sim/lib/knowledge/application/read-search-document.test.ts +++ b/apps/sim/lib/knowledge/application/read-search-document.test.ts @@ -253,6 +253,17 @@ describe('precise bounded passage expansion', () => { ) }) + it('rejects a continuation when all remaining chunks have disappeared', async () => { + mocks.chunks.mockResolvedValue({ + chunks: [], + pagination: { total: 2, hasMore: false }, + }) + await expect( + readSearchDocument.execute({ principal, input: { ...input, startChunkIndex: 7 } }) + ).rejects.toThrow('no longer available') + expect(mocks.provenance).not.toHaveBeenCalled() + }) + it('rejects positions without an anchor and stale within-chunk continuation', async () => { await expect( readSearchDocument.execute({ principal, input: { ...input, startOffset: 2 } }) diff --git a/apps/sim/lib/knowledge/application/read-search-document.ts b/apps/sim/lib/knowledge/application/read-search-document.ts index 9d5dfa8116d..dbc52ccc0b5 100644 --- a/apps/sim/lib/knowledge/application/read-search-document.ts +++ b/apps/sim/lib/knowledge/application/read-search-document.ts @@ -31,6 +31,8 @@ export interface ReadSearchDocumentInput { /** Bounds model text to at most 24KB of UTF-8, with continuation even inside a large chunk. */ const READ_PAGE_CHARACTERS = 8000 const READ_PAGE_CHUNKS = 8 +const STALE_POSITION_MESSAGE = + 'The passage position is no longer available; search again or read from the chunk start' /** Reads enabled indexed passages with the same document scope and ACLs as search. */ export const readSearchDocument = defineAuthorizedKnowledgeUseCase({ @@ -100,6 +102,9 @@ export const readSearchDocument = defineAuthorizedKnowledgeUseCase({ ) ) if (page.pagination.total === 0) throw new OrchestrationError('not_found', 'Document not found') + if (input.startChunkIndex !== undefined && page.chunks.length === 0) { + throw new OrchestrationError('validation', STALE_POSITION_MESSAGE) + } const provenance = await measureSearchStage('result_provenance', () => importKnowledgeSearchResultSecretProvenance({ registry: input.resultSecretRegistry, @@ -131,10 +136,7 @@ export const readSearchDocument = defineAuthorizedKnowledgeUseCase({ (projectedChunks[0]?.chunkIndex !== input.startChunkIndex || input.startOffset >= projectedChunks[0].content.length) ) { - throw new OrchestrationError( - 'validation', - 'The passage position is no longer available; search again or read from the chunk start' - ) + throw new OrchestrationError('validation', STALE_POSITION_MESSAGE) } let remaining = READ_PAGE_CHARACTERS const chunks: ReadSearchDocumentResult['chunks'] = [] diff --git a/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts b/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts index 32ecacf0059..301766ff25f 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts @@ -63,6 +63,7 @@ describe.each([ queueTableRows(member, [{ role: 'member' }]) expect(await operation.execute({ principal, input })).toEqual({ results: [], + retrieval: { status: 'complete', timedOutLegs: [] }, query: 'policy', knowledgeBases: [], })