fix(api): decrement knowledgebase counters on SDK re-parse / stop-parse (#17236)

api/apps/restful_apis/chunk_api.py

@@ -33,6 +33,7 @@ from api.db.joint_services.tenant_model_service import (
 )
 from api.db.db_models import Document, Task
 from api.db.services.doc_metadata_service import DocMetadataService
+from api.db.services.document_counter_service import release_reparse_counters
 from api.db.services.document_service import DocumentService
 from api.db.services.file2document_service import File2DocumentService
 from api.db.services.knowledgebase_service import KnowledgebaseService
@@ -179,6 +180,22 @@ def _enrich_chunks_with_document_metadata(chunks: list[dict], metadata_fields=No
     enrich_chunks_with_document_metadata(chunks, metadata_fields)
 
 
+def _release_doc_counters(doc):
+    """Roll back the document's and knowledgebase's chunk/token/duration counters
+    so a re-parse starts from zero. Callers that delete a document's chunks must
+    do this, otherwise the removed counts stay in the knowledgebase total. The
+    release re-reads the row under a lock (see release_reparse_counters) so it is
+    safe against a worker still parsing the document. Returns an error result if
+    the document is gone, else None.
+    """
+    try:
+        release_reparse_counters(doc.id)
+    except LookupError:
+        logging.exception("Failed to release counters for document %s in knowledgebase %s", doc.id, doc.kb_id)
+        return get_error_data_result(message=f"Document {doc.id} not found")
+    return None
+
+
 @manager.route("/datasets/<dataset_id>/chunks", methods=["POST"])  # noqa: F821
 @login_required
 @add_tenant_id_to_kwargs
@@ -211,7 +228,7 @@ async def parse(tenant_id, dataset_id):
             continue
         if not doc:
             return get_error_data_result(message=f"you don't own the document {id}")
-        info = {"run": "1", "progress": 0, "progress_msg": "", "chunk_num": 0, "token_num": 0}
+        info = {"run": "1", "progress": 0, "progress_msg": ""}
         if (
             DocumentService.filter_update(
                 [
@@ -223,6 +240,8 @@ async def parse(tenant_id, dataset_id):
             == 0
         ):
             return get_error_data_result("Can't parse document that is currently being processed")
+        if err := _release_doc_counters(doc[0]):
+            return err
         index_name = search.index_name(dataset_tenant_id)
         if settings.docStoreConn.index_exist(index_name, doc[0].kb_id):
             settings.docStoreConn.delete({"doc_id": id}, index_name, doc[0].kb_id)
@@ -283,8 +302,10 @@ async def stop_parsing(tenant_id, dataset_id):
                 data={"error_code": DOC_STOP_PARSING_INVALID_STATE_ERROR_CODE},
             )
         cancel_all_task_of(id)
-        info = {"run": "2", "progress": 0, "chunk_num": 0}
+        info = {"run": "2", "progress": 0}
         DocumentService.update_by_id(id, info)
+        if err := _release_doc_counters(doc[0]):
+            return err
         index_name = search.index_name(dataset_tenant_id)
         if settings.docStoreConn.index_exist(index_name, doc[0].kb_id):
             settings.docStoreConn.delete({"doc_id": doc[0].id}, index_name, doc[0].kb_id)

api/apps/restful_apis/document_api.py

@@ -40,6 +40,7 @@ from api.db import VALID_FILE_TYPES, FileType
 from api.db.db_models import API4Conversation, DB
 from api.db.services import duplicate_name
 from api.db.services.doc_metadata_service import DocMetadataService
+from api.db.services.document_counter_service import release_reparse_counters
 from api.db.db_models import Task
 from api.db.services.document_service import DocumentService
 from api.db.services.file2document_service import File2DocumentService
@@ -1734,13 +1735,22 @@ async def stop_parse_documents(tenant_id, dataset_id):
                     continue
 
                 cancel_all_task_of(doc_id)
+                # Release the document's partial chunk/token counts from the
+                # knowledgebase aggregate under the row lock (see
+                # release_reparse_counters). This is the sole counter adjustment,
+                # so the status update below must not touch chunk_num.
+                try:
+                    release_reparse_counters(doc_id)
+                except LookupError:
+                    logging.exception("Failed to release counters for document %s during stop-parse", doc_id)
+                    errors.append(f"Document not found: {doc_id}")
+                    continue
                 cancel_doc_msg = f"\n{datetime.now().strftime('%H:%M:%S')} Task stopped by user."
                 DocumentService.update_by_id(
                     doc_id,
                     {
                         "run": str(TaskStatus.CANCEL.value),
                         "progress": 0,
-                        "chunk_num": 0,
                         "progress_msg": (doc.progress_msg or "") + cancel_doc_msg,
                     },
                 )

api/apps/services/document_api_service.py

@@ -15,6 +15,7 @@
 #
 import logging
 
+from api.db.services.document_counter_service import release_reparse_counters
 from api.db.services.document_service import DocumentService
 from api.db.services.file2document_service import File2DocumentService
 from api.db.services.file_service import FileService
@@ -128,22 +129,14 @@ def reset_document_for_reparse(doc, tenant_id, parser_id=None, pipeline_id=None)
     if not e:
         return get_error_data_result(message="document not found")
 
-    # Update document statistics before deleting all document rows. Pipeline
-    # compilation rows may exist even when token_num is zero, so the doc-store
-    # cleanup must not be gated by the document counters.
-    if doc.token_num > 0:
-        try:
-            e = DocumentService.increment_chunk_num(
-                doc.id,
-                doc.kb_id,
-                doc.token_num * -1,
-                doc.chunk_num * -1,
-                doc.process_duration * -1,
-            )
-        except LookupError:
-            return get_error_data_result(message="document not found")
-        if not e:
-            return get_error_data_result(message="document not found")
+    # Release the document's chunk/token/duration counters from the knowledgebase
+    # aggregate under a row lock before clearing the chunks. release_reparse_counters
+    # guards the zero case internally, so the doc-store cleanup below still runs for
+    # pipeline compilation rows that exist even when token_num is zero.
+    try:
+        release_reparse_counters(doc.id)
+    except LookupError:
+        return get_error_data_result(message="Document not found!")
     settings.docStoreConn.delete({"doc_id": doc.id}, search.index_name(tenant_id), doc.kb_id)
 
     # Delete chunk images

api/db/services/document_counter_service.py

@@ -0,0 +1,51 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+import logging
+
+from api.db.db_models import DB, Document
+from api.db.services.document_service import DocumentService
+
+
+def release_reparse_counters(doc_id):
+    """Roll back a document's chunk, token, and duration counters and the owning
+    knowledgebase's chunk/token totals so a re-parse starts from zero.
+
+    The counters are re-read under a ``FOR UPDATE`` row lock in the same
+    transaction as the decrement, so the release subtracts the row's committed
+    value at release time rather than a request-time snapshot a concurrent worker
+    may have already moved past. ``increment_chunk_num`` updates both ledgers
+    together. This does not serialize a worker that writes its final counts in a
+    separate transaction after the release commits; fully closing the
+    stop-parse-during-parse race needs a worker-side cancel check.
+
+    Raises ``LookupError`` if the document row no longer exists so callers can
+    surface a not-found result.
+    """
+    with DB.atomic():
+        fresh = Document.select().where(Document.id == doc_id).for_update().first()
+        if fresh is None:
+            raise LookupError(doc_id)
+        if not (fresh.token_num or fresh.chunk_num or fresh.process_duration):
+            logging.debug("release_reparse_counters: nothing to release for document %s", doc_id)
+            return
+        DocumentService.increment_chunk_num(fresh.id, fresh.kb_id, -fresh.token_num, -fresh.chunk_num, -fresh.process_duration)
+        logging.debug(
+            "release_reparse_counters: released document %s (token=%s chunk=%s duration=%s)",
+            doc_id,
+            fresh.token_num,
+            fresh.chunk_num,
+            fresh.process_duration,
+        )

test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py

@@ -14,6 +14,8 @@
 #  limitations under the License.
 #
 import asyncio
+import contextlib
+import functools
 import inspect
 import importlib.util
 import sys
@@ -85,6 +87,7 @@ class _DummyDoc:
         doc_type=FileType.OTHER,
         status=True,
         run=0,
+        progress_msg="",
     ):
         self.id = doc_id
         self.kb_id = kb_id
@@ -97,6 +100,7 @@ class _DummyDoc:
         self.type = doc_type
         self.status = status
         self.run = run
+        self.progress_msg = progress_msg
 
     def to_dict(self):
         return {
@@ -137,7 +141,19 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"):
     monkeypatch.setitem(sys.modules, "common", common_pkg)
 
     apps_mod = ModuleType("api.apps")
-    apps_mod.login_required = lambda func: func
+
+    def _login_required(func=None, **_kwargs):
+        # Real login_required is used both bare (chunk_api) and as a factory with
+        # auth_types=[...] (document_api); the mock must pass through both forms.
+        if func is None:
+            return lambda inner: inner
+        return func
+
+    apps_mod.login_required = _login_required
+    apps_mod.current_user = SimpleNamespace(id="tenant-1")
+    apps_mod.AUTH_JWT = None
+    apps_mod.AUTH_API = None
+    apps_mod.AUTH_BETA = None
     monkeypatch.setitem(sys.modules, "api.apps", apps_mod)
 
     common_settings_mod = ModuleType("common.settings")
@@ -152,6 +168,7 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"):
         return func(*args, **kwargs)
 
     common_misc_utils_mod.thread_pool_exec = _thread_pool_exec
+    common_misc_utils_mod.get_uuid = lambda: "uuid-1"
     monkeypatch.setitem(sys.modules, "common.misc_utils", common_misc_utils_mod)
 
     common_string_utils_mod = ModuleType("common.string_utils")
@@ -180,21 +197,58 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"):
         def is_null(self, value=True):
             return _FakeExpr()
 
+    class _StubFreshDoc:
+        id = "doc-1"
+        kb_id = "kb-1"
+        token_num = 2
+        chunk_num = 1
+        process_duration = 0.0
+
+    class _StubDocQuery:
+        # Stands in for Document.select().where(...).for_update().first().
+        def where(self, *_args, **_kwargs):
+            return self
+
+        def for_update(self):
+            return self
+
+        def first(self):
+            return _StubDocumentModel.fresh_doc
+
     class _StubDocumentModel:
         id = _FakeField()
         run = _FakeField()
+        # The row re-read under lock by _release_doc_counters; tests that assert on
+        # the decrement override this to mirror the document's current counters.
+        fresh_doc = _StubFreshDoc()
+
+        @classmethod
+        def select(cls, *_args, **_kwargs):
+            return _StubDocQuery()
 
     class _StubTaskModel:
         doc_id = _FakeField()
 
+    class _AnyFieldMeta(type):
+        def __getattr__(cls, _name):
+            return _FakeField()
+
+    class _StubModel(metaclass=_AnyFieldMeta):
+        pass
+
     db_models_mod = ModuleType("api.db.db_models")
     db_models_mod.APIToken = SimpleNamespace(query=lambda **_kwargs: [])
     db_models_mod.Document = _StubDocumentModel
     db_models_mod.Task = _StubTaskModel
+    db_models_mod.DB = SimpleNamespace(atomic=lambda: contextlib.nullcontext(), connection_context=lambda: lambda fn: fn)
+    # Transitively-loaded real services import assorted model classes (File,
+    # Knowledgebase, UserTenant, ...); hand them a permissive stub on demand.
+    db_models_mod.__getattr__ = lambda _name: _StubModel
     monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod)
 
     services_pkg = ModuleType("api.db.services")
     services_pkg.__path__ = [str(repo_root / "api" / "db" / "services")]
+    services_pkg.duplicate_name = lambda _query, name="", **_kwargs: name
     monkeypatch.setitem(sys.modules, "api.db.services", services_pkg)
 
     doc_metadata_service_mod = ModuleType("api.db.services.doc_metadata_service")
@@ -207,10 +261,12 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"):
     document_service_mod = ModuleType("api.db.services.document_service")
     document_service_mod.DocumentService = SimpleNamespace(
         query=lambda **_kwargs: [],
+        accessible=lambda *_args, **_kwargs: True,
         filter_update=lambda *_args, **_kwargs: 0,
         get_by_id=lambda *_args, **_kwargs: (False, None),
         update_by_id=lambda *_args, **_kwargs: True,
         decrement_chunk_num=lambda *_args, **_kwargs: None,
+        increment_chunk_num=lambda *_args, **_kwargs: True,
         get_embd_id=lambda *_args, **_kwargs: "",
         get_tenant_embd_id=lambda *_args, **_kwargs: None,
     )
@@ -233,13 +289,38 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"):
     monkeypatch.setitem(sys.modules, "api.db.services.knowledgebase_service", knowledgebase_service_mod)
 
     task_service_mod = ModuleType("api.db.services.task_service")
-    task_service_mod.TaskService = SimpleNamespace(filter_delete=lambda *_args, **_kwargs: None)
+    task_service_mod.TaskService = SimpleNamespace(filter_delete=lambda *_args, **_kwargs: None, query=lambda **_kwargs: [])
     task_service_mod.cancel_all_task_of = lambda *_args, **_kwargs: None
     task_service_mod.queue_tasks = lambda *_args, **_kwargs: None
     monkeypatch.setitem(sys.modules, "api.db.services.task_service", task_service_mod)
 
+    file_service_mod = ModuleType("api.db.services.file_service")
+    file_service_mod.FileService = SimpleNamespace()
+    monkeypatch.setitem(sys.modules, "api.db.services.file_service", file_service_mod)
+
+    canvas_service_mod = ModuleType("api.db.services.canvas_service")
+    canvas_service_mod.UserCanvasService = SimpleNamespace()
+    monkeypatch.setitem(sys.modules, "api.db.services.canvas_service", canvas_service_mod)
+
+    # document_api imports check_kb_team_permission; stub it so the real module
+    # and its user_service/common_service (peewee) import chain stay out.
+    check_team_permission_mod = ModuleType("api.common.check_team_permission")
+    check_team_permission_mod.check_kb_team_permission = lambda *_args, **_kwargs: True
+    check_team_permission_mod.check_file_team_permission = lambda *_args, **_kwargs: True
+    monkeypatch.setitem(sys.modules, "api.common.check_team_permission", check_team_permission_mod)
+
     api_utils_mod = ModuleType("api.utils.api_utils")
-    api_utils_mod.add_tenant_id_to_kwargs = lambda func: func
+
+    def _add_tenant_id_to_kwargs(func):
+        # Mirror the real decorator's functools.wraps so tests can reach the raw
+        # handler through ``route.__wrapped__`` and pass ``tenant_id`` explicitly.
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    api_utils_mod.add_tenant_id_to_kwargs = _add_tenant_id_to_kwargs
     api_utils_mod.check_duplicate_ids = lambda ids, _kind="item": (ids, [])
     api_utils_mod.construct_json_result = lambda code=0, message="success", data=None: {"code": code, "message": message, "data": data}
     api_utils_mod.get_error_data_result = lambda message="Sorry! Data missing!", code=102: {"code": code, "message": message}
@@ -248,6 +329,9 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"):
         key: value for key, value in {"code": code, "message": message, "data": data, "total": total}.items() if value is not None
     }
     api_utils_mod.server_error_response = lambda e: {"code": 500, "message": str(e)}
+    api_utils_mod.get_data_error_result = lambda message="Sorry! Data missing!", code=102: {"code": code, "message": message}
+    api_utils_mod.get_error_argument_result = lambda message="": {"code": 101, "message": message}
+    api_utils_mod.get_json_result = lambda code=0, message="success", data=None, **_kwargs: {"code": code, "message": message, "data": data}
     monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod)
 
     image_utils_mod = ModuleType("api.utils.image_utils")
@@ -278,6 +362,7 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"):
     common_metadata_utils_mod = ModuleType("common.metadata_utils")
     common_metadata_utils_mod.convert_conditions = lambda conditions: conditions
     common_metadata_utils_mod.meta_filter = lambda *_args, **_kwargs: []
+    common_metadata_utils_mod.turn2jsonschema = lambda *_args, **_kwargs: {}
     monkeypatch.setitem(sys.modules, "common.metadata_utils", common_metadata_utils_mod)
 
     rag_app_tag_mod = ModuleType("rag.app.tag")
@@ -470,10 +555,19 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"):
         # Return mock tenant with default model configurations
         return _MockModelConfig2(tenant_id, "chat-model").to_dict()
 
+    def _split_model_name(model_name):
+        parts = model_name.rsplit("@", 2)
+        if len(parts) == 3:
+            return parts[0], parts[1], parts[2]
+        if len(parts) == 2:
+            return parts[0], "default", parts[1]
+        return parts[0], "", ""
+
     tenant_model_service_mod.get_model_config_by_id = _get_model_config_by_id
     tenant_model_service_mod.get_model_config_from_provider_instance = _get_model_config_from_provider_instance
     tenant_model_service_mod.resolve_model_config = _get_model_config_from_provider_instance
     tenant_model_service_mod.get_tenant_default_model_by_type = _get_tenant_default_model_by_type
+    tenant_model_service_mod.split_model_name = _split_model_name
     monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod)
 
     if module_basename == "document_api":
@@ -493,6 +587,29 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"):
         document_api_service_mod.update_document_status_only = lambda *_args, **_kwargs: None
         document_api_service_mod.reset_document_for_reparse = lambda *_args, **_kwargs: None
         monkeypatch.setitem(sys.modules, "api.apps.services.document_api_service", document_api_service_mod)
+    else:
+        # chunk_api imports structure_graph_common from api.apps.services; stub
+        # it in sys.modules so the real module and its transitive imports are
+        # not loaded.
+        stub_apps_services = ModuleType("api.apps.services")
+        monkeypatch.setitem(sys.modules, "api.apps.services", stub_apps_services)
+
+        sgc_mod = ModuleType("api.apps.services.structure_graph_common")
+
+        async def _sgc_keyword_subgraph(*_args, **_kwargs):
+            return {}, [], []
+
+        async def _sgc_build_bucket(*_args, **_kwargs):
+            return [], []
+
+        sgc_mod.keyword_subgraph = _sgc_keyword_subgraph
+        sgc_mod.build_bucket = _sgc_build_bucket
+        monkeypatch.setitem(sys.modules, "api.apps.services.structure_graph_common", sgc_mod)
+
+    # document_counter_service is a real module (release_reparse_counters); evict
+    # any cached copy so it re-imports against this test's freshly-stubbed
+    # db_models / document_service rather than a prior test's stubs.
+    monkeypatch.delitem(sys.modules, "api.db.services.document_counter_service", raising=False)
 
     module_path = repo_root / "api" / "apps" / "restful_apis" / f"{module_basename}.py"
     spec = importlib.util.spec_from_file_location("test_doc_sdk_routes_unit", module_path)
@@ -552,80 +669,63 @@ class TestDocRoutesUnit:
         assert "length of 5" in str(exc_info.value)
 
     def test_download_and_download_doc_errors(self, monkeypatch):
-        module = _load_doc_module(monkeypatch)
+        module = _load_doc_module(monkeypatch, module_basename="document_api")
         _patch_send_file(monkeypatch, module)
         _patch_storage(monkeypatch, module, file_stream=b"")
-        res = _run(module.download.__wrapped__("tenant-1", "ds-1", ""))
-        assert res["message"] == "Specify document_id please."
-        monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [])
-        res = _run(module.download.__wrapped__("tenant-1", "ds-1", "doc-1"))
-        assert "do not own the dataset" in res["message"]
 
-        monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [1])
+        # download(dataset_id, document_id)
+        res = _run(module.download("ds-1", ""))
+        assert res["message"] == "Specify document_id please."
+
+        monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False)
+        res = _run(module.download("ds-1", "doc-1"))
+        assert res["message"] == "Document not found!"
+
+        monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True)
+        monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: False)
+        res = _run(module.download("ds-1", "doc-1"))
+        assert res["message"] == "Document not found!"
+
+        monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: True)
         monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [])
-        res = _run(module.download.__wrapped__("tenant-1", "ds-1", "doc-1"))
+        res = _run(module.download("ds-1", "doc-1"))
         assert "not own the document" in res["message"]
 
         monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()])
         monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n"))
-        res = _run(module.download.__wrapped__("tenant-1", "ds-1", "doc-1"))
+        res = _run(module.download("ds-1", "doc-1"))
         assert res["message"] == "This file is empty."
 
-        monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer"}))
-        res = _run(module.download_doc("doc-1"))
-        assert "Authorization is not valid" in res["message"]
-
-        monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer token"}))
-        monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [])
-        res = _run(module.download_doc("doc-1"))
-        assert "API key is invalid" in res["message"]
-
-        monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1"), SimpleNamespace(tenant_id="tenant-2")])
-        res = _run(module.download_doc("doc-1"))
-        assert "API key configuration is ambiguous" in res["message"]
-
-        monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")])
-        res = _run(module.download_doc(""))
+        # download_document(document_id)
+        res = _run(module.download_document(""))
         assert res["message"] == "Specify document_id please."
 
+        monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: False)
+        res = _run(module.download_document("doc-1"))
+        assert res["message"] == "Document not found!"
+
+        monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: True)
         monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [])
-        res = _run(module.download_doc("doc-1"))
+        res = _run(module.download_document("doc-1"))
         assert "not own the document" in res["message"]
 
         monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()])
-        kb_query_calls = []
-
-        def _deny_kb_query(**kwargs):
-            kb_query_calls.append(kwargs)
-            return []
-
-        monkeypatch.setattr(module.KnowledgebaseService, "query", _deny_kb_query)
-        monkeypatch.setattr(
-            module.File2DocumentService,
-            "get_storage_address",
-            lambda **_kwargs: (_ for _ in ()).throw(AssertionError("storage lookup must not run before tenant authorization")),
-        )
-        res = _run(module.download_doc("doc-1"))
-        assert res["message"] == "You do not have access to this document."
-        assert kb_query_calls == [{"id": "kb-1", "tenant_id": "tenant-1"}]
-
-        monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [1])
-        monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n"))
-        _patch_storage(monkeypatch, module, file_stream=b"")
-        res = _run(module.download_doc("doc-1"))
+        res = _run(module.download_document("doc-1"))
         assert res["message"] == "This file is empty."
 
         _patch_storage(monkeypatch, module, file_stream=b"abc")
-        res = _run(module.download_doc("doc-1"))
+        res = _run(module.download_document("doc-1"))
         assert res["filename"] == "doc.txt"
 
     def test_download_mimetype_from_filename(self, monkeypatch):
         module = _load_doc_module(monkeypatch, module_basename="document_api")
         _patch_send_file(monkeypatch, module)
         _patch_storage(monkeypatch, module, file_stream=b"pdf-bytes")
+        monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True)
+        monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: True)
         monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc(name="report.pdf", doc_type=FileType.PDF)])
         monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n"))
-        res = _run(module.download.__wrapped__("ds-1", "doc-1"))
+        res = _run(module.download("ds-1", "doc-1"))
         assert res["filename"] == "report.pdf"
         assert res["mimetype"] == "application/pdf"
 
@@ -636,6 +736,7 @@ class TestDocRoutesUnit:
         assert "don't own the dataset" in res["message"]
 
         monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True)
+        monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", pipeline_id=None)))
         monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]}))
         monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, []))
         toggle_doc = _ToggleBoolDocList(_DummyDoc(progress=0))
@@ -670,6 +771,58 @@ class TestDocRoutesUnit:
         assert res["code"] == module.RetCode.DATA_ERROR
         assert "Duplicate document ids" in res["message"]
 
+    def test_parse_and_stop_decrement_kb_counters(self, monkeypatch):
+        # Both routes delete the document's chunks, so the knowledgebase aggregate
+        # must drop by the document's current counters. Zeroing only the document
+        # row leaves that amount stranded in the KB total on every re-parse.
+        module = _load_doc_module(monkeypatch)
+        # _release_doc_counters re-reads the row under lock; mirror the document's
+        # current counters so the decrement is driven by that fresh read.
+        monkeypatch.setattr(module.Document, "fresh_doc", SimpleNamespace(id="doc-1", kb_id="kb-1", token_num=70, chunk_num=7, process_duration=1.5))
+        monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True)
+        monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", pipeline_id=None)))
+        monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]}))
+        monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, []))
+        monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _id: (True, _DummyDoc()))
+        monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n"))
+        monkeypatch.setattr(module.TaskService, "filter_delete", lambda *_args, **_kwargs: None)
+        monkeypatch.setattr(module, "queue_tasks", lambda *_args, **_kwargs: None)
+        monkeypatch.setattr(module, "cancel_all_task_of", lambda *_args, **_kwargs: None)
+        _patch_docstore(monkeypatch, module, delete=lambda *_args, **_kwargs: None)
+
+        decrements = []
+        updates = []
+        update_by_id_payloads = []
+
+        def _capture_filter_update(_conditions, info):
+            updates.append(info)
+            return 1
+
+        monkeypatch.setattr(module.DocumentService, "update_by_id", lambda _id, info: update_by_id_payloads.append(info) or True)
+
+        monkeypatch.setattr(module.DocumentService, "increment_chunk_num", lambda *args: decrements.append(args))
+        monkeypatch.setattr(module.DocumentService, "filter_update", _capture_filter_update)
+
+        monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc(token_num=70, chunk_num=7, process_duration=1.5)])
+        assert _run(module.parse.__wrapped__("tenant-1", "ds-1"))["code"] == 0
+        assert decrements == [("doc-1", "kb-1", -70, -7, -1.5)]
+        # The document update must not zero the counters itself; that is exactly
+        # what strands the difference in the KB aggregate.
+        assert "chunk_num" not in updates[0]
+        assert "token_num" not in updates[0]
+
+        decrements.clear()
+        monkeypatch.setattr(
+            module.DocumentService,
+            "query",
+            lambda **_kwargs: [_DummyDoc(run=module.TaskStatus.RUNNING.value, token_num=70, chunk_num=7, process_duration=1.5)],
+        )
+        assert _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1"))["code"] == 0
+        assert decrements == [("doc-1", "kb-1", -70, -7, -1.5)]
+        # stop_parsing must not zero the counters in its own document update either.
+        assert "chunk_num" not in update_by_id_payloads[0]
+        assert "token_num" not in update_by_id_payloads[0]
+
     def test_stop_parsing_branches(self, monkeypatch):
         module = _load_doc_module(monkeypatch)
         monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False)
@@ -677,6 +830,7 @@ class TestDocRoutesUnit:
         assert "don't own the dataset" in res["message"]
 
         monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True)
+        monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", pipeline_id=None)))
         monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({}))
         res = _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1"))
         assert "`document_ids` is required" in res["message"]
@@ -733,7 +887,7 @@ class TestDocRoutesUnit:
         monkeypatch.setattr(
             module.KnowledgebaseService,
             "get_by_id",
-            lambda _id: (True, SimpleNamespace(tenant_id=owner_tenant)),
+            lambda _id: (True, SimpleNamespace(tenant_id=owner_tenant, pipeline_id=None)),
         )
         monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]}))
         monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, []))
@@ -769,7 +923,7 @@ class TestDocRoutesUnit:
         monkeypatch.setattr(
             module.KnowledgebaseService,
             "get_by_id",
-            lambda _id: (True, SimpleNamespace(tenant_id=owner_tenant)),
+            lambda _id: (True, SimpleNamespace(tenant_id=owner_tenant, pipeline_id=None)),
         )
         monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]}))
         monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, []))
@@ -796,6 +950,7 @@ class TestDocRoutesUnit:
         module = _load_doc_module(monkeypatch, module_basename="document_api")
         updated = []
         deleted = []
+        decrements = []
 
         monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True)
         monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]}))
@@ -808,6 +963,14 @@ class TestDocRoutesUnit:
         )
         monkeypatch.setattr(module.TaskService, "query", lambda **_kwargs: [SimpleNamespace(progress=0.5)])
         monkeypatch.setattr(module, "cancel_all_task_of", lambda *_args, **_kwargs: None)
+        # release_reparse_counters re-reads the row under lock and decrements the
+        # KB by the document's partial counts before chunk_num is zeroed below.
+        monkeypatch.setattr(
+            sys.modules["api.db.db_models"].Document,
+            "fresh_doc",
+            SimpleNamespace(id="doc-1", kb_id="kb-1", token_num=70, chunk_num=7, process_duration=1.5),
+        )
+        monkeypatch.setattr(module.DocumentService, "increment_chunk_num", lambda *args: decrements.append(args))
         monkeypatch.setattr(
             module.DocumentService,
             "update_by_id",
@@ -824,16 +987,19 @@ class TestDocRoutesUnit:
 
         assert res["code"] == 0
         assert res["data"]["success_count"] == 1
-        assert updated == [
-            (
-                "doc-1",
-                {
-                    "run": module.TaskStatus.CANCEL.value,
-                    "progress": 0,
-                    "chunk_num": 0,
-                },
-            )
-        ]
+        assert len(updated) == 1
+        updated_doc_id, updated_info = updated[0]
+        assert updated_doc_id == "doc-1"
+        assert updated_info["run"] == str(module.TaskStatus.CANCEL.value)
+        assert updated_info["progress"] == 0
+        # release_reparse_counters is the sole counter adjustment; the status
+        # update must not set chunk_num itself (that would strand KB counts).
+        assert "chunk_num" not in updated_info
+        # progress_msg carries a timestamped cancellation marker, so match loosely.
+        assert "Task stopped by user." in updated_info["progress_msg"]
+        # The partial chunk/token counts are released from the KB aggregate via the
+        # row re-read under lock.
+        assert decrements == [("doc-1", "kb-1", -70, -7, -1.5)]
         assert deleted == [({"doc_id": "doc-1"}, module.search.index_name("tenant-1"), "kb-1")]
 
         deleted.clear()

test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py

@@ -15,6 +15,7 @@
 #
 
 import asyncio
+import contextlib
 import inspect
 import importlib.util
 import sys
@@ -223,6 +224,14 @@ def _load_chunk_module(monkeypatch):
     constants_mod.LLMType = _DummyLLMType
     constants_mod.ParserType = _DummyParserType
     constants_mod.PAGERANK_FLD = "pagerank_flt"
+    constants_mod.TaskStatus = SimpleNamespace(
+        UNSTART=SimpleNamespace(value="0"),
+        RUNNING=SimpleNamespace(value="1"),
+        CANCEL=SimpleNamespace(value="2"),
+        DONE=SimpleNamespace(value="3"),
+        FAIL=SimpleNamespace(value="4"),
+        SCHEDULE=SimpleNamespace(value="5"),
+    )
     monkeypatch.setitem(sys.modules, "common.constants", constants_mod)
 
     string_utils_mod = ModuleType("common.string_utils")
@@ -232,8 +241,28 @@ def _load_chunk_module(monkeypatch):
 
     metadata_utils_mod = ModuleType("common.metadata_utils")
     metadata_utils_mod.apply_meta_data_filter = lambda *_args, **_kwargs: {}
+    metadata_utils_mod.convert_conditions = lambda *_args, **_kwargs: {}
+    metadata_utils_mod.meta_filter = lambda *_args, **_kwargs: {}
     monkeypatch.setitem(sys.modules, "common.metadata_utils", metadata_utils_mod)
 
+    doc_store_base_mod = ModuleType("common.doc_store.doc_store_base")
+    doc_store_base_mod.OrderByExpr = type("OrderByExpr", (), {})
+    monkeypatch.setitem(sys.modules, "common.doc_store", ModuleType("common.doc_store"))
+    monkeypatch.setitem(sys.modules, "common.doc_store.doc_store_base", doc_store_base_mod)
+
+    tag_feature_utils_mod = ModuleType("common.tag_feature_utils")
+    tag_feature_utils_mod.validate_tag_features = lambda *_args, **_kwargs: None
+    monkeypatch.setitem(sys.modules, "common.tag_feature_utils", tag_feature_utils_mod)
+
+    pagination_utils_mod = ModuleType("api.utils.pagination_utils")
+    pagination_utils_mod.validate_rest_api_page_size = lambda *_args, **_kwargs: (1, 30)
+    monkeypatch.setitem(sys.modules, "api.utils.pagination_utils", pagination_utils_mod)
+
+    reference_metadata_utils_mod = ModuleType("api.utils.reference_metadata_utils")
+    reference_metadata_utils_mod.enrich_chunks_with_document_metadata = lambda chunks, *_args, **_kwargs: chunks
+    reference_metadata_utils_mod.resolve_reference_metadata_preferences = lambda *_args, **_kwargs: {}
+    monkeypatch.setitem(sys.modules, "api.utils.reference_metadata_utils", reference_metadata_utils_mod)
+
     misc_utils_mod = ModuleType("common.misc_utils")
 
     async def _thread_pool_exec(func):
@@ -293,6 +322,7 @@ def _load_chunk_module(monkeypatch):
     api_utils_mod.add_tenant_id_to_kwargs = lambda func: func
     api_utils_mod.check_duplicate_ids = lambda ids, _kind: (list(dict.fromkeys(ids)), [] if len(ids) == len(set(ids)) else [f"Duplicate {_kind} ids"])
     api_utils_mod.get_request_json = lambda: _AwaitableValue({})
+    api_utils_mod.construct_json_result = lambda code=0, message="success", data=None: {"code": code, "message": message, "data": data}
     monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod)
 
     image_utils_mod = ModuleType("api.utils.image_utils")
@@ -312,8 +342,94 @@ def _load_chunk_module(monkeypatch):
     tenant_model_service_mod.get_model_config_from_provider_instance = lambda *_args, **_kwargs: {"llm_name": "embed", "model_type": "embedding"}
     tenant_model_service_mod.resolve_model_config = lambda *_args, **_kwargs: {"llm_name": "embed", "model_type": "embedding"}
     tenant_model_service_mod.get_tenant_default_model_by_type = lambda *_args, **_kwargs: {"llm_name": "chat", "model_type": "chat"}
+    tenant_model_service_mod.split_model_name = lambda model_name: (model_name.rsplit("@", 2) + ["", ""])[:3]
     monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod)
 
+    # chunk_api imports structure_graph_common from api.apps.services; stub it in
+    # sys.modules so the real module and its transitive imports are not loaded.
+    stub_apps_services = ModuleType("api.apps.services")
+    monkeypatch.setitem(sys.modules, "api.apps.services", stub_apps_services)
+
+    sgc_mod = ModuleType("api.apps.services.structure_graph_common")
+
+    async def _sgc_keyword_subgraph(*_args, **_kwargs):
+        return {}, [], []
+
+    async def _sgc_build_bucket(*_args, **_kwargs):
+        return [], []
+
+    sgc_mod.keyword_subgraph = _sgc_keyword_subgraph
+    sgc_mod.build_bucket = _sgc_build_bucket
+    monkeypatch.setitem(sys.modules, "api.apps.services.structure_graph_common", sgc_mod)
+
+    # chunk_api imports DB, Document, Task from db_models; stub them so the real
+    # module (which pulls in quart_auth) is never loaded.
+    class _FakeExpr:
+        def __or__(self, other):
+            return self
+
+        def __and__(self, other):
+            return self
+
+    class _FakeField:
+        def __eq__(self, other):
+            return _FakeExpr()
+
+        def __ne__(self, other):
+            return _FakeExpr()
+
+        def is_null(self, value=True):
+            return _FakeExpr()
+
+    class _StubFreshDoc:
+        id = "doc-1"
+        kb_id = "kb-1"
+        token_num = 2
+        chunk_num = 1
+        process_duration = 0.0
+
+    class _StubDocQuery:
+        def where(self, *_args, **_kwargs):
+            return self
+
+        def for_update(self):
+            return self
+
+        def first(self):
+            return _StubDocumentModel.fresh_doc
+
+    class _StubDocumentModel:
+        id = _FakeField()
+        run = _FakeField()
+        fresh_doc = _StubFreshDoc()
+
+        @classmethod
+        def select(cls, *_args, **_kwargs):
+            return _StubDocQuery()
+
+    class _StubTaskModel:
+        doc_id = _FakeField()
+
+    db_models_mod = ModuleType("api.db.db_models")
+    db_models_mod.Document = _StubDocumentModel
+    db_models_mod.Task = _StubTaskModel
+    db_models_mod.DB = SimpleNamespace(atomic=lambda: contextlib.nullcontext())
+    monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod)
+
+    file2document_service_mod = ModuleType("api.db.services.file2document_service")
+    file2document_service_mod.File2DocumentService = SimpleNamespace(get_storage_address=lambda **_kwargs: ("", ""))
+    monkeypatch.setitem(sys.modules, "api.db.services.file2document_service", file2document_service_mod)
+
+    task_service_mod = ModuleType("api.db.services.task_service")
+    task_service_mod.TaskService = SimpleNamespace(filter_delete=lambda *_args, **_kwargs: None)
+    task_service_mod.cancel_all_task_of = lambda *_args, **_kwargs: None
+    task_service_mod.queue_tasks = lambda *_args, **_kwargs: None
+    monkeypatch.setitem(sys.modules, "api.db.services.task_service", task_service_mod)
+
+    document_counter_service_mod = ModuleType("api.db.services.document_counter_service")
+    document_counter_service_mod.release_reparse_counters = lambda *_args, **_kwargs: None
+    monkeypatch.setitem(sys.modules, "api.db.services.document_counter_service", document_counter_service_mod)
+
     document_service_mod = ModuleType("api.db.services.document_service")
 
     class _DocumentService:
@@ -499,6 +615,14 @@ def _load_chunk_api_module(monkeypatch):
     module.manager = _DummyManager()
     monkeypatch.setitem(sys.modules, module_name, module)
     spec.loader.exec_module(module)
+    # chunk_api imports these inside the chunk-write helpers; re-expose the shared
+    # stubs so tests can reach them via module.<name> (setattr / method patching).
+    module.rag_tokenizer = sys.modules["rag.nlp"].rag_tokenizer
+    module.beAdoc = sys.modules["rag.app.qa"].beAdoc
+    module.rmPrefix = sys.modules["rag.app.qa"].rmPrefix
+    module.label_question = sys.modules["rag.app.tag"].label_question
+    module.cross_languages = sys.modules["rag.prompts.generator"].cross_languages
+    module.keyword_extraction = sys.modules["rag.prompts.generator"].keyword_extraction
     return module