fix: build telemetry providers lazily so comfy install can upgrade pydantic_core (BE-3289)

comfy_cli/tracking.py

@@ -6,17 +6,19 @@
 import logging as logginglib
 import os
 import sys
+import threading
 import uuid
-from typing import Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol
 
 import typer
-from mixpanel import Mixpanel
-from posthog import Posthog
 
 from comfy_cli import constants, logging, ui
 from comfy_cli.config_manager import ConfigManager
 from comfy_cli.workspace_manager import WorkspaceManager
 
+if TYPE_CHECKING:
+    from posthog import Posthog
+
 # Ignore logs from urllib3 that Mixpanel/PostHog use.
 logginglib.getLogger("urllib3").setLevel(logginglib.ERROR)
 
@@ -139,7 +141,12 @@ def flush(self) -> None: ...
 
 class MixpanelProvider:
     def __init__(self, token: str):
-        self.client = Mixpanel(token) if token else None
+        self.client = None
+        if token:
+            # Imported here, not at module scope: see `_get_providers`.
+            from mixpanel import Mixpanel
+
+            self.client = Mixpanel(token)
         self.enabled = self.client is not None
 
     def track(self, event_name: str, distinct_id: str | None, properties: dict[str, Any]) -> None:
@@ -165,6 +172,9 @@ def __init__(self, token: str, host: str):
         self.enabled = False
         if not token:
             return
+        # Imported here, not at module scope: see `_get_providers`.
+        from posthog import Posthog
+
         # disable_geoip=False lets PostHog enrich events with IP-derived location.
         self.client = Posthog(project_api_key=token, host=host, disable_geoip=False)
         self.enabled = True
@@ -186,10 +196,47 @@ def flush(self) -> None:
         self.client.flush()
 
 
-PROVIDERS: list[TelemetryProvider] = [
-    MixpanelProvider(MIXPANEL_TOKEN),
-    PostHogProvider(POSTHOG_TOKEN, POSTHOG_HOST),
-]
+# Built on the first send, not at import. `mixpanel` pulls in pydantic (and thus
+# the compiled pydantic_core extension) for its feature-flags module, and importing
+# it here meant every `comfy` process held that .pyd open for its whole lifetime.
+# On Windows a loaded DLL cannot be replaced, so `comfy install` — which shells out
+# to `uv pip install` against its own interpreter's environment — could not upgrade
+# pydantic_core and died with "Access is denied (os error 5)", leaving the package
+# half-removed and every later `comfy` invocation unable to import it.
+# Deferring the import means a run that never sends telemetry (no consent, or
+# DO_NOT_TRACK) never loads the extension at all. See tests/comfy_cli/test_tracking_lazy_import.py.
+PROVIDERS: list[TelemetryProvider] | None = None
+
+# Building at import used to be single-shot courtesy of the import lock; keep that
+# guarantee now that it happens on demand. Racing builds would strand a PostHog
+# client — and its unflushed queue — in the list that lost.
+_PROVIDERS_LOCK = threading.Lock()
+
+
+def _get_providers() -> list[TelemetryProvider]:
+    global PROVIDERS
+    if PROVIDERS is None:
+        with _PROVIDERS_LOCK:
+            if PROVIDERS is None:
+                # Construction runs the deferred SDK imports, so it can fail on exactly
+                # the broken dependency tree this deferral exists to avoid (a half-removed
+                # pydantic_core raises ImportError). Telemetry is best-effort: degrade to a
+                # no-op rather than take the user's command down with us. Each provider is
+                # built independently so one unusable SDK doesn't silence the other, and the
+                # result is cached either way — including [] — so a doomed import isn't
+                # retried on every later event.
+                built = []
+                for name, factory in (
+                    ("MixpanelProvider", lambda: MixpanelProvider(MIXPANEL_TOKEN)),
+                    ("PostHogProvider", lambda: PostHogProvider(POSTHOG_TOKEN, POSTHOG_HOST)),
+                ):
+                    try:
+                        built.append(factory())
+                    except Exception as e:  # noqa: BLE001
+                        logging.warning(f"Failed to initialize telemetry provider {name}, skipping it: {e}")
+                PROVIDERS = built
+    return PROVIDERS
+
 
 app = typer.Typer()
 
@@ -215,7 +262,7 @@ def _dispatch(
     passive telemetry, env-only for feedback).
     """
     properties = {**properties, "cli_version": cli_version, "tracing_id": tracing_id}
-    for provider in PROVIDERS:
+    for provider in _get_providers():
         provider_event_name = (
             mixpanel_name if (mixpanel_name is not None and isinstance(provider, MixpanelProvider)) else event_name
         )
@@ -408,7 +455,18 @@ def init_tracking(enable_tracking: bool):
 
 
 def _flush_all_providers() -> None:
-    for provider in PROVIDERS:
+    # Deliberately reads PROVIDERS rather than calling _get_providers(): a run that
+    # never sent anything has nothing to drain, and constructing providers from an
+    # atexit hook would import the SDKs we just went to the trouble of deferring.
+    # Taking the lock (without building) makes an in-flight build resolve first —
+    # otherwise we could read None mid-construction and exit without draining the
+    # racing thread's PostHog queue. Released before flushing so a slow network
+    # drain doesn't block a concurrent send.
+    with _PROVIDERS_LOCK:
+        providers = PROVIDERS
+    if providers is None:
+        return
+    for provider in providers:
         try:
             provider.flush()
         except Exception as e:  # noqa: BLE001

tests/comfy_cli/test_tracking_lazy_import.py

@@ -0,0 +1,139 @@
+"""The telemetry SDKs must not be imported just by starting the CLI (BE-3289).
+
+`mixpanel` depends on pydantic, which loads the compiled `pydantic_core`
+extension. A `comfy` process that has that extension loaded holds the .pyd/.so
+open for its whole lifetime, and on Windows an open DLL cannot be replaced — so
+`comfy install`, which shells out to `uv pip install` against its own
+interpreter's environment, could not upgrade pydantic_core:
+
+    error: failed to remove file `...\\pydantic_core\\_pydantic_core.cp312-win_amd64.pyd`:
+    Access is denied. (os error 5)
+
+uv exited 2 having already unlinked the distribution's metadata, so every later
+`comfy` invocation died with `ImportError: cannot import name '__version__' from
+'pydantic_core' (unknown location)`.
+
+These run in subprocesses on purpose: the assertion is about what a *fresh*
+interpreter loads, and the pytest process has already imported plenty on its own.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+# The SDKs themselves plus the transitive dependency that actually did the damage.
+_FORBIDDEN = ("mixpanel", "posthog", "pydantic", "pydantic_core")
+
+
+def _modules_after(source: str) -> set[str]:
+    """Run *source* in a fresh interpreter; return its top-level sys.modules names."""
+    script = f"{source}\nimport sys, json\nprint(json.dumps(sorted(m.split('.')[0] for m in sys.modules)))"
+    proc = subprocess.run(
+        [sys.executable, "-c", script],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    import json
+
+    return set(json.loads(proc.stdout.strip().splitlines()[-1]))
+
+
+def test_importing_tracking_does_not_load_telemetry_sdks():
+    loaded = _modules_after("import comfy_cli.tracking")
+    assert not loaded & set(_FORBIDDEN), (
+        f"comfy_cli.tracking imported {sorted(loaded & set(_FORBIDDEN))} at module scope; "
+        "these must stay lazy so `comfy install` can upgrade pydantic_core on Windows"
+    )
+
+
+def test_importing_the_cli_entrypoint_does_not_load_telemetry_sdks():
+    # cmdline is what `comfy` actually runs, and it is the import chain in the
+    # BE-3289 traceback (cmdline -> tracking -> mixpanel -> pydantic).
+    loaded = _modules_after("import comfy_cli.cmdline")
+    assert not loaded & set(_FORBIDDEN), (
+        f"comfy_cli.cmdline imported {sorted(loaded & set(_FORBIDDEN))} at module scope; "
+        "these must stay lazy so `comfy install` can upgrade pydantic_core on Windows"
+    )
+
+
+def test_atexit_flush_does_not_construct_providers():
+    """The atexit hook must not be the thing that imports the SDKs at exit."""
+    loaded = _modules_after("import comfy_cli.tracking as t\nt._flush_all_providers()")
+    assert not loaded & set(_FORBIDDEN)
+
+
+def test_providers_are_built_on_first_dispatch():
+    """The deferral must be a deferral, not a removal — telemetry still works."""
+    loaded = _modules_after(
+        "import comfy_cli.tracking as t\n"
+        "assert t.PROVIDERS is None, 'providers built before any send'\n"
+        "providers = t._get_providers()\n"
+        "assert t.PROVIDERS is providers, 'accessor did not cache onto the module global'\n"
+        "assert len(providers) == 2, f'expected mixpanel+posthog, got {providers!r}'\n"
+        "assert t._get_providers() is providers, 'accessor rebuilt providers on second call'\n"
+    )
+    # Constructing them is precisely what pulls the SDKs in.
+    assert {"mixpanel", "posthog"} <= loaded
+
+
+def test_broken_sdk_import_degrades_instead_of_crashing_the_command():
+    """Deferring the import moves it onto the user's command path, so the failure it
+    used to hit at startup can now surface mid-run. A half-removed pydantic_core is
+    exactly the state this change exists to prevent, but a user can still land in it
+    from an earlier bad install — telemetry must go quiet, not take the command down.
+    """
+    _modules_after(
+        "import builtins\n"
+        "_real = builtins.__import__\n"
+        "def boom(name, *a, **k):\n"
+        "    if name in ('mixpanel', 'posthog'):\n"
+        "        raise ImportError(\"cannot import name '__version__' from 'pydantic_core'\")\n"
+        "    return _real(name, *a, **k)\n"
+        "builtins.__import__ = boom\n"
+        "import comfy_cli.tracking as t\n"
+        "t.MIXPANEL_TOKEN = 'tok'\n"
+        "t.POSTHOG_TOKEN = 'tok'\n"
+        # The send path must survive: _get_providers() is called in _dispatch's loop
+        # header, outside the per-provider try/except that only guards .track().
+        "t._dispatch('evt', {}, distinct_id='abc')\n"
+        "assert t._get_providers() == [], f'expected no providers, got {t._get_providers()!r}'\n"
+        # Cached, so a doomed import isn't retried on every later event.
+        "assert t.PROVIDERS == []\n"
+        "t._flush_all_providers()\n"
+    )
+
+
+def test_one_broken_sdk_does_not_silence_the_other():
+    """Providers are built independently: an unusable PostHog SDK must not throw away
+    a perfectly good Mixpanel provider (and vice versa)."""
+    _modules_after(
+        "import builtins\n"
+        "_real = builtins.__import__\n"
+        "def boom(name, *a, **k):\n"
+        "    if name == 'posthog':\n"
+        "        raise ImportError('posthog is broken')\n"
+        "    return _real(name, *a, **k)\n"
+        "builtins.__import__ = boom\n"
+        "import comfy_cli.tracking as t\n"
+        "t.MIXPANEL_TOKEN = 'tok'\n"
+        "t.POSTHOG_TOKEN = 'tok'\n"
+        "providers = t._get_providers()\n"
+        "names = [type(p).__name__ for p in providers]\n"
+        "assert names == ['MixpanelProvider'], f'lost the healthy provider: {names!r}'\n"
+    )
+
+
+def test_concurrent_first_send_builds_one_set_of_providers():
+    """Module-level init was single-shot via the import lock; the lazy accessor
+    has to keep that. A racing build would strand a PostHog client, and its
+    unflushed event queue, in the list that lost."""
+    _modules_after(
+        "import comfy_cli.tracking as t\n"
+        "from concurrent.futures import ThreadPoolExecutor\n"
+        "with ThreadPoolExecutor(max_workers=8) as ex:\n"
+        "    got = [f.result() for f in [ex.submit(t._get_providers) for _ in range(8)]]\n"
+        "assert all(g is got[0] for g in got), 'threads saw different provider lists'\n"
+        "assert t.PROVIDERS is got[0]\n"
+    )

tests/e2e/verify_tracking_live.py

@@ -62,7 +62,10 @@ def _send_smoketest_event(nonce: str) -> tuple[str, str]:
     if tracking._telemetry_disabled_by_env():
         _die("telemetry is opted out via DO_NOT_TRACK / COMFY_NO_TELEMETRY — unset it to verify")
 
-    posthog_providers = [p for p in tracking.PROVIDERS if isinstance(p, tracking.PostHogProvider) and p.enabled]
+    # _get_providers(), not PROVIDERS: providers are built on first send, so the
+    # module global is still None until something asks for them.
+    providers = tracking._get_providers()
+    posthog_providers = [p for p in providers if isinstance(p, tracking.PostHogProvider) and p.enabled]
     if not posthog_providers:
         _die("no enabled PostHog provider — check POSTHOG_API_KEY token")