From 2181579586f82cc094c740c0dfecce3a7b0d7e3e Mon Sep 17 00:00:00 2001
From: Amjith Ramanujam <amjith@openai.com>
Date: Thu, 8 May 2025 16:29:51 -0700
Subject: [PATCH] Refactor click instrumentation.

---
 cli_telemetry/instrumentation/__init__.py | 21 ++++++++++++
 cli_telemetry/instrumentation/click.py    | 40 +++++++++++++++++++++++
 cli_telemetry/telemetry.py                | 39 ++--------------------
 3 files changed, 63 insertions(+), 37 deletions(-)
 create mode 100644 cli_telemetry/instrumentation/__init__.py
 create mode 100644 cli_telemetry/instrumentation/click.py

diff --git a/cli_telemetry/instrumentation/__init__.py b/cli_telemetry/instrumentation/__init__.py
new file mode 100644
index 0000000..a340e2b
--- /dev/null
+++ b/cli_telemetry/instrumentation/__init__.py
@@ -0,0 +1,21 @@
+"""
+Automatic instrumentation registry for supported libraries.
+"""
+import os
+
+from .click import auto_instrument_click
+
+def init_auto_instrumentation() -> None:
+    """
+    Initialize automatic instrumentation for all supported libraries.
+    Future instrumentation modules should be added here.
+    """
+    # Click instrumentation
+    if "CLI_TELEMETRY_DISABLE_CLICK_INSTRUMENTATION" in os.environ:
+        return
+    try:
+        auto_instrument_click()
+    except Exception:
+        # Safely ignore instrumentation errors
+        pass
+
diff --git a/cli_telemetry/instrumentation/click.py b/cli_telemetry/instrumentation/click.py
new file mode 100644
index 0000000..5de8e70
--- /dev/null
+++ b/cli_telemetry/instrumentation/click.py
@@ -0,0 +1,40 @@
+"""
+Instrumentation for Click commands to auto-wrap their invocation in telemetry spans.
+"""
+
+from ..telemetry import Span, add_tag
+
+def auto_instrument_click():
+    """Monkeypatch click.Command and click.Group to auto-wrap in telemetry spans."""
+    import os
+
+    if os.environ.get("CLI_TELEMETRY_DISABLE_CLICK_PATCH") == "1":
+        return  # User has opted out of monkeypatching
+
+    try:
+        import click
+
+        def telemetry_wrapper(original_invoke):
+            def invoke_with_span(self, ctx):
+                # Assumes init_telemetry() already called by app
+                with Span(self.name, attributes={"cli.command": ctx.command_path}) as span:
+                    try:
+                        for param in self.params:
+                            if param.name in ctx.params:
+                                add_tag(f"args.{param.name}", ctx.params[param.name])
+                    except Exception:
+                        pass
+                    return original_invoke(self, ctx)
+            return invoke_with_span
+
+        # Avoid double-patching
+        if not getattr(click.Command, "_telemetry_patched", False):
+            click.Command.invoke = telemetry_wrapper(click.Command.invoke)
+            click.Command._telemetry_patched = True
+
+        if not getattr(click.Group, "_telemetry_patched", False):
+            click.Group.invoke = telemetry_wrapper(click.Group.invoke)
+            click.Group._telemetry_patched = True
+
+    except ImportError:
+        pass  # No click? No telemetry sadness.
diff --git a/cli_telemetry/telemetry.py b/cli_telemetry/telemetry.py
index 9b53f06..99ce7b3 100644
--- a/cli_telemetry/telemetry.py
+++ b/cli_telemetry/telemetry.py
@@ -69,6 +69,7 @@ def _init_db_file(db_file: str) -> None:
 
 
 def init_telemetry(service_name: str, db_path: Optional[str] = None, user_id_file: Optional[str] = None) -> None:
+    from .instrumentation import init_auto_instrumentation
     """
     Initialize trace ID, user‐ID file, and SQLite DB.
     If db_path/user_id_file are provided, uses those; otherwise defaults to:
@@ -81,9 +82,7 @@ def init_telemetry(service_name: str, db_path: Optional[str] = None, user_id_fil
         if _initialized:
             return
 
-        # monkeypatch click if allowed
-        if os.environ.get("CLI_TELEMETRY_DISABLE_CLICK_PATCH"):
-            _monkeypatch_click()
+        init_auto_instrumentation()
 
         # determine base path under XDG_DATA_HOME
         xdg = os.environ.get("XDG_DATA_HOME", os.path.expanduser("~/.local/share"))
@@ -272,37 +271,3 @@ def end_session() -> None:
             pass
 
 
-def _monkeypatch_click():
-    """Monkeypatch click.Command and click.Group to auto-wrap in telemetry spans."""
-    import os
-
-    if os.environ.get("CLI_TELEMETRY_DISABLE_CLICK_PATCH") == "1":
-        return  # User has opted out of monkeypatching
-
-    try:
-        import click
-
-        def telemetry_wrapper(original_invoke):
-            def invoke_with_span(self, ctx):
-                # Assumes init_telemetry() already called by app
-                with Span(self.name, attributes={"cli.command": ctx.command_path}) as span:
-                    try:
-                        for param in self.params:
-                            if param.name in ctx.params:
-                                add_tag(f"args.{param.name}", ctx.params[param.name])
-                    except Exception:
-                        pass
-                    return original_invoke(self, ctx)
-            return invoke_with_span
-
-        # Avoid double-patching
-        if not getattr(click.Command, "_telemetry_patched", False):
-            click.Command.invoke = telemetry_wrapper(click.Command.invoke)
-            click.Command._telemetry_patched = True
-
-        if not getattr(click.Group, "_telemetry_patched", False):
-            click.Group.invoke = telemetry_wrapper(click.Group.invoke)
-            click.Group._telemetry_patched = True
-
-    except ImportError:
-        pass  # No click? No telemetry sadness.