refactor(api): read billing subscription and usage through the generated client (#45131)

Co-authored-by: hjlarry <hjlarry@163.com>

api/commands/retention.py

@@ -140,7 +140,7 @@ def _filter_paid_workflow_archive_tenant_ids(tenant_ids: list[str]) -> tuple[lis
     paid_tenant_ids = [
         tenant_id
         for tenant_id in tenant_ids
-        if plans.get(tenant_id) and plans[tenant_id].get("plan") in (CloudPlan.PROFESSIONAL, CloudPlan.TEAM)
+        if (plan := plans.get(tenant_id)) and plan.plan in (CloudPlan.PROFESSIONAL, CloudPlan.TEAM)
     ]
     unpaid_tenant_ids = sorted(set(tenant_ids) - set(paid_tenant_ids))
     return paid_tenant_ids, unpaid_tenant_ids

api/controllers/console/workspace/model_providers.py

@@ -400,10 +400,10 @@ class ModelProviderPaymentCheckoutUrlApi(Resource):
         if provider != "anthropic":
             raise ValueError(f"provider name {provider} is invalid")
         # pyrefly: ignore [deprecated]
-        data = BillingService.get_model_provider_payment_link(
+        link = BillingService.get_model_provider_payment_link(
             provider_name=provider,
             tenant_id=current_tenant_id,
             account_id=current_user.id,
             prefilled_email=current_user.email,
         )
-        return dump_response(ModelProviderPaymentCheckoutUrlResponse, data)
+        return dump_response(ModelProviderPaymentCheckoutUrlResponse, {"payment_link": link.url})

api/controllers/console/wraps.py

@@ -5,7 +5,6 @@ from collections.abc import Callable
 from functools import wraps
 from typing import Any, Concatenate, Protocol, cast, overload
 
-import httpx
 from flask import abort, request
 from pydantic import BaseModel, ValidationError
 from sqlalchemy import select
@@ -26,8 +25,9 @@ from models import Account
 from models.account import AccountStatus
 from models.dataset import RateLimitLog
 from models.model import DifySetup
-from services.billing_service import BillingService, _BillingHTTPStatusError
+from services.billing_service import BillingService
 from services.entities.feature_entities import LicenseStatus
+from services.errors.billing import BillingError
 from services.feature_service import FeatureService
 from services.operation_service import OperationService, UtmInfo
 from services.system_feature_service import SystemFeatureService
@@ -168,16 +168,10 @@ def is_cloud_edition_billing_paid_plan(tenant_id: str) -> bool:
     """
 
     try:
-        billing_info = BillingService.get_info(tenant_id, exclude_vector_space=True)
-    except (
-        _BillingHTTPStatusError,
-        httpx.RequestError,
-        json.JSONDecodeError,
-        UnicodeDecodeError,
-        ValidationError,
-    ) as exc:
+        subscription = BillingService.get_info(tenant_id, exclude_vector_space=True).subscription
+    except BillingError as exc:
         raise ServiceUnavailable("Billing entitlement is temporarily unavailable.") from exc
-    return billing_info["subscription"]["plan"] in (
+    return subscription.plan in (
         CloudPlan.PROFESSIONAL,
         CloudPlan.TEAM,
     )

api/controllers/service_api/app/completion.py

@@ -420,7 +420,7 @@ class ChatApi(Resource):
             and dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD
         ):
             billing_info = BillingService.get_info(app_model.tenant_id, exclude_vector_space=True)
-            if billing_info["subscription"]["plan"] == CloudPlan.SANDBOX:
+            if billing_info.subscription.plan == CloudPlan.SANDBOX:
                 raise WorkflowVersionExecutionNotAllowedError()
 
         external_trace_id = get_external_trace_id(request)

api/controllers/service_api/app/workflow.py

@@ -494,7 +494,7 @@ class WorkflowRunByIdApi(Resource):
 
         if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD:
             billing_info = BillingService.get_info(app_model.tenant_id, exclude_vector_space=True)
-            if billing_info["subscription"]["plan"] == CloudPlan.SANDBOX:
+            if billing_info.subscription.plan == CloudPlan.SANDBOX:
                 raise WorkflowVersionExecutionNotAllowedError()
 
         payload = WorkflowRunPayload.model_validate(omit_trace_session_id_from_payload(service_api_ns.payload) or {})

api/openapi/markdown/console-openapi.md

@@ -35173,7 +35173,7 @@ The subscription constructor of the trigger provider
 
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
-| interval | string |  | Yes |
+| interval | string, <br>**Available values:** "", "month", "year" | *Enum:* `""`, `"month"`, `"year"` | Yes |
 | plan | [CloudPlan](#cloudplan) |  | Yes |
 
 #### SubscriptionQuery

api/services/billing_client/transport.py

@@ -2,23 +2,55 @@
 
 Reuses the pooled client, auth header, and connection retry of `services.billing_service`.
 Network failures, timeouts, throttling, and 5xx raise `BillingUpstreamUnavailableError`;
-any other non-200 status or a body that fails model validation raises
+any other non-200 status or an enforced response that fails model validation raises
 `BillingUpstreamInvalidResponseError`.
+
+Response validation is observed before it is enforced. For a path not in
+`ENFORCED_RESPONSE_PATHS`, a body that fails validation is logged with `logger.exception`
+(reported to Sentry) and rebuilt from the raw payload: fields that validate are coerced,
+the rest keep their raw value, and missing fields read as `None`. Migrated single-response
+reads and quota operations are enforced before their results reach business decisions.
+Batch subscription plans remain observed here; `BillingService.get_plan_bulk` validates
+each item before returning it.
 """
 
+import json
+import logging
 import os
+import types
 from collections.abc import Mapping
 from enum import StrEnum
-from typing import Literal
+from functools import cache
+from typing import Any, Literal, Union, get_args, get_origin
 from urllib.parse import quote
 
 import httpx
-from pydantic import BaseModel, ConfigDict, ValidationError
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed
 
 from core.helper.http_client_pooling import get_pooled_http_client
 from services.errors.billing import BillingUpstreamInvalidResponseError, BillingUpstreamUnavailableError
 
+logger = logging.getLogger(__name__)
+
+ENFORCED_RESPONSE_PATHS: frozenset[str] = frozenset(
+    {
+        "/subscription/info",
+        "/subscription/vector-space",
+        "/subscription/knowledge-rate-limit",
+        "/subscription/payment-link",
+        "/subscription/cleanup/whitelist",
+        "/model-provider/payment-link",
+        "/invoices",
+        "/quota/info",
+        "/quota/reserve",
+        "/quota/commit",
+        "/quota/release",
+        "/quota/balance",
+        "/quota/consume-capped",
+    }
+)
+
 type HTTPMethod = Literal["GET", "POST", "PUT", "DELETE"]
 type QueryValue = str | int | float | bool | list[str]
 
@@ -78,9 +110,9 @@ def request[ResponseT: BillingModel](
 ) -> ResponseT:
     resolved_path = path.format_map({name: quote(value, safe="") for name, value in path_params.items()})
     url = f"{_base_url(upstream)}{resolved_path}"
-    json = body.model_dump(mode="json") if body is not None else None
+    request_json = body.model_dump(mode="json") if body is not None else None
     try:
-        response = _send(method, url, params=params, json=json)
+        response = _send(method, url, params=params, json=request_json)
     except httpx.RequestError as error:
         raise BillingUpstreamUnavailableError(f"{method} {path} failed") from error
 
@@ -90,7 +122,50 @@ def request[ResponseT: BillingModel](
             raise BillingUpstreamUnavailableError(message)
         raise BillingUpstreamInvalidResponseError(message)
 
+    invalid_body = f"{method} {path} returned an invalid body"
     try:
-        return response_model.model_validate_json(response.content)
+        payload = json.loads(response.content)
+    except ValueError as error:
+        raise BillingUpstreamInvalidResponseError(invalid_body) from error
+    try:
+        return response_model.model_validate(payload)
     except ValidationError as error:
-        raise BillingUpstreamInvalidResponseError(f"{method} {path} returned an invalid body") from error
+        if path in ENFORCED_RESPONSE_PATHS or not isinstance(payload, dict):
+            raise BillingUpstreamInvalidResponseError(invalid_body) from error
+        logger.exception("Billing response failed validation: %s %s", method, path)
+        return _salvage_model(response_model, payload)
+
+
+def _salvage_model[ModelT: BillingModel](model: type[ModelT], payload: Mapping[str, Any]) -> ModelT:
+    values: dict[str, Any] = {}
+    for name, field in model.model_fields.items():
+        wire_name = field.alias or name
+        raw = payload.get(wire_name, payload.get(name))
+        values[name] = _salvage_value(field.annotation, raw)
+    return model.model_construct(**values)
+
+
+def _salvage_value(annotation: Any, raw: object) -> object:
+    if raw is None:
+        return None
+    origin = get_origin(annotation)
+    if origin in (Union, types.UnionType):
+        annotation = next(arg for arg in get_args(annotation) if arg is not type(None))
+        origin = get_origin(annotation)
+    if isinstance(annotation, type) and issubclass(annotation, BillingModel) and isinstance(raw, Mapping):
+        return _salvage_model(annotation, raw)
+    if origin is dict and isinstance(raw, Mapping):
+        value_type = get_args(annotation)[1]
+        return {key: _salvage_value(value_type, value) for key, value in raw.items()}
+    if origin is list and isinstance(raw, list):
+        item_type = get_args(annotation)[0]
+        return [_salvage_value(item_type, item) for item in raw]
+    try:
+        return _adapter(annotation).validate_python(raw)
+    except ValidationError:
+        return raw
+
+
+@cache
+def _adapter(annotation: Any) -> TypeAdapter[Any]:
+    return TypeAdapter(annotation)

api/services/billing_service.py

@@ -11,8 +11,8 @@ from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fix
 from werkzeug.exceptions import InternalServerError
 
 from core.helper.http_client_pooling import get_pooled_http_client
-from enums import CloudPlan
 from extensions.ext_redis import redis_client
+from services.billing_client import billing
 from services.billing_portal_service import BillingPortalLink
 from services.compliance_download_service import ComplianceDownloadLink
 from services.errors.billing import (
@@ -40,13 +40,6 @@ class _BillingHTTPStatusError(ValueError):
         self.status_code = status_code
 
 
-class SubscriptionPlan(TypedDict):
-    """Tenant subscriptionplan information."""
-
-    plan: str
-    expiration_date: int
-
-
 class MessageResponseDict(TypedDict):
     message: str
 
@@ -68,9 +61,6 @@ class EducationAutocompleteResponseDict(TypedDict):
     has_next: bool
 
 
-_billing_portal_link_adapter = TypeAdapter(BillingPortalLink)
-
-
 _compliance_download_link_adapter = TypeAdapter(ComplianceDownloadLink)
 
 
@@ -125,7 +115,7 @@ class TenantFeatureQuotaInfo(TypedDict):
     """Response of /quota/info.
 
     NOTE (hj24):
-    - Same convention as BillingInfo: billing may return int fields as str,
+    - Billing may return int fields as str,
       always keep non-strict mode to auto-coerce.
     """
 
@@ -136,55 +126,6 @@ class TenantFeatureQuotaInfo(TypedDict):
 _tenant_feature_quota_info_adapter = TypeAdapter(TenantFeatureQuotaInfo)
 
 
-class _BillingQuota(TypedDict):
-    size: int
-    limit: int
-
-
-class _VectorSpaceQuota(TypedDict):
-    size: float
-    limit: int
-    usage_unknown: NotRequired[bool]
-
-
-class _BillingSubscription(TypedDict):
-    plan: str
-    interval: str
-    education: bool
-
-
-class BillingInfo(TypedDict):
-    """Response of /subscription/info.
-
-    NOTE (hj24):
-    - Fields not listed here (e.g. trigger_event, api_rate_limit) are stripped by TypeAdapter.validate_python()
-    - To ensure the precision, billing may convert fields like int as str, be careful when use TypeAdapter:
-        1. validate_python in non-strict mode will coerce it to the expected type
-        2. In strict mode, it will raise ValidationError
-        3. To preserve compatibility, always keep non-strict mode here and avoid strict mode
-    """
-
-    subscription: _BillingSubscription
-    members: _BillingQuota
-    apps: _BillingQuota
-    vector_space: NotRequired[_VectorSpaceQuota]
-    documents_upload_quota: _BillingQuota
-    annotation_quota_limit: _BillingQuota
-    can_replace_logo: bool
-    model_load_balancing_enabled: bool
-    knowledge_pipeline_publish_enabled: bool
-    next_credit_reset_date: NotRequired[int]
-
-
-_billing_info_adapter = TypeAdapter(BillingInfo)
-_vector_space_quota_adapter = TypeAdapter(_VectorSpaceQuota)
-
-
-class KnowledgeRateLimitDict(TypedDict):
-    limit: int
-    subscription_plan: str
-
-
 class LangContentDict(TypedDict):
     lang: str
     title: str
@@ -221,25 +162,14 @@ class BillingService:
     _PLAN_CACHE_TTL = 600
 
     @classmethod
-    def get_info(cls, tenant_id: str, exclude_vector_space: bool = False) -> BillingInfo:
-        params = {"tenant_id": tenant_id}
-        if exclude_vector_space:
-            params["exclude_vector_space"] = "true"
-
-        billing_info = cls._send_request("GET", "/subscription/info", params=params)
-        if exclude_vector_space and billing_info.get("vector_space") is None:
-            # Unset proto message fields can be serialized as null; the light billing contract treats it as absent.
-            billing_info.pop("vector_space", None)
-        return _billing_info_adapter.validate_python(billing_info)
+    def get_info(cls, tenant_id: str, exclude_vector_space: bool = False) -> billing.GetSubscriptionInfoReply:
+        return billing.subscription_get_subscription_info(
+            tenant_id=tenant_id, exclude_vector_space=exclude_vector_space
+        )
 
     @classmethod
-    def get_vector_space(cls, tenant_id: str, bypass_cache: bool = False) -> _VectorSpaceQuota:
-        params = {"tenant_id": tenant_id}
-        if bypass_cache:
-            params["bypass_cache"] = "true"
-        return _vector_space_quota_adapter.validate_python(
-            cls._send_request("GET", "/subscription/vector-space", params=params)
-        )
+    def get_vector_space(cls, tenant_id: str, bypass_cache: bool = False) -> billing.VectorSpaceQuota:
+        return billing.subscription_get_subscription_vector_space(tenant_id=tenant_id, bypass_cache=bypass_cache)
 
     @classmethod
     def invalidate_vector_space_cache(cls, tenant_id: str) -> None:
@@ -346,38 +276,31 @@ class BillingService:
         )
 
     @classmethod
-    def get_knowledge_rate_limit(cls, tenant_id: str) -> KnowledgeRateLimitDict:
-        params = {"tenant_id": tenant_id}
-
-        knowledge_rate_limit = cls._send_request("GET", "/subscription/knowledge-rate-limit", params=params)
-
-        return {
-            "limit": knowledge_rate_limit.get("limit", 10),
-            "subscription_plan": knowledge_rate_limit.get("subscription_plan", CloudPlan.SANDBOX),
-        }
+    def get_knowledge_rate_limit(cls, tenant_id: str) -> billing.GetKnowledgeRateLimitReply:
+        return billing.subscription_get_knowledge_rate_limit(tenant_id=tenant_id)
 
     @classmethod
     def get_subscription(
         cls, plan: str, interval: str, prefilled_email: str = "", tenant_id: str = ""
     ) -> BillingPortalLink:
-        params = {"plan": plan, "interval": interval, "prefilled_email": prefilled_email, "tenant_id": tenant_id}
-        return cls._send_billing_portal_request("/subscription/payment-link", params=params)
+        reply = billing.subscription_get_subscription_payment_link(
+            plan=plan, interval=interval, prefilled_email=prefilled_email, tenant_id=tenant_id
+        )
+        return {"url": reply.url}
 
     @classmethod
     @deprecated("Only used by the deprecated model-provider checkout endpoint.")
-    def get_model_provider_payment_link(cls, provider_name: str, tenant_id: str, account_id: str, prefilled_email: str):
-        params = {
-            "provider_name": provider_name,
-            "tenant_id": tenant_id,
-            "account_id": account_id,
-            "prefilled_email": prefilled_email,
-        }
-        return cls._send_request("GET", "/model-provider/payment-link", params=params)
+    def get_model_provider_payment_link(
+        cls, provider_name: str, tenant_id: str, account_id: str, prefilled_email: str
+    ) -> billing.GetModelProviderPaymentLinkReply:
+        return billing.model_provider_get_model_provider_payment_link(
+            provider_name=provider_name, tenant_id=tenant_id, account_id=account_id, prefilled_email=prefilled_email
+        )
 
     @classmethod
     def get_invoices(cls, prefilled_email: str = "", tenant_id: str = "") -> BillingPortalLink:
-        params = {"prefilled_email": prefilled_email, "tenant_id": tenant_id}
-        return cls._send_billing_portal_request("/invoices", params=params)
+        reply = billing.invoice_get_invoices(prefilled_email=prefilled_email, tenant_id=tenant_id)
+        return {"url": reply.url}
 
     @classmethod
     def _send_quota_request(
@@ -425,29 +348,6 @@ class BillingService:
             raise ValueError(f"Unable to process delete request {url}. Please try again later or contact support.")
         return response.json()
 
-    @classmethod
-    def _send_billing_portal_request(
-        cls,
-        endpoint: str,
-        *,
-        params: dict[str, str],
-    ) -> BillingPortalLink:
-        try:
-            response = cls._send_request("GET", endpoint, params=params)
-            return _billing_portal_link_adapter.validate_python(response)
-        except _BillingHTTPStatusError as error:
-            if error.status_code in {httpx.codes.REQUEST_TIMEOUT, httpx.codes.TOO_MANY_REQUESTS} or (
-                error.status_code >= 500
-            ):
-                raise BillingUpstreamUnavailableError from error
-            raise BillingUpstreamInvalidResponseError from error
-        except httpx.RequestError as error:
-            raise BillingUpstreamUnavailableError from error
-        except (json.JSONDecodeError, UnicodeDecodeError, ValidationError) as error:
-            raise BillingUpstreamInvalidResponseError from error
-        except ValueError as error:
-            raise RuntimeError("Unexpected billing service value error") from error
-
     @classmethod
     def delete_account(cls, account_id: str) -> MessageResponseDict:
         """Delete account."""
@@ -555,32 +455,33 @@ class BillingService:
         redis_client.delete(f"tenant:{tenant_id}:billing_info")
 
     @classmethod
-    def get_plan_bulk(cls, tenant_ids: Sequence[str]) -> dict[str, SubscriptionPlan]:
+    def get_plan_bulk(cls, tenant_ids: Sequence[str]) -> dict[str, billing.SubscriptionPlanItem]:
         """
-        Bulk fetch billing subscription plan via billing API.
-        Payload: {"tenant_ids": ["t1", "t2", ...]} (max 200 per request)
+        Bulk fetch billing subscription plan via billing API, at most 200 tenants per request.
+        Invalid items are skipped so callers only receive validated plans.
         Returns:
-            Mapping of tenant_id -> {plan: str, expiration_date: int}
+            Mapping of tenant_id -> subscription plan
         """
-        results: dict[str, SubscriptionPlan] = {}
-        subscription_adapter = TypeAdapter(SubscriptionPlan)
+        results: dict[str, billing.SubscriptionPlanItem] = {}
 
         chunk_size = 200
         for i in range(0, len(tenant_ids), chunk_size):
             chunk = tenant_ids[i : i + chunk_size]
             try:
-                resp = cls._send_request("POST", "/subscription/plan/batch", json={"tenant_ids": chunk})
-                data = resp.get("data", {})
-
-                for tenant_id, plan in data.items():
+                reply = billing.subscription_batch_get_subscription_plan(
+                    body=billing.BatchGetSubscriptionPlanReq(tenant_ids=list(chunk))
+                )
+                for tenant_id, item in (reply.data or {}).items():
                     try:
-                        subscription_plan = subscription_adapter.validate_python(plan)
-                        results[tenant_id] = subscription_plan
-                    except Exception:
+                        # Observed responses can contain constructed models; validating an instance would trust it.
+                        payload = (
+                            item.model_dump(warnings=False) if isinstance(item, billing.SubscriptionPlanItem) else item
+                        )
+                        results[tenant_id] = billing.SubscriptionPlanItem.model_validate(payload)
+                    except ValidationError:
                         logger.exception(
                             "get_plan_bulk: failed to validate subscription plan for tenant(%s)", tenant_id
                         )
-                        continue
             except Exception:
                 logger.exception("get_plan_bulk: failed to fetch billing info batch for tenants: %s", chunk)
                 continue
@@ -592,22 +493,20 @@ class BillingService:
         return f"{cls._PLAN_CACHE_KEY_PREFIX}{tenant_id}"
 
     @classmethod
-    def get_plan_bulk_with_cache(cls, tenant_ids: Sequence[str]) -> dict[str, SubscriptionPlan]:
+    def get_plan_bulk_with_cache(cls, tenant_ids: Sequence[str]) -> dict[str, billing.SubscriptionPlanItem]:
         """
         Bulk fetch billing subscription plan with cache to reduce billing API loads in batch job scenarios.
 
         NOTE: if you want to high data consistency, use get_plan_bulk instead.
 
         Returns:
-            Mapping of tenant_id -> {plan: str, expiration_date: int}
+            Mapping of tenant_id -> subscription plan
         """
-        tenant_plans: dict[str, SubscriptionPlan] = {}
+        tenant_plans: dict[str, billing.SubscriptionPlanItem] = {}
 
         if not tenant_ids:
             return tenant_plans
 
-        subscription_adapter = TypeAdapter(SubscriptionPlan)
-
         # Step 1: Batch fetch from Redis cache using mget
         redis_keys = [cls._make_plan_cache_key(tenant_id) for tenant_id in tenant_ids]
         try:
@@ -624,15 +523,7 @@ class BillingService:
             for tenant_id, cached_value in zip(tenant_ids, cached_values, strict=False):
                 if cached_value:
                     try:
-                        # Redis returns bytes, decode to string and parse JSON
-                        json_str = cached_value.decode("utf-8") if isinstance(cached_value, bytes) else cached_value
-                        plan_dict = json.loads(json_str)
-                        # NOTE (hj24): New billing versions may return timestamp as str, and validate_python
-                        # in non-strict mode will coerce it to the expected int type.
-                        # To preserve compatibility, always keep non-strict mode here and avoid strict mode.
-                        subscription_plan = subscription_adapter.validate_python(plan_dict)
-                        # NOTE END
-                        tenant_plans[tenant_id] = subscription_plan
+                        tenant_plans[tenant_id] = billing.SubscriptionPlanItem.model_validate_json(cached_value)
                     except Exception:
                         logger.exception(
                             "get_plan_bulk_with_cache: process tenant(%s) failed, add to cache misses", tenant_id
@@ -655,7 +546,7 @@ class BillingService:
             bulk_plans = BillingService.get_plan_bulk(cache_misses)
 
             if bulk_plans:
-                plans_to_cache: dict[str, SubscriptionPlan] = {}
+                plans_to_cache: dict[str, billing.SubscriptionPlanItem] = {}
 
                 for tenant_id, subscription_plan in bulk_plans.items():
                     tenant_plans[tenant_id] = subscription_plan
@@ -667,9 +558,7 @@ class BillingService:
                         pipe = redis_client.pipeline()
                         for tenant_id, subscription_plan in plans_to_cache.items():
                             redis_key = cls._make_plan_cache_key(tenant_id)
-                            # Serialize dict to JSON string
-                            json_str = json.dumps(subscription_plan)
-                            pipe.setex(redis_key, cls._PLAN_CACHE_TTL, json_str)
+                            pipe.setex(redis_key, cls._PLAN_CACHE_TTL, subscription_plan.model_dump_json())
                         pipe.execute()
 
                         logger.info(
@@ -683,12 +572,8 @@ class BillingService:
 
     @classmethod
     def get_expired_subscription_cleanup_whitelist(cls) -> Sequence[str]:
-        resp = cls._send_request("GET", "/subscription/cleanup/whitelist")
-        data = resp.get("data", [])
-        tenant_whitelist = []
-        for item in data:
-            tenant_whitelist.append(item["tenant_id"])
-        return tenant_whitelist
+        reply = billing.subscription_list_subscription_cleanup_whitelist()
+        return [item.tenant_id for item in reply.data or []]
 
     @classmethod
     def get_account_notification(cls, account_id: str) -> AccountNotificationDict:

api/services/clear_free_plan_tenant_expired_logs.py

@@ -372,7 +372,7 @@ class ClearFreePlanTenantExpiredLogs:
             try:
                 if (
                     dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD
-                    or BillingService.get_info(tenant_id)["subscription"]["plan"] == CloudPlan.SANDBOX
+                    or BillingService.get_info(tenant_id).subscription.plan == CloudPlan.SANDBOX
                 ):
                     # only process sandbox tenant
                     with sessionmaker(db.engine).begin() as session:

api/services/entities/feature_entities.py

@@ -1,6 +1,7 @@
 """Feature query results and policy values shared by their consumers."""
 
 from enum import StrEnum
+from typing import Literal
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -13,7 +14,7 @@ class FeatureResponseModel(BaseModel):
 
 class SubscriptionModel(FeatureResponseModel):
     plan: CloudPlan = CloudPlan.SANDBOX
-    interval: str = ""
+    interval: Literal["", "month", "year"] = ""
 
 
 class BillingModel(FeatureResponseModel):
@@ -176,7 +177,7 @@ class FeatureModel(FeatureResponseModel):
 class KnowledgeRateLimitModel(FeatureResponseModel):
     enabled: bool = False
     limit: int = 10
-    subscription_plan: str = ""
+    subscription_plan: Literal["", "sandbox", "professional", "team"] = ""
 
 
 class SystemFeatureModel(FeatureResponseModel):

api/services/feature_service.py

@@ -1,7 +1,8 @@
 from configs import dify_config
 from enums import CloudPlan, DeploymentEdition, HostedTrialProvider
 from services.account_education_service import EDUCATION_EDITIONS
-from services.billing_service import BillingInfo, BillingService
+from services.billing_client import billing
+from services.billing_service import BillingService
 from services.enterprise.enterprise_service import EnterpriseService
 from services.entities import feature_entities
 from services.knowledge_fs.feature_service import get_knowledge_fs_features
@@ -31,7 +32,7 @@ class FeatureService:
             return CloudPlan.SANDBOX
 
         billing_info = BillingService.get_info(tenant_id, exclude_vector_space=True)
-        return CloudPlan(billing_info["subscription"]["plan"])
+        return CloudPlan(billing_info.subscription.plan)
 
     @classmethod
     def get_features(cls, tenant_id: str, exclude_vector_space: bool = False) -> feature_entities.FeatureModel:
@@ -68,9 +69,9 @@ class FeatureService:
             billing_vector_space = BillingService.get_vector_space(tenant_id)
             # NOTE: billing API returns vector_space.size as float (e.g. 0.0),
             # but feature API keeps LimitationModel.size as int for compatibility.
-            vector_space.size = int(billing_vector_space["size"])
-            vector_space.limit = billing_vector_space["limit"]
-            vector_space.usage_unknown = billing_vector_space.get("usage_unknown", False)
+            vector_space.size = int(billing_vector_space.size)
+            vector_space.limit = billing_vector_space.limit
+            vector_space.usage_unknown = billing_vector_space.usage_unknown or False
 
         return vector_space
 
@@ -80,8 +81,8 @@ class FeatureService:
         if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and tenant_id:
             knowledge_rate_limit.enabled = True
             limit_info = BillingService.get_knowledge_rate_limit(tenant_id)
-            knowledge_rate_limit.limit = limit_info.get("limit", 10)
-            knowledge_rate_limit.subscription_plan = limit_info.get("subscription_plan", CloudPlan.SANDBOX)
+            knowledge_rate_limit.limit = limit_info.limit
+            knowledge_rate_limit.subscription_plan = limit_info.subscription_plan
         return knowledge_rate_limit
 
     @classmethod
@@ -158,12 +159,11 @@ class FeatureService:
         else:
             billing_info = BillingService.get_info(tenant_id)
 
+        features.billing.subscription.plan = CloudPlan(billing_info.subscription.plan)
+        features.billing.subscription.interval = billing_info.subscription.interval
+        features.education.activated = billing_info.subscription.education
         features_usage_info = BillingService.get_quota_info(tenant_id)
 
-        features.billing.subscription.plan = CloudPlan(billing_info["subscription"]["plan"])
-        features.billing.subscription.interval = billing_info["subscription"]["interval"]
-        features.education.activated = billing_info["subscription"].get("education", False)
-
         if features.billing.subscription.plan != CloudPlan.SANDBOX:
             features.webapp_copyright_enabled = True
         else:
@@ -179,46 +179,35 @@ class FeatureService:
             features.api_rate_limit.limit = features_usage_info["api_rate_limit"]["limit"]
             features.api_rate_limit.reset_date = features_usage_info["api_rate_limit"].get("reset_date", -1)
 
-        if "members" in billing_info:
-            features.members.size = billing_info["members"]["size"]
-            features.members.limit = billing_info["members"]["limit"]
+        features.members.size = billing_info.members.size
+        features.members.limit = billing_info.members.limit
 
-        if "apps" in billing_info:
-            features.apps.size = billing_info["apps"]["size"]
-            features.apps.limit = billing_info["apps"]["limit"]
+        features.apps.size = billing_info.apps.size
+        features.apps.limit = billing_info.apps.limit
 
         if not exclude_vector_space:
             assert features.vector_space is not None
             cls._fulfill_vector_space_from_billing_info(features.vector_space, billing_info)
 
-        if "documents_upload_quota" in billing_info:
-            features.documents_upload_quota.size = billing_info["documents_upload_quota"]["size"]
-            features.documents_upload_quota.limit = billing_info["documents_upload_quota"]["limit"]
+        features.documents_upload_quota.size = billing_info.documents_upload_quota.size
+        features.documents_upload_quota.limit = billing_info.documents_upload_quota.limit
 
-        if "annotation_quota_limit" in billing_info:
-            features.annotation_quota_limit.size = billing_info["annotation_quota_limit"]["size"]
-            features.annotation_quota_limit.limit = billing_info["annotation_quota_limit"]["limit"]
+        features.annotation_quota_limit.size = billing_info.annotation_quota_limit.size
+        features.annotation_quota_limit.limit = billing_info.annotation_quota_limit.limit
 
-        if "can_replace_logo" in billing_info:
-            features.can_replace_logo = billing_info["can_replace_logo"]
-
-        if "model_load_balancing_enabled" in billing_info:
-            features.model_load_balancing_enabled = billing_info["model_load_balancing_enabled"]
-
-        if "knowledge_pipeline_publish_enabled" in billing_info:
-            features.knowledge_pipeline.publish_enabled = billing_info["knowledge_pipeline_publish_enabled"]
-
-        if "next_credit_reset_date" in billing_info:
-            features.next_credit_reset_date = billing_info["next_credit_reset_date"]
+        features.can_replace_logo = billing_info.can_replace_logo
+        features.model_load_balancing_enabled = billing_info.model_load_balancing_enabled
+        features.knowledge_pipeline.publish_enabled = billing_info.knowledge_pipeline_publish_enabled
+        features.next_credit_reset_date = billing_info.next_credit_reset_date
 
     @classmethod
     def _fulfill_vector_space_from_billing_info(
-        cls, vector_space: feature_entities.LimitationModel, billing_info: BillingInfo
+        cls, vector_space: feature_entities.LimitationModel, billing_info: billing.GetSubscriptionInfoReply
     ):
-        if "vector_space" not in billing_info:
+        if billing_info.vector_space is None:
             return
 
         # NOTE: billing API returns vector_space.size as float (e.g. 0.0),
         # but feature API keeps LimitationModel.size as int for compatibility.
-        vector_space.size = int(billing_info["vector_space"]["size"])
-        vector_space.limit = billing_info["vector_space"]["limit"]
+        vector_space.size = int(billing_info.vector_space.size)
+        vector_space.limit = billing_info.vector_space.limit

api/services/network_access_group_gateway.py

@@ -21,7 +21,7 @@ from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fix
 
 from core.helper.http_client_pooling import get_pooled_http_client
 from enums import CloudPlan
-from services.billing_service import BillingService, _BillingHTTPStatusError
+from services.billing_service import BillingService
 from services.entities.network_access_group_entities import (
     NetworkAccessAppConfig,
     NetworkAccessAppStatusSnapshot,
@@ -30,6 +30,7 @@ from services.entities.network_access_group_entities import (
     NetworkAccessGroup,
     NetworkAccessGroupList,
 )
+from services.errors.billing import BillingError
 from services.network_access_group_service import (
     NetworkAccessGroupControlPlane,
     NetworkAccessGroupEntitlement,
@@ -447,13 +448,7 @@ class BillingNetworkAccessGroupEntitlementGateway(NetworkAccessGroupEntitlement)
     @override
     def is_paid_plan(self, workspace_id: str) -> bool:
         try:
-            billing_info = BillingService.get_info(workspace_id, exclude_vector_space=True)
-        except (
-            httpx.RequestError,
-            json.JSONDecodeError,
-            UnicodeDecodeError,
-            ValidationError,
-            _BillingHTTPStatusError,
-        ) as exc:
+            subscription = BillingService.get_info(workspace_id, exclude_vector_space=True).subscription
+        except BillingError as exc:
             raise NetworkAccessGroupEntitlementUnavailableError from exc
-        return billing_info["subscription"]["plan"] in (CloudPlan.PROFESSIONAL, CloudPlan.TEAM)
+        return subscription.plan in (CloudPlan.PROFESSIONAL, CloudPlan.TEAM)

api/services/retention/conversation/messages_clean_policy.py

@@ -6,7 +6,8 @@ from typing import Protocol, override
 
 from configs import dify_config
 from enums import CloudPlan, DeploymentEdition
-from services.billing_service import BillingService, SubscriptionPlan
+from services.billing_client.billing import SubscriptionPlanItem
+from services.billing_service import BillingService
 
 logger = logging.getLogger(__name__)
 
@@ -72,7 +73,7 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
 
     def __init__(
         self,
-        plan_provider: Callable[[Sequence[str]], dict[str, SubscriptionPlan]],
+        plan_provider: Callable[[Sequence[str]], dict[str, SubscriptionPlanItem]],
         graceful_period_days: int = 21,
         tenant_whitelist: Sequence[str] | None = None,
         current_timestamp: int | None = None,
@@ -119,7 +120,7 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
         self,
         messages: Sequence[SimpleMessage],
         app_to_tenant: dict[str, str],
-        tenant_plans: dict[str, SubscriptionPlan],
+        tenant_plans: dict[str, SubscriptionPlanItem],
     ) -> list[str]:
         """
         Filter messages that should be deleted based on sandbox plan expiration.
@@ -160,8 +161,8 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
             if not tenant_plan:
                 continue
 
-            plan = str(tenant_plan["plan"])
-            expiration_date = int(tenant_plan["expiration_date"])
+            plan = tenant_plan.plan
+            expiration_date = tenant_plan.expiration_date
 
             # Only process sandbox plans
             if plan != CloudPlan.SANDBOX:

api/services/retention/workflow_run/archive_paid_plan_workflow_run.py

@@ -549,7 +549,7 @@ class WorkflowRunArchiver:
         # Filter to paid tenants (any plan except SANDBOX)
         paid = set()
         for tid, info in bulk_info.items():
-            if info and info.get("plan") in (CloudPlan.PROFESSIONAL, CloudPlan.TEAM):
+            if info and info.plan in (CloudPlan.PROFESSIONAL, CloudPlan.TEAM):
                 paid.add(tid)
 
         return paid

api/services/retention/workflow_run/clear_free_plan_expired_workflow_run_logs.py

@@ -25,7 +25,8 @@ from repositories.api_workflow_run_repository import (
 )
 from repositories.factory import DifyAPIRepositoryFactory
 from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository
-from services.billing_service import BillingService, SubscriptionPlan
+from services.billing_client.billing import SubscriptionPlanItem
+from services.billing_service import BillingService
 
 logger = logging.getLogger(__name__)
 
@@ -502,7 +503,7 @@ class WorkflowRunCleanup:
                 logger.warning("Missing billing info for tenant %s in bulk resp; treating as non-free", tenant_id)
                 continue
 
-            if info.get("plan") != CloudPlan.SANDBOX:
+            if info.plan != CloudPlan.SANDBOX:
                 continue
 
             if self._is_within_grace_period(tenant_id, info):
@@ -522,12 +523,11 @@ class WorkflowRunCleanup:
             logger.exception("Failed to parse expiration timestamp for tenant %s", tenant_id)
             return None
 
-    def _is_within_grace_period(self, tenant_id: str, info: SubscriptionPlan) -> bool:
+    def _is_within_grace_period(self, tenant_id: str, info: SubscriptionPlanItem) -> bool:
         if self.free_plan_grace_period_days <= 0:
             return False
 
-        expiration_value = info.get("expiration_date", -1)
-        expiration_at = self._expiration_datetime(tenant_id, expiration_value)
+        expiration_at = self._expiration_datetime(tenant_id, info.expiration_date)
         if expiration_at is None:
             return False
 

api/services/vector_space_admission_service.py

@@ -366,10 +366,10 @@ 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):
+            if vector_space.usage_unknown:
                 raise ValueError("Vector storage usage is unknown")
-            current_usage_mb = float(vector_space["size"])
-            plan_limit_mb = int(vector_space["limit"])
+            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:
@@ -441,7 +441,7 @@ class VectorSpaceAdmissionService:
 
         plan = None
         try:
-            plan = CloudPlan(billing_info["subscription"]["plan"])
+            plan = CloudPlan(billing_info.subscription.plan)
         except ValueError:
             logger.warning(
                 "Skipping TiDB vector-space admission for unknown plan tenant_id=%s",

api/services/workflow/queue_dispatcher.py

@@ -95,7 +95,7 @@ class QueueDispatcherManager:
         if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD:
             try:
                 billing_info = BillingService.get_info(tenant_id)
-                plan = billing_info.get("subscription", {}).get("plan", "sandbox")
+                plan = billing_info.subscription.plan
             except Exception:
                 # If billing service fails, default to sandbox
                 plan = "sandbox"

api/services/workflow_service.py

@@ -737,7 +737,7 @@ class WorkflowService:
         # billing check
         if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD:
             limit_info = BillingService.get_info(app_model.tenant_id)
-            if limit_info["subscription"]["plan"] == CloudPlan.SANDBOX:
+            if limit_info.subscription.plan == CloudPlan.SANDBOX:
                 # Check trigger node count limit for SANDBOX plan
                 trigger_node_count = sum(
                     1

api/services/workspace/gateways.py

@@ -329,7 +329,7 @@ class DeploymentWorkspacePlanGateway(WorkspacePlanGateway):
         for workspace_id in ids:
             tenant_plan = bulk_plans.get(workspace_id)
             if tenant_plan:
-                resolved[workspace_id] = tenant_plan["plan"] or CloudPlan.SANDBOX
+                resolved[workspace_id] = tenant_plan.plan or CloudPlan.SANDBOX
                 continue
 
             features = FeatureService.get_features(workspace_id, exclude_vector_space=True)
@@ -354,9 +354,7 @@ class DeploymentWorkspaceFeatureGateway(WorkspaceFeatureGateway):
         if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD:
             return EffectiveCreditPool()
         billing = BillingService.get_info(workspace_id, exclude_vector_space=True)
-        return self._credit_pool(
-            workspace_id, CloudPlan(billing["subscription"]["plan"]), billing.get("next_credit_reset_date")
-        )
+        return self._credit_pool(workspace_id, CloudPlan(billing.subscription.plan), billing.next_credit_reset_date)
 
     @staticmethod
     def _credit_pool(workspace_id: str, plan: CloudPlan, reset_date: int | None) -> EffectiveCreditPool:

api/tests/e2e_tests/containers/services/test_billing_service.py

@@ -5,7 +5,9 @@ import pytest
 from flask import Flask
 
 from extensions.ext_redis import redis_client
+from services.billing_client import billing
 from services.billing_service import BillingService
+from tests.helpers import billing_builders as builders
 
 
 class TestBillingServiceGetPlanBulkWithCache:
@@ -33,15 +35,15 @@ class TestBillingServiceGetPlanBulkWithCache:
             if keys:
                 redis_client.delete(*keys)
 
-    def _create_test_plan_data(self, plan: str = "sandbox", expiration_date: int = 1735689600):
-        """Helper to create test SubscriptionPlan data."""
-        return {"plan": plan, "expiration_date": expiration_date}
+    def _create_test_plan_data(
+        self, plan: str = "sandbox", expiration_date: int = 1735689600
+    ) -> billing.SubscriptionPlanItem:
+        return builders.plan_item(plan, expiration_date)
 
-    def _set_cache(self, tenant_id: str, plan_data: dict, ttl: int = 600):
+    def _set_cache(self, tenant_id: str, plan_data: billing.SubscriptionPlanItem, ttl: int = 600):
         """Helper to set cache data in Redis."""
         cache_key = BillingService._make_plan_cache_key(tenant_id)
-        json_str = json.dumps(plan_data)
-        redis_client.setex(cache_key, ttl, json_str)
+        redis_client.setex(cache_key, ttl, plan_data.model_dump_json())
 
     def _get_cache(self, tenant_id: str):
         """Helper to get cache data from Redis."""
@@ -74,12 +76,12 @@ class TestBillingServiceGetPlanBulkWithCache:
 
             # Assert
             assert len(result) == 3
-            assert result["tenant-1"]["plan"] == "sandbox"
-            assert result["tenant-1"]["expiration_date"] == 1735689600
-            assert result["tenant-2"]["plan"] == "professional"
-            assert result["tenant-2"]["expiration_date"] == 1767225600
-            assert result["tenant-3"]["plan"] == "team"
-            assert result["tenant-3"]["expiration_date"] == 1798761600
+            assert result["tenant-1"].plan == "sandbox"
+            assert result["tenant-1"].expiration_date == 1735689600
+            assert result["tenant-2"].plan == "professional"
+            assert result["tenant-2"].expiration_date == 1767225600
+            assert result["tenant-3"].plan == "team"
+            assert result["tenant-3"].expiration_date == 1798761600
 
             # Verify API was not called
             mock_get_plan_bulk.assert_not_called()
@@ -100,8 +102,8 @@ class TestBillingServiceGetPlanBulkWithCache:
 
             # Assert
             assert len(result) == 2
-            assert result["tenant-1"]["plan"] == "sandbox"
-            assert result["tenant-2"]["plan"] == "professional"
+            assert result["tenant-1"].plan == "sandbox"
+            assert result["tenant-2"].plan == "professional"
 
             # Verify API was called with correct tenant_ids
             mock_get_plan_bulk.assert_called_once_with(tenant_ids)
@@ -113,8 +115,8 @@ class TestBillingServiceGetPlanBulkWithCache:
             assert cached_2 is not None
 
             # Verify cache content
-            cached_data_1 = json.loads(cached_1)
-            cached_data_2 = json.loads(cached_2)
+            cached_data_1 = billing.SubscriptionPlanItem.model_validate_json(cached_1)
+            cached_data_2 = billing.SubscriptionPlanItem.model_validate_json(cached_2)
             assert cached_data_1 == expected_plans["tenant-1"]
             assert cached_data_2 == expected_plans["tenant-2"]
 
@@ -142,9 +144,9 @@ class TestBillingServiceGetPlanBulkWithCache:
 
             # Assert
             assert len(result) == 3
-            assert result["tenant-1"]["plan"] == "sandbox"
-            assert result["tenant-2"]["plan"] == "professional"
-            assert result["tenant-3"]["plan"] == "team"
+            assert result["tenant-1"].plan == "sandbox"
+            assert result["tenant-2"].plan == "professional"
+            assert result["tenant-3"].plan == "team"
 
             # Verify API was called only for missing tenant
             mock_get_plan_bulk.assert_called_once_with(["tenant-3"])
@@ -152,7 +154,7 @@ class TestBillingServiceGetPlanBulkWithCache:
             # Verify tenant-3 data was written to cache
             cached_3 = self._get_cache("tenant-3")
             assert cached_3 is not None
-            cached_data_3 = json.loads(cached_3)
+            cached_data_3 = billing.SubscriptionPlanItem.model_validate_json(cached_3)
             assert cached_data_3 == missing_plan["tenant-3"]
 
     def test_get_plan_bulk_with_cache_redis_mget_failure(self, flask_app_with_containers: Flask):
@@ -174,8 +176,8 @@ class TestBillingServiceGetPlanBulkWithCache:
 
             # Assert
             assert len(result) == 2
-            assert result["tenant-1"]["plan"] == "sandbox"
-            assert result["tenant-2"]["plan"] == "professional"
+            assert result["tenant-1"].plan == "sandbox"
+            assert result["tenant-2"].plan == "professional"
 
             # Verify API was called for all tenants (fallback)
             mock_get_plan_bulk.assert_called_once_with(tenant_ids)
@@ -211,9 +213,9 @@ class TestBillingServiceGetPlanBulkWithCache:
 
             # Assert
             assert len(result) == 3
-            assert result["tenant-1"]["plan"] == "sandbox"  # From cache
-            assert result["tenant-2"]["plan"] == "professional"  # From API (fallback)
-            assert result["tenant-3"]["plan"] == "team"  # From API
+            assert result["tenant-1"].plan == "sandbox"  # From cache
+            assert result["tenant-2"].plan == "professional"  # From API (fallback)
+            assert result["tenant-3"].plan == "team"  # From API
 
             # Verify API was called for tenant-2 and tenant-3
             mock_get_plan_bulk.assert_called_once_with(["tenant-2", "tenant-3"])
@@ -221,10 +223,10 @@ class TestBillingServiceGetPlanBulkWithCache:
             # Verify tenant-2's invalid JSON was replaced with correct data in cache
             cached_2 = self._get_cache("tenant-2")
             assert cached_2 is not None
-            cached_data_2 = json.loads(cached_2)
+            cached_data_2 = billing.SubscriptionPlanItem.model_validate_json(cached_2)
             assert cached_data_2 == expected_plans["tenant-2"]
-            assert cached_data_2["plan"] == "professional"
-            assert cached_data_2["expiration_date"] == 1767225600
+            assert cached_data_2.plan == "professional"
+            assert cached_data_2.expiration_date == 1767225600
 
             # Verify tenant-2 cache has correct TTL
             cache_key_2_new = BillingService._make_plan_cache_key("tenant-2")
@@ -235,11 +237,11 @@ class TestBillingServiceGetPlanBulkWithCache:
             # Verify tenant-3 data was also written to cache
             cached_3 = self._get_cache("tenant-3")
             assert cached_3 is not None
-            cached_data_3 = json.loads(cached_3)
+            cached_data_3 = billing.SubscriptionPlanItem.model_validate_json(cached_3)
             assert cached_data_3 == expected_plans["tenant-3"]
 
     def test_get_plan_bulk_with_cache_invalid_plan_data_in_cache(self, flask_app_with_containers: Flask):
-        """Test fallback to API when cache data doesn't match SubscriptionPlan schema."""
+        """Test fallback to API when cache data doesn't match the plan model."""
         with flask_app_with_containers.app_context():
             # Arrange
             tenant_ids = ["tenant-1", "tenant-2", "tenant-3"]
@@ -264,9 +266,9 @@ class TestBillingServiceGetPlanBulkWithCache:
 
             # Assert
             assert len(result) == 3
-            assert result["tenant-1"]["plan"] == "sandbox"  # From cache
-            assert result["tenant-2"]["plan"] == "professional"  # From API (fallback)
-            assert result["tenant-3"]["plan"] == "team"  # From API
+            assert result["tenant-1"].plan == "sandbox"  # From cache
+            assert result["tenant-2"].plan == "professional"  # From API (fallback)
+            assert result["tenant-3"].plan == "team"  # From API
 
             # Verify API was called for tenant-2 and tenant-3
             mock_get_plan_bulk.assert_called_once_with(["tenant-2", "tenant-3"])
@@ -294,8 +296,8 @@ class TestBillingServiceGetPlanBulkWithCache:
 
             # Assert - Function should still return correct result despite pipeline failure
             assert len(result) == 2
-            assert result["tenant-1"]["plan"] == "sandbox"
-            assert result["tenant-2"]["plan"] == "professional"
+            assert result["tenant-1"].plan == "sandbox"
+            assert result["tenant-2"].plan == "professional"
 
             # Verify pipeline was attempted
             mock_pipeline.assert_called_once()
@@ -349,8 +351,8 @@ class TestBillingServiceGetPlanBulkWithCache:
 
             # Assert
             assert len(result) == 2
-            assert result["tenant-1"]["plan"] == "sandbox"
-            assert result["tenant-2"]["plan"] == "professional"
+            assert result["tenant-1"].plan == "sandbox"
+            assert result["tenant-2"].plan == "professional"
 
             # Verify API was called for both tenants (tenant-1 expired, tenant-2 missing)
             mock_get_plan_bulk.assert_called_once_with(tenant_ids)

api/tests/e2e_tests/containers/services/test_feature_service.py

@@ -5,6 +5,7 @@ from faker import Faker
 from sqlalchemy.orm import Session
 
 from enums import CloudPlan, DeploymentEdition
+from services.billing_client import billing
 from services.entities.feature_entities import (
     FeatureModel,
     KnowledgeRateLimitModel,
@@ -15,6 +16,7 @@ from services.entities.feature_entities import (
 )
 from services.feature_service import FeatureService
 from services.system_feature_service import SystemFeatureService
+from tests.helpers import billing_builders as builders
 
 
 class TestFeatureService:
@@ -29,21 +31,20 @@ class TestFeatureService:
             patch("services.system_feature_service.EnterpriseService", new=mock_enterprise_service),
         ):
             # Setup default mock returns for BillingService
-            mock_billing_service.get_info.return_value = {
-                "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "monthly", "education": True},
-                "members": {"size": 5, "limit": 10},
-                "apps": {"size": 3, "limit": 20},
-                "vector_space": {"size": 2, "limit": 10},
-                "documents_upload_quota": {"size": 15, "limit": 100},
-                "annotation_quota_limit": {"size": 8, "limit": 50},
-                "can_replace_logo": True,
-                "model_load_balancing_enabled": True,
-            }
+            mock_billing_service.get_info.return_value = builders.subscription_info(
+                subscription=builders.subscription_plan("professional", interval="month", education=True),
+                members=billing.SubscriptionQuota(size=5, limit=10),
+                apps=billing.SubscriptionQuota(size=3, limit=20),
+                vector_space=builders.vector_space(size=2, limit=10),
+                documents_upload_quota=billing.SubscriptionQuota(size=15, limit=100),
+                annotation_quota_limit=billing.SubscriptionQuota(size=8, limit=50),
+                can_replace_logo=True,
+                model_load_balancing_enabled=True,
+            )
 
-            mock_billing_service.get_knowledge_rate_limit.return_value = {
-                "limit": 100,
-                "subscription_plan": CloudPlan.PROFESSIONAL,
-            }
+            mock_billing_service.get_knowledge_rate_limit.return_value = billing.GetKnowledgeRateLimitReply(
+                subscription_plan="professional", limit=100
+            )
 
             # Setup default mock returns for EnterpriseService
             mock_enterprise_service.get_workspace_info.return_value = {
@@ -115,7 +116,7 @@ class TestFeatureService:
 
             # Verify billing features
             assert result.billing.subscription.plan == CloudPlan.PROFESSIONAL
-            assert result.billing.subscription.interval == "monthly"
+            assert result.billing.subscription.interval == "month"
             assert result.education.activated is True
 
             # Verify member limitations
@@ -175,16 +176,16 @@ class TestFeatureService:
             mock_config.EDUCATION_ENABLED = False
 
             # Set mock return value inside the patch context
-            mock_external_service_dependencies["billing_service"].get_info.return_value = {
-                "subscription": {"plan": CloudPlan.SANDBOX, "interval": "monthly", "education": False},
-                "members": {"size": 1, "limit": 3},
-                "apps": {"size": 1, "limit": 5},
-                "vector_space": {"size": 1, "limit": 2},
-                "documents_upload_quota": {"size": 5, "limit": 20},
-                "annotation_quota_limit": {"size": 2, "limit": 10},
-                "can_replace_logo": False,
-                "model_load_balancing_enabled": False,
-            }
+            mock_external_service_dependencies["billing_service"].get_info.return_value = builders.subscription_info(
+                subscription=builders.subscription_plan("sandbox", interval="month", education=False),
+                members=billing.SubscriptionQuota(size=1, limit=3),
+                apps=billing.SubscriptionQuota(size=1, limit=5),
+                vector_space=builders.vector_space(size=1, limit=2),
+                documents_upload_quota=billing.SubscriptionQuota(size=5, limit=20),
+                annotation_quota_limit=billing.SubscriptionQuota(size=2, limit=10),
+                can_replace_logo=False,
+                model_load_balancing_enabled=False,
+            )
 
             # Act: Execute the method under test
             result = FeatureService.get_features(tenant_id)
@@ -701,62 +702,6 @@ class TestFeatureService:
             # Verify no billing service calls
             mock_external_service_dependencies["billing_service"].get_info.assert_not_called()
 
-    def test_get_features_partial_billing_info(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test feature retrieval with partial billing information.
-
-        This test verifies:
-        - Proper handling of partial billing data
-        - Correct fallback to default values
-        - Proper billing service integration
-        - Return value correctness and structure
-        """
-        # Arrange: Setup partial billing info mock with proper config
-        tenant_id = self._create_test_tenant_id()
-
-        with patch("services.feature_service.dify_config") as mock_config:
-            mock_config.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD
-            mock_config.CAN_REPLACE_LOGO = True
-            mock_config.MODEL_LB_ENABLED = False
-            mock_config.EDUCATION_ENABLED = False
-
-            mock_external_service_dependencies["billing_service"].get_info.return_value = {
-                "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "yearly"},
-                # Missing members, apps, vector_space, etc.
-            }
-
-            # Act: Execute the method under test
-            result = FeatureService.get_features(tenant_id)
-
-        # Assert: Verify the expected outcomes
-        assert result is not None
-        assert isinstance(result, FeatureModel)
-
-        # Verify billing features
-        assert result.billing.subscription.plan == CloudPlan.PROFESSIONAL
-        assert result.billing.subscription.interval == "yearly"
-
-        # Verify default values for missing billing info
-        assert result.members.size == 0
-        assert result.members.limit == 1
-        assert result.apps.size == 0
-        assert result.apps.limit == 10
-        assert result.vector_space.size == 0
-        assert result.vector_space.limit == 5
-        assert result.documents_upload_quota.size == 0
-        assert result.documents_upload_quota.limit == 50
-        assert result.annotation_quota_limit.size == 0
-        assert result.annotation_quota_limit.limit == 10
-
-        # Verify paid plan behavior.
-        assert result.webapp_copyright_enabled is True
-        assert result.is_allow_transfer_workspace is True
-
-        # Verify mock interactions
-        mock_external_service_dependencies["billing_service"].get_info.assert_called_once_with(tenant_id)
-
     def test_get_features_edge_case_vector_space(
         self, db_session_with_containers: Session, mock_external_service_dependencies
     ):
@@ -778,11 +723,11 @@ class TestFeatureService:
             mock_config.MODEL_LB_ENABLED = False
             mock_config.EDUCATION_ENABLED = False
 
-            mock_external_service_dependencies["billing_service"].get_info.return_value = {
-                "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "monthly"},
-                "vector_space": {"size": 0, "limit": 0},
-                "apps": {"size": 5, "limit": 10},
-            }
+            mock_external_service_dependencies["billing_service"].get_info.return_value = builders.subscription_info(
+                subscription=builders.subscription_plan("professional", interval="month"),
+                vector_space=builders.vector_space(size=0, limit=0),
+                apps=billing.SubscriptionQuota(size=5, limit=10),
+            )
 
             # Act: Execute the method under test
             result = FeatureService.get_features(tenant_id)
@@ -891,11 +836,11 @@ class TestFeatureService:
             mock_config.MODEL_LB_ENABLED = False
             mock_config.EDUCATION_ENABLED = False
 
-            mock_external_service_dependencies["billing_service"].get_info.return_value = {
-                "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "yearly"},
-                "members": {"size": 10, "limit": 10},
-                "vector_space": {"size": 3, "limit": 5},
-            }
+            mock_external_service_dependencies["billing_service"].get_info.return_value = builders.subscription_info(
+                subscription=builders.subscription_plan("professional", interval="year"),
+                members=billing.SubscriptionQuota(size=10, limit=10),
+                vector_space=builders.vector_space(size=3, limit=5),
+            )
 
             # Act: Execute the method under test
             result = FeatureService.get_features(tenant_id)
@@ -1207,14 +1152,14 @@ class TestFeatureService:
             mock_config.MODEL_LB_ENABLED = False
             mock_config.EDUCATION_ENABLED = False
 
-            mock_external_service_dependencies["billing_service"].get_info.return_value = {
-                "subscription": {"plan": CloudPlan.TEAM, "interval": "yearly"},
-                "members": {"size": 0, "limit": 0},
-                "apps": {"size": 0, "limit": -1},
-                "vector_space": {"size": 0, "limit": 999999},
-                "documents_upload_quota": {"size": 0, "limit": 0},
-                "annotation_quota_limit": {"size": 0, "limit": 1},
-            }
+            mock_external_service_dependencies["billing_service"].get_info.return_value = builders.subscription_info(
+                subscription=builders.subscription_plan("team", interval="year"),
+                members=billing.SubscriptionQuota(size=0, limit=0),
+                apps=billing.SubscriptionQuota(size=0, limit=-1),
+                vector_space=builders.vector_space(size=0, limit=999999),
+                documents_upload_quota=billing.SubscriptionQuota(size=0, limit=0),
+                annotation_quota_limit=billing.SubscriptionQuota(size=0, limit=1),
+            )
 
             # Act: Execute the method under test
             result = FeatureService.get_features(tenant_id)
@@ -1310,14 +1255,14 @@ class TestFeatureService:
         """
         # Arrange: Setup edge case education mock
         tenant_id = self._create_test_tenant_id()
-        mock_external_service_dependencies["billing_service"].get_info.return_value = {
-            "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "semester", "education": True},
-            "members": {"size": 100, "limit": 200},
-            "apps": {"size": 50, "limit": 100},
-            "vector_space": {"size": 20, "limit": 50},
-            "documents_upload_quota": {"size": 500, "limit": 1000},
-            "annotation_quota_limit": {"size": 200, "limit": 500},
-        }
+        mock_external_service_dependencies["billing_service"].get_info.return_value = builders.subscription_info(
+            subscription=builders.subscription_plan("professional", interval="year", education=True),
+            members=billing.SubscriptionQuota(size=100, limit=200),
+            apps=billing.SubscriptionQuota(size=50, limit=100),
+            vector_space=builders.vector_space(size=20, limit=50),
+            documents_upload_quota=billing.SubscriptionQuota(size=500, limit=1000),
+            annotation_quota_limit=billing.SubscriptionQuota(size=200, limit=500),
+        )
 
         with patch("services.feature_service.dify_config") as mock_config:
             mock_config.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD
@@ -1523,10 +1468,10 @@ class TestFeatureService:
             mock_config.MODEL_LB_ENABLED = False
             mock_config.EDUCATION_ENABLED = False
 
-            mock_external_service_dependencies["billing_service"].get_info.return_value = {
-                "subscription": {"plan": CloudPlan.TEAM, "interval": "yearly"},
-                "annotation_quota_limit": {"size": 999, "limit": 1000},
-            }
+            mock_external_service_dependencies["billing_service"].get_info.return_value = builders.subscription_info(
+                subscription=builders.subscription_plan("team", interval="year"),
+                annotation_quota_limit=billing.SubscriptionQuota(size=999, limit=1000),
+            )
 
             # Act: Execute the method under test
             result = FeatureService.get_features(tenant_id)
@@ -1577,13 +1522,10 @@ class TestFeatureService:
             mock_config.MODEL_LB_ENABLED = False
             mock_config.EDUCATION_ENABLED = False
 
-            mock_external_service_dependencies["billing_service"].get_info.return_value = {
-                "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "monthly"},
-                "documents_upload_quota": {
-                    "size": 0,  # Edge case: zero current size
-                    "limit": 0,  # Edge case: zero limit
-                },
-            }
+            mock_external_service_dependencies["billing_service"].get_info.return_value = builders.subscription_info(
+                subscription=builders.subscription_plan("professional", interval="month"),
+                documents_upload_quota=billing.SubscriptionQuota(size=0, limit=0),
+            )
 
             # Act: Execute the method under test
             result = FeatureService.get_features(tenant_id)
@@ -1685,13 +1627,9 @@ class TestFeatureService:
             mock_config.MODEL_LB_ENABLED = False
             mock_config.EDUCATION_ENABLED = False
 
-            mock_external_service_dependencies["billing_service"].get_info.return_value = {
-                "subscription": {
-                    "plan": CloudPlan.PROFESSIONAL,
-                    "interval": "monthly",
-                    "education": False,  # Education explicitly disabled
-                },
-            }
+            mock_external_service_dependencies["billing_service"].get_info.return_value = builders.subscription_info(
+                subscription=builders.subscription_plan("professional", interval="month", education=False)
+            )
 
             # Act: Execute the method under test
             result = FeatureService.get_features(tenant_id)

api/tests/e2e_tests/containers/services/test_messages_clean_service.py

@@ -42,6 +42,7 @@ from services.retention.conversation.messages_clean_policy import (
     create_message_clean_policy,
 )
 from services.retention.conversation.messages_clean_service import MessagesCleanService
+from tests.helpers import billing_builders as builders
 
 
 class TestMessagesCleanServiceIntegration:
@@ -451,15 +452,9 @@ class TestMessagesCleanServiceIntegration:
 
         plan_map = {}
         for tenant in sandbox_tenants:
-            plan_map[tenant.id] = {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": expired_15_days_ago,
-            }
+            plan_map[tenant.id] = builders.plan_item("sandbox", expired_15_days_ago)
         for tenant in paid_tenants:
-            plan_map[tenant.id] = {
-                "plan": CloudPlan.PROFESSIONAL,
-                "expiration_date": now_timestamp + (365 * 24 * 60 * 60),  # Active for 1 year
-            }
+            plan_map[tenant.id] = builders.plan_item("professional", now_timestamp + (365 * 24 * 60 * 60))
 
         with patch("services.billing_service.BillingService.get_plan_bulk") as mock_billing:
             mock_billing.return_value = plan_map
@@ -525,12 +520,7 @@ class TestMessagesCleanServiceIntegration:
             message_ids.append(msg.id)
 
         with patch("services.billing_service.BillingService.get_plan_bulk") as mock_billing:
-            mock_billing.return_value = {
-                tenant.id: {
-                    "plan": CloudPlan.SANDBOX,
-                    "expiration_date": -1,  # No previous subscription
-                }
-            }
+            mock_billing.return_value = {tenant.id: builders.plan_item("sandbox", -1)}
 
             # Act - Use small batch size to trigger multiple batches
             end_before = datetime.datetime.now() - datetime.timedelta(days=30)
@@ -569,12 +559,7 @@ class TestMessagesCleanServiceIntegration:
             message_ids.append(msg.id)
 
         with patch("services.billing_service.BillingService.get_plan_bulk") as mock_billing:
-            mock_billing.return_value = {
-                tenant.id: {
-                    "plan": CloudPlan.SANDBOX,
-                    "expiration_date": -1,  # No previous subscription
-                }
-            }
+            mock_billing.return_value = {tenant.id: builders.plan_item("sandbox", -1)}
 
             # Act
             end_before = datetime.datetime.now() - datetime.timedelta(days=30)
@@ -625,14 +610,8 @@ class TestMessagesCleanServiceIntegration:
 
         # Only tenant[0] is confirmed as sandbox, tenant[1] is professional, tenant[2] is missing
         partial_plan_map = {
-            tenants_data[0]["tenant"].id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": -1,  # No previous subscription
-            },
-            tenants_data[1]["tenant"].id: {
-                "plan": CloudPlan.PROFESSIONAL,
-                "expiration_date": now_timestamp + (365 * 24 * 60 * 60),  # Active for 1 year
-            },
+            tenants_data[0]["tenant"].id: builders.plan_item("sandbox", -1),
+            tenants_data[1]["tenant"].id: builders.plan_item("professional", now_timestamp + (365 * 24 * 60 * 60)),
             # tenants_data[2] is missing from response
         }
 
@@ -760,12 +739,7 @@ class TestMessagesCleanServiceIntegration:
 
         # Mock billing service
         with patch("services.billing_service.BillingService.get_plan_bulk") as mock_billing:
-            mock_billing.return_value = {
-                tenant.id: {
-                    "plan": CloudPlan.SANDBOX,
-                    "expiration_date": -1,  # No previous subscription
-                }
-            }
+            mock_billing.return_value = {tenant.id: builders.plan_item("sandbox", -1)}
 
             # Act - Clean with specific time range [2024-01-10, 2024-01-20)
             policy = create_message_clean_policy(graceful_period_days=21)
@@ -859,26 +833,11 @@ class TestMessagesCleanServiceIntegration:
 
         # Mock billing service with all scenarios
         plan_map = {
-            tenant1.id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": expired_5_days_ago,
-            },
-            tenant2.id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": expired_10_days_ago,
-            },
-            tenant3.id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": -1,
-            },
-            tenant4.id: {
-                "plan": CloudPlan.PROFESSIONAL,
-                "expiration_date": future_expiration,
-            },
-            tenant5.id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": expired_exactly_8_days_ago,
-            },
+            tenant1.id: builders.plan_item("sandbox", expired_5_days_ago),
+            tenant2.id: builders.plan_item("sandbox", expired_10_days_ago),
+            tenant3.id: builders.plan_item("sandbox", -1),
+            tenant4.id: builders.plan_item("professional", future_expiration),
+            tenant5.id: builders.plan_item("sandbox", expired_exactly_8_days_ago),
         }
 
         with patch("services.billing_service.BillingService.get_plan_bulk") as mock_billing:
@@ -939,18 +898,9 @@ class TestMessagesCleanServiceIntegration:
 
         # Mock billing service - all tenants are sandbox with no subscription
         plan_map = {
-            tenants_data[0]["tenant"].id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": -1,
-            },
-            tenants_data[1]["tenant"].id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": -1,
-            },
-            tenants_data[2]["tenant"].id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": -1,
-            },
+            tenants_data[0]["tenant"].id: builders.plan_item("sandbox", -1),
+            tenants_data[1]["tenant"].id: builders.plan_item("sandbox", -1),
+            tenants_data[2]["tenant"].id: builders.plan_item("sandbox", -1),
         }
 
         # Setup whitelist - tenant0 and tenant1 are whitelisted, tenant2 is not
@@ -1021,12 +971,7 @@ class TestMessagesCleanServiceIntegration:
         db_session_with_containers.commit()
 
         with patch("services.billing_service.BillingService.get_plan_bulk") as mock_billing:
-            mock_billing.return_value = {
-                tenant.id: {
-                    "plan": CloudPlan.SANDBOX,
-                    "expiration_date": -1,
-                }
-            }
+            mock_billing.return_value = {tenant.id: builders.plan_item("sandbox", -1)}
 
             # Act - Use from_days to clean messages older than 30 days
             policy = create_message_clean_policy(graceful_period_days=21)
@@ -1075,14 +1020,8 @@ class TestMessagesCleanServiceIntegration:
 
         # Mock billing service
         plan_map = {
-            tenant1.id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": expired_30_days_ago,  # Beyond grace period
-            },
-            tenant2.id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": expired_10_days_ago,  # Within grace period
-            },
+            tenant1.id: builders.plan_item("sandbox", expired_30_days_ago),
+            tenant2.id: builders.plan_item("sandbox", expired_10_days_ago),
         }
 
         # Setup whitelist - only tenant1 is whitelisted
@@ -1134,12 +1073,7 @@ class TestMessagesCleanServiceIntegration:
             msg_ids.append(msg.id)
 
         # Mock billing service
-        plan_map = {
-            tenant.id: {
-                "plan": CloudPlan.SANDBOX,
-                "expiration_date": -1,
-            }
-        }
+        plan_map = {tenant.id: builders.plan_item("sandbox", -1)}
 
         # Setup empty whitelist (default behavior from fixture)
         mock_whitelist.return_value = []

api/tests/helpers/billing_builders.py

@@ -0,0 +1,41 @@
+"""Generated billing models with defaults, so tests state only the fields they care about."""
+
+from services.billing_client import billing
+
+
+def subscription_plan(plan: str = "sandbox", **fields: object) -> billing.BillingSubscriptionPlan:
+    defaults: dict[str, object] = {
+        "interval": "",
+        "education": False,
+        "status": "",
+        "cancel_at_period_end": False,
+        "paid_period_end": 0,
+    }
+    return billing.BillingSubscriptionPlan.model_validate({"plan": plan} | defaults | fields)
+
+
+def subscription_info(plan: str = "sandbox", **fields: object) -> billing.GetSubscriptionInfoReply:
+    defaults: dict[str, object] = {
+        "enabled": True,
+        "subscription": subscription_plan(plan),
+        "members": billing.SubscriptionQuota(size=0, limit=1),
+        "apps": billing.SubscriptionQuota(size=0, limit=10),
+        "vector_space": None,
+        "knowledge_rate_limit": billing.KnowledgeRateLimit(size=0, limit=10),
+        "documents_upload_quota": billing.SubscriptionQuota(size=0, limit=50),
+        "annotation_quota_limit": billing.SubscriptionQuota(size=0, limit=10),
+        "docs_processing": "standard",
+        "can_replace_logo": False,
+        "model_load_balancing_enabled": False,
+        "knowledge_pipeline_publish_enabled": False,
+        "next_credit_reset_date": 0,
+    }
+    return billing.GetSubscriptionInfoReply.model_validate(defaults | fields)
+
+
+def vector_space(size: float = 0.0, limit: int = 0, usage_unknown: bool | None = None) -> billing.VectorSpaceQuota:
+    return billing.VectorSpaceQuota(size=size, limit=limit, usage_unknown=usage_unknown)
+
+
+def plan_item(plan: str = "sandbox", expiration_date: int = -1) -> billing.SubscriptionPlanItem:
+    return billing.SubscriptionPlanItem.model_validate({"plan": plan, "expiration_date": expiration_date})

api/tests/unit_tests/controllers/console/app/test_app_import.py

@@ -39,7 +39,7 @@ from services.agent.roster_package_importer import RosterAgentPackageImportResul
 from services.app.console_gateway import AppTransferGateway, EnterpriseConsoleAppAccess
 from services.app_dsl_service import AppDslService, PendingData
 from services.app_package_service import AppPackageService, PreparedAppPackage
-from services.billing_service import BillingService, _BillingHTTPStatusError
+from services.billing_client import transport
 from services.enterprise.enterprise_service import EnterpriseService
 from services.entities.dsl_entities import AppImportPackage, AppImportParams, DslImportWarning, Import, ImportStatus
 from services.entities.feature_entities import FeatureModel, LimitationModel
@@ -869,14 +869,14 @@ class TestAppImportQuotaCases:
             )
         else:
 
-            def failed_billing(*_args: object, **_kwargs: object) -> dict[str, str]:
+            def failed_billing(*_args: object, **_kwargs: object) -> httpx.Response:
                 if failure == "timeout":
                     raise httpx.ReadTimeout("Private billing timeout details")
                 if failure == "upstream_503":
-                    raise _BillingHTTPStatusError("Private billing failure details", 503)
-                return {"subscription": "Private invalid billing response"}
+                    return httpx.Response(503, text="Private billing failure details")
+                return httpx.Response(200, json={"subscription": "Private invalid billing response"})
 
-            monkeypatch.setattr(BillingService, "_send_request", failed_billing)
+            monkeypatch.setattr(transport, "_send", failed_billing)
 
         app = Flask(__name__)
         app.config["TESTING"] = True

api/tests/unit_tests/controllers/console/test_workflow_run_archive.py

@@ -23,11 +23,12 @@ from controllers.console.workflow_run_archive import (
     WorkflowRunArchiveMonthApi,
     WorkflowRunArchivesApi,
 )
-from enums import CloudPlan, DeploymentEdition
+from enums import DeploymentEdition
 from extensions.ext_application_services import ApplicationServices
 from libs.login import AccountWithTenant
 from machinery.context import RequestContext
 from models import Account, TenantAccountRole
+from services.billing_client import billing
 from services.retention.workflow_run.archive_deletion_task import (
     WorkflowRunArchiveDeletionStatus,
     WorkflowRunArchiveDeletionTask,
@@ -39,6 +40,7 @@ from services.retention.workflow_run.archive_log_service import (
     WorkflowRunArchiveNotFoundError,
     WorkflowRunArchiveService,
 )
+from tests.helpers import billing_builders as builders
 from tests.helpers.unit.config_override import apply_config_overrides
 from tests.helpers.unit.model_factories import make_account
 
@@ -162,12 +164,10 @@ def test_workflow_run_archive_endpoint_allows_admitted_role_when_rbac_is_enabled
     monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: account_with_tenant)
     monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", lambda: account_with_tenant)
 
-    def get_billing_info(tenant_id: str, *, exclude_vector_space: bool) -> dict[str, object]:
+    def get_billing_info(tenant_id: str, *, exclude_vector_space: bool) -> billing.GetSubscriptionInfoReply:
         assert tenant_id == "tenant-1"
         assert exclude_vector_space
-        return {
-            "subscription": {"plan": CloudPlan.TEAM},
-        }
+        return builders.subscription_info("team")
 
     monkeypatch.setattr(
         "controllers.console.wraps.BillingService.get_info",
@@ -203,7 +203,7 @@ def test_workflow_run_archive_role_permissions_when_rbac_is_disabled(
     monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", lambda: identity)
     monkeypatch.setattr(
         "controllers.console.wraps.BillingService.get_info",
-        lambda _tenant_id, **_kwargs: {"subscription": {"plan": CloudPlan.TEAM}},
+        lambda _tenant_id, **_kwargs: builders.subscription_info("team"),
     )
     task = object()
     archives = object()

api/tests/unit_tests/controllers/console/test_wraps.py

@@ -3,7 +3,6 @@ from types import SimpleNamespace
 from typing import override
 from unittest.mock import MagicMock, patch
 
-import httpx
 import pytest
 from flask import Flask
 from flask_login import LoginManager, UserMixin
@@ -50,6 +49,8 @@ from models import Account, DifySetup
 from models.account import AccountStatus, TenantAccountRole
 from models.dataset import RateLimitLog
 from services.entities.feature_entities import LicenseStatus
+from services.errors.billing import BillingUpstreamUnavailableError
+from tests.helpers import billing_builders as builders
 from tests.helpers.unit.config_override import config_overrides_context
 
 
@@ -815,7 +816,7 @@ class TestBillingPaidPlanRequired:
         def paid_view():
             return "paid_success"
 
-        billing_info = {"subscription": {"plan": plan}}
+        billing_info = builders.subscription_info(plan)
         with (
             patch(
                 "controllers.console.wraps.current_account_with_tenant",
@@ -836,7 +837,9 @@ class TestBillingPaidPlanRequired:
         def paid_view():
             return "paid_success"
 
-        billing_info = {"subscription": {"plan": plan}}
+        billing_info = builders.subscription_info(
+            subscription=builders.subscription_plan().model_copy(update={"plan": plan})
+        )
         with app.test_request_context():
             with (
                 patch(
@@ -866,7 +869,7 @@ class TestBillingPaidPlanRequired:
                 ),
                 patch(
                     "controllers.console.wraps.BillingService.get_info",
-                    side_effect=httpx.ConnectError("billing unavailable"),
+                    side_effect=BillingUpstreamUnavailableError("billing unavailable"),
                 ),
                 pytest.raises(HTTPException) as exc_info,
             ):

api/tests/unit_tests/controllers/console/workspace/test_model_providers.py

@@ -47,6 +47,7 @@ from graphon.model_runtime.errors.validate import CredentialsValidateFailedError
 from models import Account
 from models.account import TenantAccountRole
 from models.provider import ProviderType
+from services.billing_client import billing
 from services.entities.model_provider_entities import (
     CustomConfigurationResponse,
     CustomConfigurationStatus,
@@ -703,7 +704,7 @@ class TestModelProviderPaymentCheckoutUrlApi:
             app.test_request_context("/"),
             patch(
                 "controllers.console.workspace.model_providers.BillingService.get_model_provider_payment_link",
-                return_value={"payment_link": "https://payment.example.com/provider"},
+                return_value=billing.GetModelProviderPaymentLinkReply(url="https://payment.example.com/provider"),
             ) as get_model_provider_payment_link,
         ):
             result = method(api, "tenant1", user, provider="anthropic")

api/tests/unit_tests/controllers/service_api/app/test_completion.py

@@ -7,6 +7,7 @@ from unittest.mock import MagicMock, Mock, patch
 import pytest
 from flask import Flask
 from pydantic import ValidationError
+from pytest_mock import MockerFixture
 from sqlalchemy import Engine
 from sqlalchemy.orm import Session
 from werkzeug.exceptions import BadRequest, NotFound
@@ -47,6 +48,7 @@ from services.conversation_service import ConversationService
 from services.errors.app import IsDraftWorkflowError, WorkflowIdFormatError, WorkflowNotFoundError
 from services.errors.conversation import ConversationNotExistsError
 from services.errors.llm import InvokeRateLimitError
+from tests.helpers import billing_builders as builders
 from tests.helpers.unit.config_override import apply_config_overrides
 from tests.helpers.unit.model_factories import make_app, make_end_user
 
@@ -620,14 +622,15 @@ class TestCompletionCases:
             assert exc_info.value.description == "rate limit reached"
 
         def test_rejects_sandbox_plan_workflow_version(
-            self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session
+            self, app: Flask, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, orm_session: Session
         ) -> None:
             sys.modules["controllers.service_api.app.completion"]
             apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD)
 
-            billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}})
+            billing_get_info = mocker.patch.object(
+                BillingService, "get_info", return_value=builders.subscription_info("sandbox")
+            )
             generate = Mock()
-            monkeypatch.setattr(BillingService, "get_info", billing_get_info)
             monkeypatch.setattr(AppGenerateService, "generate", generate)
 
             api = ChatApi()
@@ -662,6 +665,7 @@ class TestCompletionCases:
             self,
             app: Flask,
             monkeypatch: pytest.MonkeyPatch,
+            mocker: MockerFixture,
             orm_session: Session,
             deployment_edition: DeploymentEdition,
             plan: CloudPlan,
@@ -670,9 +674,10 @@ class TestCompletionCases:
             completion_module = sys.modules["controllers.service_api.app.completion"]
             apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=deployment_edition)
 
-            billing_get_info = Mock(return_value={"subscription": {"plan": plan}})
+            billing_get_info = mocker.patch.object(
+                BillingService, "get_info", return_value=builders.subscription_info(plan)
+            )
             generate = Mock(return_value={"result": "ok"})
-            monkeypatch.setattr(BillingService, "get_info", billing_get_info)
             monkeypatch.setattr(AppGenerateService, "generate", generate)
             monkeypatch.setattr(completion_module.helper, "compact_generate_response", lambda response: response)
 

api/tests/unit_tests/controllers/service_api/app/test_workflow.py

@@ -7,8 +7,10 @@ from inspect import unwrap
 from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock
 
+import httpx
 import pytest
 from flask import Flask
+from pytest_mock import MockerFixture
 from sqlalchemy.engine import Engine
 from sqlalchemy.orm import Session, sessionmaker
 from werkzeug.exceptions import BadRequest, NotFound
@@ -40,6 +42,7 @@ from models.model import App, AppMode, EndUser
 from models.workflow import WorkflowAppLog, WorkflowAppLogCreatedFrom, WorkflowRun, WorkflowType
 from repositories.workflow_app_log_query_repository import WorkflowAppLogQueryRepository
 from services.app_generate_service import AppGenerateService
+from services.billing_client import transport
 from services.billing_service import BillingService
 from services.errors.app import (
     IsDraftWorkflowError,
@@ -48,8 +51,10 @@ from services.errors.app import (
 from services.errors.app import (
     TriggerWorkflowServiceModeUnavailableError as TriggerWorkflowServiceModeUnavailableServiceError,
 )
+from services.errors.billing import BillingUpstreamInvalidResponseError
 from services.errors.llm import InvokeRateLimitError
 from services.workflow_app_log_query_service import WorkflowAppLogQueryService
+from tests.helpers import billing_builders as builders
 
 """
 Unit tests for Service API Workflow controllers.
@@ -457,6 +462,7 @@ class TestWorkflowCases:
             self,
             app: Flask,
             monkeypatch: pytest.MonkeyPatch,
+            mocker: MockerFixture,
             sqlite_session: Session,
             config_overrides: Callable[..., None],
             api_class: type[WorkflowRunApi | WorkflowRunByIdApi],
@@ -466,9 +472,7 @@ class TestWorkflowCases:
             error_code: str,
         ) -> None:
             config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD)
-            monkeypatch.setattr(
-                BillingService, "get_info", Mock(return_value={"subscription": {"plan": CloudPlan.PROFESSIONAL}})
-            )
+            mocker.patch.object(BillingService, "get_info", return_value=builders.subscription_info("professional"))
             generate = Mock(side_effect=source_error("rate limit reached"))
             monkeypatch.setattr(AppGenerateService, "generate", generate)
 
@@ -519,15 +523,17 @@ class TestWorkflowCases:
             self,
             app: Flask,
             monkeypatch: pytest.MonkeyPatch,
+            mocker: MockerFixture,
             sqlite_session: Session,
             config_overrides: Callable[..., None],
         ) -> None:
             config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD)
             sys.modules["controllers.service_api.app.workflow"]
 
-            billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}})
+            billing_get_info = mocker.patch.object(
+                BillingService, "get_info", return_value=builders.subscription_info("sandbox")
+            )
             generate = Mock(return_value={"result": "ok"})
-            monkeypatch.setattr(BillingService, "get_info", billing_get_info)
             monkeypatch.setattr(AppGenerateService, "generate", generate)
 
             api = WorkflowRunApi()
@@ -546,6 +552,37 @@ class TestWorkflowCases:
             generate.assert_called_once()
 
     class TestWorkflowRunByIdApi:
+        def test_missing_billing_plan_cannot_bypass_sandbox_restriction(
+            self,
+            app: Flask,
+            monkeypatch: pytest.MonkeyPatch,
+            sqlite_session: Session,
+            config_overrides: Callable[..., None],
+        ) -> None:
+            config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD)
+            payload = builders.subscription_info().model_dump()
+            del payload["subscription"]["plan"]
+            monkeypatch.setenv("BILLING_API_URL", "http://billing/v1")
+
+            def unexpected_generation(*_args: object, **_kwargs: object) -> None:
+                pytest.fail("Invalid billing data must be rejected before workflow execution")
+
+            monkeypatch.setattr(AppGenerateService, "generate", unexpected_generation)
+            api = WorkflowRunByIdApi()
+            with httpx.Client(transport=httpx.MockTransport(lambda _: httpx.Response(200, json=payload))) as client:
+                monkeypatch.setattr(transport, "_http_client", client)
+                with (
+                    app.test_request_context("/workflows/w1/run", method="POST", json={"inputs": {}}),
+                    pytest.raises(BillingUpstreamInvalidResponseError),
+                ):
+                    unwrap(api.post)(
+                        api,
+                        session=sqlite_session,
+                        app_model=_make_app_model(),
+                        end_user=_make_end_user(),
+                        workflow_id="w1",
+                    )
+
         def test_trigger_workflow_version_returns_stable_unavailable_error(
             self,
             app: Flask,
@@ -577,15 +614,17 @@ class TestWorkflowCases:
             self,
             app: Flask,
             monkeypatch: pytest.MonkeyPatch,
+            mocker: MockerFixture,
             sqlite_session: Session,
             config_overrides: Callable[..., None],
         ) -> None:
             config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD)
             sys.modules["controllers.service_api.app.workflow"]
 
-            billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}})
+            billing_get_info = mocker.patch.object(
+                BillingService, "get_info", return_value=builders.subscription_info("sandbox")
+            )
             generate = Mock()
-            monkeypatch.setattr(BillingService, "get_info", billing_get_info)
             monkeypatch.setattr(AppGenerateService, "generate", generate)
 
             api = WorkflowRunByIdApi()
@@ -627,6 +666,7 @@ class TestWorkflowCases:
             self,
             app: Flask,
             monkeypatch: pytest.MonkeyPatch,
+            mocker: MockerFixture,
             deployment_edition: DeploymentEdition,
             plan: CloudPlan,
             sqlite_session: Session,
@@ -634,9 +674,10 @@ class TestWorkflowCases:
         ) -> None:
             config_overrides(DEPLOYMENT_EDITION=deployment_edition)
 
-            billing_get_info = Mock(return_value={"subscription": {"plan": plan}})
+            billing_get_info = mocker.patch.object(
+                BillingService, "get_info", return_value=builders.subscription_info(plan)
+            )
             generate = Mock(return_value={"result": "ok"})
-            monkeypatch.setattr(BillingService, "get_info", billing_get_info)
             monkeypatch.setattr(AppGenerateService, "generate", generate)
 
             api = WorkflowRunByIdApi()

api/tests/unit_tests/services/billing_client/test_billing.py

@@ -42,13 +42,23 @@ def test_request_body_is_serialized_from_the_model(respond_with: RespondWith) ->
     assert body.model_validate_json(requests[0].content) == body
 
 
-def test_missing_required_field_is_an_invalid_response(respond_with: RespondWith) -> None:
+def test_missing_required_balance_field_is_an_invalid_response(respond_with: RespondWith) -> None:
     respond_with(httpx.Response(200, json={"available": "1"}))
 
     with pytest.raises(BillingUpstreamInvalidResponseError):
         billing.quota_get_balance(tenant_id="t1", feature_key="api_rate_limit", bucket="")
 
 
+def test_invalid_reservation_id_is_rejected_before_business_decisions(respond_with: RespondWith) -> None:
+    respond_with(httpx.Response(200, json={"reservation_id": 123, "available": 99, "reserved": 1}))
+    body = billing.QuotaReserveReq(
+        tenant_id="t1", feature_key="trigger_event", request_id="req", amount=1, meta={}, bucket=""
+    )
+
+    with pytest.raises(BillingUpstreamInvalidResponseError):
+        billing.quota_reserve(body=body)
+
+
 def test_omitted_optional_field_parses_as_none(respond_with: RespondWith) -> None:
     respond_with(httpx.Response(200, json={"size": 1.25, "limit": "50"}))
 

api/tests/unit_tests/services/billing_client/test_transport.py

@@ -1,8 +1,10 @@
+import logging
+
 import httpx
 import pytest
 import tenacity
 
-from services.billing_client import transport
+from services.billing_client import billing, transport
 from services.billing_client.transport import BillingModel, Upstream
 from services.errors.billing import BillingUpstreamInvalidResponseError, BillingUpstreamUnavailableError
 from tests.unit_tests.services.billing_client.conftest import RespondWith
@@ -85,3 +87,106 @@ def test_upstream_selects_base_url(
 
     assert _get(upstream).ok is True
     assert str(requests[0].url) == expected
+
+
+def test_invalid_body_is_reported_and_kept_raw_while_observed(
+    respond_with: RespondWith, caplog: pytest.LogCaptureFixture
+) -> None:
+    respond_with(
+        httpx.Response(
+            200,
+            json={
+                "enabled": True,
+                "subscription": {"plan": "enterprise", "interval": "month", "education": "yes"},
+                "members": {"size": "3", "limit": 5},
+                "vector_space": None,
+                "next_credit_reset_date": "1767225600",
+            },
+        )
+    )
+
+    with caplog.at_level(logging.ERROR, logger=transport.__name__):
+        info = transport.request(
+            Upstream.BILLING,
+            "GET",
+            "/probe",
+            response_model=billing.GetSubscriptionInfoReply,
+            path_params={},
+            params={},
+            body=None,
+        )
+
+    assert "GET /probe" in caplog.text
+    assert caplog.records[0].exc_info is not None
+    assert info.subscription is not None
+    assert info.subscription.plan == "enterprise"
+    assert info.subscription.interval == "month"
+    assert info.subscription.education is True
+    assert info.subscription.paid_period_end is None
+    assert info.members == billing.SubscriptionQuota(size=3, limit=5)
+    assert info.vector_space is None
+    assert info.next_credit_reset_date == 1767225600
+    assert info.apps is None
+
+
+@pytest.mark.parametrize(
+    ("path", "response_model"),
+    [
+        ("/subscription/info", billing.GetSubscriptionInfoReply),
+        ("/subscription/vector-space", billing.VectorSpaceQuota),
+        ("/subscription/knowledge-rate-limit", billing.GetKnowledgeRateLimitReply),
+        ("/subscription/payment-link", billing.GetSubscriptionPaymentLinkReply),
+        ("/subscription/cleanup/whitelist", billing.ListSubscriptionCleanupWhitelistReply),
+        ("/model-provider/payment-link", billing.GetModelProviderPaymentLinkReply),
+        ("/invoices", billing.GetInvoicesReply),
+    ],
+)
+def test_migrated_reads_reject_missing_required_fields(
+    respond_with: RespondWith, path: str, response_model: type[BillingModel]
+) -> None:
+    respond_with(httpx.Response(200, json={}))
+
+    with pytest.raises(BillingUpstreamInvalidResponseError):
+        transport.request(
+            Upstream.BILLING, "GET", path, response_model=response_model, path_params={}, params={}, body=None
+        )
+
+
+@pytest.mark.parametrize(
+    ("method", "path", "response_model"),
+    [
+        ("GET", "/quota/info", billing.QuotaGetInfoReply),
+        ("POST", "/quota/reserve", billing.QuotaReserveReply),
+        ("POST", "/quota/commit", billing.QuotaCommitReply),
+        ("POST", "/quota/release", billing.QuotaReleaseReply),
+        ("GET", "/quota/balance", billing.QuotaGetBalanceReply),
+        ("POST", "/quota/consume-capped", billing.QuotaConsumeCappedReply),
+    ],
+)
+def test_quota_responses_reject_missing_required_fields(
+    respond_with: RespondWith, method: transport.HTTPMethod, path: str, response_model: type[BillingModel]
+) -> None:
+    respond_with(httpx.Response(200, json={}))
+
+    with pytest.raises(BillingUpstreamInvalidResponseError):
+        transport.request(
+            Upstream.BILLING_QUOTA, method, path, response_model=response_model, path_params={}, params={}, body=None
+        )
+
+
+def test_invalid_body_is_an_invalid_response_once_enforced(
+    respond_with: RespondWith, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(transport, "ENFORCED_RESPONSE_PATHS", frozenset({"/probe"}))
+    respond_with(httpx.Response(200, json={"ok": "maybe"}))
+
+    with pytest.raises(BillingUpstreamInvalidResponseError):
+        _get(Upstream.BILLING)
+
+
+@pytest.mark.parametrize("content", [b"not json", b"[]"])
+def test_non_object_body_is_an_invalid_response(respond_with: RespondWith, content: bytes) -> None:
+    respond_with(httpx.Response(200, content=content))
+
+    with pytest.raises(BillingUpstreamInvalidResponseError):
+        _get(Upstream.BILLING)

api/tests/unit_tests/services/knowledge/documents/test_adapters.py

@@ -651,7 +651,7 @@ class TestDocumentCreationCases:
             outside_transaction()
             raise StopBeforeNetworkError()
 
-        monkeypatch.setattr("services.billing_service._http_client.request", request)
+        monkeypatch.setattr("services.billing_client.transport._http_client.request", request)
         with pytest.raises(StopBeforeNetworkError):
             operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, settings())
         outside_transaction()

api/tests/unit_tests/services/knowledge_fs/test_vector_store.py

@@ -26,6 +26,7 @@ from services.knowledge_fs.vector_store import (
 )
 from services.knowledge_fs.vector_store_elasticsearch import ElasticsearchVectorStore
 from services.knowledge_fs.vector_store_weaviate import WeaviateVectorStore
+from tests.helpers import billing_builders as builders
 
 if TYPE_CHECKING:
     from elasticsearch import Elasticsearch
@@ -34,6 +35,7 @@ if TYPE_CHECKING:
 _KNOWLEDGE_FS_VECTOR_BACKENDS_A, _KNOWLEDGE_FS_VECTOR_BACKENDS_B, _KNOWLEDGE_FS_VECTOR_BACKENDS_GENERATION = (
     f"00000000-0000-4000-8000-{n:012d}" for n in (1, 2, 3)
 )
+
 _KNOWLEDGE_FS_VECTOR_BACKENDS_SCOPE: dict[str, object] = {
     "tenant_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A,
     "knowledge_space_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_B,
@@ -1382,7 +1384,7 @@ class TestKnowledgeFsVectorStoreCases:
             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},
+                return_value=builders.vector_space(size=0, limit=limit, usage_unknown=True),
             ),
             patch.object(service, "_reserve_projected_usage") as reserve,
         ):

api/tests/unit_tests/services/retention/conversation/test_messages_clean_policy.py

@@ -1,13 +1,19 @@
 import datetime
 from unittest.mock import MagicMock, patch
 
+import httpx
+import pytest
+
 from enums import DeploymentEdition
+from services.billing_client import transport
+from services.billing_service import BillingService
 from services.retention.conversation.messages_clean_policy import (
     BillingDisabledPolicy,
     BillingSandboxPolicy,
     SimpleMessage,
     create_message_clean_policy,
 )
+from tests.helpers import billing_builders as builders
 
 MODULE = "services.retention.conversation.messages_clean_policy"
 
@@ -46,6 +52,22 @@ class TestBillingSandboxPolicy:
         policy = self._policy({})
         assert policy.filter_message_ids([], {"app1": "t1"}) == []
 
+    def test_malformed_billing_item_does_not_interrupt_other_tenants(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        payload = {
+            "data": {
+                "t1": {"plan": "sandbox", "expiration_date": "-1"},
+                "t2": {"plan": "sandbox"},
+            }
+        }
+        policy = BillingSandboxPolicy(plan_provider=BillingService.get_plan_bulk)
+        monkeypatch.setenv("BILLING_API_URL", "http://billing/v1")
+        with httpx.Client(transport=httpx.MockTransport(lambda _: httpx.Response(200, json=payload))) as client:
+            monkeypatch.setattr(transport, "_http_client", client)
+
+            result = policy.filter_message_ids([_msg("m1", "a1"), _msg("m2", "a2")], {"a1": "t1", "a2": "t2"})
+
+        assert result == ["m1"]
+
     def test_empty_app_to_tenant_returns_empty(self):
         policy = self._policy({})
         assert policy.filter_message_ids([_msg("m1", "app1")], {}) == []
@@ -56,14 +78,14 @@ class TestBillingSandboxPolicy:
         assert policy.filter_message_ids(msgs, {"app1": "t1"}) == []
 
     def test_non_sandbox_tenant_skipped(self):
-        plans = {"t1": {"plan": "professional", "expiration_date": 0}}
+        plans = {"t1": builders.plan_item("professional", 0)}
         policy = self._policy(plans)
         msgs = [_msg("m1", "app1")]
 
         assert policy.filter_message_ids(msgs, {"app1": "t1"}) == []
 
     def test_sandbox_no_previous_subscription_deletes(self):
-        plans = {"t1": {"plan": "sandbox", "expiration_date": -1}}
+        plans = {"t1": builders.plan_item("sandbox", -1)}
         policy = self._policy(plans)
         msgs = [_msg("m1", "app1")]
 
@@ -72,7 +94,7 @@ class TestBillingSandboxPolicy:
     def test_sandbox_expired_beyond_grace_period_deletes(self):
         now = 1_000_000_000
         expired_long_ago = now - (22 * 24 * 60 * 60)  # 22 days ago > 21 day grace
-        plans = {"t1": {"plan": "sandbox", "expiration_date": expired_long_ago}}
+        plans = {"t1": builders.plan_item("sandbox", expired_long_ago)}
         policy = self._policy(plans, now=now)
         msgs = [_msg("m1", "app1")]
 
@@ -81,21 +103,21 @@ class TestBillingSandboxPolicy:
     def test_sandbox_within_grace_period_kept(self):
         now = 1_000_000_000
         expired_recently = now - (10 * 24 * 60 * 60)  # 10 days ago < 21 day grace
-        plans = {"t1": {"plan": "sandbox", "expiration_date": expired_recently}}
+        plans = {"t1": builders.plan_item("sandbox", expired_recently)}
         policy = self._policy(plans, now=now)
         msgs = [_msg("m1", "app1")]
 
         assert policy.filter_message_ids(msgs, {"app1": "t1"}) == []
 
     def test_whitelisted_tenant_skipped(self):
-        plans = {"t1": {"plan": "sandbox", "expiration_date": -1}}
+        plans = {"t1": builders.plan_item("sandbox", -1)}
         policy = self._policy(plans, whitelist=["t1"])
         msgs = [_msg("m1", "app1")]
 
         assert policy.filter_message_ids(msgs, {"app1": "t1"}) == []
 
     def test_message_without_tenant_mapping_skipped(self):
-        plans = {"t1": {"plan": "sandbox", "expiration_date": -1}}
+        plans = {"t1": builders.plan_item("sandbox", -1)}
         policy = self._policy(plans)
         msgs = [_msg("m1", "unmapped_app")]
 
@@ -103,8 +125,8 @@ class TestBillingSandboxPolicy:
 
     def test_mixed_tenants_only_sandbox_deleted(self):
         plans = {
-            "t_sandbox": {"plan": "sandbox", "expiration_date": -1},
-            "t_pro": {"plan": "professional", "expiration_date": 0},
+            "t_sandbox": builders.plan_item("sandbox", -1),
+            "t_pro": builders.plan_item("professional", 0),
         }
         policy = self._policy(plans)
         msgs = [_msg("m1", "app_sandbox"), _msg("m2", "app_pro")]

api/tests/unit_tests/services/retention/conversation/test_messages_clean_service.py

@@ -11,7 +11,7 @@ import pytest
 from sqlalchemy import Engine, delete, func, select
 
 from core.db.session_factory import session_factory
-from enums import CloudPlan, DeploymentEdition
+from enums import DeploymentEdition
 from models import Tenant
 from models.enums import FeedbackFromSource, FeedbackRating
 from models.model import (
@@ -32,6 +32,7 @@ from services.retention.conversation.messages_clean_policy import (
     create_message_clean_policy,
 )
 from services.retention.conversation.messages_clean_service import MessagesCleanService
+from tests.helpers import billing_builders as builders
 
 """SQLite-backed coverage for message retention cleanup."""
 
@@ -440,7 +441,7 @@ class TestMessagesCleanServiceCases:
                 make_simple_message("msg2", "app2"),
             ]
             app_to_tenant = {}  # No mapping
-            tenant_plans = {"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}}
+            tenant_plans = {"tenant1": builders.plan_item("sandbox", -1)}
             plan_provider = make_plan_provider(tenant_plans)
 
             policy = BillingSandboxPolicy(
@@ -488,9 +489,9 @@ class TestMessagesCleanServiceCases:
             ]
             app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant3"}
             tenant_plans = {
-                "tenant1": {"plan": CloudPlan.PROFESSIONAL, "expiration_date": -1},
-                "tenant2": {"plan": CloudPlan.TEAM, "expiration_date": -1},
-                "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},  # Only this one
+                "tenant1": builders.plan_item("professional", -1),
+                "tenant2": builders.plan_item("team", -1),
+                "tenant3": builders.plan_item("sandbox", -1),  # Only this one
             }
             plan_provider = make_plan_provider(tenant_plans)
 
@@ -516,9 +517,9 @@ class TestMessagesCleanServiceCases:
             ]
             app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant3"}
             tenant_plans = {
-                "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
-                "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
-                "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
+                "tenant1": builders.plan_item("sandbox", -1),
+                "tenant2": builders.plan_item("sandbox", -1),
+                "tenant3": builders.plan_item("sandbox", -1),
             }
             plan_provider = make_plan_provider(tenant_plans)
             tenant_whitelist = ["tenant1", "tenant3"]
@@ -545,8 +546,8 @@ class TestMessagesCleanServiceCases:
             ]
             app_to_tenant = {"app1": "tenant1", "app2": "tenant2"}
             tenant_plans = {
-                "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
-                "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
+                "tenant1": builders.plan_item("sandbox", -1),
+                "tenant2": builders.plan_item("sandbox", -1),
             }
             plan_provider = make_plan_provider(tenant_plans)
 
@@ -577,9 +578,9 @@ class TestMessagesCleanServiceCases:
             ]
             app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant3"}
             tenant_plans = {
-                "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_1_day_ago},
-                "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_5_days_ago},
-                "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_7_days_ago},
+                "tenant1": builders.plan_item("sandbox", expired_1_day_ago),
+                "tenant2": builders.plan_item("sandbox", expired_5_days_ago),
+                "tenant3": builders.plan_item("sandbox", expired_7_days_ago),
             }
             plan_provider = make_plan_provider(tenant_plans)
 
@@ -604,7 +605,7 @@ class TestMessagesCleanServiceCases:
             messages = [make_simple_message("msg1", "app1")]
             app_to_tenant = {"app1": "tenant1"}
             tenant_plans = {
-                "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_exactly_8_days_ago},
+                "tenant1": builders.plan_item("sandbox", expired_exactly_8_days_ago),
             }
             plan_provider = make_plan_provider(tenant_plans)
 
@@ -633,8 +634,8 @@ class TestMessagesCleanServiceCases:
             ]
             app_to_tenant = {"app1": "tenant1", "app2": "tenant2"}
             tenant_plans = {
-                "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_9_days_ago},
-                "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_30_days_ago},
+                "tenant1": builders.plan_item("sandbox", expired_9_days_ago),
+                "tenant2": builders.plan_item("sandbox", expired_30_days_ago),
             }
             plan_provider = make_plan_provider(tenant_plans)
 
@@ -655,7 +656,7 @@ class TestMessagesCleanServiceCases:
             # Arrange
             messages: list[SimpleMessage] = []
             app_to_tenant = {"app1": "tenant1"}
-            plan_provider = make_plan_provider({"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}})
+            plan_provider = make_plan_provider({"tenant1": builders.plan_item("sandbox", -1)})
 
             policy = BillingSandboxPolicy(
                 plan_provider=plan_provider,
@@ -721,11 +722,11 @@ class TestMessagesCleanServiceCases:
                 # app5 has no mapping
             }
             tenant_plans = {
-                "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
-                "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": sandbox_expired_old},
-                "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": sandbox_expired_recent},
-                "tenant4": {"plan": CloudPlan.TEAM, "expiration_date": future_expiration},
-                "tenant7": {"plan": CloudPlan.SANDBOX, "expiration_date": sandbox_expired_old},
+                "tenant1": builders.plan_item("sandbox", -1),
+                "tenant2": builders.plan_item("sandbox", sandbox_expired_old),
+                "tenant3": builders.plan_item("sandbox", sandbox_expired_recent),
+                "tenant4": builders.plan_item("team", future_expiration),
+                "tenant7": builders.plan_item("sandbox", sandbox_expired_old),
                 # tenant6 has no plan
             }
             plan_provider = make_plan_provider(tenant_plans)

api/tests/unit_tests/services/retention/workflow_run/test_archive_paid_plan_workflow_run.py

@@ -27,6 +27,7 @@ from services.retention.workflow_run.archive_paid_plan_workflow_run import (
     WorkflowRunArchiver,
 )
 from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_FORMAT, ARCHIVE_BUNDLE_SCHEMA_VERSION
+from tests.helpers import billing_builders as builders
 
 """SQLite and in-process coverage for workflow-run bundle archiving."""
 
@@ -333,9 +334,9 @@ class TestWorkflowRunArchiverSqliteCases:
             archiver = WorkflowRunArchiver(days=90)
 
             mock_bulk = {
-                "t1": {"plan": "professional"},
-                "t2": {"plan": "sandbox"},
-                "t3": {"plan": "team"},
+                "t1": builders.plan_item("professional"),
+                "t2": builders.plan_item("sandbox"),
+                "t3": builders.plan_item("team"),
             }
 
             with (

api/tests/unit_tests/services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py

@@ -6,11 +6,12 @@ from unittest.mock import MagicMock, patch
 import pytest
 from sqlalchemy.orm import Session
 
-from enums import CloudPlan, DeploymentEdition
+from enums import DeploymentEdition
 from repositories.api_workflow_run_repository import WorkflowRunCleanupRef
-from services.billing_service import SubscriptionPlan
+from services.billing_client.billing import SubscriptionPlanItem
 from services.retention.workflow_run import clear_free_plan_expired_workflow_run_logs as cleanup_module
 from services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup
+from tests.helpers import billing_builders as builders
 
 """
 Unit tests for WorkflowRunCleanup service.
@@ -172,25 +173,25 @@ class TestClearFreePlanExpiredWorkflowRunLogsCases:
     class TestIsWithinGracePeriod:
         def test_zero_grace_period_returns_false(self, cleanup):
             cleanup.free_plan_grace_period_days = 0
-            assert cleanup._is_within_grace_period("t1", {"expiration_date": 9999999999}) is False
+            assert cleanup._is_within_grace_period("t1", builders.plan_item(expiration_date=9999999999)) is False
 
         def test_within_grace_period(self, cleanup):
             cleanup.free_plan_grace_period_days = 30
             # expired just 1 day ago
             expired = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=1)
             ts = int(expired.timestamp())
-            assert cleanup._is_within_grace_period("t1", {"expiration_date": ts}) is True
+            assert cleanup._is_within_grace_period("t1", builders.plan_item(expiration_date=ts)) is True
 
         def test_outside_grace_period(self, cleanup):
             cleanup.free_plan_grace_period_days = 5
             # expired 100 days ago
             expired = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=100)
             ts = int(expired.timestamp())
-            assert cleanup._is_within_grace_period("t1", {"expiration_date": ts}) is False
+            assert cleanup._is_within_grace_period("t1", builders.plan_item(expiration_date=ts)) is False
 
         def test_missing_expiration_date_returns_false(self, cleanup):
             cleanup.free_plan_grace_period_days = 30
-            assert cleanup._is_within_grace_period("t1", {"expiration_date": -1}) is False
+            assert cleanup._is_within_grace_period("t1", builders.plan_item(expiration_date=-1)) is False
 
     # ---------------------------------------------------------------------------
     # _get_cleanup_whitelist
@@ -249,8 +250,8 @@ class TestClearFreePlanExpiredWorkflowRunLogsCases:
                 "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService"
             ) as bs:
                 bs.get_plan_bulk_with_cache.return_value = {
-                    "t1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
-                    "t2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
+                    "t1": builders.plan_item("sandbox", -1),
+                    "t2": builders.plan_item("sandbox", -1),
                 }
                 result = c._filter_free_tenants(["t1", "t2"])
             assert "t1" not in result
@@ -264,7 +265,7 @@ class TestClearFreePlanExpiredWorkflowRunLogsCases:
                 "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService"
             ) as bs:
                 bs.get_plan_bulk_with_cache.return_value = {
-                    "t1": {"plan": "professional", "expiration_date": -1},
+                    "t1": builders.plan_item("professional", -1),
                 }
                 result = c._filter_free_tenants(["t1"])
             assert result == set()
@@ -545,8 +546,8 @@ class FakeRepo:
         return result
 
 
-def plan_info(plan: str, expiration: int) -> SubscriptionPlan:
-    return SubscriptionPlan(plan=plan, expiration_date=expiration)
+def plan_info(plan: str, expiration: int) -> SubscriptionPlanItem:
+    return builders.plan_item(plan, expiration)
 
 
 def create_cleanup(
@@ -577,7 +578,7 @@ class TestClearFreePlanExpiredWorkflowRunLogsCasesServices:
     def test_filter_free_tenants_outside_cloud_edition(monkeypatch: pytest.MonkeyPatch) -> None:
         cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10)
 
-        def fail_bulk(_: list[str]) -> dict[str, SubscriptionPlan]:
+        def fail_bulk(_: list[str]) -> dict[str, SubscriptionPlanItem]:
             raise RuntimeError("should not call")
 
         monkeypatch.setattr(cleanup_module.BillingService, "get_plan_bulk_with_cache", staticmethod(fail_bulk))
@@ -623,7 +624,7 @@ class TestClearFreePlanExpiredWorkflowRunLogsCasesServices:
         within_grace_ts = int((now - datetime.timedelta(days=10)).timestamp())
         outside_grace_ts = int((now - datetime.timedelta(days=90)).timestamp())
 
-        def fake_bulk(_: list[str]) -> dict[str, SubscriptionPlan]:
+        def fake_bulk(_: list[str]) -> dict[str, SubscriptionPlanItem]:
             return {
                 "recently_downgraded": plan_info("sandbox", within_grace_ts),
                 "long_sandbox": plan_info("sandbox", outside_grace_ts),
@@ -731,7 +732,7 @@ class TestClearFreePlanExpiredWorkflowRunLogsCasesServices:
 
         billing_calls: list[list[str]] = []
 
-        def fake_bulk(tenant_ids: list[str]) -> dict[str, SubscriptionPlan]:
+        def fake_bulk(tenant_ids: list[str]) -> dict[str, SubscriptionPlanItem]:
             billing_calls.append(tenant_ids)
             return {
                 "t_free": plan_info("sandbox", -1),

api/tests/unit_tests/services/test_billing_service.py

@@ -23,12 +23,13 @@ import pytest
 import tenacity
 from werkzeug.exceptions import InternalServerError
 
-from enums import CloudPlan
+from services.billing_client import billing, transport
 from services.billing_service import BillingService, _BillingHTTPStatusError
 from services.errors.billing import (
     BillingUpstreamInvalidResponseError,
     BillingUpstreamUnavailableError,
 )
+from tests.helpers import billing_builders as builders
 
 TENANT_ID = "11111111-1111-1111-1111-111111111111"
 ACCOUNT_ID = "33333333-3333-3333-3333-333333333333"
@@ -331,243 +332,122 @@ class TestBillingServiceSendRequest:
         assert retry_sleep.call_args_list == [call(2)] * 4
 
 
-class TestBillingServicePortalRequest:
-    def test_sends_get_request(self) -> None:
-        params = {"tenant_id": "tenant-1"}
-        with patch.object(
-            BillingService,
-            "_send_request",
-            return_value={"url": "https://example.com", "ignored": True},
-        ) as send_request:
-            result = BillingService._send_billing_portal_request("/test", params=params)
-
-        assert result == {"url": "https://example.com"}
-        send_request.assert_called_once_with("GET", "/test", params=params)
-
-    def test_invalid_response_shape_is_invalid_upstream_response(self) -> None:
-        with patch.object(BillingService, "_send_request", return_value={}):
-            with pytest.raises(BillingUpstreamInvalidResponseError):
-                BillingService._send_billing_portal_request("/test", params={})
-
-    @pytest.mark.parametrize(
-        "status_code",
-        [httpx.codes.BAD_REQUEST, httpx.codes.UNAUTHORIZED, httpx.codes.NOT_FOUND],
-    )
-    def test_terminal_http_response_is_invalid_upstream_response(self, status_code: int) -> None:
-        with patch.object(
-            BillingService,
-            "_send_request",
-            side_effect=_BillingHTTPStatusError("request failed", status_code),
-        ):
-            with pytest.raises(BillingUpstreamInvalidResponseError):
-                BillingService._send_billing_portal_request("/test", params={})
-
-    @pytest.mark.parametrize(
-        "status_code",
-        [httpx.codes.REQUEST_TIMEOUT, httpx.codes.TOO_MANY_REQUESTS, httpx.codes.INTERNAL_SERVER_ERROR],
-    )
-    def test_retryable_http_response_is_unavailable(self, status_code: int) -> None:
-        with patch.object(
-            BillingService,
-            "_send_request",
-            side_effect=_BillingHTTPStatusError("request failed", status_code),
-        ):
-            with pytest.raises(BillingUpstreamUnavailableError):
-                BillingService._send_billing_portal_request("/test", params={})
-
-    def test_transport_failure_is_unavailable(self) -> None:
-        with patch.object(
-            BillingService,
-            "_send_request",
-            side_effect=httpx.RequestError("network error"),
-        ):
-            with pytest.raises(BillingUpstreamUnavailableError):
-                BillingService._send_billing_portal_request("/test", params={})
-
-    @pytest.mark.parametrize(
-        "decode_error",
-        [
-            json.JSONDecodeError("Expecting value", "", 0),
-            UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid start byte"),
-        ],
-    )
-    def test_invalid_payload_is_invalid_upstream_response(self, decode_error: Exception) -> None:
-        with patch.object(BillingService, "_send_request", side_effect=decode_error):
-            with pytest.raises(BillingUpstreamInvalidResponseError):
-                BillingService._send_billing_portal_request("/test", params={})
-
-    def test_unknown_error_is_not_reclassified(self) -> None:
-        with patch.object(BillingService, "_send_request", side_effect=RuntimeError("programming error")):
-            with pytest.raises(RuntimeError, match="programming error"):
-                BillingService._send_billing_portal_request("/test", params={})
-
-    def test_unknown_value_error_is_not_exposed_as_invalid_request(self) -> None:
-        original_error = ValueError("programming error")
-        with patch.object(BillingService, "_send_request", side_effect=original_error):
-            with pytest.raises(RuntimeError, match="Unexpected billing service value error") as exc_info:
-                BillingService._send_billing_portal_request("/test", params={})
-
-        assert exc_info.value.__cause__ is original_error
-
-
 class TestBillingServiceSubscriptionInfo:
-    """Unit tests for subscription tier and billing info retrieval.
+    """Subscription and usage reads delegate to the generated billing client."""
 
-    Tests cover:
-    - Billing information retrieval
-    - Knowledge base rate limits with default and custom values
-    - Payment link generation for subscriptions and model providers
-    - Invoice retrieval
-    """
+    def test_get_info_returns_the_generated_reply(self):
+        reply = builders.subscription_info("team", vector_space=builders.vector_space(size=1.5, limit=50))
 
-    @pytest.fixture
-    def mock_send_request(self):
-        """Mock _send_request method."""
-        with patch.object(BillingService, "_send_request") as mock:
-            yield mock
+        with patch.object(billing, "subscription_get_subscription_info", return_value=reply) as get_info:
+            result = BillingService.get_info(TENANT_ID, exclude_vector_space=True)
 
-    def test_get_info_success(self, mock_send_request):
-        """Test successful retrieval of billing information."""
-        # Arrange
-        tenant_id = "tenant-123"
-        expected_response = {
-            "subscription": {"plan": "professional", "interval": "month", "education": False},
-            "members": {"size": 1, "limit": 50},
-            "apps": {"size": 1, "limit": 200},
-            "vector_space": {"size": 0.0, "limit": 20480},
-            "documents_upload_quota": {"size": 0, "limit": 1000},
-            "annotation_quota_limit": {"size": 0, "limit": 5000},
-            "can_replace_logo": True,
-            "model_load_balancing_enabled": True,
-            "knowledge_pipeline_publish_enabled": True,
-            "next_credit_reset_date": 1775952000,
-        }
-        mock_send_request.return_value = {
-            **expected_response,
-            "docs_processing": "top-priority",
-            "knowledge_rate_limit": {"limit": 1000},
-        }
+        assert result is reply
+        get_info.assert_called_once_with(tenant_id=TENANT_ID, exclude_vector_space=True)
 
-        # Act
-        result = BillingService.get_info(tenant_id)
+    @pytest.mark.parametrize("bypass_cache", [False, True])
+    def test_get_vector_space_forwards_bypass_cache(self, bypass_cache: bool):
+        reply = builders.vector_space(size=12.5, limit=50, usage_unknown=True)
 
-        # Assert
-        assert result == expected_response
-        mock_send_request.assert_called_once_with("GET", "/subscription/info", params={"tenant_id": tenant_id})
+        with patch.object(billing, "subscription_get_subscription_vector_space", return_value=reply) as get:
+            result = BillingService.get_vector_space(TENANT_ID, bypass_cache=bypass_cache)
 
-    def test_get_info_exclude_vector_space(self, mock_send_request):
-        """When requested, get_info asks billing to skip vector_space."""
-        # Arrange
-        tenant_id = "tenant-123"
-        expected_response = {
-            "subscription": {"plan": "professional", "interval": "month", "education": False},
-            "members": {"size": 1, "limit": 50},
-            "apps": {"size": 1, "limit": 200},
-            "documents_upload_quota": {"size": 0, "limit": 1000},
-            "annotation_quota_limit": {"size": 0, "limit": 5000},
-            "can_replace_logo": True,
-            "model_load_balancing_enabled": True,
-            "knowledge_pipeline_publish_enabled": True,
-        }
-        mock_send_request.return_value = expected_response
+        assert result is reply
+        get.assert_called_once_with(tenant_id=TENANT_ID, bypass_cache=bypass_cache)
 
-        # Act
-        result = BillingService.get_info(tenant_id, exclude_vector_space=True)
+    def test_get_knowledge_rate_limit_returns_the_generated_reply(self):
+        reply = billing.GetKnowledgeRateLimitReply(subscription_plan="professional", limit=100)
 
-        # Assert
-        assert "vector_space" not in result
-        mock_send_request.assert_called_once_with(
-            "GET",
-            "/subscription/info",
-            params={"tenant_id": tenant_id, "exclude_vector_space": "true"},
+        with patch.object(billing, "subscription_get_knowledge_rate_limit", return_value=reply) as get:
+            result = BillingService.get_knowledge_rate_limit(TENANT_ID)
+
+        assert result is reply
+        get.assert_called_once_with(tenant_id=TENANT_ID)
+
+    def test_get_subscription_returns_the_portal_link(self):
+        reply = billing.GetSubscriptionPaymentLinkReply(url="https://payment.example.com/checkout")
+
+        with patch.object(billing, "subscription_get_subscription_payment_link", return_value=reply) as get:
+            result = BillingService.get_subscription("professional", "month", "user@example.com", TENANT_ID)
+
+        assert result == {"url": "https://payment.example.com/checkout"}
+        get.assert_called_once_with(
+            plan="professional", interval="month", prefilled_email="user@example.com", tenant_id=TENANT_ID
         )
 
-    def test_get_info_exclude_vector_space_normalizes_null_field(self, mock_send_request):
-        """When billing serializes skipped vector_space as null, get_info treats it as absent."""
-        # Arrange
-        tenant_id = "tenant-123"
-        expected_response = {
-            "subscription": {"plan": "professional", "interval": "month", "education": False},
-            "members": {"size": 1, "limit": 50},
-            "apps": {"size": 1, "limit": 200},
-            "vector_space": None,
-            "documents_upload_quota": {"size": 0, "limit": 1000},
-            "annotation_quota_limit": {"size": 0, "limit": 5000},
-            "can_replace_logo": True,
-            "model_load_balancing_enabled": True,
-            "knowledge_pipeline_publish_enabled": True,
+    def test_get_invoices_returns_the_portal_link(self):
+        reply = billing.GetInvoicesReply(url="https://billing.example.com/invoices")
+
+        with patch.object(billing, "invoice_get_invoices", return_value=reply) as get:
+            result = BillingService.get_invoices("user@example.com", TENANT_ID)
+
+        assert result == {"url": "https://billing.example.com/invoices"}
+        get.assert_called_once_with(prefilled_email="user@example.com", tenant_id=TENANT_ID)
+
+    def test_get_model_provider_payment_link_returns_the_generated_reply(self):
+        reply = billing.GetModelProviderPaymentLinkReply(url="https://payment.example.com/provider")
+
+        with patch.object(billing, "model_provider_get_model_provider_payment_link", return_value=reply) as get:
+            # pyrefly: ignore [deprecated]
+            result = BillingService.get_model_provider_payment_link("anthropic", TENANT_ID, ACCOUNT_ID, "a@b.com")
+
+        assert result is reply
+        get.assert_called_once_with(
+            provider_name="anthropic", tenant_id=TENANT_ID, account_id=ACCOUNT_ID, prefilled_email="a@b.com"
+        )
+
+    def test_get_plan_bulk_chunks_requests_and_skips_failed_batches(self):
+        tenant_ids = [f"tenant-{i}" for i in range(450)]
+
+        def batch(*, body: billing.BatchGetSubscriptionPlanReq) -> billing.BatchGetSubscriptionPlanReply:
+            if body.tenant_ids[0] == "tenant-200":
+                raise BillingUpstreamUnavailableError("POST /subscription/plan/batch failed")
+            return billing.BatchGetSubscriptionPlanReply(
+                data={tenant_id: builders.plan_item("team", 1735689600) for tenant_id in body.tenant_ids}
+            )
+
+        with patch.object(billing, "subscription_batch_get_subscription_plan", side_effect=batch) as get:
+            result = BillingService.get_plan_bulk(tenant_ids)
+
+        assert [len(c.kwargs["body"].tenant_ids) for c in get.call_args_list] == [200, 200, 50]
+        assert set(result) == set(tenant_ids[:200] + tenant_ids[400:])
+        assert result["tenant-0"] == builders.plan_item("team", 1735689600)
+
+    def test_get_plan_bulk_skips_invalid_items_from_http_response(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        payload = {
+            "data": {
+                "valid": {"plan": "sandbox", "expiration_date": "-1"},
+                "missing_expiration": {"plan": "sandbox"},
+                "invalid_expiration": {"plan": "sandbox", "expiration_date": "invalid"},
+                "missing_plan": {"expiration_date": -1},
+                "invalid_item": "not-an-object",
+                "null_item": None,
+            }
         }
-        mock_send_request.return_value = expected_response
+        monkeypatch.setenv("BILLING_API_URL", "http://billing/v1")
+        with httpx.Client(transport=httpx.MockTransport(lambda _: httpx.Response(200, json=payload))) as client:
+            monkeypatch.setattr(transport, "_http_client", client)
 
-        # Act
-        result = BillingService.get_info(tenant_id, exclude_vector_space=True)
+            plans = BillingService.get_plan_bulk(list(payload["data"]))
 
-        # Assert
-        assert "vector_space" not in result
-        mock_send_request.assert_called_once_with(
-            "GET",
-            "/subscription/info",
-            params={"tenant_id": tenant_id, "exclude_vector_space": "true"},
+        assert plans == {"valid": builders.plan_item("sandbox", -1)}
+
+    def test_get_expired_subscription_cleanup_whitelist_returns_tenant_ids(self):
+        reply = billing.ListSubscriptionCleanupWhitelistReply(
+            data=[
+                billing.SubscriptionCleanupWhitelist(
+                    id=f"wl-{tenant_id}",
+                    tenant_id=tenant_id,
+                    contact="ops@example.com",
+                    expired_at="2026-01-01T00:00:00Z",
+                    created_at="2025-01-01T00:00:00Z",
+                    updated_at="2025-01-01T00:00:00Z",
+                )
+                for tenant_id in ("tenant-a", "tenant-b")
+            ]
         )
 
-    def test_get_vector_space_success(self, mock_send_request):
-        """Test successful retrieval of vector-space usage and limit."""
-        # Arrange
-        tenant_id = "tenant-123"
-        expected_response = {"size": 5120.75, "limit": 20480}
-        mock_send_request.return_value = expected_response
-
-        # Act
-        result = BillingService.get_vector_space(tenant_id)
-
-        # Assert
-        assert result == expected_response
-        mock_send_request.assert_called_once_with(
-            "GET",
-            "/subscription/vector-space",
-            params={"tenant_id": tenant_id},
-        )
-
-    def test_get_vector_space_preserves_unknown_usage(self, mock_send_request):
-        tenant_id = "tenant-123"
-        expected_response = {"size": 0.0, "limit": 50, "usage_unknown": True}
-        mock_send_request.return_value = expected_response
-
-        result = BillingService.get_vector_space(tenant_id)
-
-        assert result == expected_response
-
-    def test_get_info_preserves_unknown_vector_space_usage(self, mock_send_request):
-        tenant_id = "tenant-123"
-        expected_response = {
-            "subscription": {"plan": "sandbox", "interval": "", "education": False},
-            "members": {"size": 1, "limit": 1},
-            "apps": {"size": 1, "limit": 10},
-            "vector_space": {"size": 0.0, "limit": 50, "usage_unknown": True},
-            "documents_upload_quota": {"size": 1, "limit": 50},
-            "annotation_quota_limit": {"size": 0, "limit": 10},
-            "can_replace_logo": False,
-            "model_load_balancing_enabled": False,
-            "knowledge_pipeline_publish_enabled": False,
-        }
-        mock_send_request.return_value = expected_response
-
-        result = BillingService.get_info(tenant_id)
-
-        assert result["vector_space"]["usage_unknown"] is True
-
-    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"},
-        )
+        with patch.object(billing, "subscription_list_subscription_cleanup_whitelist", return_value=reply):
+            assert BillingService.get_expired_subscription_cleanup_whitelist() == ["tenant-a", "tenant-b"]
 
     def test_invalidate_vector_space_cache_bypasses_cache(self):
         tenant_id = "tenant-123"
@@ -597,96 +477,6 @@ class TestBillingServiceSubscriptionInfo:
             params={"tenant_id": tenant_id, "feature_key": "credit_pool", "bucket": "trial"},
         )
 
-    def test_get_knowledge_rate_limit_with_defaults(self, mock_send_request):
-        """Test knowledge rate limit retrieval with default values."""
-        # Arrange
-        tenant_id = "tenant-456"
-        mock_send_request.return_value = {}
-
-        # Act
-        result = BillingService.get_knowledge_rate_limit(tenant_id)
-
-        # Assert
-        assert result["limit"] == 10  # Default limit
-        assert result["subscription_plan"] == CloudPlan.SANDBOX  # Default plan
-        mock_send_request.assert_called_once_with(
-            "GET", "/subscription/knowledge-rate-limit", params={"tenant_id": tenant_id}
-        )
-
-    def test_get_knowledge_rate_limit_with_custom_values(self, mock_send_request):
-        """Test knowledge rate limit retrieval with custom values."""
-        # Arrange
-        tenant_id = "tenant-789"
-        mock_send_request.return_value = {"limit": 100, "subscription_plan": CloudPlan.PROFESSIONAL}
-
-        # Act
-        result = BillingService.get_knowledge_rate_limit(tenant_id)
-
-        # Assert
-        assert result["limit"] == 100
-        assert result["subscription_plan"] == CloudPlan.PROFESSIONAL
-
-    def test_get_subscription_payment_link(self):
-        """Test subscription payment link generation."""
-        # Arrange
-        plan = "professional"
-        interval = "month"
-        email = "user@example.com"
-        tenant_id = "tenant-123"
-        expected_response = {"url": "https://payment.example.com/checkout"}
-        with patch.object(
-            BillingService, "_send_billing_portal_request", return_value=expected_response
-        ) as send_request:
-            result = BillingService.get_subscription(plan, interval, email, tenant_id)
-
-        # Assert
-        assert result == expected_response
-        send_request.assert_called_once_with(
-            "/subscription/payment-link",
-            params={"plan": plan, "interval": interval, "prefilled_email": email, "tenant_id": tenant_id},
-        )
-
-    def test_get_model_provider_payment_link(self, mock_send_request):
-        """Test model provider payment link generation."""
-        # Arrange
-        provider_name = "openai"
-        tenant_id = "tenant-123"
-        account_id = "account-456"
-        email = "user@example.com"
-        expected_response = {"payment_link": "https://payment.example.com/provider"}
-        mock_send_request.return_value = expected_response
-
-        # Act
-        result = BillingService.get_model_provider_payment_link(provider_name, tenant_id, account_id, email)
-
-        # Assert
-        assert result == expected_response
-        mock_send_request.assert_called_once_with(
-            "GET",
-            "/model-provider/payment-link",
-            params={
-                "provider_name": provider_name,
-                "tenant_id": tenant_id,
-                "account_id": account_id,
-                "prefilled_email": email,
-            },
-        )
-
-    def test_get_invoices(self):
-        """Test invoice retrieval."""
-        # Arrange
-        email = "user@example.com"
-        tenant_id = "tenant-123"
-        expected_response = {"url": "https://payment.example.com/invoices"}
-        with patch.object(
-            BillingService, "_send_billing_portal_request", return_value=expected_response
-        ) as send_request:
-            result = BillingService.get_invoices(email, tenant_id)
-
-        # Assert
-        assert result == expected_response
-        send_request.assert_called_once_with("/invoices", params={"prefilled_email": email, "tenant_id": tenant_id})
-
 
 class TestBillingServiceUsageCalculation:
     """Unit tests for usage calculation and credit management.
@@ -1224,510 +1014,6 @@ class TestBillingServiceCacheManagement:
         mock_redis_client.delete.assert_called_once_with(expected_key)
 
 
-class TestBillingServiceEdgeCases:
-    """Unit tests for edge cases and error scenarios.
-
-    Tests cover:
-    - Empty responses from billing API
-    - Malformed JSON responses
-    - Boundary conditions for rate limits
-    - Multiple subscription tiers
-    """
-
-    @pytest.fixture
-    def mock_send_request(self):
-        """Mock _send_request method."""
-        with patch.object(BillingService, "_send_request") as mock:
-            yield mock
-
-    def test_get_info_empty_response(self, mock_send_request):
-        """Empty response from billing API should raise ValidationError due to missing required fields."""
-        from pydantic import ValidationError
-
-        tenant_id = "tenant-empty"
-        mock_send_request.return_value = {}
-
-        with pytest.raises(ValidationError):
-            BillingService.get_info(tenant_id)
-
-    def test_get_knowledge_rate_limit_all_subscription_tiers(self, mock_send_request):
-        """Test knowledge rate limit for all subscription tiers."""
-        # Test SANDBOX tier
-        mock_send_request.return_value = {"limit": 10, "subscription_plan": CloudPlan.SANDBOX}
-        result = BillingService.get_knowledge_rate_limit("tenant-sandbox")
-        assert result["subscription_plan"] == CloudPlan.SANDBOX
-        assert result["limit"] == 10
-
-        # Test PROFESSIONAL tier
-        mock_send_request.return_value = {"limit": 100, "subscription_plan": CloudPlan.PROFESSIONAL}
-        result = BillingService.get_knowledge_rate_limit("tenant-pro")
-        assert result["subscription_plan"] == CloudPlan.PROFESSIONAL
-        assert result["limit"] == 100
-
-        # Test TEAM tier
-        mock_send_request.return_value = {"limit": 500, "subscription_plan": CloudPlan.TEAM}
-        result = BillingService.get_knowledge_rate_limit("tenant-team")
-        assert result["subscription_plan"] == CloudPlan.TEAM
-        assert result["limit"] == 500
-
-    def test_get_subscription_with_empty_optional_params(self, mock_send_request):
-        """Test subscription payment link with empty optional parameters."""
-        # Arrange
-        plan = "professional"
-        interval = "year"
-        expected_response = {"url": "https://payment.example.com/checkout"}
-        mock_send_request.return_value = expected_response
-
-        # Act - empty email and tenant_id
-        result = BillingService.get_subscription(plan, interval, "", "")
-
-        # Assert
-        assert result == expected_response
-        mock_send_request.assert_called_once_with(
-            "GET",
-            "/subscription/payment-link",
-            params={"plan": plan, "interval": interval, "prefilled_email": "", "tenant_id": ""},
-        )
-
-    def test_get_invoices_with_empty_params(self, mock_send_request):
-        """Test invoice retrieval with empty parameters."""
-        # Arrange
-        expected_response = {"url": "https://payment.example.com/invoices"}
-        mock_send_request.return_value = expected_response
-
-        # Act
-        result = BillingService.get_invoices("", "")
-
-        # Assert
-        assert result == expected_response
-
-
-class TestBillingServiceSubscriptionOperations:
-    """Unit tests for subscription operations in BillingService.
-
-    Tests cover:
-    - Bulk plan retrieval with chunking
-    - Expired subscription cleanup whitelist retrieval
-    """
-
-    @pytest.fixture
-    def mock_send_request(self):
-        """Mock _send_request method."""
-        with patch.object(BillingService, "_send_request") as mock:
-            yield mock
-
-    def test_get_plan_bulk_with_empty_list(self, mock_send_request):
-        """Test bulk plan retrieval with empty tenant list."""
-        # Arrange
-        tenant_ids = []
-
-        # Act
-        result = BillingService.get_plan_bulk(tenant_ids)
-
-        # Assert
-        assert result == {}
-        mock_send_request.assert_not_called()
-
-    def test_get_plan_bulk_with_chunking(self, mock_send_request):
-        """Test bulk plan retrieval with more than 200 tenants (chunking logic)."""
-        # Arrange - 250 tenants to test chunking (chunk_size = 200)
-        tenant_ids = [f"tenant-{i}" for i in range(250)]
-
-        # First chunk: tenants 0-199
-        first_chunk_response = {
-            "data": {f"tenant-{i}": {"plan": "sandbox", "expiration_date": 1735689600} for i in range(200)}
-        }
-
-        # Second chunk: tenants 200-249
-        second_chunk_response = {
-            "data": {f"tenant-{i}": {"plan": "professional", "expiration_date": 1767225600} for i in range(200, 250)}
-        }
-
-        mock_send_request.side_effect = [first_chunk_response, second_chunk_response]
-
-        # Act
-        result = BillingService.get_plan_bulk(tenant_ids)
-
-        # Assert
-        assert len(result) == 250
-        assert result["tenant-0"]["plan"] == "sandbox"
-        assert result["tenant-199"]["plan"] == "sandbox"
-        assert result["tenant-200"]["plan"] == "professional"
-        assert result["tenant-249"]["plan"] == "professional"
-        assert mock_send_request.call_count == 2
-
-        # Verify first chunk call
-        first_call = mock_send_request.call_args_list[0]
-        assert first_call[0][0] == "POST"
-        assert first_call[0][1] == "/subscription/plan/batch"
-        assert len(first_call[1]["json"]["tenant_ids"]) == 200
-
-        # Verify second chunk call
-        second_call = mock_send_request.call_args_list[1]
-        assert len(second_call[1]["json"]["tenant_ids"]) == 50
-
-    def test_get_plan_bulk_with_partial_batch_failure(self, mock_send_request):
-        """Test bulk plan retrieval when one batch fails but others succeed."""
-        # Arrange - 250 tenants, second batch will fail
-        tenant_ids = [f"tenant-{i}" for i in range(250)]
-
-        # First chunk succeeds
-        first_chunk_response = {
-            "data": {f"tenant-{i}": {"plan": "sandbox", "expiration_date": 1735689600} for i in range(200)}
-        }
-
-        # Second chunk fails - need to create a mock that raises when called
-        def side_effect_func[**P](*args: P.args, **kwargs: P.kwargs):
-            if mock_send_request.call_count == 1:
-                return first_chunk_response
-            else:
-                raise ValueError("API error")
-
-        mock_send_request.side_effect = side_effect_func
-
-        # Act
-        result = BillingService.get_plan_bulk(tenant_ids)
-
-        # Assert - should only have data from first batch
-        assert len(result) == 200
-        assert result["tenant-0"]["plan"] == "sandbox"
-        assert result["tenant-199"]["plan"] == "sandbox"
-        assert "tenant-200" not in result
-        assert mock_send_request.call_count == 2
-
-    def test_get_plan_bulk_with_all_batches_failing(self, mock_send_request):
-        """Test bulk plan retrieval when all batches fail."""
-        # Arrange
-        tenant_ids = [f"tenant-{i}" for i in range(250)]
-
-        # All chunks fail
-        def side_effect_func[**P](*args: P.args, **kwargs: P.kwargs):
-            raise ValueError("API error")
-
-        mock_send_request.side_effect = side_effect_func
-
-        # Act
-        result = BillingService.get_plan_bulk(tenant_ids)
-
-        # Assert - should return empty dict
-        assert result == {}
-        assert mock_send_request.call_count == 2
-
-    def test_get_plan_bulk_with_exactly_200_tenants(self, mock_send_request):
-        """Test bulk plan retrieval with exactly 200 tenants (boundary condition)."""
-        # Arrange
-        tenant_ids = [f"tenant-{i}" for i in range(200)]
-        mock_send_request.return_value = {
-            "data": {f"tenant-{i}": {"plan": "sandbox", "expiration_date": 1735689600} for i in range(200)}
-        }
-
-        # Act
-        result = BillingService.get_plan_bulk(tenant_ids)
-
-        # Assert
-        assert len(result) == 200
-        assert mock_send_request.call_count == 1
-
-    def test_get_plan_bulk_with_empty_data_response(self, mock_send_request):
-        """Test bulk plan retrieval with empty data in response."""
-        # Arrange
-        tenant_ids = ["tenant-1", "tenant-2"]
-        mock_send_request.return_value = {"data": {}}
-
-        # Act
-        result = BillingService.get_plan_bulk(tenant_ids)
-
-        # Assert
-        assert result == {}
-
-    def test_get_plan_bulk_converts_string_expiration_date_to_int(self, mock_send_request):
-        """Test bulk plan retrieval converts string expiration_date to int."""
-        # Arrange
-        tenant_ids = ["tenant-1"]
-        mock_send_request.return_value = {
-            "data": {
-                "tenant-1": {"plan": "sandbox", "expiration_date": "1735689600"},
-            }
-        }
-
-        # Act
-        result = BillingService.get_plan_bulk(tenant_ids)
-
-        # Assert
-        assert "tenant-1" in result
-        assert isinstance(result["tenant-1"]["expiration_date"], int)
-        assert result["tenant-1"]["expiration_date"] == 1735689600
-
-    def test_get_plan_bulk_with_invalid_tenant_plan_skipped(self, mock_send_request, caplog: pytest.LogCaptureFixture):
-        """Test bulk plan retrieval when one tenant has invalid plan data (should skip that tenant)."""
-        # Arrange
-        tenant_ids = ["tenant-valid-1", "tenant-invalid", "tenant-valid-2"]
-
-        # Response with one invalid tenant plan (missing expiration_date) and two valid ones
-        mock_send_request.return_value = {
-            "data": {
-                "tenant-valid-1": {"plan": "sandbox", "expiration_date": 1735689600},
-                "tenant-invalid": {"plan": "professional"},  # Missing expiration_date field
-                "tenant-valid-2": {"plan": "team", "expiration_date": 1767225600},
-            }
-        }
-
-        # Act
-        with caplog.at_level(logging.ERROR, logger="services.billing_service"):
-            result = BillingService.get_plan_bulk(tenant_ids)
-
-        # Assert - should only contain valid tenants
-        assert len(result) == 2
-        assert "tenant-valid-1" in result
-        assert "tenant-valid-2" in result
-        assert "tenant-invalid" not in result
-
-        # Verify valid tenants have correct data
-        assert result["tenant-valid-1"]["plan"] == "sandbox"
-        assert result["tenant-valid-1"]["expiration_date"] == 1735689600
-        assert result["tenant-valid-2"]["plan"] == "team"
-        assert result["tenant-valid-2"]["expiration_date"] == 1767225600
-
-        # Verify exception was logged for the invalid tenant
-        exception_records = [r for r in caplog.records if r.levelname == "ERROR"]
-        assert len(exception_records) == 1
-        formatted = exception_records[0].getMessage()
-        assert "get_plan_bulk: failed to validate subscription plan for tenant" in formatted
-        assert "tenant-invalid" in formatted
-
-    def test_get_expired_subscription_cleanup_whitelist_success(self, mock_send_request):
-        """Test successful retrieval of expired subscription cleanup whitelist."""
-        # Arrange
-        api_response = [
-            {
-                "created_at": "2025-10-16T01:56:17",
-                "tenant_id": "36bd55ec-2ea9-4d75-a9ea-1f26aeb4ffe6",
-                "contact": "example@dify.ai",
-                "id": "36bd55ec-2ea9-4d75-a9ea-1f26aeb4ffe5",
-                "expired_at": "2026-01-01T01:56:17",
-                "updated_at": "2025-10-16T01:56:17",
-            },
-            {
-                "created_at": "2025-10-16T02:00:00",
-                "tenant_id": "tenant-2",
-                "contact": "test@example.com",
-                "id": "whitelist-id-2",
-                "expired_at": "2026-02-01T00:00:00",
-                "updated_at": "2025-10-16T02:00:00",
-            },
-            {
-                "created_at": "2025-10-16T03:00:00",
-                "tenant_id": "tenant-3",
-                "contact": "another@example.com",
-                "id": "whitelist-id-3",
-                "expired_at": "2026-03-01T00:00:00",
-                "updated_at": "2025-10-16T03:00:00",
-            },
-        ]
-        mock_send_request.return_value = {"data": api_response}
-
-        # Act
-        result = BillingService.get_expired_subscription_cleanup_whitelist()
-
-        # Assert - should return only tenant_ids
-        assert result == ["36bd55ec-2ea9-4d75-a9ea-1f26aeb4ffe6", "tenant-2", "tenant-3"]
-        assert len(result) == 3
-        assert result[0] == "36bd55ec-2ea9-4d75-a9ea-1f26aeb4ffe6"
-        assert result[1] == "tenant-2"
-        assert result[2] == "tenant-3"
-        mock_send_request.assert_called_once_with("GET", "/subscription/cleanup/whitelist")
-
-    def test_get_expired_subscription_cleanup_whitelist_empty_list(self, mock_send_request):
-        """Test retrieval of empty cleanup whitelist."""
-        # Arrange
-        mock_send_request.return_value = {"data": []}
-
-        # Act
-        result = BillingService.get_expired_subscription_cleanup_whitelist()
-
-        # Assert
-        assert result == []
-        assert len(result) == 0
-
-
-class TestBillingServiceIntegrationScenarios:
-    """Integration-style tests simulating real-world usage scenarios.
-
-    These tests combine multiple service methods to test common workflows:
-    - Complete subscription upgrade flow
-    - Rate limit boundary testing
-    """
-
-    @pytest.fixture
-    def mock_send_request(self):
-        """Mock _send_request method."""
-        with patch.object(BillingService, "_send_request") as mock:
-            yield mock
-
-    def test_subscription_upgrade_workflow(self, mock_send_request):
-        """Test complete subscription upgrade workflow."""
-        # Arrange
-        tenant_id = "tenant-upgrade"
-
-        # Step 1: Get current billing info
-        mock_send_request.return_value = {
-            "subscription": {"plan": "sandbox", "interval": "", "education": False},
-            "members": {"size": 0, "limit": 1},
-            "apps": {"size": 0, "limit": 5},
-            "vector_space": {"size": 0.0, "limit": 50},
-            "documents_upload_quota": {"size": 0, "limit": 50},
-            "annotation_quota_limit": {"size": 0, "limit": 10},
-            "can_replace_logo": False,
-            "model_load_balancing_enabled": False,
-            "knowledge_pipeline_publish_enabled": False,
-        }
-        current_info = BillingService.get_info(tenant_id)
-        assert current_info["subscription"]["plan"] == "sandbox"
-
-        # Step 2: Get payment link for upgrade
-        mock_send_request.return_value = {"url": "https://payment.example.com/upgrade"}
-        payment_link = BillingService.get_subscription("professional", "month", "user@example.com", tenant_id)
-        assert "url" in payment_link
-
-        # Step 3: Verify new rate limits after upgrade
-        mock_send_request.return_value = {"limit": 100, "subscription_plan": CloudPlan.PROFESSIONAL}
-        rate_limit = BillingService.get_knowledge_rate_limit(tenant_id)
-        assert rate_limit["subscription_plan"] == CloudPlan.PROFESSIONAL
-        assert rate_limit["limit"] == 100
-
-
-class TestBillingServiceSubscriptionInfoDataType:
-    """Unit tests for data type coercion in BillingService.get_info
-
-    1. Verifies the get_info returns correct Python types for numeric fields
-    2. Ensure the compatibility regardless of what results the upstream billing API returns
-    """
-
-    @pytest.fixture
-    def mock_send_request(self):
-        with patch.object(BillingService, "_send_request") as mock:
-            yield mock
-
-    @pytest.fixture
-    def normal_billing_response(self) -> dict:
-        return {
-            "subscription": {
-                "plan": "team",
-                "interval": "year",
-                "education": False,
-            },
-            "members": {"size": 10, "limit": 50},
-            "apps": {"size": 80, "limit": 200},
-            "vector_space": {"size": 5120.75, "limit": 20480},
-            "documents_upload_quota": {"size": 450, "limit": 1000},
-            "annotation_quota_limit": {"size": 1200, "limit": 5000},
-            "can_replace_logo": True,
-            "model_load_balancing_enabled": True,
-            "knowledge_pipeline_publish_enabled": True,
-            "next_credit_reset_date": 1745971200,
-        }
-
-    @pytest.fixture
-    def string_billing_response(self) -> dict:
-        return {
-            "subscription": {
-                "plan": "team",
-                "interval": "year",
-                "education": False,
-            },
-            "members": {"size": "10", "limit": "50"},
-            "apps": {"size": "80", "limit": "200"},
-            "vector_space": {"size": 5120.75, "limit": "20480"},
-            "documents_upload_quota": {"size": "450", "limit": "1000"},
-            "annotation_quota_limit": {"size": "1200", "limit": "5000"},
-            "can_replace_logo": True,
-            "model_load_balancing_enabled": True,
-            "knowledge_pipeline_publish_enabled": True,
-            "next_credit_reset_date": "1745971200",
-        }
-
-    @staticmethod
-    def _assert_billing_info_types(result: dict):
-        assert isinstance(result["subscription"]["plan"], str)
-        assert isinstance(result["subscription"]["interval"], str)
-        assert isinstance(result["subscription"]["education"], bool)
-
-        assert isinstance(result["members"]["size"], int)
-        assert isinstance(result["members"]["limit"], int)
-
-        assert isinstance(result["apps"]["size"], int)
-        assert isinstance(result["apps"]["limit"], int)
-
-        if "vector_space" in result:
-            assert isinstance(result["vector_space"]["size"], float)
-            assert isinstance(result["vector_space"]["limit"], int)
-            if "usage_unknown" in result["vector_space"]:
-                assert isinstance(result["vector_space"]["usage_unknown"], bool)
-
-        assert isinstance(result["documents_upload_quota"]["size"], int)
-        assert isinstance(result["documents_upload_quota"]["limit"], int)
-
-        assert isinstance(result["annotation_quota_limit"]["size"], int)
-        assert isinstance(result["annotation_quota_limit"]["limit"], int)
-
-        assert isinstance(result["can_replace_logo"], bool)
-        assert isinstance(result["model_load_balancing_enabled"], bool)
-        assert isinstance(result["knowledge_pipeline_publish_enabled"], bool)
-        if "next_credit_reset_date" in result:
-            assert isinstance(result["next_credit_reset_date"], int)
-
-    def test_get_info_with_normal_types(self, mock_send_request, normal_billing_response):
-        """When the billing API returns native numeric types, get_info should preserve them."""
-        mock_send_request.return_value = normal_billing_response
-
-        result = BillingService.get_info("tenant-type-test")
-
-        self._assert_billing_info_types(result)
-        mock_send_request.assert_called_once_with("GET", "/subscription/info", params={"tenant_id": "tenant-type-test"})
-
-    def test_get_info_with_string_types(self, mock_send_request, string_billing_response):
-        """When the billing API returns numeric values as strings, get_info should coerce them."""
-        mock_send_request.return_value = string_billing_response
-
-        result = BillingService.get_info("tenant-type-test")
-
-        self._assert_billing_info_types(result)
-        mock_send_request.assert_called_once_with("GET", "/subscription/info", params={"tenant_id": "tenant-type-test"})
-
-    def test_get_info_without_optional_fields(self, mock_send_request, string_billing_response):
-        """NotRequired fields can be absent without raising."""
-        del string_billing_response["next_credit_reset_date"]
-        del string_billing_response["vector_space"]
-        mock_send_request.return_value = string_billing_response
-
-        result = BillingService.get_info("tenant-type-test")
-
-        assert "next_credit_reset_date" not in result
-        assert "vector_space" not in result
-        self._assert_billing_info_types(result)
-
-    def test_get_info_with_extra_fields(self, mock_send_request, string_billing_response):
-        """Undefined fields are silently stripped by validate_python."""
-        string_billing_response["new_feature"] = "something"
-        mock_send_request.return_value = string_billing_response
-
-        result = BillingService.get_info("tenant-type-test")
-
-        # extra fields are dropped by TypeAdapter on TypedDict
-        assert "new_feature" not in result
-        self._assert_billing_info_types(result)
-
-    def test_get_info_missing_required_field_raises(self, mock_send_request, string_billing_response):
-        """Missing a required field should raise ValidationError."""
-        from pydantic import ValidationError
-
-        del string_billing_response["members"]
-        mock_send_request.return_value = string_billing_response
-
-        with pytest.raises(ValidationError):
-            BillingService.get_info("tenant-type-test")
-
-
 def test_pooled_billing_client_carries_bounded_timeout() -> None:
     """Regression for #39874: the pooled billing client must carry a
     read/connect timeout so a stalled Stripe / cloud-billing proxy

api/tests/unit_tests/services/test_clear_free_plan_tenant_expired_logs.py

@@ -11,7 +11,7 @@ from sqlalchemy import event
 from sqlalchemy.engine import Engine
 from sqlalchemy.orm import Session
 
-from enums import CloudPlan, DeploymentEdition
+from enums import DeploymentEdition
 from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus
 from graphon.file import FileTransferMethod, FileType
 from models.account import Tenant
@@ -45,7 +45,9 @@ from models.workflow import (
     WorkflowType,
 )
 from services import clear_free_plan_tenant_expired_logs as service_module
+from services.billing_client import billing
 from services.clear_free_plan_tenant_expired_logs import ClearFreePlanTenantExpiredLogs
+from tests.helpers import billing_builders as builders
 
 
 @pytest.fixture(autouse=True)
@@ -545,12 +547,12 @@ def test_process_with_tenant_ids_filters_by_plan_and_logs_errors(
     _configure_process_boundaries(monkeypatch, sqlite_engine)
     monkeypatch.setattr(service_module.click, "echo", MagicMock())
 
-    def fake_get_info(tenant_id: str) -> dict[str, dict[str, str]]:
+    def fake_get_info(tenant_id: str) -> billing.GetSubscriptionInfoReply:
         if tenant_id == "tenant-sandbox":
-            return {"subscription": {"plan": CloudPlan.SANDBOX}}
+            return builders.subscription_info("sandbox")
         if tenant_id == "tenant-fail":
             raise RuntimeError("billing failure")
-        return {"subscription": {"plan": "team"}}
+        return builders.subscription_info("team")
 
     monkeypatch.setattr(service_module.BillingService, "get_info", staticmethod(fake_get_info))
     process_tenant = MagicMock(side_effect=RuntimeError("cleanup failure"))

api/tests/unit_tests/services/test_feature_service.py

@@ -1,17 +1,18 @@
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import cast
 from unittest.mock import Mock, patch
 
 import pytest
+from pytest_mock import MockerFixture
 
 from enums import CloudPlan, DeploymentEdition, HostedTrialProvider
 from services import feature_service as feature_service_module
 from services.account_education_service import EDUCATION_EDITIONS
-from services.billing_service import BillingInfo
+from services.billing_client import billing
 from services.entities.feature_entities import FeatureModel, LimitationModel
 from services.feature_service import FeatureService
 from services.system_feature_service import SystemFeatureService
+from tests.helpers import billing_builders as builders
 
 
 def test_education_feature_is_disabled_by_default() -> None:
@@ -121,7 +122,7 @@ def test_get_vector_space_converts_billing_float_size(
     monkeypatch.setattr(
         feature_service_module.BillingService,
         "get_vector_space",
-        lambda _tenant_id: {"size": 5120.75, "limit": 20480},
+        lambda _tenant_id: builders.vector_space(size=5120.75, limit=20480),
     )
 
     result = FeatureService.get_vector_space("tenant-1")
@@ -138,7 +139,7 @@ def test_get_vector_space_preserves_unknown_usage(
     monkeypatch.setattr(
         feature_service_module.BillingService,
         "get_vector_space",
-        lambda _tenant_id: {"size": 0.0, "limit": 50, "usage_unknown": True},
+        lambda _tenant_id: builders.vector_space(size=0.0, limit=50, usage_unknown=True),
     )
 
     result = FeatureService.get_vector_space("tenant-1")
@@ -160,7 +161,7 @@ def test_get_vector_space_preserves_unknown_usage(
     ],
 )
 def test_get_knowledge_file_size_limit(
-    monkeypatch: pytest.MonkeyPatch,
+    mocker: MockerFixture,
     config_overrides: Callable[..., None],
     deployment_edition: DeploymentEdition,
     tenant_id: str | None,
@@ -172,12 +173,9 @@ def test_get_knowledge_file_size_limit(
         UPLOAD_FILE_SIZE_LIMIT=15,
         KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN=50,
     )
-    get_info = Mock(
-        return_value={
-            "subscription": {"plan": plan},
-        }
+    get_info = mocker.patch.object(
+        feature_service_module.BillingService, "get_info", return_value=builders.subscription_info(plan)
     )
-    monkeypatch.setattr(feature_service_module.BillingService, "get_info", get_info)
 
     assert FeatureService.get_knowledge_file_size_limit(tenant_id) == expected
 
@@ -198,9 +196,7 @@ def test_paid_knowledge_file_size_limit_never_reduces_default(
     monkeypatch.setattr(
         feature_service_module.BillingService,
         "get_info",
-        lambda *_args, **_kwargs: {
-            "subscription": {"plan": CloudPlan.PROFESSIONAL},
-        },
+        lambda *_args, **_kwargs: builders.subscription_info("professional"),
     )
 
     assert FeatureService.get_knowledge_file_size_limit("tenant-1") == 100
@@ -321,16 +317,17 @@ def test_get_trial_models_filters_providers_by_workspace_plan(
 
 def test_get_features_exclude_vector_space_sets_vector_space_to_none(config_overrides: Callable[..., None]) -> None:
     tenant_id = "tenant-id"
-    billing_info = {
-        "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "monthly", "education": False},
-        "members": {"size": 1, "limit": 10},
-        "apps": {"size": 2, "limit": 20},
-        "documents_upload_quota": {"size": 3, "limit": 100},
-        "annotation_quota_limit": {"size": 4, "limit": 50},
-        "can_replace_logo": True,
-        "model_load_balancing_enabled": True,
-        "knowledge_pipeline_publish_enabled": True,
-    }
+    billing_info = builders.subscription_info(
+        subscription=builders.subscription_plan("professional", interval="month", education=True),
+        members=billing.SubscriptionQuota(size=1, limit=10),
+        apps=billing.SubscriptionQuota(size=2, limit=20),
+        documents_upload_quota=billing.SubscriptionQuota(size=3, limit=100),
+        annotation_quota_limit=billing.SubscriptionQuota(size=4, limit=50),
+        can_replace_logo=True,
+        model_load_balancing_enabled=True,
+        knowledge_pipeline_publish_enabled=True,
+        next_credit_reset_date=1775001600,
+    )
 
     config_overrides(
         DEPLOYMENT_EDITION=DeploymentEdition.CLOUD,
@@ -345,6 +342,17 @@ def test_get_features_exclude_vector_space_sets_vector_space_to_none(config_over
         features = FeatureService.get_features(tenant_id, exclude_vector_space=True)
 
     assert features.vector_space is None
+    assert features.billing.subscription.plan == CloudPlan.PROFESSIONAL
+    assert features.billing.subscription.interval == "month"
+    assert features.education.activated is True
+    assert (features.members.size, features.members.limit) == (1, 10)
+    assert (features.apps.size, features.apps.limit) == (2, 20)
+    assert (features.documents_upload_quota.size, features.documents_upload_quota.limit) == (3, 100)
+    assert (features.annotation_quota_limit.size, features.annotation_quota_limit.limit) == (4, 50)
+    assert features.can_replace_logo is True
+    assert features.model_load_balancing_enabled is True
+    assert features.knowledge_pipeline.publish_enabled is True
+    assert features.next_credit_reset_date == 1775001600
     get_info.assert_called_once_with(tenant_id, exclude_vector_space=True)
 
 
@@ -353,7 +361,7 @@ def test_full_features_keep_treating_unknown_vector_usage_as_zero() -> None:
 
     FeatureService._fulfill_vector_space_from_billing_info(
         vector_space,
-        cast(BillingInfo, {"vector_space": {"size": 0.0, "limit": 50, "usage_unknown": True}}),
+        builders.subscription_info(vector_space=builders.vector_space(size=0.0, limit=50, usage_unknown=True)),
     )
 
     assert vector_space.size == 0

api/tests/unit_tests/services/test_network_access_group_gateway.py

@@ -6,7 +6,7 @@ import httpx
 import pytest
 
 from enums import CloudPlan
-from services.billing_service import _BillingHTTPStatusError
+from services.errors.billing import BillingUpstreamInvalidResponseError, BillingUpstreamUnavailableError
 from services.network_access_group_gateway import (
     BillingNetworkAccessGroupEntitlementGateway,
     NetworkAccessGroupGateway,
@@ -16,6 +16,7 @@ from services.network_access_group_service import (
     NetworkAccessGroupInvalidResponseError,
     NetworkAccessGroupUpstreamError,
 )
+from tests.helpers import billing_builders as builders
 
 TENANT_ID = "11111111-1111-4111-8111-111111111111"
 ACCOUNT_ID = "22222222-2222-4222-8222-222222222222"
@@ -449,7 +450,7 @@ def test_billing_entitlement_gateway_accepts_paid_plans(plan: CloudPlan) -> None
     gateway = BillingNetworkAccessGroupEntitlementGateway()
     with patch(
         "services.network_access_group_gateway.BillingService.get_info",
-        return_value={"subscription": {"plan": plan}},
+        return_value=builders.subscription_info(plan),
     ) as get_info:
         assert gateway.is_paid_plan(TENANT_ID) is True
 
@@ -460,30 +461,19 @@ def test_billing_entitlement_gateway_rejects_unpaid_plan() -> None:
     gateway = BillingNetworkAccessGroupEntitlementGateway()
     with patch(
         "services.network_access_group_gateway.BillingService.get_info",
-        return_value={"subscription": {"plan": CloudPlan.SANDBOX}},
+        return_value=builders.subscription_info("sandbox"),
     ):
         assert gateway.is_paid_plan(TENANT_ID) is False
 
 
-def test_billing_entitlement_gateway_maps_unavailable_billing_service() -> None:
+@pytest.mark.parametrize(
+    "error",
+    [BillingUpstreamUnavailableError("unavailable"), BillingUpstreamInvalidResponseError("invalid body")],
+)
+def test_billing_entitlement_gateway_maps_billing_upstream_errors(error: Exception) -> None:
     gateway = BillingNetworkAccessGroupEntitlementGateway()
     with (
-        patch(
-            "services.network_access_group_gateway.BillingService.get_info",
-            side_effect=httpx.ConnectError("unavailable"),
-        ),
-        pytest.raises(NetworkAccessGroupEntitlementUnavailableError),
-    ):
-        gateway.is_paid_plan(TENANT_ID)
-
-
-def test_billing_entitlement_gateway_maps_billing_http_status_error() -> None:
-    gateway = BillingNetworkAccessGroupEntitlementGateway()
-    with (
-        patch(
-            "services.network_access_group_gateway.BillingService.get_info",
-            side_effect=_BillingHTTPStatusError("unavailable", 503),
-        ),
+        patch("services.network_access_group_gateway.BillingService.get_info", side_effect=error),
         pytest.raises(NetworkAccessGroupEntitlementUnavailableError),
     ):
         gateway.is_paid_plan(TENANT_ID)

api/tests/unit_tests/services/test_vector_space_admission_service.py

@@ -14,6 +14,7 @@ from core.rag.index_processor.constant.index_type import IndexStructureType, Ind
 from core.rag.models.document import AttachmentDocument, ChildDocument, Document
 from enums import CloudPlan, DeploymentEdition
 from models.dataset import Dataset
+from services.billing_client import billing
 from services.vector_space_admission_service import (
     VECTOR_SPACE_ADMISSION_ERROR_CODE,
     VectorSpaceAdmissionError,
@@ -27,6 +28,7 @@ from services.vector_space_admission_service import (
     get_vector_space_admission_error_fields,
     parse_vector_space_estimate_limits,
 )
+from tests.helpers import billing_builders as builders
 from tests.helpers.unit.model_factories import make_dataset
 
 _MEBIBYTE = 1024 * 1024
@@ -119,7 +121,7 @@ def _check_estimate(
         ),
         patch(
             "services.vector_space_admission_service.BillingService.get_vector_space",
-            return_value={"size": usage_mb, "limit": plan_limit_mb},
+            return_value=builders.vector_space(size=usage_mb, limit=plan_limit_mb),
         ),
         patch("services.vector_space_admission_service.redis_client", redis),
     ):
@@ -375,7 +377,7 @@ def test_usage_lookup_is_refreshed_for_each_document(sqlite_session: Session) ->
         ),
         patch(
             "services.vector_space_admission_service.BillingService.get_vector_space",
-            side_effect=[{"size": 40.0, "limit": 50}, {"size": 50.0, "limit": 50}],
+            side_effect=[builders.vector_space(size=40.0, limit=50), builders.vector_space(size=50.0, limit=50)],
         ) as get_vector_space,
         patch("services.vector_space_admission_service.redis_client", redis),
     ):
@@ -628,7 +630,7 @@ def test_billing_plan_lookup_excludes_vector_space_and_is_cached() -> None:
     service = VectorSpaceAdmissionService()
     with patch(
         "services.vector_space_admission_service.BillingService.get_info",
-        return_value={"subscription": {"plan": "professional"}},
+        return_value=builders.subscription_info("professional"),
     ) as get_info:
         assert service._get_plan("tenant-1") == CloudPlan.PROFESSIONAL
         assert service._get_plan("tenant-1") == CloudPlan.PROFESSIONAL
@@ -641,7 +643,9 @@ def test_unknown_billing_plan_is_not_logged(caplog: pytest.LogCaptureFixture) ->
     upstream_plan = "unexpected-private-plan"
     with patch(
         "services.vector_space_admission_service.BillingService.get_info",
-        return_value={"subscription": {"plan": upstream_plan}},
+        return_value=builders.subscription_info(
+            subscription=builders.subscription_plan().model_copy(update={"plan": upstream_plan})
+        ),
     ):
         assert service._get_plan("tenant-1") is None
 
@@ -675,7 +679,13 @@ def test_unavailable_billing_quota_fails_before_embedding_or_reservation(
             "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(
+            "services.vector_space_admission_service.BillingService.get_vector_space",
+            # Observe-mode validation keeps invalid values raw and reads missing fields as None.
+            return_value=billing.VectorSpaceQuota.model_construct(
+                size=quota.get("size"), limit=quota.get("limit"), usage_unknown=quota.get("usage_unknown")
+            ),
+        ),
         patch.object(service, "_get_embedding_dimension") as dimension,
         patch.object(service, "_reserve_projected_usage") as reserve,
     ):

api/tests/unit_tests/services/test_workflow_service.py

@@ -81,6 +81,7 @@ from services.workflow_service import (
     _rebuild_single_file,
     _setup_variable_pool,
 )
+from tests.helpers import billing_builders as builders
 from tests.helpers.unit.config_override import apply_config_overrides
 from tests.helpers.unit.model_factories import make_account, make_app
 from tests.workflow_test_utils import build_test_graph_init_params
@@ -1503,7 +1504,7 @@ class TestWorkflowServiceCases:
             with (
                 patch("services.workflow_service.BillingService") as mock_billing_service,
             ):
-                mock_billing_service.get_info.return_value = {"subscription": {"plan": "sandbox"}}
+                mock_billing_service.get_info.return_value = builders.subscription_info("sandbox")
 
                 with pytest.raises(TriggerNodeLimitExceededError):
                     workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account)

api/tests/unit_tests/services/workflow/test_queue_dispatcher.py

@@ -11,6 +11,7 @@ from services.workflow.queue_dispatcher import (
     SandboxQueueDispatcher,
     TeamQueueDispatcher,
 )
+from tests.helpers import billing_builders as builders
 
 
 class TestQueuePriority:
@@ -44,7 +45,7 @@ class TestQueueDispatcherManager:
 
     @patch("services.workflow.queue_dispatcher.BillingService")
     def test_cloud_edition_professional_plan(self, mock_billing: MagicMock) -> None:
-        mock_billing.get_info.return_value = {"subscription": {"plan": "professional"}}
+        mock_billing.get_info.return_value = builders.subscription_info("professional")
 
         dispatcher = QueueDispatcherManager.get_dispatcher("tenant-1")
 
@@ -52,7 +53,7 @@ class TestQueueDispatcherManager:
 
     @patch("services.workflow.queue_dispatcher.BillingService")
     def test_cloud_edition_team_plan(self, mock_billing: MagicMock) -> None:
-        mock_billing.get_info.return_value = {"subscription": {"plan": "team"}}
+        mock_billing.get_info.return_value = builders.subscription_info("team")
 
         dispatcher = QueueDispatcherManager.get_dispatcher("tenant-1")
 
@@ -60,7 +61,7 @@ class TestQueueDispatcherManager:
 
     @patch("services.workflow.queue_dispatcher.BillingService")
     def test_cloud_edition_sandbox_plan(self, mock_billing: MagicMock) -> None:
-        mock_billing.get_info.return_value = {"subscription": {"plan": "sandbox"}}
+        mock_billing.get_info.return_value = builders.subscription_info("sandbox")
 
         dispatcher = QueueDispatcherManager.get_dispatcher("tenant-1")
 
@@ -68,7 +69,9 @@ class TestQueueDispatcherManager:
 
     @patch("services.workflow.queue_dispatcher.BillingService")
     def test_cloud_edition_unknown_plan_defaults_to_sandbox(self, mock_billing: MagicMock) -> None:
-        mock_billing.get_info.return_value = {"subscription": {"plan": "enterprise"}}
+        mock_billing.get_info.return_value = builders.subscription_info(
+            subscription=builders.subscription_plan().model_copy(update={"plan": "enterprise"})
+        )
 
         dispatcher = QueueDispatcherManager.get_dispatcher("tenant-1")
 
@@ -91,7 +94,7 @@ class TestQueueDispatcherManager:
 
     @patch("services.workflow.queue_dispatcher.BillingService")
     def test_missing_subscription_key_defaults_to_sandbox(self, mock_billing: MagicMock) -> None:
-        mock_billing.get_info.return_value = dict[str, object]()
+        mock_billing.get_info.return_value = builders.subscription_info().model_copy(update={"subscription": None})
 
         dispatcher = QueueDispatcherManager.get_dispatcher("tenant-1")
 

api/tests/unit_tests/services/workspace/test_gateways.py

@@ -23,6 +23,7 @@ from services.workspace.gateways import (
     DeploymentWorkspacePlanGateway,
     WorkspaceFileGateway,
 )
+from tests.helpers import billing_builders as builders
 from tests.helpers.unit.account_domain import AccountDomain
 from tests.helpers.unit.rbac_fakes import RBACTransport
 
@@ -262,7 +263,7 @@ class TestWorkspaceCreditPoolCases:
     def test_effective_credit_pool(
         monkeypatch: pytest.MonkeyPatch,
         config_overrides: Callable[..., None],
-        plan: CloudPlan | None,
+        plan: CloudPlan,
         paid_quota: int | None,
         paid_used: int,
         trial_quota: int | None,
@@ -273,7 +274,7 @@ class TestWorkspaceCreditPoolCases:
         monkeypatch.setattr(
             gateways.BillingService,
             "get_info",
-            lambda *_a, **_k: {"subscription": {"plan": plan}, "next_credit_reset_date": 1775001600},
+            lambda *_a, **_k: builders.subscription_info(plan, next_credit_reset_date=1775001600),
         )
         paid = CreditPoolBalance("tenant", "paid", paid_quota, paid_used) if paid_quota is not None else None
         trial = CreditPoolBalance("tenant", "trial", trial_quota, 40) if trial_quota is not None else None
@@ -337,7 +338,7 @@ class TestWorkspaceCreditPoolCases:
     @pytest.mark.parametrize(
         ("edition", "bulk", "expected", "fallback_ids"),
         [
-            (DeploymentEdition.CLOUD, {"w1": {"plan": "team"}}, {"w1": "team", "w2": "professional"}, ["w2"]),
+            (DeploymentEdition.CLOUD, {"w1": builders.plan_item("team")}, {"w1": "team", "w2": "professional"}, ["w2"]),
             (DeploymentEdition.CLOUD, {}, {"w1": "professional", "w2": "professional"}, ["w1", "w2"]),
             (DeploymentEdition.COMMUNITY, {}, {"w1": "professional", "w2": "professional"}, ["w1", "w2"]),
             (DeploymentEdition.ENTERPRISE, {}, {"w1": "sandbox", "w2": "sandbox"}, []),

packages/contracts/generated/api/console/features/types.gen.ts

@@ -72,7 +72,7 @@ export type LicenseLimitationModel = {
 }
 
 export type SubscriptionModel = {
-  interval: string
+  interval: '' | 'month' | 'year'
   plan: CloudPlan
 }
 

packages/contracts/generated/api/console/features/zod.gen.ts

@@ -87,7 +87,7 @@ export const zCloudPlan = z.enum(['professional', 'sandbox', 'team'])
  * SubscriptionModel
  */
 export const zSubscriptionModel = z.object({
-  interval: z.string().default(''),
+  interval: z.enum(['', 'month', 'year']).default(''),
   plan: zCloudPlan.default('sandbox'),
 })