mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-09-23 01:46:37 +08:00
fix(rag): retain interrupted ingestion state and engine identity
This commit is contained in:
RockChinQ
@@ -0,0 +1,406 @@
|
||||
"""B1 lifecycle regressions; isolated databases, no external engines or customer data."""
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.integration.persistence import test_rag_document_identity as identity
|
||||
from tests.integration.persistence.test_rag_document_identity import CONTEXT, read_file_row, seed
|
||||
from langbot.pkg.core.entities import LifecycleControlScope
|
||||
from langbot.pkg.core.taskmgr import AsyncTaskManager, TaskContext
|
||||
from langbot.pkg.rag.knowledge.kbmgr import RuntimeKnowledgeBase
|
||||
|
||||
# Reuse the real isolated-database fixtures without duplicating their lifecycle.
|
||||
database = identity.database
|
||||
runtime = identity.runtime
|
||||
|
||||
|
||||
def attach_task_manager(runtime):
|
||||
runtime.ap.event_loop = asyncio.get_running_loop()
|
||||
runtime.ap.instance_config = SimpleNamespace(data={})
|
||||
runtime.ap.task_mgr = AsyncTaskManager(runtime.ap)
|
||||
return runtime.ap.task_mgr
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_ingestion_retains_recovery_resources(runtime):
|
||||
"""Local cancellation is interrupted observation, not remote quiescence."""
|
||||
file = await seed(runtime)
|
||||
entered, terminal = asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def ingest(*_):
|
||||
entered.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
terminal.set()
|
||||
|
||||
runtime.ap.plugin_connector.call_rag_ingest.side_effect = ingest
|
||||
manager = attach_task_manager(runtime)
|
||||
wrapper = manager.create_user_task(
|
||||
runtime._store_file_task(CONTEXT, file, TaskContext.new()),
|
||||
kind='knowledge-operation',
|
||||
name=f'knowledge-store-file-{file.file_name}',
|
||||
instance_uuid=CONTEXT.instance_uuid,
|
||||
workspace_uuid=CONTEXT.workspace_uuid,
|
||||
placement_generation=CONTEXT.placement_generation,
|
||||
)
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 5)
|
||||
manager.cancel_by_scope(LifecycleControlScope.APPLICATION)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await wrapper.task
|
||||
assert terminal.is_set()
|
||||
assert wrapper.task.cancelled()
|
||||
assert (await read_file_row(runtime))['status'] == 'interrupted'
|
||||
runtime.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
|
||||
with pytest.raises(RuntimeError, match='interrupted'):
|
||||
await runtime.delete_file(CONTEXT, file.uuid)
|
||||
runtime.ap.plugin_connector.call_rag_delete_document.assert_not_awaited()
|
||||
assert await read_file_row(runtime) is not None
|
||||
finally:
|
||||
if not wrapper.task.done():
|
||||
wrapper.task.cancel()
|
||||
await asyncio.gather(wrapper.task, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('status', ['pending', 'processing'])
|
||||
async def test_recreated_host_marks_abandoned_rows_interrupted(runtime, status):
|
||||
"""Fresh Host state recovers observation without authorizing remote deletion."""
|
||||
await seed(runtime, status=status)
|
||||
new_ap = SimpleNamespace(**vars(runtime.ap))
|
||||
recreated = RuntimeKnowledgeBase(new_ap, runtime.knowledge_base_entity, CONTEXT)
|
||||
manager = attach_task_manager(recreated)
|
||||
assert manager.get_all_tasks() == []
|
||||
await recreated.initialize()
|
||||
assert (await read_file_row(recreated))['status'] == 'interrupted'
|
||||
with pytest.raises(RuntimeError, match='interrupted'):
|
||||
await recreated.delete_file(CONTEXT, 'host-id')
|
||||
recreated.ap.plugin_connector.call_rag_delete_document.assert_not_awaited()
|
||||
recreated.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_b1_live_task_still_blocks_delete_after_runtime_object_recreation(runtime):
|
||||
file = await seed(runtime)
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def ingest(*_):
|
||||
entered.set()
|
||||
await release.wait()
|
||||
return {'document_id': file.uuid, 'status': 'completed'}
|
||||
|
||||
runtime.ap.plugin_connector.call_rag_ingest.side_effect = ingest
|
||||
manager = attach_task_manager(runtime)
|
||||
wrapper = manager.create_user_task(
|
||||
runtime._store_file_task(CONTEXT, file, Mock()),
|
||||
kind='knowledge-operation',
|
||||
name=f'knowledge-store-file-{file.file_name}',
|
||||
instance_uuid=CONTEXT.instance_uuid,
|
||||
workspace_uuid=CONTEXT.workspace_uuid,
|
||||
placement_generation=CONTEXT.placement_generation,
|
||||
)
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 5)
|
||||
recreated = RuntimeKnowledgeBase(runtime.ap, runtime.knowledge_base_entity, CONTEXT)
|
||||
await recreated.initialize()
|
||||
assert not wrapper.task.done()
|
||||
with pytest.raises(RuntimeError, match='ingest'):
|
||||
await recreated.delete_file(CONTEXT, file.uuid)
|
||||
recreated.ap.plugin_connector.call_rag_delete_document.assert_not_awaited()
|
||||
assert (await read_file_row(recreated))['status'] == 'processing'
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(wrapper.task, 5)
|
||||
await recreated.delete_file(CONTEXT, file.uuid)
|
||||
assert await read_file_row(recreated) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('failure', [TimeoutError('timeout'), ConnectionError('disconnect')])
|
||||
async def test_transport_failure_is_interrupted_not_failed(runtime, failure):
|
||||
file = await seed(runtime)
|
||||
runtime.ap.plugin_connector.call_rag_ingest.side_effect = failure
|
||||
with pytest.raises(type(failure)):
|
||||
await runtime._store_file_task(CONTEXT, file, Mock())
|
||||
assert (await read_file_row(runtime))['status'] == 'interrupted'
|
||||
runtime.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recovered_row_cannot_be_replayed_by_delayed_old_task(runtime):
|
||||
file = await seed(runtime)
|
||||
await runtime.initialize()
|
||||
with pytest.raises(Exception):
|
||||
await runtime._store_file_task(CONTEXT, file, Mock())
|
||||
assert (await read_file_row(runtime))['status'] == 'interrupted'
|
||||
runtime.ap.plugin_connector.call_rag_ingest.assert_not_awaited()
|
||||
runtime.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconciliation_preserves_other_scopes_and_terminal_states(runtime):
|
||||
await seed(runtime, status='processing')
|
||||
await seed(runtime, file_id='other-kb', kb='kb-other', status='processing')
|
||||
await seed(runtime, file_id='other-workspace', workspace='workspace-b', kb='kb-b', status='processing')
|
||||
for status in ('completed', 'failed', 'interrupted'):
|
||||
await seed(runtime, file_id=status, status=status)
|
||||
await runtime.initialize()
|
||||
assert (await read_file_row(runtime))['status'] == 'interrupted'
|
||||
for file_id in ('other-kb', 'other-workspace'):
|
||||
assert (await read_file_row(runtime, file_id))['status'] == 'processing'
|
||||
for status in ('completed', 'failed', 'interrupted'):
|
||||
assert (await read_file_row(runtime, status))['status'] == status
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancellation_after_result_retains_identity(runtime):
|
||||
file = await seed(runtime)
|
||||
original = runtime._set_file_status
|
||||
|
||||
async def cancel_completion(context, uuid, status, **kwargs):
|
||||
if status == 'completed':
|
||||
raise asyncio.CancelledError()
|
||||
return await original(context, uuid, status, **kwargs)
|
||||
|
||||
runtime._set_file_status = cancel_completion
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await runtime._store_file_task(CONTEXT, file, Mock())
|
||||
row = await read_file_row(runtime)
|
||||
assert row['status'] == 'interrupted'
|
||||
assert row['engine_document_id'] == 'upstream-id'
|
||||
runtime.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancellation_after_completion_commit_does_not_downgrade(runtime):
|
||||
file = await seed(runtime)
|
||||
original = runtime._set_file_status
|
||||
|
||||
async def cancel_after_commit(context, uuid, status, **kwargs):
|
||||
result = await original(context, uuid, status, **kwargs)
|
||||
if status == 'completed':
|
||||
raise asyncio.CancelledError()
|
||||
return result
|
||||
|
||||
runtime._set_file_status = cancel_after_commit
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await runtime._store_file_task(CONTEXT, file, Mock())
|
||||
row = await read_file_row(runtime)
|
||||
assert row['status'] == 'completed'
|
||||
assert row['engine_document_id'] == 'upstream-id'
|
||||
runtime.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
|
||||
|
||||
|
||||
def make_service(runtime):
|
||||
from langbot.pkg.api.http.service.knowledge import KnowledgeService
|
||||
|
||||
runtime.ap.rag_mgr = SimpleNamespace(
|
||||
get_knowledge_base_details=AsyncMock(return_value={'uuid': 'kb-a'}),
|
||||
get_knowledge_base_by_uuid=AsyncMock(return_value=runtime),
|
||||
delete_knowledge_base=AsyncMock(),
|
||||
)
|
||||
return KnowledgeService(runtime.ap)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_list_recovers_abandoned_task_without_restart(runtime):
|
||||
await seed(runtime, status='processing')
|
||||
files = await make_service(runtime).get_files_by_knowledge_base(CONTEXT, 'kb-a')
|
||||
assert files[0]['status'] == 'interrupted'
|
||||
assert (await read_file_row(runtime))['status'] == 'interrupted'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('status', ['pending', 'processing', 'interrupted'])
|
||||
async def test_bulk_kb_delete_cannot_discard_unsettled_ingestion(runtime, status):
|
||||
from langbot.pkg.entity.persistence.rag import Chunk
|
||||
|
||||
async with runtime.ap.persistence_mgr.get_db_engine().begin() as conn:
|
||||
await conn.run_sync(lambda sync: Chunk.__table__.create(sync))
|
||||
await seed(runtime, status=status)
|
||||
service = make_service(runtime)
|
||||
with pytest.raises(RuntimeError, match='ingestion'):
|
||||
await service.delete_knowledge_base(CONTEXT, 'kb-a')
|
||||
assert await read_file_row(runtime) is not None
|
||||
runtime.ap.rag_mgr.delete_knowledge_base.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queued_task_cancelled_before_start_recovers_on_file_list(runtime):
|
||||
runtime.ap.storage_mgr.exists_scoped_object_key = AsyncMock(return_value=True)
|
||||
manager = attach_task_manager(runtime)
|
||||
await runtime.store_file(CONTEXT, 'upload.txt')
|
||||
wrapper = manager.get_all_tasks()[0]
|
||||
wrapper.task.cancel()
|
||||
await asyncio.gather(wrapper.task, return_exceptions=True)
|
||||
assert wrapper.task.cancelled()
|
||||
files = await make_service(runtime).get_files_by_knowledge_base(CONTEXT, 'kb-a')
|
||||
assert files[0]['status'] == 'interrupted'
|
||||
runtime.ap.plugin_connector.call_rag_ingest.assert_not_awaited()
|
||||
runtime.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('disconnect', [False, True], ids=['cancel-waiter', 'close-host-connection'])
|
||||
async def test_actual_sdk_late_write_preserves_interrupted_state(runtime, tmp_path, monkeypatch, disconnect):
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot_plugin.entities.io.actions.enums import RuntimeToPluginAction
|
||||
from langbot_plugin.runtime.io.handler import ActionResponse
|
||||
from tests.integration.plugin.test_rag_file_transfer_protocol import BINDING, protocol_stack
|
||||
|
||||
context = ExecutionContext(instance_uuid='instance-a', workspace_uuid='workspace-a', placement_generation=7)
|
||||
runtime.execution_context = context
|
||||
file = await seed(runtime)
|
||||
async with protocol_stack(tmp_path, monkeypatch, 'shared', BINDING) as stack:
|
||||
entered, release, finished = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
selected = SimpleNamespace(_runtime_plugin_handler=stack.bridge)
|
||||
monkeypatch.setattr(stack.runtime.plugin_mgr, '_get_connected_rag_plugin', lambda *_: (selected, 'engine'))
|
||||
vectors = set()
|
||||
|
||||
@stack.plugin.action(RuntimeToPluginAction.INGEST_DOCUMENT)
|
||||
async def ingest(data):
|
||||
entered.set()
|
||||
await release.wait()
|
||||
vectors.add('opaque-upstream-id')
|
||||
finished.set()
|
||||
return ActionResponse.success({'document_id': 'opaque-upstream-id', 'status': 'completed'})
|
||||
|
||||
async def send_ingest(_plugin_id, data):
|
||||
with stack.core.installation_scope(BINDING):
|
||||
return await stack.core.rag_ingest_document('tester', 'engine', data)
|
||||
|
||||
runtime.ap.plugin_connector.call_rag_ingest = send_ingest
|
||||
task = asyncio.create_task(runtime._store_file_task(context, file, Mock()))
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 5)
|
||||
if disconnect:
|
||||
await stack.core.close()
|
||||
else:
|
||||
task.cancel()
|
||||
outcome = await asyncio.gather(task, return_exceptions=True)
|
||||
assert isinstance(outcome[0], BaseException)
|
||||
assert (await read_file_row(runtime))['status'] == 'interrupted'
|
||||
assert not finished.is_set()
|
||||
with pytest.raises(RuntimeError, match='interrupted'):
|
||||
await runtime.delete_file(context, file.uuid)
|
||||
runtime.ap.plugin_connector.call_rag_delete_document.assert_not_awaited()
|
||||
release.set()
|
||||
await asyncio.wait_for(finished.wait(), 5)
|
||||
assert vectors == {'opaque-upstream-id'}
|
||||
assert (await read_file_row(runtime))['status'] == 'interrupted'
|
||||
assert (await read_file_row(runtime))['engine_document_id'] is None
|
||||
runtime.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
|
||||
finally:
|
||||
release.set()
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_late_observed_identity_survives_recovery_without_false_completion(runtime):
|
||||
file = await seed(runtime)
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def ingest(*_):
|
||||
entered.set()
|
||||
await release.wait()
|
||||
return {'document_id': 'late-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)
|
||||
# Independent Host observation, not shared process state.
|
||||
new_ap = SimpleNamespace(**{k: v for k, v in vars(runtime.ap).items() if not k.startswith('_knowledge_')})
|
||||
recreated = RuntimeKnowledgeBase(new_ap, runtime.knowledge_base_entity, CONTEXT)
|
||||
await recreated.initialize()
|
||||
release.set()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
row = await read_file_row(runtime)
|
||||
assert row['status'] == 'interrupted'
|
||||
assert row['engine_document_id'] == 'late-id'
|
||||
runtime.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parser_transport_failure_retains_upload(runtime):
|
||||
file = await seed(runtime)
|
||||
runtime.ap.storage_mgr.load_scoped_object_key = AsyncMock(return_value=b'file')
|
||||
runtime.ap.plugin_connector.call_parser = AsyncMock(side_effect=TimeoutError())
|
||||
with pytest.raises(TimeoutError):
|
||||
await runtime._store_file_task(CONTEXT, file, Mock(), parser_plugin_id='author/parser')
|
||||
assert (await read_file_row(runtime))['status'] == 'interrupted'
|
||||
runtime.ap.storage_mgr.delete_scoped_object_key.assert_not_awaited()
|
||||
runtime.ap.plugin_connector.call_rag_ingest.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_kb_delete_serializes_with_new_upload_admission(runtime):
|
||||
from langbot.pkg.entity.persistence.rag import Chunk
|
||||
|
||||
async with runtime.ap.persistence_mgr.get_db_engine().begin() as conn:
|
||||
await conn.run_sync(lambda sync: Chunk.__table__.create(sync))
|
||||
runtime.ap.storage_mgr.exists_scoped_object_key = AsyncMock(return_value=True)
|
||||
attach_task_manager(runtime)
|
||||
service = make_service(runtime)
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def pause_delete(*_):
|
||||
entered.set()
|
||||
await release.wait()
|
||||
|
||||
runtime.ap.rag_mgr.delete_knowledge_base.side_effect = pause_delete
|
||||
deletion = asyncio.create_task(service.delete_knowledge_base(CONTEXT, 'kb-a'))
|
||||
upload = None
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 5)
|
||||
upload = asyncio.create_task(runtime.store_file(CONTEXT, 'upload.txt'))
|
||||
await asyncio.sleep(0.1)
|
||||
assert not upload.done(), 'upload admission must wait for the in-progress KB deletion'
|
||||
runtime.ap.plugin_connector.call_rag_ingest.assert_not_awaited()
|
||||
release.set()
|
||||
await deletion
|
||||
with pytest.raises(Exception):
|
||||
await upload # FK rejects admission after the KB was deleted.
|
||||
runtime.ap.plugin_connector.call_rag_ingest.assert_not_awaited()
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(deletion, *([upload] if upload else []), return_exceptions=True)
|
||||
await asyncio.gather(*(wrapper.task for wrapper in runtime.ap.task_mgr.get_all_tasks()), return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('provider_name', ['LocalStorageProvider', 'S3StorageProvider'])
|
||||
async def test_retention_cleanup_preserves_ingestion_recovery_upload(runtime, tmp_path, provider_name):
|
||||
from langbot.pkg.api.http.service.maintenance import MaintenanceService
|
||||
|
||||
await seed(runtime, status='interrupted')
|
||||
retained = tmp_path / 'retained.txt'
|
||||
expired = tmp_path / 'expired.txt'
|
||||
retained.write_text('recovery source')
|
||||
expired.write_text('unreferenced upload')
|
||||
provider = type(provider_name, (), {})()
|
||||
provider.delete = AsyncMock()
|
||||
runtime.ap.storage_mgr.storage_provider = provider
|
||||
candidates = [
|
||||
{'key': 'upload.txt', 'path': str(retained)},
|
||||
{'key': 'unreferenced.txt', 'path': str(expired)},
|
||||
]
|
||||
service = MaintenanceService(runtime.ap)
|
||||
service._expired_local_upload_candidates = Mock(return_value=candidates)
|
||||
service._expired_s3_upload_candidates = AsyncMock(return_value=candidates)
|
||||
count = await service._cleanup_expired_uploaded_files(CONTEXT, 7)
|
||||
assert count == 1
|
||||
if provider_name == 'LocalStorageProvider':
|
||||
assert retained.read_text() == 'recovery source'
|
||||
assert not expired.exists()
|
||||
else:
|
||||
provider.delete.assert_awaited_once_with('unreferenced.txt')
|
||||
Reference in New Issue
Block a user