opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-08 02:43:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(knowledge-fs): finalize async source imports

Browse Source
This commit is contained in:
Stephen Zhou 2026-08-20 12:01:24 +08:00
parent 4880c5c771
commit d061d5ed3e
No known key found for this signature in database
11 changed files with 297 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
api/extensions/ext_celery.py
View File

@ -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),

3
api/services/knowledge_fs/source_import_commit_service.py
View File

@ -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,

64
api/tasks/knowledge_fs_source_import_tasks.py
View File

@ -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

2
api/tests/unit_tests/extensions/test_knowledge_fs_celery.py
View File

@ -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"]

2
api/tests/unit_tests/services/test_knowledge_fs_source_import_commit_service.py
View File

@ -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()

52
api/tests/unit_tests/tasks/test_knowledge_fs_source_import_tasks.py
View File

@ -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()

56
web/features/new-rag/__tests__/connected-source-setup.spec.tsx
View File

@ -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])

47
web/features/new-rag/__tests__/sources-page.spec.tsx
View File

@ -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(<SourcesPage knowledgeSpaceId="space-1" />)
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 = {

99
web/features/new-rag/connected-source-setup.tsx
View File

@ -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

29
web/features/new-rag/source-models.ts
View File

@ -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<string, unknown>
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)
)

4
web/features/new-rag/sources-page.tsx
View File

@ -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(