Sophia/crates/sophia-server/src/ws.rs

//! WebSocket bridge from the simulation to a single browser client.
//!
//! Wire protocol:
//! - `text` frames: JSON-serialized [`SimEvent`] (Hello, EngramCreated,
//!   BBoxUpdated). One message per frame.
//! - `binary` frames: position frames, encoded as
//!   `[tag u32 LE = 0x01][t_ms u32 LE][n u32 LE][n × (x f32 LE, y f32 LE, z f32 LE)]`
//!   — 12-byte header followed by the float region. The header is encoded
//!   as three little-endian u32s rather than a tighter (u8, u32, u32) so the
//!   float region starts at a 4-byte-aligned offset, which lets the browser
//!   wrap it as a `Float32Array` view without copying. (`Float32Array`
//!   requires its byte offset to be a multiple of 4 and throws otherwise.)
//!   One frame at ~20 Hz, in `instance_idx` order.
//!
//! See `docs/system-analysis.md` Risk #1 in §13 plan for why we use binary
//! frames here.

use axum::extract::ws::{Message, WebSocket};
use bytes::{BufMut, BytesMut};
use futures_util::{SinkExt, StreamExt};
use sophia_core::{GalaxyId, PositionFrame, SimEvent};
use sophia_sim::SimHandle;
use tokio::sync::broadcast;
use tracing::{debug, warn};

const POS_FRAME_TAG: u32 = 0x01;

pub async fn run_galaxy_socket(socket: WebSocket, sim: SimHandle, galaxy: GalaxyId) {
    let (mut sender, mut receiver) = socket.split();

    let (info, snapshots, synapses, mut bus) = match sim.subscribe(galaxy).await {
        Ok(t) => t,
        Err(e) => {
            warn!("ws subscribe failed: {e}");
            let _ = sender
                .send(Message::Text(
                    serde_json::json!({ "type": "error", "message": e.to_string() }).to_string(),
                ))
                .await;
            return;
        }
    };

    // Hello: tell the client about the galaxy + every existing engram + synapse.
    let hello = SimEvent::Hello { galaxy: info, engrams: snapshots, synapses };
    if let Err(e) = send_text(&mut sender, &hello).await {
        debug!("ws hello send failed: {e}");
        return;
    }

    loop {
        tokio::select! {
            // Inbound: ignore messages for now (Stage 1 has no client→server).
            // Just drain so the socket stays alive and we notice closure.
            msg = receiver.next() => {
                match msg {
                    Some(Ok(Message::Close(_))) | None => break,
                    Some(Err(e)) => { debug!("ws recv err: {e}"); break; }
                    Some(Ok(_)) => {}
                }
            }
            // Outbound: forward sim events to the client.
            ev = bus.recv() => {
                match ev {
                    Ok(SimEvent::PositionFrame(f)) => {
                        if let Err(e) = sender.send(Message::Binary(encode_position_frame(&f))).await {
                            debug!("ws send pos frame failed: {e}");
                            break;
                        }
                    }
                    Ok(other) => {
                        if let Err(e) = send_text(&mut sender, &other).await {
                            debug!("ws send text failed: {e}");
                            break;
                        }
                    }
                    Err(broadcast::error::RecvError::Closed) => break,
                    Err(broadcast::error::RecvError::Lagged(n)) => {
                        warn!("ws lagged by {n} events; client will catch up on next frame");
                    }
                }
            }
        }
    }
}

async fn send_text(
    sender: &mut futures_util::stream::SplitSink<WebSocket, Message>,
    ev: &SimEvent,
) -> anyhow::Result<()> {
    let body = serde_json::to_string(ev)?;
    sender.send(Message::Text(body)).await?;
    Ok(())
}

fn encode_position_frame(f: &PositionFrame) -> Vec<u8> {
    let n = f.positions.len();
    let mut buf = BytesMut::with_capacity(12 + n * 12);
    buf.put_u32_le(POS_FRAME_TAG);
    buf.put_u32_le(f.t_ms);
    buf.put_u32_le(n as u32);
    for [x, y, z] in &f.positions {
        buf.put_f32_le(*x);
        buf.put_f32_le(*y);
        buf.put_f32_le(*z);
    }
    buf.to_vec()
}