diff --git a/astrbot/core/utils/pip_installer.py b/astrbot/core/utils/pip_installer.py index 97e9653d6b..8aad8db75a 100644 --- a/astrbot/core/utils/pip_installer.py +++ b/astrbot/core/utils/pip_installer.py @@ -5,12 +5,14 @@ import importlib.util import io import logging +import ntpath import os import re import shlex import sys import threading from collections import deque +from collections.abc import Mapping from dataclasses import dataclass from urllib.parse import urlparse @@ -30,6 +32,9 @@ _DISTLIB_FINDER_PATCH_ATTEMPTED = False _SITE_PACKAGES_IMPORT_LOCK = threading.RLock() +_PIP_IN_PROCESS_ENV_LOCK = threading.RLock() +_WINDOWS_UNC_PATH_PREFIXES = ("\\\\?\\UNC\\", "\\??\\UNC\\") +_WINDOWS_EXTENDED_PATH_PREFIXES = ("\\\\?\\", "\\??\\") _PIP_FAILURE_PATTERNS = { "error_prefix": re.compile(r"^\s*error:", re.IGNORECASE), "user_requested": re.compile(r"\bthe user requested\b", re.IGNORECASE), @@ -235,6 +240,120 @@ def _run_pip_main_streaming(pip_main, args: list[str]) -> tuple[int, list[str]]: return result_code, stream.lines +@contextlib.contextmanager +def _temporary_environ(updates: Mapping[str, str]): + if not updates: + yield + return + + missing = object() + previous_values = {key: os.environ.get(key, missing) for key in updates} + + try: + os.environ.update(updates) + yield + finally: + for key, previous_value in previous_values.items(): + if previous_value is missing: + os.environ.pop(key, None) + else: + assert isinstance(previous_value, str) + os.environ[key] = previous_value + + +def _run_pip_main_with_temporary_environ( + pip_main, + args: list[str], +) -> tuple[int, list[str]]: + # os.environ is process-wide; serialize reading current INCLUDE/LIB values + # together with the temporary mutation window around the in-process pip + # invocation. + with _PIP_IN_PROCESS_ENV_LOCK: + env_updates = _build_packaged_windows_runtime_build_env(base_env=os.environ) + if not env_updates: + return _run_pip_main_streaming(pip_main, args) + + with _temporary_environ(env_updates): + return _run_pip_main_streaming(pip_main, args) + + +def _normalize_windows_native_build_path(path: str) -> str: + """Normalize a Windows path returned by native APIs or sys.executable. + + Extended UNC prefixes are converted back to the standard ``\\server`` form, + other extended prefixes are stripped, and the remaining path is normalized. + """ + normalized = path.replace("/", "\\") + + # Extended UNC: \\?\UNC\server\share\... -> \\server\share\... + for prefix in _WINDOWS_UNC_PATH_PREFIXES: + if normalized.startswith(prefix): + return ntpath.normpath(f"\\\\{normalized[len(prefix) :]}") + + # Other extended prefixes are stripped before normalizing the path. + for prefix in _WINDOWS_EXTENDED_PATH_PREFIXES: + if normalized.startswith(prefix): + normalized = normalized[len(prefix) :] + break + + return ntpath.normpath(normalized) + + +def _get_case_insensitive_env_value( + env: Mapping[str, str], + upper_to_key: Mapping[str, str], + name: str, +) -> str | None: + direct = env.get(name) + if direct is not None: + return direct + + existing_key = upper_to_key.get(name.upper()) + if existing_key is not None: + return env.get(existing_key) + + return None + + +def _build_packaged_windows_runtime_build_env( + *, + base_env: Mapping[str, str] | None = None, +) -> dict[str, str]: + if sys.platform != "win32" or not is_packaged_desktop_runtime(): + return {} + + base_env = os.environ if base_env is None else base_env + + runtime_executable = _normalize_windows_native_build_path(sys.executable) + runtime_dir = ntpath.dirname(runtime_executable) + if not runtime_dir: + return {} + + include_dir = _normalize_windows_native_build_path( + ntpath.join(runtime_dir, "include") + ) + libs_dir = _normalize_windows_native_build_path(ntpath.join(runtime_dir, "libs")) + include_exists = os.path.isdir(include_dir) + libs_exists = os.path.isdir(libs_dir) + + if not (include_exists or libs_exists): + return {} + + upper_to_key = {key.upper(): key for key in base_env} + env_updates: dict[str, str] = {} + + if include_exists: + existing = _get_case_insensitive_env_value(base_env, upper_to_key, "INCLUDE") + env_updates["INCLUDE"] = ( + f"{include_dir};{existing}" if existing else include_dir + ) + if libs_exists: + existing = _get_case_insensitive_env_value(base_env, upper_to_key, "LIB") + env_updates["LIB"] = f"{libs_dir};{existing}" if existing else libs_dir + + return env_updates + + def _matches_pip_failure_pattern(line: str, *pattern_names: str) -> bool: names = pattern_names or tuple(_PIP_FAILURE_PATTERNS) return any(_PIP_FAILURE_PATTERNS[name].search(line) for name in names) @@ -931,7 +1050,9 @@ async def _run_pip_in_process(self, args: list[str]) -> int: original_handlers = list(logging.getLogger().handlers) try: result_code, output_lines = await asyncio.to_thread( - _run_pip_main_streaming, pip_main, args + _run_pip_main_with_temporary_environ, + pip_main, + args, ) finally: _cleanup_added_root_handlers(original_handlers) diff --git a/tests/test_pip_installer.py b/tests/test_pip_installer.py index d61507f4c5..adbd174e32 100644 --- a/tests/test_pip_installer.py +++ b/tests/test_pip_installer.py @@ -1,4 +1,5 @@ import asyncio +import ntpath import threading from unittest.mock import AsyncMock @@ -9,6 +10,15 @@ from astrbot.core.utils import requirements_utils from astrbot.core.utils.pip_installer import PipInstaller +WINDOWS_RUNTIME_ROOT = ntpath.join(r"C:\astrbot-test", "backend", "python") +WINDOWS_RUNTIME_EXECUTABLE = ntpath.join(WINDOWS_RUNTIME_ROOT, "python.exe") +WINDOWS_PACKAGED_RUNTIME_EXECUTABLE = f"\\\\?\\{WINDOWS_RUNTIME_EXECUTABLE}" +WINDOWS_RUNTIME_INCLUDE_DIR = ntpath.join(WINDOWS_RUNTIME_ROOT, "include") +WINDOWS_RUNTIME_LIBS_DIR = ntpath.join(WINDOWS_RUNTIME_ROOT, "libs") +EXISTING_WINDOWS_INCLUDE_DIR = ntpath.join(r"C:\toolchain", "include") +EXISTING_WINDOWS_LIB_DIR = ntpath.join(r"C:\toolchain", "lib") +_ENV_MISSING = object() + def _make_run_pip_mock( code: int = 0, @@ -24,6 +34,64 @@ async def run_pip(*args, **kwargs): return AsyncMock(side_effect=run_pip) +def _configure_run_pip_in_process_capture( + monkeypatch, + *, + platform: str, + packaged_runtime: bool, + runtime_executable: str = WINDOWS_PACKAGED_RUNTIME_EXECUTABLE, + include_value: str | object = _ENV_MISSING, + lib_value: str | object = _ENV_MISSING, + existing_runtime_dirs: set[str] | None = None, +) -> dict[str, str | None]: + observed_env: dict[str, str | None] = {} + + def fake_pip_main(args): + del args + observed_env["INCLUDE"] = pip_installer_module.os.environ.get("INCLUDE") + observed_env["LIB"] = pip_installer_module.os.environ.get("LIB") + return 0 + + if packaged_runtime: + monkeypatch.setenv("ASTRBOT_DESKTOP_CLIENT", "1") + else: + monkeypatch.delenv("ASTRBOT_DESKTOP_CLIENT", raising=False) + + if include_value is _ENV_MISSING: + monkeypatch.delenv("INCLUDE", raising=False) + else: + monkeypatch.setenv("INCLUDE", include_value) + + if lib_value is _ENV_MISSING: + monkeypatch.delenv("LIB", raising=False) + else: + monkeypatch.setenv("LIB", lib_value) + + monkeypatch.setattr(pip_installer_module.sys, "platform", platform) + monkeypatch.setattr(pip_installer_module.sys, "executable", runtime_executable) + + if existing_runtime_dirs is not None: + monkeypatch.setattr( + pip_installer_module.os.path, + "isdir", + lambda path: path in existing_runtime_dirs, + ) + + monkeypatch.setattr( + "astrbot.core.utils.pip_installer._get_pip_main", + lambda: fake_pip_main, + ) + return observed_env + + +@pytest.fixture +def configure_run_pip_in_process_capture(monkeypatch): + def _configure(**kwargs): + return _configure_run_pip_in_process_capture(monkeypatch, **kwargs) + + return _configure + + @pytest.mark.asyncio async def test_install_targets_site_packages_for_desktop_client(monkeypatch, tmp_path): monkeypatch.setenv("ASTRBOT_DESKTOP_CLIENT", "1") @@ -226,6 +294,316 @@ def fake_pip_main(args): ] +@pytest.mark.parametrize( + ("path", "expected"), + [ + ( + WINDOWS_RUNTIME_EXECUTABLE, + WINDOWS_RUNTIME_EXECUTABLE, + ), + ( + WINDOWS_PACKAGED_RUNTIME_EXECUTABLE, + WINDOWS_RUNTIME_EXECUTABLE, + ), + ( + f"\\??\\{WINDOWS_RUNTIME_EXECUTABLE}", + WINDOWS_RUNTIME_EXECUTABLE, + ), + ( + r"\\?\UNC\server\share\include", + r"\\server\share\include", + ), + ( + r"\??\UNC\server\share\libs", + r"\\server\share\libs", + ), + ( + r"\\server\share\include", + r"\\server\share\include", + ), + ( + "C:/astrbot-test/backend/python/libs", + WINDOWS_RUNTIME_LIBS_DIR, + ), + ], +) +def test_normalize_windows_native_build_path_variants(path, expected): + assert pip_installer_module._normalize_windows_native_build_path(path) == expected + + +def test_temporary_environ_restores_previous_values(monkeypatch): + monkeypatch.setenv("INCLUDE", EXISTING_WINDOWS_INCLUDE_DIR) + monkeypatch.delenv("LIB", raising=False) + + with pip_installer_module._temporary_environ( + { + "INCLUDE": WINDOWS_RUNTIME_INCLUDE_DIR, + "LIB": WINDOWS_RUNTIME_LIBS_DIR, + } + ): + assert pip_installer_module.os.environ["INCLUDE"] == WINDOWS_RUNTIME_INCLUDE_DIR + assert pip_installer_module.os.environ["LIB"] == WINDOWS_RUNTIME_LIBS_DIR + + assert pip_installer_module.os.environ["INCLUDE"] == EXISTING_WINDOWS_INCLUDE_DIR + assert "LIB" not in pip_installer_module.os.environ + + +def test_build_packaged_windows_runtime_build_env_uses_base_env_snapshot( + monkeypatch, +): + snapshot_include = ntpath.join(r"C:\snapshot-toolchain", "include") + snapshot_lib = ntpath.join(r"C:\snapshot-toolchain", "lib") + process_include = ntpath.join(r"C:\process-toolchain", "include") + process_lib = ntpath.join(r"C:\process-toolchain", "lib") + + monkeypatch.setenv("ASTRBOT_DESKTOP_CLIENT", "1") + monkeypatch.setenv("INCLUDE", process_include) + monkeypatch.setenv("LIB", process_lib) + monkeypatch.setattr(pip_installer_module.sys, "platform", "win32") + monkeypatch.setattr( + pip_installer_module.sys, + "executable", + WINDOWS_PACKAGED_RUNTIME_EXECUTABLE, + ) + monkeypatch.setattr( + pip_installer_module.os.path, + "isdir", + lambda path: path in {WINDOWS_RUNTIME_INCLUDE_DIR, WINDOWS_RUNTIME_LIBS_DIR}, + ) + + env_updates = pip_installer_module._build_packaged_windows_runtime_build_env( + base_env={ + "INCLUDE": snapshot_include, + "LIB": snapshot_lib, + } + ) + + assert env_updates == { + "INCLUDE": f"{WINDOWS_RUNTIME_INCLUDE_DIR};{snapshot_include}", + "LIB": f"{WINDOWS_RUNTIME_LIBS_DIR};{snapshot_lib}", + } + + +def test_build_packaged_windows_runtime_build_env_matches_snapshot_keys_case_insensitively( + monkeypatch, +): + snapshot_include = ntpath.join(r"C:\snapshot-toolchain", "include") + snapshot_lib = ntpath.join(r"C:\snapshot-toolchain", "lib") + + monkeypatch.setenv("ASTRBOT_DESKTOP_CLIENT", "1") + monkeypatch.setattr(pip_installer_module.sys, "platform", "win32") + monkeypatch.setattr( + pip_installer_module.sys, + "executable", + WINDOWS_PACKAGED_RUNTIME_EXECUTABLE, + ) + monkeypatch.setattr( + pip_installer_module.os.path, + "isdir", + lambda path: path in {WINDOWS_RUNTIME_INCLUDE_DIR, WINDOWS_RUNTIME_LIBS_DIR}, + ) + + env_updates = pip_installer_module._build_packaged_windows_runtime_build_env( + base_env={ + "include": snapshot_include, + "lib": snapshot_lib, + } + ) + + assert env_updates == { + "INCLUDE": f"{WINDOWS_RUNTIME_INCLUDE_DIR};{snapshot_include}", + "LIB": f"{WINDOWS_RUNTIME_LIBS_DIR};{snapshot_lib}", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("include_exists", "libs_exists"), + [ + (True, True), + (True, False), + (False, True), + ], +) +async def test_run_pip_in_process_injects_windows_runtime_build_env( + configure_run_pip_in_process_capture, include_exists, libs_exists +): + existing_runtime_dirs = set() + if include_exists: + existing_runtime_dirs.add(WINDOWS_RUNTIME_INCLUDE_DIR) + if libs_exists: + existing_runtime_dirs.add(WINDOWS_RUNTIME_LIBS_DIR) + + observed_env = configure_run_pip_in_process_capture( + platform="win32", + packaged_runtime=True, + include_value=EXISTING_WINDOWS_INCLUDE_DIR, + lib_value=EXISTING_WINDOWS_LIB_DIR, + existing_runtime_dirs=existing_runtime_dirs, + ) + + installer = PipInstaller("") + result = await installer._run_pip_in_process(["install", "demo-package"]) + + assert result == 0 + expected_include = EXISTING_WINDOWS_INCLUDE_DIR + expected_lib = EXISTING_WINDOWS_LIB_DIR + if include_exists: + expected_include = ( + f"{WINDOWS_RUNTIME_INCLUDE_DIR};{EXISTING_WINDOWS_INCLUDE_DIR}" + ) + if libs_exists: + expected_lib = f"{WINDOWS_RUNTIME_LIBS_DIR};{EXISTING_WINDOWS_LIB_DIR}" + assert observed_env == {"INCLUDE": expected_include, "LIB": expected_lib} + assert pip_installer_module.os.environ["INCLUDE"] == EXISTING_WINDOWS_INCLUDE_DIR + assert pip_installer_module.os.environ["LIB"] == EXISTING_WINDOWS_LIB_DIR + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("include_exists", "libs_exists"), + [ + (True, True), + (True, False), + (False, True), + ], +) +async def test_run_pip_in_process_injects_windows_runtime_build_env_without_existing_paths( + configure_run_pip_in_process_capture, include_exists, libs_exists +): + existing_runtime_dirs = set() + if include_exists: + existing_runtime_dirs.add(WINDOWS_RUNTIME_INCLUDE_DIR) + if libs_exists: + existing_runtime_dirs.add(WINDOWS_RUNTIME_LIBS_DIR) + + observed_env = configure_run_pip_in_process_capture( + platform="win32", + packaged_runtime=True, + existing_runtime_dirs=existing_runtime_dirs, + ) + + installer = PipInstaller("") + result = await installer._run_pip_in_process(["install", "demo-package"]) + + assert result == 0 + assert observed_env == { + "INCLUDE": WINDOWS_RUNTIME_INCLUDE_DIR if include_exists else None, + "LIB": WINDOWS_RUNTIME_LIBS_DIR if libs_exists else None, + } + if include_exists: + assert ";" not in observed_env["INCLUDE"] + if libs_exists: + assert ";" not in observed_env["LIB"] + assert "INCLUDE" not in pip_installer_module.os.environ + assert "LIB" not in pip_installer_module.os.environ + + +@pytest.mark.asyncio +async def test_run_pip_in_process_does_not_inject_when_runtime_dirs_missing( + configure_run_pip_in_process_capture, +): + observed_env = configure_run_pip_in_process_capture( + platform="win32", + packaged_runtime=True, + include_value=EXISTING_WINDOWS_INCLUDE_DIR, + lib_value=EXISTING_WINDOWS_LIB_DIR, + existing_runtime_dirs=set(), + ) + + installer = PipInstaller("") + result = await installer._run_pip_in_process(["install", "demo-package"]) + + assert result == 0 + assert observed_env == { + "INCLUDE": EXISTING_WINDOWS_INCLUDE_DIR, + "LIB": EXISTING_WINDOWS_LIB_DIR, + } + assert pip_installer_module.os.environ["INCLUDE"] == EXISTING_WINDOWS_INCLUDE_DIR + assert pip_installer_module.os.environ["LIB"] == EXISTING_WINDOWS_LIB_DIR + + +@pytest.mark.asyncio +async def test_run_pip_in_process_uses_latest_env_when_building_runtime_paths( + monkeypatch, + configure_run_pip_in_process_capture, +): + updated_include = ntpath.join(r"C:\new-toolchain", "include") + updated_lib = ntpath.join(r"C:\new-toolchain", "lib") + observed_env = configure_run_pip_in_process_capture( + platform="win32", + packaged_runtime=True, + include_value=EXISTING_WINDOWS_INCLUDE_DIR, + lib_value=EXISTING_WINDOWS_LIB_DIR, + existing_runtime_dirs={ + WINDOWS_RUNTIME_INCLUDE_DIR, + WINDOWS_RUNTIME_LIBS_DIR, + }, + ) + + async def fake_to_thread(func, *args): + pip_installer_module.os.environ["INCLUDE"] = updated_include + pip_installer_module.os.environ["LIB"] = updated_lib + return func(*args) + + monkeypatch.setattr(pip_installer_module.asyncio, "to_thread", fake_to_thread) + + installer = PipInstaller("") + result = await installer._run_pip_in_process(["install", "demo-package"]) + + assert result == 0 + assert observed_env == { + "INCLUDE": f"{WINDOWS_RUNTIME_INCLUDE_DIR};{updated_include}", + "LIB": f"{WINDOWS_RUNTIME_LIBS_DIR};{updated_lib}", + } + assert pip_installer_module.os.environ["INCLUDE"] == updated_include + assert pip_installer_module.os.environ["LIB"] == updated_lib + + +@pytest.mark.asyncio +async def test_run_pip_in_process_does_not_modify_env_on_non_windows( + configure_run_pip_in_process_capture, +): + existing_include = "/toolchain/include" + existing_lib = "/toolchain/lib" + observed_env = configure_run_pip_in_process_capture( + platform="linux", + packaged_runtime=True, + include_value=existing_include, + lib_value=existing_lib, + ) + + installer = PipInstaller("") + result = await installer._run_pip_in_process(["install", "demo-package"]) + + assert result == 0 + assert observed_env == {"INCLUDE": existing_include, "LIB": existing_lib} + assert pip_installer_module.os.environ["INCLUDE"] == existing_include + assert pip_installer_module.os.environ["LIB"] == existing_lib + + +@pytest.mark.asyncio +async def test_run_pip_in_process_does_not_inject_env_when_not_packaged( + configure_run_pip_in_process_capture, +): + observed_env = configure_run_pip_in_process_capture( + platform="win32", + packaged_runtime=False, + existing_runtime_dirs={ + WINDOWS_RUNTIME_INCLUDE_DIR, + WINDOWS_RUNTIME_LIBS_DIR, + }, + ) + + installer = PipInstaller("") + result = await installer._run_pip_in_process(["install", "demo-package"]) + + assert result == 0 + assert observed_env == {"INCLUDE": None, "LIB": None} + assert "INCLUDE" not in pip_installer_module.os.environ + assert "LIB" not in pip_installer_module.os.environ + + @pytest.mark.asyncio async def test_run_pip_in_process_classifies_nonstandard_conflict_output(monkeypatch): def fake_pip_main(args):