From 851104b133921344eb3815337f8e54f6ea30f4d4 Mon Sep 17 00:00:00 2001 From: JosephTian876 Date: Tue, 1 Sep 2026 05:00:49 +0800 Subject: [PATCH] fix: make plugin imports rollback-safe --- astrbot/core/provider/register.py | 31 ++++ astrbot/core/star/star_manager.py | 146 +++++++++++---- tests/test_plugin_manager.py | 293 ++++++++++++++++++++++++++++++ 3 files changed, 436 insertions(+), 34 deletions(-) diff --git a/astrbot/core/provider/register.py b/astrbot/core/provider/register.py index 409f22f376..a8d7ae33eb 100644 --- a/astrbot/core/provider/register.py +++ b/astrbot/core/provider/register.py @@ -51,3 +51,34 @@ def decorator(cls): return cls return decorator + + +def unregister_provider_adapters_by_module(module_path_prefix: str) -> list[str]: + """Unregister provider adapters declared by one plugin module tree. + + Args: + module_path_prefix: Complete plugin module prefix to unregister. + + Returns: + Adapter type names removed from the registries. + """ + to_remove = [] + for metadata in provider_registry: + module_path = getattr(metadata.cls_type, "__module__", None) + if module_path and ( + module_path == module_path_prefix + or module_path.startswith(f"{module_path_prefix}.") + ): + to_remove.append(metadata) + + for metadata in to_remove: + provider_registry.remove(metadata) + if provider_cls_map.get(metadata.type) is metadata: + provider_cls_map.pop(metadata.type, None) + logger.debug( + "Model provider unregistered: %s (from module %s)", + metadata.type, + getattr(metadata.cls_type, "__module__", None), + ) + + return [metadata.type for metadata in to_remove] diff --git a/astrbot/core/star/star_manager.py b/astrbot/core/star/star_manager.py index c9ef92e5ce..6c5acddf34 100644 --- a/astrbot/core/star/star_manager.py +++ b/astrbot/core/star/star_manager.py @@ -31,7 +31,10 @@ from astrbot.core.config.astrbot_config import AstrBotConfig from astrbot.core.config.default import VERSION from astrbot.core.platform.register import unregister_platform_adapters_by_module -from astrbot.core.provider.register import llm_tools +from astrbot.core.provider.register import ( + llm_tools, + unregister_provider_adapters_by_module, +) from astrbot.core.utils.astrbot_path import ( get_astrbot_config_path, get_astrbot_path, @@ -455,6 +458,9 @@ async def _import_plugin_with_dependency_recovery( try: return __import__(path, fromlist=[module_str]) except ModuleNotFoundError as import_exc: + # The failed import may already have executed decorators. Clear only + # this plugin's partial registrations before retrying in-process. + self._cleanup_plugin_state(root_dir_name, reserved) if recovery_state.mode in { ImportDependencyRecoveryMode.PRELOAD_AND_RECOVER, ImportDependencyRecoveryMode.RECOVER_ON_FAILURE, @@ -468,6 +474,7 @@ async def _import_plugin_with_dependency_recovery( ) if recovered_module is not None: return recovered_module + self._cleanup_plugin_state(root_dir_name, reserved) elif ( recovery_state.mode is ImportDependencyRecoveryMode.REINSTALL_ON_FAILURE ): @@ -480,6 +487,7 @@ async def _import_plugin_with_dependency_recovery( ) await self._check_plugin_dept_update(target_plugin=root_dir_name) + self._cleanup_plugin_state(root_dir_name, reserved) return __import__(path, fromlist=[module_str]) @staticmethod @@ -731,6 +739,25 @@ def _is_plugin_module_path(module_path: str | None, module_prefix: str) -> bool: ) ) + @staticmethod + def _plugin_root_module_path(plugin_module_path: str) -> str: + """Return the import prefix shared by every module in one plugin. + + Args: + plugin_module_path: Plugin main or child module path. + + Returns: + Plugin package import prefix. + """ + parts = plugin_module_path.split(".") + for marker in ("plugins", "builtin_stars"): + if marker not in parts: + continue + marker_index = parts.index(marker) + if marker_index + 1 < len(parts): + return ".".join(parts[: marker_index + 2]) + return plugin_module_path + def _purge_modules( self, module_patterns: list[str] | None = None, @@ -802,14 +829,19 @@ def _cleanup_plugin_state(self, dir_name: str, is_reserved: bool = False) -> Non # 清理工具 for tool in list(llm_tools.func_list): - handler_module_path = getattr(tool, "handler_module_path", None) - if self._is_plugin_module_path(handler_module_path, module_prefix): + if any( + self._is_plugin_llm_tool(concrete_tool, module_prefix) + for concrete_tool in self._iter_concrete_llm_tools(tool) + ): llm_tools.func_list.remove(tool) logger.info(f"Removed tool: {tool.name}") for adapter_name in unregister_platform_adapters_by_module(module_prefix): logger.info(f"Removed platform adapter: {adapter_name}") + for adapter_name in unregister_provider_adapters_by_module(module_prefix): + logger.info(f"Removed provider adapter: {adapter_name}") + def _build_failed_plugin_record( self, *, @@ -877,14 +909,15 @@ def _rebuild_failed_plugin_info(self) -> None: @staticmethod def _iter_concrete_llm_tools(func_tool: FunctionTool) -> Iterable[FunctionTool]: - """Return concrete function tools that may belong to a plugin. + """Return a registered tool and any tools nested in a handoff. Args: func_tool: A registered function tool, possibly a handoff tool. Returns: - The concrete function tools to inspect for plugin ownership. + Function tools to inspect for plugin ownership. """ + yield func_tool if isinstance(func_tool, HandoffTool): agent = getattr(func_tool, "agent", None) tools = getattr(agent, "tools", None) if agent else None @@ -892,10 +925,10 @@ def _iter_concrete_llm_tools(func_tool: FunctionTool) -> Iterable[FunctionTool]: if isinstance(tool, FunctionTool): yield tool return - yield func_tool - @staticmethod + @classmethod def _is_plugin_llm_tool( + cls, func_tool: FunctionTool, plugin_module_path: str | None, ) -> bool: @@ -909,16 +942,56 @@ def _is_plugin_llm_tool( Whether the tool belongs to the plugin module. """ module_path = getattr(func_tool, "handler_module_path", None) + raw_handler = ( + func_tool.handler.func + if isinstance(func_tool.handler, functools.partial) + else func_tool.handler + ) + raw_handler_module_path = getattr(raw_handler, "__module__", None) return bool( plugin_module_path - and module_path and ( - module_path == plugin_module_path - or module_path.startswith(f"{plugin_module_path}.") + cls._is_plugin_module_path(module_path, plugin_module_path) + or cls._is_plugin_module_path( + raw_handler_module_path, + plugin_module_path, + ) + ) + and not (module_path or "").endswith( + ("astrbot.builtin_stars", "data.plugins") ) - and not module_path.endswith(("astrbot.builtin_stars", "data.plugins")) ) + @classmethod + def _claim_llm_tools_for_plugin( + cls, + plugin_module_path: str, + ) -> None: + """Assign plugin-owned tools to the plugin's canonical Star module. + + Args: + plugin_module_path: Canonical Star module path for the plugin. + + Returns: + None. + """ + module_prefix = cls._plugin_root_module_path(plugin_module_path) + for func_tool in llm_tools.func_list: + for tool in cls._iter_concrete_llm_tools(func_tool): + raw_handler = ( + tool.handler.func + if isinstance(tool.handler, functools.partial) + else tool.handler + ) + if cls._is_plugin_module_path( + getattr(raw_handler, "__module__", None), + module_prefix, + ) or cls._is_plugin_module_path( + tool.handler_module_path, + module_prefix, + ): + tool.handler_module_path = plugin_module_path + @classmethod def _iter_plugin_llm_tools( cls, @@ -1124,13 +1197,16 @@ async def load( # 尝试导入模块 try: - module = await self._import_plugin_with_dependency_recovery( - path=path, - module_str=module_str, - root_dir_name=root_dir_name, - requirements_path=requirements_path, - reserved=reserved, - ) + try: + module = await self._import_plugin_with_dependency_recovery( + path=path, + module_str=module_str, + root_dir_name=root_dir_name, + requirements_path=requirements_path, + reserved=reserved, + ) + finally: + self._claim_llm_tools_for_plugin(path) except Exception as e: error_trace = traceback.format_exc() logger.error(error_trace) @@ -1273,19 +1349,21 @@ async def load( # Apply the same idempotent binding lifecycle to LLM tools. for func_tool in llm_tools.func_list: for ft in self._iter_concrete_llm_tools(func_tool): - if ft.handler and ( - getattr(ft.handler, "__module__", None) - == metadata.module_path - or ( - isinstance(ft.handler, functools.partial) - and ft.handler_module_path == metadata.module_path + raw_handler = ( + ft.handler.func + if isinstance(ft.handler, functools.partial) + else ft.handler + ) + if raw_handler and ( + self._is_plugin_module_path( + getattr(raw_handler, "__module__", None), + metadata.module_path, ) - ): - raw_handler = ( - ft.handler.func - if isinstance(ft.handler, functools.partial) - else ft.handler + or self._is_plugin_llm_tool( + ft, + metadata.module_path, ) + ): ft.handler_module_path = metadata.module_path ft.handler = raw_handler if ( @@ -1430,6 +1508,8 @@ async def load( except Exception: logger.error(traceback.format_exc()) + self.failed_plugin_dict.pop(root_dir_name, None) + except BaseException as e: logger.error(f"----- Failed to load plugin {root_dir_name} -----") errors = traceback.format_exc() @@ -1905,11 +1985,9 @@ async def _unbind_plugin(self, plugin_name: str, plugin_module_path: str) -> Non # llm_tools 中移除该插件的工具函数绑定 to_remove = [] for func_tool in llm_tools.func_list: - mp = func_tool.handler_module_path - if ( - mp - and mp.startswith(plugin_module_path) - and not mp.endswith(("astrbot.builtin_stars", "data.plugins")) + if any( + self._is_plugin_llm_tool(tool, plugin_module_path) + for tool in self._iter_concrete_llm_tools(func_tool) ): to_remove.append(func_tool) for func_tool in to_remove: diff --git a/tests/test_plugin_manager.py b/tests/test_plugin_manager.py index 5cb8747ebb..f88a12b157 100644 --- a/tests/test_plugin_manager.py +++ b/tests/test_plugin_manager.py @@ -994,6 +994,299 @@ async def mock_sync_command_configs(): assert plugin_name in plugin_manager_pm.failed_plugin_dict +@pytest.mark.asyncio +async def test_partial_import_rolls_back_only_its_provider_registration( + plugin_manager_pm: PluginManager, + monkeypatch, +): + """A failed import removes only registrations owned by that plugin.""" + from astrbot.core.provider.register import ( + provider_cls_map, + provider_registry, + register_provider_adapter, + ) + + _clear_star_runtime_state() + registry_before = list(provider_registry) + map_before = dict(provider_cls_map) + plugin_name = "partial_provider_plugin" + provider_type = "partial_provider_adapter" + core_provider_type = "partial_provider_core_sentinel" + sibling_provider_type = "partial_provider_sibling_sentinel" + plugin_root = Path(plugin_manager_pm.plugin_store_path).parents[1] + plugin_path = Path(plugin_manager_pm.plugin_store_path) / plugin_name + plugin_path.mkdir(parents=True) + (plugin_path / "metadata.yaml").write_text( + yaml.dump( + { + "name": plugin_name, + "author": "AstrBot Team", + "desc": "Partial provider test plugin", + "version": "1.0.0", + } + ), + encoding="utf-8", + ) + (plugin_path / "main.py").write_text( + f"""from astrbot.core.provider.register import register_provider_adapter + + +@register_provider_adapter("{provider_type}", "Partial test provider") +class PartialProvider: + pass + + +raise RuntimeError("import failed after registration") +""", + encoding="utf-8", + ) + + class CoreSentinel: + pass + + class SiblingSentinel: + pass + + CoreSentinel.__module__ = "astrbot.core.provider.sources.test_sentinel" + SiblingSentinel.__module__ = f"data.plugins.{plugin_name}_extra.main" + register_provider_adapter(core_provider_type, "Core sentinel")(CoreSentinel) + register_provider_adapter(sibling_provider_type, "Sibling sentinel")( + SiblingSentinel + ) + core_metadata = provider_cls_map[core_provider_type] + sibling_metadata = provider_cls_map[sibling_provider_type] + + async def mock_global_get(key, default=None): + del key + return default + + async def mock_sync_command_configs(): + return None + + monkeypatch.syspath_prepend(str(plugin_root)) + monkeypatch.setattr(star_manager_module.sp, "global_get", mock_global_get) + monkeypatch.setattr( + star_manager_module, + "sync_command_configs", + mock_sync_command_configs, + ) + + try: + success, error = await plugin_manager_pm.load(specified_dir_name=plugin_name) + + assert success is False + assert error is not None + assert "import failed after registration" in error + assert provider_type not in provider_cls_map + assert all(metadata.type != provider_type for metadata in provider_registry) + assert provider_cls_map[core_provider_type] is core_metadata + assert provider_cls_map[sibling_provider_type] is sibling_metadata + assert plugin_name in plugin_manager_pm.failed_plugin_dict + + (plugin_path / "main.py").write_text( + f"""from astrbot.api.star import Star +from astrbot.core.provider.register import register_provider_adapter + + +@register_provider_adapter("{provider_type}", "Partial test provider") +class PartialProvider: + pass + + +class Main(Star): + pass +""", + encoding="utf-8", + ) + + success, error = await plugin_manager_pm.load(specified_dir_name=plugin_name) + + assert success is True + assert error is None + assert sum(md.type == provider_type for md in provider_registry) == 1 + assert plugin_name not in plugin_manager_pm.failed_plugin_dict + finally: + plugin_manager_pm._cleanup_plugin_state(plugin_name) + provider_registry[:] = registry_before + provider_cls_map.clear() + provider_cls_map.update(map_before) + _clear_star_runtime_state() + + +@pytest.mark.asyncio +async def test_dependency_recovery_reimports_partial_decorators_once( + plugin_manager_pm: PluginManager, + monkeypatch, +): + """An in-process dependency retry first rolls back decorator side effects.""" + from astrbot.core.provider.register import provider_cls_map, provider_registry + + _clear_star_runtime_state() + registry_before = list(provider_registry) + map_before = dict(provider_cls_map) + plugin_name = "dependency_recovery_provider_plugin" + provider_type = "dependency_recovery_provider_adapter" + missing_dependency = "astrbot_test_recovered_dependency" + plain_tool_name = "dependency_recovery_plain_tool" + agent_name = "dependency_recovery_agent" + agent_tool_name = "dependency_recovery_agent_tool" + module_path = f"data.plugins.{plugin_name}.main" + plugin_root = Path(plugin_manager_pm.plugin_store_path).parents[1] + plugin_path = Path(plugin_manager_pm.plugin_store_path) / plugin_name + plugin_path.mkdir(parents=True) + (plugin_path / "metadata.yaml").write_text( + yaml.dump( + { + "name": plugin_name, + "author": "AstrBot Team", + "desc": "Provider dependency recovery test plugin", + "version": "1.0.0", + } + ), + encoding="utf-8", + ) + (plugin_path / "requirements.txt").write_text( + f"{missing_dependency}\n", + encoding="utf-8", + ) + (plugin_path / "tools.py").write_text( + f"""from astrbot.core.star.register.star_handler import register_agent, register_llm_tool + + +@register_llm_tool(name="{plain_tool_name}") +async def plain_tool(self): + pass + + +@register_agent(name="{agent_name}", instruction="Recovery helper") +async def helper_agent(self): + pass + + +@helper_agent.llm_tool(name="{agent_tool_name}") +async def agent_tool(self): + pass +""", + encoding="utf-8", + ) + (plugin_path / "main.py").write_text( + f"""from astrbot.api.star import Star +from astrbot.core.provider.register import register_provider_adapter + + +@register_provider_adapter("{provider_type}", "Recovered provider") +class RecoveredProvider: + pass + + +from . import tools +import {missing_dependency} + + +class Main(Star): + pass +""", + encoding="utf-8", + ) + llm_tools = cast(Any, star_manager_module.llm_tools) + original_func_list = llm_tools.func_list + llm_tools.func_list = list(original_func_list) + + async def unrelated_dynamic_tool(): + return None + + unrelated_dynamic_tool.__module__ = "external.runtime_tools" + unrelated_tool = star_manager_module.FunctionTool( + name="unrelated_dynamic_tool", + description="Registered outside the plugin import", + parameters={"type": "object", "properties": {}}, + handler=unrelated_dynamic_tool, + ) + llm_tools.func_list.append(unrelated_tool) + dependency_installed = False + + async def mock_global_get(key, default=None): + del key + return default + + async def mock_sync_command_configs(): + return None + + async def mock_check_plugin_dept_update(*, target_plugin=None): + nonlocal dependency_installed + assert target_plugin == plugin_name + dependency_installed = True + monkeypatch.setitem( + star_manager_module.sys.modules, + missing_dependency, + ModuleType(missing_dependency), + ) + + recovery_state = star_manager_module.ImportDependencyRecoveryState( + star_manager_module.ImportDependencyRecoveryMode.REINSTALL_ON_FAILURE, + MissingRequirementsPlan( + missing_names=frozenset({missing_dependency}), + install_lines=(missing_dependency,), + version_mismatch_names=frozenset({missing_dependency}), + ), + ) + + monkeypatch.syspath_prepend(str(plugin_root)) + monkeypatch.setattr(star_manager_module.sp, "global_get", mock_global_get) + monkeypatch.setattr( + star_manager_module, + "sync_command_configs", + mock_sync_command_configs, + ) + monkeypatch.setattr( + plugin_manager_pm, + "_resolve_import_dependency_recovery_state", + lambda *args, **kwargs: recovery_state, + ) + monkeypatch.setattr( + plugin_manager_pm, + "_check_plugin_dept_update", + mock_check_plugin_dept_update, + ) + + try: + success, error = await plugin_manager_pm.load(specified_dir_name=plugin_name) + + assert success is True + assert error is None + assert dependency_installed is True + assert sum(md.type == provider_type for md in provider_registry) == 1 + plain_tools = [ + tool for tool in llm_tools.func_list if tool.name == plain_tool_name + ] + handoff_tools = [ + tool + for tool in llm_tools.func_list + if tool.name == f"transfer_to_{agent_name}" + ] + assert len(plain_tools) == 1 + assert len(handoff_tools) == 1 + nested_tools = [ + tool + for tool in handoff_tools[0].agent.tools + if tool.name == agent_tool_name + ] + assert len(nested_tools) == 1 + for tool in [plain_tools[0], handoff_tools[0], nested_tools[0]]: + assert tool.handler_module_path == module_path + assert isinstance(tool.handler, functools.partial) + assert tool.handler.func.__module__ == f"data.plugins.{plugin_name}.tools" + assert unrelated_tool.handler_module_path is None + assert unrelated_tool.active is True + finally: + plugin_manager_pm._cleanup_plugin_state(plugin_name) + provider_registry[:] = registry_before + provider_cls_map.clear() + provider_cls_map.update(map_before) + llm_tools.func_list = original_func_list + _clear_star_runtime_state() + + @pytest.mark.asyncio async def test_ensure_plugin_requirements_reraises_cancelled_error( plugin_manager_pm: PluginManager, local_updater: Path, monkeypatch