mirror of
https://github.com/langgenius/dify.git
synced 2026-09-05 00:31:19 +08:00
feat(api): bootstrap Tokener for new tenants
This commit is contained in:
Yansong Zhang
parent
fbe1c60e3f
commit
0c51942b66
@ -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,
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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
|
||||
@ -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"] = {
|
||||
|
||||
@ -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")
|
||||
@ -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",
|
||||
|
||||
84
api/models/tokener.py
Normal file
84
api/models/tokener.py
Normal file
@ -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,
|
||||
)
|
||||
@ -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:
|
||||
|
||||
@ -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}
|
||||
|
||||
509
api/tasks/bootstrap_tokener_tenant_task.py
Normal file
509
api/tasks/bootstrap_tokener_tenant_task.py
Normal file
@ -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
|
||||
@ -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()
|
||||
|
||||
@ -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
|
||||
@ -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"
|
||||
)
|
||||
|
||||
@ -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]
|
||||
)
|
||||
|
||||
@ -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
|
||||
323
api/tests/unit_tests/tasks/test_bootstrap_tokener_tenant_task.py
Normal file
323
api/tests/unit_tests/tasks/test_bootstrap_tokener_tenant_task.py
Normal file
@ -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
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user