feat(sdk): allow disabling metrics programmatically

packages/traceloop-sdk/tests/test_sdk_initialization.py

@@ -7,6 +7,7 @@
 from traceloop.sdk import Traceloop
 from traceloop.sdk.config import is_content_tracing_enabled
 from traceloop.sdk.decorators import workflow
+from traceloop.sdk.instruments import Instruments
 from traceloop.sdk.tracing.tracing import TracerWrapper
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor, BatchSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
@@ -324,6 +325,111 @@ def test_passing_both_kwargs_raises_type_error(isolated_tracer_wrapper):
             use_legacy_attributes=True,
         )
 
+
+@pytest.mark.parametrize(
+    ("metrics_enabled", "environment_value", "expected_metrics_enabled"),
+    [
+        (None, None, True),
+        (None, "true", True),
+        (None, "false", False),
+        (True, "false", True),
+        (False, "true", False),
+    ],
+)
+def test_metrics_enabled_precedence(
+    monkeypatch,
+    metrics_enabled,
+    environment_value,
+    expected_metrics_enabled,
+):
+    """An explicit init value overrides the environment; None preserves it."""
+    if environment_value is None:
+        monkeypatch.delenv("TRACELOOP_METRICS_ENABLED", raising=False)
+    else:
+        monkeypatch.setenv("TRACELOOP_METRICS_ENABLED", environment_value)
+
+    metrics_exporter = object()
+    with (
+        patch("traceloop.sdk.TracerWrapper"),
+        patch("traceloop.sdk.ImageUploader"),
+        patch("traceloop.sdk.MetricsWrapper") as metrics_wrapper,
+    ):
+        Traceloop.init(
+            exporter=InMemorySpanExporter(),
+            metrics_exporter=metrics_exporter,
+            metrics_enabled=metrics_enabled,
+        )
+
+    assert metrics_wrapper.called is expected_metrics_enabled
+    if expected_metrics_enabled:
+        metrics_wrapper.assert_called_once_with(exporter=metrics_exporter)
+
+
+@pytest.mark.parametrize(
+    ("metrics_enabled", "environment_value"),
+    [(False, "true"), (True, "false")],
+)
+def test_metrics_enabled_is_visible_during_instrumentation(
+    monkeypatch,
+    isolated_tracer_wrapper,
+    metrics_enabled,
+    environment_value,
+):
+    """Instrumentors must observe the explicit override while they initialize."""
+    from opentelemetry.instrumentation.openai.utils import is_metrics_enabled
+
+    monkeypatch.setenv("TRACELOOP_METRICS_ENABLED", environment_value)
+    observed_metrics_enabled = []
+    metrics_exporter = object()
+
+    def observe_metrics_enabled(*args, **kwargs):
+        observed_metrics_enabled.append(is_metrics_enabled())
+        return True
+
+    with (
+        patch(
+            "traceloop.sdk.tracing.tracing.init_openai_instrumentor",
+            side_effect=observe_metrics_enabled,
+        ),
+        patch("traceloop.sdk.MetricsWrapper") as metrics_wrapper,
+    ):
+        Traceloop.init(
+            exporter=InMemorySpanExporter(),
+            metrics_exporter=metrics_exporter,
+            metrics_enabled=metrics_enabled,
+            instruments={Instruments.OPENAI},
+            disable_batch=True,
+        )
+
+    assert observed_metrics_enabled == [metrics_enabled]
+    assert metrics_wrapper.called is metrics_enabled
+    if metrics_enabled:
+        metrics_wrapper.assert_called_once_with(exporter=metrics_exporter)
+    assert os.environ["TRACELOOP_METRICS_ENABLED"] == environment_value
+
+
+@pytest.mark.parametrize("environment_value", [None, "true"])
+def test_metrics_enabled_restores_environment_when_tracer_initialization_fails(
+    monkeypatch,
+    environment_value,
+):
+    """The process environment must be restored even if tracer setup fails."""
+    if environment_value is None:
+        monkeypatch.delenv("TRACELOOP_METRICS_ENABLED", raising=False)
+    else:
+        monkeypatch.setenv("TRACELOOP_METRICS_ENABLED", environment_value)
+
+    with patch("traceloop.sdk.TracerWrapper", side_effect=RuntimeError("boom")):
+        with pytest.raises(RuntimeError, match="boom"):
+            Traceloop.init(
+                exporter=InMemorySpanExporter(),
+                metrics_enabled=False,
+                disable_batch=True,
+            )
+
+    assert os.environ.get("TRACELOOP_METRICS_ENABLED") == environment_value
+
+
 def test_use_attributes_defaults_to_true(isolated_tracer_wrapper):
     """When use_attributes is not passed, instrumentors keep the default spec-compliant
     behavior (emit prompts/completions as span attributes)."""

packages/traceloop-sdk/traceloop/sdk/__init__.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import threading
 import warnings
 from pathlib import Path
 
@@ -35,6 +36,10 @@
 from traceloop.sdk.associations.associations import AssociationProperty as AssociationProperty
 
 
+# Serializes the temporary TRACELOOP_METRICS_ENABLED override in Traceloop.init()
+_metrics_enabled_lock = threading.Lock()
+
+
 class Traceloop:
     AUTO_CREATED_KEY_PATH = str(
         Path.home() / ".cache" / "traceloop" / "auto_created_key"
@@ -74,6 +79,7 @@ def init(
         use_attributes: Optional[bool] = None,
         use_legacy_attributes: Optional[bool] = None,
         trace_content: Optional[bool] = None,
+        metrics_enabled: Optional[bool] = None,
     ) -> Optional[Client]:
         """Initialize Traceloop tracing, metrics, and instrumentation.
 
@@ -107,6 +113,11 @@ def init(
                 when ``enabled=False`` or when tracing is disabled via
                 ``TRACELOOP_TRACING_ENABLED`` — both checks happen before this
                 override runs.
+            metrics_enabled: Enables or disables metric exporting. An explicit value
+                takes precedence over ``TRACELOOP_METRICS_ENABLED``. If ``None``, the
+                environment variable is used and metrics default to enabled. Like the
+                environment variable, it only applies when metrics are first set up, so
+                a later ``init()`` call does not stop metrics that are already running.
         """
         if use_attributes is not None and use_legacy_attributes is not None:
             raise TypeError(
@@ -205,22 +216,36 @@ def init(
         TracerWrapper.set_static_params(
             resource_attributes, enable_content_tracing, api_endpoint, headers
         )
-        Traceloop.__tracer_wrapper = TracerWrapper(
-            disable_batch=disable_batch,
-            processor=processor,
-            propagator=propagator,
-            exporter=exporter,
-            sampler=sampler,
-            should_enrich_metrics=should_enrich_metrics,
-            image_uploader=image_uploader or ImageUploader(api_endpoint, api_key),
-            instruments=instruments,
-            block_instruments=block_instruments,
-            span_postprocess_callback=span_postprocess_callback,
-            use_attributes=use_attributes,
-        )
-
-        metrics_disabled_by_config = not is_metrics_enabled()
+        with _metrics_enabled_lock:
+            metrics_enabled_by_config = is_metrics_enabled() if metrics_enabled is None else metrics_enabled
+            previous_metrics_enabled = os.environ.get("TRACELOOP_METRICS_ENABLED")
+            if metrics_enabled is not None:
+                os.environ["TRACELOOP_METRICS_ENABLED"] = str(metrics_enabled).lower()
+            try:
+                Traceloop.__tracer_wrapper = TracerWrapper(
+                    disable_batch=disable_batch,
+                    processor=processor,
+                    propagator=propagator,
+                    exporter=exporter,
+                    sampler=sampler,
+                    should_enrich_metrics=should_enrich_metrics,
+                    image_uploader=image_uploader or ImageUploader(api_endpoint, api_key),
+                    instruments=instruments,
+                    block_instruments=block_instruments,
+                    span_postprocess_callback=span_postprocess_callback,
+                    use_attributes=use_attributes,
+                )
+            finally:
+                if metrics_enabled is not None:
+                    if previous_metrics_enabled is None:
+                        os.environ.pop("TRACELOOP_METRICS_ENABLED", None)
+                    else:
+                        os.environ["TRACELOOP_METRICS_ENABLED"] = previous_metrics_enabled
+
+        metrics_disabled_by_config = not metrics_enabled_by_config
         has_custom_spans_pipeline = processor or exporter
+        # A custom trace pipeline still needs a custom metrics exporter, regardless
+        # of metrics_enabled, to avoid sending metrics to an unintended endpoint.
         custom_trace_without_custom_metrics = has_custom_spans_pipeline and not metrics_exporter
 
         if metrics_disabled_by_config or custom_trace_without_custom_metrics: