fix(http): isolate parent span per request, extend user agent span over transport, and record telemetry errors

packages/http/httpx/kiota_http/middleware/__init__.py

@@ -1,7 +1,7 @@
 from .async_kiota_transport import AsyncKiotaTransport
 from .body_inspection_handler import BodyInspectionHandler
 from .headers_inspection_handler import HeadersInspectionHandler
-from .middleware import REQUEST_OPTIONS_KEY, BaseMiddleware, MiddlewarePipeline
+from .middleware import PARENT_SPAN_KEY, REQUEST_OPTIONS_KEY, BaseMiddleware, MiddlewarePipeline
 from .parameters_name_decoding_handler import ParametersNameDecodingHandler
 from .redirect_handler import RedirectHandler
 from .retry_handler import RetryHandler
@@ -12,6 +12,7 @@
     "AsyncKiotaTransport",
     "BodyInspectionHandler",
     "HeadersInspectionHandler",
+    "PARENT_SPAN_KEY",
     "REQUEST_OPTIONS_KEY",
     "BaseMiddleware",
     "MiddlewarePipeline",

packages/http/httpx/kiota_http/middleware/middleware.py

@@ -11,6 +11,7 @@
 tracer = trace.get_tracer(ObservabilityOptions.get_tracer_instrumentation_name(), VERSION)
 
 REQUEST_OPTIONS_KEY = "kiota_request_options"
+PARENT_SPAN_KEY = "parent_span"
 
 
 class MiddlewarePipeline():
@@ -37,7 +38,11 @@ async def send(self, request: httpx.Request) -> httpx.Response:
         if self._first_middleware is not None:
             return await self._first_middleware.send(request, self._transport)
         # No middleware in pipeline, remove the request options and send the request
-        request.extensions.pop(REQUEST_OPTIONS_KEY, None)
+        if hasattr(request, "extensions") and isinstance(request.extensions, dict):
+            options = request.extensions.get(REQUEST_OPTIONS_KEY)
+            if options and (req_parent_span := options.get("parent_span", None)):
+                request.extensions[PARENT_SPAN_KEY] = req_parent_span
+            request.extensions.pop(REQUEST_OPTIONS_KEY, None)
         return await self._transport.handle_async_request(request)
 
 
@@ -56,6 +61,9 @@ async def send(
         if self.next is None:
             # Remove request options if there's no other middleware in the chain.
             if hasattr(request, "extensions") and isinstance(request.extensions, dict):
+                options = request.extensions.get(REQUEST_OPTIONS_KEY)
+                if options and (req_parent_span := options.get("parent_span", None)):
+                    request.extensions[PARENT_SPAN_KEY] = req_parent_span
                 request.extensions.pop(REQUEST_OPTIONS_KEY, None)
             response = await transport.handle_async_request(request)
             response.request = request
@@ -69,8 +77,15 @@ def _create_observability_span(self, request: httpx.Request, span_name: str) ->
         options = None
         if hasattr(request, "extensions") and isinstance(request.extensions, dict):
             options = request.extensions.get(REQUEST_OPTIONS_KEY)
-        if options and (parent_span := options.get("parent_span", None)):
-            self.parent_span = parent_span
-        if self.parent_span is None:
+        parent_span = self.parent_span
+        if options and (req_parent_span := options.get("parent_span", None)):
+            parent_span = req_parent_span
+            if hasattr(request, "extensions") and isinstance(request.extensions, dict):
+                request.extensions[PARENT_SPAN_KEY] = req_parent_span
+        elif hasattr(request, "extensions") and isinstance(
+            request.extensions, dict
+        ) and (req_parent_span := request.extensions.get(PARENT_SPAN_KEY)):
+            parent_span = req_parent_span
+        if parent_span is None:
             return tracer.start_span(span_name)
-        return tracer.start_span(span_name, trace.set_span_in_context(self.parent_span))
+        return tracer.start_span(span_name, trace.set_span_in_context(parent_span))

packages/http/httpx/kiota_http/middleware/url_replace_handler.py

@@ -34,6 +34,8 @@ async def send(
         Returns:
             Response: The response object.
         """
+        if request is None:
+            raise TypeError("request cannot be null")
         _enable_span = self._create_observability_span(request, "UrlReplaceHandler_send")
         try:
             current_options = self._get_current_options(request)
@@ -45,6 +47,9 @@ async def send(
                 request.url = httpx.URL(url_string)
                 _enable_span.set_attribute(URL_FULL, str(request.url))
             return await super().send(request, transport)
+        except Exception as exc:
+            _enable_span.record_exception(exc)
+            raise
         finally:
             _enable_span.end()
 

packages/http/httpx/kiota_http/middleware/user_agent_handler.py

@@ -22,14 +22,21 @@ async def send(self, request: Request, transport: AsyncBaseTransport) -> Respons
         Checks if the request has a User-Agent header and updates it if the
         platform config allows.
         """
+        if request is None:
+            raise TypeError("request cannot be null")
         current_options = self._get_current_options(request)
         _span = self._create_observability_span(request, "UserAgentHandler_send")
-        if current_options and current_options.is_enabled:
-            _span.set_attribute("com.microsoft.kiota.handler.useragent.enable", True)
-            value = f"{current_options.product_name}/{current_options.product_version}"
-            self._update_user_agent(request, value)
-        _span.end()
-        return await super().send(request, transport)
+        try:
+            if current_options and current_options.is_enabled:
+                _span.set_attribute("com.microsoft.kiota.handler.useragent.enable", True)
+                value = f"{current_options.product_name}/{current_options.product_version}"
+                self._update_user_agent(request, value)
+            return await super().send(request, transport)
+        except Exception as exc:
+            _span.record_exception(exc)
+            raise
+        finally:
+            _span.end()
 
     def _get_current_options(self, request: Request) -> UserAgentHandlerOption:
         """Returns the options to use for the request.Overries default options if

packages/http/httpx/tests/middleware_tests/test_base_middleware.py

@@ -3,7 +3,7 @@
 import pytest
 from opentelemetry import trace
 
-from kiota_http.middleware import REQUEST_OPTIONS_KEY, BaseMiddleware, MiddlewarePipeline
+from kiota_http.middleware import PARENT_SPAN_KEY, REQUEST_OPTIONS_KEY, BaseMiddleware, MiddlewarePipeline
 from kiota_http.middleware import middleware as middleware_module
 
 
@@ -30,19 +30,21 @@ def test_span_created(request_info, otel_tracer):
         span = middleware._create_observability_span(request_info, "test_span_created")
         span.end()
         after_state = outside_span.is_recording()
-        assert(before_state == after_state)
+        assert (before_state == after_state)
 
 
 def test_span_without_parent_span_uses_the_current_span(span_exporter):
     middleware = BaseMiddleware()
     with middleware_module.tracer.start_as_current_span("outer"):
-        middleware._create_observability_span(httpx.Request("GET", "https://example.com"), "inner").end()
+        middleware._create_observability_span(httpx.Request("GET", "https://example.com"),
+                                              "inner").end()
     inner, outer = span_exporter.get_finished_spans()
     assert inner.parent.span_id == outer.context.span_id
 
 
 @pytest.mark.asyncio
 async def test_pipeline_without_middleware_sends_the_request():
+
     def request_handler(request: httpx.Request):
         return httpx.Response(200)
 
@@ -51,3 +53,54 @@ def request_handler(request: httpx.Request):
         response = await MiddlewarePipeline(transport).send(request)
     assert response.status_code == 200
     assert REQUEST_OPTIONS_KEY not in request.extensions
+
+
+def test_parent_span_option_does_not_mutate_middleware_instance(span_exporter):
+    middleware = BaseMiddleware()
+    assert middleware.parent_span is None
+
+    with middleware_module.tracer.start_as_current_span("parent-span-1") as span1:
+        req1 = httpx.Request(
+            "GET", "https://example.com", extensions={REQUEST_OPTIONS_KEY: {
+                "parent_span": span1
+            }}
+        )
+        child1 = middleware._create_observability_span(req1, "child1")
+        child1.end()
+
+    assert middleware.parent_span is None
+
+    req2 = httpx.Request("GET", "https://example.com")
+    child2 = middleware._create_observability_span(req2, "child2")
+    child2.end()
+    assert middleware.parent_span is None
+
+    spans = span_exporter.get_finished_spans()
+    child1_span = [s for s in spans if s.name == "child1"][0]
+    parent1_span = [s for s in spans if s.name == "parent-span-1"][0]
+    assert child1_span.parent.span_id == parent1_span.context.span_id
+
+
+@pytest.mark.asyncio
+async def test_base_middleware_send_preserves_parent_span_key_during_transport_cleanup():
+
+    def request_handler(request: httpx.Request):
+        assert REQUEST_OPTIONS_KEY not in request.extensions
+        assert PARENT_SPAN_KEY in request.extensions
+        return httpx.Response(200)
+
+    middleware = BaseMiddleware()
+    span = middleware_module.tracer.start_span("test-parent")
+    request = httpx.Request(
+        "GET",
+        "https://example.com",
+        extensions={REQUEST_OPTIONS_KEY: {
+            "parent_span": span
+        }},
+    )
+    async with httpx.MockTransport(request_handler) as transport:
+        response = await middleware.send(request, transport)
+    assert response.status_code == 200
+    assert middleware.parent_span is None
+    assert REQUEST_OPTIONS_KEY not in request.extensions
+    assert request.extensions.get(PARENT_SPAN_KEY) is span

packages/http/httpx/tests/middleware_tests/test_retry_handler.py

@@ -4,7 +4,14 @@
 import httpx
 import pytest
 
-from kiota_http.middleware import RetryHandler
+from kiota_http.middleware import middleware as middleware_module
+from kiota_http.middleware import (
+    PARENT_SPAN_KEY,
+    REQUEST_OPTIONS_KEY,
+    MiddlewarePipeline,
+    RetryHandler,
+    UserAgentHandler,
+)
 from kiota_http.middleware.options import RetryHandlerOption
 
 BASE_URL = 'https://httpbin.org'
@@ -349,13 +356,15 @@ def request_handler(request: httpx.Request):
 
 @pytest.mark.asyncio
 async def test_retry_spans_end(span_exporter):
+
     def request_handler(request: httpx.Request):
         if RETRY_ATTEMPT in request.headers:
             return httpx.Response(200)
         return httpx.Response(SERVICE_UNAVAILABLE)
 
     async with httpx.MockTransport(request_handler) as transport:
-        await RetryHandler(RetryHandlerOption(delay=0)).send(httpx.Request("GET", BASE_URL), transport)
+        await RetryHandler(RetryHandlerOption(delay=0)
+                           ).send(httpx.Request("GET", BASE_URL), transport)
     assert [span.name for span in span_exporter.get_finished_spans()] == [
         "RetryHandler_send", "RetryHandler_send - attempt 0", "RetryHandler_send - attempt 1"
     ]
@@ -376,3 +385,62 @@ def test_get_retry_after_invalid_value():
     response = httpx.Response(SERVICE_UNAVAILABLE, headers={RETRY_AFTER: "soon"})
     assert RetryHandler()._get_retry_after(response) is None
 
+
+@pytest.mark.asyncio
+async def test_retry_503_to_200_preserves_parent_span_across_attempts_and_downstream_spans(
+    span_exporter
+):
+    """
+    Ensure RetryHandler preserves the request-local parent span across 503->200
+    retries for both attempt spans and rerun downstream middleware spans without
+    mutating shared middleware instances.
+    """
+    attempts = 0
+
+    def request_handler(request: httpx.Request):
+        nonlocal attempts
+        attempts += 1
+        if attempts == 1:
+            return httpx.Response(SERVICE_UNAVAILABLE)
+        return httpx.Response(200)
+
+    retry_handler = RetryHandler(RetryHandlerOption(delay=0.0, max_retries=2))
+    retry_handler.backoff_factor = 0.0
+    retry_handler.backoff_max = 0
+
+    user_agent_handler = UserAgentHandler()
+
+    pipeline = MiddlewarePipeline(httpx.MockTransport(request_handler))
+    pipeline.add_middleware(retry_handler)
+    pipeline.add_middleware(user_agent_handler)
+
+    with middleware_module.tracer.start_as_current_span("explicit-parent-span") as parent_span:
+        pass
+
+    request = httpx.Request(
+        "GET",
+        BASE_URL,
+        extensions={REQUEST_OPTIONS_KEY: {
+            "parent_span": parent_span
+        }},
+    )
+
+    response = await pipeline.send(request)
+
+    assert response.status_code == 200
+    assert attempts == 2
+    assert retry_handler.parent_span is None
+    assert user_agent_handler.parent_span is None
+
+    spans = span_exporter.get_finished_spans()
+    expected_parent_id = parent_span.get_span_context().span_id
+
+    attempt_0 = next(s for s in spans if s.name == "RetryHandler_send - attempt 0")
+    attempt_1 = next(s for s in spans if s.name == "RetryHandler_send - attempt 1")
+    assert attempt_0.parent.span_id == expected_parent_id
+    assert attempt_1.parent.span_id == expected_parent_id
+
+    ua_spans = [s for s in spans if s.name == "UserAgentHandler_send"]
+    assert len(ua_spans) == 2
+    for ua_span in ua_spans:
+        assert ua_span.parent.span_id == expected_parent_id

packages/http/httpx/tests/middleware_tests/test_url_replace_handler.py

@@ -84,3 +84,27 @@
             await handler.send(request, transport)
 
     assert [span.name for span in span_exporter.get_finished_spans()] == ["UrlReplaceHandler_send"]
+
+@pytest.mark.asyncio
+async def test_url_replace_handler_null_request_raises():
+    handler = UrlReplaceHandler()
+    with pytest.raises(TypeError, match="request cannot be null"):
+        await handler.send(None, httpx.MockTransport(lambda req: httpx.Response(200)))
+
+
+@pytest.mark.asyncio
+async def test_url_replace_handler_span_records_exception_on_failure(span_exporter):
+    def request_handler(request):
+        raise httpx.ReadError("connection reset")
+
+    handler = UrlReplaceHandler()
+    request = httpx.Request("GET", ORIGINAL_URL)
+    async with httpx.MockTransport(request_handler) as transport:
+        with pytest.raises(httpx.ReadError):
+            await handler.send(request, transport)
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].name == "UrlReplaceHandler_send"
+    assert len(spans[0].events) > 0
+    assert spans[0].events[0].name == "exception"

packages/http/httpx/tests/middleware_tests/test_user_agent_handler.py

@@ -42,3 +42,44 @@
         response = await UserAgentHandler().send(request, transport)
 
     assert response.request.headers["User-Agent"] == "custom/9"
+
+@pytest.mark.asyncio
+async def test_user_agent_handler_null_request_raises():
+    handler = UserAgentHandler()
+    with pytest.raises(TypeError, match="request cannot be null"):
+        await handler.send(None, httpx.MockTransport(lambda req: httpx.Response(200)))
+
+
+@pytest.mark.asyncio
+async def test_user_agent_handler_span_encloses_transport(span_exporter):
+    def request_handler(request):
+        assert len(span_exporter.get_finished_spans()) == 0
+        return httpx.Response(200)
+
+    handler = UserAgentHandler()
+    request = httpx.Request("GET", "https://example.com")
+    async with httpx.MockTransport(request_handler) as transport:
+        resp = await handler.send(request, transport)
+
+    assert resp.status_code == 200
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].name == "UserAgentHandler_send"
+
+
+@pytest.mark.asyncio
+async def test_user_agent_handler_span_records_exception_on_transport_failure(span_exporter):
+    def request_handler(request):
+        raise httpx.ReadError("connection failed")
+
+    handler = UserAgentHandler()
+    request = httpx.Request("GET", "https://example.com")
+    async with httpx.MockTransport(request_handler) as transport:
+        with pytest.raises(httpx.ReadError):
+            await handler.send(request, transport)
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].name == "UserAgentHandler_send"
+    assert len(spans[0].events) > 0
+    assert spans[0].events[0].name == "exception"