From 73ff55fd4c4703429cd49035acc73b7079f3b645 Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Sat, 24 Oct 2026 16:08:22 +0530 Subject: [PATCH 1/2] fix(http): isolate parent span per request, extend user agent span over transport, and record telemetry errors --- .../httpx/kiota_http/middleware/middleware.py | 9 ++-- .../middleware/url_replace_handler.py | 5 +++ .../middleware/user_agent_handler.py | 19 ++++++--- .../middleware_tests/test_base_middleware.py | 21 ++++++++++ .../test_url_replace_handler.py | 24 +++++++++++ .../test_user_agent_handler.py | 41 +++++++++++++++++++ 6 files changed, 109 insertions(+), 10 deletions(-) diff --git a/packages/http/httpx/kiota_http/middleware/middleware.py b/packages/http/httpx/kiota_http/middleware/middleware.py index 17142649..0346f81d 100644 --- a/packages/http/httpx/kiota_http/middleware/middleware.py +++ b/packages/http/httpx/kiota_http/middleware/middleware.py @@ -69,8 +69,9 @@ 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 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)) diff --git a/packages/http/httpx/kiota_http/middleware/url_replace_handler.py b/packages/http/httpx/kiota_http/middleware/url_replace_handler.py index d1efb866..207f58e2 100644 --- a/packages/http/httpx/kiota_http/middleware/url_replace_handler.py +++ b/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() diff --git a/packages/http/httpx/kiota_http/middleware/user_agent_handler.py b/packages/http/httpx/kiota_http/middleware/user_agent_handler.py index 28e32300..e72d354f 100644 --- a/packages/http/httpx/kiota_http/middleware/user_agent_handler.py +++ b/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 diff --git a/packages/http/httpx/tests/middleware_tests/test_base_middleware.py b/packages/http/httpx/tests/middleware_tests/test_base_middleware.py index ec1983e4..cffd98e3 100644 --- a/packages/http/httpx/tests/middleware_tests/test_base_middleware.py +++ b/packages/http/httpx/tests/middleware_tests/test_base_middleware.py @@ -51,3 +51,24 @@ 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 diff --git a/packages/http/httpx/tests/middleware_tests/test_url_replace_handler.py b/packages/http/httpx/tests/middleware_tests/test_url_replace_handler.py index c19cfecf..e19fec1a 100644 --- a/packages/http/httpx/tests/middleware_tests/test_url_replace_handler.py +++ b/packages/http/httpx/tests/middleware_tests/test_url_replace_handler.py @@ -84,3 +84,27 @@ def request_handler(request): 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" diff --git a/packages/http/httpx/tests/middleware_tests/test_user_agent_handler.py b/packages/http/httpx/tests/middleware_tests/test_user_agent_handler.py index f1d8f4a8..ff9d9445 100644 --- a/packages/http/httpx/tests/middleware_tests/test_user_agent_handler.py +++ b/packages/http/httpx/tests/middleware_tests/test_user_agent_handler.py @@ -42,3 +42,44 @@ async def test_request_option_sets_the_product(): 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" From a73da334a6948b53674424b1019fe91ed52e89af Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Mon, 26 Oct 2026 01:21:25 +0530 Subject: [PATCH 2/2] fix(http): preserve request parent span across retries and transport cleanup Preserve explicit parent_span in request.extensions[PARENT_SPAN_KEY] across transport cleanup so RetryHandler attempts and downstream middleware reruns retain trace context without mutating shared middleware instances. Add 503->200 regression test in test_retry_handler.py verifying attempt and downstream span parent IDs match the explicit parent, and unit test in test_base_middleware.py. --- .../httpx/kiota_http/middleware/__init__.py | 3 +- .../httpx/kiota_http/middleware/middleware.py | 16 ++++- .../middleware_tests/test_base_middleware.py | 40 +++++++++-- .../middleware_tests/test_retry_handler.py | 72 ++++++++++++++++++- 4 files changed, 123 insertions(+), 8 deletions(-) diff --git a/packages/http/httpx/kiota_http/middleware/__init__.py b/packages/http/httpx/kiota_http/middleware/__init__.py index 2c4760ff..9555b155 100644 --- a/packages/http/httpx/kiota_http/middleware/__init__.py +++ b/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", diff --git a/packages/http/httpx/kiota_http/middleware/middleware.py b/packages/http/httpx/kiota_http/middleware/middleware.py index 0346f81d..ece95fcf 100644 --- a/packages/http/httpx/kiota_http/middleware/middleware.py +++ b/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 @@ -72,6 +80,12 @@ def _create_observability_span(self, request: httpx.Request, span_name: str) -> 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(parent_span)) diff --git a/packages/http/httpx/tests/middleware_tests/test_base_middleware.py b/packages/http/httpx/tests/middleware_tests/test_base_middleware.py index cffd98e3..16ad53c6 100644 --- a/packages/http/httpx/tests/middleware_tests/test_base_middleware.py +++ b/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) @@ -52,12 +54,17 @@ def request_handler(request: httpx.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}}) + req1 = httpx.Request( + "GET", "https://example.com", extensions={REQUEST_OPTIONS_KEY: { + "parent_span": span1 + }} + ) child1 = middleware._create_observability_span(req1, "child1") child1.end() @@ -72,3 +79,28 @@ def test_parent_span_option_does_not_mutate_middleware_instance(span_exporter): 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 diff --git a/packages/http/httpx/tests/middleware_tests/test_retry_handler.py b/packages/http/httpx/tests/middleware_tests/test_retry_handler.py index 1cbe07ed..5bfc3488 100644 --- a/packages/http/httpx/tests/middleware_tests/test_retry_handler.py +++ b/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