fix: compensate partial knowledge base upload failures (#9007) (#9246)

* fix: roll back document rows when FAISS insert fails

Validate embedding shape/dimension before local writes. If DocumentStorage
commits but FAISS write fails, best-effort delete the partial FAISS ids and
document rows so the two stores do not diverge.

* fix: clean up partial KB upload state before metadata commit

Track metadata_committed right after commit succeeds. On failure before
metadata commit, roll back chunks/vectors first, then residual KB rows and
media files. After metadata is committed, keep the document even if stats
or refresh fails.

* test: cover knowledge base upload atomicity and rollback

Add unit tests for insert_batch FAISS failure rollback, dimension validation
before document writes, upload cleanup order, no-rollback after metadata
commit, and real DocumentStorage+FAISS no-orphan cases.

* fix: log media directory cleanup failures after upload rollback

Best-effort rmdir of the per-document media directory should still emit a
warning when removal fails so incomplete cleanup is diagnosable.

astrbot/core/db/vec_db/faiss_impl/vec_db.py

@@ -135,26 +135,8 @@ class FaissVecDB(BaseVecDB):
                 },
             )
 
-        # 使用 DocumentStorage 的批量插入方法
-        int_ids = await self.document_storage.insert_documents_batch(
-            ids,
-            contents,
-            metadatas,
-        )
-        if len(int_ids) != content_count:
-            raise KnowledgeBaseUploadError(
-                stage="storage",
-                user_message=(
-                    f"存储失败：写入文档索引后返回的内部 ID 数量与文本分块数量不一致"
-                    f"（期望 {content_count}，实际 {len(int_ids)}）。"
-                ),
-                details={
-                    "expected_contents": content_count,
-                    "actual_int_ids": len(int_ids),
-                },
-            )
-
-        # 批量插入向量到 FAISS
+        # Validate vector format/dimension before any local write so validation
+        # failures never leave a half-written document store.
         try:
             vectors_array = np.asarray(vectors, dtype=np.float32)
         except (TypeError, ValueError) as exc:
@@ -187,7 +169,31 @@ class FaissVecDB(BaseVecDB):
                     "actual_dimension": int(vectors_array.shape[1]),
                 },
             )
-        await self.embedding_storage.insert_batch(vectors_array, int_ids)
+
+        # DocumentStorage and FAISS cannot share a transaction. If FAISS
+        # write fails after documents are committed, compensate by deleting.
+        int_ids = await self.document_storage.insert_documents_batch(
+            ids,
+            contents,
+            metadatas,
+        )
+        try:
+            if len(int_ids) != content_count:
+                raise KnowledgeBaseUploadError(
+                    stage="storage",
+                    user_message=(
+                        f"存储失败：写入文档索引后返回的内部 ID 数量与文本分块数量不一致"
+                        f"（期望 {content_count}，实际 {len(int_ids)}）。"
+                    ),
+                    details={
+                        "expected_contents": content_count,
+                        "actual_int_ids": len(int_ids),
+                    },
+                )
+            await self.embedding_storage.insert_batch(vectors_array, int_ids)
+        except Exception:
+            await self._rollback_partial_insert(ids=ids, int_ids=int_ids)
+            raise
         return int_ids
 
     async def retrieve(
@@ -255,6 +261,39 @@ class FaissVecDB(BaseVecDB):
 
         return top_k_results
 
+    async def _rollback_partial_insert(
+        self,
+        ids: list[str],
+        int_ids: list[int],
+    ) -> None:
+        """Best-effort cleanup after a partial ``insert_batch`` write.
+
+        DocumentStorage commits independently from FAISS. When FAISS insertion
+        fails after documents are written — or after vectors were added in
+        memory but not flushed — remove those rows/ids so the two stores do
+        not diverge.
+
+        Args:
+            ids: Chunk UUID strings written to DocumentStorage.
+            int_ids: Internal integer ids returned by DocumentStorage, also
+                used as FAISS ids when any vectors may have been added.
+        """
+        if int_ids:
+            try:
+                await self.embedding_storage.delete(int_ids)
+            except Exception as exc:
+                logger.warning(
+                    f"Failed to roll back FAISS vectors for partial insert: {exc}",
+                )
+
+        for doc_id in ids:
+            try:
+                await self.document_storage.delete_document_by_doc_id(doc_id)
+            except Exception as exc:
+                logger.warning(
+                    f"Failed to roll back document storage entry {doc_id}: {exc}",
+                )
+
     async def delete(self, doc_id: str) -> None:
         """删除一条文档块（chunk）"""
         # 获得对应的 int id

astrbot/core/knowledge_base/kb_helper.py

@@ -221,28 +221,35 @@ class KBHelper:
         progress_callback=None,
         pre_chunked_text: list[str] | None = None,
     ) -> KBDocument:
-        """上传并处理文档（带原子性保证和失败清理）
+        """Upload and process a document with compensating cleanup on failure.
 
-        流程:
-        1. 保存原始文件
-        2. 解析文档内容
-        3. 提取多媒体资源
-        4. 分块处理
-        5. 生成向量并存储
-        6. 保存元数据（事务）
-        7. 更新统计
+        Flow:
+        1. Parse document content
+        2. Extract media resources
+        3. Chunk text
+        4. Generate embeddings and store them (chunk text DB + FAISS)
+        5. Persist document metadata (KBDocument / KBMedia)
+        6. Refresh stats
+
+        Multi-store writes cannot share a transaction. Failures before metadata
+        commit best-effort roll back written chunks/vectors/media. After
+        metadata is committed, only report stats-refresh errors and keep the
+        document.
 
         Args:
-            progress_callback: 进度回调函数，接收参数 (stage, current, total)
-                - stage: 当前阶段 ('parsing', 'chunking', 'embedding')
-                - current: 当前进度
-                - total: 总数
-
+            progress_callback: Progress callback ``(stage, current, total)``.
+                - stage: Current stage (``parsing``, ``chunking``, ``embedding``)
+                - current: Current progress
+                - total: Total units
         """
         await self._ensure_vec_db()
         doc_id = str(uuid.uuid4())
         media_paths: list[Path] = []
         file_size = 0
+        # Only roll back chunks/vectors/media when metadata has not been
+        # committed yet. After commit (e.g. stats refresh failure) the document
+        # is already user-visible and must not be fully undone.
+        metadata_committed = False
 
         # file_path = self.kb_files_dir / f"{doc_id}.{file_type}"
         # async with aiofiles.open(file_path, "wb") as f:
@@ -397,7 +404,11 @@ class KBHelper:
                 raise KnowledgeBaseUploadError(
                     stage="storage",
                     user_message=("存储失败：文本块已生成，但写入知识库索引时出错。"),
-                    details={"file_name": file_name},
+                    details={
+                        "file_name": file_name,
+                        "doc_id": doc_id,
+                        "cause": str(exc),
+                    },
                 ) from exc
 
             # 保存文档的元数据
@@ -419,11 +430,21 @@ class KBHelper:
                         for media in saved_media:
                             session.add(media)
                         await session.commit()
-
+                    # Mark committed immediately after commit succeeds. A later
+                    # refresh failure must not trigger full upload rollback.
+                    metadata_committed = True
                     await session.refresh(doc)
             except KnowledgeBaseUploadError:
                 raise
             except Exception as exc:
+                if metadata_committed:
+                    raise KnowledgeBaseUploadError(
+                        stage="metadata",
+                        user_message=(
+                            "元数据更新失败：文档已上传，但文档记录刷新失败。"
+                        ),
+                        details={"file_name": file_name, "doc_id": doc_id},
+                    ) from exc
                 raise KnowledgeBaseUploadError(
                     stage="metadata",
                     user_message=(
@@ -453,18 +474,84 @@ class KBHelper:
                 logger.warning(f"上传文档失败: {e}", extra={"details": e.details})
             else:
                 logger.error(f"上传文档失败: {e}", exc_info=True)
-            # if file_path.exists():
-            #     file_path.unlink()
 
-            for media_path in media_paths:
-                try:
-                    if media_path.exists():
-                        media_path.unlink()
-                except Exception as me:
-                    logger.warning(f"清理多媒体文件失败 {media_path}: {me}")
+            if not metadata_committed:
+                await self._cleanup_failed_upload(
+                    doc_id=doc_id, media_paths=media_paths
+                )
 
             raise
 
+    async def _cleanup_failed_upload(
+        self,
+        doc_id: str,
+        media_paths: list[Path],
+    ) -> None:
+        """Best-effort compensating cleanup after a failed upload.
+
+        Multi-store writes (media files, chunk/FTS rows, FAISS vectors, KB
+        metadata) cannot share a single transaction. On failure before the KB
+        document row is committed, remove any partial state keyed by ``doc_id``.
+
+        Cleanup order intentionally differs from user-facing document deletion:
+        chunk/vector data is removed first, then residual KB metadata rows.
+        This avoids the "metadata gone, orphans remain" window that
+        ``delete_document_by_id`` can leave when vector deletion fails.
+
+        Args:
+            doc_id: Pre-generated document id used for this upload attempt.
+            media_paths: Media files written to disk during this attempt.
+        """
+        from sqlalchemy import delete
+        from sqlmodel import col
+
+        # 1) chunks + vectors first (most common orphan after partial insert)
+        vec_db = getattr(self, "vec_db", None)
+        if vec_db is not None:
+            try:
+                await vec_db.delete_documents(
+                    metadata_filters={"kb_doc_id": doc_id},
+                )
+            except Exception as ve:
+                logger.warning(
+                    f"Failed to roll back chunks/vectors for failed upload "
+                    f"(doc_id={doc_id}): {ve}",
+                )
+
+        # 2) residual KBDocument / KBMedia rows only (normally none yet)
+        try:
+            async with self.kb_db.get_db() as session, session.begin():
+                await session.execute(
+                    delete(KBMedia).where(col(KBMedia.doc_id) == doc_id),
+                )
+                await session.execute(
+                    delete(KBDocument).where(col(KBDocument.doc_id) == doc_id),
+                )
+        except Exception as exc:
+            logger.warning(
+                f"Failed to roll back document metadata for failed upload "
+                f"(doc_id={doc_id}): {exc}",
+            )
+
+        # 3) media files on disk
+        for media_path in media_paths:
+            try:
+                if media_path.exists():
+                    media_path.unlink()
+            except Exception as me:
+                logger.warning(f"Failed to clean up media file {media_path}: {me}")
+
+        # 4) empty media directory for this doc
+        try:
+            media_dir = self.kb_medias_dir / doc_id
+            if media_dir.exists() and media_dir.is_dir():
+                media_dir.rmdir()
+        except Exception as de:
+            logger.warning(
+                f"Failed to remove media directory after failed upload "
+                f"(doc_id={doc_id}): {de}",
+            )
+
     async def list_documents(
         self,
         offset: int = 0,

tests/unit/test_kb_upload_atomicity.py

@@ -0,0 +1,574 @@
+"""
+Unit tests for knowledge base upload atomicity / compensating rollback.
+
+Covers:
+1. insert_batch rolls back document rows when FAISS insert fails
+2. upload_document cleans up chunks/vectors/media on storage failure
+3. upload_document cleans up after metadata failure (post insert_batch)
+4. upload_document does NOT roll back after metadata is committed
+5. Real DocumentStorage + EmbeddingStorage leave no orphans on FAISS failure
+6. Dimension validation failures never write to DocumentStorage
+
+These tests use lazy imports and a ProviderManager stub to avoid circular
+import issues in the astrbot core module chain (same pattern as
+test_kb_manager_resilience.py).
+"""
+
+import sys
+import types
+from contextlib import asynccontextmanager
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from astrbot.core.db.vec_db.faiss_impl.vec_db import FaissVecDB
+from astrbot.core.exceptions import KnowledgeBaseUploadError
+from astrbot.core.knowledge_base.models import KnowledgeBase
+
+
+@pytest.fixture
+def stub_provider_manager_module():
+    """Stub provider manager module to avoid circular imports in unit tests."""
+    original_module = sys.modules.get("astrbot.core.provider.manager")
+    stub_module = types.ModuleType("astrbot.core.provider.manager")
+
+    class ProviderManager: ...
+
+    setattr(stub_module, "ProviderManager", ProviderManager)
+    sys.modules["astrbot.core.provider.manager"] = stub_module
+
+    # Drop already-imported modules that transitively need ProviderManager so
+    # they re-import against the stub.
+    to_drop = [
+        name
+        for name in list(sys.modules)
+        if name.startswith("astrbot.core.knowledge_base.kb_helper")
+        or name.startswith("astrbot.core.knowledge_base.kb_mgr")
+    ]
+    for name in to_drop:
+        sys.modules.pop(name, None)
+
+    try:
+        yield
+    finally:
+        if original_module is not None:
+            sys.modules["astrbot.core.provider.manager"] = original_module
+        else:
+            sys.modules.pop("astrbot.core.provider.manager", None)
+
+
+def _make_vec_db() -> FaissVecDB:
+    vec_db = FaissVecDB.__new__(FaissVecDB)
+    vec_db.embedding_provider = AsyncMock()
+    vec_db.document_storage = AsyncMock()
+    vec_db.embedding_storage = AsyncMock()
+    vec_db.embedding_storage.dimension = 2
+    return vec_db
+
+
+def _import_kb_helper():
+    from astrbot.core.knowledge_base.kb_helper import KBHelper
+
+    return KBHelper
+
+
+def _failing_get_db():
+    @asynccontextmanager
+    async def _cm():
+        raise RuntimeError("kb.db locked")
+        yield  # pragma: no cover
+
+    return _cm
+
+
+def _successful_get_db(session):
+    @asynccontextmanager
+    async def _cm():
+        yield session
+
+    return _cm
+
+
+def _successful_begin():
+    @asynccontextmanager
+    async def _cm():
+        yield None
+
+    return _cm
+
+
+def _session_with_begin(execute_side_effect=None):
+    """Session mock that supports ``async with session.begin()``."""
+    session = MagicMock()
+    session.begin = _successful_begin()
+    session.add = MagicMock()
+    session.commit = AsyncMock()
+    session.refresh = AsyncMock()
+    session.execute = AsyncMock(side_effect=execute_side_effect)
+    return session
+
+
+async def _make_real_vec_db(tmp_path: Path, dim: int = 4) -> FaissVecDB:
+    """Build a FaissVecDB backed by real DocumentStorage + EmbeddingStorage."""
+    doc_path = str(tmp_path / "doc.db")
+    index_path = str(tmp_path / "index.faiss")
+
+    embedding_provider = MagicMock()
+    # get_dim is sync in EmbeddingProvider; must return a plain int for FAISS.
+    embedding_provider.get_dim = MagicMock(return_value=dim)
+    embedding_provider.get_embeddings_batch = AsyncMock()
+    embedding_provider.get_embedding = AsyncMock()
+
+    vec_db = FaissVecDB(
+        doc_store_path=doc_path,
+        index_store_path=index_path,
+        embedding_provider=embedding_provider,
+    )
+    await vec_db.initialize()
+    return vec_db
+
+
+@pytest.mark.asyncio
+async def test_insert_batch_rolls_back_documents_when_faiss_fails() -> None:
+    """FAISS write failure should delete rows already committed to DocumentStorage."""
+    vec_db = _make_vec_db()
+    vec_db.embedding_provider.get_embeddings_batch.return_value = [
+        [0.1, 0.2],
+        [0.3, 0.4],
+    ]
+    vec_db.document_storage.insert_documents_batch.return_value = [11, 12]
+    vec_db.embedding_storage.insert_batch.side_effect = RuntimeError(
+        "faiss write failed",
+    )
+
+    with pytest.raises(RuntimeError, match="faiss write failed"):
+        await FaissVecDB.insert_batch(
+            vec_db,
+            contents=["chunk-1", "chunk-2"],
+            metadatas=[{"kb_doc_id": "doc-1"}, {"kb_doc_id": "doc-1"}],
+            ids=["c1", "c2"],
+        )
+
+    vec_db.embedding_storage.delete.assert_awaited_once_with([11, 12])
+    assert vec_db.document_storage.delete_document_by_doc_id.await_count == 2
+    vec_db.document_storage.delete_document_by_doc_id.assert_any_await("c1")
+    vec_db.document_storage.delete_document_by_doc_id.assert_any_await("c2")
+
+
+@pytest.mark.asyncio
+async def test_insert_batch_rolls_back_on_int_id_count_mismatch() -> None:
+    """Mismatched int_id count after document insert should roll back those docs."""
+    vec_db = _make_vec_db()
+    vec_db.embedding_provider.get_embeddings_batch.return_value = [
+        [0.1, 0.2],
+        [0.3, 0.4],
+    ]
+    vec_db.document_storage.insert_documents_batch.return_value = [11]  # mismatch
+
+    with pytest.raises(KnowledgeBaseUploadError) as exc_info:
+        await FaissVecDB.insert_batch(
+            vec_db,
+            contents=["chunk-1", "chunk-2"],
+            metadatas=[{}, {}],
+            ids=["c1", "c2"],
+        )
+
+    assert "内部 ID 数量" in str(exc_info.value)
+    vec_db.embedding_storage.insert_batch.assert_not_awaited()
+    vec_db.embedding_storage.delete.assert_awaited_once_with([11])
+    assert vec_db.document_storage.delete_document_by_doc_id.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_insert_batch_dimension_mismatch_does_not_write_documents() -> None:
+    """Dimension validation must fail before DocumentStorage is written."""
+    vec_db = _make_vec_db()
+    vec_db.embedding_storage.dimension = 4
+    vec_db.embedding_provider.get_embeddings_batch.return_value = [
+        [0.1, 0.2],  # wrong dim
+        [0.3, 0.4],
+    ]
+
+    with pytest.raises(KnowledgeBaseUploadError) as exc_info:
+        await FaissVecDB.insert_batch(
+            vec_db,
+            contents=["chunk-1", "chunk-2"],
+            metadatas=[{}, {}],
+            ids=["c1", "c2"],
+        )
+
+    assert "维度" in str(exc_info.value)
+    vec_db.document_storage.insert_documents_batch.assert_not_awaited()
+    vec_db.embedding_storage.insert_batch.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_insert_batch_real_storage_rolls_back_on_faiss_failure(
+    tmp_path: Path,
+) -> None:
+    """Real DocumentStorage must not keep orphan chunks when FAISS write fails."""
+    vec_db = await _make_real_vec_db(tmp_path, dim=4)
+    vec_db.embedding_provider.get_embeddings_batch.return_value = [
+        [0.1, 0.2, 0.3, 0.4],
+        [0.5, 0.6, 0.7, 0.8],
+    ]
+
+    original_insert = vec_db.embedding_storage.insert_batch
+
+    async def boom(vectors, ids):
+        # Simulate failure after in-memory add would happen: force raise before
+        # success so DocumentStorage rows must be cleaned up.
+        raise RuntimeError("simulated faiss disk write failure")
+
+    vec_db.embedding_storage.insert_batch = boom  # type: ignore[method-assign]
+
+    kb_doc_id = "kb-doc-real-1"
+    with pytest.raises(RuntimeError, match="simulated faiss"):
+        await vec_db.insert_batch(
+            contents=["alpha chunk", "beta chunk"],
+            metadatas=[
+                {"kb_id": "kb-1", "kb_doc_id": kb_doc_id, "chunk_index": 0},
+                {"kb_id": "kb-1", "kb_doc_id": kb_doc_id, "chunk_index": 1},
+            ],
+            ids=["chunk-a", "chunk-b"],
+        )
+
+    remaining = await vec_db.document_storage.get_documents(
+        metadata_filters={"kb_doc_id": kb_doc_id},
+        offset=None,
+        limit=None,
+    )
+    assert remaining == []
+    assert await vec_db.count_documents(metadata_filter={"kb_doc_id": kb_doc_id}) == 0
+
+    # Restore for clean close; index should still be empty / consistent.
+    vec_db.embedding_storage.insert_batch = original_insert  # type: ignore[method-assign]
+    await vec_db.close()
+
+
+@pytest.mark.asyncio
+async def test_insert_batch_real_storage_dimension_mismatch_leaves_no_docs(
+    tmp_path: Path,
+) -> None:
+    """Real storage: wrong embedding dim must not leave any document rows."""
+    vec_db = await _make_real_vec_db(tmp_path, dim=4)
+    vec_db.embedding_provider.get_embeddings_batch.return_value = [
+        [0.1, 0.2],  # dim 2 != 4
+        [0.3, 0.4],
+    ]
+
+    with pytest.raises(KnowledgeBaseUploadError):
+        await vec_db.insert_batch(
+            contents=["alpha", "beta"],
+            metadatas=[
+                {"kb_id": "kb-1", "kb_doc_id": "doc-dim", "chunk_index": 0},
+                {"kb_id": "kb-1", "kb_doc_id": "doc-dim", "chunk_index": 1},
+            ],
+            ids=["c1", "c2"],
+        )
+
+    remaining = await vec_db.document_storage.get_documents(
+        metadata_filters={"kb_doc_id": "doc-dim"},
+        offset=None,
+        limit=None,
+    )
+    assert remaining == []
+    await vec_db.close()
+
+
+@pytest.mark.asyncio
+async def test_upload_document_cleans_up_on_storage_failure(
+    tmp_path: Path,
+    stub_provider_manager_module,
+) -> None:
+    """Storage failure should clean media and request chunk/vector rollback."""
+    KBHelper = _import_kb_helper()
+
+    helper = KBHelper.__new__(KBHelper)
+    helper.kb = KnowledgeBase(
+        kb_name="Test KB",
+        description="",
+        embedding_provider_id="emb",
+    )
+    helper.kb_db = MagicMock()
+    helper.vec_db = AsyncMock()
+    helper.kb_medias_dir = tmp_path / "medias"
+    helper.kb_medias_dir.mkdir()
+    helper.chunker = AsyncMock()
+    helper.chunker.chunk = AsyncMock(return_value=["hello world"])
+
+    media_file = helper.kb_medias_dir / "will-be-set" / "img.png"
+
+    async def fake_save_media(**kwargs):
+        nonlocal media_file
+        doc_id = kwargs["doc_id"]
+        media_dir = helper.kb_medias_dir / doc_id
+        media_dir.mkdir(parents=True, exist_ok=True)
+        media_file = media_dir / "img.png"
+        media_file.write_bytes(b"fake-image")
+        media = MagicMock()
+        media.file_path = str(media_file)
+        return media
+
+    helper._save_media = AsyncMock(side_effect=fake_save_media)
+    helper.vec_db.insert_batch.side_effect = RuntimeError("embedding provider down")
+    helper.vec_db.delete_documents = AsyncMock()
+    helper.kb_db.get_db = _successful_get_db(_session_with_begin())
+
+    parse_result = MagicMock()
+    parse_result.text = "hello world"
+    parse_result.media = [
+        MagicMock(
+            media_type="image",
+            file_name="img.png",
+            content=b"fake-image",
+            mime_type="image/png",
+        ),
+    ]
+
+    with (
+        patch(
+            "astrbot.core.knowledge_base.kb_helper.select_parser",
+            new=AsyncMock(
+                return_value=MagicMock(parse=AsyncMock(return_value=parse_result)),
+            ),
+        ),
+        patch.object(helper, "_ensure_vec_db", new=AsyncMock()),
+        pytest.raises(KnowledgeBaseUploadError) as exc_info,
+    ):
+        await helper.upload_document(
+            file_name="demo.txt",
+            file_content=b"hello world",
+            file_type="txt",
+        )
+
+    assert exc_info.value.stage == "storage"
+    assert "cause" in exc_info.value.details
+    helper.vec_db.delete_documents.assert_awaited()
+    assert not media_file.exists()
+
+
+@pytest.mark.asyncio
+async def test_upload_document_cleans_up_on_metadata_failure(
+    stub_provider_manager_module,
+) -> None:
+    """Metadata commit failure after insert_batch should delete written chunks."""
+    KBHelper = _import_kb_helper()
+
+    helper = KBHelper.__new__(KBHelper)
+    helper.kb = KnowledgeBase(
+        kb_name="Test KB",
+        description="",
+        embedding_provider_id="emb",
+    )
+    helper.kb_db = MagicMock()
+    helper.vec_db = AsyncMock()
+    helper.kb_medias_dir = Path("/tmp/kb-medias-unused")
+    helper.chunker = AsyncMock()
+
+    helper.vec_db.insert_batch = AsyncMock()
+    helper.vec_db.delete_documents = AsyncMock()
+    helper.kb_db.get_db = _failing_get_db()
+
+    with (
+        patch.object(helper, "_ensure_vec_db", new=AsyncMock()),
+        pytest.raises(KnowledgeBaseUploadError) as exc_info,
+    ):
+        await helper.upload_document(
+            file_name="demo.txt",
+            file_content=None,
+            file_type="txt",
+            pre_chunked_text=["chunk a", "chunk b"],
+        )
+
+    assert exc_info.value.stage == "metadata"
+    helper.vec_db.insert_batch.assert_awaited_once()
+    helper.vec_db.delete_documents.assert_awaited()
+
+
+@pytest.mark.asyncio
+async def test_upload_document_skips_rollback_after_metadata_commit(
+    stub_provider_manager_module,
+) -> None:
+    """Stats refresh failure after metadata commit must not roll back the doc."""
+    KBHelper = _import_kb_helper()
+
+    helper = KBHelper.__new__(KBHelper)
+    helper.kb = KnowledgeBase(
+        kb_name="Test KB",
+        description="",
+        embedding_provider_id="emb",
+    )
+    helper.kb_db = MagicMock()
+    helper.vec_db = AsyncMock()
+    helper.kb_medias_dir = Path("/tmp/kb-medias-unused")
+    helper.chunker = AsyncMock()
+    helper.vec_db.insert_batch = AsyncMock()
+    helper.vec_db.delete_documents = AsyncMock()
+    helper.kb_db.update_kb_stats = AsyncMock(side_effect=RuntimeError("stats fail"))
+    helper.refresh_kb = AsyncMock()
+    helper.refresh_document = AsyncMock()
+
+    session = _session_with_begin()
+    helper.kb_db.get_db = _successful_get_db(session)
+
+    with (
+        patch.object(helper, "_ensure_vec_db", new=AsyncMock()),
+        pytest.raises(KnowledgeBaseUploadError) as exc_info,
+    ):
+        await helper.upload_document(
+            file_name="demo.txt",
+            file_content=None,
+            file_type="txt",
+            pre_chunked_text=["chunk a"],
+        )
+
+    assert exc_info.value.stage == "metadata"
+    assert "统计信息刷新失败" in str(exc_info.value)
+    helper.vec_db.delete_documents.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_upload_document_skips_rollback_when_refresh_fails_after_commit(
+    stub_provider_manager_module,
+) -> None:
+    """If commit succeeds but session.refresh fails, do not roll back."""
+    KBHelper = _import_kb_helper()
+
+    helper = KBHelper.__new__(KBHelper)
+    helper.kb = KnowledgeBase(
+        kb_name="Test KB",
+        description="",
+        embedding_provider_id="emb",
+    )
+    helper.kb_db = MagicMock()
+    helper.vec_db = AsyncMock()
+    helper.kb_medias_dir = Path("/tmp/kb-medias-unused")
+    helper.chunker = AsyncMock()
+    helper.vec_db.insert_batch = AsyncMock()
+    helper.vec_db.delete_documents = AsyncMock()
+
+    session = _session_with_begin()
+    session.refresh = AsyncMock(side_effect=RuntimeError("refresh failed"))
+    helper.kb_db.get_db = _successful_get_db(session)
+
+    with (
+        patch.object(helper, "_ensure_vec_db", new=AsyncMock()),
+        pytest.raises(KnowledgeBaseUploadError) as exc_info,
+    ):
+        await helper.upload_document(
+            file_name="demo.txt",
+            file_content=None,
+            file_type="txt",
+            pre_chunked_text=["chunk a"],
+        )
+
+    assert exc_info.value.stage == "metadata"
+    assert "文档记录刷新失败" in str(exc_info.value)
+    helper.vec_db.delete_documents.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_cleanup_failed_upload_deletes_vectors_before_metadata(
+    tmp_path: Path,
+    stub_provider_manager_module,
+) -> None:
+    """Rollback should delete vectors first, then metadata rows, then media."""
+    KBHelper = _import_kb_helper()
+
+    helper = KBHelper.__new__(KBHelper)
+    helper.vec_db = AsyncMock()
+    helper.kb_db = MagicMock()
+    helper.kb_medias_dir = tmp_path / "medias"
+    helper.kb_medias_dir.mkdir()
+
+    call_order: list[str] = []
+
+    async def track_delete_documents(**kwargs):
+        call_order.append("vectors")
+
+    helper.vec_db.delete_documents = AsyncMock(side_effect=track_delete_documents)
+
+    async def track_execute(*args, **kwargs):
+        if "vectors" in call_order and "metadata" not in call_order:
+            call_order.append("metadata")
+        return None
+
+    helper.kb_db.get_db = _successful_get_db(
+        _session_with_begin(execute_side_effect=track_execute),
+    )
+
+    media = helper.kb_medias_dir / "doc-x" / "a.png"
+    media.parent.mkdir()
+    media.write_bytes(b"x")
+
+    await helper._cleanup_failed_upload(doc_id="doc-x", media_paths=[media])
+
+    assert call_order[0] == "vectors"
+    assert "metadata" in call_order
+    assert not media.exists()
+    helper.vec_db.delete_documents.assert_awaited_once_with(
+        metadata_filters={"kb_doc_id": "doc-x"},
+    )
+
+
+@pytest.mark.asyncio
+async def test_cleanup_failed_upload_is_best_effort(
+    tmp_path: Path,
+    stub_provider_manager_module,
+) -> None:
+    """Rollback path failures must not raise; media cleanup still runs."""
+    KBHelper = _import_kb_helper()
+
+    helper = KBHelper.__new__(KBHelper)
+    helper.kb_db = MagicMock()
+    helper.vec_db = AsyncMock()
+    helper.vec_db.delete_documents.side_effect = RuntimeError("vec delete failed")
+    helper.kb_db.get_db = _failing_get_db()
+    helper.kb_medias_dir = tmp_path / "medias"
+    helper.kb_medias_dir.mkdir()
+
+    media = helper.kb_medias_dir / "doc-x" / "a.png"
+    media.parent.mkdir()
+    media.write_bytes(b"x")
+
+    # Should not raise
+    await helper._cleanup_failed_upload(doc_id="doc-x", media_paths=[media])
+    assert not media.exists()
+
+
+@pytest.mark.asyncio
+async def test_cleanup_failed_upload_real_vec_db_by_kb_doc_id(
+    tmp_path: Path,
+    stub_provider_manager_module,
+) -> None:
+    """Cleanup with a real vec_db removes chunks keyed by kb_doc_id."""
+    KBHelper = _import_kb_helper()
+    vec_db = await _make_real_vec_db(tmp_path, dim=4)
+    vec_db.embedding_provider.get_embeddings_batch.return_value = [
+        [0.1, 0.2, 0.3, 0.4],
+        [0.5, 0.6, 0.7, 0.8],
+    ]
+
+    kb_doc_id = "upload-doc-cleanup"
+    await vec_db.insert_batch(
+        contents=["one", "two"],
+        metadatas=[
+            {"kb_id": "kb-1", "kb_doc_id": kb_doc_id, "chunk_index": 0},
+            {"kb_id": "kb-1", "kb_doc_id": kb_doc_id, "chunk_index": 1},
+        ],
+        ids=["u1", "u2"],
+    )
+    assert await vec_db.count_documents(metadata_filter={"kb_doc_id": kb_doc_id}) == 2
+
+    helper = KBHelper.__new__(KBHelper)
+    helper.vec_db = vec_db
+    helper.kb_db = MagicMock()
+    helper.kb_medias_dir = tmp_path / "medias"
+    helper.kb_medias_dir.mkdir()
+    helper.kb_db.get_db = _successful_get_db(_session_with_begin())
+
+    await helper._cleanup_failed_upload(doc_id=kb_doc_id, media_paths=[])
+    assert await vec_db.count_documents(metadata_filter={"kb_doc_id": kb_doc_id}) == 0
+    await vec_db.close()