Merge remote-tracking branch 'origin/main' into merge-main-develop-strands

Author: TimPansino

newrelic/api/asgi_application.py

@@ -132,10 +132,20 @@ async def send_inject_browser_agent(self, message):
 
         message_type = message["type"]
         if message_type == "http.response.start" and not self.initial_message:
-            headers = list(message.get("headers", ()))
+            # message["headers"] may be a generator, and consuming it via process_response will leave the original
+            # application with no headers. Fix this by preserving them in a list before consuming them.
+            if "headers" in message:
+                message["headers"] = headers = list(message["headers"])
+            else:
+                headers = []
+
+            # Check if we should insert the HTML snippet based on the headers.
+            # Currently if there are no headers this will always be False, but call the function
+            # anyway in case this logic changes in the future.
             if not self.should_insert_html(headers):
                 await self.abort()
                 return
+
             message["headers"] = headers
             self.initial_message = message
         elif message_type == "http.response.body" and self.initial_message:
@@ -232,7 +242,13 @@ async def send(self, event):
             finally:
                 self.__exit__(*sys.exc_info())
         elif event["type"] == "http.response.start":
-            self.process_response(event["status"], event.get("headers", ()))
+            # event["headers"] may be a generator, and consuming it via process_response will leave the original
+            # ASGI application with no headers. Fix this by preserving them in a list before consuming them.
+            if "headers" in event:
+                event["headers"] = headers = list(event["headers"])
+            else:
+                headers = []
+            self.process_response(event["status"], headers)
         return await self._send(event)
 
 

newrelic/config.py

@@ -2084,6 +2084,10 @@ def _process_module_builtin_defaults():
         "asyncio.base_events", "newrelic.hooks.coroutines_asyncio", "instrument_asyncio_base_events"
     )
 
+    _process_module_definition("asyncio.events", "newrelic.hooks.coroutines_asyncio", "instrument_asyncio_events")
+
+    _process_module_definition("asyncio.runners", "newrelic.hooks.coroutines_asyncio", "instrument_asyncio_runners")
+
     _process_module_definition(
         "langchain_core.runnables.base",
         "newrelic.hooks.mlmodel_langchain",
@@ -2671,8 +2675,6 @@ def _process_module_builtin_defaults():
         "langchain_core.callbacks.manager", "newrelic.hooks.mlmodel_langchain", "instrument_langchain_callbacks_manager"
     )
 
-    _process_module_definition("asyncio.events", "newrelic.hooks.coroutines_asyncio", "instrument_asyncio_events")
-
     _process_module_definition("asgiref.sync", "newrelic.hooks.adapter_asgiref", "instrument_asgiref_sync")
 
     _process_module_definition(

newrelic/hooks/coroutines_asyncio.py

@@ -16,36 +16,73 @@
 from newrelic.core.trace_cache import trace_cache
 
 
-def remove_from_cache(task):
+def remove_from_cache_callback(task):
     cache = trace_cache()
     cache.task_stop(task)
 
 
-def propagate_task_context(task):
+def wrap_create_task(task):
     trace_cache().task_start(task)
-    task.add_done_callback(remove_from_cache)
+    task.add_done_callback(remove_from_cache_callback)
     return task
 
 
-def _bind_loop(loop, *args, **kwargs):
+def _instrument_event_loop(loop):
+    if loop and hasattr(loop, "create_task") and not hasattr(loop.create_task, "__wrapped__"):
+        wrap_out_function(loop, "create_task", wrap_create_task)
+
+
+def _bind_set_event_loop(loop, *args, **kwargs):
     return loop
 
 
-def wrap_create_task(wrapped, instance, args, kwargs):
-    loop = _bind_loop(*args, **kwargs)
+def wrap_set_event_loop(wrapped, instance, args, kwargs):
+    loop = _bind_set_event_loop(*args, **kwargs)
 
-    if loop and not hasattr(loop.create_task, "__wrapped__"):
-        wrap_out_function(loop, "create_task", propagate_task_context)
+    _instrument_event_loop(loop)
 
     return wrapped(*args, **kwargs)
 
 
+def wrap__lazy_init(wrapped, instance, args, kwargs):
+    result = wrapped(*args, **kwargs)
+    # This logic can be used for uvloop, but should
+    # work for any valid custom loop factory.
+
+    # A custom loop_factory will be used to create
+    # a new event loop instance.  It will then run
+    # the main() coroutine on this event loop.  Once
+    # this coroutine is complete, the event loop will
+    # be stopped and closed.
+
+    # The new loop that is created and set as the
+    # running loop of the duration of the run() call.
+    # When the coroutine starts, it runs in the context
+    # that was active when run() was called.  Any tasks
+    # created within this coroutine on this new event
+    # loop will inherit that context.
+
+    # Note: The loop created by loop_factory is never
+    # set as the global current loop for the thread,
+    # even while it is running.
+    loop = instance._loop
+    _instrument_event_loop(loop)
+
+    return result
+
+
 def instrument_asyncio_base_events(module):
-    wrap_out_function(module, "BaseEventLoop.create_task", propagate_task_context)
+    wrap_out_function(module, "BaseEventLoop.create_task", wrap_create_task)
 
 
 def instrument_asyncio_events(module):
     if hasattr(module, "_BaseDefaultEventLoopPolicy"):  # Python >= 3.14
-        wrap_function_wrapper(module, "_BaseDefaultEventLoopPolicy.set_event_loop", wrap_create_task)
-    else:  # Python <= 3.13
-        wrap_function_wrapper(module, "BaseDefaultEventLoopPolicy.set_event_loop", wrap_create_task)
+        wrap_function_wrapper(module, "_BaseDefaultEventLoopPolicy.set_event_loop", wrap_set_event_loop)
+    elif hasattr(module, "BaseDefaultEventLoopPolicy"):  # Python <= 3.13
+        wrap_function_wrapper(module, "BaseDefaultEventLoopPolicy.set_event_loop", wrap_set_event_loop)
+
+
+# For Python >= 3.11
+def instrument_asyncio_runners(module):
+    if hasattr(module, "Runner") and hasattr(module.Runner, "_lazy_init"):
+        wrap_function_wrapper(module, "Runner._lazy_init", wrap__lazy_init)

newrelic/hooks/external_aiobotocore.py

@@ -149,6 +149,17 @@ async def wrap_client__make_api_call(wrapped, instance, args, kwargs):
             bedrock_attrs = extract_bedrock_converse_attrs(
                 args[1], response, response_headers, model, span_id, trace_id
             )
+
+            if response_streaming:
+                # Wrap EventStream object here to intercept __iter__ method instead of instrumenting class.
+                # This class is used in numerous other services in botocore, and would cause conflicts.
+                response["stream"] = stream = AsyncEventStreamWrapper(response["stream"])
+                stream._nr_ft = ft or None
+                stream._nr_bedrock_attrs = bedrock_attrs or {}
+                stream._nr_model_extractor = stream_extractor or None
+                stream._nr_is_converse = True
+                return response
+
         else:
             bedrock_attrs = {
                 "request_id": response_headers.get("x-amzn-requestid"),