urllib instrumentation + span stauts error updates

Author: sohankshirsagar

drift/core/content_type_utils.py

@@ -17,7 +17,7 @@
     "application/vnd.api+json": DecodedType.JSON,
     # Plain Text (ALLOWED)
     "text/plain": DecodedType.PLAIN_TEXT,
-    # HTML (BLOCKED)
+    # HTML
     "text/html": DecodedType.HTML,
     "application/xhtml+xml": DecodedType.HTML,
     # CSS (BLOCKED)
@@ -113,7 +113,7 @@
 
 # Only JSON and plain text are acceptable (matches Node SDK)
 # All other content types will cause trace blocking
-ACCEPTABLE_DECODED_TYPES = {DecodedType.JSON, DecodedType.PLAIN_TEXT}
+ACCEPTABLE_DECODED_TYPES = {DecodedType.JSON, DecodedType.PLAIN_TEXT, DecodedType.HTML}
 
 
 def get_decoded_type(content_type: str | None) -> DecodedType | None:

drift/core/drift_sdk.py

@@ -406,6 +406,16 @@ def _init_auto_instrumentations(self) -> None:
         except ImportError:
             pass
 
+        try:
+            import urllib.request
+
+            from ..instrumentation.urllib import UrllibInstrumentation
+
+            _ = UrllibInstrumentation()
+            logger.debug("urllib instrumentation initialized")
+        except ImportError:
+            pass
+
         # Initialize PostgreSQL instrumentation before Django
         # Instrument BOTH psycopg2 and psycopg if available
         # This allows apps to use either or both

drift/core/mock_utils.py

@@ -183,7 +183,9 @@ def find_mock_response_sync(
         mock_response = sdk.request_mock_sync(mock_request)
 
         if not mock_response or not mock_response.found:
-            logger.debug(f"No matching mock found for {trace_id} with input value: {input_value}")
+            logger.debug(
+                f"No matching mock found for {trace_id} with input value: {input_value}, input schema: {input_schema_merges}, input schema hash: {outbound_span.input_schema_hash}, input value hash: {outbound_span.input_value_hash}"
+            )
             return None
 
         logger.debug(f"Found mock response for {trace_id}")

drift/core/trace_blocking_manager.py

@@ -204,7 +204,7 @@ def should_block_span(span: CleanSpanData) -> bool:
     """Check if a span should be blocked due to size or server error status.
 
     Blocks the trace if:
-    1. The span is a SERVER span with ERROR status (e.g., HTTP >= 300)
+    1. The span is a SERVER span with ERROR status (e.g., HTTP >= 400)
     2. The span exceeds the maximum size limit (1MB)
 
     This matches Node SDK behavior in TdSpanExporter.ts.
@@ -221,7 +221,7 @@ def should_block_span(span: CleanSpanData) -> bool:
     span_name = span.name
     blocking_manager = TraceBlockingManager.get_instance()
 
-    # Check 1: Block SERVER spans with ERROR status (e.g., HTTP >= 300)
+    # Check 1: Block SERVER spans with ERROR status (e.g., HTTP >= 400)
     if span.kind == SpanKind.SERVER and span.status.code == StatusCode.ERROR:
         logger.debug(f"Blocking trace {trace_id} - server span '{span_name}' has error status")
         blocking_manager.block_trace(trace_id, reason="server_error")

drift/instrumentation/django/middleware.py

@@ -390,8 +390,8 @@ def dict_to_schema_merges(merges_dict):
         duration_seconds = duration_ns // 1_000_000_000
         duration_nanos = duration_ns % 1_000_000_000
 
-        # Match Node SDK: >= 300 is considered an error (redirects, client errors, server errors)
-        if status_code >= 300:
+        # Match Node SDK: >= 400 is considered an error
+        if status_code >= 400:
             status = SpanStatus(code=StatusCode.ERROR, message=f"HTTP {status_code}")
         else:
             status = SpanStatus(code=StatusCode.OK, message="")

drift/instrumentation/fastapi/instrumentation.py

@@ -530,8 +530,8 @@ def _finalize_span(
     TuskDrift.get_instance()
 
     status_code = response_data.get("status_code", 200)
-    # Match Node SDK: >= 300 is considered an error (redirects, client errors, server errors)
-    if status_code >= 300:
+    # Match Node SDK: >= 400 is considered an error
+    if status_code >= 400:
         span_info.span.set_status(Status(OTelStatusCode.ERROR, f"HTTP {status_code}"))
     else:
         span_info.span.set_status(Status(OTelStatusCode.OK))

drift/instrumentation/pyjwt/instrumentation.py

@@ -22,9 +22,7 @@
 class PyJWTInstrumentation(InstrumentationBase):
     """Patches PyJWT to disable verification in REPLAY mode."""
 
-    def __init__(
-        self, mode: TuskDriftMode = TuskDriftMode.DISABLED, enabled: bool = True
-    ) -> None:
+    def __init__(self, mode: TuskDriftMode = TuskDriftMode.DISABLED, enabled: bool = True) -> None:
         self.mode = mode
         should_enable = enabled and mode == TuskDriftMode.REPLAY
 

drift/instrumentation/urllib/__init__.py

@@ -0,0 +1,5 @@
+"""urllib.request instrumentation module."""
+
+from .instrumentation import UrllibInstrumentation
+
+__all__ = ["UrllibInstrumentation"]

drift/instrumentation/urllib/e2e-tests/.tusk/config.yaml

@@ -0,0 +1,27 @@
+version: 1
+
+service:
+  id: "urllib-e2e-test-id"
+  name: "urllib-e2e-test"
+  port: 8000
+  start:
+    command: "python src/app.py"
+  readiness_check:
+    command: "curl -f http://localhost:8000/health"
+    timeout: 45s
+    interval: 5s
+
+tusk_api:
+  url: "http://localhost:8000"
+
+test_execution:
+  concurrent_limit: 10
+  batch_size: 10
+  timeout: 30s
+
+recording:
+  sampling_rate: 1.0
+  export_spans: false
+
+replay:
+  enable_telemetry: false

drift/instrumentation/urllib/e2e-tests/Dockerfile

@@ -0,0 +1,21 @@
+FROM python-e2e-base:latest
+
+# Copy SDK source for editable install
+COPY . /sdk
+
+# Copy test files
+COPY drift/instrumentation/urllib/e2e-tests /app
+
+WORKDIR /app
+
+# Install dependencies (requirements.txt uses -e /sdk for SDK)
+RUN pip install -q -r requirements.txt
+
+# Make entrypoint executable
+RUN chmod +x entrypoint.py
+
+# Create .tusk directories
+RUN mkdir -p /app/.tusk/traces /app/.tusk/logs
+
+# Run entrypoint
+ENTRYPOINT ["python", "entrypoint.py"]