From 5e8b20999c34ce023f2e3a66dce020592faa08ce Mon Sep 17 00:00:00 2001 From: beemines <182060364+beemines@users.noreply.github.com> Date: Mon, 14 Sep 2026 08:12:45 +0800 Subject: [PATCH 1/2] fix: keep failed knowledge-base document deletions retryable --- astrbot/core/knowledge_base/kb_db_sqlite.py | 16 ++- tests/unit/test_kb_delete_retry.py | 125 ++++++++++++++++++++ 2 files changed, 137 insertions(+), 4 deletions(-) create mode 100644 tests/unit/test_kb_delete_retry.py diff --git a/astrbot/core/knowledge_base/kb_db_sqlite.py b/astrbot/core/knowledge_base/kb_db_sqlite.py index 081a6ab9b4..36efb44b68 100644 --- a/astrbot/core/knowledge_base/kb_db_sqlite.py +++ b/astrbot/core/knowledge_base/kb_db_sqlite.py @@ -320,7 +320,18 @@ async def get_documents_with_metadata_batch( return metadata_map async def delete_document_by_id(self, doc_id: str, vec_db: "FaissVecDB") -> None: - """删除单个文档及其相关数据（包括多媒体记录）""" + """Delete stored chunks before removing their document and media records. + + Keep the document listed if vector or chunk deletion fails so the user + can retry. The stores do not share a transaction; a later metadata + failure also leaves the document available for an idempotent retry. + + Args: + doc_id: ID of the document to delete. + vec_db: Vector store containing this document's chunks. + """ + await vec_db.delete_documents(metadata_filters={"kb_doc_id": doc_id}) + async with self.get_db() as session, session.begin(): # 删除多媒体记录 delete_media_stmt = delete(KBMedia).where(col(KBMedia.doc_id) == doc_id) @@ -330,9 +341,6 @@ async def delete_document_by_id(self, doc_id: str, vec_db: "FaissVecDB") -> None delete_stmt = delete(KBDocument).where(col(KBDocument.doc_id) == doc_id) await session.execute(delete_stmt) - # 在 vec db 中删除相关向量 - await vec_db.delete_documents(metadata_filters={"kb_doc_id": doc_id}) - # ===== 多媒体查询 ===== async def list_media_by_doc(self, doc_id: str) -> list[KBMedia]: diff --git a/tests/unit/test_kb_delete_retry.py b/tests/unit/test_kb_delete_retry.py new file mode 100644 index 0000000000..dc7c2b20fa --- /dev/null +++ b/tests/unit/test_kb_delete_retry.py @@ -0,0 +1,125 @@ +"""Document deletion remains discoverable and retryable after storage failures.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from astrbot.core.db.vec_db.faiss_impl.vec_db import FaissVecDB +from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase +from astrbot.core.knowledge_base.models import KBDocument, KBMedia, KnowledgeBase + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure_stage", ["index_save", "chunk_delete", "metadata_delete"] +) +async def test_failed_document_delete_remains_listed_and_can_be_retried( + tmp_path, monkeypatch, failure_stage +): + kb_db = KBSQLiteDatabase(str(tmp_path / "kb.db")) + provider = MagicMock() + provider.get_dim.return_value = 2 + provider.get_embedding = AsyncMock(return_value=[1.0, 0.0]) + vec_db = FaissVecDB( + str(tmp_path / "chunks.db"), str(tmp_path / "index.faiss"), provider + ) + await kb_db.initialize() + await kb_db.migrate_to_v1() + await vec_db.initialize() + try: + kb = KnowledgeBase(kb_name="Manuals", embedding_provider_id="test") + async with kb_db.get_db() as session, session.begin(): + session.add(kb) + await session.flush() + + documents = [] + for name in ("obsolete.txt", "current.txt"): + doc = KBDocument( + kb_id=kb.kb_id, + doc_name=name, + file_type="txt", + file_size=10, + file_path="", + chunk_count=1, + ) + async with kb_db.get_db() as session, session.begin(): + session.add(doc) + await session.flush() + session.add( + KBMedia( + kb_id=kb.kb_id, + doc_id=doc.doc_id, + media_type="image", + file_name=name + ".png", + file_path=str(tmp_path / (name + ".png")), + file_size=1, + mime_type="image/png", + ) + ) + await vec_db.insert( + name, + {"kb_id": kb.kb_id, "kb_doc_id": doc.doc_id, "chunk_index": 0}, + id=name, + ) + documents.append(doc) + target, other = documents + + # Fail at each storage boundary, including a metadata transaction that + # has deleted media rows but has not yet deleted the document row. + with monkeypatch.context() as patch: + if failure_stage == "index_save": + patch.setattr( + vec_db.embedding_storage, + "save_index", + AsyncMock(side_effect=OSError("injected storage failure")), + ) + elif failure_stage == "chunk_delete": + patch.setattr( + vec_db.document_storage, + "delete_documents", + AsyncMock(side_effect=OSError("injected storage failure")), + ) + else: + execute = AsyncSession.execute + + async def fail_document_delete(session, statement, *args, **kwargs): + if statement.is_delete and statement.table.name == "kb_documents": + raise OSError("injected storage failure") + return await execute(session, statement, *args, **kwargs) + + patch.setattr(AsyncSession, "execute", fail_document_delete) + with pytest.raises(OSError, match="injected storage failure"): + await kb_db.delete_document_by_id(target.doc_id, vec_db) + + # Read through a fresh database connection, as a refreshed document list would. + await kb_db.close() + kb_db = KBSQLiteDatabase(str(tmp_path / "kb.db")) + listed = await kb_db.list_documents_by_kb(kb.kb_id) + assert {doc.doc_id for doc in listed} == {target.doc_id, other.doc_id} + assert len(await kb_db.list_media_by_doc(target.doc_id)) == 1 + expected_chunks = 0 if failure_stage == "metadata_delete" else 1 + assert ( + await vec_db.count_documents({"kb_doc_id": target.doc_id}) + == expected_chunks + ) + + await kb_db.delete_document_by_id(target.doc_id, vec_db) + assert await kb_db.get_document_by_id(target.doc_id) is None + assert await kb_db.list_media_by_doc(target.doc_id) == [] + assert await vec_db.count_documents({"kb_doc_id": target.doc_id}) == 0 + assert await kb_db.get_document_by_id(other.doc_id) is not None + assert len(await kb_db.list_media_by_doc(other.doc_id)) == 1 + assert await vec_db.count_documents({"kb_doc_id": other.doc_id}) == 1 + + # A restart sees the completed deletion in both chunk and vector stores. + await vec_db.close() + vec_db = FaissVecDB( + str(tmp_path / "chunks.db"), str(tmp_path / "index.faiss"), provider + ) + await vec_db.initialize() + assert vec_db.embedding_storage.index.ntotal == 1 + assert await vec_db.count_documents({"kb_doc_id": target.doc_id}) == 0 + finally: + await vec_db.close() + await kb_db.close() From c3102950ffd4b919bf1a206e522a5dfbcfdb6dcc Mon Sep 17 00:00:00 2001 From: beemines <182060364+beemines@users.noreply.github.com> Date: Mon, 14 Sep 2026 09:19:26 +0800 Subject: [PATCH 2/2] test: synchronize watchdog assertions with actual stack capture --- tests/unit/test_event_loop_diagnostics.py | 40 +++++++++++++++-------- 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/tests/unit/test_event_loop_diagnostics.py b/tests/unit/test_event_loop_diagnostics.py index 4d3e9510d9..3f92247d8a 100644 --- a/tests/unit/test_event_loop_diagnostics.py +++ b/tests/unit/test_event_loop_diagnostics.py @@ -1,6 +1,5 @@ import asyncio import threading -import time import pytest @@ -52,11 +51,19 @@ async def test_event_loop_watchdog_stops_worker_thread(): @pytest.mark.asyncio -async def test_event_loop_watchdog_writes_rotating_log(tmp_path): +async def test_event_loop_watchdog_writes_rotating_log(tmp_path, monkeypatch): """The watchdog should write to and rotate its log file.""" log_path = tmp_path / "logs" / "event_loop_watchdog.log" log_path.parent.mkdir() log_path.write_text("x" * 8, encoding="utf-8") + sampled = threading.Event() + print_stack = diagnostics.traceback.print_stack + + def record_stack(*args, **kwargs): + print_stack(*args, **kwargs) + sampled.set() + + monkeypatch.setattr(diagnostics.traceback, "print_stack", record_stack) task = asyncio.create_task( diagnostics.event_loop_watchdog( @@ -66,15 +73,20 @@ async def test_event_loop_watchdog_writes_rotating_log(tmp_path): max_bytes=4, ) ) - await asyncio.sleep(0) - time.sleep(0.05) # noqa: ASYNC251 - Intentionally block the event loop. - await asyncio.sleep(0.02) - task.cancel() - await asyncio.gather(task, return_exceptions=True) + try: + await asyncio.sleep(0) + # Keep this coroutine on the event-loop stack until the real dump is + # sampled, instead of assuming the worker runs within a fixed 50ms. + assert sampled.wait(timeout=5), "Watchdog did not capture a thread stack" + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) log_content = log_path.read_text(encoding="utf-8") assert "Event loop stalled for" in log_content - assert "test_event_loop_diagnostics.py" in log_content + main_thread_dump = log_content.split("\nThread", 2)[1] + assert "test_event_loop_diagnostics.py" in main_thread_dump + assert "test_event_loop_watchdog_writes_rotating_log" in main_thread_dump assert ( log_path.with_name("event_loop_watchdog.log.1").read_text(encoding="utf-8") == "x" * 8 @@ -105,10 +117,12 @@ def flaky_open(path, max_bytes): dump_path=log_path, ) ) - await asyncio.sleep(0) - time.sleep(0.06) # noqa: ASYNC251 - Intentionally block the event loop. - assert dumped.is_set() - task.cancel() - await asyncio.gather(task, return_exceptions=True) + try: + await asyncio.sleep(0) + # The worker must retry while the event loop is still blocked. + assert dumped.wait(timeout=5), "Watchdog did not retry the failed dump" + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) assert attempts >= 2