fix(knowledge-fs): stabilize initial source creation

This commit is contained in:
Stephen Zhou 2026-08-10 22:35:24 +08:00
parent e4e9271fbe
commit a2194300cc
No known key found for this signature in database
5 changed files with 441 additions and 52 deletions

View File

@ -32,7 +32,7 @@ from services.knowledge_fs.product_dto import (
KnowledgeFSSourceUpdatePayload,
KnowledgeFSSourceWorkflowImportPayload,
)
from services.knowledge_fs.product_remote import KnowledgeFSProductResourceNotFoundError
from services.knowledge_fs.product_remote import KnowledgeFSProductRemoteError, KnowledgeFSProductResourceNotFoundError
from services.knowledge_fs.runtime import get_knowledge_fs_runtime
_LEGACY_WEBSITE_PLUGIN_IDS = {
@ -554,6 +554,10 @@ def _run_initial_source_task(
},
)
raise task.retry(exc=exc)
except KnowledgeFSProductResourceNotFoundError:
raise
except KnowledgeFSProductRemoteError as exc:
raise task.retry(exc=exc)
@shared_task(bind=True, queue="knowledge_fs_lifecycle", max_retries=180, default_retry_delay=2)
@ -614,6 +618,10 @@ def import_initial_website_source(
},
)
raise self.retry(exc=exc)
except KnowledgeFSProductResourceNotFoundError:
raise
except KnowledgeFSProductRemoteError as exc:
raise self.retry(exc=exc)
__all__ = [

View File

@ -10,7 +10,7 @@ from services.knowledge_fs.product_dto import (
KnowledgeFSInitialOnlineDriveSourcePayload,
KnowledgeFSInitialWebsiteSourcePayload,
)
from services.knowledge_fs.product_remote import KnowledgeFSProductResourceNotFoundError
from services.knowledge_fs.product_remote import KnowledgeFSProductRemoteError, KnowledgeFSProductResourceNotFoundError
from tasks.knowledge_fs_initial_source_tasks import (
KnowledgeFSInitialSourceNotReadyError,
import_initial_source,
@ -693,3 +693,50 @@ def test_initial_source_task_validates_discriminated_connector_payload() -> None
parsed_payload = start_import.call_args.kwargs["payload"]
assert isinstance(parsed_payload, KnowledgeFSInitialOnlineDocumentSourcePayload)
assert parsed_payload.credential_id == "notion-credential-1"
def test_initial_source_task_retries_transient_remote_error() -> None:
serialized_payload = _drive_payload().model_dump(mode="json", by_alias=True)
remote_error = KnowledgeFSProductRemoteError("temporary outage")
retry_error = RuntimeError("retry requested")
with (
patch(
"tasks.knowledge_fs_initial_source_tasks.start_initial_source_import",
side_effect=remote_error,
),
patch.object(import_initial_source, "retry", side_effect=retry_error) as retry,
pytest.raises(RuntimeError, match="retry requested"),
):
import_initial_source.run(
tenant_id="tenant-1",
account_id="account-1",
control_space_id="control-1",
operation_id="operation-1",
payload=serialized_payload,
workflow_id="workflow-1",
)
retry.assert_called_once_with(exc=remote_error)
def test_initial_source_task_does_not_retry_authoritative_missing_resource() -> None:
serialized_payload = _drive_payload().model_dump(mode="json", by_alias=True)
missing_resource = KnowledgeFSProductResourceNotFoundError("workflow was not found")
with (
patch(
"tasks.knowledge_fs_initial_source_tasks.start_initial_source_import",
side_effect=missing_resource,
),
patch.object(import_initial_source, "retry") as retry,
pytest.raises(KnowledgeFSProductResourceNotFoundError, match="workflow was not found"),
):
import_initial_source.run(
tenant_id="tenant-1",
account_id="account-1",
control_space_id="control-1",
operation_id="operation-1",
payload=serialized_payload,
workflow_id="workflow-1",
)
retry.assert_not_called()

View File

@ -239,6 +239,37 @@ const firecrawlDatasourceAuth = {
provider: 'firecrawl',
}
const jinaDatasourcePlugin = {
...firecrawlDatasourcePlugin,
declaration: {
...firecrawlDatasourcePlugin.declaration,
identity: {
...firecrawlDatasourcePlugin.declaration.identity,
label: { en_US: 'Jina Reader', zh_Hans: 'Jina Reader' },
name: 'jinareader',
},
},
plugin_id: 'langgenius/jina_datasource',
plugin_unique_identifier: 'langgenius/jina_datasource:1.0.0@local',
provider: 'jinareader',
}
const jinaDatasourceAuth = {
...firecrawlDatasourceAuth,
credentials_list: [
{
...firecrawlDatasourceAuth.credentials_list[0],
id: 'jina-credential-1',
name: 'Default Jina Reader',
},
],
label: { en_US: 'Jina Reader' },
name: 'jinareader',
plugin_id: 'langgenius/jina_datasource',
plugin_unique_identifier: 'langgenius/jina_datasource:1.0.0@local',
provider: 'jinareader',
}
const notionDatasourcePlugin = {
...firecrawlDatasourcePlugin,
declaration: {
@ -1107,6 +1138,54 @@ describe('CreateKnowledgePage', () => {
expect(screen.getByText('Getting started')).toBeInTheDocument()
})
it('creates an initial source from a synchronous Jina Reader preview', async () => {
const user = userEvent.setup()
navigationMock.startMode = 'source'
datasourceQueryMock.plugins.data = [firecrawlDatasourcePlugin, jinaDatasourcePlugin]
datasourceQueryMock.auth.data = { result: [firecrawlDatasourceAuth, jinaDatasourceAuth] }
serviceMock.createCrawl.mockResolvedValueOnce({
data: {
content: '# Dify introduction',
description: 'Introduction',
title: 'Dify introduction',
url: 'https://docs.dify.ai/introduction',
},
})
renderPage()
await fillRequiredFields(user)
await user.click(screen.getByRole('radio', { name: 'Jina Reader' }))
await user.type(
screen.getByPlaceholderText('dataset.newKnowledge.rootUrlPlaceholder'),
'https://docs.dify.ai/introduction',
)
await user.type(
screen.getByPlaceholderText('dataset.newKnowledge.sourceNamePlaceholder'),
'Dify introduction',
)
await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.crawlAndPreview' }))
await user.click(await screen.findByRole('checkbox', { name: 'Dify introduction' }))
expect(serviceMock.getCrawlStatus).not.toHaveBeenCalled()
await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.createTitle' }))
await waitFor(() => expect(serviceMock.create).toHaveBeenCalledOnce())
expect(serviceMock.create).toHaveBeenCalledWith({
body: expect.objectContaining({
initial_source: expect.objectContaining({
credentialId: 'jina-credential-1',
pluginId: 'langgenius/jina_datasource',
provider: 'jinareader',
selection: [
expect.objectContaining({
source_url: 'https://docs.dify.ai/introduction',
title: 'Dify introduction',
}),
],
}),
}),
})
})
it('shows and can stop an ongoing website crawl', async () => {
const user = userEvent.setup()
navigationMock.startMode = 'source'
@ -1338,6 +1417,203 @@ describe('CreateKnowledgePage', () => {
})
})
it('uses the online-drive transport when creating with Google Docs', async () => {
const user = userEvent.setup()
navigationMock.startMode = 'source'
datasourceQueryMock.plugins.data = [
firecrawlDatasourcePlugin,
notionDatasourcePlugin,
googleDriveDatasourcePlugin,
]
datasourceQueryMock.auth.data = {
result: [firecrawlDatasourceAuth, notionDatasourceAuth, googleDriveDatasourceAuth],
}
serviceMock.previewInitialSource.mockResolvedValue({
files: [
{
bucket: null,
id: 'doc-1',
mime_type: 'application/vnd.google-apps.document',
name: 'Launch plan',
provider_item_id: '["","doc-1"]',
size: 1024,
type: 'application/vnd.google-apps.document',
},
],
kind: 'online_drive',
next_page_parameters: null,
})
renderPage()
await fillRequiredFields(user)
await user.click(screen.getByRole('radio', { name: 'dataset.newKnowledge.onlineDocuments' }))
await user.click(screen.getByRole('radio', { name: 'Google Docs' }))
await user.type(
screen.getByPlaceholderText('dataset.newKnowledge.sourceNamePlaceholder'),
'Team docs',
)
await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.preview' }))
await waitFor(() =>
expect(serviceMock.previewInitialSource).toHaveBeenCalledWith({
body: expect.objectContaining({ kind: 'online_drive' }),
}),
)
await user.click(await screen.findByRole('checkbox', { name: 'Launch plan' }))
await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.createTitle' }))
await waitFor(() => expect(serviceMock.create).toHaveBeenCalledOnce())
expect(serviceMock.create).toHaveBeenCalledWith({
body: expect.objectContaining({
initial_source: expect.objectContaining({
kind: 'online_drive',
selection: [expect.objectContaining({ id: 'doc-1' })],
}),
}),
})
})
it('preserves root pagination while expanding a drive folder', async () => {
const user = userEvent.setup()
navigationMock.startMode = 'source'
datasourceQueryMock.plugins.data = [firecrawlDatasourcePlugin, googleDriveDatasourcePlugin]
datasourceQueryMock.auth.data = {
result: [firecrawlDatasourceAuth, googleDriveDatasourceAuth],
}
serviceMock.previewInitialSource.mockImplementation(
({ body }: { body: { parameters: Record<string, unknown> } }) => {
if (body.parameters.prefix === 'folder-1') {
return Promise.resolve({
files: [
{
bucket: null,
id: 'child-1',
mime_type: 'application/pdf',
name: 'Folder child.pdf',
provider_item_id: '["","child-1"]',
size: 128,
type: 'application/pdf',
},
],
kind: 'online_drive',
next_page_parameters: null,
})
}
if (body.parameters.next_page_parameters) {
return Promise.resolve({
files: [
{
bucket: null,
id: 'root-2',
mime_type: 'application/pdf',
name: 'Second root file.pdf',
provider_item_id: '["","root-2"]',
size: 256,
type: 'application/pdf',
},
],
kind: 'online_drive',
next_page_parameters: null,
})
}
return Promise.resolve({
files: [
{
bucket: null,
id: 'folder-1',
mime_type: null,
name: 'Plans',
provider_item_id: '["","folder-1"]',
size: 0,
type: 'folder',
},
],
kind: 'online_drive',
next_page_parameters: { cursor: 'root-next' },
})
},
)
renderPage()
await user.click(screen.getByRole('radio', { name: 'dataset.newKnowledge.onlineDrive' }))
await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.preview' }))
const folderButton = await screen.findByRole('button', { name: 'Plans' })
expect(folderButton).toHaveAttribute('aria-expanded', 'false')
await user.click(folderButton)
expect(folderButton).toHaveAttribute('aria-expanded', 'true')
expect(await screen.findByRole('checkbox', { name: 'Folder child.pdf' })).toBeInTheDocument()
await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.loadMore' }))
expect(
await screen.findByRole('checkbox', { name: 'Second root file.pdf' }),
).toBeInTheDocument()
expect(serviceMock.previewInitialSource).toHaveBeenLastCalledWith({
body: expect.objectContaining({
parameters: { next_page_parameters: { cursor: 'root-next' } },
}),
})
})
it('limits a paginated drive selection to the backend maximum', async () => {
const user = userEvent.setup()
navigationMock.startMode = 'source'
datasourceQueryMock.plugins.data = [firecrawlDatasourcePlugin, googleDriveDatasourcePlugin]
datasourceQueryMock.auth.data = {
result: [firecrawlDatasourceAuth, googleDriveDatasourceAuth],
}
const files = Array.from({ length: 200 }, (_, index) => ({
bucket: null,
id: `file-${index + 1}`,
mime_type: 'application/pdf',
name: `File ${index + 1}.pdf`,
provider_item_id: `["","file-${index + 1}"]`,
size: 128,
type: 'application/pdf',
}))
serviceMock.previewInitialSource
.mockResolvedValueOnce({
files,
kind: 'online_drive',
next_page_parameters: { cursor: 'next' },
})
.mockResolvedValueOnce({
files: [
{
bucket: null,
id: 'file-201',
mime_type: 'application/pdf',
name: 'File 201.pdf',
provider_item_id: '["","file-201"]',
size: 128,
type: 'application/pdf',
},
],
kind: 'online_drive',
next_page_parameters: null,
})
renderPage()
await fillRequiredFields(user)
await user.click(screen.getByRole('radio', { name: 'dataset.newKnowledge.onlineDrive' }))
await user.type(
screen.getByPlaceholderText('dataset.newKnowledge.sourceNamePlaceholder'),
'Drive archive',
)
await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.preview' }))
await user.click(await screen.findByRole('button', { name: 'dataset.newKnowledge.loadMore' }))
await screen.findByRole('checkbox', { name: 'File 201.pdf' })
await user.click(screen.getByRole('checkbox', { name: 'dataset.newKnowledge.selectAll' }))
expect(screen.getByRole('checkbox', { name: 'File 201.pdf' })).toHaveAttribute(
'aria-disabled',
'true',
)
expect(screen.getByText(/^dataset\.newKnowledge\.pagesSelected/)).toHaveTextContent('200')
await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.createTitle' }))
await waitFor(() => expect(serviceMock.create).toHaveBeenCalledOnce())
expect(serviceMock.create.mock.calls[0]?.[0].body.initial_source.selection).toHaveLength(200)
})
it('requires a selected website preview page before creating with an initial source', async () => {
const user = userEvent.setup()
navigationMock.startMode = 'source'

View File

@ -41,6 +41,10 @@ type NextPageRequest = {
prefix?: string
}
const MAX_SELECTION = 200
const ROOT_PAGE_SCOPE = 'root'
const SELECTION_LIMIT_ID = 'create-connected-source-selection-limit'
function isDriveContainer(file: PreviewFile) {
return /bucket|directory|folder|workspace/i.test(file.type)
}
@ -80,7 +84,11 @@ export function CreateConnectedSourceSetup({
const [loadingMore, setLoadingMore] = useState(false)
const [error, setError] = useState(false)
const [previewed, setPreviewed] = useState(false)
const [nextPageRequest, setNextPageRequest] = useState<NextPageRequest | null>()
const [nextPageRequests, setNextPageRequests] = useState<Map<string, NextPageRequest>>(
() => new Map(),
)
const driveTransport = providerOption.providerType === 'online_drive'
const selectionAtLimit = selected.size >= MAX_SELECTION
const selectableResources = useMemo(
() =>
resources.filter(
@ -113,7 +121,7 @@ export function CreateConnectedSourceSetup({
pluginId: providerOption.plugin.plugin_id,
provider: providerOption.plugin.provider,
}
if (draft.sourceType === 'onlineDocuments') {
if (!driveTransport) {
onInitialSourceChange({
...binding,
kind: 'online_document',
@ -158,6 +166,7 @@ export function CreateConnectedSourceSetup({
}, [
credential.id,
draft,
driveTransport,
onInitialSourceChange,
providerOption.datasource.identity.name,
providerOption.plugin.plugin_id,
@ -189,7 +198,7 @@ export function CreateConnectedSourceSetup({
body: {
credentialId: credential.id,
datasource: providerOption.datasource.identity.name,
kind: draft.sourceType === 'onlineDocuments' ? 'online_document' : 'online_drive',
kind: driveTransport ? 'online_drive' : 'online_document',
parameters: {
...(bucket ? { bucket } : {}),
...(prefix ? { prefix } : {}),
@ -199,39 +208,43 @@ export function CreateConnectedSourceSetup({
provider: providerOption.plugin.provider,
},
})
const nextResources: PreviewResource[] =
draft.sourceType === 'onlineDocuments'
? (response.documents ?? []).map((document) => ({
depth,
document,
key: `document:${document.provider_item_id}`,
kind: 'document' as const,
parentKey,
}))
: (response.files ?? []).map((file) => ({
depth,
file,
key: `file:${file.provider_item_id}`,
kind: 'file' as const,
parentKey,
}))
const nextResources: PreviewResource[] = driveTransport
? (response.files ?? []).map((file) => ({
depth,
file,
key: `file:${file.provider_item_id}`,
kind: 'file' as const,
parentKey,
}))
: (response.documents ?? []).map((document) => ({
depth,
document,
key: `document:${document.provider_item_id}`,
kind: 'document' as const,
parentKey,
}))
setResources((current) => {
const next = new Map((append ? current : []).map((resource) => [resource.key, resource]))
for (const resource of nextResources) next.set(resource.key, resource)
return [...next.values()]
})
if (parentKey) setExpanded((current) => new Set(current).add(parentKey))
setNextPageRequest(
response.next_page_parameters
? {
bucket,
depth,
nextPage: response.next_page_parameters,
parentKey,
prefix,
}
: null,
)
setNextPageRequests((current) => {
const next = append ? new Map(current) : new Map<string, NextPageRequest>()
const scope = parentKey ?? ROOT_PAGE_SCOPE
if (response.next_page_parameters) {
next.set(scope, {
bucket,
depth,
nextPage: response.next_page_parameters,
parentKey,
prefix,
})
} else {
next.delete(scope)
}
return next
})
setPreviewed(true)
} catch {
setError(true)
@ -240,14 +253,14 @@ export function CreateConnectedSourceSetup({
setLoadingMore(false)
}
},
[credential.id, draft.sourceType, providerOption],
[credential.id, driveTransport, providerOption],
)
const toggle = (key: string) => {
setSelected((current) => {
const next = new Set(current)
if (next.has(key)) next.delete(key)
else next.add(key)
else if (next.size < MAX_SELECTION) next.add(key)
return next
})
}
@ -259,7 +272,10 @@ export function CreateConnectedSourceSetup({
selectableResources.every((resource) => current.has(resource.key))
for (const resource of selectableResources) {
if (allSelected) next.delete(resource.key)
else next.add(resource.key)
else {
if (next.size >= MAX_SELECTION) break
next.add(resource.key)
}
}
return next
})
@ -281,6 +297,9 @@ export function CreateConnectedSourceSetup({
prefix: resource.file.id || undefined,
})
}
const visibleNextPageRequests = [...nextPageRequests.entries()].filter(
([scope]) => scope === ROOT_PAGE_SCOPE || expanded.has(scope),
)
return (
<div className="space-y-4">
@ -332,11 +351,21 @@ export function CreateConnectedSourceSetup({
<div className="flex items-center gap-2 border-b border-divider-subtle px-3 py-2">
<Checkbox
aria-label={t(($) => $['newKnowledge.selectAll'])}
aria-describedby={selectionAtLimit ? SELECTION_LIMIT_ID : undefined}
checked={
selectableResources.length > 0 &&
selectableResources.every((resource) => selected.has(resource.key))
}
disabled={disabled || !selectableResources.length}
disabled={
disabled ||
!selectableResources.length ||
(selectionAtLimit &&
!selectableResources.every((resource) => selected.has(resource.key)))
}
indeterminate={
selected.size > 0 &&
!selectableResources.every((resource) => selected.has(resource.key))
}
onCheckedChange={toggleAll}
/>
<span className="system-xs-medium text-text-secondary">
@ -344,8 +373,13 @@ export function CreateConnectedSourceSetup({
? t(($) => $['newKnowledge.selectPagesToSync'])
: t(($) => $['newKnowledge.selectFilesAndFolders'])}
</span>
<span className="ml-auto system-xs-regular text-text-tertiary">
<span role="status" className="ml-auto system-xs-regular text-text-tertiary">
{t(($) => $['newKnowledge.pagesSelected'], { count: selected.size })}
{selectionAtLimit && (
<span id={SELECTION_LIMIT_ID} className="ml-2 text-text-destructive">
{t(($) => $['newKnowledge.maxPages'])}: {MAX_SELECTION}
</span>
)}
</span>
</div>
<ul className="max-h-64 overflow-y-auto p-1.5">
@ -364,12 +398,13 @@ export function CreateConnectedSourceSetup({
size="small"
className="size-5 px-0"
aria-label={resourceName(resource)}
aria-expanded={expanded.has(resource.key)}
disabled={disabled || loadingMore}
onClick={() => expandContainer(resource)}
>
<span
aria-hidden
className={`i-ri-arrow-right-s-line size-4 transition-transform ${
className={`i-ri-arrow-right-s-line size-4 transition-transform motion-reduce:transition-none ${
expanded.has(resource.key) ? 'rotate-90' : ''
}`}
/>
@ -377,8 +412,13 @@ export function CreateConnectedSourceSetup({
) : (
<Checkbox
aria-label={resourceName(resource)}
aria-describedby={
selectionAtLimit && !selected.has(resource.key)
? SELECTION_LIMIT_ID
: undefined
}
checked={selected.has(resource.key)}
disabled={disabled}
disabled={disabled || (selectionAtLimit && !selected.has(resource.key))}
onCheckedChange={() => toggle(resource.key)}
/>
)}
@ -390,17 +430,25 @@ export function CreateConnectedSourceSetup({
)
})}
</ul>
{nextPageRequest && (
<div className="border-t border-divider-subtle px-3 py-2 text-center">
<Button
type="button"
size="small"
variant="ghost"
loading={loadingMore}
onClick={() => void requestPreview({ append: true, ...nextPageRequest })}
>
{t(($) => $['newKnowledge.loadMore'])}
</Button>
{visibleNextPageRequests.length > 0 && (
<div className="flex flex-wrap justify-center gap-2 border-t border-divider-subtle px-3 py-2 text-center">
{visibleNextPageRequests.map(([scope, request]) => {
const parent =
scope === ROOT_PAGE_SCOPE ? undefined : resources.find(({ key }) => key === scope)
return (
<Button
key={scope}
type="button"
size="small"
variant="ghost"
loading={loadingMore}
onClick={() => void requestPreview({ append: true, ...request })}
>
{t(($) => $['newKnowledge.loadMore'])}
{parent ? ` · ${resourceName(parent)}` : ''}
</Button>
)
})}
</div>
)}
</section>

View File

@ -66,8 +66,12 @@ type LocalCrawlState = 'error' | 'idle' | 'running' | 'stopped' | 'success'
type InitialSource = NonNullable<KnowledgeFsSpaceCreatePayload['initial_source']>
function crawlPages(response: Record<string, unknown>): CrawlResultItem[] {
if (!Array.isArray(response.data)) return []
return response.data.flatMap((item) => {
const items = Array.isArray(response.data)
? response.data
: response.data && typeof response.data === 'object'
? [response.data]
: []
return items.flatMap((item) => {
if (!item || typeof item !== 'object') return []
const page = item as Record<string, unknown>
const sourceUrl =
@ -262,6 +266,12 @@ export function CreateSourceSetup({
},
url: draft.rootUrl,
})) as Record<string, unknown>
const synchronousPages = crawlPages(created)
if (synchronousPages.length) {
setPreviewPages(synchronousPages)
setCrawlState('success')
return
}
const jobId = typeof created.job_id === 'string' ? created.job_id : undefined
if (!jobId) throw new Error('Website crawl did not return a job id')