fix(rag): retain interrupted ingestion state and engine identity

2026-09-23 01:46:37 +08:00 · 2026-09-21 07:47:31 +00:00
parent 52f5699533
commit 381bb3f852
11 changed files with 1162 additions and 83 deletions
@@ -0,0 +1,386 @@
+"""Real database regressions for Host/engine identity; no live plugin or customer data."""
+
+import asyncio
+import os
+import uuid
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+import pytest_asyncio
+import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from langbot.pkg.api.http.context import ExecutionContext
+from langbot.pkg.entity.persistence.base import Base
+from langbot.pkg.entity.persistence.rag import File, KnowledgeBase
+from langbot.pkg.entity.persistence.user import User
+from langbot.pkg.entity.persistence.workspace import Workspace
+from langbot.pkg.persistence.alembic_runner import (
+    get_alembic_current,
+    run_alembic_downgrade,
+    run_alembic_stamp,
+    run_alembic_upgrade,
+)
+from langbot.pkg.persistence.mgr import PersistenceManager
+from langbot.pkg.rag.knowledge.kbmgr import RuntimeKnowledgeBase
+from langbot.pkg.workspace.errors import WorkspaceNotFoundError
+
+OLD_HEAD = '0024_passkey_credentials'
+NEW_HEAD = '0025_rag_document_identity'
+CONTEXT = ExecutionContext(instance_uuid='instance-a', workspace_uuid='workspace-a', placement_generation=5)
+
+
+@pytest_asyncio.fixture(params=['sqlite', 'postgres'])
+async def database(request, tmp_path):
+    admin = None
+    schema = 'ke_identity_' + uuid.uuid4().hex
+    if request.param == 'postgres':
+        url = os.environ.get('TEST_POSTGRES_URL')
+        if not url:
+            pytest.skip('TEST_POSTGRES_URL is required for disposable PostgreSQL tests')
+        admin = create_async_engine(url)
+        async with admin.begin() as conn:
+            await conn.execute(sa.text(f'CREATE SCHEMA {schema}'))
+        engine = create_async_engine(url, connect_args={'server_settings': {'search_path': schema}})
+    else:
+        engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "identity.db"}')
+
+        @sa.event.listens_for(engine.sync_engine, 'connect')
+        def enable_foreign_keys(connection, _):
+            connection.execute('PRAGMA foreign_keys=ON')
+
+    try:
+        yield engine
+    finally:
+        await engine.dispose()
+        if admin is not None:
+            async with admin.begin() as conn:
+                await conn.execute(sa.text(f'DROP SCHEMA {schema} CASCADE'))
+            await admin.dispose()
+
+
+async def create_schema(engine, *, legacy=False):
+    # Only the fixture-owned dependency closure; never all imported application tables.
+    async with engine.begin() as conn:
+        await conn.run_sync(
+            lambda sync: Base.metadata.create_all(
+                sync, tables=[User.__table__, Workspace.__table__, KnowledgeBase.__table__]
+            )
+        )
+        if legacy:
+            # Exact pre-0025 File shape, NOT current metadata stamped with an old revision.
+            await conn.execute(
+                sa.text("""CREATE TABLE knowledge_base_files (
+                uuid VARCHAR(255) PRIMARY KEY UNIQUE,
+                workspace_uuid VARCHAR(36) NOT NULL REFERENCES workspaces(uuid) ON DELETE CASCADE,
+                kb_id VARCHAR(255), file_name VARCHAR, extension VARCHAR, created_at TIMESTAMP, status VARCHAR,
+                CONSTRAINT uq_knowledge_base_files_workspace_uuid UNIQUE (workspace_uuid, uuid),
+                CONSTRAINT fk_knowledge_base_files_workspace_kb FOREIGN KEY (workspace_uuid, kb_id)
+                    REFERENCES knowledge_bases(workspace_uuid, uuid) ON DELETE CASCADE
+            )""")
+            )
+            await conn.execute(
+                sa.text(
+                    'CREATE INDEX ix_knowledge_base_files_workspace_kb ON knowledge_base_files (workspace_uuid, kb_id)'
+                )
+            )
+        else:
+            await conn.run_sync(lambda sync: File.__table__.create(sync))
+        for workspace in ('workspace-a', 'workspace-b'):
+            await conn.execute(
+                sa.insert(Workspace).values(
+                    uuid=workspace,
+                    instance_uuid='instance-a',
+                    name=workspace,
+                    slug=workspace,
+                    source='cloud_projection',
+                )
+            )
+        for kb, workspace in [('kb-a', 'workspace-a'), ('kb-other', 'workspace-a'), ('kb-b', 'workspace-b')]:
+            await conn.execute(
+                sa.insert(KnowledgeBase).values(
+                    uuid=kb,
+                    workspace_uuid=workspace,
+                    name=kb,
+                    knowledge_engine_plugin_id='author/engine',
+                    collection_id=kb,
+                )
+            )
+
+
+@pytest_asyncio.fixture
+async def runtime(database):
+    await create_schema(database)
+    ap = SimpleNamespace(
+        logger=Mock(),
+        workspace_service=SimpleNamespace(
+            get_execution_binding=AsyncMock(
+                return_value=SimpleNamespace(
+                    instance_uuid='instance-a',
+                    placement_generation=5,
+                )
+            )
+        ),
+        storage_mgr=SimpleNamespace(
+            require_scoped_object_key=Mock(),
+            size_scoped_object_key=AsyncMock(return_value=12),
+            delete_scoped_object_key=AsyncMock(),
+        ),
+        plugin_connector=SimpleNamespace(
+            require_workspace_context=AsyncMock(side_effect=lambda context: context),
+            call_rag_ingest=AsyncMock(return_value={'document_id': 'upstream-id', 'status': 'processing'}),
+            call_rag_delete_document=AsyncMock(return_value=True),
+        ),
+    )
+    ap.persistence_mgr = PersistenceManager(ap)
+    ap.persistence_mgr.db = SimpleNamespace(get_engine=lambda: database)
+    kb = KnowledgeBase(uuid='kb-a', workspace_uuid='workspace-a', name='kb', knowledge_engine_plugin_id='author/engine')
+    return RuntimeKnowledgeBase(ap, kb, CONTEXT)
+
+
+async def seed(runtime, *, status='pending', file_id='host-id', workspace='workspace-a', kb='kb-a'):
+    values = dict(
+        uuid=file_id, workspace_uuid=workspace, kb_id=kb, file_name='upload.txt', extension='txt', status=status
+    )
+    await runtime.ap.persistence_mgr.execute_async(sa.insert(File).values(**values))
+    return File(**values)
+
+
+async def read_file_row(runtime, file_id='host-id'):
+    # A fresh connection proves committed state rather than an identity-map/mock result.
+    async with runtime.ap.persistence_mgr.get_db_engine().connect() as conn:
+        row = (await conn.execute(sa.select(File).where(File.uuid == file_id))).first()
+        return None if row is None else dict(row._mapping)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'engine_id', ['dify-upstream', 'fastgpt-collection', 'ragflow-upstream', 'host-id', ' opaque ID ']
+)
+async def test_ingest_persists_exact_engine_identity_and_restart_delete(runtime, engine_id):
+    file = await seed(runtime)
+    runtime.ap.plugin_connector.call_rag_ingest.return_value['document_id'] = engine_id
+    await runtime._store_file_task(CONTEXT, file, Mock())
+    row = await read_file_row(runtime)
+    assert row['uuid'] == 'host-id'
+    assert row.get('engine_document_id') == engine_id
+    assert row['status'] == 'completed'
+    restarted = RuntimeKnowledgeBase(runtime.ap, runtime.knowledge_base_entity, CONTEXT)
+    await restarted.delete_file(CONTEXT, 'host-id')
+    runtime.ap.plugin_connector.call_rag_delete_document.assert_awaited_once_with('author/engine', engine_id, 'kb-a')
+    assert await read_file_row(runtime) is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('response', [False, None, 0, 1, 'true', {}])
+async def test_delete_without_explicit_confirmation_preserves_row(runtime, response):
+    await seed(runtime, status='completed')
+    runtime.ap.plugin_connector.call_rag_delete_document.return_value = response
+    with pytest.raises(RuntimeError, match='delet'):
+        await runtime.delete_file(CONTEXT, 'host-id')
+    assert await read_file_row(runtime) is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('failure', ['exception', 'missing_plugin'])
+async def test_delete_unavailable_retains_row(runtime, failure):
+    await seed(runtime, status='completed')
+    if failure == 'exception':
+        runtime.ap.plugin_connector.call_rag_delete_document.side_effect = RuntimeError('upstream offline')
+    else:
+        runtime.knowledge_base_entity.knowledge_engine_plugin_id = None
+    with pytest.raises(RuntimeError, match='delet'):
+        await runtime.delete_file(CONTEXT, 'host-id')
+    assert await read_file_row(runtime) is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('status', ['pending', 'processing'])
+async def test_delete_rejects_inflight_ingestion(runtime, status):
+    await seed(runtime, status=status)
+    with pytest.raises(RuntimeError, match='ingest'):
+        await runtime.delete_file(CONTEXT, 'host-id')
+    runtime.ap.plugin_connector.call_rag_delete_document.assert_not_awaited()
+    assert await read_file_row(runtime) is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('document_id', [None, '', '   ', 123, [], {}])
+async def test_invalid_response_identity_cannot_complete(runtime, document_id):
+    file = await seed(runtime)
+    runtime.ap.plugin_connector.call_rag_ingest.return_value = {'status': 'completed', 'document_id': document_id}
+    with pytest.raises(ValueError, match='document_id'):
+        await runtime._store_file_task(CONTEXT, file, Mock())
+    assert (await read_file_row(runtime))['status'] == 'interrupted'
+
+
+@pytest.mark.asyncio
+async def test_missing_response_identity_cannot_complete(runtime):
+    file = await seed(runtime)
+    runtime.ap.plugin_connector.call_rag_ingest.return_value = {'status': 'completed'}
+    with pytest.raises(ValueError, match='document_id'):
+        await runtime._store_file_task(CONTEXT, file, Mock())
+    assert (await read_file_row(runtime))['status'] == 'interrupted'
+
+
+@pytest.mark.asyncio
+async def test_failed_ingestion_retains_returned_identity_for_cleanup(runtime):
+    file = await seed(runtime)
+    runtime.ap.plugin_connector.call_rag_ingest.return_value = {
+        'document_id': 'uploaded-before-parsing-failed',
+        'status': 'failed',
+        'error_message': 'parsing failed',
+    }
+    with pytest.raises(Exception, match='parsing failed'):
+        await runtime._store_file_task(CONTEXT, file, Mock())
+    row = await read_file_row(runtime)
+    assert row['status'] == 'failed'
+    assert row.get('engine_document_id') == 'uploaded-before-parsing-failed'
+    await runtime.delete_file(CONTEXT, 'host-id')
+    runtime.ap.plugin_connector.call_rag_delete_document.assert_awaited_once_with(
+        'author/engine', 'uploaded-before-parsing-failed', 'kb-a'
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('status', ['completed', 'failed'])
+async def test_legacy_rows_use_host_id_only_with_confirmed_delete(runtime, status):
+    await seed(runtime, status=status)
+    await runtime.delete_file(CONTEXT, 'host-id')
+    runtime.ap.plugin_connector.call_rag_delete_document.assert_awaited_once_with('author/engine', 'host-id', 'kb-a')
+    assert await read_file_row(runtime) is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('workspace,kb', [('workspace-a', 'kb-other'), ('workspace-b', 'kb-b')])
+async def test_status_update_and_delete_do_not_touch_other_scope(runtime, workspace, kb):
+    await seed(runtime, workspace=workspace, kb=kb)
+    assert not await runtime._set_file_status(CONTEXT, 'host-id', 'processing')
+    with pytest.raises(WorkspaceNotFoundError):
+        await runtime.delete_file(CONTEXT, 'host-id')
+    assert (await read_file_row(runtime))['status'] == 'pending'
+    runtime.ap.plugin_connector.call_rag_delete_document.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_delete_rechecks_generation_after_plugin_response(runtime):
+    await seed(runtime, status='completed')
+
+    async def delete(*_):
+        runtime.ap.workspace_service.get_execution_binding.side_effect = WorkspaceNotFoundError('stale placement')
+        return True
+
+    runtime.ap.plugin_connector.call_rag_delete_document.side_effect = delete
+    with pytest.raises(WorkspaceNotFoundError):
+        await runtime.delete_file(CONTEXT, 'host-id')
+    assert await read_file_row(runtime) is not None
+
+
+@pytest.mark.asyncio
+async def test_ingest_rechecks_generation_before_mapping_write(runtime):
+    file = await seed(runtime)
+
+    async def ingest(*_):
+        runtime.ap.workspace_service.get_execution_binding.side_effect = WorkspaceNotFoundError('stale placement')
+        return {'document_id': 'upstream-id', 'status': 'completed'}
+
+    runtime.ap.plugin_connector.call_rag_ingest.side_effect = ingest
+    with pytest.raises(WorkspaceNotFoundError):
+        await runtime._store_file_task(CONTEXT, file, Mock())
+    row = await read_file_row(runtime)
+    assert row['status'] == 'processing'
+    assert row.get('engine_document_id') is None
+    runtime.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_concurrent_delete_cannot_remove_ingestion_tracking(runtime):
+    file = await seed(runtime)
+    entered, release = asyncio.Event(), asyncio.Event()
+
+    async def ingest(*_):
+        entered.set()
+        await release.wait()
+        return {'document_id': 'upstream-id', 'status': 'completed'}
+
+    runtime.ap.plugin_connector.call_rag_ingest.side_effect = ingest
+    task = asyncio.create_task(runtime._store_file_task(CONTEXT, file, Mock()))
+    try:
+        await asyncio.wait_for(entered.wait(), 5)
+        with pytest.raises(RuntimeError, match='ingest'):
+            await runtime.delete_file(CONTEXT, 'host-id')
+    finally:
+        release.set()
+        await task
+    assert (await read_file_row(runtime)).get('engine_document_id') == 'upstream-id'
+
+
+@pytest.mark.asyncio
+async def test_identity_and_completion_are_one_atomic_write(runtime):
+    file = await seed(runtime)
+    engine = runtime.ap.persistence_mgr.get_db_engine()
+    attempts = []
+
+    def reject_completion(_conn, _cursor, statement, parameters, _context, _many):
+        if statement.startswith('UPDATE knowledge_base_files') and 'completed' in parameters:
+            attempts.append(statement)
+            raise RuntimeError('fixture mapping write failure')
+
+    sa.event.listen(engine.sync_engine, 'before_cursor_execute', reject_completion)
+    try:
+        with pytest.raises(RuntimeError, match='mapping write failure'):
+            await runtime._store_file_task(CONTEXT, file, Mock())
+    finally:
+        sa.event.remove(engine.sync_engine, 'before_cursor_execute', reject_completion)
+    assert len(attempts) == 1
+    assert 'engine_document_id=' in attempts[0]
+    row = await read_file_row(runtime)
+    assert row['status'] != 'completed'
+    assert row.get('engine_document_id') == 'upstream-id'
+
+
+@pytest.mark.asyncio
+async def test_populated_legacy_migration_roundtrip(database):
+    await create_schema(database, legacy=True)
+    async with database.begin() as conn:
+        await conn.execute(
+            sa.text("""INSERT INTO knowledge_base_files
+            (uuid, workspace_uuid, kb_id, file_name, extension, status)
+            VALUES ('legacy', 'workspace-a', 'kb-a', 'original.txt', 'txt', 'completed')""")
+        )
+        assert 'engine_document_id' not in await conn.run_sync(
+            lambda sync: {col['name'] for col in sa.inspect(sync).get_columns('knowledge_base_files')}
+        )
+    await run_alembic_stamp(database, OLD_HEAD)
+    await run_alembic_upgrade(database)
+    async with database.connect() as conn:
+        columns = await conn.run_sync(lambda sync: sa.inspect(sync).get_columns('knowledge_base_files'))
+        assert 'engine_document_id' in {col['name'] for col in columns}
+        assert next(col for col in columns if col['name'] == 'engine_document_id')['nullable']
+        row = (await conn.execute(sa.text('SELECT * FROM knowledge_base_files'))).mappings().one()
+        assert row['uuid'] == 'legacy' and row['status'] == 'completed'
+        assert row['engine_document_id'] is None
+    assert await get_alembic_current(database) == NEW_HEAD
+    await run_alembic_upgrade(database)
+    await run_alembic_stamp(database, OLD_HEAD)
+    await run_alembic_upgrade(database)
+    await run_alembic_downgrade(database, OLD_HEAD)
+    async with database.connect() as conn:
+        assert 'engine_document_id' not in await conn.run_sync(
+            lambda sync: {col['name'] for col in sa.inspect(sync).get_columns('knowledge_base_files')}
+        )
+        assert (await conn.execute(sa.text('SELECT uuid FROM knowledge_base_files'))).scalar_one() == 'legacy'
+    await run_alembic_upgrade(database)
+
+
+@pytest.mark.asyncio
+async def test_fresh_metadata_then_migration_is_idempotent(database):
+    await create_schema(database)
+    await run_alembic_stamp(database, OLD_HEAD)
+    await run_alembic_upgrade(database)
+    assert await get_alembic_current(database) == NEW_HEAD
+    async with database.connect() as conn:
+        assert 'engine_document_id' in await conn.run_sync(
+            lambda sync: {col['name'] for col in sa.inspect(sync).get_columns('knowledge_base_files')}
+        )