Pass through disconnection on ws.receive()

Author: nnarayen

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "truss"
-version = "0.11.1rc3"
+version = "0.11.1rc15"
 description = "A seamless bridge from model development to model delivery"
 authors = [
     { name = "Pankaj Gupta", email = "[email protected]" },

truss-chains/truss_chains/remote_chainlet/utils.py

@@ -588,8 +588,18 @@ async def close(self, code: int = 1000, reason: Optional[str] = None) -> None:
         await self._websocket.close(code=code, reason=reason)
 
     async def receive(self) -> Union[str, bytes]:
+        try:
+            import fastapi
+        except ImportError:
+            raise utils.make_optional_import_error("fastapi")
+
         message = await self._websocket.receive()
-        if message.get("text"):
+
+        if message.get("type") == "websocket.disconnect":
+            # NB(nikhil): Mimics FastAPI `_raise_on_disconnect`, since otherwise the user has no
+            # way of detecting that the client disconnected.
+            raise fastapi.WebSocketDisconnect(message["code"], message.get("reason"))
+        elif message.get("text"):
             return typing.cast(str, message["text"])
         else:
             return typing.cast(bytes, message["bytes"])

truss/templates/control/requirements.txt

@@ -6,7 +6,7 @@ loguru>=0.7.2
 python-json-logger>=2.0.2
 tenacity>=8.1.0
  # To avoid divergence, this should follow the latest release.
-truss==0.11.1rc3
+truss==0.11.1rc15
 uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0