fix redirects

Author: sohankshirsagar

drift/instrumentation/httpx/e2e-tests/src/test_requests.py

@@ -100,28 +100,16 @@ def make_request(method, endpoint, **kwargs):
     # Async sequential chained requests
     make_request("GET", "/api/async/chain")
 
-    # ==========================================================================
-    # Bug Hunting Tests - Confirmed Instrumentation Bugs
-    # These tests expose bugs in the httpx instrumentation
-    # ==========================================================================
-    print("\n--- Bug Hunting Tests (Confirmed Bugs) ---\n")
-
-    # BUG: Streaming response - triggers unpatched socket calls during replay
     make_request("GET", "/test/streaming")
 
-    # BUG: Top-level stream - triggers unpatched socket calls during replay
     make_request("GET", "/test/toplevel-stream")
 
-    # BUG: Multipart file upload - triggers unpatched socket calls during replay
     make_request("POST", "/test/multipart-files")
 
-    # BUG: Follow redirects - replay returns wrong final_url and redirect_count
-    make_request("GET", "/test/follow-redirects")
-
-    # BUG: AsyncClient.send() - triggers unpatched socket calls during replay
     make_request("GET", "/test/async-send")
 
-    # BUG: Async streaming - triggers unpatched socket calls during replay
     make_request("GET", "/test/async-stream")
 
+    make_request("GET", "/test/follow-redirects")
+
     print("\nAll requests completed successfully")

drift/instrumentation/httpx/instrumentation.py

@@ -1229,17 +1229,38 @@ def _create_mock_response(self, httpx_module: Any, mock_data: dict[str, Any], me
         else:
             content = json.dumps(body).encode("utf-8")
 
+        # Determine final URL - use from mock data if present (for redirect handling)
+        final_url = mock_data.get("finalUrl", url)
+
+        # Build history list if redirects were recorded
+        history = []
+        history_count = int(mock_data.get("historyCount", 0))
+        if history_count > 0:
+            # Create minimal placeholder Response objects for history
+            # These represent the intermediate redirect responses
+            for i in range(history_count):
+                redirect_request = httpx_module.Request(
+                    method.upper(), url if i == 0 else f"redirect_{i}"
+                )
+                redirect_response = httpx_module.Response(
+                    status_code=302,  # Standard redirect status
+                    content=b"",
+                    request=redirect_request,
+                )
+                history.append(redirect_response)
+
         # Create httpx.Response
-        # httpx.Response requires a request object
-        request = httpx_module.Request(method.upper(), url)
+        # httpx.Response requires a request object - use final URL so response.url is correct
+        request = httpx_module.Request(method.upper(), final_url)
         response = httpx_module.Response(
             status_code=status_code,
             headers=headers,
             content=content,
             request=request,
+            history=history,
         )
 
-        logger.debug(f"Created mock httpx response: {status_code} for {url}")
+        logger.debug(f"Created mock httpx response: {status_code} for {final_url}")
         return response
 
     def _finalize_span_from_request(
@@ -1331,6 +1352,14 @@ def _finalize_span_from_request(
                     output_value["body"] = response_body_base64
                     output_value["bodySize"] = response_body_size
 
+                # Capture redirect information for proper replay
+                final_url = str(response.url)
+                if final_url != url:  # Only store if redirects occurred
+                    output_value["finalUrl"] = final_url
+
+                if response.history:
+                    output_value["historyCount"] = len(response.history)
+
                 if response.status_code >= 400:
                     status = SpanStatus(
                         code=StatusCode.ERROR,
@@ -1515,6 +1544,14 @@ def _finalize_span(
                     output_value["body"] = response_body_base64
                     output_value["bodySize"] = response_body_size
 
+                # Capture redirect information for proper replay
+                final_url = str(response.url)
+                if final_url != url:  # Only store if redirects occurred
+                    output_value["finalUrl"] = final_url
+
+                if response.history:
+                    output_value["historyCount"] = len(response.history)
+
                 if response.status_code >= 400:
                     status = SpanStatus(
                         code=StatusCode.ERROR,