WebSockets

Enable real-time, streaming, bidirectional communication using WebSockets for Truss models and Chains.

Use WebSockets to send and receive messages over a persistent connection to a model or Chain. Each connection stays on one replica, so your model or Chain can keep state between messages.

WebSockets in Truss models

A Truss WebSocket model implements a single websocket method in place of the usual preprocess, predict, and postprocess methods. All input and output flows through the WebSocket object itself, not through arguments or return values. load still works as it does for HTTP models.

1. Initialize your Truss:

   Terminal

   truss init websocket-model
   cd websocket-model
   

2. Replace the predict method in model/model.py with a websocket method. For example:

   model/model.py

   import fastapi
   
   class Model:
       async def websocket(self, websocket: fastapi.WebSocket):
           try:
               while True:
                   message = await websocket.receive_text()
                   await websocket.send_text(f"WS obtained: {message}")
           except fastapi.WebSocketDisconnect:
               pass
   

3. Set runtime.transport.kind to websocket in config.yaml:

   config.yaml

   runtime:
     transport:
       kind: websocket
   

4. Deploy the model:

   Baseten CLI

   baseten model push
   

This creates a published deployment. For live-reload during development, use baseten model push --watch.

For more information, see truss init and baseten model push.

Constraints and behavior

• Your websocket method controls message exchange. The example loops until the client disconnects. To close the connection from the server, call await websocket.close().
• Define websocket as an async method with one WebSocket argument.
• Don’t implement predict, preprocess, or postprocess alongside websocket. Truss rejects models that combine these methods.
• Baseten accepts the connection for you, so don’t call websocket.accept(). You can close the connection yourself when you’re done; otherwise Baseten closes it after your websocket method returns.

Call the model

Set BASETEN_API_KEY to your API key. Replace {MODEL_ID} with your model ID, then use websocat to connect:

Terminal

websocat -H="Authorization: Bearer $BASETEN_API_KEY" \
   "wss://model-{MODEL_ID}.api.baseten.co/environments/production/websocket"

Enter Hello. The model replies with WS obtained: Hello. Press Ctrl+C to close the connection.

The path depends on the environment or deployment you’re calling:

• Environment: wss://model-{MODEL_ID}.api.baseten.co/environments/{ENVIRONMENT_NAME}/websocket
• Deployment: wss://model-{MODEL_ID}.api.baseten.co/deployment/{DEPLOYMENT_ID}/websocket
• Regional environment: wss://model-{MODEL_ID}-{ENV_NAME}.api.baseten.co/websocket. See Regional environments.
• Regional deployment: wss://model-{MODEL_ID}-region-{REGION_SLUG}.api.baseten.co/websocket. See Regional deployments.

See the WebSocket endpoint reference for full details.

WebSockets in Chains

Chains wrap WebSockets in a reduced WebSocketProtocol object. Processing happens in run_remote as usual, but inputs and outputs both flow through the WebSocket itself using async send_* and receive_* methods (text, bytes, and json variants). A convenience receive method handles both str and bytes.

Example chainlet

chainlet.py

import fastapi
import truss_chains as chains

class Dependency(chains.ChainletBase):
    async def run_remote(self, name: str) -> str:
        return f"Hello from dependency, {name}."

@chains.mark_entrypoint
class WSEntrypoint(chains.ChainletBase):
    def __init__(self, dependency=chains.depends(Dependency)):
        self._dependency = dependency

    async def run_remote(self, websocket: chains.WebSocketProtocol) -> None:
        try:
            while True:
                message = await websocket.receive_text()
                if message == "dep":
                    response = await self._dependency.run_remote("WSEntrypoint")
                else:
                    response = f"You said: {message}"
                await websocket.send_text(response)
        except fastapi.WebSocketDisconnect:
            print("Disconnected.")

Constraints and behavior

• Define run_remote as an async method with a WebSocketProtocol argument. This protocol exposes a subset of fastapi.WebSocket; you can’t call accept(). Baseten has already accepted the connection by the time your chainlet runs.
• run_remote accepts no other arguments when using WebSockets.
• The return type must be None. Send any data back to the client through the WebSocket instead.
• WebSockets are only supported on the entrypoint chainlet, not on dependencies.
• Unlike Truss models, Chains don’t require you to set runtime.transport.kind.

Call the Chain

Save the example as chainlet.py, then deploy the Chain to its production environment:

Terminal

truss chains push chainlet.py --environment production

With BASETEN_API_KEY set, replace {CHAIN_ID} with your Chain ID and use websocat to connect:

Terminal

websocat -H="Authorization: Bearer $BASETEN_API_KEY" \
   "wss://chain-{CHAIN_ID}.api.baseten.co/environments/production/websocket"

Enter dep. The Chain replies with Hello from dependency, WSEntrypoint.. Press Ctrl+C to close the connection.

Like models, chains accept WebSocket connections on either a deployment or environment path. For regional environments, use wss://chain-{CHAIN_ID}-{ENV_NAME}.api.baseten.co/websocket. See Regional environments.

See the WebSocket endpoint reference for full details.

WebSockets with custom servers

Deploy a WebSocket server from a custom Docker image using the docker_server configuration. This fits when you already have a WebSocket server packaged as a container, or when you need a runtime Baseten’s managed images don’t provide.

Configuration

Set the following in config.yaml:

config.yaml

base_image:
  image: bryanzhang2/custom_ws:v0.0.4
docker_server:
  start_command: /app/server
  readiness_endpoint: /health
  liveness_endpoint: /health
  predict_endpoint: /v1/websocket
  server_port: 8081
model_name: custom_ws
runtime:
  transport:
    kind: "websocket"

Set predict_endpoint to your server’s WebSocket route and runtime.transport.kind to websocket. Use HTTP routes for the readiness and liveness probes. See the docker_server configuration reference for field definitions.

Deploy from the directory containing config.yaml:

Terminal

baseten model push

Call the model

With BASETEN_API_KEY set, replace {MODEL_ID} with your model ID and use websocat to connect:

Terminal

websocat -H="Authorization: Bearer $BASETEN_API_KEY" \
   "wss://model-{MODEL_ID}.api.baseten.co/environments/production/websocket"

Baseten routes the connection to the predict_endpoint path on your server.

For more on custom server deployment, see Custom servers.

Deployment and concurrency considerations

Scheduling

An established WebSocket connection stays on its replica and counts toward that replica’s concurrency while open, including idle periods. Configure autoscaling for the number of simultaneous connections your deployment needs. Close idle connections when your application no longer needs them.

Connection lifetime

Connections can close when the client or server ends the session, a replica restarts, or a network failure occurs. Design clients to reconnect and restore any application state they need. A new connection can reach a different replica.

Concurrency changes

Lowering the concurrency target doesn’t close existing connections. Open WebSockets keep running until they close naturally, even if a replica ends up above the new target.

For example, if a replica holds 10 active WebSockets and you change concurrency_target from 10 to 5, Baseten leaves all 10 open. They drain as clients or servers close their connections.

Promotion

You can promote a WebSocket model or chain to an environment through the REST API or UI, the same way you promote HTTP deployments. For more information, see Environments.

On promotion, Baseten routes new connections to the new deployment, but existing connections stay on the previous deployment until they terminate. This means older deployments can take longer to scale down than HTTP deployments: their connections outlive the promotion.

Maximum message size

Baseten enforces a 100 MiB limit on individual messages sent over a WebSocket. Both clients and models are capped at 100 MiB per outgoing message. There’s no cap on the total data sent over a connection’s lifetime.

Monitoring

Use the Metrics page to monitor connection volume, duration, and transferred bytes. WebSocket measurements describe connections rather than individual messages.

Inference volume

The Metrics page tracks inference volume as the number of connections per minute. Baseten publishes each data point after the connection closes, so every point carries the status the connection ended with.

Two families of status codes appear for WebSocket deployments:

• HTTP status codes for connections that failed before the WebSocket upgrade completed.
• WebSocket close codes for connections that completed the upgrade and later closed.

HTTP status codes

Code	Label	What it means
408	Request timeout	The connection request timed out or was canceled before the upgrade completed.
504	Gateway timeout	Baseten couldn’t complete the connection to the deployment within its timeout or retry limits. Check deployment readiness and capacity.

WebSocket close codes

The Metrics page reports the following WebSocket close codes. See the IANA registry for registered code definitions.

Code	Label	What it means
1000	Normal closure	Either side closed the connection cleanly. This is normal, expected traffic.
1001	Going away	One side is going away, for example a replica restarting or a browser navigating away.
1002	Protocol error	One side sent a frame that violates the WebSocket protocol.
1003	Unsupported data	One side received a frame type it cannot accept, for example binary data on a text-only endpoint.
1005	No status received	The connection closed without a status code. Reserved and not sent on the wire.
1006	Abnormal closure	The connection ended without a recorded close frame. This can happen after a network failure or server termination.
1007	Invalid frame payload data	A message payload was inconsistent with its declared type, for example non-UTF-8 bytes in a text frame.
1008	Policy violation	One side received a message that violates its policy.
1009	Message too big	A message exceeded the receiver’s size limit. Check the client library’s limit and Baseten’s maximum message size.
1010	Mandatory extension missing	The client expected a WebSocket extension that the server did not negotiate.
1011	Internal error	The server side hit an unexpected error that forced the connection to close. Check your model logs.
1012	Service restart	The server is restarting.
1013	Try again later	The server is temporarily overloaded.
1014	Bad gateway	A server acting as a gateway received an invalid response from an upstream server.
1015	TLS handshake failure	Reserved and not sent on the wire.

For managed Truss models and Chains, these cases help diagnose handler problems:

• Truss closes with 1000 when the handler returns normally.
• Truss closes with 1003 when the deployment has no WebSocket handler. Check that you deployed a model with a websocket method or a Chain with a WebSocket entrypoint.
• Truss closes with 1011 when the handler raises an uncaught exception. Check the deployment logs for the error.

Custom servers control their own handler and close behavior. For protocol definitions, see CloseEvent codes on MDN.

The Metrics page shows HTTP 408 and 504 and the WebSocket close codes listed here. It excludes other status codes. The recorded status describes the upstream connection and can differ from the close code your client observes.

End-to-end connection duration

Baseten measures duration from when it starts handling the connection request until the request ends, including time before the WebSocket upgrade completes. The Metrics page reports p50, p90, p95, and p99 after the connection ends.

For speech applications, a connection can contain multiple utterances. Measure first partial transcript, finalization delay, or first audio in your client. Connection duration doesn’t measure these events. See Voice inference performance for timing definitions and concurrency measurements.

Connection input and output size

Cumulative message payload bytes transferred over the connection’s lifetime, reported at p50, p90, p95, and p99:

• Connection input size: Bytes sent by the client to the server.
• Connection output size: Bytes sent by the server to the client.

Connect to a speech model

For real-time transcription with a model from the Model Library, start with the speech quickstart. Each model defines its audio format, message types, and completion signals. Follow the transcription WebSocket reference for the Qwen3-ASR or Whisper protocol.

Send audio and receive transcripts concurrently. After sending the end-of-audio message, keep receiving until the model confirms completion. For Qwen3-ASR, wait for a final transcription with is_end_of_audio_flush: true. A normal WebSocket close alone doesn’t confirm that all audio was transcribed.

Handle model error messages as failures and use receive timeouts to avoid waiting indefinitely. The quickstart client demonstrates completion handling, timeouts, and task cleanup.