fix: set packaged Windows runtime build env for pip native builds

astrbot/core/utils/pip_installer.py

@@ -5,6 +5,7 @@
 import importlib.util
 import io
 import logging
+import ntpath
 import os
 import re
 import shlex
@@ -30,6 +31,7 @@
 
 _DISTLIB_FINDER_PATCH_ATTEMPTED = False
 _SITE_PACKAGES_IMPORT_LOCK = threading.RLock()
+_PIP_IN_PROCESS_ENV_LOCK = threading.RLock()
 _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 +237,83 @@ def _run_pip_main_streaming(pip_main, args: list[str]) -> tuple[int, list[str]]:
     return result_code, stream.lines
 
 
+def _run_pip_main_with_temporary_environ(
+    pip_main,
+    args: list[str],
+    env_updates: dict[str, str],
+) -> tuple[int, list[str]]:
+    # os.environ is process-wide; serialize temporary mutations around the
+    # in-process pip invocation and keep the mutation window inside the worker.
+    with _PIP_IN_PROCESS_ENV_LOCK:
+        with _temporary_environ(env_updates):
+            return _run_pip_main_streaming(pip_main, args)
+
+
+def _normalize_windows_native_build_path(path: str) -> str:
+    normalized = path.replace("/", "\\")
+
+    for prefix in ("\\\\?\\UNC\\", "\\??\\UNC\\"):
+        if normalized.startswith(prefix):
+            return ntpath.normpath(f"\\\\{normalized[len(prefix) :]}")
+
+    for prefix in ("\\\\?\\", "\\??\\"):
+        if normalized.startswith(prefix):
+            normalized = normalized[len(prefix) :]
+            break
+
+    return ntpath.normpath(normalized)
+
+
+def _prepend_windows_env_path(name: str, path: str) -> str:
+    existing = os.environ.get(name)
+    if existing:
+        return f"{path};{existing}"
+    return path
+
+
+def _build_packaged_windows_runtime_build_env() -> dict[str, str]:
+    if sys.platform != "win32" or not is_packaged_desktop_runtime():
+        return {}
+
+    runtime_executable = _normalize_windows_native_build_path(sys.executable)
+    runtime_dir = ntpath.dirname(runtime_executable)
+    if not runtime_dir:
+        return {}
+
+    env_updates: dict[str, str] = {}
+    include_dir = _normalize_windows_native_build_path(
+        ntpath.join(runtime_dir, "include")
+    )
+    libs_dir = _normalize_windows_native_build_path(ntpath.join(runtime_dir, "libs"))
+
+    if os.path.isdir(include_dir):
+        env_updates["INCLUDE"] = _prepend_windows_env_path("INCLUDE", include_dir)
+    if os.path.isdir(libs_dir):
+        env_updates["LIB"] = _prepend_windows_env_path("LIB", libs_dir)
+
+    return env_updates
+
+
+@contextlib.contextmanager
+def _temporary_environ(updates: dict[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:
+                os.environ[key] = previous_value
+
+
 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)
@@ -927,11 +1006,15 @@ def prefer_installed_dependencies(self, requirements_path: str) -> None:
     async def _run_pip_in_process(self, args: list[str]) -> int:
         pip_main = _get_pip_main()
         _patch_distlib_finder_for_frozen_runtime()
+        build_env = _build_packaged_windows_runtime_build_env()
 
         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,
+                build_env,
             )
         finally:
             _cleanup_added_root_handlers(original_handlers)

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,56 @@ 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.mark.asyncio
 async def test_install_targets_site_packages_for_desktop_client(monkeypatch, tmp_path):
     monkeypatch.setenv("ASTRBOT_DESKTOP_CLIENT", "1")
@@ -226,6 +286,145 @@ def fake_pip_main(args):
     ]
 
 
+@pytest.mark.parametrize(
+    ("path", "expected"),
+    [
+        (
+            WINDOWS_PACKAGED_RUNTIME_EXECUTABLE,
+            WINDOWS_RUNTIME_EXECUTABLE,
+        ),
+        (
+            r"\\?\UNC\server\share\include",
+            r"\\server\share\include",
+        ),
+        (
+            r"\??\UNC\server\share\libs",
+            r"\\server\share\libs",
+        ),
+        (
+            "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
+
+
+@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(
+    monkeypatch, 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(
+        monkeypatch,
+        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
+async def test_run_pip_in_process_injects_windows_runtime_build_env_without_existing_paths(
+    monkeypatch,
+):
+    observed_env = _configure_run_pip_in_process_capture(
+        monkeypatch,
+        platform="win32",
+        packaged_runtime=True,
+        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": WINDOWS_RUNTIME_INCLUDE_DIR,
+        "LIB": WINDOWS_RUNTIME_LIBS_DIR,
+    }
+    assert ";" not in observed_env["INCLUDE"]
+    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_modify_env_on_non_windows(monkeypatch):
+    existing_include = "/toolchain/include"
+    existing_lib = "/toolchain/lib"
+    observed_env = _configure_run_pip_in_process_capture(
+        monkeypatch,
+        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(monkeypatch):
+    observed_env = _configure_run_pip_in_process_capture(
+        monkeypatch,
+        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):