fix: capped responses keep elapsed, next_request and default_encoding; capped redirect and auth hops reuse their connection (#179)

Author: lesnik512

docs/adr/0009-the-body-cap-counts-decoded-bytes.md

@@ -4,6 +4,8 @@
 status. `Content-Length` may reject early but never admit early, because it reports the compressed
 size and a 1000:1 gzip bomb passes a header check by construction. It replaced an error-only
 pre-check, since memory exhaustion has no status code, and it never applies to caller-driven
-`stream()` iteration, whose purpose is bodies too large to buffer. The rebuild through
-`httpx2.Response(content=...)` uses only public API and loses `response.elapsed` when a cap is set.
-That was accepted over reaching into httpx2 internals.
+`stream()` iteration, whose purpose is bodies too large to buffer. The response is rebuilt
+through `httpx2.Response(content=...)` with public API only: `elapsed`, `next_request` and
+`default_encoding` are copied over, and `num_bytes_downloaded`, which has no public setter, reads 0.
+Counting bytes in a wrapper around the response stream would avoid the rebuild; that is deferred to
+[#170](https://github.com/modern-python/httpware/issues/170).

docs/errors.md

@@ -200,7 +200,9 @@ Unlike `DecodeError`, this error is raised before the request is sent.
 
 ## `ResponseTooLargeError`
 
-Both clients accept `max_response_body_bytes: int | None = None`. By default there is no limit. When it is set, a response body larger than the cap raises `ResponseTooLargeError` instead of being returned, whatever the status: a `200` trips it as easily as a `500`. The cap counts decoded bytes, after decompression. It applies to `send()` and the verb methods, and to the error body that `stream()` reads before raising a `StatusError`. Bytes you read yourself while iterating a `stream()` are never capped, except with an `auth` that sets `requires_response_body`: that auth needs every body, so `stream()` reads the final one under the cap before yielding it. With a cap set, httpware runs the client's `auth` flow and, with `follow_redirects=True`, follows the redirects itself, and caps only the final response. It closes each intermediate response, such as a redirect or a `DigestAuth` challenge, without reading its body, so the responses in `response.history` have no content. An `auth` that sets `requires_response_body`, like a token refresh that parses the token endpoint's JSON, gets each response read under the cap instead. Client `auth` is sent to the first URL only. `httpx2` keeps its `Authorization` header on a redirect within the same origin or from `http` to `https` on the same host, and drops it otherwise.
+Both clients accept `max_response_body_bytes: int | None = None`. By default there is no limit. When it is set, a response body larger than the cap raises `ResponseTooLargeError` instead of being returned, whatever the status: a `200` trips it as easily as a `500`. The cap counts decoded bytes, after decompression. It applies to `send()` and the verb methods, and to the error body that `stream()` reads before raising a `StatusError`. Bytes you read yourself while iterating a `stream()` are never capped, except with an `auth` that sets `requires_response_body`: that auth needs every body, so `stream()` reads the final one under the cap before yielding it. With a cap set, httpware runs the client's `auth` flow and, with `follow_redirects=True`, follows the redirects itself, and caps only the final response. It reads the body of each intermediate response, such as a redirect or a `DigestAuth` challenge, and throws it away, so the connection can be reused and the responses in `response.history` have no content. An intermediate body declared or read past the cap is not an error: httpware stops reading it and closes that connection. An `auth` that sets `requires_response_body`, like a token refresh that parses the token endpoint's JSON, gets each response read under the cap instead. Client `auth` is sent to the first URL only. `httpx2` keeps its `Authorization` header on a redirect within the same origin or from `http` to `https` on the same host, and drops it otherwise.
+
+A response read under the cap is rebuilt from its decoded body. It keeps `elapsed`, `next_request` and the client's `default_encoding`. It loses the `Content-Encoding` and `Transfer-Encoding` headers, its `Content-Length` is the decoded size, and `num_bytes_downloaded` is `0`.
 
 `ResponseTooLargeError` carries:
 

src/httpware/_internal/body_cap.py

@@ -1,5 +1,6 @@
 """Response-body cap enforcement: validate, read-capped (sync + async)."""
 
+import contextlib
 import typing
 from collections.abc import Mapping
 from http import HTTPStatus
@@ -98,24 +99,29 @@ def _reject_if_declared_length_exceeds_cap(response: httpx2.Response, cap: int)
 
 
 def _build_buffered_response(response: httpx2.Response, content: bytes, request: httpx2.Request) -> httpx2.Response:
-    """Rebuild a buffered Response from the original streaming `response` and the final decoded `content`."""
-    return httpx2.Response(
+    """Rebuild a buffered Response from the closed streaming `response` and its final decoded `content`."""
+    buffered = httpx2.Response(
         status_code=response.status_code,
         headers=_buffered_headers(response.headers),
         content=content,
         request=request,
         extensions=_safe_extensions(response.extensions),
         history=response.history,
+        default_encoding=response.default_encoding,
     )
+    buffered.next_request = response.next_request
+    with contextlib.suppress(RuntimeError):  # httpx2 has no elapsed for a stream it never closed
+        buffered.elapsed = response.elapsed
+    return buffered
 
 
 def _read_capped(response: httpx2.Response, cap: int, request: httpx2.Request) -> httpx2.Response:
     """Buffer a streaming sync `response` under `cap` decoded bytes; return a buffered Response.
 
     Raises `ResponseTooLargeError` (reason="declared") if the declared
     Content-Length already exceeds `cap` — before any byte is read — and
-    (reason="streamed") if the decoded body crosses `cap` mid-read. Does not
-    close `response`; the caller owns the stream lifecycle.
+    (reason="streamed") if the decoded body crosses `cap` mid-read. Closes
+    `response` once its whole body is read; on any other path the caller closes it.
     """
     if not _response_has_body(request.method, response.status_code):
         response.read()  # empty body; preserve the original response (and its headers)
@@ -127,6 +133,7 @@ def _read_capped(response: httpx2.Response, cap: int, request: httpx2.Request) -
         raise ResponseTooLargeError(
             status_code=response.status_code, limit=cap, content_length=content_length, reason="streamed"
         ) from None
+    response.close()
     return _build_buffered_response(response, content, request)
 
 
@@ -143,4 +150,36 @@ async def _read_capped_async(response: httpx2.Response, cap: int, request: httpx
             raise ResponseTooLargeError(
                 status_code=response.status_code, limit=cap, content_length=content_length, reason="streamed"
             )
+    await response.aclose()
     return _build_buffered_response(response, bytes(buf), request)
+
+
+def _fits_declared_length(response: httpx2.Response, cap: int) -> bool:
+    """Whether `response` may carry a body within `cap`: no body, or no Content-Length above it."""
+    if not _response_has_body(response.request.method, response.status_code):
+        return True
+    content_length = _parse_content_length(response.headers.get("content-length"))
+    return content_length is None or content_length <= cap
+
+
+def _discard_within_cap(response: httpx2.Response, cap: int) -> None:
+    """Read and discard `response`'s body, stopping once it is declared or read past `cap` decoded bytes."""
+    if not _fits_declared_length(response, cap):
+        return
+    read = 0
+    for chunk in response.iter_bytes():
+        read += len(chunk)
+        if read > cap:
+            return
+
+
+async def _discard_within_cap_async(response: httpx2.Response, cap: int) -> None:
+    """Async mirror of `_discard_within_cap`."""
+    if not _fits_declared_length(response, cap):
+        return
+    read = 0
+    async with contextlib.aclosing(response.aiter_bytes()) as chunks:
+        async for chunk in chunks:
+            read += len(chunk)
+            if read > cap:
+                return

src/httpware/client.py

@@ -11,7 +11,13 @@
 import typing_extensions
 
 from httpware._internal import import_checker
-from httpware._internal.body_cap import _read_capped, _read_capped_async, _validate_max_response_body_bytes
+from httpware._internal.body_cap import (
+    _discard_within_cap,
+    _discard_within_cap_async,
+    _read_capped,
+    _read_capped_async,
+    _validate_max_response_body_bytes,
+)
 from httpware._internal.exception_mapping import (
     _httpx2_exception_mapper,
     _httpx2_exception_mapper_sync,
@@ -165,12 +171,22 @@ async def _read_capped_and_close_async(streaming: httpx2.Response, cap: int) ->
         await streaming.aclose()
 
 
+async def _discard_and_close_async(response: httpx2.Response, cap: int) -> None:
+    """Close an intermediate `response`, first reading an unread body that fits `cap` so its connection is reused."""
+    try:
+        if not response.is_stream_consumed:
+            await _discard_within_cap_async(response, cap)
+    finally:
+        await response.aclose()
+
+
 async def _send_redirect_hops_async(
     client: httpx2.AsyncClient,
     request: httpx2.Request,
     prior_history: list[httpx2.Response],
+    cap: int,
 ) -> httpx2.Response:
-    """Send `request` streaming, following redirects hop by hop without reading intermediate bodies.
+    """Send `request` streaming, following redirects hop by hop and keeping no intermediate body.
 
     Histories and the `max_redirects` count include `prior_history`, as in httpx2.
     """
@@ -182,13 +198,13 @@ async def _send_redirect_hops_async(
         response.history = list(hops)
         if not client.follow_redirects or response.next_request is None:
             return response
-        await response.aclose()
+        await _discard_and_close_async(response, cap)
         hops.append(response)
         request = response.next_request
 
 
 async def _send_capped_async(client: httpx2.AsyncClient, request: httpx2.Request, cap: int) -> httpx2.Response:
-    """Send `request` streaming, driving the client's auth flow and redirects without reading intermediate bodies.
+    """Send `request` streaming, driving the client's auth flow and redirects and keeping no intermediate body.
 
     An auth that sets `requires_response_body` gets each response buffered under `cap` instead.
     """
@@ -198,7 +214,7 @@ async def _send_capped_async(client: httpx2.AsyncClient, request: httpx2.Request
     try:
         request = await anext(flow)
         while True:
-            response = await _send_redirect_hops_async(client, request, history)
+            response = await _send_redirect_hops_async(client, request, history, cap)
             if auth.requires_response_body:
                 response = await _read_capped_and_close_async(response, cap)
             try:
@@ -208,7 +224,7 @@ async def _send_capped_async(client: httpx2.AsyncClient, request: httpx2.Request
             except BaseException:
                 await response.aclose()
                 raise
-            await response.aclose()
+            await _discard_and_close_async(response, cap)
             response.history = list(history)
             history.append(response)
             request = next_request
@@ -252,10 +268,20 @@ def _read_capped_and_close(streaming: httpx2.Response, cap: int) -> httpx2.Respo
         streaming.close()
 
 
+def _discard_and_close(response: httpx2.Response, cap: int) -> None:
+    """Sync mirror of `_discard_and_close_async`."""
+    try:
+        if not response.is_stream_consumed:
+            _discard_within_cap(response, cap)
+    finally:
+        response.close()
+
+
 def _send_redirect_hops(
     client: httpx2.Client,
     request: httpx2.Request,
     prior_history: list[httpx2.Response],
+    cap: int,
 ) -> httpx2.Response:
     """Sync mirror of `_send_redirect_hops_async`."""
     hops = list(prior_history)
@@ -266,7 +292,7 @@ def _send_redirect_hops(
         response.history = list(hops)
         if not client.follow_redirects or response.next_request is None:
             return response
-        response.close()
+        _discard_and_close(response, cap)
         hops.append(response)
         request = response.next_request
 
@@ -279,7 +305,7 @@ def _send_capped(client: httpx2.Client, request: httpx2.Request, cap: int) -> ht
     try:
         request = next(flow)
         while True:
-            response = _send_redirect_hops(client, request, history)
+            response = _send_redirect_hops(client, request, history, cap)
             if auth.requires_response_body:
                 response = _read_capped_and_close(response, cap)
             try:
@@ -289,7 +315,7 @@ def _send_capped(client: httpx2.Client, request: httpx2.Request, cap: int) -> ht
             except BaseException:
                 response.close()
                 raise
-            response.close()
+            _discard_and_close(response, cap)
             response.history = list(history)
             history.append(response)
             request = next_request

tests/test_client_body_cap.py

@@ -5,8 +5,10 @@
 coverage lives in tests/test_client_stream*.py.
 """
 
+import datetime as dt
 import gzip
 from collections.abc import AsyncIterator
+from http import HTTPStatus
 
 import httpx2
 import pytest
@@ -204,3 +206,74 @@ def handler(request: httpx2.Request) -> httpx2.Response:  # noqa: ARG001
     request = client.build_request("GET", "https://example.test/x")
     assert (await client.send(request)).content == body
     await client.aclose()
+
+
+# ---- response attributes kept with or without a cap ----
+
+
+def _latin1_body(request: httpx2.Request) -> httpx2.Response:  # noqa: ARG001
+    return httpx2.Response(HTTPStatus.OK, headers={"content-type": "text/plain"}, content="café".encode("latin-1"))
+
+
+def _streamed_body(request: httpx2.Request) -> httpx2.Response:  # noqa: ARG001
+    return httpx2.Response(HTTPStatus.OK, content=iter([b"hello"]))
+
+
+async def _hello() -> AsyncIterator[bytes]:
+    yield b"hello"
+
+
+def _streamed_body_async(request: httpx2.Request) -> httpx2.Response:  # noqa: ARG001
+    return httpx2.Response(HTTPStatus.OK, content=_hello())
+
+
+def _redirect(request: httpx2.Request) -> httpx2.Response:  # noqa: ARG001
+    return httpx2.Response(HTTPStatus.FOUND, headers={"location": "/next"})
+
+
+@pytest.mark.parametrize("cap", [None, 1024])
+def test_sync_response_keeps_elapsed_with_or_without_a_cap(cap: int | None) -> None:
+    with Client(transport=httpx2.MockTransport(_streamed_body), max_response_body_bytes=cap) as client:
+        response = client.get("https://example.test/x")
+    assert response.elapsed >= dt.timedelta(0)
+
+
+@pytest.mark.parametrize("cap", [None, 1024])
+def test_sync_response_keeps_next_request_with_or_without_a_cap(cap: int | None) -> None:
+    with Client(transport=httpx2.MockTransport(_redirect), max_response_body_bytes=cap) as client:
+        response = client.get("https://example.test/x")
+    assert response.next_request is not None
+    assert str(response.next_request.url) == "https://example.test/next"
+
+
+@pytest.mark.parametrize("cap", [None, 1024])
+def test_sync_response_keeps_default_encoding_with_or_without_a_cap(cap: int | None) -> None:
+    with Client(
+        transport=httpx2.MockTransport(_latin1_body), default_encoding="latin-1", max_response_body_bytes=cap
+    ) as client:
+        response = client.get("https://example.test/x")
+    assert response.text == "café"
+
+
+@pytest.mark.parametrize("cap", [None, 1024])
+async def test_async_response_keeps_elapsed_with_or_without_a_cap(cap: int | None) -> None:
+    async with AsyncClient(transport=httpx2.MockTransport(_streamed_body_async), max_response_body_bytes=cap) as client:
+        response = await client.get("https://example.test/x")
+    assert response.elapsed >= dt.timedelta(0)
+
+
+@pytest.mark.parametrize("cap", [None, 1024])
+async def test_async_response_keeps_next_request_with_or_without_a_cap(cap: int | None) -> None:
+    async with AsyncClient(transport=httpx2.MockTransport(_redirect), max_response_body_bytes=cap) as client:
+        response = await client.get("https://example.test/x")
+    assert response.next_request is not None
+    assert str(response.next_request.url) == "https://example.test/next"
+
+
+@pytest.mark.parametrize("cap", [None, 1024])
+async def test_async_response_keeps_default_encoding_with_or_without_a_cap(cap: int | None) -> None:
+    async with AsyncClient(
+        transport=httpx2.MockTransport(_latin1_body), default_encoding="latin-1", max_response_body_bytes=cap
+    ) as client:
+        response = await client.get("https://example.test/x")
+    assert response.text == "café"