fix(telemetry): bound network I/O so telemetry can't outlive the run envelope (BE-3403)

comfy_cli/tracking.py

@@ -6,10 +6,13 @@
 import logging as logginglib
 import os
 import sys
+import threading
+import time
 import uuid
 from typing import Any, Protocol
 
 import typer
+from mixpanel import Consumer as MixpanelConsumer
 from mixpanel import Mixpanel
 from posthog import Posthog
 
@@ -139,7 +142,15 @@ def flush(self) -> None: ...
 
 class MixpanelProvider:
     def __init__(self, token: str):
-        self.client = Mixpanel(token) if token else None
+        # mixpanel-python's default Consumer uses request_timeout=None → an
+        # unbounded, synchronous requests.post on the main thread, so a
+        # blackholed telemetry endpoint (accepts TCP, never responds) hangs the
+        # CLI indefinitely (BE-3354/BE-3403). track() sends inline on the calling
+        # thread and flush() is a no-op, so this bound is the ONLY thing guarding
+        # the hot event path — it isn't covered by the atexit daemon deadline.
+        # retry_limit=1 (default is 4 with backoff) keeps a blackholed send to a
+        # single ~10s attempt instead of ~40s+ across retries.
+        self.client = Mixpanel(token, consumer=MixpanelConsumer(request_timeout=10, retry_limit=1)) if token else None
         self.enabled = self.client is not None
 
     def track(self, event_name: str, distinct_id: str | None, properties: dict[str, Any]) -> None:
@@ -166,7 +177,18 @@ def __init__(self, token: str, host: str):
         if not token:
             return
         # disable_geoip=False lets PostHog enrich events with IP-derived location.
-        self.client = Posthog(project_api_key=token, host=host, disable_geoip=False)
+        # max_retries/timeout tighten the consumer drain budget from the posthog
+        # 7.x defaults (3 × 15s ≈ 50s worst case) to ~21s, so the atexit flush
+        # can't linger on a blackholed endpoint after the terminal envelope is
+        # already on stdout (BE-3354/BE-3403).
+        self.client = Posthog(project_api_key=token, host=host, disable_geoip=False, max_retries=1, timeout=10)
+        # Posthog's constructor registers its own atexit.register(self.join),
+        # which runs self.join() synchronously on the main thread at shutdown —
+        # independently of _flush_all_providers and NOT bounded by its daemon
+        # deadline. Against a blackholed endpoint that join can still block ~21s
+        # after the terminal envelope is on stdout, defeating this change. Drop it
+        # so our bounded flush is the only shutdown drain path (BE-3403).
+        atexit.unregister(self.client.join)
         self.enabled = True
 
     def track(self, event_name: str, distinct_id: str | None, properties: dict[str, Any]) -> None:
@@ -407,12 +429,44 @@ def init_tracking(enable_tracking: bool):
         track_event("install")
 
 
+def _flush_one(provider: TelemetryProvider) -> None:
+    try:
+        provider.flush()
+    except Exception as e:  # noqa: BLE001
+        logging.warning(f"Failed to flush telemetry provider {type(provider).__name__}: {e}")
+
+
+_FLUSH_DEADLINE_SECONDS = 5.0
+
+
 def _flush_all_providers() -> None:
+    # Telemetry is best-effort by contract: a blackholed endpoint (accepts TCP,
+    # never responds) must never let this atexit hook wedge every consumer of the
+    # CLI's stdout after the terminal envelope is already emitted (BE-3329/BE-3403).
+    # Start every provider's flush in a daemon thread, then join them all against
+    # a SINGLE shared deadline so total exit delay stays ~5s regardless of how
+    # many providers there are (a per-provider join would make it 5s × N).
+    # Dropping in-flight events beats hanging the process. t.start()/t.join() are
+    # wrapped defensively so nothing this hook does can raise and print a
+    # traceback to stderr after the terminal envelope.
+    deadline = time.monotonic() + _FLUSH_DEADLINE_SECONDS
+    threads: list[tuple[threading.Thread, TelemetryProvider]] = []
     for provider in PROVIDERS:
+        t = threading.Thread(target=_flush_one, args=(provider,), daemon=True)
+        try:
+            t.start()
+        except RuntimeError as e:  # e.g. a thread-creation race during shutdown
+            logging.warning(f"could not start telemetry flush for {type(provider).__name__}: {e}")
+            continue
+        threads.append((t, provider))
+    for t, provider in threads:
         try:
-            provider.flush()
-        except Exception as e:  # noqa: BLE001
-            logging.warning(f"Failed to flush telemetry provider {type(provider).__name__}: {e}")
+            t.join(timeout=max(0.0, deadline - time.monotonic()))
+        except RuntimeError as e:  # pragma: no cover - defensive
+            logging.warning(f"telemetry flush join failed for {type(provider).__name__}: {e}")
+            continue
+        if t.is_alive():
+            logging.warning(f"telemetry flush timed out for {type(provider).__name__}; dropping in-flight events")
 
 
 atexit.register(_flush_all_providers)

tests/comfy_cli/test_tracking_providers.py

@@ -6,6 +6,7 @@
 ``workflow_run_id`` on the canonical execution lifecycle events.
 """
 
+import threading
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -216,6 +217,42 @@ def test_mixpanel_with_empty_token_is_disabled(self):
         provider = MixpanelProvider("")
         assert provider.enabled is False
 
+    def test_mixpanel_client_has_bounded_request_timeout(self):
+        """Regression guard for BE-3354/BE-3403: mixpanel-python's default
+        Consumer uses request_timeout=None → an unbounded requests.post that
+        hangs the CLI forever on a blackholed endpoint. The provider must build
+        its Mixpanel client with an explicit 10s consumer timeout so this can't
+        silently regress to None. retry_limit must also be 1: the Consumer
+        default is 4 with backoff, which would let a blackholed send run ~40s+
+        across retries even with a 10s per-attempt timeout — and this send is
+        synchronous on the main thread, not covered by the atexit deadline."""
+        provider = MixpanelProvider("token-mp")
+        assert provider.client is not None
+        consumer = provider.client._consumer
+        assert consumer._request_timeout == 10
+        # retry_limit is threaded into the session's urllib3 Retry.total.
+        adapter = consumer._session.get_adapter("https://api.mixpanel.com")
+        assert adapter.max_retries.total == 1
+
+    def test_posthog_unregisters_its_own_atexit_join(self):
+        """Regression guard for BE-3403: Posthog's constructor registers its own
+        ``atexit.register(self.join)``, which flushes synchronously on the main
+        thread at shutdown, unbounded by ``_flush_all_providers``' 5s daemon
+        deadline. Against a blackholed endpoint that join can block ~21s after
+        the terminal envelope. The provider must unregister it so the bounded
+        flush is the only shutdown drain path."""
+        import comfy_cli.tracking as tracking_mod
+
+        fake_client = MagicMock()
+        with (
+            patch.object(tracking_mod, "Posthog", return_value=fake_client),
+            patch.object(tracking_mod, "atexit") as fake_atexit,
+        ):
+            provider = PostHogProvider("phc_test", "https://t.comfy.org")
+
+        assert provider.enabled is True
+        fake_atexit.unregister.assert_called_once_with(fake_client.join)
+
     def test_posthog_track_skips_when_distinct_id_is_none(self, tracking_with_two_providers):
         tracking_mod, _, ph_provider = tracking_with_two_providers
         with patch.object(tracking_mod, "user_id", None):
@@ -296,3 +333,35 @@ def test_flush_swallows_provider_errors(self):
             tracking_mod._flush_all_providers()
 
         p2.flush.assert_called_once()
+
+    def test_flush_returns_before_deadline_when_a_provider_hangs(self):
+        """A provider whose flush() blocks (e.g. a blackholed telemetry endpoint)
+        must not wedge the atexit hook past its per-provider deadline. The hook
+        runs each flush in a daemon thread and joins with a ~5s timeout, so a
+        60s-hanging provider is abandoned rather than allowed to hang the CLI
+        (BE-3354/BE-3403). Bounds the total at well under the 60s hang."""
+        import time
+
+        import comfy_cli.tracking as tracking_mod
+
+        release = threading.Event()
+
+        class _HangingProvider:
+            def flush(self):
+                # Blocks until the test tears down; the deadline must fire first.
+                release.wait(timeout=60)
+
+        fast = MagicMock()
+        try:
+            with patch.object(tracking_mod, "PROVIDERS", [_HangingProvider(), fast]):
+                start = time.monotonic()
+                tracking_mod._flush_all_providers()
+                elapsed = time.monotonic() - start
+
+            # Deadline is 5s per provider; allow generous slack but stay far
+            # under the 60s the hanging provider would otherwise cost.
+            assert elapsed < 8.0, f"flush did not honor the deadline (took {elapsed:.1f}s)"
+            # The healthy provider after the hanging one is still drained.
+            fast.flush.assert_called_once()
+        finally:
+            release.set()