youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-05-25 03:17:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

add plugin auto upgrade by category

Browse Source
This commit is contained in:
hjlarry
2026-05-18 15:09:48 +08:00
parent 83ab854d33
commit 00e9abdf79
14 changed files with 855 additions and 205 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
api/commands/__init__.py
View File

@ -4,6 +4,7 @@ CLI command modules extracted from `commands.py`.
from .account import create_tenant, reset_email, reset_password
from .plugin import (
backfill_plugin_auto_upgrade,
extract_plugins,
extract_unique_plugins,
install_plugins,
@ -38,6 +39,7 @@ from .vector import (
__all__ = [
"add_qdrant_index",
"archive_workflow_runs",
"backfill_plugin_auto_upgrade",
"clean_expired_messages",
"clean_workflow_runs",
"cleanup_orphaned_draft_variables",

109
api/commands/plugin.py
View File

@ -1,10 +1,11 @@
import json
import logging
import time
from typing import Any, cast
import click
from pydantic import TypeAdapter
from sqlalchemy import delete, select
from sqlalchemy import delete, func, select
from sqlalchemy.engine import CursorResult
from configs import dify_config
@ -14,11 +15,13 @@ from core.plugin.impl.plugin import PluginInstaller
from core.tools.utils.system_encryption import encrypt_system_params
from extensions.ext_database import db
from models import Tenant
from models.account import TenantPluginAutoUpgradeStrategy
from models.oauth import DatasourceOauthParamConfig, DatasourceProvider
from models.provider_ids import DatasourceProviderID, ToolProviderID
from models.source import DataSourceApiKeyAuthBinding, DataSourceOauthBinding
from models.tools import ToolOAuthSystemClient
from services.plugin.data_migration import PluginDataMigration
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
from services.plugin.plugin_migration import PluginMigration
from services.plugin.plugin_service import PluginService
@ -402,6 +405,110 @@ def migrate_data_for_plugin():
click.echo(click.style("Migrate data for plugin completed.", fg="green"))
def _candidate_auto_upgrade_strategy_tenant_ids_stmt(limit: int | None = None):
category_count = len(TenantPluginAutoUpgradeStrategy.PluginCategory)
stmt = (
select(TenantPluginAutoUpgradeStrategy.tenant_id)
.group_by(TenantPluginAutoUpgradeStrategy.tenant_id)
.having(func.count(func.distinct(TenantPluginAutoUpgradeStrategy.category)) < category_count)
.order_by(TenantPluginAutoUpgradeStrategy.tenant_id)
)
if limit is not None:
stmt = stmt.limit(limit)
return stmt
def _count_auto_upgrade_strategy_tenant_ids(limit: int | None) -> int:
candidate_stmt = _candidate_auto_upgrade_strategy_tenant_ids_stmt(limit).subquery()
return db.session.scalar(select(func.count()).select_from(candidate_stmt)) or 0
def _iter_auto_upgrade_strategy_tenant_ids(limit: int | None):
stmt = _candidate_auto_upgrade_strategy_tenant_ids_stmt(limit).execution_options(yield_per=1000)
yield from db.session.scalars(stmt)
@click.command(
"backfill-plugin-auto-upgrade",
help="Backfill category-scoped plugin auto-upgrade strategies and normalize plugin lists.",
)
@click.option("--tenant-id", multiple=True, help="Tenant ID to backfill. Can be passed multiple times.")
@click.option("--limit", type=int, default=None, help="Maximum number of candidate tenants to process.")
@click.option("--batch-size", type=int, default=500, show_default=True, help="Progress reporting batch size.")
@click.option("--dry-run", is_flag=True, help="Only print candidate tenant count.")
def backfill_plugin_auto_upgrade(
tenant_id: tuple[str, ...],
limit: int | None,
batch_size: int,
dry_run: bool,
):
"""
Backfill historical auto-upgrade strategies after the category column exists.
Missing category rows are created from the tenant's tool/default row. Pure default
strategies become latest for model plugins and fix-only for all other categories.
Tenants with include/exclude plugin IDs are split
by installed plugin category using plugin daemon metadata.
"""
start_at = time.perf_counter()
candidate_count = len(tenant_id) if tenant_id else _count_auto_upgrade_strategy_tenant_ids(limit)
click.echo(click.style(f"Found {candidate_count} candidate tenants.", fg="yellow"))
if dry_run:
elapsed = time.perf_counter() - start_at
click.echo(click.style(f"Dry run completed. elapsed={elapsed:.2f}s", fg="green"))
return
tenant_ids = list(tenant_id) if tenant_id else _iter_auto_upgrade_strategy_tenant_ids(limit)
backfilled_count = 0
created_count = 0
normalized_count = 0
skipped_count = 0
failed_count = 0
for index, current_tenant_id in enumerate(tenant_ids, start=1):
try:
result = PluginAutoUpgradeService.backfill_strategy_categories(
current_tenant_id,
)
except Exception as e:
failed_count += 1
click.echo(click.style(f"Failed tenant {current_tenant_id}: {str(e)}", fg="red"))
continue
if result.created_count > 0:
backfilled_count += 1
created_count += result.created_count
elif not result.normalized:
skipped_count += 1
if result.normalized:
normalized_count += 1
if batch_size > 0 and index % batch_size == 0:
click.echo(
click.style(
f"Processed {index}/{candidate_count} tenants. "
f"backfilled={backfilled_count}, created_rows={created_count}, "
f"normalized={normalized_count}, skipped={skipped_count}, failed={failed_count}, "
f"elapsed={time.perf_counter() - start_at:.2f}s",
fg="yellow",
)
)
elapsed = time.perf_counter() - start_at
click.echo(
click.style(
f"Backfill plugin auto-upgrade strategy categories completed. "
f"backfilled={backfilled_count}, created_rows={created_count}, "
f"normalized={normalized_count}, skipped={skipped_count}, failed={failed_count}, "
f"elapsed={elapsed:.2f}s",
fg="green",
)
)
@click.command("extract-plugins", help="Extract plugins.")
@click.option("--output_file", prompt=True, help="The file to store the extracted plugins.", default="plugins.jsonl")
@click.option("--workers", prompt=True, help="The number of workers to extract plugins.", default=10)

161
api/controllers/console/workspace/plugin.py
View File

@ -1,15 +1,15 @@
import io
from collections.abc import Mapping
from typing import Any, Literal
from typing import Any, Literal, TypedDict
from flask import request, send_file
from flask_restx import Resource
from pydantic import BaseModel, Field
from pydantic import BaseModel, ConfigDict, Field
from werkzeug.datastructures import FileStorage
from werkzeug.exceptions import Forbidden
from configs import dify_config
from controllers.common.schema import register_enum_models, register_schema_models
from controllers.common.schema import query_params_from_model, register_enum_models, register_schema_models
from controllers.console import console_ns
from controllers.console.workspace import plugin_permission_required
from controllers.console.wraps import account_initialization_required, is_admin_or_owner_required, setup_required
@ -23,6 +23,14 @@ from services.plugin.plugin_permission_service import PluginPermissionService
from services.plugin.plugin_service import PluginService
class AutoUpgradeSettingsResponse(TypedDict):
strategy_setting: TenantPluginAutoUpgradeStrategy.StrategySetting
upgrade_time_of_day: int
upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode
exclude_plugins: list[str]
include_plugins: list[str]
class ParserList(BaseModel):
page: int = Field(default=1, ge=1, description="Page number")
page_size: int = Field(default=256, ge=1, le=256, description="Page size (1-256)")
@ -86,8 +94,8 @@ class ParserUninstall(BaseModel):
class ParserPermissionChange(BaseModel):
install_permission: TenantPluginPermission.InstallPermission
debug_permission: TenantPluginPermission.DebugPermission
install_permission: TenantPluginPermission.InstallPermission = TenantPluginPermission.InstallPermission.EVERYONE
debug_permission: TenantPluginPermission.DebugPermission = TenantPluginPermission.DebugPermission.EVERYONE
class ParserDynamicOptions(BaseModel):
@ -123,13 +131,22 @@ class PluginAutoUpgradeSettingsPayload(BaseModel):
include_plugins: list[str] = Field(default_factory=list)
class ParserPreferencesChange(BaseModel):
permission: PluginPermissionSettingsPayload
class ParserAutoUpgradeChange(BaseModel):
model_config = ConfigDict(extra="forbid")
category: TenantPluginAutoUpgradeStrategy.PluginCategory
auto_upgrade: PluginAutoUpgradeSettingsPayload
class ParserAutoUpgradeFetch(BaseModel):
category: TenantPluginAutoUpgradeStrategy.PluginCategory
class ParserExcludePlugin(BaseModel):
model_config = ConfigDict(extra="forbid")
plugin_id: str
category: TenantPluginAutoUpgradeStrategy.PluginCategory
class ParserReadme(BaseModel):
@ -156,7 +173,8 @@ register_schema_models(
ParserPermissionChange,
ParserDynamicOptions,
ParserDynamicOptionsWithCredentials,
ParserPreferencesChange,
ParserAutoUpgradeChange,
ParserAutoUpgradeFetch,
ParserExcludePlugin,
ParserReadme,
)
@ -164,12 +182,36 @@ register_schema_models(
register_enum_models(
console_ns,
TenantPluginPermission.DebugPermission,
TenantPluginAutoUpgradeStrategy.PluginCategory,
TenantPluginAutoUpgradeStrategy.UpgradeMode,
TenantPluginAutoUpgradeStrategy.StrategySetting,
TenantPluginPermission.InstallPermission,
)
def _default_auto_upgrade_settings(
tenant_id: str,
category: TenantPluginAutoUpgradeStrategy.PluginCategory,
) -> AutoUpgradeSettingsResponse:
return {
"strategy_setting": PluginAutoUpgradeService.default_strategy_setting_for_category(category),
"upgrade_time_of_day": PluginAutoUpgradeService.default_upgrade_time_of_day(tenant_id),
"upgrade_mode": TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
"exclude_plugins": [],
"include_plugins": [],
}
def _auto_upgrade_settings_to_dict(strategy: TenantPluginAutoUpgradeStrategy) -> AutoUpgradeSettingsResponse:
return {
"strategy_setting": strategy.strategy_setting,
"upgrade_time_of_day": strategy.upgrade_time_of_day,
"upgrade_mode": strategy.upgrade_mode,
"exclude_plugins": strategy.exclude_plugins,
"include_plugins": strategy.include_plugins,
}
def _read_upload_content(file: FileStorage, max_size: int) -> bytes:
"""
Read the uploaded file and validate its actual size before delegating to the plugin service.
@ -617,11 +659,13 @@ class PluginChangePermissionApi(Resource):
tenant_id = current_tenant_id
return {
"success": PluginPermissionService.change_permission(
tenant_id, args.install_permission, args.debug_permission
)
}
set_permission_result = PluginPermissionService.change_permission(
tenant_id, args.install_permission, args.debug_permission
)
if not set_permission_result:
return jsonable_encoder({"success": False, "message": "Failed to set permission"})
return jsonable_encoder({"success": True})
@console_ns.route("/workspaces/current/plugin/permission/fetch")
@ -710,9 +754,9 @@ class PluginFetchDynamicSelectOptionsWithCredentialsApi(Resource):
return jsonable_encoder({"options": options})
@console_ns.route("/workspaces/current/plugin/preferences/change")
class PluginChangePreferencesApi(Resource):
@console_ns.expect(console_ns.models[ParserPreferencesChange.__name__])
@console_ns.route("/workspaces/current/plugin/auto-upgrade/change")
class PluginChangeAutoUpgradeApi(Resource):
@console_ns.expect(console_ns.models[ParserAutoUpgradeChange.__name__])
@setup_required
@login_required
@account_initialization_required
@ -721,38 +765,17 @@ class PluginChangePreferencesApi(Resource):
if not user.is_admin_or_owner:
raise Forbidden()
args = ParserPreferencesChange.model_validate(console_ns.payload)
permission = args.permission
install_permission = permission.install_permission
debug_permission = permission.debug_permission
args = ParserAutoUpgradeChange.model_validate(console_ns.payload)
auto_upgrade = args.auto_upgrade
strategy_setting = auto_upgrade.strategy_setting
upgrade_time_of_day = auto_upgrade.upgrade_time_of_day
upgrade_mode = auto_upgrade.upgrade_mode
exclude_plugins = auto_upgrade.exclude_plugins
include_plugins = auto_upgrade.include_plugins
# set permission
set_permission_result = PluginPermissionService.change_permission(
tenant_id,
install_permission,
debug_permission,
)
if not set_permission_result:
return jsonable_encoder({"success": False, "message": "Failed to set permission"})
# set auto upgrade strategy
set_auto_upgrade_strategy_result = PluginAutoUpgradeService.change_strategy(
tenant_id,
strategy_setting,
upgrade_time_of_day,
upgrade_mode,
exclude_plugins,
include_plugins,
auto_upgrade.strategy_setting,
auto_upgrade.upgrade_time_of_day,
auto_upgrade.upgrade_mode,
auto_upgrade.exclude_plugins,
auto_upgrade.include_plugins,
category=args.category,
)
if not set_auto_upgrade_strategy_result:
return jsonable_encoder({"success": False, "message": "Failed to set auto upgrade strategy"})
@ -760,46 +783,32 @@ class PluginChangePreferencesApi(Resource):
return jsonable_encoder({"success": True})
@console_ns.route("/workspaces/current/plugin/preferences/fetch")
class PluginFetchPreferencesApi(Resource):
@console_ns.route("/workspaces/current/plugin/auto-upgrade/fetch")
class PluginFetchAutoUpgradeApi(Resource):
@console_ns.doc(params=query_params_from_model(ParserAutoUpgradeFetch))
@setup_required
@login_required
@account_initialization_required
def get(self):
_, tenant_id = current_account_with_tenant()
permission = PluginPermissionService.get_permission(tenant_id)
permission_dict = {
"install_permission": TenantPluginPermission.InstallPermission.EVERYONE,
"debug_permission": TenantPluginPermission.DebugPermission.EVERYONE,
}
args = ParserAutoUpgradeFetch.model_validate(request.args.to_dict(flat=True))
auto_upgrade = PluginAutoUpgradeService.get_strategy(tenant_id, args.category)
auto_upgrade_dict = (
_auto_upgrade_settings_to_dict(auto_upgrade)
if auto_upgrade
else _default_auto_upgrade_settings(tenant_id, args.category)
)
if permission:
permission_dict["install_permission"] = permission.install_permission
permission_dict["debug_permission"] = permission.debug_permission
auto_upgrade = PluginAutoUpgradeService.get_strategy(tenant_id)
auto_upgrade_dict = {
"strategy_setting": TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED,
"upgrade_time_of_day": 0,
"upgrade_mode": TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
"exclude_plugins": [],
"include_plugins": [],
}
if auto_upgrade:
auto_upgrade_dict = {
"strategy_setting": auto_upgrade.strategy_setting,
"upgrade_time_of_day": auto_upgrade.upgrade_time_of_day,
"upgrade_mode": auto_upgrade.upgrade_mode,
"exclude_plugins": auto_upgrade.exclude_plugins,
"include_plugins": auto_upgrade.include_plugins,
return jsonable_encoder(
{
"category": args.category,
"auto_upgrade": auto_upgrade_dict,
}
return jsonable_encoder({"permission": permission_dict, "auto_upgrade": auto_upgrade_dict})
)
@console_ns.route("/workspaces/current/plugin/preferences/autoupgrade/exclude")
@console_ns.route("/workspaces/current/plugin/auto-upgrade/exclude")
class PluginAutoUpgradeExcludePluginApi(Resource):
@console_ns.expect(console_ns.models[ParserExcludePlugin.__name__])
@setup_required
@ -811,7 +820,9 @@ class PluginAutoUpgradeExcludePluginApi(Resource):
args = ParserExcludePlugin.model_validate(console_ns.payload)
return jsonable_encoder({"success": PluginAutoUpgradeService.exclude_plugin(tenant_id, args.plugin_id)})
return jsonable_encoder(
{"success": PluginAutoUpgradeService.exclude_plugin(tenant_id, args.plugin_id, args.category)}
)
@console_ns.route("/workspaces/current/plugin/readme")

2
api/extensions/ext_commands.py
View File

@ -5,6 +5,7 @@ def init_app(app: DifyApp):
from commands import (
add_qdrant_index,
archive_workflow_runs,
backfill_plugin_auto_upgrade,
clean_expired_messages,
clean_workflow_runs,
cleanup_orphaned_draft_variables,
@ -49,6 +50,7 @@ def init_app(app: DifyApp):
fix_app_site_missing,
migrate_data_for_plugin,
migrate_member_roles_to_rbac,
backfill_plugin_auto_upgrade,
extract_plugins,
extract_unique_plugins,
install_plugins,

42
api/migrations/versions/2026_05_15_1200-f6a7b8c9d012_add_plugin_auto_upgrade_category.py Normal file
View File

@ -0,0 +1,42 @@
"""add plugin auto upgrade category
Revision ID: f6a7b8c9d012
Revises: a4f2d8c9b731
Create Date: 2026-05-15 12:00:00.000000
"""
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision = "f6a7b8c9d012"
down_revision = "a4f2d8c9b731"
branch_labels = None
depends_on = None
LEGACY_CATEGORY = "tool"
UNIQUE_CONSTRAINT_NAME = "unique_tenant_plugin_auto_upgrade_strategy"
UPGRADE_TIME_INDEX_NAME = "idx_tenant_plugin_auto_upgrade_strategy_time"
STRATEGY_TABLE_NAME = "tenant_plugin_auto_upgrade_strategies"
def upgrade():
with op.batch_alter_table(STRATEGY_TABLE_NAME, schema=None) as batch_op:
batch_op.add_column(
sa.Column("category", sa.String(length=32), server_default=LEGACY_CATEGORY, nullable=False)
)
batch_op.drop_constraint(UNIQUE_CONSTRAINT_NAME, type_="unique")
batch_op.create_unique_constraint(UNIQUE_CONSTRAINT_NAME, ["tenant_id", "category"])
batch_op.create_index(UPGRADE_TIME_INDEX_NAME, ["upgrade_time_of_day"])
def downgrade():
op.execute(sa.text(f"DELETE FROM {STRATEGY_TABLE_NAME} WHERE category != '{LEGACY_CATEGORY}'"))
with op.batch_alter_table(STRATEGY_TABLE_NAME, schema=None) as batch_op:
batch_op.drop_index(UPGRADE_TIME_INDEX_NAME)
batch_op.drop_constraint(UNIQUE_CONSTRAINT_NAME, type_="unique")
batch_op.drop_column("category")
batch_op.create_unique_constraint(UNIQUE_CONSTRAINT_NAME, ["tenant_id"])

17
api/models/account.py
View File

@ -401,6 +401,14 @@ class TenantPluginPermission(TypeBase):
class TenantPluginAutoUpgradeStrategy(TypeBase):
class PluginCategory(enum.StrEnum):
TOOL = "tool"
MODEL = "model"
EXTENSION = "extension"
AGENT_STRATEGY = "agent-strategy"
DATASOURCE = "datasource"
TRIGGER = "trigger"
class StrategySetting(enum.StrEnum):
DISABLED = "disabled"
FIX_ONLY = "fix_only"
@ -414,13 +422,20 @@ class TenantPluginAutoUpgradeStrategy(TypeBase):
__tablename__ = "tenant_plugin_auto_upgrade_strategies"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="tenant_plugin_auto_upgrade_strategy_pkey"),
sa.UniqueConstraint("tenant_id", name="unique_tenant_plugin_auto_upgrade_strategy"),
sa.UniqueConstraint("tenant_id", "category", name="unique_tenant_plugin_auto_upgrade_strategy"),
sa.Index("idx_tenant_plugin_auto_upgrade_strategy_time", "upgrade_time_of_day"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
category: Mapped[PluginCategory] = mapped_column(
EnumText(PluginCategory, length=32),
nullable=False,
server_default="tool",
default=PluginCategory.TOOL,
)
strategy_setting: Mapped[StrategySetting] = mapped_column(
EnumText(StrategySetting, length=16),
nullable=False,

1
api/schedule/check_upgradable_plugin_task.py
View File

@ -73,6 +73,7 @@ def check_upgradable_plugin_task():
strategy.upgrade_mode,
strategy.exclude_plugins,
strategy.include_plugins,
strategy.category,
)
# Only sleep if batch_interval_time > 0.0001 AND current batch is not the last one

21
api/services/account_service.py
View File

@ -64,6 +64,7 @@ from services.errors.account import (
from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
from services.enterprise.rbac_service import ListOption, RBACService
from services.feature_service import FeatureService
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
from tasks.delete_account_task import delete_account_task
from tasks.mail_account_deletion_task import send_account_deletion_verification_code
from tasks.mail_change_mail_task import (
@ -1111,15 +1112,17 @@ class TenantService:
db.session.add(tenant)
db.session.commit()
plugin_upgrade_strategy = TenantPluginAutoUpgradeStrategy(
tenant_id=tenant.id,
strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
upgrade_time_of_day=0,
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
exclude_plugins=[],
include_plugins=[],
)
db.session.add(plugin_upgrade_strategy)
for category in TenantPluginAutoUpgradeStrategy.PluginCategory:
plugin_upgrade_strategy = TenantPluginAutoUpgradeStrategy(
tenant_id=tenant.id,
category=category,
strategy_setting=PluginAutoUpgradeService.default_strategy_setting_for_category(category),
upgrade_time_of_day=PluginAutoUpgradeService.default_upgrade_time_of_day(tenant.id),
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
exclude_plugins=[],
include_plugins=[],
)
db.session.add(plugin_upgrade_strategy)
db.session.commit()
tenant.encrypt_public_key = generate_key_pair(tenant.id)

402
api/services/plugin/plugin_auto_upgrade_service.py
View File

@ -1,18 +1,296 @@
"""Manage tenant plugin auto-upgrade strategies.
The storage is category-scoped: each tenant can have one strategy per plugin
category. Public mutation helpers require an explicit category so callers do
not accidentally overwrite every plugin type with one workspace-level policy.
"""
import logging
from dataclasses import dataclass
from hashlib import sha256
from sqlalchemy import select
from sqlalchemy.orm import Session
from core.db.session_factory import session_factory
from core.plugin.impl.plugin import PluginInstaller
from models.account import TenantPluginAutoUpgradeStrategy
logger = logging.getLogger(__name__)
PluginCategory = TenantPluginAutoUpgradeStrategy.PluginCategory
PLUGIN_CATEGORIES = tuple(PluginCategory)
SECONDS_PER_DAY = 24 * 60 * 60
AUTO_UPGRADE_CHECK_SLOT_SECONDS = 15 * 60
AUTO_UPGRADE_CHECK_SLOT_COUNT = SECONDS_PER_DAY // AUTO_UPGRADE_CHECK_SLOT_SECONDS
AUTO_UPGRADE_CHECK_SLOT_OFFSET_SECONDS = AUTO_UPGRADE_CHECK_SLOT_SECONDS // 2
@dataclass(frozen=True)
class PluginAutoUpgradeBackfillResult:
created_count: int
normalized: bool
class PluginAutoUpgradeService:
@staticmethod
def get_strategy(tenant_id: str) -> TenantPluginAutoUpgradeStrategy | None:
with session_factory.create_session() as session:
return session.scalar(
select(TenantPluginAutoUpgradeStrategy)
.where(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id)
.limit(1)
def default_strategy_setting_for_category(
category: PluginCategory,
) -> TenantPluginAutoUpgradeStrategy.StrategySetting:
if category == PluginCategory.MODEL:
return TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST
return TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY
@staticmethod
def default_upgrade_time_of_day(tenant_id: str) -> int:
"""Spread default checks across 15-minute slots by tenant."""
hash_input = tenant_id.encode()
slot = int.from_bytes(sha256(hash_input).digest()[:8], "big") % AUTO_UPGRADE_CHECK_SLOT_COUNT
return slot * AUTO_UPGRADE_CHECK_SLOT_SECONDS + AUTO_UPGRADE_CHECK_SLOT_OFFSET_SECONDS
@staticmethod
def _coerce_category(category: object) -> PluginCategory | None:
"""Accept daemon enum/string categories and ignore unknown values."""
category_value = getattr(category, "value", category)
if category_value is None:
return None
try:
return PluginCategory(str(category_value))
except ValueError:
return None
@staticmethod
def _get_installed_plugin_categories(tenant_id: str) -> dict[str, PluginCategory]:
"""Build a plugin_id -> category map for splitting legacy include/exclude lists."""
installed_plugins = PluginInstaller().list_plugins(tenant_id)
plugin_categories: dict[str, PluginCategory] = {}
for plugin in installed_plugins:
plugin_category = PluginAutoUpgradeService._coerce_category(plugin.declaration.category)
if plugin_category is not None:
plugin_categories[plugin.plugin_id] = plugin_category
return plugin_categories
@staticmethod
def _filter_plugin_ids_for_category(
plugin_ids: list[str],
category: PluginCategory,
plugin_categories: dict[str, PluginCategory],
) -> list[str]:
return [plugin_id for plugin_id in plugin_ids if plugin_categories.get(plugin_id) == category]
@staticmethod
def _log_unknown_plugin_ids(
tenant_id: str,
field_name: str,
plugin_ids: list[str],
plugin_categories: dict[str, PluginCategory],
) -> None:
unknown_plugin_ids = [plugin_id for plugin_id in plugin_ids if plugin_id not in plugin_categories]
if not unknown_plugin_ids:
return
logger.warning(
"Skipped unknown plugin IDs while backfilling plugin auto-upgrade strategies: "
"tenant_id=%s, field=%s, plugin_ids=%s",
tenant_id,
field_name,
unknown_plugin_ids,
)
@staticmethod
def _has_default_strategy(strategy: TenantPluginAutoUpgradeStrategy) -> bool:
return (
strategy.strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY
and strategy.upgrade_time_of_day == 0
and strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE
and not strategy.exclude_plugins
and not strategy.include_plugins
)
@staticmethod
def _strategy_setting_for_category(
source_strategy: TenantPluginAutoUpgradeStrategy,
category: PluginCategory,
source_has_default_strategy: bool,
) -> TenantPluginAutoUpgradeStrategy.StrategySetting:
# Only pure legacy defaults adopt the new model=latest default. User-edited
# strategies keep their original setting across all categories.
if source_has_default_strategy:
return PluginAutoUpgradeService.default_strategy_setting_for_category(category)
return source_strategy.strategy_setting
@staticmethod
def _upgrade_time_of_day_for_category(
tenant_id: str,
source_strategy: TenantPluginAutoUpgradeStrategy,
source_has_default_strategy: bool,
) -> int:
# Pure legacy defaults are spread by tenant so all default rows do not
# concentrate in the same scheduler window. User-edited schedules keep their time.
if source_has_default_strategy:
return PluginAutoUpgradeService.default_upgrade_time_of_day(tenant_id)
return source_strategy.upgrade_time_of_day
@staticmethod
def backfill_strategy_categories(
tenant_id: str,
) -> PluginAutoUpgradeBackfillResult:
"""Create missing category strategies and split include/exclude lists when needed.
The historical row is treated as the workspace-level source strategy.
New category rows copy it first, then plugin lists are narrowed by real
plugin category when the source strategy contains include/exclude IDs.
"""
with session_factory.create_session() as session, session.begin():
strategies = list(
session.scalars(
select(TenantPluginAutoUpgradeStrategy).where(
TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id
)
).all()
)
if not strategies:
return PluginAutoUpgradeBackfillResult(created_count=0, normalized=False)
# Schema migration marks the historical workspace-level row as tool.
source_strategy = next(
(strategy for strategy in strategies if strategy.category == PluginCategory.TOOL),
strategies[0],
)
source_has_default_strategy = PluginAutoUpgradeService._has_default_strategy(source_strategy)
strategies_by_category = {strategy.category: strategy for strategy in strategies}
exclude_plugins = source_strategy.exclude_plugins
include_plugins = source_strategy.include_plugins
should_split_plugin_lists = bool(exclude_plugins or include_plugins)
# Query daemon only for tenants that actually customized plugin lists.
plugin_categories = (
PluginAutoUpgradeService._get_installed_plugin_categories(tenant_id)
if should_split_plugin_lists
else {}
)
if should_split_plugin_lists:
PluginAutoUpgradeService._log_unknown_plugin_ids(
tenant_id,
"exclude_plugins",
exclude_plugins,
plugin_categories,
)
PluginAutoUpgradeService._log_unknown_plugin_ids(
tenant_id,
"include_plugins",
include_plugins,
plugin_categories,
)
created_count = 0
for category in PLUGIN_CATEGORIES:
strategy = strategies_by_category.get(category)
if strategy is None:
# Start from the legacy workspace-level behavior before narrowing lists.
strategy = TenantPluginAutoUpgradeStrategy(
tenant_id=tenant_id,
category=category,
strategy_setting=PluginAutoUpgradeService._strategy_setting_for_category(
source_strategy, category, source_has_default_strategy
),
upgrade_time_of_day=PluginAutoUpgradeService._upgrade_time_of_day_for_category(
tenant_id, source_strategy, source_has_default_strategy
),
upgrade_mode=source_strategy.upgrade_mode,
exclude_plugins=source_strategy.exclude_plugins.copy(),
include_plugins=source_strategy.include_plugins.copy(),
)
session.add(strategy)
created_count += 1
elif source_has_default_strategy:
strategy.strategy_setting = PluginAutoUpgradeService.default_strategy_setting_for_category(
strategy.category
)
strategy.upgrade_time_of_day = PluginAutoUpgradeService.default_upgrade_time_of_day(tenant_id)
if not should_split_plugin_lists:
continue
# Narrow include/exclude lists to the current category after all rows exist.
strategy.exclude_plugins = PluginAutoUpgradeService._filter_plugin_ids_for_category(
exclude_plugins,
strategy.category,
plugin_categories,
)
strategy.include_plugins = PluginAutoUpgradeService._filter_plugin_ids_for_category(
include_plugins,
strategy.category,
plugin_categories,
)
return PluginAutoUpgradeBackfillResult(created_count=created_count, normalized=should_split_plugin_lists)
@staticmethod
def _get_strategy(
session: Session,
tenant_id: str,
category: PluginCategory,
) -> TenantPluginAutoUpgradeStrategy | None:
return session.scalar(
select(TenantPluginAutoUpgradeStrategy)
.where(
TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id,
TenantPluginAutoUpgradeStrategy.category == category,
)
.limit(1)
)
@staticmethod
def get_strategy(
tenant_id: str,
category: PluginCategory,
) -> TenantPluginAutoUpgradeStrategy | None:
with session_factory.create_session() as session:
return PluginAutoUpgradeService._get_strategy(session, tenant_id, category)
@staticmethod
def get_strategies(tenant_id: str) -> list[TenantPluginAutoUpgradeStrategy]:
with session_factory.create_session() as session:
return list(
session.scalars(
select(TenantPluginAutoUpgradeStrategy).where(
TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id
)
).all()
)
@staticmethod
def _change_strategy(
session: Session,
tenant_id: str,
category: PluginCategory,
strategy_setting: TenantPluginAutoUpgradeStrategy.StrategySetting,
upgrade_time_of_day: int,
upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode,
exclude_plugins: list[str],
include_plugins: list[str],
) -> None:
exist_strategy = PluginAutoUpgradeService._get_strategy(session, tenant_id, category)
if not exist_strategy:
strategy = TenantPluginAutoUpgradeStrategy(
tenant_id=tenant_id,
category=category,
strategy_setting=strategy_setting,
upgrade_time_of_day=upgrade_time_of_day,
upgrade_mode=upgrade_mode,
exclude_plugins=exclude_plugins,
include_plugins=include_plugins,
)
session.add(strategy)
else:
exist_strategy.strategy_setting = strategy_setting
exist_strategy.upgrade_time_of_day = upgrade_time_of_day
exist_strategy.upgrade_mode = upgrade_mode
exist_strategy.exclude_plugins = exclude_plugins
exist_strategy.include_plugins = include_plugins
@staticmethod
def change_strategy(
@ -22,64 +300,72 @@ class PluginAutoUpgradeService:
upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode,
exclude_plugins: list[str],
include_plugins: list[str],
category: PluginCategory,
) -> bool:
with session_factory.create_session() as session, session.begin():
exist_strategy = session.scalar(
select(TenantPluginAutoUpgradeStrategy)
.where(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id)
.limit(1)
PluginAutoUpgradeService._change_strategy(
session,
tenant_id=tenant_id,
category=category,
strategy_setting=strategy_setting,
upgrade_time_of_day=upgrade_time_of_day,
upgrade_mode=upgrade_mode,
exclude_plugins=exclude_plugins,
include_plugins=include_plugins,
)
if not exist_strategy:
strategy = TenantPluginAutoUpgradeStrategy(
tenant_id=tenant_id,
strategy_setting=strategy_setting,
upgrade_time_of_day=upgrade_time_of_day,
upgrade_mode=upgrade_mode,
exclude_plugins=exclude_plugins,
include_plugins=include_plugins,
)
session.add(strategy)
else:
exist_strategy.strategy_setting = strategy_setting
exist_strategy.upgrade_time_of_day = upgrade_time_of_day
exist_strategy.upgrade_mode = upgrade_mode
exist_strategy.exclude_plugins = exclude_plugins
exist_strategy.include_plugins = include_plugins
return True
@staticmethod
def exclude_plugin(tenant_id: str, plugin_id: str) -> bool:
with session_factory.create_session() as session, session.begin():
exist_strategy = session.scalar(
select(TenantPluginAutoUpgradeStrategy)
.where(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id)
.limit(1)
def _exclude_plugin(
session: Session,
tenant_id: str,
category: PluginCategory,
plugin_id: str,
) -> None:
"""Remove one plugin from automatic updates for a single category strategy."""
exist_strategy = PluginAutoUpgradeService._get_strategy(session, tenant_id, category)
if not exist_strategy:
PluginAutoUpgradeService._change_strategy(
session,
tenant_id,
category,
TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
0,
TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
[plugin_id],
[],
)
if not exist_strategy:
# create for this tenant
PluginAutoUpgradeService.change_strategy(
tenant_id,
TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
0,
TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
[plugin_id],
[],
)
return True
else:
if exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE:
if plugin_id not in exist_strategy.exclude_plugins:
new_exclude_plugins = exist_strategy.exclude_plugins.copy()
new_exclude_plugins.append(plugin_id)
exist_strategy.exclude_plugins = new_exclude_plugins
elif exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL:
if plugin_id in exist_strategy.include_plugins:
new_include_plugins = exist_strategy.include_plugins.copy()
new_include_plugins.remove(plugin_id)
exist_strategy.include_plugins = new_include_plugins
elif exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL:
exist_strategy.upgrade_mode = TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE
exist_strategy.exclude_plugins = [plugin_id]
else:
if exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE:
# In exclude mode, disabling one plugin means adding it to exclude_plugins.
if plugin_id not in exist_strategy.exclude_plugins:
new_exclude_plugins = exist_strategy.exclude_plugins.copy()
new_exclude_plugins.append(plugin_id)
exist_strategy.exclude_plugins = new_exclude_plugins
elif exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL:
# In partial mode, disabling one plugin means removing it from include_plugins.
if plugin_id in exist_strategy.include_plugins:
new_include_plugins = exist_strategy.include_plugins.copy()
new_include_plugins.remove(plugin_id)
exist_strategy.include_plugins = new_include_plugins
elif exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL:
# In all mode, switch to exclude mode so only this plugin is skipped.
exist_strategy.upgrade_mode = TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE
exist_strategy.exclude_plugins = [plugin_id]
return True
@staticmethod
def exclude_plugin(
tenant_id: str,
plugin_id: str,
category: PluginCategory,
) -> bool:
with session_factory.create_session() as session, session.begin():
PluginAutoUpgradeService._exclude_plugin(
session,
tenant_id,
category,
plugin_id,
)
return True

37
api/tasks/process_tenant_plugin_autoupgrade_check_task.py
View File

@ -7,7 +7,7 @@ import click
from celery import shared_task
from core.plugin.entities.marketplace import MarketplacePluginSnapshot
from core.plugin.entities.plugin import PluginInstallationSource
from core.plugin.entities.plugin import PluginInstallation, PluginInstallationSource
from core.plugin.impl.plugin import PluginInstaller
from extensions.ext_redis import redis_client
from models.account import TenantPluginAutoUpgradeStrategy
@ -15,6 +15,7 @@ from services.plugin.plugin_service import PluginService
logger = logging.getLogger(__name__)
PluginCategory = TenantPluginAutoUpgradeStrategy.PluginCategory
RETRY_TIMES_OF_ONE_PLUGIN_IN_ONE_TENANT = 3
CACHE_REDIS_KEY_PREFIX = "plugin_autoupgrade_check_task:cached_plugin_snapshot:"
CACHE_REDIS_TTL = 60 * 60 # 1 hour
@ -72,6 +73,25 @@ def marketplace_batch_fetch_plugin_manifests(
return result
def _normalize_category(category: PluginCategory | str | None) -> str | None:
if category is None:
return None
if isinstance(category, PluginCategory):
return category.value
return str(category)
def _plugin_matches_category(plugin: PluginInstallation, category: str | None) -> bool:
"""Return whether an installed plugin should be checked by a category strategy."""
if category is None:
return True
declaration = getattr(plugin, "declaration", None)
plugin_category = getattr(declaration, "category", None)
plugin_category_value = getattr(plugin_category, "value", plugin_category)
return plugin_category_value == category
@shared_task(queue="plugin")
def process_tenant_plugin_autoupgrade_check_task(
tenant_id: str,
@ -80,13 +100,15 @@ def process_tenant_plugin_autoupgrade_check_task(
upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode,
exclude_plugins: list[str],
include_plugins: list[str],
category: PluginCategory | str | None = None,
):
try:
manager = PluginInstaller()
category_value = _normalize_category(category)
click.echo(
click.style(
f"Checking upgradable plugin for tenant: {tenant_id}",
f"Checking upgradable plugin for tenant: {tenant_id}, category: {category_value or 'all'}",
fg="green",
)
)
@ -102,7 +124,11 @@ def process_tenant_plugin_autoupgrade_check_task(
all_plugins = manager.list_plugins(tenant_id)
for plugin in all_plugins:
if plugin.source == PluginInstallationSource.Marketplace and plugin.plugin_id in include_plugins:
if (
plugin.source == PluginInstallationSource.Marketplace
and plugin.plugin_id in include_plugins
and _plugin_matches_category(plugin, category_value)
):
plugin_ids.append(
(
plugin.plugin_id,
@ -117,7 +143,9 @@ def process_tenant_plugin_autoupgrade_check_task(
plugin_ids = [
(plugin.plugin_id, plugin.version, plugin.plugin_unique_identifier)
for plugin in all_plugins
if plugin.source == PluginInstallationSource.Marketplace and plugin.plugin_id not in exclude_plugins
if plugin.source == PluginInstallationSource.Marketplace
and plugin.plugin_id not in exclude_plugins
and _plugin_matches_category(plugin, category_value)
]
elif upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL:
all_plugins = manager.list_plugins(tenant_id)
@ -125,6 +153,7 @@ def process_tenant_plugin_autoupgrade_check_task(
(plugin.plugin_id, plugin.version, plugin.plugin_unique_identifier)
for plugin in all_plugins
if plugin.source == PluginInstallationSource.Marketplace
and _plugin_matches_category(plugin, category_value)
]
if not plugin_ids:

35
api/tests/integration_tests/services/plugin/test_plugin_lifecycle.py
View File

@ -7,6 +7,8 @@ from models.account import TenantPluginAutoUpgradeStrategy, TenantPluginPermissi
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
from services.plugin.plugin_permission_service import PluginPermissionService
PLUGIN_CATEGORY = TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL
@pytest.fixture
def tenant(flask_req_ctx):
@ -71,7 +73,7 @@ class TestPluginPermissionLifecycle:
class TestPluginAutoUpgradeLifecycle:
def test_get_returns_none_for_new_tenant(self, tenant):
assert PluginAutoUpgradeService.get_strategy(tenant) is None
assert PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY) is None
def test_change_creates_row(self, tenant):
result = PluginAutoUpgradeService.change_strategy(
@ -81,10 +83,11 @@ class TestPluginAutoUpgradeLifecycle:
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL,
exclude_plugins=[],
include_plugins=[],
category=PLUGIN_CATEGORY,
)
assert result is True
strategy = PluginAutoUpgradeService.get_strategy(tenant)
strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY)
assert strategy is not None
assert strategy.strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST
assert strategy.upgrade_time_of_day == 3
@ -97,6 +100,7 @@ class TestPluginAutoUpgradeLifecycle:
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL,
exclude_plugins=[],
include_plugins=[],
category=PLUGIN_CATEGORY,
)
PluginAutoUpgradeService.change_strategy(
tenant,
@ -105,9 +109,10 @@ class TestPluginAutoUpgradeLifecycle:
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL,
exclude_plugins=[],
include_plugins=["plugin-a"],
category=PLUGIN_CATEGORY,
)
strategy = PluginAutoUpgradeService.get_strategy(tenant)
strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY)
assert strategy is not None
assert strategy.strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST
assert strategy.upgrade_time_of_day == 12
@ -115,9 +120,9 @@ class TestPluginAutoUpgradeLifecycle:
assert strategy.include_plugins == ["plugin-a"]
def test_exclude_plugin_creates_strategy_when_none_exists(self, tenant):
PluginAutoUpgradeService.exclude_plugin(tenant, "my-plugin")
PluginAutoUpgradeService.exclude_plugin(tenant, "my-plugin", PLUGIN_CATEGORY)
strategy = PluginAutoUpgradeService.get_strategy(tenant)
strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY)
assert strategy is not None
assert strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE
assert "my-plugin" in strategy.exclude_plugins
@ -130,10 +135,11 @@ class TestPluginAutoUpgradeLifecycle:
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
exclude_plugins=["existing"],
include_plugins=[],
category=PLUGIN_CATEGORY,
)
PluginAutoUpgradeService.exclude_plugin(tenant, "new-plugin")
PluginAutoUpgradeService.exclude_plugin(tenant, "new-plugin", PLUGIN_CATEGORY)
strategy = PluginAutoUpgradeService.get_strategy(tenant)
strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY)
assert strategy is not None
assert "existing" in strategy.exclude_plugins
assert "new-plugin" in strategy.exclude_plugins
@ -146,10 +152,11 @@ class TestPluginAutoUpgradeLifecycle:
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
exclude_plugins=["same-plugin"],
include_plugins=[],
category=PLUGIN_CATEGORY,
)
PluginAutoUpgradeService.exclude_plugin(tenant, "same-plugin")
PluginAutoUpgradeService.exclude_plugin(tenant, "same-plugin", PLUGIN_CATEGORY)
strategy = PluginAutoUpgradeService.get_strategy(tenant)
strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY)
assert strategy is not None
assert strategy.exclude_plugins.count("same-plugin") == 1
@ -161,10 +168,11 @@ class TestPluginAutoUpgradeLifecycle:
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL,
exclude_plugins=[],
include_plugins=["p1", "p2"],
category=PLUGIN_CATEGORY,
)
PluginAutoUpgradeService.exclude_plugin(tenant, "p1")
PluginAutoUpgradeService.exclude_plugin(tenant, "p1", PLUGIN_CATEGORY)
strategy = PluginAutoUpgradeService.get_strategy(tenant)
strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY)
assert strategy is not None
assert "p1" not in strategy.include_plugins
assert "p2" in strategy.include_plugins
@ -177,10 +185,11 @@ class TestPluginAutoUpgradeLifecycle:
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL,
exclude_plugins=[],
include_plugins=[],
category=PLUGIN_CATEGORY,
)
PluginAutoUpgradeService.exclude_plugin(tenant, "excluded-plugin")
PluginAutoUpgradeService.exclude_plugin(tenant, "excluded-plugin", PLUGIN_CATEGORY)
strategy = PluginAutoUpgradeService.get_strategy(tenant)
strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY)
assert strategy is not None
assert strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE
assert "excluded-plugin" in strategy.exclude_plugins

86
api/tests/unit_tests/controllers/console/workspace/test_plugin.py
View File

@ -9,12 +9,13 @@ from werkzeug.exceptions import Forbidden
from controllers.console.workspace.plugin import (
PluginAssetApi,
PluginAutoUpgradeExcludePluginApi,
PluginChangeAutoUpgradeApi,
PluginChangePermissionApi,
PluginChangePreferencesApi,
PluginDebuggingKeyApi,
PluginDeleteAllInstallTaskItemsApi,
PluginDeleteInstallTaskApi,
PluginDeleteInstallTaskItemApi,
PluginFetchAutoUpgradeApi,
PluginFetchDynamicSelectOptionsApi,
PluginFetchDynamicSelectOptionsWithCredentialsApi,
PluginFetchInstallTaskApi,
@ -22,7 +23,6 @@ from controllers.console.workspace.plugin import (
PluginFetchManifestApi,
PluginFetchMarketplacePkgApi,
PluginFetchPermissionApi,
PluginFetchPreferencesApi,
PluginIconApi,
PluginInstallFromGithubApi,
PluginInstallFromMarketplaceApi,
@ -901,18 +901,15 @@ class TestPluginFetchDynamicSelectOptionsWithCredentialsApi:
assert result == ({"code": "plugin_error", "message": "error"}, 400)
class TestPluginChangePreferencesApi:
class TestPluginChangeAutoUpgradeApi:
def test_success(self, app: Flask):
api = PluginChangePreferencesApi()
api = PluginChangeAutoUpgradeApi()
method = unwrap(api.post)
user = MagicMock(is_admin_or_owner=True)
payload = {
"permission": {
"install_permission": TenantPluginPermission.InstallPermission.EVERYONE,
"debug_permission": TenantPluginPermission.DebugPermission.EVERYONE,
},
"category": TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL.value,
"auto_upgrade": {
"strategy_setting": TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
"upgrade_time_of_day": 0,
@ -925,24 +922,53 @@ class TestPluginChangePreferencesApi:
with (
app.test_request_context("/", json=payload),
patch("controllers.console.workspace.plugin.current_account_with_tenant", return_value=(user, "t1")),
patch("controllers.console.workspace.plugin.PluginPermissionService.change_permission", return_value=True),
patch("controllers.console.workspace.plugin.PluginAutoUpgradeService.change_strategy", return_value=True),
patch(
"controllers.console.workspace.plugin.PluginAutoUpgradeService.change_strategy", return_value=True
) as change,
):
result = method(api)
assert result["success"] is True
change.assert_called_once()
def test_permission_fail(self, app: Flask):
api = PluginChangePreferencesApi()
def test_success_with_model_category_auto_upgrade(self, app: Flask):
api = PluginChangeAutoUpgradeApi()
method = unwrap(api.post)
user = MagicMock(is_admin_or_owner=True)
payload = {
"permission": {
"install_permission": TenantPluginPermission.InstallPermission.EVERYONE,
"debug_permission": TenantPluginPermission.DebugPermission.EVERYONE,
"category": TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL.value,
"auto_upgrade": {
"strategy_setting": TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST,
"upgrade_time_of_day": 3600,
"upgrade_mode": TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL,
"exclude_plugins": [],
"include_plugins": [],
},
}
with (
app.test_request_context("/", json=payload),
patch("controllers.console.workspace.plugin.current_account_with_tenant", return_value=(user, "t1")),
patch(
"controllers.console.workspace.plugin.PluginAutoUpgradeService.change_strategy", return_value=True
) as change,
):
result = method(api)
assert result["success"] is True
change.assert_called_once()
assert change.call_args.kwargs["category"] == TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL
def test_auto_upgrade_fail(self, app: Flask):
api = PluginChangeAutoUpgradeApi()
method = unwrap(api.post)
user = MagicMock(is_admin_or_owner=True)
payload = {
"category": TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL.value,
"auto_upgrade": {
"strategy_setting": TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
"upgrade_time_of_day": 0,
@ -955,24 +981,20 @@ class TestPluginChangePreferencesApi:
with (
app.test_request_context("/", json=payload),
patch("controllers.console.workspace.plugin.current_account_with_tenant", return_value=(user, "t1")),
patch("controllers.console.workspace.plugin.PluginPermissionService.change_permission", return_value=False),
patch("controllers.console.workspace.plugin.PluginAutoUpgradeService.change_strategy", return_value=False),
):
result = method(api)
assert result["success"] is False
class TestPluginFetchPreferencesApi:
class TestPluginFetchAutoUpgradeApi:
def test_success(self, app: Flask):
api = PluginFetchPreferencesApi()
api = PluginFetchAutoUpgradeApi()
method = unwrap(api.get)
permission = MagicMock(
install_permission=TenantPluginPermission.InstallPermission.EVERYONE,
debug_permission=TenantPluginPermission.DebugPermission.EVERYONE,
)
auto_upgrade = MagicMock(
category=TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL,
strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
upgrade_time_of_day=1,
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
@ -981,19 +1003,19 @@ class TestPluginFetchPreferencesApi:
)
with (
app.test_request_context("/"),
app.test_request_context(
f"/?category={TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL.value}"
),
patch("controllers.console.workspace.plugin.current_account_with_tenant", return_value=(None, "t1")),
patch(
"controllers.console.workspace.plugin.PluginPermissionService.get_permission", return_value=permission
),
patch(
"controllers.console.workspace.plugin.PluginAutoUpgradeService.get_strategy", return_value=auto_upgrade
"controllers.console.workspace.plugin.PluginAutoUpgradeService.get_strategy",
return_value=auto_upgrade,
),
):
result = method(api)
assert "permission" in result
assert "auto_upgrade" in result
assert result["category"] == TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL
assert result["auto_upgrade"]["upgrade_time_of_day"] == 1
class TestPluginAutoUpgradeExcludePluginApi:
@ -1001,7 +1023,7 @@ class TestPluginAutoUpgradeExcludePluginApi:
api = PluginAutoUpgradeExcludePluginApi()
method = unwrap(api.post)
payload = {"plugin_id": "p"}
payload = {"plugin_id": "p", "category": TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL.value}
with (
app.test_request_context("/", json=payload),
@ -1016,7 +1038,7 @@ class TestPluginAutoUpgradeExcludePluginApi:
api = PluginAutoUpgradeExcludePluginApi()
method = unwrap(api.post)
payload = {"plugin_id": "p"}
payload = {"plugin_id": "p", "category": TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL.value}
with (
app.test_request_context("/", json=payload),

113
api/tests/unit_tests/services/plugin/test_plugin_auto_upgrade_service.py
View File

@ -1,8 +1,10 @@
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from models.account import TenantPluginAutoUpgradeStrategy
MODULE = "services.plugin.plugin_auto_upgrade_service"
PLUGIN_CATEGORY = TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL
def _patched_session():
@ -25,7 +27,7 @@ class TestGetStrategy:
with p1:
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
result = PluginAutoUpgradeService.get_strategy("t1")
result = PluginAutoUpgradeService.get_strategy("t1", PLUGIN_CATEGORY)
assert result is strategy
@ -36,7 +38,7 @@ class TestGetStrategy:
with p1:
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
result = PluginAutoUpgradeService.get_strategy("t1")
result = PluginAutoUpgradeService.get_strategy("t1", PLUGIN_CATEGORY)
assert result is None
@ -57,6 +59,7 @@ class TestChangeStrategy:
TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL,
[],
[],
category=PLUGIN_CATEGORY,
)
assert result is True
@ -77,6 +80,7 @@ class TestChangeStrategy:
TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL,
["p1"],
["p2"],
category=PLUGIN_CATEGORY,
)
assert result is True
@ -96,17 +100,19 @@ class TestExcludePlugin:
p1,
patch(f"{MODULE}.select"),
patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls,
patch(f"{MODULE}.PluginAutoUpgradeService.change_strategy") as cs,
):
strat_cls.StrategySetting.FIX_ONLY = "fix_only"
strat_cls.UpgradeMode.EXCLUDE = "exclude"
cs.return_value = True
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
result = PluginAutoUpgradeService.exclude_plugin("t1", "plugin-1")
result = PluginAutoUpgradeService.exclude_plugin(
"t1",
"plugin-1",
PLUGIN_CATEGORY,
)
assert result is True
cs.assert_called_once()
session.add.assert_called_once()
def test_appends_to_exclude_list_in_exclude_mode(self):
p1, session = _patched_session()
@ -121,7 +127,7 @@ class TestExcludePlugin:
strat_cls.UpgradeMode.ALL = "all"
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
result = PluginAutoUpgradeService.exclude_plugin("t1", "p-new")
result = PluginAutoUpgradeService.exclude_plugin("t1", "p-new", PLUGIN_CATEGORY)
assert result is True
assert existing.exclude_plugins == ["p-existing", "p-new"]
@ -139,7 +145,7 @@ class TestExcludePlugin:
strat_cls.UpgradeMode.ALL = "all"
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
result = PluginAutoUpgradeService.exclude_plugin("t1", "p1")
result = PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY)
assert result is True
assert existing.include_plugins == ["p2"]
@ -156,7 +162,7 @@ class TestExcludePlugin:
strat_cls.UpgradeMode.ALL = "all"
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
result = PluginAutoUpgradeService.exclude_plugin("t1", "p1")
result = PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY)
assert result is True
assert existing.upgrade_mode == "exclude"
@ -175,6 +181,93 @@ class TestExcludePlugin:
strat_cls.UpgradeMode.ALL = "all"
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
PluginAutoUpgradeService.exclude_plugin("t1", "p1")
PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY)
assert existing.exclude_plugins == ["p1"]
class TestBackfillStrategyCategories:
def test_creates_default_missing_categories_without_fetching_daemon(self):
p1, session = _patched_session()
tool_strategy = SimpleNamespace(
category=TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL,
strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
upgrade_time_of_day=0,
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
exclude_plugins=[],
include_plugins=[],
)
session.scalars.return_value.all.return_value = [tool_strategy]
installer = MagicMock()
with p1, patch(f"{MODULE}.PluginInstaller", return_value=installer):
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
result = PluginAutoUpgradeService.backfill_strategy_categories("t1")
expected_time = PluginAutoUpgradeService.default_upgrade_time_of_day("t1")
assert result.created_count == len(TenantPluginAutoUpgradeStrategy.PluginCategory) - 1
assert result.normalized is False
installer.list_plugins.assert_not_called()
assert tool_strategy.upgrade_time_of_day == expected_time
created_strategies = [call.args[0] for call in session.add.call_args_list]
model_strategy = next(
strategy
for strategy in created_strategies
if strategy.category == TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL
)
assert model_strategy.strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST
assert model_strategy.upgrade_time_of_day == expected_time
def test_creates_missing_categories_and_splits_known_plugins(self):
p1, session = _patched_session()
tool_strategy = SimpleNamespace(
category=TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL,
strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
upgrade_time_of_day=0,
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
exclude_plugins=["tool-plugin", "model-plugin", "unknown-plugin"],
include_plugins=["model-plugin", "tool-plugin"],
)
model_strategy = SimpleNamespace(
category=TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL,
strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
upgrade_time_of_day=0,
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
exclude_plugins=["tool-plugin", "model-plugin", "unknown-plugin"],
include_plugins=["model-plugin", "tool-plugin"],
)
session.scalars.return_value.all.return_value = [tool_strategy, model_strategy]
installed_plugins = [
SimpleNamespace(
plugin_id="tool-plugin",
declaration=SimpleNamespace(category=TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL),
),
SimpleNamespace(
plugin_id="model-plugin",
declaration=SimpleNamespace(category=TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL),
),
]
installer = MagicMock()
installer.list_plugins.return_value = installed_plugins
with p1, patch(f"{MODULE}.PluginInstaller", return_value=installer), patch(f"{MODULE}.logger") as logger:
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
result = PluginAutoUpgradeService.backfill_strategy_categories("t1")
assert result.created_count == len(TenantPluginAutoUpgradeStrategy.PluginCategory) - 2
assert result.normalized is True
assert session.add.call_count == len(TenantPluginAutoUpgradeStrategy.PluginCategory) - 2
assert tool_strategy.exclude_plugins == ["tool-plugin"]
assert tool_strategy.include_plugins == ["tool-plugin"]
assert model_strategy.exclude_plugins == ["model-plugin"]
assert model_strategy.include_plugins == ["model-plugin"]
logger.warning.assert_called_once_with(
"Skipped unknown plugin IDs while backfilling plugin auto-upgrade strategies: "
"tenant_id=%s, field=%s, plugin_ids=%s",
"t1",
"exclude_plugins",
["unknown-plugin"],
)

32
api/tests/unit_tests/tasks/test_process_tenant_plugin_autoupgrade_check_task.py
View File

@ -4,19 +4,25 @@ from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from core.plugin.entities.marketplace import MarketplacePluginSnapshot
from core.plugin.entities.plugin import PluginInstallationSource
from core.plugin.entities.plugin import PluginCategory, PluginInstallationSource
from models.account import TenantPluginAutoUpgradeStrategy
MODULE = "tasks.process_tenant_plugin_autoupgrade_check_task"
def _make_plugin(plugin_id: str, version: str, source=PluginInstallationSource.Marketplace):
def _make_plugin(
plugin_id: str,
version: str,
source=PluginInstallationSource.Marketplace,
category: PluginCategory = PluginCategory.Tool,
):
"""Build a minimal stand-in for a PluginInstallation entry returned by manager.list_plugins."""
return SimpleNamespace(
plugin_id=plugin_id,
version=version,
plugin_unique_identifier=f"{plugin_id}:{version}@deadbeef",
source=source,
declaration=SimpleNamespace(category=category),
)
@ -39,6 +45,7 @@ def _run_task(
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL,
exclude_plugins=None,
include_plugins=None,
category=None,
):
"""
Execute the celery task synchronously with mocks for the plugin manager,
@ -72,6 +79,7 @@ def _run_task(
upgrade_mode,
exclude_plugins or [],
include_plugins or [],
category,
)
return upgrade_mock, upgrade_calls
@ -246,6 +254,26 @@ class TestUpgradeMode:
assert upgrade_mock.call_count == 1
assert calls[0][1] == plugins[0].plugin_unique_identifier
def test_category_strategy_only_upgrades_matching_category(self):
plugins = [
_make_plugin("acme/model-provider", "1.0.0", category=PluginCategory.Model),
_make_plugin("acme/tool-provider", "1.0.0", category=PluginCategory.Tool),
]
manifests = [
_make_manifest("acme/model-provider", "1.0.1"),
_make_manifest("acme/tool-provider", "1.0.1"),
]
upgrade_mock, calls = _run_task(
plugins=plugins,
manifests=manifests,
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL,
category=TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL,
)
upgrade_mock.assert_called_once()
assert calls[0][1] == plugins[0].plugin_unique_identifier
class TestErrorIsolation:
def test_one_plugin_failure_does_not_block_others(self):