Better code propagation, better Chains websocket API #1918

nnarayen · 2025-09-09T03:58:27Z

🚀 What

Heavily based off #1912, adds two small improvements:

Propagates error codes bidirectionally instead of masking with 1000
Improves the API for dealing with websockets in Chains, exposing the more generic receive() -> str | bytes

💻 How

🔬 Testing

Unit tests
Testing with 0.11.1rc1

nnarayen · 2025-09-09T04:06:47Z

truss-chains/truss_chains/public_types.py


    async def close(self, code: int = 1000, reason: Optional[str] = None) -> None: ...

+    async def receive(self) -> Union[str, bytes]: ...


Context for this change is here: https://basetenlabs.slack.com/archives/C06CZ3RSXRU/p1757360869605229

nnarayen · 2025-09-09T04:07:30Z

truss/templates/control/control/endpoints.py

 control_app = APIRouter()


+class CloseableWebsocket(Protocol):


Maybe overkill, and mentioned through comments below, but it's unfortunate that the incoming websocket (FastAPI) and outbound websockets (httpx) are different types with incompatible APIs.

We get around that separately below, but luckily close is common, so I wrapped in a Protocol here

nnarayen · 2025-09-09T04:09:54Z

truss/templates/control/control/endpoints.py

+) -> None:
+    while True:
+        message = await client_ws.receive()
+        if message.get("type") == "websocket.disconnect":


The docs seem to suggest that the client closing the connection could either:

Trigger an StartletteWebSocketDisconnect

Send an explicit websocket.disconnect message, that we have to handle

To be safe, we re-raise a StartletteWebSocketDisconnect here and handle below

nnarayen · 2025-09-09T04:10:30Z

truss/templates/control/control/endpoints.py

+        await _safe_close_ws(server_ws, logger)
+
+
+async def _attempt_websocket_proxy(


This refactor was mainly to improve on the 6 layer indent for proxy_ws before, but there are some material changes above

nnarayen · 2025-09-09T04:10:52Z

truss/templates/control/control/endpoints.py

+        else:
+            logger.warning(f"Ungraceful websocket close: {exc}")
+    finally:
+        await _safe_close_ws(client_ws, logger)


We have protections on double close, but I want this in the finally just to make sure everything is cleaned up in all casese

Had this question as well, you might want to add a comment above just in case!

nnarayen · 2025-09-09T04:11:49Z

truss/templates/control/control/endpoints.py

+):
+    logger = client_ws.app.state.logger
+    try:
+        async with asyncio.TaskGroup() as tg:  # type: ignore[attr-defined]


Changing from asyncio.gather to TaskGroup means the other task gets automatically cancelled when one raises an exception, compared to running until it attempts to receive on a closed websocket

ndeepak-baseten · 2025-09-09T17:53:08Z

truss/templates/control/control/endpoints.py

+async def _safe_close_ws(
+    ws: CloseableWebsocket,
+    logger: logging.Logger,
+    code: int = 1000,


Cosmetic: here and in other places, could we use a constant like WEBSOCKET_NORMAL_CLOSE instead of 1000? Perhaps the websocket library already offers one.

ndeepak-baseten · 2025-09-09T18:04:32Z

truss/templates/control/control/endpoints.py

+        else:
+            logger.warning(f"Ungraceful websocket close: {exc}")
+    finally:
+        await _safe_close_ws(client_ws, logger)


Had this question as well, you might want to add a comment above just in case!

ndeepak-baseten · 2025-09-09T18:05:59Z

truss/templates/control/control/endpoints.py

+            tg.create_task(forward_to_client(client_ws, server_ws))
+            tg.create_task(forward_to_server(client_ws, server_ws))
+    except ExceptionGroup as eg:  # type: ignore[name-defined] # noqa: F821
+        exc = eg.exceptions[0]  # NB(nikhil): Only care about the first one.


Why the first one? Is it because the remaining ones are consequences of the first one? Worth mentioning that as well.

ndeepak-baseten · 2025-09-09T18:08:27Z

truss/templates/control/control/endpoints.py

-                        await asyncio.gather(forward_to_client(), forward_to_server())
-                    finally:
-                        await _safe_close_ws(client_ws, logger)
+                await _attempt_websocket_proxy(client_ws, proxy_client, logger)


Thanks for breaking this code block down into smaller functions!

ndeepak-baseten · 2025-09-09T18:10:37Z

truss/templates/server/truss_server.py

                    await ws.accept()
                    await self._model.websocket(ws)
-                    await _safe_close_websocket(ws, None, status_code=1000)
+                    await _safe_close_websocket(ws, status_code=1000, reason=None)


If we always specify arguments, do we still need default values for _safe_close_websocket()? It makes the code DRY-er and more explicit.

Unfortunately on L182 I have a failsafe close, but I can make code without a default!

ndeepak-baseten · 2025-09-09T18:10:53Z

truss/templates/server/truss_server.py

                except Exception:
                    await _safe_close_websocket(
-                        ws, errors.MODEL_ERROR_MESSAGE, status_code=1011
+                        ws, status_code=1011, reason=errors.MODEL_ERROR_MESSAGE


1011 might benefit from a constant string too.

ndeepak-baseten · 2025-09-09T18:14:56Z

truss/tests/templates/control/control/test_endpoints.py

    ):
-        await proxy_ws(client_ws)
+        proxy_task = asyncio.create_task(proxy_ws(client_ws))
+        await asyncio.sleep(0.5)


In the interest of reducing test flakiness, is there any pattern we could use instead of a sleep()?

I think I actually don't need this - removed!

nnarayen · 2025-09-10T03:19:29Z

truss-chains/truss_chains/remote_chainlet/utils.py

        await self._websocket.close(code=code, reason=reason)

+    async def receive(self) -> Union[str, bytes]:
+        try:


This is pulled from lines above, since we need to expose WebsocketDisconnect to end users.

This wrapper exists to try and hide FastAPI, but per a comment above we already knew we were leaking details via Exceptions. I'd rather not come up with a separate protocol for how users would have to think about catching websocket disconnects (especially since it would be different in Truss), so for now I just pass it through

squidarth · 2025-09-10T16:32:08Z

@nnarayen should we change the docs here?

nnarayen · 2025-09-10T16:48:04Z

@nnarayen should we change the docs here?

Good call! https://github.com/basetenlabs/docs.baseten.co/pull/263

nnarayen force-pushed the nikhil/propagate-error-codes branch 3 times, most recently from 9cf8b4c to 1a6cd89 Compare September 9, 2025 04:00

nnarayen commented Sep 9, 2025

View reviewed changes

nnarayen force-pushed the nikhil/propagate-error-codes branch from 1a6cd89 to 04bf0e5 Compare September 9, 2025 04:08

nnarayen commented Sep 9, 2025

View reviewed changes

nnarayen requested review from bryzhang and ndeepak-baseten September 9, 2025 17:04

ndeepak-baseten reviewed Sep 9, 2025

View reviewed changes

ndeepak-baseten approved these changes Sep 9, 2025

View reviewed changes

nnarayen force-pushed the nikhil/propagate-error-codes branch 2 times, most recently from bf9d00f to c3ab8e7 Compare September 9, 2025 22:18

nnarayen added 3 commits September 9, 2025 15:19

Better code propagation, better Websocket websocket API

abbfea3

Address comments

cc1d114

Push new version of control server

802e090

nnarayen force-pushed the nikhil/propagate-error-codes branch 4 times, most recently from 4d328b1 to 0f94794 Compare September 9, 2025 23:20

Pass through disconnection on ws.receive()

7bd86ad

nnarayen force-pushed the nikhil/propagate-error-codes branch from 0f94794 to 7bd86ad Compare September 10, 2025 03:11

nnarayen commented Sep 10, 2025

View reviewed changes

nnarayen merged commit ee8aed0 into main Sep 10, 2025
31 checks passed

nnarayen deleted the nikhil/propagate-error-codes branch September 10, 2025 03:22

nnarayen changed the title ~~Better code propagation, better Websocket websocket API~~ Better code propagation, better Chains websocket API Sep 10, 2025


		async def close(self, code: int = 1000, reason: Optional[str] = None) -> None: ...

		async def receive(self) -> Union[str, bytes]: ...

		control_app = APIRouter()


		class CloseableWebsocket(Protocol):

		await _safe_close_ws(server_ws, logger)


		async def _attempt_websocket_proxy(

Better code propagation, better Chains websocket API #1918

Better code propagation, better Chains websocket API #1918

Uh oh!

Conversation

nnarayen commented Sep 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🚀 What

💻 How

🔬 Testing

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nnarayen Sep 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

squidarth commented Sep 10, 2025

Uh oh!

nnarayen commented Sep 10, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nnarayen commented Sep 9, 2025 •

edited

Loading

nnarayen Sep 9, 2025 •

edited

Loading