From 0c51942b6689ca71203ed07fe9563ea4e1b003fa Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Wed, 2 Sep 2026 17:26:20 +0800 Subject: [PATCH] feat(api): bootstrap Tokener for new tenants --- api/configs/feature/__init__.py | 47 ++ api/events/event_handlers/__init__.py | 2 + ...e_tokener_bootstrap_when_tenant_created.py | 23 + api/extensions/ext_celery.py | 9 + ...2b4d5e6_add_tenant_tokener_integrations.py | 63 +++ api/models/__init__.py | 3 + api/models/tokener.py | 84 +++ api/services/account_service.py | 8 + api/services/billing_service.py | 77 +++ api/tasks/bootstrap_tokener_tenant_task.py | 509 ++++++++++++++++++ .../unit_tests/configs/test_dify_config.py | 19 + ...e_tokener_bootstrap_when_tenant_created.py | 52 ++ .../unit_tests/extensions/test_celery_ssl.py | 11 + .../services/test_billing_service.py | 72 ++- .../services/test_tokener_tenant_creation.py | 54 ++ .../test_bootstrap_tokener_tenant_task.py | 323 +++++++++++ docker/envs/core-services/shared.env.example | 12 + 17 files changed, 1367 insertions(+), 1 deletion(-) create mode 100644 api/events/event_handlers/queue_tokener_bootstrap_when_tenant_created.py create mode 100644 api/migrations/versions/2026_09_02_1730-c3f1a2b4d5e6_add_tenant_tokener_integrations.py create mode 100644 api/models/tokener.py create mode 100644 api/tasks/bootstrap_tokener_tenant_task.py create mode 100644 api/tests/unit_tests/events/event_handlers/test_queue_tokener_bootstrap_when_tenant_created.py create mode 100644 api/tests/unit_tests/services/test_tokener_tenant_creation.py create mode 100644 api/tests/unit_tests/tasks/test_bootstrap_tokener_tenant_task.py diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 3fd5cfaac6e..5e552798c41 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -318,6 +318,41 @@ class PluginConfig(BaseSettings): default="", ) + TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED: bool = Field( + description="Provision and configure the managed Tokener model provider for newly created tenants", + default=False, + ) + + TOKENER_BILLING_API_URL: str = Field( + description="Internal dify-saas billing base URL used only by the Tokener bootstrap client", + default="", + ) + + TOKENER_PLUGIN_UNIQUE_IDENTIFIER: str = Field( + description="Pinned package identifier used for the managed Tokener plugin", + default="", + ) + + TOKENER_PLUGIN_INSTALL_SOURCE: Literal["marketplace", "package"] = Field( + description="Install the pinned Tokener plugin from Marketplace or a package pre-uploaded to plugin-daemon", + default="marketplace", + ) + + TOKENER_PROVIDER_NAME: str = Field( + description="Canonical model-provider name declared by the managed Tokener plugin", + default="langgenius/tokener/tokener", + ) + + TOKENER_DEFAULT_LLM_MODEL: str = Field( + description="Tokener LLM selected as the default for newly created tenants", + default="deepseek-v4-flash", + ) + + TOKENER_ENDPOINT_URL: str = Field( + description="Optional allowlisted Tokener data-plane endpoint passed to a compatible plugin package", + default="", + ) + @property def NEW_USER_DEFAULT_MODEL_LIST(self) -> list[tuple[str, str, str]]: default_models: list[tuple[str, str, str]] = [] @@ -1394,6 +1429,18 @@ class CeleryBeatConfig(BaseSettings): class CeleryScheduleTasksConfig(BaseSettings): + ENABLE_TOKENER_BOOTSTRAP_RECOVERY_TASK: bool = Field( + description="Enable periodic recovery of incomplete new-tenant Tokener bootstraps", + default=True, + ) + TOKENER_BOOTSTRAP_RECOVERY_TASK_INTERVAL: PositiveInt = Field( + description="Minimum age and periodic recovery interval for incomplete Tokener bootstraps, in minutes", + default=5, + ) + TOKENER_BOOTSTRAP_RECOVERY_BATCH_SIZE: PositiveInt = Field( + description="Maximum number of incomplete Tokener bootstraps requeued per recovery sweep", + default=100, + ) ENABLE_CONVERSATION_CLEANUP_TASK: bool = Field( description="Enable periodic recovery of soft-deleted conversation cleanup", default=True, diff --git a/api/events/event_handlers/__init__.py b/api/events/event_handlers/__init__.py index b60d86b2b54..01a7f62d7de 100644 --- a/api/events/event_handlers/__init__.py +++ b/api/events/event_handlers/__init__.py @@ -10,6 +10,7 @@ from .queue_credential_sync_when_tenant_created import handle as handle_queue_cr from .queue_default_plugin_install_when_tenant_created import ( handle as handle_queue_default_plugin_install_when_tenant_created, ) +from .queue_tokener_bootstrap_when_tenant_created import handle as handle_queue_tokener_bootstrap_when_tenant_created from .sync_plugin_trigger_when_app_created import handle as handle_sync_plugin_trigger_when_app_created from .sync_webhook_when_app_created import handle as handle_sync_webhook_when_app_created from .sync_workflow_schedule_when_app_published import handle as handle_sync_workflow_schedule_when_app_published @@ -36,6 +37,7 @@ __all__ = [ "handle_delete_tool_parameters_cache_when_sync_draft_workflow", "handle_queue_credential_sync_when_tenant_created", "handle_queue_default_plugin_install_when_tenant_created", + "handle_queue_tokener_bootstrap_when_tenant_created", "handle_sync_plugin_trigger_when_app_created", "handle_sync_webhook_when_app_created", "handle_sync_workflow_schedule_when_app_published", diff --git a/api/events/event_handlers/queue_tokener_bootstrap_when_tenant_created.py b/api/events/event_handlers/queue_tokener_bootstrap_when_tenant_created.py new file mode 100644 index 00000000000..fbb5dca69b3 --- /dev/null +++ b/api/events/event_handlers/queue_tokener_bootstrap_when_tenant_created.py @@ -0,0 +1,23 @@ +"""Queue managed Tokener provisioning after a tenant is committed.""" + +import logging + +from configs import dify_config +from events.tenant_event import tenant_was_created +from tasks.bootstrap_tokener_tenant_task import bootstrap_tokener_tenant_task + +logger = logging.getLogger(__name__) + + +@tenant_was_created.connect +def handle(sender, **kwargs) -> None: + if not dify_config.TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED: + return + + try: + bootstrap_tokener_tenant_task.delay(sender.id) + except Exception: + # Registration itself is already committed. The durable integration row + # remains pending so an operator or recovery sweep can safely requeue it. + # Do not serialize broker exception details into a registration request log. + logger.error("Failed to queue Tokener bootstrap for tenant %s", sender.id) # noqa: TRY400 diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index 724fb4774cc..e27041d1029 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -172,6 +172,7 @@ def init_app(app: DifyApp) -> Celery: "tasks.generate_summary_index_task", # summary index generation "tasks.regenerate_summary_index_task", # summary index regeneration "tasks.initialize_created_app_rbac_access_task", # app access initialization + "tasks.bootstrap_tokener_tenant_task", # managed Tokener setup for new tenants "tasks.install_default_plugins_task", # tenant default plugin installation "tasks.new_agent_beta_task", # New Agent Beta eligibility checks "tasks.refresh_billing_vector_space_task", # billing vector-space cache refresh @@ -182,6 +183,14 @@ def init_app(app: DifyApp) -> Celery: # if you add a new task, please add the switch to CeleryScheduleTasksConfig beat_schedule: dict[str, CeleryBeatScheduleEntry] = {} + if ( + dify_config.TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED + and dify_config.ENABLE_TOKENER_BOOTSTRAP_RECOVERY_TASK + ): + beat_schedule["tokener_bootstrap_recovery_sweeper"] = { + "task": "tasks.bootstrap_tokener_tenant_task.sweep_pending_tokener_integrations_task", + "schedule": timedelta(minutes=dify_config.TOKENER_BOOTSTRAP_RECOVERY_TASK_INTERVAL), + } if dify_config.ENABLE_CONVERSATION_CLEANUP_TASK: imports.append("tasks.delete_conversation_task") beat_schedule["conversation_cleanup_sweeper"] = { diff --git a/api/migrations/versions/2026_09_02_1730-c3f1a2b4d5e6_add_tenant_tokener_integrations.py b/api/migrations/versions/2026_09_02_1730-c3f1a2b4d5e6_add_tenant_tokener_integrations.py new file mode 100644 index 00000000000..fbf28eae02f --- /dev/null +++ b/api/migrations/versions/2026_09_02_1730-c3f1a2b4d5e6_add_tenant_tokener_integrations.py @@ -0,0 +1,63 @@ +"""add tenant tokener integrations + +Revision ID: c3f1a2b4d5e6 +Revises: 5578e028b2f2 +Create Date: 2026-09-02 17:30:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +import models + +# revision identifiers, used by Alembic. +revision = "c3f1a2b4d5e6" +down_revision = "5578e028b2f2" +branch_labels = None +depends_on = None + + +def upgrade(): + op.create_table( + "tenant_tokener_integrations", + sa.Column("id", models.types.StringUUID(), nullable=False), + sa.Column("tenant_id", models.types.StringUUID(), nullable=False), + sa.Column("status", sa.String(length=40), server_default=sa.text("'pending'"), nullable=False), + sa.Column("plugin_unique_identifier", sa.String(length=255), nullable=True), + sa.Column("plugin_install_task_id", sa.String(length=255), nullable=True), + sa.Column("provider_credential_id", models.types.StringUUID(), nullable=True), + sa.Column("attempt_count", sa.Integer(), server_default=sa.text("0"), nullable=False), + sa.Column("last_error_code", sa.String(length=100), nullable=True), + sa.Column("last_attempt_at", sa.DateTime(), nullable=True), + sa.Column("ready_at", sa.DateTime(), nullable=True), + sa.Column("created_at", sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False), + sa.Column("updated_at", sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False), + sa.ForeignKeyConstraint( + ["provider_credential_id"], + ["provider_credentials.id"], + name=op.f("tenant_tokener_integrations_provider_credential_id_fkey"), + ondelete="SET NULL", + ), + sa.ForeignKeyConstraint( + ["tenant_id"], + ["tenants.id"], + name=op.f("tenant_tokener_integrations_tenant_id_fkey"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("tenant_tokener_integration_pkey")), + sa.UniqueConstraint("tenant_id", name=op.f("tenant_tokener_integration_tenant_id_key")), + ) + op.create_index( + "tenant_tokener_integration_status_updated_at_idx", + "tenant_tokener_integrations", + ["status", "updated_at"], + ) + + +def downgrade(): + op.drop_index( + "tenant_tokener_integration_status_updated_at_idx", + table_name="tenant_tokener_integrations", + ) + op.drop_table("tenant_tokener_integrations") diff --git a/api/models/__init__.py b/api/models/__init__.py index 67f22c3f4da..1161ebb83a6 100644 --- a/api/models/__init__.py +++ b/api/models/__init__.py @@ -116,6 +116,7 @@ from .skill import AgentSkillBinding, Skill, SkillDraftFile, SkillFileKind, Skil from .snippet import CustomizedSnippet, SnippetType from .source import DataSourceApiKeyAuthBinding, DataSourceOauthBinding from .task import CeleryTask, CeleryTaskSet +from .tokener import TenantTokenerIntegration, TenantTokenerIntegrationStatus from .tools import ( ApiToolProvider, BuiltinToolProvider, @@ -264,6 +265,8 @@ __all__ = [ "TenantDefaultModel", "TenantPreferredModelProvider", "TenantStatus", + "TenantTokenerIntegration", + "TenantTokenerIntegrationStatus", "TidbAuthBinding", "ToolConversationVariables", "ToolFile", diff --git a/api/models/tokener.py b/api/models/tokener.py new file mode 100644 index 00000000000..164f91b1789 --- /dev/null +++ b/api/models/tokener.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from datetime import datetime +from enum import StrEnum + +import sqlalchemy as sa +from sqlalchemy import DateTime, String, func, text +from sqlalchemy.orm import Mapped, mapped_column + +from libs.uuid_utils import uuidv7 + +from .base import TypeBase +from .types import EnumText, StringUUID + + +class TenantTokenerIntegrationStatus(StrEnum): + """Durable stages for bootstrapping Tokener in a workspace.""" + + PENDING = "pending" + INSTALLING_PLUGIN = "installing_plugin" + PROVISIONING = "provisioning" + CONFIGURING_PROVIDER = "configuring_provider" + READY = "ready" + FAILED = "failed" + + +class TenantTokenerIntegration(TypeBase): + """Non-secret, replayable state for a tenant's managed Tokener setup.""" + + __tablename__ = "tenant_tokener_integrations" + __table_args__ = ( + sa.PrimaryKeyConstraint("id", name="tenant_tokener_integration_pkey"), + sa.UniqueConstraint("tenant_id", name="tenant_tokener_integration_tenant_id_key"), + sa.Index("tenant_tokener_integration_status_updated_at_idx", "status", "updated_at"), + ) + + id: Mapped[str] = mapped_column( + StringUUID, + primary_key=True, + insert_default=lambda: str(uuidv7()), + default_factory=lambda: str(uuidv7()), + init=False, + ) + tenant_id: Mapped[str] = mapped_column( + StringUUID, + sa.ForeignKey("tenants.id", ondelete="CASCADE"), + nullable=False, + ) + status: Mapped[TenantTokenerIntegrationStatus] = mapped_column( + EnumText(TenantTokenerIntegrationStatus, length=40), + nullable=False, + server_default=text("'pending'"), + default=TenantTokenerIntegrationStatus.PENDING, + ) + plugin_unique_identifier: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None) + plugin_install_task_id: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None) + provider_credential_id: Mapped[str | None] = mapped_column( + StringUUID, + sa.ForeignKey("provider_credentials.id", ondelete="SET NULL"), + nullable=True, + default=None, + ) + attempt_count: Mapped[int] = mapped_column( + sa.Integer, + nullable=False, + server_default=text("0"), + default=0, + ) + last_error_code: Mapped[str | None] = mapped_column(String(100), nullable=True, default=None) + last_attempt_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) + ready_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) + created_at: Mapped[datetime] = mapped_column( + DateTime, + nullable=False, + server_default=func.current_timestamp(), + init=False, + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime, + nullable=False, + server_default=func.current_timestamp(), + onupdate=func.current_timestamp(), + init=False, + ) diff --git a/api/services/account_service.py b/api/services/account_service.py index 6e5fda0b0ba..893a1e1ce3a 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -47,6 +47,7 @@ from models.account import ( ) from models.dataset import Dataset from models.model import App, DifySetup +from models.tokener import TenantTokenerIntegration from services.account_email import normalize_email from services.billing_service import BillingService from services.email_code_login_challenge import ( @@ -1146,6 +1147,13 @@ class TenantService: tenant = Tenant(name=name) session.add(tenant) + if dify_config.TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED: + session.add( + TenantTokenerIntegration( + tenant_id=tenant.id, + plugin_unique_identifier=dify_config.TOKENER_PLUGIN_UNIQUE_IDENTIFIER.strip() or None, + ) + ) session.commit() for category in TenantPluginAutoUpgradeCategory: diff --git a/api/services/billing_service.py b/api/services/billing_service.py index f89547f4558..0d5342d6392 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -10,6 +10,7 @@ from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fix from typing_extensions import deprecated from werkzeug.exceptions import InternalServerError +from configs import dify_config from core.helper.http_client_pooling import get_pooled_http_client from enums import CloudPlan from extensions.ext_redis import redis_client @@ -40,6 +41,15 @@ class _BillingHTTPStatusError(ValueError): self.status_code = status_code +class TokenerBootstrapUpstreamError(RuntimeError): + """Sanitized failure returned by the billing-side Tokener bootstrap API.""" + + def __init__(self, error_code: str, *, retryable: bool) -> None: + super().__init__(error_code) + self.error_code = error_code + self.retryable = retryable + + class SubscriptionPlan(TypedDict): """Tenant subscriptionplan information.""" @@ -62,6 +72,14 @@ class EducationStatusResponseDict(TypedDict): allow_refresh: bool +class TokenerTenantBootstrapResponse(TypedDict): + tenant_id: str + status: Literal["pending", "ready"] + data_plane_api_key: NotRequired[str] + retryable: bool + error_code: NotRequired[str] + + class EducationAutocompleteResponseDict(TypedDict): data: list[str] curr_page: int @@ -245,6 +263,65 @@ class BillingService: def ensure_new_agent_beta_workflow(cls, workflow_id: str) -> None: cls._send_request("POST", f"/new-agent-beta/workflows/{workflow_id}/ensure") + @classmethod + def bootstrap_tokener_tenant(cls, tenant_id: str, display_name: str) -> TokenerTenantBootstrapResponse: + """Ensure the remote Tokener org and trial allowance without logging its one-time key.""" + base_url = dify_config.TOKENER_BILLING_API_URL.strip().rstrip("/") + if not base_url: + raise TokenerBootstrapUpstreamError("tokener_billing_api_not_configured", retryable=False) + try: + payload = cls._send_request( + "POST", + f"/internal/v1/tokener/tenants/{tenant_id}/bootstrap", + json={"tenant_id": tenant_id, "display_name": display_name}, + base_url=base_url, + ) + except _BillingHTTPStatusError as error: + retryable = error.status_code in {httpx.codes.REQUEST_TIMEOUT, httpx.codes.TOO_MANY_REQUESTS} or ( + error.status_code >= 500 + ) + raise TokenerBootstrapUpstreamError("tokener_bootstrap_http_error", retryable=retryable) from None + except httpx.RequestError: + raise TokenerBootstrapUpstreamError("tokener_bootstrap_unavailable", retryable=True) from None + + if not isinstance(payload, dict): + raise TokenerBootstrapUpstreamError("tokener_bootstrap_invalid_response", retryable=False) + + raw_data_plane_api_key = payload.pop("data_plane_api_key", None) + response_tenant_id = payload.get("tenant_id") + status = payload.get("status") + retryable = payload.get("retryable", status == "pending") + raw_error_code = payload.get("error_code") + error_code = ( + raw_error_code + if isinstance(raw_error_code, str) + and raw_error_code + and len(raw_error_code) <= 100 + and all(character.isalnum() or character in "_-" for character in raw_error_code) + else None + ) + + if response_tenant_id != tenant_id or status not in {"pending", "ready"} or not isinstance(retryable, bool): + raw_data_plane_api_key = None + raise TokenerBootstrapUpstreamError("tokener_bootstrap_invalid_response", retryable=False) + + response: TokenerTenantBootstrapResponse = { + "tenant_id": tenant_id, + "status": status, + "retryable": retryable, + } + if error_code: + response["error_code"] = error_code + + if status == "ready": + if not isinstance(raw_data_plane_api_key, str) or not raw_data_plane_api_key: + raw_data_plane_api_key = None + raise TokenerBootstrapUpstreamError("tokener_bootstrap_key_missing", retryable=True) + response["data_plane_api_key"] = raw_data_plane_api_key + raw_data_plane_api_key = None + + return response + @classmethod def get_info(cls, tenant_id: str, exclude_vector_space: bool = False) -> BillingInfo: params = {"tenant_id": tenant_id} diff --git a/api/tasks/bootstrap_tokener_tenant_task.py b/api/tasks/bootstrap_tokener_tenant_task.py new file mode 100644 index 00000000000..4149224a20b --- /dev/null +++ b/api/tasks/bootstrap_tokener_tenant_task.py @@ -0,0 +1,509 @@ +"""Provision and configure the managed Tokener provider for a new tenant.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from datetime import timedelta +from typing import cast + +from celery import shared_task +from sqlalchemy import select + +from configs import dify_config +from core.db.session_factory import session_factory +from core.plugin.entities.plugin_daemon import PluginInstallTaskStatus +from core.plugin.plugin_service import PluginService +from extensions.ext_redis import redis_client +from libs.datetime_utils import naive_utc_now +from models.account import Tenant +from models.provider import Provider, ProviderCredential, ProviderType +from models.tokener import TenantTokenerIntegration, TenantTokenerIntegrationStatus +from services.billing_service import BillingService, TokenerBootstrapUpstreamError +from services.model_provider_service import ModelProviderService + +logger = logging.getLogger(__name__) + +TOKENER_BOOTSTRAP_QUEUE = "plugin" +MANAGED_TOKENER_CREDENTIAL_NAME = "__dify_managed_tokener_v1__" + +_MAX_RETRIES = 60 +_RETRY_DELAY_SECONDS = 5 +_MAX_RETRY_DELAY_SECONDS = 60 +_LOCK_TIMEOUT_SECONDS = 15 * 60 +_UNSET = object() + + +@dataclass(frozen=True, slots=True) +class _IntegrationSnapshot: + tenant_id: str + tenant_name: str + status: TenantTokenerIntegrationStatus + plugin_unique_identifier: str | None + plugin_install_task_id: str | None + + +@dataclass(frozen=True, slots=True) +class _CredentialWriteResult: + """Secret-free result returned after consuming a one-time data-plane key.""" + + credential_id: str | None = None + error_code: str | None = None + + +class _BootstrapStepError(RuntimeError): + """A secret-free task failure safe to persist and hand to Celery.""" + + def __init__( + self, + error_code: str, + *, + retryable: bool, + stage: TenantTokenerIntegrationStatus, + ) -> None: + super().__init__(error_code) + self.error_code = error_code + self.retryable = retryable + self.stage = stage + + +def _begin_attempt(tenant_id: str) -> _IntegrationSnapshot | None: + now = naive_utc_now() + with session_factory.create_session() as session, session.begin(): + integration = session.scalar( + select(TenantTokenerIntegration) + .where(TenantTokenerIntegration.tenant_id == tenant_id) + .with_for_update() + ) + if integration is None: + return None + + tenant = session.get(Tenant, tenant_id) + if tenant is None: + return None + + if integration.status != TenantTokenerIntegrationStatus.READY: + integration.attempt_count += 1 + integration.last_attempt_at = now + integration.last_error_code = None + + return _IntegrationSnapshot( + tenant_id=tenant_id, + tenant_name=tenant.name, + status=integration.status, + plugin_unique_identifier=integration.plugin_unique_identifier, + plugin_install_task_id=integration.plugin_install_task_id, + ) + + +def _update_integration( + tenant_id: str, + *, + status: TenantTokenerIntegrationStatus, + plugin_install_task_id: str | None | object = _UNSET, + provider_credential_id: str | None = None, + error_code: str | None = None, + ready: bool = False, +) -> None: + with session_factory.create_session() as session, session.begin(): + integration = session.scalar( + select(TenantTokenerIntegration) + .where(TenantTokenerIntegration.tenant_id == tenant_id) + .with_for_update() + ) + if integration is None: + return + + integration.status = status + if plugin_install_task_id is not _UNSET: + integration.plugin_install_task_id = cast(str | None, plugin_install_task_id) + integration.last_error_code = error_code + if provider_credential_id is not None: + integration.provider_credential_id = provider_credential_id + if ready: + integration.ready_at = naive_utc_now() + + +def _installed_plugin_matches(tenant_id: str, plugin_unique_identifier: str) -> bool: + plugin_id = plugin_unique_identifier.split(":", 1)[0] + installed_plugins = PluginService.list(tenant_id) + exact_match = any( + plugin.plugin_id == plugin_id and plugin.plugin_unique_identifier == plugin_unique_identifier + for plugin in installed_plugins + ) + if exact_match: + return True + + if any(plugin.plugin_id == plugin_id for plugin in installed_plugins): + raise _BootstrapStepError( + "tokener_plugin_version_conflict", + retryable=False, + stage=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + ) + return False + + +def _ensure_plugin_installed(snapshot: _IntegrationSnapshot) -> None: + plugin_unique_identifier = snapshot.plugin_unique_identifier + if not plugin_unique_identifier: + raise _BootstrapStepError( + "tokener_plugin_not_configured", + retryable=False, + stage=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + ) + + _update_integration( + snapshot.tenant_id, + status=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + plugin_install_task_id=snapshot.plugin_install_task_id, + ) + if _installed_plugin_matches(snapshot.tenant_id, plugin_unique_identifier): + _update_integration( + snapshot.tenant_id, + status=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + plugin_install_task_id=None, + ) + return + + if snapshot.plugin_install_task_id: + install_task = PluginService.fetch_install_task(snapshot.tenant_id, snapshot.plugin_install_task_id) + if install_task.status in {PluginInstallTaskStatus.Pending, PluginInstallTaskStatus.Running}: + raise _BootstrapStepError( + "tokener_plugin_install_pending", + retryable=True, + stage=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + ) + if install_task.status == PluginInstallTaskStatus.Failed: + _update_integration( + snapshot.tenant_id, + status=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + plugin_install_task_id=None, + error_code="tokener_plugin_install_failed", + ) + raise _BootstrapStepError( + "tokener_plugin_install_failed", + retryable=True, + stage=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + ) + if _installed_plugin_matches(snapshot.tenant_id, plugin_unique_identifier): + return + + if dify_config.TOKENER_PLUGIN_INSTALL_SOURCE == "package": + response = PluginService.install_from_local_pkg(snapshot.tenant_id, [plugin_unique_identifier]) + else: + response = PluginService.install_from_marketplace_pkg(snapshot.tenant_id, [plugin_unique_identifier]) + if response.all_installed: + return + if not response.task_id: + raise _BootstrapStepError( + "tokener_plugin_install_task_missing", + retryable=True, + stage=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + ) + + _update_integration( + snapshot.tenant_id, + status=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + plugin_install_task_id=response.task_id, + ) + raise _BootstrapStepError( + "tokener_plugin_install_pending", + retryable=True, + stage=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + ) + + +def _find_managed_credential(tenant_id: str) -> tuple[str | None, bool]: + provider_name = dify_config.TOKENER_PROVIDER_NAME + with session_factory.create_session() as session: + provider = session.scalar( + select(Provider).where( + Provider.tenant_id == tenant_id, + Provider.provider_name == provider_name, + Provider.provider_type == ProviderType.CUSTOM, + ) + ) + credentials = list( + session.scalars( + select(ProviderCredential) + .where( + ProviderCredential.tenant_id == tenant_id, + ProviderCredential.provider_name == provider_name, + ProviderCredential.credential_name == MANAGED_TOKENER_CREDENTIAL_NAME, + ) + .order_by(ProviderCredential.created_at, ProviderCredential.id) + ) + ) + + if not credentials: + return None, False + if provider is not None: + active = next((credential for credential in credentials if credential.id == provider.credential_id), None) + if active is not None and provider.is_valid: + return active.id, True + return credentials[0].id, False + + +def _activate_managed_credential(tenant_id: str, credential_id: str) -> None: + try: + ModelProviderService().switch_active_provider_credential( + tenant_id=tenant_id, + provider=dify_config.TOKENER_PROVIDER_NAME, + credential_id=credential_id, + ) + except Exception: + raise _BootstrapStepError( + "tokener_provider_credential_activation_failed", + retryable=True, + stage=TenantTokenerIntegrationStatus.CONFIGURING_PROVIDER, + ) from None + + +def _consume_data_plane_api_key(tenant_id: str, data_plane_api_key: str) -> _CredentialWriteResult: + """Consume a one-time key without allowing an exception to escape this frame. + + Any exception raised while validating or persisting the credential is converted + to a secret-free value before this frame returns. The caller clears its own + reference before raising the corresponding task error, keeping plaintext out + of traceback locals and Sentry events. + """ + credentials = {"api_key": data_plane_api_key} + endpoint_url = dify_config.TOKENER_ENDPOINT_URL.strip() + if endpoint_url: + credentials["endpoint_url"] = endpoint_url + try: + ModelProviderService().create_provider_credential( + tenant_id=tenant_id, + provider=dify_config.TOKENER_PROVIDER_NAME, + credentials=credentials, + credential_name=MANAGED_TOKENER_CREDENTIAL_NAME, + ) + except Exception: + return _CredentialWriteResult(error_code="tokener_provider_credential_rejected") + finally: + credentials.clear() + data_plane_api_key = "" + + try: + credential_id, active = _find_managed_credential(tenant_id) + if credential_id is None: + return _CredentialWriteResult(error_code="tokener_provider_credential_not_persisted") + if not active: + _activate_managed_credential(tenant_id, credential_id) + return _CredentialWriteResult(credential_id=credential_id) + except _BootstrapStepError as error: + return _CredentialWriteResult(error_code=error.error_code) + except Exception: + return _CredentialWriteResult(error_code="tokener_provider_credential_activation_failed") + + +def _ensure_managed_credential(snapshot: _IntegrationSnapshot) -> str: + credential_id, active = _find_managed_credential(snapshot.tenant_id) + if credential_id is not None: + if not active: + _activate_managed_credential(snapshot.tenant_id, credential_id) + return credential_id + + _update_integration( + snapshot.tenant_id, + status=TenantTokenerIntegrationStatus.PROVISIONING, + plugin_install_task_id=None, + ) + try: + response = BillingService.bootstrap_tokener_tenant(snapshot.tenant_id, snapshot.tenant_name) + except TokenerBootstrapUpstreamError as error: + raise _BootstrapStepError( + error.error_code, + retryable=error.retryable, + stage=TenantTokenerIntegrationStatus.PROVISIONING, + ) from None + + if response["status"] == "pending": + raise _BootstrapStepError( + response.get("error_code", "tokener_bootstrap_pending"), + retryable=response["retryable"], + stage=TenantTokenerIntegrationStatus.PROVISIONING, + ) + + data_plane_api_key = response.pop("data_plane_api_key", "") + write_result = _consume_data_plane_api_key(snapshot.tenant_id, data_plane_api_key) + data_plane_api_key = "" + _update_integration(snapshot.tenant_id, status=TenantTokenerIntegrationStatus.CONFIGURING_PROVIDER) + if write_result.credential_id is not None: + return write_result.credential_id + raise _BootstrapStepError( + write_result.error_code or "tokener_provider_credential_rejected", + retryable=True, + stage=TenantTokenerIntegrationStatus.CONFIGURING_PROVIDER, + ) + + +def _set_default_llm(tenant_id: str) -> None: + try: + ModelProviderService().update_default_model_of_model_type( + tenant_id=tenant_id, + model_type="llm", + provider=dify_config.TOKENER_PROVIDER_NAME, + model=dify_config.TOKENER_DEFAULT_LLM_MODEL, + ) + except Exception: + raise _BootstrapStepError( + "tokener_default_model_configuration_failed", + retryable=True, + stage=TenantTokenerIntegrationStatus.CONFIGURING_PROVIDER, + ) from None + + +def _run_bootstrap(tenant_id: str) -> None: + snapshot = _begin_attempt(tenant_id) + if snapshot is None or snapshot.status == TenantTokenerIntegrationStatus.READY: + return + + _ensure_plugin_installed(snapshot) + credential_id = _ensure_managed_credential(snapshot) + _set_default_llm(snapshot.tenant_id) + _update_integration( + snapshot.tenant_id, + status=TenantTokenerIntegrationStatus.READY, + plugin_install_task_id=None, + provider_credential_id=credential_id, + ready=True, + ) + + +@shared_task( + queue=TOKENER_BOOTSTRAP_QUEUE, + bind=True, + max_retries=_MAX_RETRIES, + default_retry_delay=_RETRY_DELAY_SECONDS, + acks_late=True, + reject_on_worker_lost=True, +) +def bootstrap_tokener_tenant_task(self, tenant_id: str) -> None: + """Drive the idempotent bootstrap while serializing work per tenant.""" + if not dify_config.TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED: + return + + try: + lock = redis_client.lock( + f"tokener:new-tenant-bootstrap:{tenant_id}", + timeout=_LOCK_TIMEOUT_SECONDS, + blocking_timeout=0, + thread_local=False, + ) + acquired = lock.acquire(blocking=False) + except Exception: + lock_error = _BootstrapStepError( + "tokener_bootstrap_lock_unavailable", + retryable=True, + stage=TenantTokenerIntegrationStatus.PENDING, + ) + logger.warning("Tokener bootstrap lock is unavailable for tenant %s; scheduling retry", tenant_id) + raise self.retry(exc=lock_error, countdown=_RETRY_DELAY_SECONDS) from None + + if not acquired: + raise self.retry( + exc=_BootstrapStepError( + "tokener_bootstrap_locked", + retryable=True, + stage=TenantTokenerIntegrationStatus.PENDING, + ), + countdown=_RETRY_DELAY_SECONDS, + ) + + try: + _run_bootstrap(tenant_id) + except _BootstrapStepError as step_error: + exhausted = self.request.retries >= _MAX_RETRIES + _update_integration( + tenant_id, + status=TenantTokenerIntegrationStatus.FAILED if exhausted or not step_error.retryable else step_error.stage, + error_code=step_error.error_code, + ) + if step_error.retryable and not exhausted: + countdown = min(_RETRY_DELAY_SECONDS * (2**self.request.retries), _MAX_RETRY_DELAY_SECONDS) + logger.warning( + "Tokener bootstrap will retry for tenant %s, error_code=%s, retry=%s/%s", + tenant_id, + step_error.error_code, + self.request.retries + 1, + _MAX_RETRIES, + ) + raise self.retry(exc=step_error, countdown=countdown) + + # Deliberately omit exception serialization from this secret-adjacent task. + logger.error( # noqa: TRY400 + "Tokener bootstrap stopped for tenant %s, error_code=%s", + tenant_id, + step_error.error_code, + ) + raise + except Exception: + sanitized_error = _BootstrapStepError( + "tokener_bootstrap_internal_error", + retryable=True, + stage=TenantTokenerIntegrationStatus.FAILED, + ) + exhausted = self.request.retries >= _MAX_RETRIES + _update_integration( + tenant_id, + status=TenantTokenerIntegrationStatus.FAILED, + error_code=sanitized_error.error_code, + ) + if not exhausted: + countdown = min(_RETRY_DELAY_SECONDS * (2**self.request.retries), _MAX_RETRY_DELAY_SECONDS) + logger.warning( + "Tokener bootstrap hit an internal error for tenant %s; scheduling retry %s/%s", + tenant_id, + self.request.retries + 1, + _MAX_RETRIES, + ) + raise self.retry(exc=sanitized_error, countdown=countdown) from None + # Deliberately omit exception serialization from this secret-adjacent task. + logger.error("Tokener bootstrap retry budget exhausted for tenant %s", tenant_id) # noqa: TRY400 + raise sanitized_error from None + finally: + try: + lock.release() + except Exception: + logger.warning("Tokener bootstrap lock expired before release for tenant %s", tenant_id) + + +@shared_task(queue=TOKENER_BOOTSTRAP_QUEUE) +def sweep_pending_tokener_integrations_task() -> int: + """Recover bootstraps whose initial broker dispatch or retry chain was lost.""" + if not dify_config.TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED: + return 0 + + stale_before = naive_utc_now() - timedelta( + minutes=dify_config.TOKENER_BOOTSTRAP_RECOVERY_TASK_INTERVAL, + ) + recoverable_statuses = ( + TenantTokenerIntegrationStatus.PENDING, + TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + TenantTokenerIntegrationStatus.PROVISIONING, + TenantTokenerIntegrationStatus.CONFIGURING_PROVIDER, + ) + with session_factory.create_session() as session: + tenant_ids = list( + session.scalars( + select(TenantTokenerIntegration.tenant_id) + .where( + TenantTokenerIntegration.status.in_(recoverable_statuses), + TenantTokenerIntegration.updated_at < stale_before, + ) + .order_by(TenantTokenerIntegration.updated_at, TenantTokenerIntegration.id) + .limit(dify_config.TOKENER_BOOTSTRAP_RECOVERY_BATCH_SIZE) + ) + ) + + dispatched = 0 + for tenant_id in tenant_ids: + try: + bootstrap_tokener_tenant_task.delay(tenant_id) + except Exception: + # Do not attach broker exception details; the next beat will retry. + logger.error("Failed to recover Tokener bootstrap for tenant %s", tenant_id) # noqa: TRY400 + else: + dispatched += 1 + return dispatched diff --git a/api/tests/unit_tests/configs/test_dify_config.py b/api/tests/unit_tests/configs/test_dify_config.py index 8e9547ad68f..c918e583dd4 100644 --- a/api/tests/unit_tests/configs/test_dify_config.py +++ b/api/tests/unit_tests/configs/test_dify_config.py @@ -208,6 +208,25 @@ def test_new_user_default_models_reject_duplicate_model_types() -> None: _ = config.NEW_USER_DEFAULT_MODEL_LIST +def test_tokener_new_tenant_bootstrap_config_is_parsed() -> None: + config = _make_config( + TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED="true", + TOKENER_BILLING_API_URL="http://dify-saas-billing:8081", + TOKENER_PLUGIN_UNIQUE_IDENTIFIER="langgenius/tokener:0.1.2@checksum", + TOKENER_PLUGIN_INSTALL_SOURCE="package", + TOKENER_ENDPOINT_URL="https://api-staging.tokener.dev/v1", + TOKENER_BOOTSTRAP_RECOVERY_TASK_INTERVAL="7", + TOKENER_BOOTSTRAP_RECOVERY_BATCH_SIZE="50", + ) + + assert config.TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED is True + assert config.TOKENER_BILLING_API_URL == "http://dify-saas-billing:8081" + assert config.TOKENER_PLUGIN_INSTALL_SOURCE == "package" + assert config.TOKENER_ENDPOINT_URL == "https://api-staging.tokener.dev/v1" + assert config.TOKENER_BOOTSTRAP_RECOVERY_TASK_INTERVAL == 7 + assert config.TOKENER_BOOTSTRAP_RECOVERY_BATCH_SIZE == 50 + + def test_http_timeout_defaults(): """Test that HTTP timeout defaults are correctly set""" config = _make_config() diff --git a/api/tests/unit_tests/events/event_handlers/test_queue_tokener_bootstrap_when_tenant_created.py b/api/tests/unit_tests/events/event_handlers/test_queue_tokener_bootstrap_when_tenant_created.py new file mode 100644 index 00000000000..1d27067c4a7 --- /dev/null +++ b/api/tests/unit_tests/events/event_handlers/test_queue_tokener_bootstrap_when_tenant_created.py @@ -0,0 +1,52 @@ +import logging +from unittest.mock import MagicMock + +import pytest + +from events.event_handlers import queue_tokener_bootstrap_when_tenant_created as handler_module +from models.account import Tenant +from tests.unit_tests.config_override import apply_config_overrides + + +def _tenant() -> Tenant: + tenant = Tenant(name="Tokener tenant") + tenant.id = "tenant-1" + return tenant + + +def test_handle_skips_when_tokener_bootstrap_is_disabled(monkeypatch: pytest.MonkeyPatch) -> None: + delay = MagicMock() + apply_config_overrides(monkeypatch, TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED=False) + monkeypatch.setattr(handler_module.bootstrap_tokener_tenant_task, "delay", delay) + + handler_module.handle(_tenant()) + + delay.assert_not_called() + + +def test_handle_queues_tokener_bootstrap(monkeypatch: pytest.MonkeyPatch) -> None: + delay = MagicMock() + apply_config_overrides(monkeypatch, TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED=True) + monkeypatch.setattr(handler_module.bootstrap_tokener_tenant_task, "delay", delay) + + handler_module.handle(_tenant()) + + delay.assert_called_once_with("tenant-1") + + +def test_handle_keeps_registration_successful_when_broker_is_unavailable( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + apply_config_overrides(monkeypatch, TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED=True) + monkeypatch.setattr( + handler_module.bootstrap_tokener_tenant_task, + "delay", + MagicMock(side_effect=ConnectionError("broker unavailable")), + ) + + with caplog.at_level(logging.ERROR, logger=handler_module.logger.name): + handler_module.handle(_tenant()) + + assert "Failed to queue Tokener bootstrap for tenant tenant-1" in caplog.text + assert "broker unavailable" not in caplog.text diff --git a/api/tests/unit_tests/extensions/test_celery_ssl.py b/api/tests/unit_tests/extensions/test_celery_ssl.py index c5df0a3b7ff..8a6cb0dcaee 100644 --- a/api/tests/unit_tests/extensions/test_celery_ssl.py +++ b/api/tests/unit_tests/extensions/test_celery_ssl.py @@ -162,6 +162,8 @@ class TestCelerySSLConfiguration: # Mock all the scheduler configs mock_config.CELERY_BEAT_SCHEDULER_TIME = 1 + mock_config.TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED = False + mock_config.ENABLE_TOKENER_BOOTSTRAP_RECOVERY_TASK = False mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = False mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False @@ -212,6 +214,9 @@ class TestCelerySSLConfiguration: mock_config.CELERY_TASK_ANNOTATIONS = {} mock_config.CELERY_BEAT_SCHEDULER_TIME = 1 + mock_config.TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED = True + mock_config.ENABLE_TOKENER_BOOTSTRAP_RECOVERY_TASK = True + mock_config.TOKENER_BOOTSTRAP_RECOVERY_TASK_INTERVAL = 7 mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = True mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False @@ -245,7 +250,13 @@ class TestCelerySSLConfiguration: assert celery_app.conf["broker_transport_options"]["global_keyprefix"] == "enterprise-a:" assert celery_app.conf["result_backend_transport_options"]["global_keyprefix"] == "enterprise-a:" assert "tasks.collect_agent_resources_task" in celery_app.conf["imports"] + assert "tasks.bootstrap_tokener_tenant_task" in celery_app.conf["imports"] assert "tasks.delete_conversation_task" in celery_app.conf["imports"] + recovery_schedule = celery_app.conf["beat_schedule"]["tokener_bootstrap_recovery_sweeper"] + assert recovery_schedule["task"] == ( + "tasks.bootstrap_tokener_tenant_task.sweep_pending_tokener_integrations_task" + ) + assert recovery_schedule["schedule"].total_seconds() == 7 * 60 assert celery_app.conf["beat_schedule"]["conversation_cleanup_sweeper"]["task"] == ( "tasks.delete_conversation_task.sweep_deleted_conversations" ) diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index cf587aed3fc..54a2a0afe0c 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/api/tests/unit_tests/services/test_billing_service.py @@ -23,11 +23,12 @@ from werkzeug.exceptions import InternalServerError from enums import CloudPlan from models import Account, Tenant -from services.billing_service import BillingService, _BillingHTTPStatusError +from services.billing_service import BillingService, TokenerBootstrapUpstreamError, _BillingHTTPStatusError from services.errors.billing import ( BillingUpstreamInvalidResponseError, BillingUpstreamUnavailableError, ) +from tests.unit_tests.config_override import config_overrides_context TENANT_ID = "11111111-1111-1111-1111-111111111111" ACCOUNT_ID = "33333333-3333-3333-3333-333333333333" @@ -209,6 +210,75 @@ class TestBillingServiceSendRequest: with pytest.raises(json.JSONDecodeError): BillingService.ensure_new_agent_beta_revision("revision-1") + def test_tokener_bootstrap_uses_separate_base_url_and_sends_matching_tenant_id(self): + payload = { + "tenant_id": TENANT_ID, + "status": "ready", + "data_plane_api_key": "one-time-secret", + "retryable": False, + } + with ( + config_overrides_context(TOKENER_BILLING_API_URL="http://dify-saas-billing:8081/"), + patch.object(BillingService, "_send_request", return_value=payload) as send_request, + ): + result = BillingService.bootstrap_tokener_tenant(TENANT_ID, "Test Tenant") + + assert result == { + "tenant_id": TENANT_ID, + "status": "ready", + "data_plane_api_key": "one-time-secret", + "retryable": False, + } + assert "data_plane_api_key" not in payload + send_request.assert_called_once_with( + "POST", + f"/internal/v1/tokener/tenants/{TENANT_ID}/bootstrap", + json={"tenant_id": TENANT_ID, "display_name": "Test Tenant"}, + base_url="http://dify-saas-billing:8081", + ) + + def test_tokener_bootstrap_rejects_missing_key_without_exposing_payload(self): + payload = {"tenant_id": TENANT_ID, "status": "ready", "retryable": False, "other": "secret-value"} + with ( + config_overrides_context(TOKENER_BILLING_API_URL="http://dify-saas-billing:8081"), + patch.object(BillingService, "_send_request", return_value=payload), + pytest.raises(TokenerBootstrapUpstreamError) as exc_info, + ): + BillingService.bootstrap_tokener_tenant(TENANT_ID, "Test Tenant") + + assert exc_info.value.error_code == "tokener_bootstrap_key_missing" + assert exc_info.value.retryable is True + assert "secret-value" not in str(exc_info.value) + + def test_tokener_bootstrap_invalid_response_traceback_does_not_retain_key(self): + one_time_key = "key-that-must-not-enter-sentry" + payload = { + "tenant_id": "wrong-tenant", + "status": "ready", + "retryable": False, + "data_plane_api_key": one_time_key, + } + with ( + config_overrides_context(TOKENER_BILLING_API_URL="http://dify-saas-billing:8081"), + patch.object(BillingService, "_send_request", return_value=payload), + pytest.raises(TokenerBootstrapUpstreamError) as exc_info, + ): + BillingService.bootstrap_tokener_tenant(TENANT_ID, "Test Tenant") + + assert "data_plane_api_key" not in payload + traceback_cursor = exc_info.value.__traceback__ + checked_parser_frame = False + while traceback_cursor is not None: + frame = traceback_cursor.tb_frame + if ( + frame.f_code.co_filename.endswith("services/billing_service.py") + and frame.f_code.co_name == "bootstrap_tokener_tenant" + ): + checked_parser_frame = True + assert one_time_key not in repr(frame.f_locals) + traceback_cursor = traceback_cursor.tb_next + assert checked_parser_frame is True + @pytest.mark.parametrize( "status_code", [httpx.codes.BAD_REQUEST, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.NOT_FOUND] ) diff --git a/api/tests/unit_tests/services/test_tokener_tenant_creation.py b/api/tests/unit_tests/services/test_tokener_tenant_creation.py new file mode 100644 index 00000000000..82bfc13f407 --- /dev/null +++ b/api/tests/unit_tests/services/test_tokener_tenant_creation.py @@ -0,0 +1,54 @@ +from unittest.mock import MagicMock, patch + +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session + +from models.tokener import TenantTokenerIntegration, TenantTokenerIntegrationStatus +from services.account_service import TenantService +from tests.unit_tests.config_override import apply_config_overrides + + +def test_create_tenant_persists_tokener_integration_in_initial_commit( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + apply_config_overrides( + monkeypatch, + TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED=True, + TOKENER_PLUGIN_UNIQUE_IDENTIFIER="langgenius/tokener:0.1.2@checksum", + ) + + with ( + patch("services.account_service.SystemFeatureService.is_workspace_creation_allowed", return_value=True), + patch("services.account_service.generate_key_pair", return_value="public-key"), + patch("services.credit_pool_service.CreditPoolService.create_default_pool"), + ): + tenant = TenantService.create_tenant("Tokener tenant", session=sqlite_session) + + integration = sqlite_session.scalar( + select(TenantTokenerIntegration).where(TenantTokenerIntegration.tenant_id == tenant.id) + ) + assert integration is not None + assert integration.status == TenantTokenerIntegrationStatus.PENDING + assert integration.plugin_unique_identifier == "langgenius/tokener:0.1.2@checksum" + assert integration.attempt_count == 0 + + +def test_create_tenant_does_not_persist_tokener_integration_when_disabled( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + apply_config_overrides(monkeypatch, TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED=False) + + with ( + patch("services.account_service.SystemFeatureService.is_workspace_creation_allowed", return_value=True), + patch("services.account_service.generate_key_pair", return_value="public-key"), + patch("services.credit_pool_service.CreditPoolService.create_default_pool", MagicMock()), + ): + tenant = TenantService.create_tenant("Legacy tenant", session=sqlite_session) + + integration = sqlite_session.scalar( + select(TenantTokenerIntegration).where(TenantTokenerIntegration.tenant_id == tenant.id) + ) + assert integration is None diff --git a/api/tests/unit_tests/tasks/test_bootstrap_tokener_tenant_task.py b/api/tests/unit_tests/tasks/test_bootstrap_tokener_tenant_task.py new file mode 100644 index 00000000000..ff89e0e6a2e --- /dev/null +++ b/api/tests/unit_tests/tasks/test_bootstrap_tokener_tenant_task.py @@ -0,0 +1,323 @@ +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from celery.exceptions import Retry +from sqlalchemy import select +from sqlalchemy.orm import Session + +from libs.datetime_utils import naive_utc_now +from models.account import Tenant +from models.tokener import TenantTokenerIntegration, TenantTokenerIntegrationStatus +from tasks import bootstrap_tokener_tenant_task as task_module +from tests.unit_tests.config_override import apply_config_overrides + + +def _persist_integration( + session: Session, + *, + status: TenantTokenerIntegrationStatus = TenantTokenerIntegrationStatus.PENDING, + install_task_id: str | None = None, +) -> TenantTokenerIntegration: + tenant = Tenant(name="Tokener tenant") + session.add(tenant) + session.flush() + integration = TenantTokenerIntegration( + tenant_id=tenant.id, + status=status, + plugin_unique_identifier="langgenius/tokener:0.1.2@checksum", + plugin_install_task_id=install_task_id, + ) + session.add(integration) + session.commit() + return integration + + +def _snapshot(integration: TenantTokenerIntegration) -> task_module._IntegrationSnapshot: + return task_module._IntegrationSnapshot( + tenant_id=integration.tenant_id, + tenant_name="Tokener tenant", + status=integration.status, + plugin_unique_identifier=integration.plugin_unique_identifier, + plugin_install_task_id=integration.plugin_install_task_id, + ) + + +def test_task_uses_plugin_queue_and_late_acknowledgement() -> None: + task = task_module.bootstrap_tokener_tenant_task + + assert task.queue == task_module.TOKENER_BOOTSTRAP_QUEUE + assert task.acks_late is True + assert task.reject_on_worker_lost is True + + +def test_task_persists_sanitized_retry_state_and_releases_lock( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + integration = _persist_integration(sqlite_session) + apply_config_overrides(monkeypatch, TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED=True) + bootstrap_error = task_module._BootstrapStepError( + "tokener_plugin_install_pending", + retryable=True, + stage=TenantTokenerIntegrationStatus.INSTALLING_PLUGIN, + ) + monkeypatch.setattr(task_module, "_run_bootstrap", MagicMock(side_effect=bootstrap_error)) + lock = MagicMock() + lock.acquire.return_value = True + monkeypatch.setattr(task_module.redis_client, "lock", MagicMock(return_value=lock)) + retry = MagicMock(side_effect=Retry()) + monkeypatch.setattr(task_module.bootstrap_tokener_tenant_task, "retry", retry) + + with pytest.raises(Retry): + task_module.bootstrap_tokener_tenant_task.run(integration.tenant_id) + + retry.assert_called_once_with(exc=bootstrap_error, countdown=task_module._RETRY_DELAY_SECONDS) + lock.release.assert_called_once_with() + sqlite_session.expire_all() + persisted = sqlite_session.get(TenantTokenerIntegration, integration.id) + assert persisted is not None + assert persisted.status == TenantTokenerIntegrationStatus.INSTALLING_PLUGIN + assert persisted.last_error_code == "tokener_plugin_install_pending" + + +def test_lock_backend_failure_is_sanitized_and_retried(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED=True) + monkeypatch.setattr( + task_module.redis_client, + "lock", + MagicMock(side_effect=ConnectionError("redis connection contained sensitive diagnostics")), + ) + retry = MagicMock(side_effect=Retry()) + monkeypatch.setattr(task_module.bootstrap_tokener_tenant_task, "retry", retry) + + with pytest.raises(Retry): + task_module.bootstrap_tokener_tenant_task.run("tenant-1") + + retry.assert_called_once() + retry_error = retry.call_args.kwargs["exc"] + assert isinstance(retry_error, task_module._BootstrapStepError) + assert retry_error.error_code == "tokener_bootstrap_lock_unavailable" + assert "sensitive diagnostics" not in str(retry_error) + + +def test_run_bootstrap_marks_integration_ready( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + integration = _persist_integration(sqlite_session) + install = MagicMock() + ensure_credential = MagicMock(return_value="credential-1") + set_default = MagicMock() + monkeypatch.setattr(task_module, "_ensure_plugin_installed", install) + monkeypatch.setattr(task_module, "_ensure_managed_credential", ensure_credential) + monkeypatch.setattr(task_module, "_set_default_llm", set_default) + + task_module._run_bootstrap(integration.tenant_id) + + sqlite_session.expire_all() + persisted = sqlite_session.scalar( + select(TenantTokenerIntegration).where(TenantTokenerIntegration.tenant_id == integration.tenant_id) + ) + assert persisted is not None + assert persisted.status == TenantTokenerIntegrationStatus.READY + assert persisted.provider_credential_id == "credential-1" + assert persisted.ready_at is not None + assert persisted.attempt_count == 1 + install.assert_called_once() + ensure_credential.assert_called_once() + set_default.assert_called_once_with(integration.tenant_id) + + +def test_run_bootstrap_is_noop_after_ready(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + integration = _persist_integration(sqlite_session, status=TenantTokenerIntegrationStatus.READY) + install = MagicMock() + monkeypatch.setattr(task_module, "_ensure_plugin_installed", install) + + task_module._run_bootstrap(integration.tenant_id) + + install.assert_not_called() + sqlite_session.expire_all() + persisted = sqlite_session.get(TenantTokenerIntegration, integration.id) + assert persisted is not None + assert persisted.attempt_count == 0 + + +def test_package_install_persists_daemon_task_id_for_retry( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + integration = _persist_integration(sqlite_session) + apply_config_overrides(monkeypatch, TOKENER_PLUGIN_INSTALL_SOURCE="package") + monkeypatch.setattr(task_module.PluginService, "list", MagicMock(return_value=[])) + local_install = MagicMock(return_value=SimpleNamespace(all_installed=False, task_id="plugin-task-1")) + marketplace_install = MagicMock() + monkeypatch.setattr(task_module.PluginService, "install_from_local_pkg", local_install) + monkeypatch.setattr(task_module.PluginService, "install_from_marketplace_pkg", marketplace_install) + + with pytest.raises(task_module._BootstrapStepError) as exc_info: + task_module._ensure_plugin_installed(_snapshot(integration)) + + assert exc_info.value.error_code == "tokener_plugin_install_pending" + local_install.assert_called_once_with(integration.tenant_id, [integration.plugin_unique_identifier]) + marketplace_install.assert_not_called() + sqlite_session.expire_all() + persisted = sqlite_session.get(TenantTokenerIntegration, integration.id) + assert persisted is not None + assert persisted.plugin_install_task_id == "plugin-task-1" + + +def test_existing_managed_credential_skips_remote_bootstrap( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + integration = _persist_integration(sqlite_session) + find_credential = MagicMock(return_value=("credential-1", True)) + remote_bootstrap = MagicMock() + monkeypatch.setattr(task_module, "_find_managed_credential", find_credential) + monkeypatch.setattr(task_module.BillingService, "bootstrap_tokener_tenant", remote_bootstrap) + + credential_id = task_module._ensure_managed_credential(_snapshot(integration)) + + assert credential_id == "credential-1" + remote_bootstrap.assert_not_called() + + +def test_one_time_key_is_removed_from_response_after_credential_creation( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + integration = _persist_integration(sqlite_session) + response = { + "tenant_id": integration.tenant_id, + "status": "ready", + "data_plane_api_key": "one-time-secret", + "retryable": False, + } + monkeypatch.setattr(task_module, "_find_managed_credential", MagicMock(return_value=(None, False))) + monkeypatch.setattr( + task_module.BillingService, + "bootstrap_tokener_tenant", + MagicMock(return_value=response), + ) + consume_key = MagicMock(return_value=task_module._CredentialWriteResult(credential_id="credential-1")) + monkeypatch.setattr(task_module, "_consume_data_plane_api_key", consume_key) + + credential_id = task_module._ensure_managed_credential(_snapshot(integration)) + + assert credential_id == "credential-1" + consume_key.assert_called_once_with(integration.tenant_id, "one-time-secret") + assert "data_plane_api_key" not in response + + +def test_consume_data_plane_api_key_includes_configured_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + apply_config_overrides( + monkeypatch, + TOKENER_PROVIDER_NAME="langgenius/tokener/tokener", + TOKENER_ENDPOINT_URL="https://api-staging.tokener.dev/v1", + ) + provider_service = MagicMock() + captured_credentials: list[dict[str, str]] = [] + provider_service.create_provider_credential.side_effect = lambda **kwargs: captured_credentials.append( + dict(kwargs["credentials"]) + ) + monkeypatch.setattr(task_module, "ModelProviderService", MagicMock(return_value=provider_service)) + monkeypatch.setattr(task_module, "_find_managed_credential", MagicMock(return_value=("credential-1", True))) + + result = task_module._consume_data_plane_api_key("tenant-1", "one-time-secret") + + assert result == task_module._CredentialWriteResult(credential_id="credential-1") + provider_service.create_provider_credential.assert_called_once() + assert captured_credentials == [ + { + "api_key": "one-time-secret", + "endpoint_url": "https://api-staging.tokener.dev/v1", + } + ] + + +def test_credential_failure_traceback_frames_do_not_contain_one_time_key( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + integration = _persist_integration(sqlite_session) + one_time_key = "one-time-key-that-must-not-enter-sentry" + response = { + "tenant_id": integration.tenant_id, + "status": "ready", + "data_plane_api_key": one_time_key, + "retryable": False, + } + monkeypatch.setattr(task_module, "_find_managed_credential", MagicMock(return_value=(None, False))) + monkeypatch.setattr(task_module.BillingService, "bootstrap_tokener_tenant", MagicMock(return_value=response)) + provider_service = MagicMock() + provider_service.create_provider_credential.side_effect = RuntimeError(f"provider rejected {one_time_key}") + monkeypatch.setattr(task_module, "ModelProviderService", MagicMock(return_value=provider_service)) + + with pytest.raises(task_module._BootstrapStepError) as exc_info: + task_module._ensure_managed_credential(_snapshot(integration)) + + assert exc_info.value.error_code == "tokener_provider_credential_rejected" + assert one_time_key not in str(exc_info.value) + assert "data_plane_api_key" not in response + traceback_cursor = exc_info.value.__traceback__ + checked_production_frame = False + while traceback_cursor is not None: + frame = traceback_cursor.tb_frame + if frame.f_code.co_filename.endswith("tasks/bootstrap_tokener_tenant_task.py"): + checked_production_frame = True + assert one_time_key not in repr(frame.f_locals) + traceback_cursor = traceback_cursor.tb_next + assert checked_production_frame is True + + +def test_recovery_sweeper_requeues_only_stale_incomplete_integrations( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + stale_pending = _persist_integration(sqlite_session) + recent_pending = _persist_integration(sqlite_session) + ready = _persist_integration(sqlite_session, status=TenantTokenerIntegrationStatus.READY) + stale_at = naive_utc_now() - timedelta(minutes=10) + stale_pending.updated_at = stale_at + ready.updated_at = stale_at + sqlite_session.commit() + apply_config_overrides( + monkeypatch, + TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED=True, + TOKENER_BOOTSTRAP_RECOVERY_TASK_INTERVAL=5, + TOKENER_BOOTSTRAP_RECOVERY_BATCH_SIZE=100, + ) + delay = MagicMock() + monkeypatch.setattr(task_module.bootstrap_tokener_tenant_task, "delay", delay) + + dispatched = task_module.sweep_pending_tokener_integrations_task.run() + + assert dispatched == 1 + delay.assert_called_once_with(stale_pending.tenant_id) + assert recent_pending.tenant_id != stale_pending.tenant_id + + +def test_recovery_sweeper_retries_a_previous_broker_dispatch_failure( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + stale_pending = _persist_integration(sqlite_session) + stale_pending.updated_at = naive_utc_now() - timedelta(minutes=10) + sqlite_session.commit() + apply_config_overrides( + monkeypatch, + TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED=True, + TOKENER_BOOTSTRAP_RECOVERY_TASK_INTERVAL=5, + TOKENER_BOOTSTRAP_RECOVERY_BATCH_SIZE=100, + ) + delay = MagicMock(side_effect=ConnectionError("broker unavailable")) + monkeypatch.setattr(task_module.bootstrap_tokener_tenant_task, "delay", delay) + + assert task_module.sweep_pending_tokener_integrations_task.run() == 0 + delay.side_effect = None + assert task_module.sweep_pending_tokener_integrations_task.run() == 1 + assert delay.call_count == 2 diff --git a/docker/envs/core-services/shared.env.example b/docker/envs/core-services/shared.env.example index f5ae7604d0c..29ce87e3527 100644 --- a/docker/envs/core-services/shared.env.example +++ b/docker/envs/core-services/shared.env.example @@ -87,6 +87,18 @@ NEW_USER_DEFAULT_PLUGIN_IDS= # Comma-separated model_type:provider:model entries assigned after default plugins finish installing. # Example: llm:langgenius/openai/openai:gpt-4o-mini,text-embedding:langgenius/openai/openai:text-embedding-3-small NEW_USER_DEFAULT_MODELS= +# Managed Tokener bootstrap for newly created tenants. Keep disabled unless a pinned +# Tokener marketplace package and the billing-side bootstrap API are deployed. +TOKENER_NEW_TENANT_BOOTSTRAP_ENABLED=false +TOKENER_BILLING_API_URL= +TOKENER_PLUGIN_UNIQUE_IDENTIFIER= +TOKENER_PLUGIN_INSTALL_SOURCE=marketplace +TOKENER_PROVIDER_NAME=langgenius/tokener/tokener +TOKENER_DEFAULT_LLM_MODEL=deepseek-v4-flash +TOKENER_ENDPOINT_URL= +ENABLE_TOKENER_BOOTSTRAP_RECOVERY_TASK=true +TOKENER_BOOTSTRAP_RECOVERY_TASK_INTERVAL=5 +TOKENER_BOOTSTRAP_RECOVERY_BATCH_SIZE=100 ENDPOINT_URL_TEMPLATE=http://localhost/e/{hook_id} LOG_LEVEL=INFO LOG_OUTPUT_FORMAT=text