From 381bb3f852d437bc422ffa907e3e7528655097cc Mon Sep 17 00:00:00 2001 From: RockChinQ <45992437+RockChinQ@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:47:31 +0000 Subject: [PATCH] fix(rag): retain interrupted ingestion state and engine identity --- .../references/langrag-knowledge-base.md | 13 + src/langbot/pkg/api/http/service/knowledge.py | 24 +- .../pkg/api/http/service/maintenance.py | 19 + src/langbot/pkg/entity/persistence/rag.py | 6 +- .../versions/0025_rag_document_identity.py | 28 ++ src/langbot/pkg/rag/knowledge/kbmgr.py | 270 +++++++++--- .../persistence/test_rag_document_identity.py | 386 +++++++++++++++++ .../test_rag_interrupted_ingestion.py | 406 ++++++++++++++++++ .../plugin/test_rag_cancellation_protocol.py | 64 +++ tests/unit_tests/rag/test_file_storage.py | 10 +- tests/unit_tests/rag/test_kbmgr.py | 19 +- 11 files changed, 1162 insertions(+), 83 deletions(-) create mode 100644 src/langbot/pkg/persistence/alembic/versions/0025_rag_document_identity.py create mode 100644 tests/integration/persistence/test_rag_document_identity.py create mode 100644 tests/integration/persistence/test_rag_interrupted_ingestion.py create mode 100644 tests/integration/plugin/test_rag_cancellation_protocol.py diff --git a/skills/skills/langbot-testing/references/langrag-knowledge-base.md b/skills/skills/langbot-testing/references/langrag-knowledge-base.md index 00e0fb816..1d78dd82a 100644 --- a/skills/skills/langbot-testing/references/langrag-knowledge-base.md +++ b/skills/skills/langbot-testing/references/langrag-knowledge-base.md @@ -60,6 +60,19 @@ fixtures/rag/sentinel-doc.txt - Retrieve Test returns the uploaded document with the sentinel text. - Browser console has no unexpected errors. +## Document Lifecycle Acceptance + +- Keep the public Host file UUID from upload/listing when calling file deletion. Core stores the engine-returned `document_id` separately in the server-owned `engine_document_id` column and sends it to the engine; do not overwrite the Host UUID or put this mapping in creation settings. +- Wait for the ingestion task to finish before deletion. Pending/processing files reject deletion to avoid losing the tracking row while an upstream document is still being created. +- Verify both the Host file-list readback and absence of the sentinel upstream. Core only removes its row after an explicit engine `True`; `False`, missing configuration, runtime errors, and unconfirmed absence remain failures with the row retained. A connector's `False` is not proof that the upstream document is absent. +- Failed ingestion with an acknowledged engine ID retains it for cleanup. Historical failed files with no mapping still use the Host UUID fallback, subject to confirmed deletion. New unacknowledged/malformed responses are `interrupted`, not confirmed failures. +- `interrupted` means Core cannot establish the remote ingestion outcome (cancellation, disconnect, timeout, lost acknowledgement, or abandoned pending/processing work). It is not proof of failure, successful ingestion, or remote quiescence. Core preserves the tracking row, any known engine ID, and its source upload; retention cleanup also protects pending/processing/interrupted uploads. +- Runtime loading and normal file listing recover abandoned rows to `interrupted`. Reloading a KB object must preserve genuinely live tasks. Delayed old tasks cannot replay interrupted rows or overwrite them with completion; an observed late engine ID is still retained. +- File deletion and whole-KB deletion reject interrupted work with operator guidance. Do not bypass this guard merely because the task list is empty or the Host restarted: the old SDK/plugin action may still write. There is deliberately no automatic retry, force-delete, or claim of a remote cancellation fence. +- Recovery procedure: preserve a DB/storage backup; identify the exact Workspace, KB, Host file UUID, engine ID (if known), and installation; inspect that plugin/upstream operation; establish that the old operation has stopped or settled; then have an operator reconcile the confirmed upstream outcome and exact mapping before cleanup or re-upload. Never invent an opaque upstream ID or blindly set `failed` to unlock deletion. Lost historical uploads/IDs cannot be reconstructed by this change. Automatic cleanup of arbitrary connector orphans is outside this recovery contract. +- SDK ingest/delete action signatures and envelopes are unchanged; old SDKs use the same conservative interrupted fallback. Existing acknowledged asynchronous-engine responses retain the legacy Host completion semantics (Host ingestion acknowledgement, not a guarantee that the upstream index is ready). +- Migration `0025_rag_document_identity` leaves historical mappings null: it cannot reconstruct IDs previously discarded, nor restore already-deleted Host rows. Those require separately authorized investigation/recovery. Downgrading removes the mapping column and loses these identities; back up before rollback. + ## Local-Agent RAG Check After retrieval passes: diff --git a/src/langbot/pkg/api/http/service/knowledge.py b/src/langbot/pkg/api/http/service/knowledge.py index 06ba047b0..572eaff6f 100644 --- a/src/langbot/pkg/api/http/service/knowledge.py +++ b/src/langbot/pkg/api/http/service/knowledge.py @@ -294,6 +294,11 @@ class KnowledgeService: workspace_uuid = require_workspace_uuid(context) if await self.get_knowledge_base(context, kb_uuid) is None: raise WorkspaceNotFoundError('Knowledge base not found') + if isinstance(context, (RequestContext, ExecutionContext)): + execution_context = self._execution_context(context) + runtime_kb = await self.ap.rag_mgr.get_knowledge_base_by_uuid(execution_context, kb_uuid) + if runtime_kb is not None: + await runtime_kb.reconcile_interrupted_ingestions(execution_context) result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(persistence_rag.File) .where(persistence_rag.File.workspace_uuid == workspace_uuid) @@ -332,11 +337,20 @@ class KnowledgeService: kb_uuid: str, ) -> None: """删除知识库""" - workspace_uuid = require_workspace_uuid(context) + require_workspace_uuid(context) if await self.get_knowledge_base(context, kb_uuid) is None: raise WorkspaceNotFoundError('Knowledge base not found') - # delete files + execution_context = self._execution_context(context) + runtime_kb = await self.ap.rag_mgr.get_knowledge_base_by_uuid(execution_context, kb_uuid) + if runtime_kb is not None: + async with runtime_kb.ingestion_admission_lock: + await self._delete_knowledge_base_records(context, kb_uuid) + else: + await self._delete_knowledge_base_records(context, kb_uuid) + + async def _delete_knowledge_base_records(self, context: RequestContext | ExecutionContext, kb_uuid: str) -> None: + workspace_uuid = require_workspace_uuid(context) # NOTE: Chunk cleanup is for legacy (pre-plugin) KBs that stored chunks locally. # For plugin-based Knowledge Engines, the Chunk table is not populated, so this is a no-op. files = await self.ap.persistence_mgr.execute_async( @@ -344,6 +358,12 @@ class KnowledgeService: .where(persistence_rag.File.workspace_uuid == workspace_uuid) .where(persistence_rag.File.kb_id == kb_uuid) ) + files = files.all() + if any(file.status in {'pending', 'processing', 'interrupted'} for file in files): + raise RuntimeError( + 'Knowledge base has active or interrupted ingestion; retain its files and reconcile ' + 'plugin/upstream state before deleting the knowledge base.' + ) for file in files: # delete chunks await self.ap.persistence_mgr.execute_async( diff --git a/src/langbot/pkg/api/http/service/maintenance.py b/src/langbot/pkg/api/http/service/maintenance.py index 3ba4e77c5..8a1e561de 100644 --- a/src/langbot/pkg/api/http/service/maintenance.py +++ b/src/langbot/pkg/api/http/service/maintenance.py @@ -13,6 +13,7 @@ import sqlalchemy from ....core import app from ....entity.persistence import bstorage as persistence_bstorage from ....entity.persistence import monitoring as persistence_monitoring +from ....entity.persistence import rag as persistence_rag from ..authz import WorkspaceRequiredError from ..context import ExecutionContext from .tenant import TenantContext, require_workspace_uuid @@ -223,6 +224,7 @@ class MaintenanceService: retention_days, True, ) + candidates = await self._exclude_ingestion_uploads(context, candidates) return await asyncio.to_thread( self._delete_local_candidates, candidates, @@ -233,6 +235,22 @@ class MaintenanceService: return 0 + async def _exclude_ingestion_uploads( + self, context: TenantContext, candidates: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + """Keep source material for active or unacknowledged remote ingestion.""" + protected = set() + for offset in range(0, len(candidates), 500): + keys = [item['key'] for item in candidates[offset : offset + 500]] + result = await self.ap.persistence_mgr.execute_async( + sqlalchemy.select(persistence_rag.File.file_name) + .where(persistence_rag.File.workspace_uuid == require_workspace_uuid(context)) + .where(persistence_rag.File.status.in_(['pending', 'processing', 'interrupted'])) + .where(persistence_rag.File.file_name.in_(keys)) + ) + protected.update(result.scalars().all()) + return [item for item in candidates if item['key'] not in protected] + async def _expired_uploaded_candidates( self, context: TenantContext, @@ -256,6 +274,7 @@ class MaintenanceService: ) -> int: provider = self.ap.storage_mgr.storage_provider candidates = await self._expired_s3_upload_candidates(context, retention_days) + candidates = await self._exclude_ingestion_uploads(context, candidates) deleted = 0 for item in candidates: await provider.delete(item['key']) diff --git a/src/langbot/pkg/entity/persistence/rag.py b/src/langbot/pkg/entity/persistence/rag.py index 8cd1592eb..27cc979de 100644 --- a/src/langbot/pkg/entity/persistence/rag.py +++ b/src/langbot/pkg/entity/persistence/rag.py @@ -80,7 +80,11 @@ class File(Base): file_name = sqlalchemy.Column(sqlalchemy.String) extension = sqlalchemy.Column(sqlalchemy.String) created_at = sqlalchemy.Column(sqlalchemy.DateTime, default=sqlalchemy.func.now()) - status = sqlalchemy.Column(sqlalchemy.String, default='pending') # pending, processing, completed, failed + status = sqlalchemy.Column( + sqlalchemy.String, default='pending' + ) # pending, processing, completed, failed, interrupted + # Server-owned engine identity; the public Host file UUID never changes. + engine_document_id = sqlalchemy.Column(sqlalchemy.Text, nullable=True) __table_args__ = ( sqlalchemy.UniqueConstraint('workspace_uuid', 'uuid', name='uq_knowledge_base_files_workspace_uuid'), diff --git a/src/langbot/pkg/persistence/alembic/versions/0025_rag_document_identity.py b/src/langbot/pkg/persistence/alembic/versions/0025_rag_document_identity.py new file mode 100644 index 000000000..29f1d7959 --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0025_rag_document_identity.py @@ -0,0 +1,28 @@ +"""Persist engine document identity separately from the public Host file UUID.""" + +from alembic import op +import sqlalchemy as sa + +revision = '0025_rag_document_identity' +down_revision = '0024_passkey_credentials' +branch_labels = None +depends_on = None + + +def upgrade() -> None: + inspector = sa.inspect(op.get_bind()) + if 'knowledge_base_files' not in inspector.get_table_names(): + return + columns = {column['name'] for column in inspector.get_columns('knowledge_base_files')} + if 'engine_document_id' not in columns: + # Legacy upstream IDs cannot be inferred from Host UUIDs or user config. + op.add_column('knowledge_base_files', sa.Column('engine_document_id', sa.Text(), nullable=True)) + + +def downgrade() -> None: + inspector = sa.inspect(op.get_bind()) + if 'knowledge_base_files' not in inspector.get_table_names(): + return + columns = {column['name'] for column in inspector.get_columns('knowledge_base_files')} + if 'engine_document_id' in columns: + op.drop_column('knowledge_base_files', 'engine_document_id') diff --git a/src/langbot/pkg/rag/knowledge/kbmgr.py b/src/langbot/pkg/rag/knowledge/kbmgr.py index 8042b3271..d34ca5a4c 100644 --- a/src/langbot/pkg/rag/knowledge/kbmgr.py +++ b/src/langbot/pkg/rag/knowledge/kbmgr.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio import io +import inspect import mimetypes import os.path import traceback @@ -43,9 +44,42 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): super().__init__(ap) self.knowledge_base_entity = knowledge_base_entity self.execution_context = execution_context + # Shared across KB object reloads, but deliberately not a remote-work fence. + self._ingestion_tasks = ap.__dict__.setdefault('_knowledge_ingestion_tasks', {}) + locks = ap.__dict__.setdefault('_knowledge_ingestion_locks', {}) + self.ingestion_admission_lock = locks.setdefault( + (execution_context.workspace_uuid, self.get_uuid()), asyncio.Lock() + ) async def initialize(self): - pass + await self.reconcile_interrupted_ingestions(self.execution_context) + + async def reconcile_interrupted_ingestions(self, execution_context: ExecutionContext) -> None: + """Persist loss of Host observation, never infer remote quiescence. + + A restarted Host cannot observe the old SDK request's outcome. Retain its + row, upload and identity for operator reconciliation; do not retry/delete. + The shared admission lock protects newly queued tasks during KB reloads. + """ + async with self.ingestion_admission_lock: + live_ids = [ + key[2] + for key, task in self._ingestion_tasks.items() + if key[:2] == (execution_context.workspace_uuid, self.get_uuid()) and not task.done() + ] + + async def reconcile(): + await self._assert_execution_context(execution_context) + await self.ap.persistence_mgr.execute_async( + sqlalchemy.update(persistence_rag.File) + .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) + .where(persistence_rag.File.kb_id == self.get_uuid()) + .where(persistence_rag.File.status.in_(['pending', 'processing'])) + .where(persistence_rag.File.uuid.not_in(live_ids)) + .values(status='interrupted') + ) + + await run_in_workspace_uow(self.ap, execution_context.workspace_uuid, reconcile) async def _assert_execution_context(self, execution_context: ExecutionContext) -> None: """Reject stale or cross-Workspace runtime access.""" @@ -99,23 +133,38 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): task_context: taskmgr.TaskContext, parser_plugin_id: str | None = None, ): - await run_in_workspace_uow( - self.ap, - execution_context.workspace_uuid, - lambda: self._assert_execution_context(execution_context), - ) - self._require_upload_object_key(execution_context, file.file_name) + key = (execution_context.workspace_uuid, self.get_uuid(), file.uuid) + current_task = asyncio.current_task() + existing = self._ingestion_tasks.get(key) + if existing is not None and existing is not current_task and not existing.done(): + raise RuntimeError('Knowledge file ingestion is already running') + self._ingestion_tasks[key] = current_task + engine_document_id = None + dispatched = False + confirmed_failure = False + cleanup_upload = False + status_visible = False try: - # set file status to processing + await run_in_workspace_uow( + self.ap, + execution_context.workspace_uuid, + lambda: self._assert_execution_context(execution_context), + ) + self._require_upload_object_key(execution_context, file.file_name) + # Claim only pending work. Recovered/terminal rows cannot be replayed. status_visible = False for retry_delay in (0.0, 0.01, 0.05, 0.1): if retry_delay: await asyncio.sleep(retry_delay) - if await self._set_file_status(execution_context, file.uuid, 'processing'): + async with self.ingestion_admission_lock: + claimed = await self._set_file_status( + execution_context, file.uuid, 'processing', expected_statuses=('pending',) + ) + if claimed: status_visible = True break if not status_visible: - raise WorkspaceNotFoundError('Knowledge file was not committed before its background task started') + raise WorkspaceNotFoundError('Knowledge file is missing, already claimed, or interrupted') task_context.set_current_action('Processing file') @@ -146,9 +195,12 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): 'metadata': {}, } await self._require_plugin_runtime_context(execution_context) + dispatched = True parsed_content = await self.ap.plugin_connector.call_parser(parser_plugin_id, parse_context, file_bytes) + dispatched = False - # Call plugin to ingest document + # From dispatch until a valid response, failure is an unknown outcome. + dispatched = True result = await self._ingest_document( execution_context, { @@ -162,57 +214,109 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): parsed_content=parsed_content, ) - # Check plugin result status + # Failed ingestion can still have created an upstream document (for + # example, upload succeeded but parsing failed). Retain that identity + # for cleanup too. Never coerce or normalize an opaque engine ID. + returned_id = result.get('document_id') + if isinstance(returned_id, str) and returned_id.strip(): + engine_document_id = returned_id if result.get('status') == 'failed': + confirmed_failure = engine_document_id is not None error_msg = result.get('error_message', 'Plugin ingestion returned failed status') raise Exception(error_msg) + if engine_document_id is None: + raise ValueError('Plugin ingestion must return a nonempty string document_id') - # set file status to completed - if not await self._set_file_status(execution_context, file.uuid, 'completed'): - raise WorkspaceNotFoundError('Knowledge file not found') + # Commit the identity and completion together, never a status-only + # success that loses the only way to address the upstream document. + if not await self._set_file_status( + execution_context, + file.uuid, + 'completed', + engine_document_id=engine_document_id, + expected_statuses=('processing',), + ): + raise WorkspaceNotFoundError('Knowledge file not found or ingestion interrupted') + cleanup_upload = True - except Exception as e: - self.ap.logger.error(f'Error storing file {file.uuid}: {e}') - traceback.print_exc() - # A stale placement is fenced from all writes, including failure - # status updates from an old background task. + except (Exception, asyncio.CancelledError) as e: + cancelled = isinstance(e, asyncio.CancelledError) + status = 'interrupted' if cancelled or (dispatched and not confirmed_failure) else 'failed' + self.ap.logger.warning(f'Knowledge ingestion {status} for file {file.uuid}') + # A cancelled RPC waiter or transport error cannot establish remote + # failure. Preserve any returned ID and never downgrade a committed + # completion after an ambiguous commit acknowledgement. try: - if not await self._set_file_status(execution_context, file.uuid, 'failed'): - raise WorkspaceNotFoundError('Knowledge file not found') + if status_visible or cancelled: + changed = await self._set_file_status( + execution_context, + file.uuid, + status, + engine_document_id=engine_document_id, + expected_statuses=('pending', 'processing') if cancelled else ('processing',), + ) + cleanup_upload = changed and status == 'failed' except Exception: self.ap.logger.warning(f'Skipping stale RAG task status update for file {file.uuid}') - raise finally: - # An old background task must not touch an upload after its - # placement generation has been fenced off. - try: - await self._assert_execution_context(execution_context) - await self.ap.storage_mgr.delete_scoped_object_key( - execution_context, - file.file_name, - expected_owner_type='upload_document', - ) - except (WorkspaceRequiredError, WorkspaceNotFoundError): - self.ap.logger.warning(f'Skipping stale RAG upload cleanup for file {file.uuid}') + if self._ingestion_tasks.get(key) is current_task: + self._ingestion_tasks.pop(key, None) + # Only release recovery material after an acknowledged terminal write. + if cleanup_upload: + try: + await run_in_workspace_uow( + self.ap, + execution_context.workspace_uuid, + lambda: self._assert_execution_context(execution_context), + ) + await self.ap.storage_mgr.delete_scoped_object_key( + execution_context, + file.file_name, + expected_owner_type='upload_document', + ) + except (WorkspaceRequiredError, WorkspaceNotFoundError): + self.ap.logger.warning(f'Skipping stale RAG upload cleanup for file {file.uuid}') async def _set_file_status( self, execution_context: ExecutionContext, file_uuid: str, status: str, + *, + engine_document_id: str | None = None, + expected_statuses: tuple[str, ...] | None = None, ) -> bool: - """Commit one detached-task status transition in its own tenant UoW.""" + """Commit one detached-task status/identity transition in its tenant UoW.""" async def update() -> bool: await self._assert_execution_context(execution_context) - result = await self.ap.persistence_mgr.execute_async( + values = {'status': status} + if engine_document_id is not None: + values['engine_document_id'] = engine_document_id + statement = ( sqlalchemy.update(persistence_rag.File) .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) + .where(persistence_rag.File.kb_id == self.knowledge_base_entity.uuid) .where(persistence_rag.File.uuid == file_uuid) - .values(status=status) + .values(**values) ) - return getattr(result, 'rowcount', 0) > 0 + if expected_statuses is not None: + statement = statement.where(persistence_rag.File.status.in_(expected_statuses)) + result = await self.ap.persistence_mgr.execute_async(statement) + changed = getattr(result, 'rowcount', 0) > 0 + if not changed and engine_document_id is not None: + # A different Host may already have recovered observation. Save + # the late identity without claiming that the attempt completed. + await self.ap.persistence_mgr.execute_async( + sqlalchemy.update(persistence_rag.File) + .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) + .where(persistence_rag.File.kb_id == self.get_uuid()) + .where(persistence_rag.File.uuid == file_uuid) + .where(persistence_rag.File.status == 'interrupted') + .values(engine_document_id=engine_document_id) + ) + return changed persistence_mgr = self.ap.persistence_mgr managed_mode = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) in { @@ -264,33 +368,43 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): file_obj = persistence_rag.File(**file_obj_data) - await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_rag.File).values(file_obj_data)) + # Serialize admission with reconciliation, not with remote ingestion. + async with self.ingestion_admission_lock: + await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_rag.File).values(file_obj_data)) + ctx = taskmgr.TaskContext.new() + coroutine = self._store_file_task( + execution_context, file_obj, task_context=ctx, parser_plugin_id=parser_plugin_id + ) + try: + wrapper = self.ap.task_mgr.create_user_task( + coroutine, + kind='knowledge-operation', + name=f'knowledge-store-file-{file_id}', + label=f'Store file {file_id}', + context=ctx, + instance_uuid=execution_context.instance_uuid, + workspace_uuid=execution_context.workspace_uuid, + placement_generation=execution_context.placement_generation, + ) + except taskmgr.TaskCapacityError: + await self.ap.persistence_mgr.execute_async( + sqlalchemy.delete(persistence_rag.File) + .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) + .where(persistence_rag.File.uuid == file_uuid) + ) + raise + key = (execution_context.workspace_uuid, kb_id, file_uuid) + self._ingestion_tasks[key] = wrapper.task - # run background task asynchronously - ctx = taskmgr.TaskContext.new() - try: - wrapper = self.ap.task_mgr.create_user_task( - self._store_file_task( - execution_context, - file_obj, - task_context=ctx, - parser_plugin_id=parser_plugin_id, - ), - kind='knowledge-operation', - name=f'knowledge-store-file-{file_id}', - label=f'Store file {file_id}', - context=ctx, - instance_uuid=execution_context.instance_uuid, - workspace_uuid=execution_context.workspace_uuid, - placement_generation=execution_context.placement_generation, - ) - except taskmgr.TaskCapacityError: - await self.ap.persistence_mgr.execute_async( - sqlalchemy.delete(persistence_rag.File) - .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) - .where(persistence_rag.File.uuid == file_uuid) - ) - raise + def forget_task(task): + # The task manager wraps the coroutine; pre-start cancellation + # need not enter that wrapper's finally block. + if inspect.getcoroutinestate(coroutine) == inspect.CORO_CREATED: + coroutine.close() + if self._ingestion_tasks.get(key) is task: + self._ingestion_tasks.pop(key, None) + + wrapper.task.add_done_callback(forget_task) return wrapper.id async def _store_zip_file( @@ -445,17 +559,37 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): async def delete_file(self, execution_context: ExecutionContext, file_id: str): await self._assert_execution_context(execution_context) result = await self.ap.persistence_mgr.execute_async( - sqlalchemy.select(persistence_rag.File.uuid) + sqlalchemy.select(persistence_rag.File) .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) .where(persistence_rag.File.kb_id == self.knowledge_base_entity.uuid) .where(persistence_rag.File.uuid == file_id) .limit(1) ) - if result.first() is None: + file = result.first() + if file is None: raise WorkspaceNotFoundError('Knowledge file not found') - await self._delete_document(execution_context, file_id) + # Pending/processing tasks may still create an upstream document. Do not + # discard their tracking row or race a delete against that creation. + if file.status in {'pending', 'processing'}: + raise RuntimeError( + 'Cannot delete a file while ingestion is pending or processing; wait for the task to finish' + ) + if file.status == 'interrupted': + raise RuntimeError( + 'Knowledge ingestion was interrupted; its remote outcome is unknown. ' + 'The file, upload and known engine identity were retained. ' + 'Check plugin/upstream state and stop or settle the old ingestion before operator reconciliation; ' + 'automatic deletion or re-ingestion is unsafe.' + ) + document_id = file.engine_document_id if file.engine_document_id is not None else file_id + if await self._delete_document(execution_context, document_id) is not True: + raise RuntimeError( + 'Knowledge engine did not confirm document deletion; the file was retained. ' + 'Check the engine configuration and plugin logs before retrying.' + ) - # Also cleanup DB record + # The plugin call may outlive the original placement generation. + await self._assert_execution_context(execution_context) await self.ap.persistence_mgr.execute_async( sqlalchemy.delete(persistence_rag.File) .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) diff --git a/tests/integration/persistence/test_rag_document_identity.py b/tests/integration/persistence/test_rag_document_identity.py new file mode 100644 index 000000000..fe569c71d --- /dev/null +++ b/tests/integration/persistence/test_rag_document_identity.py @@ -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')} + ) diff --git a/tests/integration/persistence/test_rag_interrupted_ingestion.py b/tests/integration/persistence/test_rag_interrupted_ingestion.py new file mode 100644 index 000000000..aa6b53dc6 --- /dev/null +++ b/tests/integration/persistence/test_rag_interrupted_ingestion.py @@ -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') diff --git a/tests/integration/plugin/test_rag_cancellation_protocol.py b/tests/integration/plugin/test_rag_cancellation_protocol.py new file mode 100644 index 000000000..9ae3623a0 --- /dev/null +++ b/tests/integration/plugin/test_rag_cancellation_protocol.py @@ -0,0 +1,64 @@ +"""Exercise actual installed SDK RPC cancellation, not a mocked cancel contract.""" + +import asyncio +from types import SimpleNamespace + +import pytest + +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 + + +@pytest.mark.asyncio +@pytest.mark.parametrize('disconnect', [False, True], ids=['cancel-waiter', 'close-host-connection']) +async def test_b1_host_cancellation_does_not_fence_remote_ingest(tmp_path, monkeypatch, disconnect): + async with protocol_stack(tmp_path, monkeypatch, 'shared', BINDING) as stack: + entered, release, finished = asyncio.Event(), asyncio.Event(), asyncio.Event() + vectors = set() + selected = SimpleNamespace(_runtime_plugin_handler=stack.bridge) + monkeypatch.setattr(stack.runtime.plugin_mgr, '_get_connected_rag_plugin', lambda *_: (selected, 'engine')) + + @stack.plugin.action(RuntimeToPluginAction.INGEST_DOCUMENT) + async def ingest(data): + entered.set() + await release.wait() + vectors.add('host-id') + finished.set() + return ActionResponse.success({'document_id': 'host-id', 'status': 'completed'}) + + @stack.plugin.action(RuntimeToPluginAction.DELETE_DOCUMENT) + async def delete(data): + vectors.discard(data['document_id']) + return ActionResponse.success({'success': True}) + + async def send_ingest(): + with stack.core.installation_scope(BINDING): + return await stack.core.rag_ingest_document('tester', 'engine', {}) + + local = asyncio.create_task(send_ingest()) + try: + await asyncio.wait_for(entered.wait(), 5) + if disconnect: + await stack.core.close() + result = await asyncio.gather(local, return_exceptions=True) + assert isinstance(result[0], Exception) + else: + local.cancel() + with pytest.raises(asyncio.CancelledError): + await local + assert not finished.is_set() + # The SDK has no ingest/delete barrier: a new caller can receive a + # confirmed deletion while the old plugin action can still write. + # Use the surviving Runtime->plugin connection after Host disconnect. + result = await stack.bridge.rag_delete_document('kb-a', 'host-id') + assert result == {'success': True} + assert vectors == set() + release.set() + await asyncio.wait_for(finished.wait(), 5) + assert vectors == {'host-id'} + finally: + release.set() + if not local.done(): + local.cancel() + await asyncio.gather(local, return_exceptions=True) diff --git a/tests/unit_tests/rag/test_file_storage.py b/tests/unit_tests/rag/test_file_storage.py index 7b464d502..8886e221a 100644 --- a/tests/unit_tests/rag/test_file_storage.py +++ b/tests/unit_tests/rag/test_file_storage.py @@ -90,7 +90,7 @@ class TestStoreFile: def create_user_task(coro, **kwargs): coro.close() - return SimpleNamespace(id='task-1', kwargs=kwargs) + return SimpleNamespace(id='task-1', kwargs=kwargs, task=Mock()) kb.ap.task_mgr.create_user_task = Mock(side_effect=create_user_task) @@ -279,7 +279,7 @@ class TestStoreFileTask: kb._assert_execution_context = AsyncMock(side_effect=assert_execution_context) kb._set_file_status = AsyncMock(side_effect=[True, True]) - kb._ingest_document = AsyncMock(return_value={'status': 'completed'}) + kb._ingest_document = AsyncMock(return_value={'status': 'completed', 'document_id': 'file-uuid'}) object_key = _upload_key('scoped.pdf') file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf') @@ -290,7 +290,7 @@ class TestStoreFileTask: @pytest.mark.asyncio async def test_store_file_task_marks_completed_and_cleans_storage(self): kb = _make_kb() - kb._ingest_document = AsyncMock(return_value={'status': 'completed'}) + kb._ingest_document = AsyncMock(return_value={'status': 'completed', 'document_id': 'file-uuid'}) object_key = _upload_key('test.pdf') file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf') task_context = Mock() @@ -306,7 +306,9 @@ class TestStoreFileTask: @pytest.mark.asyncio async def test_store_file_task_marks_failed_and_cleans_storage(self): kb = _make_kb() - kb._ingest_document = AsyncMock(return_value={'status': 'failed', 'error_message': 'parser failed'}) + kb._ingest_document = AsyncMock( + return_value={'status': 'failed', 'error_message': 'parser failed', 'document_id': 'file-uuid'} + ) object_key = _upload_key('bad.pdf') file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf') task_context = Mock() diff --git a/tests/unit_tests/rag/test_kbmgr.py b/tests/unit_tests/rag/test_kbmgr.py index 9699269d9..2a1334c03 100644 --- a/tests/unit_tests/rag/test_kbmgr.py +++ b/tests/unit_tests/rag/test_kbmgr.py @@ -268,7 +268,9 @@ async def test_ingestion_payload_uses_host_owned_kb_collection(): async def test_delete_file_checks_workspace_and_parent_before_plugin_call(): app = _app() runtime = RuntimeKnowledgeBase(app, _entity(), CONTEXT_A) - app.persistence_mgr.execute_async.return_value = _Result(first=('file-a',)) + app.persistence_mgr.execute_async.return_value = _Result( + first=SimpleNamespace(uuid='file-a', status='completed', engine_document_id=None) + ) await runtime.delete_file(CONTEXT_A, 'file-a') app.plugin_connector.call_rag_delete_document.assert_awaited_once_with( @@ -384,7 +386,8 @@ class TestRAGManagerCreateKnowledgeBase: ) assert manager.knowledge_bases == {} - assert app.persistence_mgr.execute_async.await_count == 2 + # Insert, interrupted-ingestion reconciliation, rollback delete. + assert app.persistence_mgr.execute_async.await_count == 3 @pytest.mark.asyncio async def test_sets_default_retrieval_settings(self): @@ -456,7 +459,9 @@ class TestRuntimeKnowledgeBaseDeleteFile: @pytest.mark.asyncio async def test_delete_file_calls_plugin_and_db(self): app = _app() - app.persistence_mgr.execute_async.return_value = _Result(first=('file-uuid',)) + app.persistence_mgr.execute_async.return_value = _Result( + first=SimpleNamespace(uuid='file-uuid', status='completed', engine_document_id=None) + ) await RuntimeKnowledgeBase(app, _entity(), CONTEXT_A).delete_file( CONTEXT_A, @@ -514,7 +519,7 @@ class TestRAGManagerLoadKnowledgeBasesFromDB: } @pytest.mark.asyncio - async def test_cloud_startup_reuses_validated_binding(self): + async def test_cloud_startup_revalidates_binding_before_recovery_write(self): class TenantUow: async def __aenter__(self): return self @@ -534,15 +539,13 @@ class TestRAGManagerLoadKnowledgeBasesFromDB: app.persistence_mgr.tenant_uow = lambda _workspace_uuid: TenantUow() app.persistence_mgr.execute_async.return_value = _Result([_entity()]) app.workspace_service.list_active_execution_bindings = AsyncMock(return_value=[binding]) - app.workspace_service.get_execution_binding = AsyncMock( - side_effect=AssertionError('startup RAG loader repeated a validated binding lookup') - ) + app.workspace_service.get_execution_binding = AsyncMock(return_value=binding) manager = RAGManager(app) await manager.load_knowledge_bases_from_db() assert set(manager.knowledge_bases) == {('workspace-a', 'kb-a')} - app.workspace_service.get_execution_binding.assert_not_awaited() + app.workspace_service.get_execution_binding.assert_awaited_once_with('workspace-a', expected_generation=5) @pytest.mark.asyncio async def test_handles_load_error_gracefully(self):