More fixes

Author: jy-tan

drift/core/adaptive_sampling.py

@@ -81,7 +81,7 @@ def __init__(
         self._admission_multiplier = 1.0
         self._state: AdaptiveSamplingState = "fixed" if config.mode == "fixed" else "healthy"
         self._paused_until_s = 0.0
-        self._last_updated_at_s = 0.0
+        self._last_updated_at_s: float | None = None
         self._last_decrease_at_s = 0.0
 
         self._prev_dropped_span_count = 0
@@ -99,7 +99,7 @@ def update(self, snapshot: AdaptiveSamplingHealthSnapshot) -> None:
                 return
 
             now_s = self._now_fn()
-            elapsed_s = 2.0 if self._last_updated_at_s == 0 else max(0.001, now_s - self._last_updated_at_s)
+            elapsed_s = 2.0 if self._last_updated_at_s is None else max(0.001, now_s - self._last_updated_at_s)
             self._last_updated_at_s = now_s
 
             decay = math.exp(-(elapsed_s * 1000.0) / 30000.0)

drift/core/mode_utils.py

@@ -180,7 +180,8 @@ def should_record_inbound_http_request(
     Returns:
         Tuple of (should_record, skip_reason):
         - should_record: True if request should be recorded
-        - skip_reason: If False, explains why ("dropped" or "not_sampled"), None otherwise
+        - skip_reason: If False, explains why ("dropped", "not_sampled",
+          "load_shed", or "critical_pause"), None otherwise
     """
     if transform_engine and transform_engine.should_drop_inbound_request(method, target, headers):
         return False, "dropped"
@@ -191,6 +192,6 @@ def should_record_inbound_http_request(
         sdk = TuskDrift.get_instance()
         decision = sdk.should_record_root_request(is_pre_app_start=is_pre_app_start)
         if not decision.should_record:
-            return False, "not_sampled"
+            return False, decision.reason
 
     return True, None

drift/instrumentation/wsgi/handler.py

@@ -250,7 +250,8 @@ def _create_and_handle_request(
         )
         if not should_record:
             logger.debug(f"[WSGI] Skipping request ({skip_reason}), path={path}")
-            response = original_wsgi_app(app, environ, start_response)
+            with suppress_recording():
+                response = original_wsgi_app(app, environ, start_response)
             return SuppressedResponseIterable(response)
 
     # Capture request body

tests/unit/test_adaptive_sampling.py

@@ -1,3 +1,4 @@
+import math
 import threading
 
 from drift.core.adaptive_sampling import (
@@ -44,6 +45,22 @@ def test_controller_load_sheds_and_pauses_on_drops():
     assert paused_decision.reason == "critical_pause"
 
 
+def test_elapsed_time_uses_zero_timestamp_as_real_value():
+    now = {"value": 0.0}
+    controller = AdaptiveSamplingController(
+        ResolvedSamplingConfig(mode="adaptive", base_rate=0.5, min_rate=0.1),
+        random_fn=lambda: 0.0,
+        now_fn=lambda: now["value"],
+    )
+
+    controller.update(AdaptiveSamplingHealthSnapshot(export_failure_count=1))
+    now["value"] = 0.5
+    controller.update(AdaptiveSamplingHealthSnapshot(export_failure_count=1))
+
+    expected_decay = math.exp(-(0.5 * 1000.0) / 30000.0)
+    assert math.isclose(controller._recent_failure_signal, expected_decay, rel_tol=1e-6)
+
+
 def test_get_decision_waits_for_controller_lock():
     controller = AdaptiveSamplingController(
         ResolvedSamplingConfig(mode="adaptive", base_rate=0.5, min_rate=0.1),

tests/unit/test_mode_utils.py

@@ -360,6 +360,7 @@ def test_returns_false_when_not_sampled(self, mocker):
         mock_drift = mocker.patch("drift.core.drift_sdk.TuskDrift")
         mock_sdk = mocker.MagicMock()
         mock_sdk.should_record_root_request.return_value.should_record = False
+        mock_sdk.should_record_root_request.return_value.reason = "not_sampled"
         mock_drift.get_instance.return_value = mock_sdk
 
         result, reason = should_record_inbound_http_request(
@@ -373,6 +374,25 @@ def test_returns_false_when_not_sampled(self, mocker):
         assert result is False
         assert reason == "not_sampled"
 
+    def test_returns_controller_reason_when_adaptive_sampling_skips(self, mocker):
+        """Should preserve adaptive controller reasons for debug logging."""
+        mock_drift = mocker.patch("drift.core.drift_sdk.TuskDrift")
+        mock_sdk = mocker.MagicMock()
+        mock_sdk.should_record_root_request.return_value.should_record = False
+        mock_sdk.should_record_root_request.return_value.reason = "critical_pause"
+        mock_drift.get_instance.return_value = mock_sdk
+
+        result, reason = should_record_inbound_http_request(
+            method="GET",
+            target="/api/users",
+            headers={},
+            transform_engine=None,
+            is_pre_app_start=False,
+        )
+
+        assert result is False
+        assert reason == "critical_pause"
+
     def test_drop_check_happens_before_sampling(self, mocker):
         """Should check drop rules before sampling."""
         mock_transform = mocker.MagicMock()

tests/unit/test_wsgi_handler.py

@@ -44,6 +44,9 @@ def test_skipped_wsgi_request_keeps_suppression_during_streaming_iteration_and_c
     )
 
     def original_wsgi_app(_app: Any, _environ: dict[str, Any], _start_response: Any) -> Iterable[bytes]:
+        from drift.core.no_recording import is_recording_suppressed
+
+        observed.append(("call", is_recording_suppressed()))
         return response
 
     def app(_environ: dict[str, Any], _start_response: Any) -> Iterable[bytes]:
@@ -72,6 +75,7 @@ def app(_environ: dict[str, Any], _start_response: Any) -> Iterable[bytes]:
     close_method()
 
     assert observed == [
+        ("call", True),
         ("iter", True),
         ("next", True),
         ("next", True),