From d061d5ed3ee78ad99097dd0558d2eaab63c29eb0 Mon Sep 17 00:00:00 2001 From: Stephen Zhou <38493346+hyoban@users.noreply.github.com> Date: Thu, 20 Aug 2026 12:01:24 +0800 Subject: [PATCH] fix(knowledge-fs): finalize async source imports --- api/extensions/ext_celery.py | 1 + .../source_import_commit_service.py | 3 +- api/tasks/knowledge_fs_source_import_tasks.py | 64 +++++++++--- .../extensions/test_knowledge_fs_celery.py | 2 + ...owledge_fs_source_import_commit_service.py | 2 +- .../test_knowledge_fs_source_import_tasks.py | 52 +++++++++- .../__tests__/connected-source-setup.spec.tsx | 56 +++++++++++ .../new-rag/__tests__/sources-page.spec.tsx | 47 ++++++++- .../new-rag/connected-source-setup.tsx | 99 +++++++++++-------- web/features/new-rag/source-models.ts | 29 ++++++ web/features/new-rag/sources-page.tsx | 4 +- 11 files changed, 297 insertions(+), 62 deletions(-) diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index 1c42a08939e..273c4365221 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -190,6 +190,7 @@ def init_app(app: DifyApp) -> Celery: if get_configured_knowledge_fs_lifecycle_worker_readiness().ready: imports.append("tasks.knowledge_fs_initial_source_tasks") imports.append("tasks.knowledge_fs_lifecycle_tasks") + imports.append("tasks.knowledge_fs_source_import_tasks") beat_schedule["knowledge_fs_lifecycle_worker"] = { "task": "tasks.knowledge_fs_lifecycle_tasks.run_knowledge_fs_lifecycle_worker", "schedule": timedelta(seconds=dify_config.KNOWLEDGE_FS_LIFECYCLE_POLL_INTERVAL_SECONDS), diff --git a/api/services/knowledge_fs/source_import_commit_service.py b/api/services/knowledge_fs/source_import_commit_service.py index b32fc7f1374..76248f597f9 100644 --- a/api/services/knowledge_fs/source_import_commit_service.py +++ b/api/services/knowledge_fs/source_import_commit_service.py @@ -151,7 +151,8 @@ def resume_committed_source_import( "syncPolicy": last_import.get("syncPolicy"), } metadata = dict(source.metadata) - metadata.pop("lastImport", None) + # updateSource merges metadata; null explicitly supersedes the terminal marker while retrying. + metadata["lastImport"] = None facade.update_source( tenant_id=tenant_id, account_id=account_id, diff --git a/api/tasks/knowledge_fs_source_import_tasks.py b/api/tasks/knowledge_fs_source_import_tasks.py index 6ed277f2733..35a69e77f3f 100644 --- a/api/tasks/knowledge_fs_source_import_tasks.py +++ b/api/tasks/knowledge_fs_source_import_tasks.py @@ -14,6 +14,11 @@ from services.knowledge_fs.product_remote import KnowledgeFSProductRemoteError, from services.knowledge_fs.runtime import get_knowledge_fs_runtime _ACTIVE_STATES = {"queued", "running", "crawling", "importing", "syncing"} +_ASYNC_IMPORT_KINDS = { + "crawl-preview-selection", + "online-document-import", + "online-drive-import", +} _PENDING_IMPORT_KEY = "pendingImport" @@ -41,12 +46,25 @@ def finalize_source_import_once( source_id=source_id, ) pending = source.metadata.get(_PENDING_IMPORT_KEY) - if not isinstance(pending, dict) or pending.get("workflowId") != workflow_id: + last_import = source.metadata.get("lastImport") + completed_import = ( + last_import + if isinstance(last_import, dict) + and last_import.get("kind") in _ASYNC_IMPORT_KINDS + and last_import.get("state") == "completed" + and last_import.get("workflowId") == workflow_id + else None + ) + if (not isinstance(pending, dict) or pending.get("workflowId") != workflow_id) and completed_import is None: return workflow_id metadata = dict(source.metadata) - metadata.pop(_PENDING_IMPORT_KEY, None) if workflow.state != "completed": + if not isinstance(pending, dict): + return workflow_id + # updateSource applies a metadata merge patch, so omission preserves the old marker. + # An explicit null is the tombstone consumed by Dify/UI readers. + metadata[_PENDING_IMPORT_KEY] = None failure = { "errorCode": workflow.last_error_code, "errorMessage": workflow.failure.message if workflow.failure is not None else None, @@ -69,7 +87,36 @@ def finalize_source_import_once( ) return workflow_id - desired_policy = KnowledgeFSDeferredSyncPolicyPayload.model_validate(pending.get("syncPolicy")) + import_metadata = completed_import or pending + if not isinstance(import_metadata, dict): + return workflow_id + if completed_import is None: + completion = { + "kind": import_metadata.get("kind"), + **( + {"previewWorkflowId": import_metadata.get("previewWorkflowId")} + if import_metadata.get("previewWorkflowId") is not None + else {} + ), + "state": "completed", + "syncPolicy": import_metadata.get("syncPolicy"), + "workflowId": workflow.id, + } + metadata[_PENDING_IMPORT_KEY] = None + source = facade.update_source( + tenant_id=tenant_id, + account_id=account_id, + control_space_id=control_space_id, + source_id=source_id, + payload=KnowledgeFSSourceUpdatePayload( + expectedVersion=source.version, + metadata={**metadata, "lastImport": completion, "preview": False}, + status="active", + ), + ) + import_metadata = completion + + desired_policy = KnowledgeFSDeferredSyncPolicyPayload.model_validate(import_metadata.get("syncPolicy")) try: current_policy = facade.get_source_sync_policy( tenant_id=tenant_id, @@ -93,17 +140,6 @@ def finalize_source_import_once( expectedSourceVersion=source.version, ), ) - facade.update_source( - tenant_id=tenant_id, - account_id=account_id, - control_space_id=control_space_id, - source_id=source_id, - payload=KnowledgeFSSourceUpdatePayload( - expectedVersion=source.version, - metadata={**metadata, "preview": False}, - status="active", - ), - ) return workflow_id diff --git a/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py b/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py index 78f4c6b031a..5dc3af70642 100644 --- a/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py +++ b/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py @@ -57,6 +57,7 @@ def test_celery_registers_initial_source_task_when_knowledge_fs_lifecycle_is_rea assert "tasks.knowledge_fs_initial_source_preview_tasks" in celery_app.conf["imports"] assert "tasks.knowledge_fs_failed_retrieval_tasks" in celery_app.conf["imports"] assert "tasks.knowledge_fs_lifecycle_tasks" in celery_app.conf["imports"] + assert "tasks.knowledge_fs_source_import_tasks" in celery_app.conf["imports"] assert "tasks.delete_conversation_task" in celery_app.conf["imports"] assert celery_app.conf["beat_schedule"]["conversation_cleanup_sweeper"] == { "task": "tasks.delete_conversation_task.sweep_deleted_conversations", @@ -84,3 +85,4 @@ def test_celery_registers_initial_source_task_when_knowledge_fs_lifecycle_is_rea assert "tasks.knowledge_fs_failed_retrieval_tasks" in preview_only_app.conf["imports"] assert "tasks.knowledge_fs_initial_source_tasks" not in preview_only_app.conf["imports"] assert "tasks.knowledge_fs_lifecycle_tasks" not in preview_only_app.conf["imports"] + assert "tasks.knowledge_fs_source_import_tasks" not in preview_only_app.conf["imports"] diff --git a/api/tests/unit_tests/services/test_knowledge_fs_source_import_commit_service.py b/api/tests/unit_tests/services/test_knowledge_fs_source_import_commit_service.py index 6067ebb014e..02de10cbd7a 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_source_import_commit_service.py +++ b/api/tests/unit_tests/services/test_knowledge_fs_source_import_commit_service.py @@ -157,6 +157,6 @@ def test_resume_committed_source_import_restores_pending_marker() -> None: update = facade.update_source.call_args.kwargs["payload"] assert update.status == "syncing" - assert "lastImport" not in update.metadata + assert update.metadata["lastImport"] is None assert update.metadata["pendingImport"]["workflowId"] == "import-1" delay.assert_called_once() diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_source_import_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_source_import_tasks.py index 2a91b46b4c0..006c538048c 100644 --- a/api/tests/unit_tests/tasks/test_knowledge_fs_source_import_tasks.py +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_source_import_tasks.py @@ -49,10 +49,25 @@ def test_finalize_source_import_waits_for_terminal_workflow() -> None: facade.get_source.assert_not_called() -def test_finalize_source_import_applies_policy_then_activates_source() -> None: +def test_finalize_source_import_activates_source_then_applies_policy_for_new_version() -> None: facade = MagicMock() - facade.get_source_workflow.return_value = SimpleNamespace(state="completed") + facade.get_source_workflow.return_value = SimpleNamespace(id="import-1", state="completed") facade.get_source.return_value = _source() + facade.update_source.return_value = SimpleNamespace( + id="source-1", + metadata={ + "lastImport": { + "kind": "crawl-preview-selection", + "previewWorkflowId": "preview-1", + "state": "completed", + "syncPolicy": {"enabled": False, "mode": "manual"}, + "workflowId": "import-1", + }, + "preview": False, + }, + status="active", + version=5, + ) facade.get_source_sync_policy.side_effect = KnowledgeFSProductResourceNotFoundError("missing") assert _run(facade) == "import-1" @@ -61,10 +76,38 @@ def test_finalize_source_import_applies_policy_then_activates_source() -> None: assert policy.enabled is False assert policy.mode == "manual" assert policy.expected_revision == 0 - assert policy.expected_source_version == 4 + assert policy.expected_source_version == 5 update = facade.update_source.call_args.kwargs["payload"] assert update.status == "active" - assert "pendingImport" not in update.metadata + assert update.metadata["pendingImport"] is None + assert update.metadata["lastImport"]["state"] == "completed" + + +def test_finalize_source_import_retries_policy_after_source_activation() -> None: + facade = MagicMock() + facade.get_source_workflow.return_value = SimpleNamespace(id="import-1", state="completed") + facade.get_source.return_value = SimpleNamespace( + id="source-1", + metadata={ + "lastImport": { + "kind": "crawl-preview-selection", + "previewWorkflowId": "preview-1", + "state": "completed", + "syncPolicy": {"enabled": True, "mode": "interval"}, + "workflowId": "import-1", + }, + "preview": False, + }, + status="active", + version=5, + ) + facade.get_source_sync_policy.side_effect = KnowledgeFSProductResourceNotFoundError("missing") + + assert _run(facade) == "import-1" + + facade.update_source.assert_not_called() + policy = facade.update_source_sync_policy.call_args.kwargs["payload"] + assert policy.expected_source_version == 5 def test_finalize_source_import_persists_failure_on_visible_source() -> None: @@ -90,4 +133,5 @@ def test_finalize_source_import_persists_failure_on_visible_source() -> None: "syncPolicy": {"enabled": False, "mode": "manual"}, "workflowId": "import-1", } + assert update.metadata["pendingImport"] is None facade.update_source_sync_policy.assert_not_called() diff --git a/web/features/new-rag/__tests__/connected-source-setup.spec.tsx b/web/features/new-rag/__tests__/connected-source-setup.spec.tsx index 358af1173b0..f381f90b412 100644 --- a/web/features/new-rag/__tests__/connected-source-setup.spec.tsx +++ b/web/features/new-rag/__tests__/connected-source-setup.spec.tsx @@ -752,6 +752,62 @@ describe('ConnectedSourceSetup', () => { expect(clientMock.deleteSource).not.toHaveBeenCalled() }) + it('reconciles an accepted import when the response is lost before navigation', async () => { + const user = userEvent.setup() + clientMock.listDatasourceAuth.mockResolvedValue({ + result: [notionDatasourceAuth([notionCredential])], + }) + clientMock.listConnections.mockResolvedValue({ + data: [connectionResponse()], + next_cursor: null, + } satisfies KnowledgeFsSourceConnectionListResponse) + clientMock.getPages.mockResolvedValue({ + next_cursor: null, + workspaces: [ + { + pages: [ + { + last_edited_time: null, + page_id: 'page-1', + page_name: 'Product roadmap', + parent_id: null, + type: 'page', + }, + ], + total: 1, + workspace_id: 'workspace-1', + workspace_name: 'Acme workspace', + }, + ], + }) + clientMock.createWorkflowImport.mockRejectedValue(new TypeError('Failed to fetch')) + clientMock.getSource.mockResolvedValue( + sourceResponse({ + metadata: { + pendingImport: { + kind: 'online-document-import', + syncPolicy: { enabled: true, mode: 'provider' }, + workflowId: 'import-run-1', + }, + preview: false, + }, + name: 'Team wiki', + status: 'disabled', + version: 5, + }), + ) + const view = renderSetup({ ...defaultDraft, sourceName: 'Team wiki' }) + + await user.click(await screen.findByRole('checkbox', { name: 'Product roadmap' })) + await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.addSource' })) + + await waitFor(() => expect(view.onCompleted).toHaveBeenCalledOnce()) + expect(screen.queryByRole('alert')).not.toBeInTheDocument() + + view.unmount() + expect(clientMock.deleteSource).not.toHaveBeenCalled() + }) + it('applies datasource parameters once instead of rebuilding the preview on every keypress', async () => { const user = userEvent.setup() clientMock.listDatasourcePlugins.mockResolvedValue([notionDatasourcePluginWithParameters]) diff --git a/web/features/new-rag/__tests__/sources-page.spec.tsx b/web/features/new-rag/__tests__/sources-page.spec.tsx index 5a450b71678..b73d536261f 100644 --- a/web/features/new-rag/__tests__/sources-page.spec.tsx +++ b/web/features/new-rag/__tests__/sources-page.spec.tsx @@ -828,9 +828,17 @@ describe('SourcesPage', () => { { items: [ source({ + metadata: { + lastImport: { + kind: 'online-document-import', + state: 'failed', + syncPolicy: { enabled: true, mode: 'provider' }, + workflowId: 'import-workflow', + }, + preview: false, + }, name: 'Failed connected source', status: 'error', - syncWorkflow: sourceWorkflow('failed'), }), ], }, @@ -856,7 +864,7 @@ describe('SourcesPage', () => { await waitFor(() => expect(clientMock.retrySourceWorkflow).toHaveBeenCalledOnce()) expect(clientMock.retrySourceWorkflow).toHaveBeenCalledWith({ - params: { control_space_id: 'space-1', run_id: 'workflow-1' }, + params: { control_space_id: 'space-1', run_id: 'import-workflow' }, }) expect(clientMock.syncSource).not.toHaveBeenCalled() expect(within(sourceRow).getByRole('status')).toHaveTextContent( @@ -864,6 +872,41 @@ describe('SourcesPage', () => { ) }) + it('shows and polls an accepted async import even when the source status is briefly disabled', () => { + const pendingSource = source({ + metadata: { + pendingImport: { + kind: 'online-document-import', + syncPolicy: { enabled: true, mode: 'provider' }, + workflowId: 'import-workflow', + }, + preview: false, + }, + name: 'Importing connected source', + status: 'disabled', + }) + sourcesQuery.data = { pages: [{ items: [pendingSource] }] } + + render() + + const sourceRow = screen.getByRole('row', { name: /Importing connected source/ }) + expect(within(sourceRow).getByRole('status')).toHaveTextContent( + 'dataset.newKnowledge.sourceStatus.syncing', + ) + const options = infiniteOptionsMock.mock.lastCall?.[0] + expect(options).toBeDefined() + if (!options) throw new Error('Expected source infinite query options') + expect( + options.refetchInterval({ + state: { + data: { + pages: [{ data: [sourceApiResponse(pendingSource)] }], + }, + }, + }), + ).toBe(3000) + }) + it('keeps polling when the source list briefly returns the stale failed retry snapshot', async () => { const user = userEvent.setup() const staleFailedWorkflow: SourceWorkflowRun = { diff --git a/web/features/new-rag/connected-source-setup.tsx b/web/features/new-rag/connected-source-setup.tsx index 88960b41a11..9a1e65f7635 100644 --- a/web/features/new-rag/connected-source-setup.tsx +++ b/web/features/new-rag/connected-source-setup.tsx @@ -42,7 +42,9 @@ import { sourceConnectionFromApi, sourceConnectionListFromApi, sourceFromApi, + sourceHasPendingAsyncImport, sourceProviderListFromApi, + sourceWorkflowFromApi, } from './source-models' import { discoverSourceProviderOptions, @@ -1309,44 +1311,62 @@ function ResourceConfiguration({ requestId: createRequestId(), } } - await (!driveTransport - ? consoleClient.knowledgeFs.spaces.byControlSpaceId.sources.bySourceId.asyncImport.post({ - body: { - items: selectedPages.map((resource) => ({ - lastEditedTime: resource.page.last_edited_time ?? undefined, - name: resource.page.page_name, - pageId: resource.page.page_id, - providerItemId: JSON.stringify([resource.groupId, resource.page.page_id]), - type: resource.page.type, - workspaceId: resource.groupId, - })), - kind: 'online-document-import', - syncPolicy: policy, - }, - headers: { 'Idempotency-Key': importRequestRef.current.requestId }, - params: { - control_space_id: knowledgeSpaceId, - source_id: finalSource.id, - }, - }) - : consoleClient.knowledgeFs.spaces.byControlSpaceId.sources.bySourceId.asyncImport.post({ - body: { - items: selectedFiles.map((resource) => ({ - bucket: resource.bucket, - id: resource.file.id, - mimeType: resource.file.type.includes('/') ? resource.file.type : undefined, - name: resource.file.name, - providerItemId: JSON.stringify([resource.bucket ?? '', resource.file.id]), - })), - kind: 'online-drive-import', - syncPolicy: policy, - }, - headers: { 'Idempotency-Key': importRequestRef.current.requestId }, - params: { - control_space_id: knowledgeSpaceId, - source_id: finalSource.id, - }, - })) + const importWorkflow = sourceWorkflowFromApi( + await (!driveTransport + ? consoleClient.knowledgeFs.spaces.byControlSpaceId.sources.bySourceId.asyncImport.post({ + body: { + items: selectedPages.map((resource) => ({ + lastEditedTime: resource.page.last_edited_time ?? undefined, + name: resource.page.page_name, + pageId: resource.page.page_id, + providerItemId: JSON.stringify([resource.groupId, resource.page.page_id]), + type: resource.page.type, + workspaceId: resource.groupId, + })), + kind: 'online-document-import', + syncPolicy: policy, + }, + headers: { 'Idempotency-Key': importRequestRef.current.requestId }, + params: { + control_space_id: knowledgeSpaceId, + source_id: finalSource.id, + }, + }) + : consoleClient.knowledgeFs.spaces.byControlSpaceId.sources.bySourceId.asyncImport.post({ + body: { + items: selectedFiles.map((resource) => ({ + bucket: resource.bucket, + id: resource.file.id, + mimeType: resource.file.type.includes('/') ? resource.file.type : undefined, + name: resource.file.name, + providerItemId: JSON.stringify([resource.bucket ?? '', resource.file.id]), + })), + kind: 'online-drive-import', + syncPolicy: policy, + }, + headers: { 'Idempotency-Key': importRequestRef.current.requestId }, + params: { + control_space_id: knowledgeSpaceId, + source_id: finalSource.id, + }, + })), + ) + const importKind = driveTransport ? 'online-drive-import' : 'online-document-import' + const committedSource: Source = { + ...finalSource, + metadata: { + ...finalSource.metadata, + pendingImport: { + kind: importKind, + syncPolicy: policy, + workflowId: importWorkflow.id, + }, + preview: false, + }, + status: 'syncing', + } + previewSourceRef.current = committedSource + setPreviewSource(committedSource) await completeSubmission() } catch { try { @@ -1363,7 +1383,8 @@ function ResourceConfiguration({ setPreviewSource(reconciledSource) if ( reconciledSource.metadata.preview === false && - reconciledSource.status !== 'disabled' + (reconciledSource.status !== 'disabled' || + sourceHasPendingAsyncImport(reconciledSource)) ) { await completeSubmission() return diff --git a/web/features/new-rag/source-models.ts b/web/features/new-rag/source-models.ts index 2e51853a9df..2c582162933 100644 --- a/web/features/new-rag/source-models.ts +++ b/web/features/new-rag/source-models.ts @@ -135,6 +135,23 @@ const SOURCE_WORKFLOW_FAILURE_STATES = new Set([ 'timed_out', 'timeout', ]) +const ASYNC_SOURCE_IMPORT_KINDS = new Set([ + 'crawl-preview-selection', + 'online-document-import', + 'online-drive-import', +]) + +function sourceImportMetadata(value: unknown) { + if (typeof value !== 'object' || value === null || Array.isArray(value)) return undefined + const metadata = value as Record + if ( + typeof metadata.kind !== 'string' || + !ASYNC_SOURCE_IMPORT_KINDS.has(metadata.kind) || + typeof metadata.workflowId !== 'string' + ) + return undefined + return metadata +} export function sourceWorkflowStatus(state: string): Source['status'] { const normalized = state.trim().toLowerCase().replaceAll('-', '_').replaceAll(' ', '_') @@ -179,6 +196,15 @@ export function initialSourceWorkflowId(source: Source) { return initialImport.workflowId } +export function sourceAsyncImportWorkflowId(source: Source) { + return (sourceImportMetadata(source.metadata.pendingImport)?.workflowId ?? + sourceImportMetadata(source.metadata.lastImport)?.workflowId) as string | undefined +} + +export function sourceHasPendingAsyncImport(source: Source) { + return sourceImportMetadata(source.metadata.pendingImport) !== undefined +} + export function sourceDisplayStatus(source: Source): SourceDisplayStatus { if (isInitialSource(source) && source.status === 'disabled' && source.metadata.preview === true) return 'initializing' @@ -192,6 +218,8 @@ export function sourceDisplayStatus(source: Source): SourceDisplayStatus { return 'error' } + if (sourceHasPendingAsyncImport(source) && source.status === 'disabled') return 'syncing' + return source.status } @@ -204,6 +232,7 @@ export function shouldHidePreviewSource(source: Source) { export function sourceNeedsPolling(source: Source) { return ( sourceDisplayStatus(source) === 'initializing' || + sourceHasPendingAsyncImport(source) || source.status === 'syncing' || sourceWorkflowIsActive(source.syncWorkflow) ) diff --git a/web/features/new-rag/sources-page.tsx b/web/features/new-rag/sources-page.tsx index 3f85458c079..aaacea31f03 100644 --- a/web/features/new-rag/sources-page.tsx +++ b/web/features/new-rag/sources-page.tsx @@ -56,6 +56,7 @@ import { initialSourceWorkflowId, isInitialSourceForOperation, shouldHidePreviewSource, + sourceAsyncImportWorkflowId, sourceDisplayStatus, sourceFromApi, sourceNeedsPolling, @@ -709,7 +710,8 @@ function SourceRow({ ) const retrySource = () => { - const retryWorkflowId = initialWorkflowId ?? syncWorkflow?.id + const retryWorkflowId = + initialWorkflowId ?? sourceAsyncImportWorkflowId(source) ?? syncWorkflow?.id if (!retryWorkflowId) return syncSource() return runAction(