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

Author: icecrasher321

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/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) => {

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)

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

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'] = []

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: [],
     })