fix: broaden telemetry kwarg redaction and _ctx filter atomically (BE-992)

comfy_cli/tracking.py

@@ -1,4 +1,5 @@
 import functools
+import json
 import logging as logginglib
 import sys
 import uuid
@@ -16,9 +17,55 @@
 MIXPANEL_TOKEN = "93aeab8962b622d431ac19800ccc9f67"
 mp = Mixpanel(MIXPANEL_TOKEN) if MIXPANEL_TOKEN else None
 
-# Kwargs whose values must never reach tracking system.
-# The key is kept (with a redacted marker) so we can still see whether the option was supplied.
-SENSITIVE_TRACKING_KEYS = frozenset({"api_key"})
+# ---------------------------------------------------------------------------
+# Sensitive-kwarg redaction and kwarg filtering for track_command()
+#
+# Historical context (BE-992): the original SENSITIVE_TRACKING_KEYS only
+# matched "api_key" by exact name, missing kwargs such as
+# set_civitai_api_token and set_hf_api_token.  Separately, the kwarg filter
+# dropped "ctx" and "context" but not "_ctx" (underscore-prefixed Click
+# Context).  These two gaps *cancelled each other out*: the un-redacted
+# credential payload always included the non-JSON-serializable _ctx object,
+# so mp.track() raised a TypeError before the event reached Mixpanel.
+#
+# Fixing the underscore-prefix filter without also broadening the credential
+# allowlist would silently activate a CWE-200 credential leak.  Both filters
+# are therefore fixed atomically here.
+# ---------------------------------------------------------------------------
+
+_SENSITIVE_SUFFIXES = ("_token", "_api_key", "_secret", "_password")
+_SENSITIVE_EXACT = frozenset({"api_key", "token", "password", "secret"})
+
+
+def _is_sensitive(name: str) -> bool:
+    """Return True if *name* looks like it holds a credential or secret.
+
+    Matched case-insensitively so future kwargs added with non-PEP-8 casing
+    (``API_KEY``, ``apiKey``, ``AuthToken``, etc.) don't bypass redaction.
+    """
+    lower = name.lower()
+    return lower in _SENSITIVE_EXACT or lower.endswith(_SENSITIVE_SUFFIXES)
+
+
+def _is_trackable(name: str, value: object) -> bool:
+    """Return True if the kwarg (name, value) pair is safe to include in a tracking event.
+
+    Excluded:
+    * ``ctx`` / ``context`` — Click context objects passed by Typer.
+    * Any name starting with ``_`` — covers ``_ctx`` and other private/internal params.
+    * Values that are not JSON-serializable — defensive guard so a single
+      unserializable kwarg doesn't silently swallow the entire event.
+    """
+    if name in ("ctx", "context"):
+        return False
+    if name.startswith("_"):
+        return False
+    try:
+        json.dumps(value)
+    except (TypeError, ValueError, OverflowError, RecursionError):
+        return False
+    return True
+
 
 # Generate a unique tracing ID per command.
 config_manager = ConfigManager()
@@ -79,12 +126,10 @@ def decorator(func):
         def wrapper(*args, **kwargs):
             command_name = f"{sub_command}:{func.__name__}" if sub_command is not None else func.__name__
 
-            # Copy kwargs to avoid mutating original dictionary
-            # Remove context and ctx from the dictionary as they are not needed for tracking and not serializable.
             filtered_kwargs = {
-                k: ("<redacted>" if v is not None else None) if k in SENSITIVE_TRACKING_KEYS else v
+                k: ("<redacted>" if v is not None else None) if _is_sensitive(k) else v
                 for k, v in kwargs.items()
-                if k != "ctx" and k != "context"
+                if _is_trackable(k, v)
             }
 
             logging.debug(f"Tracking command: {command_name} with arguments: {filtered_kwargs}")

tests/comfy_cli/test_tracking.py

@@ -84,9 +84,6 @@ def some_cmd(workflow, api_key=None):
         assert "sk-supersecret" not in str(props)
 
     def test_api_key_none_stays_none(self, tracking_module):
-        # When the user didn't pass --api-key (or set $COMFY_API_KEY), we still
-        # want to be able to see in the analytics that it was absent — not a
-        # "<redacted>" sentinel that would imply they did pass one.
         tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True")
 
         @tracking_module.track_command()
@@ -98,6 +95,84 @@ def some_cmd(workflow, api_key=None):
         _, kwargs = tracking_module.mp.track.call_args
         assert kwargs["properties"]["api_key"] is None
 
+    def test_set_civitai_api_token_is_redacted(self, tracking_module):
+        tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True")
+
+        @tracking_module.track_command("model")
+        def download(url, set_civitai_api_token=None, set_hf_api_token=None):
+            return None
+
+        download(url="https://example.com", set_civitai_api_token="civ-real-token")
+
+        tracking_module.mp.track.assert_called_once()
+        _, kwargs = tracking_module.mp.track.call_args
+        props = kwargs["properties"]
+        assert props["set_civitai_api_token"] == "<redacted>"
+        assert "civ-real-token" not in str(props)
+
+    def test_set_hf_api_token_is_redacted(self, tracking_module):
+        tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True")
+
+        @tracking_module.track_command("model")
+        def download(url, set_civitai_api_token=None, set_hf_api_token=None):
+            return None
+
+        download(url="https://example.com", set_hf_api_token="hf_real-token")
+
+        tracking_module.mp.track.assert_called_once()
+        _, kwargs = tracking_module.mp.track.call_args
+        props = kwargs["properties"]
+        assert props["set_hf_api_token"] == "<redacted>"
+        assert "hf_real-token" not in str(props)
+
+    def test_bare_token_kwarg_is_redacted(self, tracking_module):
+        tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True")
+
+        @tracking_module.track_command()
+        def some_cmd(workflow, token=None):
+            return None
+
+        some_cmd(workflow="wf.json", token="my-secret-token")
+
+        tracking_module.mp.track.assert_called_once()
+        _, kwargs = tracking_module.mp.track.call_args
+        props = kwargs["properties"]
+        assert props["token"] == "<redacted>"
+        assert "my-secret-token" not in str(props)
+
+    def test_underscore_ctx_is_excluded(self, tracking_module):
+        import click
+
+        tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True")
+
+        @tracking_module.track_command("model")
+        def download(_ctx, url, set_civitai_api_token=None):
+            return None
+
+        ctx = click.Context(click.Command("download"))
+        download(_ctx=ctx, url="https://example.com")
+
+        tracking_module.mp.track.assert_called_once()
+        _, kwargs = tracking_module.mp.track.call_args
+        props = kwargs["properties"]
+        assert "_ctx" not in props
+        assert props["url"] == "https://example.com"
+
+    def test_non_serializable_value_is_excluded(self, tracking_module):
+        tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True")
+
+        @tracking_module.track_command()
+        def some_cmd(workflow, callback=None):
+            return None
+
+        some_cmd(workflow="wf.json", callback=lambda x: x)
+
+        tracking_module.mp.track.assert_called_once()
+        _, kwargs = tracking_module.mp.track.call_args
+        props = kwargs["properties"]
+        assert "callback" not in props
+        assert props["workflow"] == "wf.json"
+
 
 class TestInitTrackingRoundTrip:
     """End-to-end: init_tracking() writes the string "False"/"True", and track_event honors it.