truss 0.11.1rc7__tar.gz → 0.11.1rc8__tar.gz

{truss-0.11.1rc7 → truss-0.11.1rc8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.1rc7
+Version: 0.11.1rc8
 Summary: A seamless bridge from model development to model delivery
 Project-URL: Repository, https://github.com/basetenlabs/truss
 Project-URL: Homepage, https://truss.baseten.co

{truss-0.11.1rc7 → truss-0.11.1rc8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "truss"
-version = "0.11.1rc7"
+version = "0.11.1rc8"
 description = "A seamless bridge from model development to model delivery"
 authors = [
     { name = "Pankaj Gupta", email = "no-reply@baseten.co" },

{truss-0.11.1rc7 → truss-0.11.1rc8}/truss/templates/control/control/endpoints.py

@@ -11,7 +11,8 @@ from starlette.requests import ClientDisconnect, Request
 from starlette.responses import Response
 from starlette.websockets import WebSocketDisconnect as StartletteWebSocketDisconnect
 from tenacity import RetryCallState, Retrying, retry_if_exception_type, wait_fixed
-from wsproto.events import BytesMessage, TextMessage
+from wsproto.connection import ConnectionState
+from wsproto.events import BytesMessage, CloseConnection, TextMessage
 
 from truss.templates.control.control.helpers.errors import (
     ModelLoadFailed,
@@ -158,7 +159,13 @@ async def forward_to_server(
 async def forward_to_client(client_ws: WebSocket, server_ws: AsyncWebSocketSession):
     while True:
         message = await server_ws.receive()
-        if isinstance(message, TextMessage):
+        if isinstance(message, CloseConnection):
+            # NB(nikhil): httpx requires explicit confirmation of the close frame.
+            if server_ws.connection.state == ConnectionState.REMOTE_CLOSING:
+                await server_ws.send(message.response())  # type: ignore[arg-type]
+
+            raise WebSocketDisconnect(code=message.code, reason=message.reason)
+        elif isinstance(message, TextMessage):
             await client_ws.send_text(message.data)
         elif isinstance(message, BytesMessage):
             await client_ws.send_bytes(message.data)