fix(api): preserve integer knowledge ingestion task IDs

src/langbot/pkg/api/http/service/knowledge.py

@@ -303,7 +303,7 @@ class KnowledgeService:
         kb_uuid: str,
         file_id: str,
         parser_plugin_id: str | None = None,
-    ) -> str:
+    ) -> int:
         """存储文件"""
         execution_context = self._execution_context(context)
         runtime_kb = await self.ap.rag_mgr.get_knowledge_base_by_uuid(execution_context, kb_uuid)

src/langbot/pkg/rag/knowledge/kbmgr.py

@@ -336,7 +336,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
         execution_context: ExecutionContext,
         file_id: str,
         parser_plugin_id: str | None = None,
-    ) -> str:
+    ) -> int:
         await self._assert_execution_context(execution_context)
         self._require_upload_object_key(execution_context, file_id)
         # pre checking
@@ -412,7 +412,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
         execution_context: ExecutionContext,
         zip_file_id: str,
         parser_plugin_id: str | None = None,
-    ) -> str:
+    ) -> int:
         """Handle ZIP file by extracting each document and storing them separately."""
         await self._assert_execution_context(execution_context)
         self._require_upload_object_key(execution_context, zip_file_id)
@@ -516,7 +516,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
             self.ap.logger.info(
                 f'Successfully processed ZIP file {zip_file_id}, extracted {len(stored_file_tasks)} files'
             )
-            return stored_file_tasks[0] if stored_file_tasks else ''
+            return stored_file_tasks[0]
         finally:
             try:
                 await self._assert_execution_context(execution_context)

tests/integration/api/test_application_resources.py

@@ -255,3 +255,107 @@ async def test_chat_management_uses_scoped_history_and_real_connection_queues(re
         == []
     )
     assert len(adapter.websocket_group_session.message_lists[pipeline]) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('extension', ['txt', 'zip'])
+async def test_knowledge_upload_returns_real_integer_task_id(resources, tmp_path, monkeypatch, extension):
+    import asyncio
+    import io
+    import zipfile
+
+    from langbot.pkg.api.context import ExecutionContext
+    from langbot.pkg.core.taskmgr import AsyncTaskManager
+    from langbot.pkg.entity.persistence.rag import File
+    from langbot.pkg.entity.persistence.workspace import Workspace, WorkspaceExecutionState
+    from langbot.pkg.storage.mgr import StorageMgr
+    from langbot.pkg.storage.providers import localstorage
+
+    api, ap, engine = resources
+    workspace = '00000000-0000-0000-0000-00000000000a'
+    request = context(workspace)
+    execution = ExecutionContext.from_request(request)
+    kb_data = {
+        'uuid': 'ingestion-kb',
+        'workspace_uuid': workspace,
+        'name': 'Ingestion',
+        'knowledge_engine_plugin_id': 'test/engine',
+        'collection_id': 'ingestion-kb',
+        'creation_settings': {},
+        'retrieval_settings': {},
+        'initialized': True,
+    }
+    async with engine.begin() as connection:
+        await connection.execute(
+            sa.insert(Workspace).values(
+                uuid=workspace,
+                instance_uuid=request.instance_uuid,
+                name='Ingestion',
+                slug='ingestion',
+                source='cloud_projection',
+            )
+        )
+        await connection.execute(
+            sa.insert(WorkspaceExecutionState).values(
+                workspace_uuid=workspace, instance_uuid=request.instance_uuid, active_generation=1, source='cloud'
+            )
+        )
+        await connection.execute(sa.insert(KnowledgeBase).values(**kb_data))
+    ap.event_loop = asyncio.get_running_loop()
+    ap.task_mgr = AsyncTaskManager(ap)
+    monkeypatch.setattr(localstorage, 'LOCAL_STORAGE_PATH', str(tmp_path / 'uploads'))
+    ap.storage_mgr = StorageMgr(ap)
+    ap.storage_mgr.storage_provider = localstorage.LocalStorageProvider(ap)
+    entered, release = asyncio.Event(), asyncio.Event()
+
+    async def ingest(_plugin_id, _payload):
+        entered.set()
+        await release.wait()
+        return {'document_id': 'engine-document', 'status': 'completed'}
+
+    ap.plugin_connector = SimpleNamespace(
+        is_enable_plugin=True,
+        require_workspace_context=AsyncMock(side_effect=lambda scope: scope),
+        list_knowledge_engines=AsyncMock(
+            return_value=[{'plugin_id': 'test/engine', 'capabilities': ['doc_ingestion']}]
+        ),
+        call_rag_ingest=AsyncMock(side_effect=ingest),
+    )
+    await ap.rag_mgr.load_knowledge_base(execution, KnowledgeBase(**kb_data))
+    content = b'A document for ingestion.'
+    if extension == 'zip':
+        buffer = io.BytesIO()
+        with zipfile.ZipFile(buffer, 'w') as archive:
+            archive.writestr('document.txt', content)
+        content = buffer.getvalue()
+    upload = await ap.storage_mgr.save_scoped(
+        execution, owner_type='upload_document', owner='ingestion-test', key='document.' + extension, value=content
+    )
+    try:
+        task_id = await api.execute('knowledge.file.store', request, {'kb_uuid': kb_data['uuid'], 'file_id': upload})
+        assert type(task_id) is int
+        task = ap.task_mgr.get_task_by_id(
+            task_id, instance_uuid=request.instance_uuid, workspace_uuid=workspace, placement_generation=1
+        )
+        assert task is not None and task.kind == 'knowledge-operation'
+        await asyncio.wait_for(entered.wait(), 5)
+        async with engine.connect() as connection:
+            pending = (
+                (await connection.execute(sa.select(File).where(File.workspace_uuid == workspace))).mappings().one()
+            )
+        assert pending['kb_id'] == kb_data['uuid'] and pending['status'] == 'processing'
+        release.set()
+        await asyncio.wait_for(task.task, 5)
+        async with engine.connect() as connection:
+            stored = (await connection.execute(sa.select(File).where(File.uuid == pending['uuid']))).mappings().one()
+        assert stored['status'] == 'completed' and stored['engine_document_id'] == 'engine-document'
+        assert not await ap.storage_mgr.exists_scoped_object_key(
+            execution, upload, expected_owner_type='upload_document'
+        )
+    finally:
+        release.set()
+        tasks = [wrapper.task for wrapper in ap.task_mgr.get_all_tasks()]
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)

tests/integration/api/test_knowledge.py

@@ -123,7 +123,7 @@ def fake_knowledge_app():
     app.knowledge_service.get_files_by_knowledge_base = AsyncMock(
         return_value=[{'uuid': 'test-file-uuid', 'filename': 'test.pdf'}]
     )
-    app.knowledge_service.store_file = AsyncMock(return_value='test-task-id')
+    app.knowledge_service.store_file = AsyncMock(return_value=17)
     app.knowledge_service.delete_file = AsyncMock()
     app.knowledge_service.retrieve_knowledge_base = AsyncMock(return_value=[{'content': 'test result', 'score': 0.95}])
 
@@ -304,7 +304,7 @@ class TestKnowledgeBaseFilesEndpoints:
         assert called.await_args.args[0].workspace_uuid == '00000000-0000-0000-0000-00000000000a'
         assert called.await_args.args[1:] == ('test-kb-uuid', 'test-file-id')
         assert called.await_args.kwargs == {'parser_plugin_id': 'test/parser'}
-        assert data['data'] == {'task_id': 'test-task-id'}
+        assert data['data'] == {'task_id': 17}
 
     @pytest.mark.asyncio
     async def test_delete_file_from_knowledge_base(self, quart_test_client, fake_knowledge_app):

tests/unit_tests/api/service/test_knowledge_service.py

@@ -265,14 +265,14 @@ async def test_file_listing_checks_parent_knowledge_base_first():
 async def test_store_and_delete_file_require_runtime_parent_and_capability():
     app = _app()
     runtime_kb = SimpleNamespace(
-        store_file=AsyncMock(return_value='task-a'),
+        store_file=AsyncMock(return_value=17),
         delete_file=AsyncMock(),
     )
     app.rag_mgr.get_knowledge_base_by_uuid.return_value = runtime_kb
     app.rag_mgr.get_knowledge_base_details.return_value = {'knowledge_engine': {'capabilities': ['doc_ingestion']}}
     service = KnowledgeService(app)
 
-    assert await service.store_file(CONTEXT, 'kb-a', 'upload.pdf', 'author/parser') == 'task-a'
+    assert await service.store_file(CONTEXT, 'kb-a', 'upload.pdf', 'author/parser') == 17
     runtime_kb.store_file.assert_awaited_once_with(CONTEXT, 'upload.pdf', parser_plugin_id='author/parser')
     await service.delete_file(CONTEXT, 'kb-a', 'file-a')
     runtime_kb.delete_file.assert_awaited_once_with(CONTEXT, 'file-a')

tests/unit_tests/rag/test_file_storage.py

@@ -90,14 +90,14 @@ class TestStoreFile:
 
         def create_user_task(coro, **kwargs):
             coro.close()
-            return SimpleNamespace(id='task-1', kwargs=kwargs, task=Mock())
+            return SimpleNamespace(id=17, kwargs=kwargs, task=Mock())
 
         kb.ap.task_mgr.create_user_task = Mock(side_effect=create_user_task)
 
         object_key = _upload_key('documents/test.pdf')
         task_id = await kb.store_file(CONTEXT, object_key)
 
-        assert task_id == 'task-1'
+        assert task_id == 17
         kb.ap.storage_mgr.storage_provider.exists.assert_awaited_once_with(object_key)
         kb.ap.persistence_mgr.execute_async.assert_awaited_once()
         call_kwargs = kb.ap.task_mgr.create_user_task.call_args.kwargs
@@ -189,12 +189,12 @@ class TestStoreZipFile:
                 }
             )
         )
-        kb.store_file = AsyncMock(side_effect=['task-pdf', 'task-txt', 'task-md', 'task-html'])
+        kb.store_file = AsyncMock(side_effect=[17, 18, 19, 20])
 
         zip_key = _upload_key('archive.zip')
         task_id = await kb._store_zip_file(CONTEXT, zip_key, parser_plugin_id='parser/plugin')
 
-        assert task_id == 'task-pdf'
+        assert task_id == 17
         assert kb.ap.storage_mgr.storage_provider.save.await_count == 4
         saved_names = [call.args[0] for call in kb.ap.storage_mgr.storage_provider.save.await_args_list]
         assert {name.rsplit('.', 1)[-1] for name in saved_names} == {'pdf', 'txt', 'md', 'html'}