fix(search): reset directly opened document caches

Author: Vikhyath Mondreti

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

@@ -5,8 +5,9 @@ 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 {
@@ -27,10 +28,11 @@ 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()
@@ -112,13 +114,22 @@ describe('personal account disconnect', () => {
           disconnectPersonalOrganizationAccountContract,
           { params: { credentialId: 'own-credential' } }
         )
-        for (const key of [own, catalog, results, ownDocument, resultOnlyDocument, chunks]) {
+        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)
-        for (const key of [otherResults, workspaceResults, otherDocument])
+        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())

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 {
@@ -45,18 +46,21 @@ 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([
+    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,4 +1,5 @@
 import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
+import { useRouter } from 'next/navigation'
 import { requestJson } from '@/lib/api/client/request'
 import {
   listSearchIntegrationsContract,
@@ -22,13 +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: (_data, { organizationId }) =>
-      Promise.all([
+    onSuccess: async (_data, { organizationId }) => {
+      await Promise.all([
         resetOrganizationSearchAccess(queryClient, organizationId),
         queryClient.invalidateQueries({ queryKey: searchIntegrationKeys.list(organizationId) }),
-      ]),
+      ])
+      router.refresh()
+    },
   })
 }

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

@@ -6,6 +6,35 @@ 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(

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

@@ -1,6 +1,4 @@
-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 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'
@@ -11,28 +9,12 @@ export async function resetOrganizationSearchAccess(
   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,
+      queryKey: [...knowledgeKeys.searches(), resourceScopeKey(scope)],
     }),
-    ...[...knowledgeBaseIds].map((id) =>
-      queryClient.resetQueries({ queryKey: knowledgeKeys.detail(id) })
-    ),
+    /** 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) }),
   ])
 }