feat: add on_response_headers_received signal to TraceConfig

CHANGES/7386.feature.rst

@@ -0,0 +1 @@
+Added on_response_headers_received signal to TraceConfig, allowing users to trace response headers as soon as they are received from the server.

aiohttp/client_reqrep.py

@@ -11,7 +11,7 @@
 from hashlib import md5, sha1, sha256
 from http.cookies import BaseCookie, SimpleCookie
 from types import MappingProxyType, TracebackType
-from typing import TYPE_CHECKING, Any, Literal, NamedTuple, TypedDict
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple, TypedDict, cast
 
 from multidict import CIMultiDict, CIMultiDictProxy, MultiDict, MultiDictProxy
 from yarl import URL, Query
@@ -460,6 +460,13 @@ async def start(self, connection: "Connection") -> "ClientResponse":
         self._headers = message.headers
         self._raw_headers = message.raw_headers
 
+        # fire headers_received trace signal
+        if self._traces:
+            for trace in self._traces:
+                await trace.send_response_headers_received(
+                    self.method, self.url, cast(CIMultiDictProxy[str], self._headers)
+                )
+
         # payload
         self.content = payload
 

aiohttp/tracing.py

@@ -2,7 +2,7 @@
 from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, overload
 
 from aiosignal import Signal
-from multidict import CIMultiDict
+from multidict import CIMultiDict, CIMultiDictProxy
 from yarl import URL
 
 from .client_reqrep import ClientResponse
@@ -30,6 +30,7 @@
     "TraceRequestChunkSentParams",
     "TraceResponseChunkReceivedParams",
     "TraceRequestHeadersSentParams",
+    "TraceResponseHeadersReceivedParams",
 )
 
 _T = TypeVar("_T", covariant=True)
@@ -97,6 +98,9 @@ def __init__(
         self._on_request_headers_sent: _TracingSignal[
             _T, TraceRequestHeadersSentParams
         ] = Signal(self)
+        self._on_response_headers_received: _TracingSignal[
+            _T, TraceResponseHeadersReceivedParams
+        ] = Signal(self)
 
         self._trace_config_ctx_factory: _Factory[_T] = trace_config_ctx_factory
 
@@ -121,6 +125,7 @@ def freeze(self) -> None:
         self._on_dns_cache_hit.freeze()
         self._on_dns_cache_miss.freeze()
         self._on_request_headers_sent.freeze()
+        self._on_response_headers_received.freeze()
 
     @property
     def on_request_start(self) -> "_TracingSignal[_T, TraceRequestStartParams]":
@@ -210,6 +215,12 @@ def on_request_headers_sent(
     ) -> "_TracingSignal[_T, TraceRequestHeadersSentParams]":
         return self._on_request_headers_sent
 
+    @property
+    def on_response_headers_received(
+        self,
+    ) -> "_TracingSignal[_T, TraceResponseHeadersReceivedParams]":
+        return self._on_response_headers_received
+
 
 @frozen_dataclass_decorator
 class TraceRequestStartParams:
@@ -330,6 +341,15 @@ class TraceRequestHeadersSentParams:
     headers: "CIMultiDict[str]"
 
 
+@frozen_dataclass_decorator
+class TraceResponseHeadersReceivedParams:
+    """Parameters sent by the `on_response_headers_received` signal"""
+
+    method: str
+    url: URL
+    headers: "CIMultiDictProxy[str]"
+
+
 class Trace:
     """Internal dependency holder class.
 
@@ -466,3 +486,12 @@ async def send_request_headers(
             self._trace_config_ctx,
             TraceRequestHeadersSentParams(method, url, headers),
         )
+
+    async def send_response_headers_received(
+        self, method: str, url: URL, headers: "CIMultiDictProxy[str]"
+    ) -> None:
+        return await self._trace_config.on_response_headers_received.send(
+            self._session,
+            self._trace_config_ctx,
+            TraceResponseHeadersReceivedParams(method, url, headers),
+        )

tests/test_tracing.py

@@ -26,6 +26,7 @@
     TraceRequestRedirectParams,
     TraceRequestStartParams,
     TraceResponseChunkReceivedParams,
+    TraceResponseHeadersReceivedParams,
 )
 
 if sys.version_info >= (3, 11):
@@ -92,6 +93,7 @@ def test_freeze(self) -> None:
         assert trace_config.on_dns_cache_hit.frozen
         assert trace_config.on_dns_cache_miss.frozen
         assert trace_config.on_request_headers_sent.frozen
+        assert trace_config.on_response_headers_received.frozen
 
 
 class TestTrace:
@@ -129,6 +131,11 @@ class TestTrace:
             ("dns_resolvehost_end", (Mock(),), TraceDnsResolveHostEndParams),
             ("dns_cache_hit", (Mock(),), TraceDnsCacheHitParams),
             ("dns_cache_miss", (Mock(),), TraceDnsCacheMissParams),
+            (
+                "response_headers_received",
+                (Mock(), Mock(), Mock()),
+                TraceResponseHeadersReceivedParams,
+            ),
         ],
     )
     async def test_send(  # type: ignore[misc]