feat(http): implement body inspection handler

packages/http/httpx/kiota_http/httpx_request_adapter.py

@@ -44,7 +44,7 @@
 
 from ._version import VERSION
 from .kiota_client_factory import KiotaClientFactory
-from .middleware import ParametersNameDecodingHandler
+from .middleware import REQUEST_OPTIONS_KEY, ParametersNameDecodingHandler
 from .middleware.options import ParametersNameDecodingHandlerOption, ResponseHandlerOption
 from .observability_options import ObservabilityOptions
 
@@ -686,18 +686,18 @@ def get_request_from_request_information(
         if self.observability_options.include_euii_attributes:
             otel_attributes.update({URL_FULL: url.geturl()})
 
+        request_options = {
+            self.observability_options.get_key(): self.observability_options,
+            "parent_span": parent_span,
+            **request_info.request_options,
+        }
         request = self._http_client.build_request(
             method=method.value,
             url=request_info.url,
             headers=request_info.request_headers,
             content=request_info.content,
+            extensions={REQUEST_OPTIONS_KEY: request_options},
         )
-        request_options = {
-            self.observability_options.get_key(): self.observability_options,
-            "parent_span": parent_span,
-            **request_info.request_options,
-        }
-        setattr(request, "options", request_options)
 
         if content_length := request.headers.get("Content-Length", None):
             otel_attributes.update({"http.request.body.size": content_length})

packages/http/httpx/kiota_http/kiota_client_factory.py

@@ -11,6 +11,7 @@
 from .middleware import (
     AsyncKiotaTransport,
     BaseMiddleware,
+    BodyInspectionHandler,
     HeadersInspectionHandler,
     MiddlewarePipeline,
     ParametersNameDecodingHandler,
@@ -19,6 +20,7 @@
     UrlReplaceHandler,
 )
 from .middleware.options import (
+    BodyInspectionHandlerOption,
     HeadersInspectionHandlerOption,
     ParametersNameDecodingHandlerOption,
     RedirectHandlerOption,
@@ -91,6 +93,7 @@ def get_default_middleware(options: Optional[dict[str, RequestOption]]) -> list[
         url_replace_handler = UrlReplaceHandler()
         user_agent_handler = UserAgentHandler()
         headers_inspection_handler = HeadersInspectionHandler()
+        body_inspection_handler = BodyInspectionHandler()
 
         if options:
             redirect_handler_options = options.get(RedirectHandlerOption.get_key())
@@ -135,11 +138,18 @@ def get_default_middleware(options: Optional[dict[str, RequestOption]]) -> list[
                     options=headers_inspection_handler_options
                 )
 
-        middleware = [
+            body_inspection_handler_options = options.get(BodyInspectionHandlerOption.get_key())
+            if body_inspection_handler_options and isinstance(
+                body_inspection_handler_options, BodyInspectionHandlerOption
+            ):
+                body_inspection_handler = BodyInspectionHandler(
+                    options=body_inspection_handler_options
+                )
+
+        return [
             redirect_handler, retry_handler, parameters_name_decoding_handler, url_replace_handler,
-            user_agent_handler, headers_inspection_handler
+            user_agent_handler, headers_inspection_handler, body_inspection_handler
         ]
-        return middleware
 
     @staticmethod
     def create_middleware_pipeline(

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

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

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

@@ -0,0 +1,116 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation. All Rights Reserved.
+# Licensed under the MIT License.
+# See License in the project root for license information.
+# ------------------------------------
+
+from typing import Optional
+
+import httpx
+
+from .middleware import REQUEST_OPTIONS_KEY, BaseMiddleware
+from .options import BodyInspectionHandlerOption
+
+BODY_INSPECTION_KEY = "com.microsoft.kiota.handler.bodyInspection.enable"
+
+
+class BodyInspectionHandler(BaseMiddleware):
+    """The Body Inspection Handler allows the developer to inspect the body of the
+    request and response.
+    """
+
+    def __init__(
+        self,
+        options: Optional[BodyInspectionHandlerOption] = None,
+    ):
+        """Create an instance of BodyInspectionHandler
+
+        Args:
+            options (BodyInspectionHandlerOption, optional): Default options to apply to the
+            handler. A new BodyInspectionHandlerOption per handler when not provided.
+        """
+        super().__init__()
+        self.options = options if options is not None else BodyInspectionHandlerOption()
+
+    async def send(
+        self, request: httpx.Request, transport: httpx.AsyncBaseTransport
+    ) -> httpx.Response:
+        """To execute the current middleware
+
+        Args:
+            request (httpx.Request): The prepared request object
+            transport (httpx.AsyncBaseTransport): The HTTP transport to use
+
+        Returns:
+            httpx.Response: The response object.
+        """
+        if request is None:
+            raise TypeError("request cannot be null")
+
+        current_options = self._get_current_options(request)
+        span = self._create_observability_span(request, "BodyInspectionHandler_send")
+        try:
+            span.set_attribute(BODY_INSPECTION_KEY, True)
+
+            if current_options and current_options.inspect_request_body:
+                content = await request.aread()
+                if content:
+                    current_options.request_body = content
+                else:
+                    current_options.request_body = None
+
+            response = await super().send(request, transport)
+
+            if current_options and current_options.inspect_response_body:
+                response_content: Optional[bytes] = None
+                # A consumed stream is inspectable only when HTTPX cached its content.
+                if hasattr(response, "_content"):
+                    response_content = response.content
+                elif not response.is_stream_consumed and not response.is_closed:
+                    num_bytes_downloaded = response.num_bytes_downloaded
+                    raw_content = b"".join([chunk async for chunk in response.aiter_raw()])
+                    self._restore_response_stream(response, raw_content, num_bytes_downloaded)
+                    response_content = await response.aread()
+                    self._restore_response_stream(response, raw_content, num_bytes_downloaded)
+                if response_content:
+                    current_options.response_body = response_content
+                else:
+                    current_options.response_body = None
+
+            return response
+        finally:
+            span.end()
+
+    def _get_current_options(self, request: httpx.Request) -> BodyInspectionHandlerOption:
+        """Returns the options to use for the request. Overrides default options if
+        request options are passed.
+
+        Args:
+            request (httpx.Request): The prepared request object
+
+        Returns:
+            BodyInspectionHandlerOption: The options to be used.
+        """
+        current_options = None
+        request_options = request.extensions.get(REQUEST_OPTIONS_KEY)
+        if request_options:
+            current_options = request_options.get(BodyInspectionHandlerOption.get_key(), None)
+        if not current_options:
+            current_options = self.options
+
+        current_options._clear_captured_bodies()
+        return current_options
+
+    @staticmethod
+    def _restore_response_stream(
+        response: httpx.Response, content: bytes, num_bytes_downloaded: int
+    ) -> None:
+        # aread() caches decoded content and a stateful decoder; discard both when rewinding.
+        if hasattr(response, "_content"):
+            del response._content
+        if hasattr(response, "_decoder"):
+            del response._decoder
+        response.stream = httpx.ByteStream(content)
+        response.is_stream_consumed = False
+        response.is_closed = False
+        response._num_bytes_downloaded = num_bytes_downloaded

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

@@ -10,7 +10,7 @@
 
 import httpx
 
-from .middleware import BaseMiddleware
+from .middleware import REQUEST_OPTIONS_KEY, BaseMiddleware
 from .options import HeadersInspectionHandlerOption
 
 HEADERS_INSPECTION_KEY = "com.microsoft.kiota.handler.headers_inspection.enable"
@@ -71,9 +71,9 @@ def _get_current_options(self, request: httpx.Request) -> HeadersInspectionHandl
             HeadersInspectionHandlerOption: The options to be used.
         """
         current_options = None
-        request_options = getattr(request, "options", None)
+        request_options = request.extensions.get(REQUEST_OPTIONS_KEY)
         if request_options:
-            current_options = request_options.get(  # type:ignore
+            current_options = request_options.get(
                 HeadersInspectionHandlerOption.get_key(), None
             )
         if current_options:

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

@@ -10,6 +10,8 @@
 
 tracer = trace.get_tracer(ObservabilityOptions.get_tracer_instrumentation_name(), VERSION)
 
+REQUEST_OPTIONS_KEY = "kiota_request_options"
+
 
 class MiddlewarePipeline():
     """MiddlewarePipeline, entry point of middleware
@@ -56,8 +58,8 @@ def __init__(self):
     async def send(self, request, transport):
         if self.next is None:
             # Remove request options if there's no other middleware in the chain.
-            if hasattr(request, "options") and request.options:
-                delattr(request, 'options')
+            if hasattr(request, "extensions") and isinstance(request.extensions, dict):
+                request.extensions.pop(REQUEST_OPTIONS_KEY, None)
             response = await transport.handle_async_request(request)
             response.request = request
             return response
@@ -68,11 +70,13 @@ def _create_observability_span(self, request, span_name: str) -> trace.Span:
         If no parent_span is found in the request, uses the parent_span in the
         object. If parent_span is None, current context will be used."""
         _span = None
-        if options := getattr(request, "options", None):
-            if parent_span := options.get("parent_span", None):
-                self.parent_span = parent_span
-                _context = trace.set_span_in_context(parent_span)
-                _span = tracer.start_span(span_name, _context)
+        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
+            _context = trace.set_span_in_context(parent_span)
+            _span = tracer.start_span(span_name, _context)
         if _span is None:
             _context = trace.set_span_in_context(self.parent_span)
             _span = tracer.start_span(span_name, _context)

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

@@ -1,3 +1,4 @@
+from .body_inspection_handler_option import BodyInspectionHandlerOption
 from .headers_inspection_handler_option import HeadersInspectionHandlerOption
 from .parameters_name_decoding_handler_option import ParametersNameDecodingHandlerOption
 from .redirect_handler_option import RedirectHandlerOption