opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-10-16 15:35:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: refresh KnowledgeFS plan usage after document changes

Browse Source
This commit is contained in:
hjlarry
2026-09-30 17:54:10 +08:00
parent e8b8cf85f2
commit cb7780a9fb
10 changed files with 445 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
web/features/new-rag/create/__tests__/page.spec.tsx
+30 -1
View File
@@ -2,6 +2,7 @@ import type { ReactNode } from 'react'
import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
import { act, fireEvent, screen, waitFor, within } from '@testing-library/react'
import userEvent from '@testing-library/user-event'
import { consoleQuery } from '@/service/console'
import { render } from '@/test/console/render'
import { createSystemFeaturesFixture } from '@/test/console/system-features'
import { CreateKnowledgePage } from '../page'
@@ -237,6 +238,10 @@ vi.mock('@/service/console', () => ({
},
},
consoleQuery: {
features: {
get: { key: () => ['features'] },
vectorSpace: { get: { key: () => ['vector-space'] } },
},
rag: {
pipelines: {
datasourcePlugins: {
@@ -1282,6 +1287,7 @@ describe('CreateKnowledgePage', () => {
it('safely resumes a downstream upload after the control space is created', async () => {
const user = userEvent.setup()
navigationMock.startMode = 'upload'
systemFeaturesStateMock.deploymentEdition = 'CLOUD'
const queryClient = new QueryClient({ defaultOptions: { mutations: { retry: false } } })
const invalidate = vi.spyOn(queryClient, 'invalidateQueries')
serviceMock.upload.mockRejectedValueOnce(new Error('upload unavailable'))
@@ -1307,6 +1313,14 @@ describe('CreateKnowledgePage', () => {
expect(invalidate).toHaveBeenCalledWith({
queryKey: ['console', 'knowledgeFs', 'listKnowledgeSpaces'],
})
expect(invalidate).not.toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
)
expect(invalidate).not.toHaveBeenCalledWith(
{ queryKey: ['vector-space'] },
{ cancelRefetch: false },
)
const nameInput = screen.getByRole('textbox', { name: 'knowledgeSpace.name' })
expect(nameInput).toBeDisabled()
expect(screen.getByRole('button', { name: /^knowledgeSpace\.permission/ })).toBeDisabled()
@@ -1318,11 +1332,22 @@ describe('CreateKnowledgePage', () => {
expect(routerMock.replace).toHaveBeenCalledWith(
'/datasets/new/e735c1dc-d2b8-4dc4-86dc-abaf2fb7d084/documents',
)
expect(invalidate).toHaveBeenCalledWith({ queryKey: ['features'] }, { cancelRefetch: false })
expect(invalidate).toHaveBeenCalledWith(
{ queryKey: ['vector-space'] },
{ cancelRefetch: false },
)
})
it('shows a quota failure after a partial upload and retries only the remaining file', async () => {
const user = userEvent.setup()
navigationMock.startMode = 'upload'
systemFeaturesStateMock.deploymentEdition = 'CLOUD'
const queryClient = new QueryClient({ defaultOptions: { mutations: { retry: false } } })
queryClient.setQueryData(consoleQuery.features.get.key(), {
documents_upload_quota: { size: 0, limit: 10 },
})
queryClient.setQueryData(consoleQuery.features.vectorSpace.get.key(), { size: 0, limit: 10 })
vi.mocked(crypto.randomUUID)
.mockReturnValueOnce('00000000-0000-4000-8000-000000000001')
.mockReturnValueOnce('00000000-0000-4000-8000-000000000002')
@@ -1339,7 +1364,7 @@ describe('CreateKnowledgePage', () => {
{ status: 413 },
),
)
renderPage()
renderPage(queryClient)
await user.upload(
screen.getByLabelText('knowledgeSpace.uploadFiles', { selector: 'input[type="file"]' }),
[new File(['first'], 'first.md'), new File(['second'], 'second.md')],
@@ -1357,6 +1382,10 @@ describe('CreateKnowledgePage', () => {
expect(alert).not.toHaveTextContent('Private billing diagnostic')
expect(screen.getAllByText('knowledgeCreate.uploadCompleted')).toHaveLength(1)
expect(routerMock.replace).not.toHaveBeenCalled()
expect(queryClient.getQueryState(consoleQuery.features.get.key())?.isInvalidated).toBe(true)
expect(
queryClient.getQueryState(consoleQuery.features.vectorSpace.get.key())?.isInvalidated,
).toBe(true)
await user.click(screen.getByRole('button', { name: 'knowledgeCreate.createTitle' }))
await waitFor(() => expect(routerMock.replace).toHaveBeenCalled())
web/features/new-rag/create/page.tsx
+14
View File
@@ -342,6 +342,20 @@ function CreateKnowledgeSession() {
setUploadPhases(
(current) => new Map([...current].filter(([, phase]) => phase === 'completed')),
)
if (
deploymentEdition === 'CLOUD' &&
validUploads.some(({ id }) => uploadProgressRef.current.get(id)?.phase === 'completed')
)
void Promise.allSettled([
queryClient.invalidateQueries(
{ queryKey: consoleQuery.features.get.key() },
{ cancelRefetch: false },
),
queryClient.invalidateQueries(
{ queryKey: consoleQuery.features.vectorSpace.get.key() },
{ cancelRefetch: false },
),
])
}
}
web/features/new-rag/documents/__tests__/page.spec.tsx
+209 -2
View File
@@ -611,6 +611,10 @@ vi.mock('@/service/console', () => ({
},
},
consoleQuery: {
features: {
get: { key: () => ['features'] },
vectorSpace: { get: { key: () => ['features', 'vector-space'] } },
},
knowledgeFs: {
spaces: {
byControlSpaceId: {
@@ -2054,9 +2058,10 @@ describe('DocumentsPage', () => {
it('confirms permanent document removal from the row action', async () => {
const user = userEvent.setup()
systemFeaturesStateMock.deploymentEdition = 'CLOUD'
documentsQuery.data = { pages: [{ items: [document()] }] }
render(<DocumentsPage knowledgeSpaceId="space-1" />)
const { rerender } = render(<DocumentsPage knowledgeSpaceId="space-1" />)
await user.click(screen.getByRole('button', { name: /knowledgeDocuments\.documentActions/ }))
await user.click(await screen.findByRole('menuitem', { name: 'common.operation.delete' }))
expect(removeDocumentMutation).not.toHaveBeenCalled()
@@ -2067,6 +2072,32 @@ describe('DocumentsPage', () => {
headers: { 'Idempotency-Key': expect.any(String) },
params: { control_space_id: 'space-1', document_id: 'document-1' },
})
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
)
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
{ queryKey: ['features', 'vector-space'] },
{ cancelRefetch: false },
)
tasksQuery.data = { pages: [{ items: [backgroundTask({ operation: 'document_delete' })] }] }
rerender(<DocumentsPage knowledgeSpaceId="space-1" />)
await waitFor(() =>
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
),
)
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features', 'vector-space'] },
{ cancelRefetch: false },
)
queryClient.invalidateQueries.mockClear()
tasksQuery.data = { pages: [{ items: [backgroundTask({ operation: 'document_delete' })] }] }
rerender(<DocumentsPage knowledgeSpaceId="space-1" />)
expect(queryClient.invalidateQueries).not.toHaveBeenCalled()
})
it('opens the upload form and consumes the one-shot URL request', async () => {
@@ -2480,6 +2511,14 @@ describe('DocumentsPage', () => {
await user.click(screen.getByRole('button', { name: 'knowledgeSpace.addDocument' }))
expect(uploadMutation.mutateAsync).toHaveBeenCalledTimes(3)
expect(getUploadQuotaFailure).not.toHaveBeenCalled()
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
)
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
{ queryKey: ['features', 'vector-space'] },
{ cancelRefetch: false },
)
expect(queryClient.invalidateQueries).toHaveBeenCalled()
const documentInvalidation = queryClient.invalidateQueries.mock.calls.find(
([options]) => options.queryKey[1] === 'documents',
@@ -2502,6 +2541,29 @@ describe('DocumentsPage', () => {
).toBe(false)
})
it('refreshes Cloud plan usage only after a staged file becomes a document', async () => {
const user = userEvent.setup()
systemFeaturesStateMock.deploymentEdition = 'CLOUD'
render(<DocumentsPage knowledgeSpaceId="space-1" />, { searchParams: '?upload=1' })
await user.upload(
screen.getByLabelText('knowledgeCreate.uploadDocuments'),
new File(['draft'], 'draft.md', { type: 'text/markdown' }),
)
await waitForDocumentFilesStaged()
expect(queryClient.invalidateQueries).not.toHaveBeenCalled()
await user.click(screen.getByRole('button', { name: 'knowledgeSpace.addDocument' }))
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
)
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features', 'vector-space'] },
{ cancelRefetch: false },
)
})
it.each([
'taskFailure.documentCountQuotaExceeded',
'taskFailure.vectorSpaceQuotaExceeded',
@@ -3036,6 +3098,7 @@ describe('DocumentsPage', () => {
'shows %s without losing the staged file or write access',
async (code, category, status, key) => {
const user = userEvent.setup()
systemFeaturesStateMock.deploymentEdition = 'CLOUD'
uploadMutation.mutateAsync.mockRejectedValueOnce(
new Response(
JSON.stringify({
@@ -3106,6 +3169,15 @@ describe('DocumentsPage', () => {
).not.toBeInTheDocument()
expect(queryClient.invalidateQueries).toHaveBeenCalled()
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
)
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features', 'vector-space'] },
{ cancelRefetch: false },
)
await user.click(screen.getByRole('button', { name: 'knowledgeCreate.retryRemainingFiles' }))
await waitFor(() =>
@@ -4286,6 +4358,7 @@ describe('DocumentsPage', () => {
it('removes selected documents through one bulk deletion request', async () => {
const user = userEvent.setup()
systemFeaturesStateMock.deploymentEdition = 'CLOUD'
documentsQuery.data = {
pages: [
{
@@ -4297,7 +4370,7 @@ describe('DocumentsPage', () => {
],
}
render(<DocumentsPage knowledgeSpaceId="space-1" />)
const { rerender } = render(<DocumentsPage knowledgeSpaceId="space-1" />)
await user.click(screen.getByRole('checkbox', { name: 'One.pdf' }))
await user.click(screen.getByRole('checkbox', { name: 'Two.pdf' }))
await user.click(screen.getByRole('button', { name: 'common.operation.remove' }))
@@ -4316,6 +4389,56 @@ describe('DocumentsPage', () => {
params: { control_space_id: 'space-1' },
})
expect(removeDocumentMutation).not.toHaveBeenCalled()
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
)
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
{ queryKey: ['features', 'vector-space'] },
{ cancelRefetch: false },
)
tasksQuery.data = {
pages: [
{
items: [
backgroundTask({
operation: 'document_delete',
progressCompleted: 2,
progressTotal: 2,
}),
],
},
],
}
rerender(<DocumentsPage knowledgeSpaceId="space-1" />)
await waitFor(() =>
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
),
)
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features', 'vector-space'] },
{ cancelRefetch: false },
)
queryClient.invalidateQueries.mockClear()
tasksQuery.data = {
pages: [
{
items: [
backgroundTask({
operation: 'document_delete',
progressCompleted: 2,
progressTotal: 2,
}),
],
},
],
}
rerender(<DocumentsPage knowledgeSpaceId="space-1" />)
expect(queryClient.invalidateQueries).not.toHaveBeenCalled()
})
it('prompts for model setup before re-indexing selected documents', async () => {
@@ -6384,6 +6507,7 @@ describe('DocumentsPage', () => {
})
it('refreshes logical documents when task-list polling reports a terminal state', async () => {
systemFeaturesStateMock.deploymentEdition = 'CLOUD'
documentsQuery.data = {
pages: [
{
@@ -6431,6 +6555,89 @@ describe('DocumentsPage', () => {
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
expect.objectContaining({ queryKey: ['knowledge-fs', 'golden-questions'] }),
)
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
)
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features', 'vector-space'] },
{ cancelRefetch: false },
)
})
it.each([1, 3])(
'refreshes Cloud plan usage after a background deletion of %s documents finishes',
async (total) => {
systemFeaturesStateMock.deploymentEdition = 'CLOUD'
const deletion = backgroundTask({
id: 'delete-documents',
operation: 'document_delete',
progressCompleted: 0,
progressPercent: 0,
progressTotal: total,
state: 'queued',
})
tasksQuery.data = { pages: [{ items: [deletion] }] }
const { rerender } = render(<DocumentsPage knowledgeSpaceId="space-1" />)
expect(queryClient.invalidateQueries).not.toHaveBeenCalled()
tasksQuery.data = { pages: [{ items: [{ ...deletion, state: 'running' }] }] }
rerender(<DocumentsPage knowledgeSpaceId="space-1" />)
expect(queryClient.invalidateQueries).not.toHaveBeenCalled()
tasksQuery.data = {
pages: [
{
items: [
{
...deletion,
progressCompleted: total,
progressPercent: 100,
state: 'succeeded',
updatedAt: '2026-07-20T10:06:00Z',
},
],
},
],
}
rerender(<DocumentsPage knowledgeSpaceId="space-1" />)
await waitFor(() =>
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
),
)
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features', 'vector-space'] },
{ cancelRefetch: false },
)
queryClient.invalidateQueries.mockClear()
rerender(<DocumentsPage knowledgeSpaceId="space-1" />)
expect(queryClient.invalidateQueries).not.toHaveBeenCalled()
},
)
it('does not refresh plan usage for completed deletions in Enterprise', async () => {
const deletion = backgroundTask({
id: 'delete-documents',
operation: 'document_delete',
state: 'running',
})
tasksQuery.data = { pages: [{ items: [deletion] }] }
const { rerender } = render(<DocumentsPage knowledgeSpaceId="space-1" />)
tasksQuery.data = { pages: [{ items: [{ ...deletion, state: 'succeeded' }] }] }
rerender(<DocumentsPage knowledgeSpaceId="space-1" />)
await waitFor(() => expect(queryClient.invalidateQueries).toHaveBeenCalled())
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
)
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
{ queryKey: ['features', 'vector-space'] },
{ cancelRefetch: false },
)
})
it('accepts an external retry newer than the terminal event version', async () => {
web/features/new-rag/documents/detail/__tests__/page.spec.tsx
+48 -10
View File
@@ -188,6 +188,13 @@ const queryClient = vi.hoisted(() => ({
invalidateQueries: vi.fn(),
removeQueries: vi.fn(),
}))
const systemFeaturesState = vi.hoisted(() => ({ deploymentEdition: 'ENTERPRISE' }))
vi.mock('@/features/system-features/state', async () => {
const { atom } = await import('jotai')
return { deploymentEditionAtom: atom(() => systemFeaturesState.deploymentEdition) }
})
const toastState = vi.hoisted(() => ({ error: vi.fn(), info: vi.fn(), success: vi.fn() }))
const virtualizerState = vi.hoisted(() => ({ scrollToIndex: vi.fn() }))
const revisionApiResponse = vi.hoisted(
@@ -667,6 +674,10 @@ vi.mock('@/service/console', () => ({
},
},
consoleQuery: {
features: {
get: { key: () => ['features'] },
vectorSpace: { get: { key: () => ['vector-space'] } },
},
knowledgeFs: {
spaces: {
byControlSpaceId: {
@@ -988,6 +999,7 @@ describe('DocumentDetailPage', () => {
next_cursor: null,
})
queryClient.invalidateQueries.mockResolvedValue(undefined)
systemFeaturesState.deploymentEdition = 'ENTERPRISE'
})
afterEach(() => {
@@ -2495,17 +2507,43 @@ describe('DocumentDetailPage', () => {
expect(screen.queryByText('knowledgeTasks.graphRepairPending')).not.toBeInTheDocument()
})
it('refreshes the document and content when an active task becomes terminal', async () => {
taskSnapshotQuery.data = task({ state: 'running' })
documentQuery.data = logicalDocument({ latestTask: taskSnapshotQuery.data })
const rendered = render(
<DocumentDetailPage documentId="document-1" knowledgeSpaceId="space-1" />,
)
taskSnapshotQuery.data = task({ state: 'succeeded' })
rendered.rerender(<DocumentDetailPage documentId="document-1" knowledgeSpaceId="space-1" />)
it.each(['CLOUD', 'ENTERPRISE'])(
'refreshes document content and applicable usage when an active task becomes terminal in %s',
async (edition) => {
systemFeaturesState.deploymentEdition = edition
taskSnapshotQuery.data = task({ state: 'running' })
documentQuery.data = logicalDocument({ latestTask: taskSnapshotQuery.data })
const rendered = render(
<DocumentDetailPage documentId="document-1" knowledgeSpaceId="space-1" />,
)
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
)
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
{ queryKey: ['vector-space'] },
{ cancelRefetch: false },
)
taskSnapshotQuery.data = task({ state: 'succeeded' })
rendered.rerender(<DocumentDetailPage documentId="document-1" knowledgeSpaceId="space-1" />)
await waitFor(() => expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(4))
})
await waitFor(() =>
expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(edition === 'CLOUD' ? 6 : 4),
)
if (edition === 'CLOUD') {
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['features'] },
{ cancelRefetch: false },
)
expect(queryClient.invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['vector-space'] },
{ cancelRefetch: false },
)
}
rendered.rerender(<DocumentDetailPage documentId="document-1" knowledgeSpaceId="space-1" />)
expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(edition === 'CLOUD' ? 6 : 4)
},
)
it('virtualizes long trees and bounds their accessible labels', () => {
chunksQuery.data = {
web/features/new-rag/documents/detail/state/workflow.ts
+24 -2
View File
@@ -9,6 +9,7 @@ import {
queryClientAtom,
} from 'jotai-tanstack-query'
import { selectAtom } from 'jotai/utils'
import { deploymentEditionAtom } from '@/features/system-features/state'
import { consoleQuery } from '@/service/console'
import {
knowledgeFsRequestFailureMessageKey,
@@ -336,6 +337,21 @@ async function invalidateDocumentWorkflow(get: Getter) {
])
}
function invalidateDocumentUsage(get: Getter) {
if (get(deploymentEditionAtom) !== 'CLOUD') return
const queryClient = get(queryClientAtom)
void Promise.allSettled([
queryClient.invalidateQueries(
{ queryKey: consoleQuery.features.get.key() },
{ cancelRefetch: false },
),
queryClient.invalidateQueries(
{ queryKey: consoleQuery.features.vectorSpace.get.key() },
{ cancelRefetch: false },
),
])
}
async function retryWritePermission(
get: Getter,
set: Setter,
@@ -417,7 +433,10 @@ export const reconcileDocumentTaskAtom = atom(null, async (get, set) => {
invalidatedTerminalTask: shouldInvalidate ? terminalTaskKey : current.invalidatedTerminalTask,
previousTaskState: latestTask?.state,
}))
if (shouldInvalidate) await invalidateDocumentWorkflow(get)
if (shouldInvalidate) {
invalidateDocumentUsage(get)
await invalidateDocumentWorkflow(get)
}
})
export const reconcileSubmittedDocumentJobAtom = atom(null, async (get, set) => {
@@ -433,7 +452,10 @@ export const reconcileSubmittedDocumentJobAtom = atom(null, async (get, set) =>
} catch {
// Browser storage recovery is optional.
}
if (terminal) await invalidateDocumentWorkflow(get)
if (terminal) {
invalidateDocumentUsage(get)
await invalidateDocumentWorkflow(get)
}
})
export const loadNextDocumentTaskPageAtom = atom(null, (get) =>
web/features/new-rag/documents/tasks/runtime-boundary.tsx
+15 -1
View File
@@ -5,6 +5,7 @@ import { useAtomValueRawSync, useSetAtom } from 'jotai'
import { useCallback, useLayoutEffect } from 'react'
import { useTranslation } from 'react-i18next'
import { toast } from '@/app/notifications'
import { deploymentEditionAtom } from '@/features/system-features/state'
import { consoleQuery } from '@/service/console'
import { documentsKnowledgeSpaceIdAtom } from '../state/inputs'
import { documentTaskPermissionGuardFactsAtom } from '../state/recovery'
@@ -18,6 +19,7 @@ import { useTaskRuntimeController } from './use-task-runtime'
export function DocumentTaskRuntimeController() {
const { t } = useTranslation(['knowledgeSpace', 'knowledgeTasks'])
const queryClient = useQueryClient()
const deploymentEdition = useAtomValueRawSync(deploymentEditionAtom)
const knowledgeSpaceId = useAtomValueRawSync(documentsKnowledgeSpaceIdAtom)
const tasksOpen = useAtomValueRawSync(documentTasksOpenAtom)
const permissionQueryFacts = useAtomValueRawSync(documentTaskPermissionGuardFactsAtom)
@@ -42,8 +44,20 @@ export function DocumentTaskRuntimeController() {
predicate: (query) => queryKeyMatchesKnowledgeSpace(query.queryKey, knowledgeSpaceId),
queryKey: consoleQuery.knowledgeFs.spaces.byControlSpaceId.goldenQuestions.get.key(),
}),
...(deploymentEdition === 'CLOUD'
? [
queryClient.invalidateQueries(
{ queryKey: consoleQuery.features.get.key() },
{ cancelRefetch: false },
),
queryClient.invalidateQueries(
{ queryKey: consoleQuery.features.vectorSpace.get.key() },
{ cancelRefetch: false },
),
]
: []),
])
}, [knowledgeSpaceId, queryClient])
}, [deploymentEdition, knowledgeSpaceId, queryClient])
const notifyTaskFailed = useCallback(
() => toast.error(t(($) => $.taskFailedNotification, { ns: 'knowledgeTasks' })),
[t],
web/features/new-rag/documents/tasks/use-task-runtime.ts
+4 -1
View File
@@ -196,8 +196,11 @@ export function useTaskRuntimeController({
for (const task of backgroundTasks) {
states.set(task.id, task.state)
const previousState = previousStates.get(task.id)
// A deletion can finish before its first polled snapshot arrives.
if (
(previousState === 'queued' || previousState === 'running') &&
(previousState === 'queued' ||
previousState === 'running' ||
(previousState === undefined && task.operation === 'document_delete')) &&
task.state !== 'queued' &&
task.state !== 'running' &&
(task.taskKind === 'document' || task.taskKind === 'document_bulk')
web/features/new-rag/documents/upload/surface.tsx
+13 -1
View File
@@ -110,8 +110,20 @@ export function DocumentUploadSurface({ children }: { children: ReactNode }) {
predicate: (query) => queryKeyMatchesKnowledgeSpace(query.queryKey, knowledgeSpaceId),
queryKey: consoleQuery.knowledgeFs.spaces.byControlSpaceId.backgroundTasks.get.key(),
}),
...(deploymentEdition === 'CLOUD'
? [
queryClient.invalidateQueries(
{ queryKey: consoleQuery.features.get.key() },
{ cancelRefetch: false },
),
queryClient.invalidateQueries(
{ queryKey: consoleQuery.features.vectorSpace.get.key() },
{ cancelRefetch: false },
),
]
: []),
])
}, [knowledgeSpaceId, queryClient])
}, [deploymentEdition, knowledgeSpaceId, queryClient])
const close = useCallback(() => {
resetProgress()
web/features/new-rag/overview/__tests__/task-sync.spec.tsx
+75 -2
View File
@@ -2,7 +2,7 @@ import type {
KnowledgeFsBackgroundTaskResponse,
KnowledgeFsOverviewAttentionResponse,
} from '@dify/contracts/api/console/knowledge-fs/types.gen'
import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
import { QueryClient, QueryClientProvider, QueryObserver } from '@tanstack/react-query'
import { act, cleanup, render, screen, waitFor } from '@testing-library/react'
import { createStore, Provider, useAtomValueRawSync } from 'jotai'
import { queryClientAtom } from 'jotai-tanstack-query'
@@ -10,7 +10,16 @@ import { OverviewTaskSync } from '../overview-task-sync'
import { overviewActivityPreviewDataAtom, overviewAttentionDataAtom } from '../state'
import { OverviewStateBoundary } from '../state-boundary'
const service = vi.hoisted(() => ({ attention: vi.fn(), activity: vi.fn(), tasks: vi.fn() }))
const service = vi.hoisted(() => ({
attention: vi.fn(),
activity: vi.fn(),
tasks: vi.fn(),
deploymentEdition: vi.fn(),
}))
vi.mock('@/features/system-features/state', async () => {
const { atom } = await import('jotai')
return { deploymentEditionAtom: atom(() => service.deploymentEdition()) }
})
vi.mock('@/service/console', () => {
const endpoint = (name: 'attention' | 'activity') => ({
get: {
@@ -29,6 +38,10 @@ vi.mock('@/service/console', () => {
})
return {
consoleQuery: {
features: {
get: { key: () => ['billing', 'features'] },
vectorSpace: { get: { key: () => ['billing', 'vector-space'] } },
},
knowledgeFs: {
spaces: {
byControlSpaceId: {
@@ -106,6 +119,9 @@ describe('overview task completion synchronization', () => {
service.activity.mockResolvedValue({ data: [], next_cursor: null })
service.attention.mockResolvedValue({ data: [] })
service.tasks.mockResolvedValue({ data: [task('running')], next_cursor: null })
service.deploymentEdition.mockReturnValue('ENTERPRISE')
client.setQueryData(['billing', 'features'], { documents_upload_quota: { size: 1, limit: 10 } })
client.setQueryData(['billing', 'vector-space'], { size: 1, limit: 10 })
})
afterEach(() => {
cleanup()
@@ -128,6 +144,9 @@ describe('overview task completion synchronization', () => {
}
it('cancels a pre-terminal snapshot and fetches failures after the last active task stops', async () => {
service.deploymentEdition.mockReturnValue('CLOUD')
const cancelQueries = vi.spyOn(client, 'cancelQueries')
const invalidateQueries = vi.spyOn(client, 'invalidateQueries')
let resolveStale!: (value: { data: KnowledgeFsOverviewAttentionResponse[] }) => void
service.attention.mockImplementationOnce(
() =>
@@ -138,6 +157,8 @@ describe('overview task completion synchronization', () => {
mount()
await waitFor(() => expect(service.attention).toHaveBeenCalledOnce())
await waitFor(() => expect(service.tasks).toHaveBeenCalledOnce())
expect(client.getQueryState(['billing', 'features'])?.isInvalidated).toBe(false)
expect(client.getQueryState(['billing', 'vector-space'])?.isInvalidated).toBe(false)
const staleSignal = service.attention.mock.calls[0]![1] as AbortSignal
service.tasks.mockResolvedValue({ data: [task('failed')], next_cursor: null })
service.attention.mockResolvedValue({ data: [failedIssue] })
@@ -145,6 +166,18 @@ describe('overview task completion synchronization', () => {
await act(() => client.invalidateQueries({ queryKey: ['overview', 'tasks', 'space-a'] }))
expect(await screen.findByText('Document processing failed')).toBeInTheDocument()
expect(staleSignal.aborted).toBe(true)
expect(client.getQueryState(['billing', 'features'])?.isInvalidated).toBe(true)
expect(client.getQueryState(['billing', 'vector-space'])?.isInvalidated).toBe(true)
expect(cancelQueries).not.toHaveBeenCalledWith({ queryKey: ['billing', 'features'] })
expect(cancelQueries).not.toHaveBeenCalledWith({ queryKey: ['billing', 'vector-space'] })
expect(invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['billing', 'features'] },
{ cancelRefetch: false },
)
expect(invalidateQueries).toHaveBeenCalledWith(
{ queryKey: ['billing', 'vector-space'] },
{ cancelRefetch: false },
)
await act(async () => {
resolveStale({ data: [] })
})
@@ -157,6 +190,44 @@ describe('overview task completion synchronization', () => {
expect(service.attention).toHaveBeenCalledTimes(calls)
})
it.each(['features', 'vector-space'])(
'preserves an active %s precheck while refreshing terminal tasks',
async (usageKey) => {
service.deploymentEdition.mockReturnValue('CLOUD')
const queryKey = ['billing', usageKey]
let resolveUsage!: (value: { size: number; limit: number }) => void
let precheckSignal!: AbortSignal
const queryFn = vi.fn(({ signal }: { signal: AbortSignal }) => {
precheckSignal = signal
return new Promise<{ size: number; limit: number }>((resolve) => {
resolveUsage = resolve
})
})
const observer = new QueryObserver(client, { queryKey, queryFn, staleTime: Infinity })
const unsubscribe = observer.subscribe(() => {})
const precheck = client
.query({ queryKey, queryFn, staleTime: 0, retry: false })
.catch((error: unknown) => error)
try {
mount()
await waitFor(() => expect(service.tasks).toHaveBeenCalledOnce())
service.tasks.mockResolvedValue({ data: [task('failed')], next_cursor: null })
service.attention.mockResolvedValue({ data: [failedIssue] })
await act(() => client.invalidateQueries({ queryKey: ['overview', 'tasks', 'space-a'] }))
expect(await screen.findByText('Document processing failed')).toBeInTheDocument()
expect(precheckSignal.aborted).toBe(false)
expect(queryFn).toHaveBeenCalledOnce()
resolveUsage({ size: 2, limit: 10 })
await expect(precheck).resolves.toEqual({ size: 2, limit: 10 })
} finally {
unsubscribe()
}
},
)
it('refreshes again for a manual retry completion while preserving another space cache', async () => {
client.setQueryData(['overview', 'attention', 'space-b'], { data: [failedIssue] })
mount()
@@ -180,5 +251,7 @@ describe('overview task completion synchronization', () => {
expect(client.getQueryData(['overview', 'attention', 'space-b'])).toEqual({
data: [failedIssue],
})
expect(client.getQueryState(['billing', 'features'])?.isInvalidated).toBe(false)
expect(client.getQueryState(['billing', 'vector-space'])?.isInvalidated).toBe(false)
})
})
web/features/new-rag/overview/state.ts
+13
View File
@@ -3,6 +3,7 @@ import { skipToken } from '@tanstack/react-query'
import { atom } from 'jotai'
import { atomWithInfiniteQuery, atomWithQuery, queryClientAtom } from 'jotai-tanstack-query'
import { atomWithLazy, selectAtom } from 'jotai/utils'
import { deploymentEditionAtom } from '@/features/system-features/state'
import { consoleQuery } from '@/service/console'
import { OVERVIEW_REFRESH_INTERVAL, overviewRefreshInterval } from './overview-format'
import { overviewQueryGroup } from './query-state'
@@ -260,6 +261,18 @@ export const reconcileOverviewAfterTasksAtom = atom(null, async (get) => {
const input = { params: { control_space_id: get(overviewKnowledgeSpaceIdAtom) } }
const overview = consoleQuery.knowledgeFs.spaces.byControlSpaceId.overview
const keys = [overview.attention.get.key({ input }), overview.activity.get.key({ input })]
if (get(deploymentEditionAtom) === 'CLOUD') {
void Promise.allSettled([
queryClient.invalidateQueries(
{ queryKey: consoleQuery.features.get.key() },
{ cancelRefetch: false },
),
queryClient.invalidateQueries(
{ queryKey: consoleQuery.features.vectorSpace.get.key() },
{ cancelRefetch: false },
),
])
}
await Promise.all(
keys.map(async (queryKey) => {
// A pre-completion request may finish late with a stale snapshot. Cancel it before the