fix: distinguish KnowledgeFS vector quota failures

api/controllers/inner_api/knowledge_fs/vectors.py

@@ -3,12 +3,13 @@
 from flask import request
 from flask_restx import Resource
 from pydantic import ValidationError
-from werkzeug.exceptions import BadRequest, Forbidden, ServiceUnavailable
+from werkzeug.exceptions import BadRequest, ServiceUnavailable
 
 from controllers.common.schema import register_response_schema_models, register_schema_models
 from controllers.inner_api import inner_api_ns
 from controllers.inner_api.wraps import knowledge_fs_inner_api_only
 from fields.base import ResponseModel
+from libs.exception import BaseHTTPException
 from libs.helper import dump_response
 from services.knowledge_fs.vector_store import (
     VectorPoint,
@@ -18,7 +19,17 @@ from services.knowledge_fs.vector_store import (
     execute_vector_request,
 )
 from services.tidb_binding_service import TidbBindingPendingError
-from services.vector_space_admission_service import VectorSpaceAdmissionError
+from services.vector_space_admission_service import VectorSpaceAdmissionError, VectorSpaceAdmissionUnavailableError
+
+
+class KnowledgeFSVectorQuotaExceededError(BaseHTTPException):
+    error_code = "vector_space_quota_exceeded"
+    code = 413
+
+
+class KnowledgeFSVectorQuotaUnavailableError(BaseHTTPException):
+    error_code = "vector_space_quota_unavailable"
+    code = 503
 
 
 class VectorMatchResponse(ResponseModel):
@@ -54,8 +65,10 @@ class KnowledgeFSVectorApi(Resource):
                 result = execute_vector_request(client, payload, check_admission=False)
         except TidbBindingPendingError:
             raise ServiceUnavailable("KnowledgeFS vector backend is being provisioned", retry_after=5) from None
+        except VectorSpaceAdmissionUnavailableError as error:
+            raise KnowledgeFSVectorQuotaUnavailableError(str(error)) from None
         except VectorSpaceAdmissionError as error:
-            raise Forbidden(str(error)) from None
+            raise KnowledgeFSVectorQuotaExceededError(str(error)) from None
         except Exception:
             # Native SDK exceptions can contain endpoint credentials or vectors.
             raise ServiceUnavailable("KnowledgeFS vector backend is unavailable") from None

api/services/vector_space_admission_service.py

@@ -43,6 +43,10 @@ class VectorSpaceAdmissionError(ValueError):
         super().__init__(message)
 
 
+class VectorSpaceAdmissionUnavailableError(VectorSpaceAdmissionError):
+    """The usage or reservation dependency cannot safely authorize a write."""
+
+
 @dataclass(frozen=True)
 class VectorStorageWorkload:
     text_points: int
@@ -189,7 +193,11 @@ def _items(value: Any, name: str) -> list[Any]:
 
 
 class VectorSpaceAdmissionService:
-    """Cloud-only pre-write guard for unusually large TiDB vector workloads."""
+    """Cloud-only pre-write guard for unusually large TiDB vector workloads.
+
+    Configured estimate thresholds reject unusually large writes and are separate
+    from Billing plan limits. Billing supplies current usage and the displayed limit.
+    """
 
     def __init__(self) -> None:
         self._dimension_by_dataset: dict[str, int] = {}
@@ -350,10 +358,14 @@ class VectorSpaceAdmissionService:
     def _get_usage_and_limit_mb(self, tenant_id: str) -> tuple[float, int]:
         try:
             vector_space = BillingService.get_vector_space(tenant_id)
+            if vector_space.get("usage_unknown", False):
+                raise ValueError("Vector storage usage is unknown")
             current_usage_mb = float(vector_space["size"])
             plan_limit_mb = int(vector_space["limit"])
+            if not math.isfinite(current_usage_mb) or current_usage_mb < 0 or plan_limit_mb < -1:
+                raise ValueError("Invalid vector storage quota")
         except Exception as error:
-            raise VectorSpaceAdmissionError(
+            raise VectorSpaceAdmissionUnavailableError(
                 "Unable to verify vector storage usage right now. Please try again later."
             ) from error
         return current_usage_mb, plan_limit_mb
@@ -413,7 +425,7 @@ class VectorSpaceAdmissionService:
 
                 return base_usage_bytes, projected_usage_bytes
         except Exception as error:
-            raise VectorSpaceAdmissionError(
+            raise VectorSpaceAdmissionUnavailableError(
                 "Unable to reserve estimated vector storage right now. Please try again later."
             ) from error
 
@@ -423,7 +435,7 @@ class VectorSpaceAdmissionService:
         try:
             billing_info = BillingService.get_info(tenant_id, exclude_vector_space=True)
         except Exception as error:
-            raise VectorSpaceAdmissionError(
+            raise VectorSpaceAdmissionUnavailableError(
                 "Unable to verify the subscription plan right now. Please try again later."
             ) from error
 
@@ -459,7 +471,7 @@ class VectorSpaceAdmissionService:
 
         embeddings = CacheEmbedding(model_instance).embed_documents([probe_text])
         if not embeddings or not embeddings[0]:
-            raise VectorSpaceAdmissionError(
+            raise VectorSpaceAdmissionUnavailableError(
                 "Unable to estimate vector storage for this document. Please try again later."
             )
 

api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_vectors.py

@@ -3,10 +3,10 @@ from contextlib import nullcontext
 from unittest.mock import patch
 
 import pytest
-from werkzeug.exceptions import BadRequest, Forbidden, NotFound, ServiceUnavailable
+from werkzeug.exceptions import BadRequest, NotFound, ServiceUnavailable
 
-from controllers.inner_api.knowledge_fs.vectors import KnowledgeFSVectorApi
-from services.vector_space_admission_service import VectorSpaceAdmissionError
+from controllers.inner_api.knowledge_fs.vectors import KnowledgeFSVectorApi, KnowledgeFSVectorQuotaExceededError
+from services.vector_space_admission_service import VectorSpaceAdmissionError, VectorSpaceAdmissionUnavailableError
 
 PAYLOAD = {
     "operation": "get",
@@ -93,7 +93,7 @@ def test_vector_controller_reports_quota_failure(app):
             side_effect=VectorSpaceAdmissionError("vector quota exceeded"),
         ),
     ):
-        with pytest.raises(Forbidden, match="vector quota exceeded"):
+        with pytest.raises(KnowledgeFSVectorQuotaExceededError, match="vector quota exceeded"):
             inspect.unwrap(handler.post)(handler)
 
 
@@ -138,11 +138,45 @@ def test_quota_rejection_cannot_allocate_a_cluster(app):
         ),
         patch("controllers.inner_api.knowledge_fs.vectors.configured_vector_client") as client,
     ):
-        with pytest.raises(Forbidden):
+        with pytest.raises(KnowledgeFSVectorQuotaExceededError):
             inspect.unwrap(handler.post)(handler)
         client.assert_not_called()
 
 
+@pytest.mark.parametrize(
+    ("admission_error", "status_code", "error_code"),
+    [
+        (VectorSpaceAdmissionError("vector quota exceeded"), 413, "vector_space_quota_exceeded"),
+        (VectorSpaceAdmissionUnavailableError("vector usage is unavailable"), 503, "vector_space_quota_unavailable"),
+    ],
+)
+def test_http_quota_failure_cannot_allocate_a_cluster(config_overrides, admission_error, status_code, error_code):
+    from flask import Flask
+    from flask_restx import Api
+
+    config_overrides(PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key")
+    http_app = Flask(__name__)
+    Api(http_app).add_resource(KnowledgeFSVectorApi, "/vectors")
+    payload = {
+        "operation": "upsert",
+        "scope": PAYLOAD["scope"],
+        "points": [
+            {"id": PAYLOAD["ids"][0], "generation_id": PAYLOAD["ids"][0], "content_hash": "a" * 64, "vector": [1, 0]}
+        ],
+    }
+    with (
+        patch("controllers.inner_api.knowledge_fs.vectors.admit_vector_request", side_effect=admission_error),
+        patch("controllers.inner_api.knowledge_fs.vectors.configured_vector_client") as client,
+        patch("controllers.inner_api.knowledge_fs.vectors.execute_vector_request") as execute,
+    ):
+        response = http_app.test_client().post("/vectors", json=payload, headers={"X-Inner-Api-Key": "trusted-key"})
+    assert response.status_code == status_code
+    assert response.get_json() == {"code": error_code, "message": str(admission_error), "status": status_code}
+    assert "Retry-After" not in response.headers
+    client.assert_not_called()
+    execute.assert_not_called()
+
+
 def test_pending_cluster_returns_explicit_retry_after_without_secrets(app):
     from services.tidb_binding_service import TidbBindingPendingError
 

api/tests/unit_tests/services/test_knowledge_fs_vector_store.py

@@ -35,6 +35,11 @@ def request(operation, **kwargs):
     return VectorRequest.model_validate({"operation": operation, "scope": SCOPE, **kwargs})
 
 
+@pytest.fixture(autouse=True)
+def _vector_store_config(config_overrides):
+    config_overrides(DEPLOYMENT_EDITION="COMMUNITY")
+
+
 @pytest.fixture
 def client():
     client = QdrantClient(":memory:")
@@ -309,7 +314,8 @@ def test_cloud_quota_uses_existing_watermark_without_embedding_calls(config_over
     with (
         patch.object(service, "_get_plan", return_value=CloudPlan.SANDBOX),
         patch.object(service, "_get_usage_and_limit_mb", return_value=(0.0, 50)),
-        patch.object(service, "_reserve_projected_usage", return_value=(0, 1024)) as reserve,
+        patch.object(service, "_get_embedding_dimension") as dimension,
+        patch.object(service, "_reserve_projected_usage", return_value=(0, 60 * 1024 * 1024)) as reserve,
     ):
         for _ in range(2):
             service.ensure_external_points_can_be_indexed(
@@ -317,8 +323,10 @@ def test_cloud_quota_uses_existing_watermark_without_embedding_calls(config_over
             )
         assert reserve.call_args.kwargs["document_id"] == "knowledgefs:stable-batch"
         assert reserve.call_args.kwargs["document_estimate_bytes"] == 2 * (3 * 8 + 3584)
+        assert reserve.call_args.kwargs["estimate_limit_bytes"] == 60 * 1024 * 1024
+        dimension.assert_not_called()
         reserve.return_value = (0, 61 * 1024 * 1024)
-        with pytest.raises(VectorSpaceAdmissionError, match="exceeding"):
+        with pytest.raises(VectorSpaceAdmissionError, match="exceeding the 50 MB limit"):
             service.ensure_external_points_can_be_indexed(
                 tenant_id=TENANT, batch_id="large", point_count=2, dimension=3
             )
@@ -343,6 +351,48 @@ def test_unmetered_workspace_does_not_contact_billing_or_reserve_quota(config_ov
         reserve.assert_not_called()
 
 
+@pytest.mark.parametrize("limit", [50, -1])
+def test_external_points_with_unknown_usage_fail_before_reservation(config_overrides, limit):
+    from enums import CloudPlan
+    from services.vector_space_admission_service import (
+        VectorSpaceAdmissionService,
+        VectorSpaceAdmissionUnavailableError,
+    )
+
+    config_overrides(DEPLOYMENT_EDITION="CLOUD", VECTOR_STORE="tidb_on_qdrant")
+    service = VectorSpaceAdmissionService()
+    with (
+        patch.object(service, "_get_plan", return_value=CloudPlan.SANDBOX),
+        patch(
+            "services.vector_space_admission_service.BillingService.get_vector_space",
+            return_value={"size": 0, "limit": limit, "usage_unknown": True},
+        ),
+        patch.object(service, "_reserve_projected_usage") as reserve,
+    ):
+        with pytest.raises(VectorSpaceAdmissionUnavailableError, match="Unable to verify vector storage usage"):
+            service.ensure_external_points_can_be_indexed(
+                tenant_id=TENANT, batch_id="unknown", point_count=2, dimension=3
+            )
+        reserve.assert_not_called()
+
+
+@pytest.mark.parametrize(("edition", "vector_store"), [("COMMUNITY", "tidb_on_qdrant"), ("CLOUD", "qdrant")])
+def test_external_points_skip_billing_outside_cloud_tidb(config_overrides, edition, vector_store):
+    from services.vector_space_admission_service import VectorSpaceAdmissionService
+
+    config_overrides(DEPLOYMENT_EDITION=edition, VECTOR_STORE=vector_store)
+    service = VectorSpaceAdmissionService()
+    with (
+        patch("services.vector_space_admission_service.BillingService.get_vector_space") as billing,
+        patch.object(service, "_reserve_projected_usage") as reserve,
+    ):
+        service.ensure_external_points_can_be_indexed(
+            tenant_id=TENANT, batch_id="unmetered", point_count=2, dimension=3
+        )
+        billing.assert_not_called()
+        reserve.assert_not_called()
+
+
 def test_incompatible_collection_is_never_recreated_or_written(client):
     collection = VectorScope.model_validate(SCOPE).collection_name
     client.delete_collection(collection)

api/tests/unit_tests/services/test_vector_space_admission_service.py

@@ -18,6 +18,7 @@ from services.vector_space_admission_service import (
     VECTOR_SPACE_ADMISSION_ERROR_CODE,
     VectorSpaceAdmissionError,
     VectorSpaceAdmissionService,
+    VectorSpaceAdmissionUnavailableError,
     VectorStorageWorkload,
     build_document_workload,
     build_pipeline_workload,
@@ -567,3 +568,84 @@ def test_unknown_billing_plan_is_not_logged(caplog: pytest.LogCaptureFixture) ->
 
     assert "unknown plan tenant_id=tenant-1" in caplog.text
     assert upstream_plan not in caplog.text
+
+
+@pytest.mark.parametrize(
+    "quota",
+    [
+        {"size": 0, "limit": 50, "usage_unknown": True},
+        {"size": 0, "limit": -1, "usage_unknown": True},
+        {"size": -0.1, "limit": 50},
+        {"size": float("nan"), "limit": 50},
+        {"size": float("inf"), "limit": 50},
+        {"size": float("-inf"), "limit": 50},
+        {"size": "invalid", "limit": 50},
+        {"size": 0, "limit": -2},
+        {"size": 0, "limit": "invalid"},
+        {"limit": 50},
+        {"size": 0},
+    ],
+)
+def test_unavailable_billing_quota_fails_before_embedding_or_reservation(
+    sqlite_session: Session, quota: dict[str, object]
+) -> None:
+    service = VectorSpaceAdmissionService()
+    with (
+        patch.object(service, "_get_plan", return_value=CloudPlan.SANDBOX),
+        patch(
+            "services.vector_space_admission_service.Vector.resolve_vector_type",
+            return_value=VectorType.TIDB_ON_QDRANT,
+        ),
+        patch("services.vector_space_admission_service.BillingService.get_vector_space", return_value=quota),
+        patch.object(service, "_get_embedding_dimension") as dimension,
+        patch.object(service, "_reserve_projected_usage") as reserve,
+    ):
+        with pytest.raises(VectorSpaceAdmissionUnavailableError, match="Unable to verify vector storage usage"):
+            service._ensure_can_write(
+                dataset=_dataset(sqlite_session),
+                document_id="document-1",
+                workload=_workload(),
+                session=sqlite_session,
+            )
+
+    dimension.assert_not_called()
+    reserve.assert_not_called()
+
+
+def test_billing_request_failure_is_unavailable() -> None:
+    with patch(
+        "services.vector_space_admission_service.BillingService.get_vector_space", side_effect=RuntimeError("offline")
+    ):
+        with pytest.raises(VectorSpaceAdmissionUnavailableError, match="Unable to verify vector storage usage"):
+            VectorSpaceAdmissionService()._get_usage_and_limit_mb("tenant-1")
+
+
+def test_watermark_failure_is_unavailable() -> None:
+    with patch("services.vector_space_admission_service.redis_client.lock", side_effect=RuntimeError("offline")):
+        with pytest.raises(VectorSpaceAdmissionUnavailableError, match="Unable to reserve estimated vector storage"):
+            VectorSpaceAdmissionService()._reserve_projected_usage(
+                tenant_id="tenant-1",
+                document_id="document-1",
+                current_usage_bytes=0,
+                document_estimate_bytes=1,
+                estimate_limit_bytes=50 * _MEBIBYTE,
+            )
+
+
+def test_watermark_is_isolated_by_tenant() -> None:
+    redis = _FakeRedis()
+    with patch("services.vector_space_admission_service.redis_client", redis):
+        for tenant_id in ("tenant-1", "tenant-2"):
+            assert VectorSpaceAdmissionService()._reserve_projected_usage(
+                tenant_id=tenant_id,
+                document_id="same-document",
+                current_usage_bytes=40 * _MEBIBYTE,
+                document_estimate_bytes=10 * _MEBIBYTE,
+                estimate_limit_bytes=50 * _MEBIBYTE,
+            ) == (40 * _MEBIBYTE, 50 * _MEBIBYTE)
+
+    assert len(redis.values) == 2
+    for tenant_id in ("tenant-1", "tenant-2"):
+        state = json.loads(redis.values[f"tenant:{tenant_id}:vector_space_estimate_watermark"])
+        assert state["projected_usage_bytes"] == 50 * _MEBIBYTE
+        assert state["document_ids"] == ["same-document"]