WebSocket: Connections and Sessions

A WebSocket is a bidirectional, long-lived connection — yet in Capi it’s modeled as a perfectly ordinary endpoint. The difference is entirely in the associated types: the Output is a live connection, the Requires capability makes the framework write the handshake and the adapter perform the upgrade, and the Decoder validates the 101 and hands the upgraded connection to you. capi_websocket itself is a plain dependency with no gating feature of its own; the websocket feature lives on the adapter that can perform the upgrade, and a client library conventionally gates its WebSocket endpoints behind a websocket feature that pulls the crate in — the client capability table says which adapters can, on which lanes.

The endpoint

#[capi]
#[derive(Debug, Clone, WireModel, ApiEndpoint)]
#[endpoint(path_template = "/v1/realtime")]
pub struct ConnectRealtime {}

impl Endpoint for ConnectRealtime {
    type ApiConfig = GrokConfig;
    type Output = TypedWsConnection<GrokConfig, RealtimeServerMessage, RealtimeClientMessage>;
    type Error = WebSocketDecodeError<RealtimeMessageError>;
    type Decoder = WebSocketDecoder;
    type Requires = WebSocket;                          // ← the whole upgrade declaration
    fn method(&self) -> Method { Method::GET }
    fn url(&self, config: &Self::ApiConfig, _context: &ApiContext) -> impl ToUrl {
        config.base_url().path(self.path_string())
    }
}

impl DecodeWebSocket for ConnectRealtime {
    type In = RealtimeServerMessage;    // what the server sends you
    type Out = RealtimeClientMessage;   // what you send the server
    fn session_spec(config: &Self::ApiConfig) -> Option<WsSessionSpec> {
        Some(text_session_spec::<Self, _>(config.json_lib()))
    }
}

type Requires = WebSocket is the entire upgrade declaration — no directive, no extra method. The WebSocket capability’s prepare_request hook writes Connection: Upgrade, Upgrade: websocket, Sec-WebSocket-Version: 13, and a fresh Sec-WebSocket-Key automatically. The URL is an ordinary https one — the capability, not the scheme, says the exchange is an upgrade. Name the two message types on DecodeWebSocket; session_spec describes the per-message schema for the fixture tooling.

Subprotocols

A service that names a subprotocol gets it from DecodeWebSocket::subprotocols, which returns the entries offered in Sec-WebSocket-Protocol in preference order:

fn subprotocols(_config: &Self::ApiConfig) -> Vec<HeaderValue> {
    vec![HeaderValue::from_static("realtime")]
}

Offer only entries the service accepts: a browser fails the connection unless the server echoes back exactly one of them.

Credentials do not belong here even though the config can reach the auth store. This hook runs while the request is assembled, and the seal resolves the credential later, at the moment of use, renewing it through its seat when it is near expiry — a token read here would be the one held before that resolution. A credential lives in the store, and the endpoint’s auth() says to offer it in the same list:

fn auth(&self) -> Auth {
    Auth::DEFAULT.append_header(SEC_WEBSOCKET_PROTOCOL, "xai-client-secret.")
}

That one is placed by the seal’s resolution, appended after the entries above, while the API’s HTTP endpoints keep drawing the same secret through their own Authorization header — see the auth model.

Using the connection

TypedWsConnection<Config, In, Out> is a typed duplex channel:

let mut conn = api.query(ConnectRealtime::new()).await?;
conn.send(RealtimeClientMessage::ResponseCreate).await?;   // Out: a unit variant
while let Some(msg) = conn.recv().await {                         // Option<Result<In, WsError>>
    match msg? { /* handle a RealtimeServerMessage */ }
}
conn.close(None).await?;

send takes your Out type, recv yields your In type, and each message is encoded / decoded through the config’s format-typed Codec via the per-message EncodeWsMessage / DecodeWsMessage traits. (A binary protocol can bypass the codec entirely.) Ping/Pong frames are handled for you; a Close routes through the close handshake. negotiated_subprotocol() reports the entry the server echoed back, when it echoed one.

The drain gotcha. recv() returns None when the connection is cleanly closed — but drain on the protocol’s terminal message, not on None. On the raw WsConnection a clean close arrives as Some(Ok(WsMessage::Close(..))) before the final None, and an abrupt transport EOF surfaces first as Some(Err(..)). On the typed connection a peer’s Close is fed through your In decoder first, so it reaches you as whatever that decoder makes of a WsMessage::Close — an error, in the usual impl — and then None. Either way, loop until your protocol’s own end-of-session signal, treating None as “the socket is gone”, not as your success condition.

The WsMessage enum carries the five frame kinds — Text, Binary, Ping, Pong, Close(Option<WsCloseFrame>) — for code that works at the raw frame level.

Split mode

To send and receive from different tasks, split the connection:

let split = conn.into_split()?;         // → TypedWsSplit { controller, recv, send }

TypedWsSplit gives you independent send and recv halves plus a WsController whose run() drives the underlying socket. This is the shape you want for a full-duplex session where a reader task and a writer task operate concurrently. It is async only: a blocking build has no tasks to hand the halves to, so there into_split() returns Err(WsSplitError::Unsupported).

Sessions and tunneled requests

A socket that carries more than one conversation at a time needs three things a bare TypedWsConnection does not provide: a way to say what an arriving frame is, a place for the state the server hands out mid-session, and correlation keys for the requests this side sends. The session feature supplies the shapes and asks the client for the meaning — as one pure trait impl, so no #[cfg(feature = "async")] twin ends up in the client.

WsSessionProtocol is the client’s entire contribution: a plain value holding whatever the server has told you, with pure methods.

impl WsSessionProtocol for NotificationProtocol {
    type Config   = MyConfig;
    type Incoming = ServerFrame;          // your existing DecodeWsMessage type
    type Outgoing = ClientFrame;          // your existing EncodeWsMessage type
    type Key      = String;               // whatever your protocol spells ids as

    fn mint_key(&mut self, sequence: u64) -> String { sequence.to_string() }

    fn classify(&self, frame: &ServerFrame) -> Classified<String> {
        match frame {
            ServerFrame::Response(rsp) => Classified::Reply { key: rsp.id.clone(), terminal: true },
            _ => Classified::Notification,
        }
    }

    fn absorb(&mut self, frame: &ServerFrame) -> Absorb<ClientFrame> {
        if let ServerFrame::Ready(ready) = frame { self.wsc = Some(ready.token.clone()); }
        Absorb::Nothing                    // or Absorb::Send(frame) for an answer you owe on sight
    }
}

classify reads the frame’s own fields — there is no table of outstanding requests to consult; the session keeps none and the caller correlates — and answers one of four arms: Reply { key, terminal } (with terminal: false while more frames under the same key are still to come, the ordinary case for a streamed answer), Request (a peer-initiated request the server expects you to answer, read out of the frame yourself), Notification (unsolicited traffic), or Discard (a frame the protocol says to drop, which recv loops past). absorb is where server-given session state lands, in your own struct’s fields, and Absorb::Send(frame) is the answer a protocol owes on sight — a pong for a ping — written before your caller sees the frame that caused it. heartbeat_frame(key) is the optional keepalive the caller paces.

WsSessionMachine holds the ordering logic once, as lane-free code a test can drive with no connection at all; WsSession<Proto> is the pump, written once per lane, each method a short delegation: send(frame), recv() (absorbs, answers, skips discards, and hands back a Received { frame, class }), send_heartbeat(), close(..), and into_parts(), which gives back the connection and the protocol value so what the session absorbed survives a reconnect. The session is an alias, pub type NotificationSession = WsSession<NotificationProtocol>;, never a newtype — a newtype would have to delegate the lane-shaped methods, and the twins would be back. Reach absorbed state through session.protocol().

WsTunnelProtocol is the seam for a socket that carries whole API requests inside its frames. Implement it and the session gains send_request(endpoint): the endpoint runs the framework’s real request pipeline — members, credential freshness, the seal, one-shot signing — the finished request is drained into a PackagedRequest, and your tunneled_frame(request, key) turns it into one of your frames. The route is described once, in the endpoint, and the frame cannot drift from it. Back comes the key the reply will echo, and session.decode_reply::<Endpt>(reply) decodes that reply to the endpoint’s own output with its typed error envelope intact, under a follow-up of the socket’s upgrade context. The body policy is the carrier’s and three-valued — splice a body already in the carrier’s format, escape one that is not (base64url, where the carrier specifies it), refuse what the carrier does not admit, checked where the content type is still in hand. tunnel_overrides() defaults to the posture the carriers share: a socket was authorized when it was opened, so a request captured into one of its frames carries no credential of its own. The packaged request and the packaged response are the capture plane, and they are not WebSocket’s: a queue payload or a batch part packages the same way.

Two things the session deliberately does not do: it schedules no keepalive (the caller sends a heartbeat between receives) and orchestrates no reconnect — each is a layer above the machine. And the connection() escape hatch bypasses absorption: frames read through it never reach absorb, so a resumption token goes missing and nothing says so. Use it for what the session does not cover, and into_parts() when you are done with the session entirely. The test framework records a session as a .req/.rsp/.meta/.ws fixture quad (replay testing).

Native vs. the browser

The handshake happens in different places by platform, and the framework hides the difference:

• Native — the client adapter performs the upgrade and validates Sec-WebSocket-Accept against the key it sent (using the standard WebSocket GUID). The decoder deliberately does not re-check it — validation already happened where the handshake did. capic_reqwest upgrades on its async lane only: a blocking build compiles, and the send fails with WebSocketRequiresAsync, because reqwest’s blocking client has no upgrade API. capic_ureq upgrades over tungstenite and hands back a message-level connection; capic_native_client performs the genuine wire upgrade on both lanes.
• The browser — capic_wasm_fetch derives ws:// or wss:// from the endpoint’s http(s) URL at its own wire boundary, the browser owns the handshake and framing entirely, and the connection reaches your decoder as the capability’s typed extension, never through the transport by hand. A refused upgrade surfaces as a connection error from the browser.

Either way, the connection arrives at your DecodeWebSocket as a typed value, and on native a refused upgrade (a 401 or redirect instead of a 101) still flows through the normal decoder path so you can surface it as an error rather than a hang. A WebSocket endpoint issues no follow-up sends of its own: the capability’s follow_up_channel is unsupported().

Next: gRPC, another capability that reshapes the transport — this time with a protobuf codec rather than a socket upgrade.