fix(api,billing): invalidate vector space cache after cleanup (#39404)

api/tests/unit_tests/events/event_handlers/test_clean_when_document_deleted.py

@@ -0,0 +1,15 @@
+from unittest.mock import patch
+
+from events.event_handlers.clean_when_document_deleted import handle
+
+
+def test_handler_dispatches_cleanup_task():
+    with patch("events.event_handlers.clean_when_document_deleted.clean_document_task.delay") as delay:
+        handle(
+            "document-1",
+            dataset_id="dataset-1",
+            doc_form="paragraph",
+            file_id="file-1",
+        )
+
+    delay.assert_called_once_with("document-1", "dataset-1", "paragraph", "file-1")

api/tests/unit_tests/services/test_billing_service.py

@@ -462,6 +462,27 @@ class TestBillingServiceSubscriptionInfo:
             params={"tenant_id": tenant_id},
         )
 
+    def test_get_vector_space_bypasses_cache(self, mock_send_request):
+        tenant_id = "tenant-123"
+        mock_send_request.return_value = {"size": 4096, "limit": 20480}
+
+        result = BillingService.get_vector_space(tenant_id, bypass_cache=True)
+
+        assert result == {"size": 4096, "limit": 20480}
+        mock_send_request.assert_called_once_with(
+            "GET",
+            "/subscription/vector-space",
+            params={"tenant_id": tenant_id, "bypass_cache": "true"},
+        )
+
+    def test_invalidate_vector_space_cache_bypasses_cache(self):
+        tenant_id = "tenant-123"
+
+        with patch.object(BillingService, "get_vector_space") as get_vector_space:
+            BillingService.invalidate_vector_space_cache(tenant_id)
+
+        get_vector_space.assert_called_once_with(tenant_id, bypass_cache=True)
+
     def test_quota_get_balance_uses_quota_request(self):
         tenant_id = "tenant-123"
         with patch.object(BillingService, "_send_quota_request") as mock_send_quota_request:

api/tests/unit_tests/tasks/test_batch_clean_document_task.py

@@ -0,0 +1,58 @@
+from unittest.mock import MagicMock, patch
+
+from tasks.batch_clean_document_task import batch_clean_document_task
+
+
+def _setup_cleanup_dependencies():
+    session = MagicMock()
+    segment = MagicMock(id="segment-1", index_node_id="node-1", content="content")
+    dataset = MagicMock(id="dataset-1", tenant_id="tenant-1")
+    session.scalars.return_value.all.return_value = [segment]
+    session.scalar.return_value = dataset
+
+    context_manager = MagicMock()
+    context_manager.__enter__.return_value = session
+    context_manager.__exit__.return_value = None
+    return session, context_manager
+
+
+def test_successful_vector_cleanup_schedules_billing_refresh():
+    _, context_manager = _setup_cleanup_dependencies()
+
+    with (
+        patch("tasks.batch_clean_document_task.session_factory.create_session", return_value=context_manager),
+        patch("tasks.batch_clean_document_task.get_image_upload_file_ids", return_value=[]),
+        patch("tasks.batch_clean_document_task.IndexProcessorFactory") as processor_factory,
+        patch("tasks.batch_clean_document_task.schedule_billing_vector_space_refresh") as schedule_refresh,
+    ):
+        batch_clean_document_task(
+            document_ids=["document-1"],
+            dataset_id="dataset-1",
+            doc_form="paragraph",
+            file_ids=[],
+        )
+
+    processor_factory.return_value.init_index_processor.return_value.clean.assert_called_once()
+    schedule_refresh.assert_called_once_with("tenant-1")
+
+
+def test_failed_vector_cleanup_does_not_schedule_billing_refresh():
+    _, context_manager = _setup_cleanup_dependencies()
+
+    with (
+        patch("tasks.batch_clean_document_task.session_factory.create_session", return_value=context_manager),
+        patch("tasks.batch_clean_document_task.get_image_upload_file_ids", return_value=[]),
+        patch("tasks.batch_clean_document_task.IndexProcessorFactory") as processor_factory,
+        patch("tasks.batch_clean_document_task.schedule_billing_vector_space_refresh") as schedule_refresh,
+    ):
+        processor_factory.return_value.init_index_processor.return_value.clean.side_effect = RuntimeError(
+            "vector cleanup failed"
+        )
+        batch_clean_document_task(
+            document_ids=["document-1"],
+            dataset_id="dataset-1",
+            doc_form="paragraph",
+            file_ids=[],
+        )
+
+    schedule_refresh.assert_not_called()

api/tests/unit_tests/tasks/test_clean_dataset_task.py

@@ -434,14 +434,15 @@ class TestIndexProcessorParameters:
         index_struct = '{"type": "paragraph"}'
 
         # Act
-        clean_dataset_task(
-            dataset_id=dataset_id,
-            tenant_id=tenant_id,
-            indexing_technique=indexing_technique,
-            index_struct=index_struct,
-            collection_binding_id=collection_binding_id,
-            doc_form=IndexStructureType.PARAGRAPH_INDEX,
-        )
+        with patch("tasks.clean_dataset_task.schedule_billing_vector_space_refresh") as schedule_refresh:
+            clean_dataset_task(
+                dataset_id=dataset_id,
+                tenant_id=tenant_id,
+                indexing_technique=indexing_technique,
+                index_struct=index_struct,
+                collection_binding_id=collection_binding_id,
+                doc_form=IndexStructureType.PARAGRAPH_INDEX,
+            )
 
         # Assert
         mock_index_processor_factory["processor"].clean.assert_called_once()
@@ -462,3 +463,28 @@ class TestIndexProcessorParameters:
         assert call_args[1]["session"] is mock_db_session.session
         assert call_args[1]["with_keywords"] is True
         assert call_args[1]["delete_child_chunks"] is True
+        schedule_refresh.assert_called_once_with(tenant_id)
+
+    def test_vector_cleanup_failure_does_not_schedule_billing_refresh(
+        self,
+        dataset_id: str,
+        tenant_id: str,
+        collection_binding_id: str,
+        mock_db_session,
+        mock_storage,
+        mock_index_processor_factory,
+        mock_get_image_upload_file_ids,
+    ):
+        mock_index_processor_factory["processor"].clean.side_effect = RuntimeError("vector cleanup failed")
+
+        with patch("tasks.clean_dataset_task.schedule_billing_vector_space_refresh") as schedule_refresh:
+            clean_dataset_task(
+                dataset_id=dataset_id,
+                tenant_id=tenant_id,
+                indexing_technique=IndexTechniqueType.HIGH_QUALITY,
+                index_struct='{"type": "paragraph"}',
+                collection_binding_id=collection_binding_id,
+                doc_form=IndexStructureType.PARAGRAPH_INDEX,
+            )
+
+        schedule_refresh.assert_not_called()

api/tests/unit_tests/tasks/test_clean_document_task.py

@@ -169,12 +169,13 @@ class TestVectorCleanupResilience:
         )
 
         # Act — must not raise out of the task even though clean() raises.
-        clean_document_task(
-            document_id=document_id,
-            dataset_id=dataset_id,
-            doc_form="paragraph",
-            file_id=None,
-        )
+        with patch("tasks.clean_document_task.schedule_billing_vector_space_refresh") as schedule_refresh:
+            clean_document_task(
+                document_id=document_id,
+                dataset_id=dataset_id,
+                doc_form="paragraph",
+                file_id=None,
+            )
 
         # Assert
         # 1. Vector cleanup was attempted.
@@ -187,6 +188,7 @@ class TestVectorCleanupResilience:
             "Step 3+ DB cleanup did not run after vector cleanup failure; "
             "this regression would re-introduce the orphan-segment bug."
         )
+        schedule_refresh.assert_not_called()
 
     def test_vector_cleanup_success_path_remains_unaffected(
         self,
@@ -229,12 +231,13 @@ class TestVectorCleanupResilience:
 
         mock_sf.create_session.side_effect = [cm1, cm2] + [_default_cm() for _ in range(10)]
 
-        clean_document_task(
-            document_id=document_id,
-            dataset_id=dataset_id,
-            doc_form="paragraph",
-            file_id=None,
-        )
+        with patch("tasks.clean_document_task.schedule_billing_vector_space_refresh") as schedule_refresh:
+            clean_document_task(
+                document_id=document_id,
+                dataset_id=dataset_id,
+                doc_form="paragraph",
+                file_id=None,
+            )
 
         assert mock_index_processor_factory["processor"].clean.call_count == 1
         # Index cleanup invoked with the expected delete_summaries / delete_child_chunks flags.
@@ -242,6 +245,7 @@ class TestVectorCleanupResilience:
         assert kwargs.get("with_keywords") is True
         assert kwargs.get("delete_child_chunks") is True
         assert kwargs.get("delete_summaries") is True
+        schedule_refresh.assert_called_once_with(tenant_id)
 
     def test_no_segments_skips_vector_cleanup(
         self,
@@ -279,13 +283,15 @@ class TestVectorCleanupResilience:
 
         mock_sf.create_session.side_effect = [cm1] + [_default_cm() for _ in range(10)]
 
-        clean_document_task(
-            document_id=document_id,
-            dataset_id=dataset_id,
-            doc_form="paragraph",
-            file_id=None,
-        )
+        with patch("tasks.clean_document_task.schedule_billing_vector_space_refresh") as schedule_refresh:
+            clean_document_task(
+                document_id=document_id,
+                dataset_id=dataset_id,
+                doc_form="paragraph",
+                file_id=None,
+            )
 
         # Vector cleanup is gated on ``index_node_ids``; when there are no
         # segments the IndexProcessorFactory path is never entered.
         mock_index_processor_factory["factory_cls"].assert_not_called()
+        schedule_refresh.assert_not_called()

api/tests/unit_tests/tasks/test_refresh_billing_vector_space_task.py

@@ -0,0 +1,45 @@
+from unittest.mock import patch
+
+import pytest
+
+from tasks.refresh_billing_vector_space_task import (
+    refresh_billing_vector_space_task,
+    schedule_billing_vector_space_refresh,
+)
+
+
+def test_refresh_invalidates_vector_space_cache():
+    with (
+        patch("tasks.refresh_billing_vector_space_task.dify_config.BILLING_ENABLED", True),
+        patch(
+            "tasks.refresh_billing_vector_space_task.BillingService.invalidate_vector_space_cache"
+        ) as invalidate_cache,
+    ):
+        refresh_billing_vector_space_task.run("tenant-1")
+
+    invalidate_cache.assert_called_once_with("tenant-1")
+
+
+def test_refresh_failure_schedules_retry():
+    error = RuntimeError("billing unavailable")
+
+    with (
+        patch("tasks.refresh_billing_vector_space_task.dify_config.BILLING_ENABLED", True),
+        patch(
+            "tasks.refresh_billing_vector_space_task.BillingService.invalidate_vector_space_cache",
+            side_effect=error,
+        ),
+        patch.object(refresh_billing_vector_space_task, "retry", side_effect=RuntimeError("retry scheduled")) as retry,
+        pytest.raises(RuntimeError, match="retry scheduled"),
+    ):
+        refresh_billing_vector_space_task.run("tenant-1")
+
+    retry.assert_called_once_with(exc=error, countdown=30)
+
+
+def test_dispatch_failure_does_not_propagate():
+    with (
+        patch("tasks.refresh_billing_vector_space_task.dify_config.BILLING_ENABLED", True),
+        patch.object(refresh_billing_vector_space_task, "delay", side_effect=RuntimeError("broker unavailable")),
+    ):
+        schedule_billing_vector_space_refresh("tenant-1")