diff --git a/apps/sim/hooks/queries/organization-accounts.test.tsx b/apps/sim/hooks/queries/organization-accounts.test.tsx index 1a45f1fd069..b5773eb720d 100644 --- a/apps/sim/hooks/queries/organization-accounts.test.tsx +++ b/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 { @@ -14,6 +15,7 @@ import { listOrganizationAccountPeopleContract, updateOrganizationAccountsContract, } from '@/lib/api/contracts/organization-accounts' +import { resourceScopeKey } from '@/lib/core/resource-scope' import { organizationAccountsKeys, useDisconnectPersonalOrganizationAccount, @@ -21,14 +23,16 @@ import { 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() diff --git a/apps/sim/hooks/queries/organization-accounts.ts b/apps/sim/hooks/queries/organization-accounts.ts index d6c08965dd8..65cf2cafef7 100644 --- a/apps/sim/hooks/queries/organization-accounts.ts +++ b/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,6 +38,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 @@ -44,20 +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([ - queryClient.invalidateQueries({ - queryKey: searchSourceKeys.list({ kind: 'organization', organizationId }), - }), + onSuccess: async () => { + await Promise.all([ + resetOrganizationSearchAccess(queryClient, organizationId), queryClient.invalidateQueries({ queryKey: organizationAccountsKeys.detail(organizationId), }), - ]), + ]) + router.refresh() + }, }) } diff --git a/apps/sim/hooks/queries/search-integrations.test.tsx b/apps/sim/hooks/queries/search-integrations.test.tsx new file mode 100644 index 00000000000..866b7b55f46 --- /dev/null +++ b/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 + function Probe() { + mutation = useUpdateSearchIntegration() + return null + } + try { + await act(async () => + root.render( + + + + ) + ) + let pending: Promise + await act(async () => { + pending = mutation.mutateAsync({ + organizationId: 'org-1', + connectorType: 'github', + approved: false, + }) + }) + await act(async () => + root.render({null}) + ) + 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() + } + } +) diff --git a/apps/sim/hooks/queries/search-integrations.ts b/apps/sim/hooks/queries/search-integrations.ts index 4fcf974086a..a8b023e6c2c 100644 --- a/apps/sim/hooks/queries/search-integrations.ts +++ b/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>({ - 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() }, }) } diff --git a/apps/sim/hooks/queries/utils/reset-organization-search-access.test.ts b/apps/sim/hooks/queries/utils/reset-organization-search-access.test.ts new file mode 100644 index 00000000000..8e6a96e4a35 --- /dev/null +++ b/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() + 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() + } +}) diff --git a/apps/sim/hooks/queries/utils/reset-organization-search-access.ts b/apps/sim/hooks/queries/utils/reset-organization-search-access.ts new file mode 100644 index 00000000000..00a2c91cb97 --- /dev/null +++ b/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) }), + ]) +}