Update telemetry handler to set context and feature flags

Author: samwelkanda

msgraph/core/middleware/telemetry.py

@@ -1,24 +1,30 @@
+import http
+import json
 import platform
 
 import httpx
-from kiota_http.middleware.middleware import BaseMiddleware
+from kiota_http.middleware import AsyncKiotaTransport, BaseMiddleware, RedirectHandler, RetryHandler
 from urllib3.util import parse_url
 
 from .._constants import SDK_VERSION
-from .._enums import NationalClouds
-from .middleware import GraphRequest
+from .._enums import FeatureUsageFlag, NationalClouds
+from .request_context import GraphRequestContext
+
+
+class GraphRequest(httpx.Request):
+    context: GraphRequestContext
 
 
 class GraphTelemetryHandler(BaseMiddleware):
     """Middleware component that attaches metadata to a Graph request in order to help
     the SDK team improve the developer experience.
     """
 
-    async def send(
-        self, request: GraphRequest, transport: httpx.AsyncBaseTransport
-    ) -> httpx.Response:
+    async def send(self, request: GraphRequest, transport: AsyncKiotaTransport):
         """Adds telemetry headers and sends the http request.
         """
+        self.set_request_context_and_feature_usage(request, transport)
+
         if self.is_graph_url(request.url):
             self._add_client_request_id_header(request)
             self._append_sdk_version_header(request)
@@ -28,6 +34,27 @@ async def send(
         response = await super().send(request, transport)
         return response
 
+    def set_request_context_and_feature_usage(
+        self, request: GraphRequest, transport: AsyncKiotaTransport
+    ) -> GraphRequest:
+
+        request_options = {}
+        options = request.headers.pop('request_options', None)
+        if options:
+            request_options = json.loads(options)
+
+        request.context = GraphRequestContext(request_options, request.headers)
+        middleware = transport.pipeline._first_middleware
+        while middleware:
+            if isinstance(middleware, RedirectHandler):
+                request.context.feature_usage = FeatureUsageFlag.REDIRECT_HANDLER_ENABLED
+            if isinstance(middleware, RetryHandler):
+                request.context.feature_usage = FeatureUsageFlag.RETRY_HANDLER_ENABLED
+
+            middleware = middleware.next
+
+        return request
+
     def is_graph_url(self, url):
         """Check if the request is made to a graph endpoint. We do not add telemetry headers to
         non-graph endpoints"""