Allow setting env var to turn off traces (#30)

* Allow setting env var to turn off traces

* Refactor disabling traces

Author: tgillam-do

gradient_adk/runtime/helpers.py

@@ -18,6 +18,12 @@
 from gradient_adk.runtime.network_interceptor import setup_digitalocean_interception
 
 
+def _is_tracing_disabled() -> bool:
+    """Check if tracing is globally disabled via DISABLE_TRACES env var."""
+    val = os.environ.get("DISABLE_TRACES", "").lower()
+    return val in ("true", "1", "yes")
+
+
 class InstrumentorProtocol(Protocol):
     """Protocol for instrumentor classes."""
 
@@ -88,6 +94,10 @@ def _ensure_tracker(self) -> Optional[DigitalOceanTracesTracker]:
         self._tracker_initialized = True
 
         try:
+            # Check if tracing is globally disabled
+            if _is_tracing_disabled():
+                return None
+
             api_token = os.environ.get("DIGITALOCEAN_API_TOKEN")
             if not api_token:
                 return None

gradient_adk/runtime/network_interceptor.py

@@ -468,6 +468,12 @@ def get_network_interceptor() -> NetworkInterceptor:
 
 
 def setup_digitalocean_interception() -> None:
+    # Check if tracing is globally disabled
+    # Import here to avoid circular imports
+    from .helpers import _is_tracing_disabled
+    if _is_tracing_disabled():
+        return
+
     intr = get_network_interceptor()
 
     # Add inference (LLM) endpoint patterns

gradient_adk/tracing.py

@@ -41,7 +41,7 @@ async def my_agent(input: dict, context: RequestContext):
 from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar
 
 from .runtime.interfaces import NodeExecution
-from .runtime.helpers import get_tracker
+from .runtime.helpers import get_tracker, _is_tracing_disabled
 from .runtime.network_interceptor import get_network_interceptor
 
 F = TypeVar("F", bound=Callable[..., Any])
@@ -178,6 +178,10 @@ def _trace_base(
     """
 
     def decorator(func: F) -> F:
+        # If tracing is disabled, return the original function unchanged
+        if _is_tracing_disabled():
+            return func
+
         span_name = name or func.__name__
 
         # Handle async generator functions (functions with `yield` that are async)
@@ -535,6 +539,9 @@ def add_llm_span(
             num_output_tokens=5,
         )
     """
+    if _is_tracing_disabled():
+        return
+
     tracker = get_tracker()
     if not tracker:
         return
@@ -613,6 +620,9 @@ def add_tool_span(
             tool_call_id="call_abc123",
         )
     """
+    if _is_tracing_disabled():
+        return
+
     tracker = get_tracker()
     if not tracker:
         return
@@ -669,6 +679,9 @@ def add_agent_span(
             metadata={"model": "gpt-4"},
         )
     """
+    if _is_tracing_disabled():
+        return
+
     tracker = get_tracker()
     if not tracker:
         return

integration_tests/run/test_adk_agents_run.py

@@ -704,3 +704,64 @@ def test_streaming_agent_with_evaluation_id_returns_single_response(
 
         finally:
             cleanup_process(process)
+
+    @pytest.mark.cli
+    def test_agent_run_with_disable_traces(self, setup_agent_in_temp):
+        """
+        Test that an agent can run successfully with DISABLE_TRACES=1.
+        Verifies:
+        - Server starts successfully with tracing disabled
+        - Health endpoint responds
+        - /run endpoint works and echoes input
+        """
+        logger = logging.getLogger(__name__)
+        temp_dir = setup_agent_in_temp
+        port = find_free_port()
+        process = None
+
+        try:
+            logger.info(f"Starting agent with DISABLE_TRACES=1 on port {port}")
+
+            # Set up environment with DISABLE_TRACES=1
+            env = os.environ.copy()
+            env["DISABLE_TRACES"] = "1"
+
+            process = subprocess.Popen(
+                [
+                    "gradient",
+                    "agent",
+                    "run",
+                    "--port",
+                    str(port),
+                    "--no-dev",
+                ],
+                cwd=temp_dir,
+                env=env,
+                start_new_session=True,
+            )
+
+            # Wait for server to be ready
+            server_ready = wait_for_server(port, timeout=30)
+            assert server_ready, "Server did not start within timeout"
+
+            # Test health endpoint
+            health_response = requests.get(f"http://localhost:{port}/health", timeout=5)
+            assert health_response.status_code == 200
+            health_data = health_response.json()
+            assert health_data["status"] == "healthy"
+            logger.info(f"Health check passed with DISABLE_TRACES=1: {health_data}")
+
+            # Test /run endpoint
+            run_response = requests.post(
+                f"http://localhost:{port}/run",
+                json={"prompt": "Hello with tracing disabled!"},
+                timeout=10,
+            )
+            assert run_response.status_code == 200
+            run_data = run_response.json()
+            assert run_data["echo"] == "Hello with tracing disabled!"
+            assert run_data["received"]["prompt"] == "Hello with tracing disabled!"
+            logger.info(f"Agent ran successfully with DISABLE_TRACES=1: {run_data}")
+
+        finally:
+            cleanup_process(process)

tests/runtime/helpers_test.py

@@ -12,6 +12,7 @@
     register_all_instrumentors,
     _register_langgraph,
     _register_pydanticai,
+    _is_tracing_disabled,
 )
 
 
@@ -290,4 +291,51 @@ def test_pydanticai_availability_check():
     check = test_registry._registrations["pydanticai"]["availability_check"]
 
     # Since we're running tests with pydantic-ai installed, it should be available
-    assert check() is True
+    assert check() is True
+
+
+# -----------------------------
+# DISABLE_TRACES Tests
+# -----------------------------
+
+
+def test_is_tracing_disabled_true_values():
+    """Test _is_tracing_disabled returns True for truthy env var values."""
+    for val in ["true", "TRUE", "True", "1", "yes", "YES", "Yes"]:
+        with patch.dict(os.environ, {"DISABLE_TRACES": val}):
+            assert _is_tracing_disabled() is True, f"Expected True for DISABLE_TRACES={val}"
+
+
+def test_is_tracing_disabled_false_values():
+    """Test _is_tracing_disabled returns False for falsy env var values."""
+    for val in ["false", "FALSE", "0", "no", "", "anything", "disabled"]:
+        with patch.dict(os.environ, {"DISABLE_TRACES": val}):
+            assert _is_tracing_disabled() is False, f"Expected False for DISABLE_TRACES={val}"
+
+
+def test_is_tracing_disabled_missing_var():
+    """Test _is_tracing_disabled returns False when env var is not set."""
+    with patch.dict(os.environ, {}, clear=True):
+        assert _is_tracing_disabled() is False
+
+
+def test_ensure_tracker_returns_none_when_tracing_disabled(fresh_registry):
+    """Test that _ensure_tracker returns None when DISABLE_TRACES=1, even with valid API token."""
+    with patch.dict(os.environ, {"DISABLE_TRACES": "1", "DIGITALOCEAN_API_TOKEN": "test-token"}):
+        result = fresh_registry._ensure_tracker()
+    
+    assert result is None
+    assert fresh_registry._tracker is None
+
+
+def test_install_all_returns_none_when_tracing_disabled():
+    """Test that install_all returns None and installs nothing when DISABLE_TRACES=1."""
+    test_registry = InstrumentorRegistry()
+    
+    with patch("gradient_adk.runtime.helpers.registry", test_registry):
+        with patch.dict(os.environ, {"DISABLE_TRACES": "1", "DIGITALOCEAN_API_TOKEN": "test-token"}):
+            register_all_instrumentors()
+            result = test_registry.install_all()
+    
+    assert result is None
+    assert test_registry.get_installed_names() == []