diff --git a/api/controllers/console/workspace/workspace.py b/api/controllers/console/workspace/workspace.py index a06b4fd195..42874e6033 100644 --- a/api/controllers/console/workspace/workspace.py +++ b/api/controllers/console/workspace/workspace.py @@ -28,7 +28,7 @@ from enums.cloud_plan import CloudPlan from extensions.ext_database import db from libs.helper import TimestampField from libs.login import current_account_with_tenant, login_required -from models.account import Tenant, TenantStatus +from models.account import Tenant, TenantCustomConfigDict, TenantStatus from services.account_service import TenantService from services.billing_service import BillingService, SubscriptionPlan from services.enterprise.enterprise_service import EnterpriseService @@ -240,8 +240,10 @@ class CustomConfigWorkspaceApi(Resource): args = WorkspaceCustomConfigPayload.model_validate(payload) tenant = db.get_or_404(Tenant, current_tenant_id) - custom_config_dict = { - "remove_webapp_brand": args.remove_webapp_brand, + custom_config_dict: TenantCustomConfigDict = { + "remove_webapp_brand": args.remove_webapp_brand + if args.remove_webapp_brand is not None + else tenant.custom_config_dict.get("remove_webapp_brand", False), "replace_webapp_logo": args.replace_webapp_logo if args.replace_webapp_logo is not None else tenant.custom_config_dict.get("replace_webapp_logo"), diff --git a/api/models/account.py b/api/models/account.py index 5960ac6564..a3074c6f63 100644 --- a/api/models/account.py +++ b/api/models/account.py @@ -2,7 +2,7 @@ import enum import json from dataclasses import field from datetime import datetime -from typing import Any, Optional +from typing import Optional, TypedDict from uuid import uuid4 import sqlalchemy as sa @@ -232,6 +232,11 @@ class TenantStatus(enum.StrEnum): ARCHIVE = "archive" +class TenantCustomConfigDict(TypedDict, total=False): + remove_webapp_brand: bool + replace_webapp_logo: str | None + + class Tenant(TypeBase): __tablename__ = "tenants" __table_args__ = (sa.PrimaryKeyConstraint("id", name="tenant_pkey"),) @@ -263,11 +268,11 @@ class Tenant(TypeBase): ) @property - def custom_config_dict(self) -> dict[str, Any]: + def custom_config_dict(self) -> TenantCustomConfigDict: return json.loads(self.custom_config) if self.custom_config else {} @custom_config_dict.setter - def custom_config_dict(self, value: dict[str, Any]) -> None: + def custom_config_dict(self, value: TenantCustomConfigDict) -> None: self.custom_config = json.dumps(value) diff --git a/api/models/tools.py b/api/models/tools.py index d8731fb8a8..02f8b5217d 100644 --- a/api/models/tools.py +++ b/api/models/tools.py @@ -356,7 +356,7 @@ class MCPToolProvider(TypeBase): return {} @property - def headers(self) -> dict[str, Any]: + def headers(self) -> dict[str, str]: if self.encrypted_headers is None: return {} try: