fix(search): preserve query attribution on fused results (#9730)

Author: waleedlatif1

apps/sim/lib/api/contracts/mothership-assistant-tools.ts

@@ -458,6 +458,8 @@ export const workspaceKnowledgeSearchResultSchema = z.object({
   sourceDate: z.string().nullable().optional(),
   sourceContainerName: z.string().optional(),
   sourceContainerUrl: z.string().optional(),
+  /** Zero-based nativeQueries entries that returned this result, retained through rank fusion. */
+  queryIndexes: z.array(z.number().int().min(0).max(8)).min(1).max(9).optional(),
   sourceDateType: z.enum(['event_start', 'message', 'modified']).optional(),
   /** The person behind the document, from its author-like tag; null when the source names none. */
   author: z.string().nullable(),

apps/sim/lib/knowledge/mcp/server.protocol.test.ts

@@ -3,6 +3,7 @@ import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js'
 import { createPersonalApiKeyPrincipal } from '@sim/testing/factories/principal.factory'
 import { apiServerRoutesMock } from '@sim/testing/mocks/api-server-routes.mock'
 import { urlsMockFns } from '@sim/testing/mocks/urls.mock'
+import { toArray, toRecord } from '@sim/utils/object'
 import { NextRequest } from 'next/server'
 import { describe, expect, it, vi } from 'vitest'
 
@@ -39,6 +40,7 @@ describe('Search MCP protocol', () => {
       content: 'Evidence',
       chunkIndex: 0,
       similarity: 0.5,
+      queryIndexes: [0],
     }
     mocks.liveSearch.mockResolvedValue({
       results: [row],
@@ -74,12 +76,19 @@ describe('Search MCP protocol', () => {
       expect(searchSchema?.properties).toHaveProperty('nativeQueries')
       expect(readSchema?.properties).toHaveProperty('startOffset')
       expect(readSchema?.properties).not.toHaveProperty('url')
-      const search = await client.callTool({ name: 'search', arguments: { query: 'release' } })
+      const search = await client.callTool({
+        name: 'search',
+        arguments: {
+          query: 'release',
+          nativeQueries: [{ provider: 'google_drive', query: "name contains 'release'" }],
+        },
+      })
       expect(search.isError).not.toBe(true)
-      expect(search.content).toEqual([{ type: 'text', text: expect.stringContaining(documentId) }])
-      expect(mocks.liveSearch).toHaveBeenCalledWith(
-        expect.objectContaining({ input: expect.objectContaining({ topK: 30 }) })
-      )
+      const searchContent = toRecord(toArray(search.content)[0])
+      if (typeof searchContent.text !== 'string') throw new Error('Missing search payload')
+      expect(JSON.parse(searchContent.text)).toMatchObject({
+        results: [{ documentId, queryIndexes: [0] }],
+      })
       const read = await client.callTool({ name: 'read_document', arguments: { documentId } })
       expect(read.isError).not.toBe(true)
       expect(read.content).toEqual([

apps/sim/lib/knowledge/mcp/server.ts

@@ -162,6 +162,7 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe
               sourceDate: row.sourceDate,
               sourceContainerName: row.sourceContainerName,
               sourceContainerUrl: row.sourceContainerUrl,
+              ...(row.queryIndexes ? { queryIndexes: row.queryIndexes } : {}),
               connectorType: row.connectorType,
               content: row.content,
               chunkIndex: row.chunkIndex,

apps/sim/lib/mothership/generated/sim-assistant-tools.generated.ts

@@ -480,6 +480,8 @@ export const workspaceKnowledgeSearchResultSchema = z.object({
   sourceDate: z.string().nullable().optional(),
   sourceContainerName: z.string().optional(),
   sourceContainerUrl: z.string().optional(),
+  /** Zero-based nativeQueries entries that returned this result, retained through rank fusion. */
+  queryIndexes: z.array(z.number().int().min(0).max(8)).min(1).max(9).optional(),
   sourceDateType: z.enum(['event_start', 'message', 'modified']).optional(),
   /** The person behind the document, from its author-like tag; null when the source names none. */
   author: z.string().nullable(),

apps/sim/lib/sim-search/live/application.ts

@@ -972,7 +972,17 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({
     })
     return {
       query: input.query,
-      results: ranked.slice(0, input.topK).map(({ result }) => result),
+      results: ranked.slice(0, input.topK).map(({ result, queries: matched }) => {
+        const queryIndexes = [
+          ...new Set(
+            matched.flatMap((index) => {
+              const queryIndex = searched[index]?.status.queryIndex
+              return queryIndex === undefined ? [] : [queryIndex]
+            })
+          ),
+        ].sort((a, b) => a - b)
+        return queryIndexes.length ? { ...result, queryIndexes } : result
+      }),
       retrieval: {
         status: retrievalStatus,
         timedOutLegs: [],

apps/sim/lib/sim-search/live/managed-mcp.integration.ts

@@ -21,6 +21,7 @@ import { sleep } from '@sim/utils/helpers'
 import { generateId } from '@sim/utils/id'
 import { eq, inArray } from 'drizzle-orm'
 import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest'
+import { workspaceKnowledgeSearchDataSchema } from '@/lib/api/contracts/mothership-assistant-tools'
 import { env } from '@/lib/core/config/env'
 import { runWithOutboundOrganization } from '@/lib/core/network/context.server'
 import { createManagedMcpConnector } from '@/lib/credential-groups/managed-mcp-service'
@@ -493,6 +494,12 @@ describe('managed Search operation sessions', () => {
         'Second topology',
         TITLE,
       ])
+      expect(found.results).toEqual(
+        expect.arrayContaining([
+          expect.objectContaining({ documentName: TITLE, queryIndexes: [0] }),
+          expect.objectContaining({ documentName: 'Second topology', queryIndexes: [1] }),
+        ])
+      )
       expect(found.retrieval.status).toBe('complete')
       expect(
         found.live?.accounts.map(({ queryIndex, status }) => ({ queryIndex, status }))
@@ -576,6 +583,12 @@ describe('managed Search operation sessions', () => {
       'Second topology',
       TITLE,
     ])
+    expect(parallel.results).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({ documentName: TITLE, queryIndexes: [0] }),
+        expect.objectContaining({ documentName: 'Second topology', queryIndexes: [1] }),
+      ])
+    )
 
     for (const options of [{ project: '42' }, { topK: 5 }, { startDate: '2026-09-01T00:00:00Z' }]) {
       expect((await browse(0, options)).live?.accounts[0]?.status).not.toBe('unavailable')
@@ -633,17 +646,27 @@ describe('managed Search operation sessions', () => {
             kind: 'lucidchart',
             accountId: actor.credentialId,
           },
+          {
+            provider: 'lucid',
+            query: 'architecture',
+            kind: 'lucidchart',
+            accountId: actor.credentialId,
+          },
         ],
       },
     })
     expect(found.results).toHaveLength(1)
+    expect(workspaceKnowledgeSearchDataSchema.parse(found).results[0]).toMatchObject({
+      queryIndexes: [1, 2],
+    })
     expect(found.live?.accounts).toEqual([
       expect.objectContaining({
         queryIndex: 0,
         folders: [{ id: '42', name: 'Architecture' }],
         nextCursor: expect.any(String),
       }),
       expect.objectContaining({ queryIndex: 1, status: 'ok' }),
+      expect.objectContaining({ queryIndex: 2, status: 'ok' }),
     ])
     expect(openTransports.size).toBe(0)
   })