diff --git a/comfy_cli/tracking.py b/comfy_cli/tracking.py index 580bceb1..4865213a 100644 --- a/comfy_cli/tracking.py +++ b/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 diff --git a/tests/comfy_cli/test_tracking_lazy_import.py b/tests/comfy_cli/test_tracking_lazy_import.py new file mode 100644 index 00000000..5ab235e7 --- /dev/null +++ b/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" + ) diff --git a/tests/e2e/verify_tracking_live.py b/tests/e2e/verify_tracking_live.py index 2928ab8a..ae0146a9 100644 --- a/tests/e2e/verify_tracking_live.py +++ b/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")