Skip to main content

rlx_moshi/stream/
protocol.rs

1/// Kyutai-style binary WebSocket message types (subset).
2#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3#[repr(u8)]
4pub enum WsMsgType {
5    Handshake = 0,
6    Audio = 1,
7    Text = 2,
8    Control = 3,
9    Error = 5,
10    Ping = 6,
11}
12
13impl WsMsgType {
14    pub fn from_u8(v: u8) -> Option<Self> {
15        match v {
16            0 => Some(Self::Handshake),
17            1 => Some(Self::Audio),
18            2 => Some(Self::Text),
19            3 => Some(Self::Control),
20            5 => Some(Self::Error),
21            6 => Some(Self::Ping),
22            _ => None,
23        }
24    }
25}
26
27pub fn encode_ws_handshake() -> Vec<u8> {
28    let mut msg = vec![WsMsgType::Handshake as u8];
29    msg.extend_from_slice(&0u32.to_le_bytes());
30    msg.extend_from_slice(&1u32.to_le_bytes());
31    msg
32}
33
34pub fn encode_ws_audio(pcm: &[f32]) -> Vec<u8> {
35    let mut msg = Vec::with_capacity(1 + pcm.len() * 4);
36    msg.push(WsMsgType::Audio as u8);
37    for &s in pcm {
38        msg.extend_from_slice(&s.to_le_bytes());
39    }
40    msg
41}
42
43pub fn encode_ws_text(text: &str) -> Vec<u8> {
44    let mut msg = Vec::with_capacity(1 + text.len());
45    msg.push(WsMsgType::Text as u8);
46    msg.extend_from_slice(text.as_bytes());
47    msg
48}
49
50/// Decode client audio message → PCM samples.
51pub fn decode_ws_message(data: &[u8]) -> anyhow::Result<Option<(WsMsgType, Vec<f32>)>> {
52    if data.is_empty() {
53        return Ok(None);
54    }
55    let Some(kind) = WsMsgType::from_u8(data[0]) else {
56        anyhow::bail!("unknown ws msg type {}", data[0]);
57    };
58    match kind {
59        WsMsgType::Audio => {
60            let payload = &data[1..];
61            ensure_audio_payload(payload)?;
62            let mut pcm = Vec::with_capacity(payload.len() / 4);
63            for chunk in payload.chunks_exact(4) {
64                pcm.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
65            }
66            Ok(Some((kind, pcm)))
67        }
68        WsMsgType::Ping => Ok(Some((kind, Vec::new()))),
69        WsMsgType::Control if data.len() >= 2 && data[1] == 0 => Ok(Some((kind, Vec::new()))),
70        _ => Ok(Some((kind, Vec::new()))),
71    }
72}
73
74fn ensure_audio_payload(payload: &[u8]) -> anyhow::Result<()> {
75    if !payload.len().is_multiple_of(4) {
76        anyhow::bail!("audio payload length {} not multiple of 4", payload.len());
77    }
78    Ok(())
79}