fix(search): bound retrieval and progressively read document evidence (#7789)

* fix(search): bound retrieval and progressively read document evidence

* fix(search): protect query output and validate document continuation

Author: icecrasher321

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,
 }))
 

apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/page.tsx

@@ -1,32 +1,37 @@
 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<SearchParams>
 }
 
 export default async function OrganizationDocumentPage({
   params,
   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({
           </p>
         ))}
         <nav className='flex gap-2' aria-label='Document pages'>
-          {offset > 0 && (
-            <ChipLink href={`${href}?offset=${Math.max(0, offset - 20)}`}>Previous</ChipLink>
+          {(position.startChunkIndex > 0 || position.startOffset > 0) && (
+            <ChipLink href={href}>Start</ChipLink>
           )}
-          {document.nextOffset !== null && (
-            <ChipLink href={`${href}?offset=${document.nextOffset}`}>Next</ChipLink>
+          {document.next && (
+            <ChipLink href={serializeDocumentReadParams(href, document.next)}>Next</ChipLink>
           )}
         </nav>
       </article>

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)

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<typeof readSearchDocumentResultSchema>

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',

apps/sim/lib/copilot/generated/tool-schemas-v1.ts

@@ -4694,16 +4694,24 @@ export const TOOL_RUNTIME_SCHEMAS: Record<string, ToolRuntimeSchemaEntry> = {
           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<string, ToolRuntimeSchemaEntry> = {
         },
         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',

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',

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,
@@ -61,6 +58,7 @@ describe('Assistant retrieval tools', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     mocks.search.mockResolvedValue({
+      retrieval: { status: 'complete', timedOutLegs: [] },
       knowledgeBases: [{ id: 'index', name: 'Enterprise Search' }],
       results: [
         {
@@ -83,7 +81,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 () => {
@@ -187,11 +185,29 @@ 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) => {
       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 +233,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 +262,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 +310,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()
   })