From 3fa8ddf63ae95c5a73f21676b1347d9c53f4d871 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 14:47:17 +0800 Subject: [PATCH] fix(plugin): close cancellation compensation gaps --- app/adapters/external/market.py | 39 ++++++++++++- app/application/plugin/install.py | 72 ++++++++++++++++------- tests/test_plugin_helper.py | 87 ++++++++++++++++++++++++++++ tests/test_plugin_install_command.py | 38 +++++++++++- 4 files changed, 212 insertions(+), 24 deletions(-) diff --git a/app/adapters/external/market.py b/app/adapters/external/market.py index 23cb67251..81c3119cc 100644 --- a/app/adapters/external/market.py +++ b/app/adapters/external/market.py @@ -1486,6 +1486,42 @@ class PluginHelper(metaclass=WeakSingleton): temp_file.write(f"{cls.__format_package_name(package_name)}>={version}\n") return Path(temp_file.name) + @classmethod + async def __async_create_runtime_constraints_file( + cls, + protected_packages: Dict[str, Version], + ) -> Path: + """创建临时约束文件，取消时等待创建收口并删除已产生的文件。""" + create_task = asyncio.create_task( + asyncio.to_thread( + cls.__create_runtime_constraints_file, + protected_packages, + ) + ) + try: + return await asyncio.shield(create_task) + except asyncio.CancelledError: + async def cleanup_created_file() -> None: + try: + created_file = await create_task + except BaseException: + return + await asyncio.to_thread(created_file.unlink, missing_ok=True) + + cleanup_task = asyncio.create_task(cleanup_created_file()) + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + continue + except Exception: + break + try: + await cleanup_task + except Exception as err: + logger.warning(f"[UV] 取消后清理运行环境约束文件失败：{err}") + raise + @staticmethod def __refresh_import_system(): """ @@ -2608,8 +2644,7 @@ class PluginHelper(metaclass=WeakSingleton): constraints_file = None if protected_packages: try: - constraints_file = await _await_thread_operation( - cls.__create_runtime_constraints_file, + constraints_file = await cls.__async_create_runtime_constraints_file( protected_packages, ) except Exception as err: diff --git a/app/application/plugin/install.py b/app/application/plugin/install.py index f858d2b9e..3c50bc547 100644 --- a/app/application/plugin/install.py +++ b/app/application/plugin/install.py @@ -72,6 +72,7 @@ class _InstallState: installed_list_persisted: bool = False runtime_touched: bool = False registrations_touched: bool = False + refresh_compensated: bool = False committed: bool = False original_plugins: list[str] = field(default_factory=list) @@ -151,6 +152,7 @@ class PluginInstallCommand: return await self._refresh_existing( plugin_id=plugin_id, repo_url=repo_url, + state=state, ) if not repo_url: return PluginInstallResult( @@ -319,6 +321,8 @@ class PluginInstallCommand: f"插件 {plugin_id} 在安装提交后被取消，Python 依赖环境可能已经改变" ) return + if state.refresh_compensated: + return if state.checkpoint is None: logger.warning( f"插件 {plugin_id} 在创建安装快照前被取消，无法执行文件补偿" @@ -362,6 +366,7 @@ class PluginInstallCommand: *, plugin_id: str, repo_url: Optional[str], + state: _InstallState, ) -> PluginInstallResult: """刷新已存在插件，不触碰包文件和已安装列表。""" if repo_url: @@ -381,33 +386,31 @@ class PluginInstallCommand: await self._plugin_reloader(plugin_id) failure_stage = "registration_refresh" await self._registration_refresher(plugin_id) - except Exception as err: - rollback_errors = [] - runtime_restored = False - registrations_restored = False - try: - await self._plugin_reloader(plugin_id) - runtime_restored = True - except Exception as rollback_err: - rollback_errors.append(f"运行态恢复失败：{rollback_err}") - if runtime_restored: + except asyncio.CancelledError: + cleanup_task = asyncio.create_task( + self._restore_refreshed_runtime(plugin_id) + ) + while not cleanup_task.done(): try: - await self._registration_refresher(plugin_id) - registrations_restored = True - except Exception as rollback_err: - rollback_errors.append(f"路由和服务注册恢复失败：{rollback_err}") + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + continue + rollback = await cleanup_task + state.refresh_compensated = True + if rollback.errors: + logger.error( + f"插件 {plugin_id} 取消刷新后的运行态补偿存在错误：" + f"{'；'.join(rollback.errors)}" + ) + raise + except Exception as err: + rollback = await self._restore_refreshed_runtime(plugin_id) result = PluginInstallResult( success=False, message=f"刷新插件运行态失败：{err}", refreshed_only=True, failure_stage=failure_stage, - rollback=PluginInstallRollback( - runtime_attempted=True, - runtime_restored=runtime_restored, - registrations_attempted=True, - registrations_restored=registrations_restored, - errors=tuple(rollback_errors), - ), + rollback=rollback, ) if isinstance(err, DatabaseWorkerOverloadedError): raise @@ -436,6 +439,33 @@ class PluginInstallCommand: report_error=report_error, ) + async def _restore_refreshed_runtime( + self, + plugin_id: str, + ) -> PluginInstallRollback: + """重新加载插件并刷新注册，使中断的运行态切换恢复到完整状态。""" + errors = [] + runtime_restored = False + registrations_restored = False + try: + await self._plugin_reloader(plugin_id) + runtime_restored = True + except Exception as err: + errors.append(f"运行态恢复失败：{err}") + if runtime_restored: + try: + await self._registration_refresher(plugin_id) + registrations_restored = True + except Exception as err: + errors.append(f"路由和服务注册恢复失败：{err}") + return PluginInstallRollback( + runtime_attempted=True, + runtime_restored=runtime_restored, + registrations_attempted=True, + registrations_restored=registrations_restored, + errors=tuple(errors), + ) + async def _failure( self, *, diff --git a/tests/test_plugin_helper.py b/tests/test_plugin_helper.py index 2495a99b1..61a45b046 100644 --- a/tests/test_plugin_helper.py +++ b/tests/test_plugin_helper.py @@ -1774,6 +1774,93 @@ demo = { index = "private" } ): asyncio.run(run_install()) + def test_constraints_created_during_cancellation_are_removed(self, tmp_path): + """约束文件创建线程收口后仍须响应取消并删除临时文件。""" + from app.adapters.external.market import PluginHelper + + helper = PluginHelper() + requirements_file = tmp_path / "requirements.txt" + requirements_file.write_text("demo-package\n", encoding="utf-8") + constraints_file = tmp_path / "runtime-constraints.txt" + created = threading.Event() + release = threading.Event() + + def create_constraints(_protected_packages): + constraints_file.write_text("fastapi==0\n", encoding="utf-8") + created.set() + release.wait(timeout=2) + return constraints_file + + async def run_install(): + task = asyncio.create_task( + helper.async_install_packages_with_fallback(requirements_file) + ) + assert await asyncio.to_thread(created.wait, 2) + task.cancel() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + with patch.object( + PluginHelper, + "_PluginHelper__get_installed_packages", + return_value={}, + ), patch.object( + PluginHelper, + "_PluginHelper__get_protected_runtime_packages", + return_value={"fastapi": "0"}, + ), patch.object( + PluginHelper, + "_PluginHelper__validate_runtime_dependency_conflicts", + return_value=(True, ""), + ), patch.object( + PluginHelper, + "_PluginHelper__create_runtime_constraints_file", + side_effect=create_constraints, + ): + asyncio.run(run_install()) + + assert not constraints_file.exists() + + def test_constraints_cleanup_failure_preserves_cancellation(self, tmp_path): + """临时文件删除失败只记录日志，不得替换调用方的取消异常。""" + from app.adapters.external.market import PluginHelper + + constraints_file = tmp_path / "runtime-constraints.txt" + created = threading.Event() + release = threading.Event() + + def create_constraints(_protected_packages): + constraints_file.write_text("fastapi==0\n", encoding="utf-8") + created.set() + release.wait(timeout=2) + return constraints_file + + async def run_create(): + task = asyncio.create_task( + PluginHelper._PluginHelper__async_create_runtime_constraints_file( + {"fastapi": Version("0")} + ) + ) + assert await asyncio.to_thread(created.wait, 2) + task.cancel() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + with patch.object( + PluginHelper, + "_PluginHelper__create_runtime_constraints_file", + side_effect=create_constraints, + ), patch.object( + Path, + "unlink", + side_effect=PermissionError("locked"), + ), patch("app.adapters.external.market.logger.warning") as warning: + asyncio.run(run_create()) + + warning.assert_called_once() + def test_install_uses_release_package_when_asset_is_available(self, monkeypatch): """ release 包可用时优先使用 zip 安装，不再额外访问文件列表。 diff --git a/tests/test_plugin_install_command.py b/tests/test_plugin_install_command.py index d1736662a..c77d8a4df 100644 --- a/tests/test_plugin_install_command.py +++ b/tests/test_plugin_install_command.py @@ -1,5 +1,5 @@ import asyncio -from unittest.mock import AsyncMock, Mock +from unittest.mock import AsyncMock, Mock, patch import pytest @@ -156,6 +156,42 @@ async def test_existing_plugin_checks_compatibility_without_reinstalling_package checkpointer.assert_not_awaited() +@pytest.mark.asyncio +async def test_cancelled_existing_plugin_refresh_restores_runtime_and_registrations(): + """已存在插件刷新被取消时，必须重新收敛运行态和注册。""" + registration_started = asyncio.Event() + calls: list[str] = [] + + async def reload_plugin(_plugin_id: str) -> None: + calls.append("reload") + + async def refresh_registrations(_plugin_id: str) -> None: + calls.append("registrations") + if calls.count("registrations") == 1: + registration_started.set() + await asyncio.Event().wait() + + with patch("app.application.plugin.install.logger.warning") as warning: + task = asyncio.create_task( + _command( + installed=["DemoPlugin"], + plugin_ids=["DemoPlugin"], + reloader=reload_plugin, + refresher=refresh_registrations, + ).execute( + plugin_id="DemoPlugin", + repo_url="https://github.com/demo/plugins", + ) + ) + await registration_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert calls == ["reload", "registrations", "reload", "registrations"] + warning.assert_not_called() + + @pytest.mark.asyncio async def test_persistence_failure_restores_package_without_touching_runtime(): """已安装列表保存失败时恢复文件，且运行态尚未开始切换。"""