ASGI: Wait for response to complete before sending disconnect message

httpx/_concurrency.py

@@ -0,0 +1,31 @@
+import typing
+
+import sniffio
+
+if typing.TYPE_CHECKING:  # pragma: no cover
+    try:
+        from typing import Protocol
+    except ImportError:
+        from typing_extensions import Protocol  # type: ignore
+
+    class Event(Protocol):
+        def set(self) -> None:
+            ...
+
+        # asyncio wait() returns True, but Trio returns None: ignore the return value.
+        async def wait(self) -> typing.Any:
+            ...
+
+        def is_set(self) -> bool:
+            ...
+
+
+def create_event() -> "Event":
+    if sniffio.current_async_library() == "trio":
+        import trio
+
+        return trio.Event()
+    else:
+        import asyncio
+
+        return asyncio.Event()

httpx/_dispatch/asgi.py

@@ -2,6 +2,7 @@
 
 import httpcore
 
+from .._concurrency import create_event
 from .._content_streams import ByteStream
 
 
@@ -76,23 +77,27 @@ async def request(
         status_code = None
         response_headers = None
         body_parts = []
+        request_complete = False
         response_started = False
-        response_complete = False
+        response_complete = create_event()
 
         headers = [] if headers is None else headers
         stream = ByteStream(b"") if stream is None else stream
 
         request_body_chunks = stream.__aiter__()
 
         async def receive() -> dict:
-            nonlocal response_complete
+            nonlocal request_complete, response_complete
 
-            if response_complete:
+            if request_complete:
+                # Simulate blocking until the response is complete and then disconnect
+                await response_complete.wait()
                 return {"type": "http.disconnect"}
 
             try:
                 body = await request_body_chunks.__anext__()
             except StopAsyncIteration:
+                request_complete = True
                 return {"type": "http.request", "body": b"", "more_body": False}
             return {"type": "http.request", "body": body, "more_body": True}
 
@@ -108,23 +113,23 @@ async def send(message: dict) -> None:
                 response_started = True
 
             elif message["type"] == "http.response.body":
-                assert not response_complete
+                assert not response_complete.is_set()
                 body = message.get("body", b"")
                 more_body = message.get("more_body", False)
 
                 if body and method != b"HEAD":
                     body_parts.append(body)
 
                 if not more_body:
-                    response_complete = True
+                    response_complete.set()
 
         try:
             await self.app(scope, receive, send)
         except Exception:
             if self.raise_app_exceptions or not response_complete:
                 raise
 
-        assert response_complete
+        assert response_complete.is_set()
         assert status_code is not None
         assert response_headers is not None
 

requirements.txt

@@ -26,6 +26,7 @@ pytest-asyncio
 pytest-trio
 pytest-cov
 trio
+trio-typing
 trustme
 uvicorn
 seed-isort-config

setup.py

@@ -61,6 +61,7 @@ def get_packages(package):
         "idna==2.*",
         "rfc3986>=1.3,<2",
         "httpcore>=0.8.1",
+        "sniffio",
     ],
     classifiers=[
         "Development Status :: 4 - Beta",

tests/test_asgi.py

@@ -69,8 +69,7 @@ async def test_asgi_exc_after_response():
         await client.get("http://www.example.org/")
 
 
-@pytest.mark.asyncio
-async def test_asgi_disconnect_after_response_complete():
+async def test_asgi_disconnect_after_response_complete(async_environment):
     disconnect = False
 
     async def read_body(scope, receive, send):