fix: preserve HTTPX response streams during instrumentation

src/harness_sdk/instrumentation/httpx/__init__.py

@@ -6,11 +6,8 @@
 from harness_sdk.plugins.control import get_control_registry
 from harness_sdk.instrumentation import BaseInstrumentorWrapper
 from harness_sdk.instrumentation.httpx.utils import (
-    decode_response_body_for_capture,
     headers_from_httpx,
     read_request_body,
-    read_response_body,
-    read_response_body_async,
     url_from_request_info,
 )
 
@@ -26,6 +23,7 @@ def __init__(self):
         logger.debug('Entering HTTPXClientInstrumentorWrapper.__init__().')
         HTTPXClientInstrumentor.__init__(self)
         BaseInstrumentorWrapper.__init__(self)
+        self._process_response_body = False
 
     def _instrument(self, **kwargs) -> None:
         '''Enable instrumentation with request/response hooks.'''
@@ -46,28 +44,20 @@ def _process_request(self, span, request_info):
 
     def _process_response(self, span, response_info):
         headers = headers_from_httpx(response_info.headers)
-        body = read_response_body(response_info.stream)
-        body = decode_response_body_for_capture(headers, body)
-        self.generic_response_handler(headers, body, span)
-
-    async def _process_response_async(self, span, response_info):
-        headers = headers_from_httpx(response_info.headers)
-        body = await read_response_body_async(response_info.stream)
-        body = decode_response_body_for_capture(headers, body)
-        self.generic_response_handler(headers, body, span)
+        self.generic_response_handler(headers, None, span)
 
     def request_hook(self, span, request_info):
         '''Capture sync client request data and run evaluation.'''
         self._process_request(span, request_info)
 
     def response_hook(self, span, request_info, response_info):  # pylint: disable=unused-argument
-        '''Capture sync client response data.'''
+        '''Capture sync client response headers.'''
         self._process_response(span, response_info)
 
     async def async_request_hook(self, span, request_info):
         '''Capture async client request data and run evaluation.'''
         self._process_request(span, request_info)
 
     async def async_response_hook(self, span, request_info, response_info):  # pylint: disable=unused-argument
-        '''Capture async client response data.'''
-        await self._process_response_async(span, response_info)
+        '''Capture async client response headers.'''
+        self._process_response(span, response_info)

src/harness_sdk/instrumentation/httpx/utils.py

@@ -1,6 +1,4 @@
 '''Helpers for extracting httpx request/response data for tracing.'''
-import gzip
-import zlib
 
 from harness_sdk.custom_logger import get_custom_logger
 
@@ -19,65 +17,6 @@ def url_from_request_info(request_info):
     return str(request_info.url)
 
 
-def _content_encoding_chain(headers):
-    '''Return Content-Encoding tokens in wire order (outer coding first).'''
-    if not headers:
-        return []
-    lower = {str(k).lower(): v for k, v in headers.items()}
-    raw = lower.get('content-encoding') or ''
-    return [t.strip().lower() for t in raw.split(',') if t.strip()]
-
-
-def _apply_content_encoding_decodings(original, chain):
-    '''Apply decodings in reverse wire order; return ``original`` on failure.'''
-    out = original
-    for encoding in reversed(chain):
-        try:
-            if encoding in ('gzip', 'x-gzip'):
-                out = gzip.decompress(out)
-            elif encoding == 'deflate':
-                try:
-                    out = zlib.decompress(out, -zlib.MAX_WBITS)
-                except zlib.error:
-                    out = zlib.decompress(out)
-            elif encoding == 'br':
-                try:
-                    import brotli  # pylint: disable=import-outside-toplevel
-                except ImportError:
-                    logger.debug('brotli not installed; skipping br decompression for capture')
-                    return original
-                out = brotli.decompress(out)
-            elif encoding in ('identity', 'compress'):
-                continue
-            else:
-                logger.debug('Unsupported content-encoding for capture: %s', encoding)
-                return original
-        except Exception:  # pylint: disable=broad-except
-            logger.debug('Response body decompression failed (%s)', encoding, exc_info=True)
-            return original
-    return out
-
-
-def decode_response_body_for_capture(headers, body):
-    '''Decode compressed response bytes for span capture (gzip/deflate/br).
-
-    Transport hooks often see the on-the-wire body while ``Content-Encoding``
-    still names the compression. httpx may also surface compressed bytes when
-    decoding is deferred. If decoding fails or an encoding is unsupported,
-    the original ``body`` is returned.
-    '''
-    if body in (None, b''):
-        return body
-    if not isinstance(body, (bytes, bytearray)):
-        return body
-
-    chain = _content_encoding_chain(headers)
-    if not chain:
-        return body
-
-    return _apply_content_encoding_decodings(bytes(body), chain)
-
-
 def _body_from_byte_stream(stream):
     internal = getattr(stream, "_stream", None)
     if isinstance(internal, (bytes, bytearray)):
@@ -99,62 +38,3 @@ def read_request_body(stream):
     except Exception:  # pylint: disable=broad-except
         logger.debug("Unable to read httpx request stream body", exc_info=True)
         return None
-
-
-def read_response_body(stream):
-    '''Read response body bytes and restore the stream for downstream consumers.'''
-    if stream is None:
-        return None
-
-    body = _body_from_byte_stream(stream)
-    if body is not None:
-        return body
-
-    httpcore_stream = getattr(stream, "_httpcore_stream", None)
-    if httpcore_stream is not None:
-        try:
-            chunks = list(httpcore_stream)
-            body = b"".join(chunks)
-            stream._httpcore_stream = chunks  # pylint: disable=protected-access
-            return body
-        except Exception:  # pylint: disable=broad-except
-            logger.debug("Unable to read httpx response stream body", exc_info=True)
-            return None
-
-    return None
-
-
-async def read_response_body_async(stream):
-    '''Read async response body bytes and restore the stream for downstream consumers.'''
-    if stream is None:
-        return None
-
-    body = _body_from_byte_stream(stream)
-    if body is not None:
-        return body
-
-    httpcore_stream = getattr(stream, "_httpcore_stream", None)
-    if httpcore_stream is not None:
-        try:
-            chunks = []
-            async for part in httpcore_stream:
-                chunks.append(part)
-            body = b"".join(chunks)
-            stream._httpcore_stream = _ReplayAsyncStream(chunks)  # pylint: disable=protected-access
-            return body
-        except Exception:  # pylint: disable=broad-except
-            logger.debug("Unable to read httpx async response stream body", exc_info=True)
-            return None
-
-    return None
-
-
-class _ReplayAsyncStream:  # pylint: disable=too-few-public-methods
-    '''Minimal async iterable used to replay captured response chunks.'''
-
-    def __init__(self, chunks):
-        self._chunks = chunks
-
-    async def __aiter__(self):
-        for chunk in self._chunks:
-            yield chunk

test/instrumentation/httpx_client/httpx_integration_test.py

@@ -9,11 +9,12 @@
 from test.instrumentation.flask.app import FlaskServer
 
 
-def _client_span(spans):
+def _client_span(spans, url):
     for span in spans:
-        if span.kind == SpanKind.CLIENT:
-            return json.loads(span.to_json())
-    raise AssertionError('No client span found')
+        span_data = json.loads(span.to_json())
+        if span.kind == SpanKind.CLIENT and span_data['attributes'].get('http.url') == url:
+            return span_data
+    raise AssertionError(f'No client span found for {url}')
 
 
 def test_httpx_client_get(agent, exporter):
@@ -38,11 +39,11 @@ def api_example():
 
         spans = exporter.get_finished_spans()
         assert spans
-        client_span = _client_span(spans)
+        client_span = _client_span(spans, url)
 
         assert client_span['attributes']['http.method'] == 'GET'
         assert client_span['attributes']['http.url'] == url
-        assert client_span['attributes']['http.response.body'] == '{ "a": "a", "xyz": "xyz" }'
+        assert 'http.response.body' not in client_span['attributes']
         assert client_span['attributes']['http.status_code'] == 200
         assert client_span['attributes']['http.response.header.tester3'] == 'tester3'
     finally:
@@ -71,15 +72,16 @@ def api_example():
 
         spans = exporter.get_finished_spans()
         assert spans
-        client_span = _client_span(spans)
+        client_span = _client_span(spans, url)
 
         assert client_span['kind'] == "SpanKind.CLIENT"
         assert client_span['attributes']['http.method'] == 'POST'
         assert client_span['attributes']['http.url'] == url
         assert client_span['attributes']['http.request.header.content-type'] == 'application/json'
         assert client_span['attributes']['http.request.body'] == '{"test":"body"}'
-        assert client_span['attributes']['http.response.body'] == '{ "a": "a", "xyz": "xyz" }'
+        assert 'http.response.body' not in client_span['attributes']
         assert client_span['attributes']['http.status_code'] == 200
+        assert client_span['attributes']['http.response.header.tester3'] == 'tester3'
     finally:
         server.shutdown()
 
@@ -114,14 +116,14 @@ def api_example():
 
     span_list = exporter.get_finished_spans()
     assert span_list
-    httpx_span = _client_span(span_list)
+    httpx_span = _client_span(span_list, url)
 
     assert httpx_span['attributes']['http.method'] == 'POST'
     assert httpx_span['attributes']['http.url'] == url
     assert httpx_span['attributes']['http.request.header.tester1'] == 'tester1'
     assert httpx_span['attributes']['http.request.header.tester2'] == 'tester2'
     assert httpx_span['attributes']['http.request.body'] == '{ "a":"b", "c": "d" }'
     assert httpx_span['attributes']['http.response.header.content-type'] == 'application/json'
-    assert httpx_span['attributes']['http.response.body'] == '{ "a": "a", "xyz": "xyz" }'
+    assert 'http.response.body' not in httpx_span['attributes']
     assert httpx_span['attributes']['http.status_code'] == 200
     server.shutdown()

test/instrumentation/httpx_client/test_httpx_hooks.py

@@ -0,0 +1,71 @@
+from types import SimpleNamespace
+
+import pytest
+
+from harness_sdk.instrumentation.httpx import HTTPXClientInstrumentorWrapper
+
+
+class _SyncStream:
+    def __init__(self):
+        self.iterations = 0
+
+    def __iter__(self):
+        self.iterations += 1
+        yield b"response"
+
+
+class _AsyncStream:
+    def __init__(self):
+        self.iterations = 0
+
+    async def __aiter__(self):
+        self.iterations += 1
+        yield b"response"
+
+
+def _wrapper_with_response_capture():
+    wrapper = HTTPXClientInstrumentorWrapper()
+    calls = []
+    wrapper.generic_response_handler = (
+        lambda headers, body, span: calls.append((headers, body, span))
+    )
+    return wrapper, calls
+
+
+def test_response_hook_does_not_consume_or_replace_stream():
+    wrapper, calls = _wrapper_with_response_capture()
+    span = object()
+    transport_stream = _SyncStream()
+    response_stream = SimpleNamespace(_httpcore_stream=transport_stream)
+    response_info = SimpleNamespace(
+        status_code=200,
+        headers={"content-type": "text/plain"},
+        stream=response_stream,
+        extensions={},
+    )
+
+    wrapper.response_hook(span, None, response_info)
+
+    assert calls == [({"content-type": "text/plain"}, None, span)]
+    assert transport_stream.iterations == 0
+    assert response_stream._httpcore_stream is transport_stream
+
+
+@pytest.mark.asyncio
+async def test_async_response_hook_does_not_consume_or_replace_stream():
+    wrapper, calls = _wrapper_with_response_capture()
+    span = object()
+    transport_stream = _AsyncStream()
+    response_stream = SimpleNamespace(_httpcore_stream=transport_stream)
+    response_info = SimpleNamespace(
+        status_code=200,
+        headers={"content-type": "text/plain"},
+        stream=response_stream,
+        extensions={},
+    )
+
+    await wrapper.async_response_hook(span, None, response_info)
+
+    assert calls == [({"content-type": "text/plain"}, None, span)]
+    assert transport_stream.iterations == 0
+    assert response_stream._httpcore_stream is transport_stream