fix: gate feature flag flushing and cover wrapper contract

Respect DD_FEATURE_FLAGS_ENABLED during import and cleanup, test the real wrapper contract with released dependencies, and distinguish expected skips in debug logs.

Environment: Datadog workspace

Co-Authored-By: Codex GPT-6 <noreply@localhost>

Author: vjfridge

datadog_lambda/config.py

@@ -85,6 +85,7 @@ def _resolve_env(self, key, default=None, cast=None, depends_on_tracing=False):
     enhanced_metrics_enabled = _get_env("DD_ENHANCED_METRICS", "true", as_bool)
 
     flush_in_thread = _get_env("DD_FLUSH_IN_THREAD", "false", as_bool)
+    feature_flags_enabled = _get_env("DD_FEATURE_FLAGS_ENABLED", "true", as_bool)
     feature_flag_flush_timeout_ms = _get_env(
         "DD_FEATURE_FLAGS_FLUSH_TIMEOUT_MS", 500.0, float
     )

datadog_lambda/feature_flags.py

@@ -17,14 +17,19 @@
 
 def flush_feature_flag_telemetry(context):
     try:
-        if _telemetry is None or getattr(_telemetry, "FLUSH_API_VERSION", None) != 1:
+        if _telemetry is None:
+            return
+        if getattr(_telemetry, "FLUSH_API_VERSION", None) != 1:
+            logger.debug("Skipping feature flag telemetry flush: unsupported tracer")
             return
         cap_ms = config.feature_flag_flush_timeout_ms
         if not math.isfinite(cap_ms) or cap_ms <= 0:
             return
         # No remaining-time information means no safe invocation budget.
-        remaining_ms = context.get_remaining_time_in_millis()
-        budget_ms = min(cap_ms, remaining_ms - SAFETY_MARGIN_MS)
+        get_remaining_ms = getattr(context, "get_remaining_time_in_millis", None)
+        if get_remaining_ms is None:
+            return
+        budget_ms = min(cap_ms, get_remaining_ms() - SAFETY_MARGIN_MS)
         if budget_ms <= 0:
             logger.debug(
                 "Skipping feature flag telemetry flush: no invocation time remains"

datadog_lambda/wrapper.py

@@ -9,7 +9,6 @@
 from importlib import import_module
 from time import time_ns
 
-from datadog_lambda.feature_flags import flush_feature_flag_telemetry
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import (
     set_cold_start,
@@ -78,6 +77,10 @@
 if config.llmobs_enabled:
     from ddtrace.llmobs import LLMObs
 
+flush_feature_flag_telemetry = None
+if config.feature_flags_enabled:
+    from datadog_lambda.feature_flags import flush_feature_flag_telemetry
+
 if config.exception_replay_enabled:
     from ddtrace.debugging._exception.replay import SpanExceptionHandler
 
@@ -212,7 +215,11 @@ def __call__(self, event, context, **kwargs):
             except Exception as e:
                 logger.error(format_err_with_traceback(e))
             finally:
-                flush_feature_flag_telemetry(context)
+                if (
+                    config.feature_flags_enabled
+                    and flush_feature_flag_telemetry is not None
+                ):
+                    flush_feature_flag_telemetry(context)
         if self.blocking_response:
             return self.blocking_response
         return self.response

tests/test_config.py

@@ -101,6 +101,7 @@ def _test_as_list(env_key, conf_key, default):
     *_test_as_bool("DD_MERGE_XRAY_TRACES", "merge_xray_traces", default=False),
     *_test_as_bool("DD_PROFILING_ENABLED", "profiling_enabled", default=False),
     *_test_as_bool("DD_LLMOBS_ENABLED", "llmobs_enabled", default=False),
+    *_test_as_bool("DD_FEATURE_FLAGS_ENABLED", "feature_flags_enabled", default=True),
     *_test_as_bool(
         "DD_EXCEPTION_REPLAY_ENABLED", "exception_replay_enabled", default=False
     ),

tests/test_feature_flags.py

@@ -1,5 +1,8 @@
 import importlib.util
 import os
+import subprocess
+import sys
+import textwrap
 import unittest
 from unittest.mock import Mock, patch
 
@@ -92,6 +95,59 @@ def test_missing_context_and_telemetry_errors_are_safe(self):
             )
             telemetry.flush.assert_called_once()
 
+    def test_missing_remaining_time_skips_without_failure_log(self):
+        for context in (None, object()):
+            with self.subTest(context=context), patch.object(
+                feature_flags, "_telemetry"
+            ) as telemetry, patch.object(feature_flags, "logger") as log:
+                telemetry.FLUSH_API_VERSION = 1
+                feature_flags.flush_feature_flag_telemetry(context)
+                telemetry.flush.assert_not_called()
+                log.debug.assert_not_called()
+
+    def test_unsupported_tracer_skip_is_logged(self):
+        with patch.object(
+            feature_flags, "_telemetry", Mock(FLUSH_API_VERSION=2)
+        ) as telemetry, patch.object(feature_flags, "logger") as log:
+            feature_flags.flush_feature_flag_telemetry(
+                Mock(get_remaining_time_in_millis=lambda: 1000)
+            )
+            telemetry.flush.assert_not_called()
+            log.debug.assert_called_once_with(
+                "Skipping feature flag telemetry flush: unsupported tracer"
+            )
+
+    def test_disabled_feature_flags_do_not_import_flush_modules(self):
+        code = textwrap.dedent("""
+            import sys
+
+            class BlockFlushImports:
+                def find_spec(self, fullname, *args):
+                    if fullname in (
+                        "datadog_lambda.feature_flags",
+                        "ddtrace.internal.openfeature._telemetry",
+                    ):
+                        raise AssertionError("unexpected import: " + fullname)
+
+            sys.meta_path.insert(0, BlockFlushImports())
+            from datadog_lambda import wrapper
+            assert not wrapper.config.feature_flags_enabled
+            assert "datadog_lambda.feature_flags" not in sys.modules
+            assert "ddtrace.internal.openfeature._telemetry" not in sys.modules
+            """)
+        result = subprocess.run(
+            [sys.executable, "-c", code],
+            env={
+                **os.environ,
+                "DD_FEATURE_FLAGS_ENABLED": "false",
+                "DD_INSTRUMENTATION_TELEMETRY_ENABLED": "false",
+            },
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
+
     def test_incomplete_diagnostic(self):
         with patch.object(feature_flags, "_telemetry") as telemetry, patch.object(
             feature_flags, "logger"

tests/test_wrapper.py

@@ -4,6 +4,7 @@
 import unittest
 import importlib
 import sys
+import types
 import pytest
 
 from unittest.mock import MagicMock, patch, call, ANY
@@ -77,6 +78,58 @@ def setUp(self):
         self.mock_dd_lambda_layer_tag = patcher.start()
         self.addCleanup(patcher.stop)
 
+    def test_real_decorator_calls_feature_flag_flush_contract(self):
+        import ddtrace.internal.openfeature as openfeature_pkg
+        from datadog_lambda import feature_flags
+
+        calls = []
+        fake = types.ModuleType("ddtrace.internal.openfeature._telemetry")
+        fake.FLUSH_API_VERSION = 1
+        fake.flush = lambda timeout: calls.append(timeout) or "completed"
+        self.addCleanup(importlib.reload, feature_flags)
+        with patch.dict(sys.modules, {fake.__name__: fake}), patch.object(
+            openfeature_pkg, "_telemetry", fake, create=True
+        ), patch.object(type(config), "feature_flags_enabled", True):
+            importlib.reload(feature_flags)
+            context = get_mock_context()
+            context.get_remaining_time_in_millis = lambda: 30000
+
+            handler = wrapper.datadog_lambda_wrapper(
+                lambda event, context: {"ok": True}
+            )
+            self.assertEqual(handler({}, context), {"ok": True})
+
+            error = ValueError("handler error")
+
+            def raising(event, context):
+                raise error
+
+            with self.assertRaises(ValueError) as caught:
+                wrapper.datadog_lambda_wrapper(raising)({}, context)
+            self.assertIs(caught.exception, error)
+        self.assertEqual(calls, [0.5, 0.5])
+
+    def test_disabled_feature_flags_skip_flush(self):
+        with patch.object(type(config), "feature_flags_enabled", False), patch.object(
+            wrapper, "flush_feature_flag_telemetry"
+        ) as flush:
+            context = get_mock_context()
+            context.get_remaining_time_in_millis = lambda: 30000
+            handler = wrapper.datadog_lambda_wrapper(
+                lambda event, context: {"ok": True}
+            )
+            self.assertEqual(handler({}, context), {"ok": True})
+
+            error = ValueError("handler error")
+
+            def raising(event, context):
+                raise error
+
+            with self.assertRaises(ValueError) as caught:
+                wrapper.datadog_lambda_wrapper(raising)({}, context)
+            self.assertIs(caught.exception, error)
+            flush.assert_not_called()
+
     @patch("datadog_lambda.config.Config.trace_enabled", False)
     def test_datadog_lambda_wrapper(self):
         @wrapper.datadog_lambda_wrapper