Checks if the client is still connected. Returns True if connected, False if disconnected.
(req_id: str, http_request: Request)
| 8 | |
| 9 | |
| 10 | async def check_client_connection(req_id: str, http_request: Request) -> bool: |
| 11 | """ |
| 12 | Checks if the client is still connected. |
| 13 | Returns True if connected, False if disconnected. |
| 14 | """ |
| 15 | try: |
| 16 | if hasattr(http_request, "_receive"): |
| 17 | try: |
| 18 | # Use a very short timeout to check for disconnect message |
| 19 | # _receive is a private Starlette/FastAPI method that returns a coroutine |
| 20 | receive_obj = http_request # type: ignore[misc] |
| 21 | receive_coro: Coroutine[Any, Any, Dict[str, Any]] = ( |
| 22 | receive_obj._receive() |
| 23 | ) # type: ignore[misc] |
| 24 | receive_task: Task[Dict[str, Any]] = asyncio.create_task(receive_coro) |
| 25 | done, pending = await asyncio.wait([receive_task], timeout=0.01) |
| 26 | |
| 27 | if done: |
| 28 | message = receive_task.result() |
| 29 | if message.get("type") == "http.disconnect": |
| 30 | return False |
| 31 | else: |
| 32 | # Cancel the task if it didn't complete immediately |
| 33 | receive_task.cancel() |
| 34 | try: |
| 35 | await receive_task |
| 36 | except asyncio.CancelledError: |
| 37 | pass |
| 38 | # If it didn't complete immediately, proceed to fallback check |
| 39 | except asyncio.CancelledError: |
| 40 | raise |
| 41 | except Exception: |
| 42 | # If checking fails, proceed to fallback |
| 43 | pass |
| 44 | |
| 45 | # Fallback to is_disconnected() if available (Starlette/FastAPI) |
| 46 | # Wrap in wait_for to prevent infinite hang in some ASGI implementations |
| 47 | if hasattr(http_request, "is_disconnected"): |
| 48 | try: |
| 49 | # Handle both sync and async versions for better mock compatibility |
| 50 | res = http_request.is_disconnected() |
| 51 | if asyncio.iscoroutine(res): |
| 52 | if await asyncio.wait_for(res, timeout=0.01): |
| 53 | return False |
| 54 | elif res: |
| 55 | return False |
| 56 | except (asyncio.TimeoutError, asyncio.CancelledError): |
| 57 | # If it times out, it's likely still connected |
| 58 | return True |
| 59 | |
| 60 | return True |
| 61 | except asyncio.CancelledError: |
| 62 | raise |
| 63 | except Exception as e: |
| 64 | # Re-raise to allow caller to log/handle |
| 65 | raise e |
| 66 | |
| 67 |