fix(search): clear cached content after disconnect (#7712)

* fix(search): clear cached content after disconnect

* fix(search): reset directly opened document caches

Author: icecrasher321

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

@@ -5,30 +5,34 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
 import { createRoot, type Root } from 'react-dom/client'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-const mocks = vi.hoisted(() => ({ request: vi.fn() }))
+const mocks = vi.hoisted(() => ({ request: vi.fn(), refresh: vi.fn() }))
 vi.mock('@/lib/api/client/request', () => ({ requestJson: mocks.request }))
+vi.mock('next/navigation', () => ({ useRouter: () => ({ refresh: mocks.refresh }) }))
 
 import { ApiClientError } from '@/lib/api/client/errors'
 import {
   disconnectPersonalOrganizationAccountContract,
   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', () => {
   it.each([true, false])(
-    'refreshes this organization only after success=%s, including after unmount',
+    'clears content and refreshes the router only after success=%s, including after unmount',
     async (success) => {
       vi.stubGlobal('IS_REACT_ACT_ENVIRONMENT', true)
       mocks.request.mockReset()
+      mocks.refresh.mockReset()
       const response = Promise.withResolvers<{ success: true }>()
       mocks.request.mockReturnValue(response.promise)
       const client = new QueryClient()
@@ -48,7 +52,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 +114,23 @@ 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,
+          otherDocument,
+        ]) {
+          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)
+        expect(mocks.refresh).toHaveBeenCalledTimes(success ? 1 : 0)
+        for (const key of [otherResults, workspaceResults])
+          expect(client.getQueryData(key)).toEqual({ content: 'previously authorized content' })
       } finally {
         await act(async () => root.unmount())
         client.clear()

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

@@ -7,6 +7,7 @@ import {
   useQuery,
   useQueryClient,
 } from '@tanstack/react-query'
+import { useRouter } from 'next/navigation'
 import { isApiClientError } from '@/lib/api/client/errors'
 import { requestJson } from '@/lib/api/client/request'
 import {
@@ -37,27 +38,29 @@ 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
 
 /** Disconnects an owned grant; indexing and source setup do not gate this operation. */
 export function useDisconnectPersonalOrganizationAccount(organizationId: string) {
   const queryClient = useQueryClient()
+  const router = useRouter()
   return useMutation({
     mutationFn: (credentialId: string) =>
       requestJson(disconnectPersonalOrganizationAccountContract, {
         params: { credentialId },
       }),
-    onSuccess: () =>
-      Promise.all([
-        queryClient.invalidateQueries({
-          queryKey: searchSourceKeys.list({ kind: 'organization', organizationId }),
-        }),
+    onSuccess: async () => {
+      await Promise.all([
+        resetOrganizationSearchAccess(queryClient, organizationId),
         queryClient.invalidateQueries({
           queryKey: organizationAccountsKeys.detail(organizationId),
         }),
-      ]),
+      ])
+      router.refresh()
+    },
   })
 }
 

apps/sim/hooks/queries/search-integrations.test.tsx

@@ -0,0 +1,69 @@
+/** @vitest-environment jsdom */
+import { act } from 'react'
+import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
+import { createRoot } from 'react-dom/client'
+import { expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => ({ request: vi.fn(), refresh: vi.fn() }))
+vi.mock('@/lib/api/client/request', () => ({ requestJson: mocks.request }))
+vi.mock('next/navigation', () => ({ useRouter: () => ({ refresh: mocks.refresh }) }))
+
+import { useUpdateSearchIntegration } from '@/hooks/queries/search-integrations'
+import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys'
+
+it.each([true, false])(
+  'clears document content and refreshes server pages after integration update success=%s',
+  async (success) => {
+    vi.stubGlobal('IS_REACT_ACT_ENVIRONMENT', true)
+    mocks.request.mockReset()
+    mocks.refresh.mockReset()
+    const response = Promise.withResolvers<{ data: { connectorType: string; approved: boolean } }>()
+    mocks.request.mockReturnValue(response.promise)
+    const client = new QueryClient()
+    const root = createRoot(document.createElement('div'))
+    const key = knowledgeKeys.document('kb-direct', 'document-direct')
+    client.setQueryData(key, { content: 'previously authorized content' })
+    let mutation: ReturnType<typeof useUpdateSearchIntegration>
+    function Probe() {
+      mutation = useUpdateSearchIntegration()
+      return null
+    }
+    try {
+      await act(async () =>
+        root.render(
+          <QueryClientProvider client={client}>
+            <Probe />
+          </QueryClientProvider>
+        )
+      )
+      let pending: Promise<unknown>
+      await act(async () => {
+        pending = mutation.mutateAsync({
+          organizationId: 'org-1',
+          connectorType: 'github',
+          approved: false,
+        })
+      })
+      await act(async () =>
+        root.render(<QueryClientProvider client={client}>{null}</QueryClientProvider>)
+      )
+      await act(async () => {
+        if (success) {
+          response.resolve({ data: { connectorType: 'github', approved: false } })
+          await pending
+        } else {
+          const rejection = expect(pending).rejects.toThrow('Try again')
+          response.reject(new Error('Try again'))
+          await rejection
+        }
+      })
+      expect(mocks.refresh).toHaveBeenCalledTimes(success ? 1 : 0)
+      if (success) expect(client.getQueryData(key)).toBeUndefined()
+      else expect(client.getQueryData(key)).toEqual({ content: 'previously authorized content' })
+    } finally {
+      await act(async () => root.unmount())
+      client.clear()
+      vi.unstubAllGlobals()
+    }
+  }
+)

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

@@ -1,15 +1,13 @@
-import { type InfiniteData, useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
+import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
+import { useRouter } from 'next/navigation'
 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
 
@@ -25,32 +23,16 @@ export function useSearchIntegrations(organizationId: string) {
 
 export function useUpdateSearchIntegration() {
   const queryClient = useQueryClient()
+  const router = useRouter()
   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([
+        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) })
-        ),
       ])
+      router.refresh()
     },
   })
 }

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

@@ -0,0 +1,79 @@
+/** @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.each([
+  { name: 'document', key: knowledgeKeys.document('kb-direct', 'document-direct') },
+  { name: 'chunks', key: knowledgeKeys.chunks('kb-direct', 'document-direct', '') },
+])('clears and cancels directly loaded $name without source or Search caches', async ({ key }) => {
+  const client = new QueryClient({ defaultOptions: { queries: { retry: false } } })
+  const response = Promise.withResolvers<{ content: string }>()
+  const aborted = vi.fn()
+  try {
+    client.setQueryData(key, { content: 'cached private content' })
+    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({ content: 'late private content' })
+    await response.promise
+    expect(client.getQueryData(key)).toBeUndefined()
+  } finally {
+    client.clear()
+  }
+})
+
+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,20 @@
+import type { QueryClient } from '@tanstack/react-query'
+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
+  await Promise.all([
+    queryClient.resetQueries({
+      queryKey: [...knowledgeKeys.searches(), resourceScopeKey(scope)],
+    }),
+    /** Document keys carry no resource scope and may exist without source or result caches. */
+    queryClient.resetQueries({ queryKey: knowledgeKeys.details() }),
+    queryClient.resetQueries({ queryKey: searchSourceKeys.list(scope) }),
+  ])
+}