fix(search): clear cached content after disconnect

Author: Vikhyath Mondreti

apps/sim/hooks/queries/organization-accounts.test.tsx

@@ -14,13 +14,15 @@ import {
   listOrganizationAccountPeopleContract,
   updateOrganizationAccountsContract,
 } from '@/lib/api/contracts/organization-accounts'
+import { resourceScopeKey } from '@/lib/core/resource-scope'
 import {
   organizationAccountsKeys,
   useDisconnectPersonalOrganizationAccount,
   useOrganizationAccountPeople,
   useUpdateOrganizationAccounts,
 } from '@/hooks/queries/organization-accounts'
 import { slackSearchKeys } from '@/hooks/queries/slack-search'
+import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys'
 import { searchSourceKeys } from '@/hooks/queries/utils/search-source-keys'
 
 describe('personal account disconnect', () => {
@@ -48,7 +50,39 @@ describe('personal account disconnect', () => {
       )
       const people = organizationAccountsKeys.people('org-1')
       const other = searchSourceKeys.list({ kind: 'organization', organizationId: 'org-2' })
-      for (const key of [own, catalog, people, other]) client.setQueryData(key, { existing: true })
+      const results = knowledgeKeys.search(
+        resourceScopeKey({ kind: 'organization', organizationId: 'org-1' }),
+        'private content'
+      )
+      const otherResults = knowledgeKeys.search(
+        resourceScopeKey({ kind: 'organization', organizationId: 'org-2' }),
+        'private content'
+      )
+      const workspaceResults = knowledgeKeys.search('workspace-1', 'private content')
+      const ownDocument = knowledgeKeys.document('kb-1', 'document-1')
+      const chunks = knowledgeKeys.chunks('kb-2', 'document-2', '')
+      const otherDocument = knowledgeKeys.document('kb-other', 'document-other')
+      const resultOnlyDocument = knowledgeKeys.document('kb-result-only', 'document-result-only')
+      const sourcePages = {
+        pages: [
+          { sources: [{ knowledgeBaseId: 'kb-1' }], nextCursor: 'next' },
+          { sources: [{ knowledgeBaseId: 'kb-2' }], nextCursor: null },
+        ],
+        pageParams: [undefined, 'next'],
+      }
+      for (const key of [own, catalog]) client.setQueryData(key, sourcePages)
+      client.setQueryData(results, [{ knowledgeBaseId: 'kb-result-only' }])
+      for (const key of [
+        people,
+        other,
+        otherResults,
+        workspaceResults,
+        ownDocument,
+        resultOnlyDocument,
+        chunks,
+        otherDocument,
+      ])
+        client.setQueryData(key, { content: 'previously authorized content' })
       try {
         await act(async () =>
           root.render(
@@ -78,9 +112,14 @@ describe('personal account disconnect', () => {
           disconnectPersonalOrganizationAccountContract,
           { params: { credentialId: 'own-credential' } }
         )
-        for (const key of [own, catalog, people])
-          expect(client.getQueryState(key)?.isInvalidated).toBe(success)
+        for (const key of [own, catalog, results, ownDocument, resultOnlyDocument, chunks]) {
+          if (success) expect(client.getQueryData(key)).toBeUndefined()
+          else expect(client.getQueryData(key)).toBeDefined()
+        }
+        expect(client.getQueryState(people)?.isInvalidated).toBe(success)
         expect(client.getQueryState(other)?.isInvalidated).toBe(false)
+        for (const key of [otherResults, workspaceResults, otherDocument])
+          expect(client.getQueryData(key)).toEqual({ content: 'previously authorized content' })
       } finally {
         await act(async () => root.unmount())
         client.clear()

apps/sim/hooks/queries/organization-accounts.ts

@@ -37,6 +37,7 @@ import {
   updateOrganizationAccountWorkspaceAccessContract,
 } from '@/lib/api/contracts/organization-accounts'
 import { slackSearchKeys } from '@/hooks/queries/slack-search'
+import { resetOrganizationSearchAccess } from '@/hooks/queries/utils/reset-organization-search-access'
 import { searchSourceKeys } from '@/hooks/queries/utils/search-source-keys'
 
 export const ORGANIZATION_ACCOUNTS_STALE_TIME = 30_000
@@ -51,9 +52,7 @@ export function useDisconnectPersonalOrganizationAccount(organizationId: string)
       }),
     onSuccess: () =>
       Promise.all([
-        queryClient.invalidateQueries({
-          queryKey: searchSourceKeys.list({ kind: 'organization', organizationId }),
-        }),
+        resetOrganizationSearchAccess(queryClient, organizationId),
         queryClient.invalidateQueries({
           queryKey: organizationAccountsKeys.detail(organizationId),
         }),

apps/sim/hooks/queries/search-integrations.ts

@@ -1,15 +1,12 @@
-import { type InfiniteData, useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
+import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
 import { requestJson } from '@/lib/api/client/request'
-import type { SearchSourcePage } from '@/lib/api/contracts/knowledge/connectors'
 import {
   listSearchIntegrationsContract,
   type UpdateSearchIntegrationBody,
   updateSearchIntegrationContract,
 } from '@/lib/api/contracts/knowledge/search-integrations'
-import { resourceScopeKey } from '@/lib/core/resource-scope'
-import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys'
+import { resetOrganizationSearchAccess } from '@/hooks/queries/utils/reset-organization-search-access'
 import { searchIntegrationKeys } from '@/hooks/queries/utils/search-integration-keys'
-import { searchSourceKeys } from '@/hooks/queries/utils/search-source-keys'
 
 export const SEARCH_INTEGRATIONS_STALE_TIME = 30_000
 
@@ -28,29 +25,10 @@ export function useUpdateSearchIntegration() {
   return useMutation({
     mutationFn: async (body: UpdateSearchIntegrationBody) =>
       (await requestJson(updateSearchIntegrationContract, { body })).data,
-    onSuccess: async (_data, { organizationId }) => {
-      const scope = { kind: 'organization', organizationId } as const
-      const pages = queryClient.getQueriesData<InfiniteData<SearchSourcePage>>({
-        queryKey: searchSourceKeys.list(scope),
-        predicate: (query) => query.queryKey[3] === 'pages',
-      })
-      const knowledgeBaseIds = new Set(
-        pages.flatMap(
-          ([, data]) =>
-            data?.pages.flatMap((page) => page.sources.map((source) => source.knowledgeBaseId)) ??
-            []
-        )
-      )
-      await Promise.all([
+    onSuccess: (_data, { organizationId }) =>
+      Promise.all([
+        resetOrganizationSearchAccess(queryClient, organizationId),
         queryClient.invalidateQueries({ queryKey: searchIntegrationKeys.list(organizationId) }),
-        queryClient.invalidateQueries({ queryKey: searchSourceKeys.list(scope) }),
-        queryClient.resetQueries({
-          queryKey: [...knowledgeKeys.searches(), resourceScopeKey(scope)],
-        }),
-        ...[...knowledgeBaseIds].map((id) =>
-          queryClient.resetQueries({ queryKey: knowledgeKeys.detail(id) })
-        ),
-      ])
-    },
+      ]),
   })
 }

apps/sim/hooks/queries/utils/reset-organization-search-access.test.ts

@@ -0,0 +1,50 @@
+/** @vitest-environment node */
+import { QueryClient } from '@tanstack/react-query'
+import { expect, it, vi } from 'vitest'
+import type { WorkspaceKnowledgeSearchResult } from '@/lib/api/contracts/knowledge/search'
+import { resourceScopeKey } from '@/lib/core/resource-scope'
+import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys'
+import { resetOrganizationSearchAccess } from '@/hooks/queries/utils/reset-organization-search-access'
+
+it('cancels an in-flight search so its late result cannot restore disconnected content', async () => {
+  const client = new QueryClient({ defaultOptions: { queries: { retry: false } } })
+  const key = knowledgeKeys.search(
+    resourceScopeKey({ kind: 'organization', organizationId: 'org-1' }),
+    'private content'
+  )
+  const response = Promise.withResolvers<WorkspaceKnowledgeSearchResult[]>()
+  const result: WorkspaceKnowledgeSearchResult = {
+    documentId: 'document-1',
+    knowledgeBaseId: 'kb-1',
+    knowledgeBaseName: 'Knowledge',
+    documentName: 'Private document',
+    sourceUrl: null,
+    connectorType: 'github',
+    sourceModifiedAt: null,
+    author: null,
+    content: 'cached private content',
+    chunkIndex: 0,
+    similarity: 1,
+  }
+  const aborted = vi.fn()
+  try {
+    client.setQueryData(key, [result])
+    const pending = client.fetchQuery({
+      queryKey: key,
+      queryFn: ({ signal }) => {
+        signal.addEventListener('abort', aborted, { once: true })
+        return response.promise
+      },
+    })
+    const rejected = expect(pending).rejects.toThrow()
+    await resetOrganizationSearchAccess(client, 'org-1')
+    await rejected
+    expect(aborted).toHaveBeenCalledOnce()
+    expect(client.getQueryData(key)).toBeUndefined()
+    response.resolve([{ ...result, content: 'late private content' }])
+    await response.promise
+    expect(client.getQueryData(key)).toBeUndefined()
+  } finally {
+    client.clear()
+  }
+})

apps/sim/hooks/queries/utils/reset-organization-search-access.ts

@@ -0,0 +1,38 @@
+import type { InfiniteData, QueryClient } from '@tanstack/react-query'
+import type { SearchSourcePage } from '@/lib/api/contracts/knowledge/connectors'
+import type { WorkspaceKnowledgeSearchResult } from '@/lib/api/contracts/knowledge/search'
+import { resourceScopeKey } from '@/lib/core/resource-scope'
+import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys'
+import { searchSourceKeys } from '@/hooks/queries/utils/search-source-keys'
+
+/** Drops cached results and document content when organization Search access changes. */
+export async function resetOrganizationSearchAccess(
+  queryClient: QueryClient,
+  organizationId: string
+) {
+  const scope = { kind: 'organization', organizationId } as const
+  const pages = queryClient.getQueriesData<InfiniteData<SearchSourcePage>>({
+    queryKey: searchSourceKeys.list(scope),
+    predicate: (query) => query.queryKey[3] === 'pages',
+  })
+  const searchKey = [...knowledgeKeys.searches(), resourceScopeKey(scope)]
+  const results = queryClient.getQueriesData<WorkspaceKnowledgeSearchResult[]>({
+    queryKey: searchKey,
+  })
+  const knowledgeBaseIds = new Set([
+    ...pages.flatMap(
+      ([, data]) =>
+        data?.pages.flatMap((page) => page.sources.map((source) => source.knowledgeBaseId)) ?? []
+    ),
+    ...results.flatMap(([, data]) => data?.map((result) => result.knowledgeBaseId) ?? []),
+  ])
+  await Promise.all([
+    queryClient.resetQueries({
+      queryKey: searchKey,
+    }),
+    ...[...knowledgeBaseIds].map((id) =>
+      queryClient.resetQueries({ queryKey: knowledgeKeys.detail(id) })
+    ),
+    queryClient.resetQueries({ queryKey: searchSourceKeys.list(scope) }),
+  ])
+}