arora_websocket/

bridge.rs

//! [`WsBridge`]: the WebSocket server driven as an Arora
//! [`Bridge`](arora_bridge::Bridge).
//!
//! Each incoming message becomes a [`BridgeCommand`] on
//! [`commands`](Bridge::commands), and the consumer — the Arora runtime —
//! reacts to the ones it cares about (writes apply to the store, reads reply).
//! Runtime state flows back out through [`send_data`](Bridge::send_data) as a
//! `values_changed` push.
//!
//! The server's handler API is kept and *built on top of* the command stream:
//! the handlers registered here simply translate a message into a command. The
//! value vocabulary is `arora_types::Value`, so the translation is structural,
//! not a conversion.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use arora_bridge::{
    Bridge, BridgeCommand, BridgeOp, BridgeResult, CommandStream, DataRequestedStream, DeviceInfo,
    DeviceInfoStream,
};
use arora_types::data::{Key, StateChange};
use arora_types::value::Value;
use async_trait::async_trait;
use futures::channel::{mpsc, oneshot};

use crate::messages::Outgoing;
use crate::server::AroraWSServer;

/// The WebSocket server as an Arora [`Bridge`].
///
/// Note: the server's `validate_paths` (on by default) checks written key
/// paths against its `Registry`, which this bridge does not populate — either
/// mirror the runtime's keys into the registry or disable `validate_paths`
/// when serving purely through the bridge.
pub struct WsBridge {
    server: Arc<AroraWSServer>,
    /// The receiving half of the command stream, handed out once by [`commands`].
    commands: Mutex<Option<mpsc::UnboundedReceiver<BridgeCommand>>>,
}

impl WsBridge {
    /// Wrap a server: register handlers that turn each incoming message into a
    /// [`BridgeCommand`] on the [`commands`](Bridge::commands) stream.
    pub async fn new(server: Arc<AroraWSServer>) -> Self {
        let (cmd_tx, cmd_rx) = mpsc::unbounded::<BridgeCommand>();

        // WriteValues -> Update. The handler is synchronous, so enqueue the
        // command and acknowledge; the runtime applies it on its next step.
        let tx = cmd_tx.clone();
        server
            .set_write_values_handler(move |values: HashMap<String, Value>| {
                let mut change = StateChange::new();
                for (path, value) in values {
                    change.set.insert(Key::from(path), Some(value));
                }
                let (reply_tx, _reply_rx) = oneshot::channel();
                tx.unbounded_send(BridgeCommand::new(BridgeOp::Update(change), reply_tx))
                    .map_err(|_| "bridge command channel closed".to_string())
            })
            .await;

        // ReadValues -> Get, awaiting the runtime's reply.
        let tx = cmd_tx.clone();
        server
            .set_read_values_handler(Arc::new(move |keys: Vec<String>| {
                let tx = tx.clone();
                Box::pin(async move {
                    let store_keys: Vec<Key> = keys.iter().cloned().map(Key::from).collect();
                    let (reply_tx, reply_rx) = oneshot::channel();
                    if tx
                        .unbounded_send(BridgeCommand::new(BridgeOp::Get(store_keys), reply_tx))
                        .is_err()
                    {
                        return HashMap::new();
                    }
                    match reply_rx.await {
                        Ok(Ok(result)) => values_from_get(&keys, result.ret),
                        _ => HashMap::new(),
                    }
                }) as _
            }))
            .await;

        Self {
            server,
            commands: Mutex::new(Some(cmd_rx)),
        }
    }
}

/// Decode a `Get` reply — an `ArrayValue` of `Option`s in request order — into a
/// `path -> value` map.
fn values_from_get(keys: &[String], ret: Value) -> HashMap<String, Value> {
    let mut out = HashMap::new();
    if let Value::ArrayValue(items) = ret {
        for (key, item) in keys.iter().zip(items) {
            if let Value::Option(Some(value)) = item {
                out.insert(key.clone(), *value);
            }
        }
    }
    out
}

#[async_trait]
impl Bridge for WsBridge {
    async fn get_device_info(&self) -> BridgeResult<Option<DeviceInfo>> {
        // A local editor connection has no device-registration concept.
        Ok(None)
    }

    async fn device_info_updated(&self) -> BridgeResult<DeviceInfoStream> {
        Ok(Box::pin(futures::stream::empty()))
    }

    async fn update_device_info(
        &self,
        info: Option<DeviceInfo>,
    ) -> BridgeResult<Option<DeviceInfo>> {
        Ok(info)
    }

    async fn data_requested(&self) -> DataRequestedStream {
        // A connected editor is a data consumer; signal interest once.
        Box::pin(futures::stream::once(async { true }))
    }

    async fn send_data(&self, data: StateChange) -> BridgeResult<()> {
        // Push the changed keys to the connected client(s).
        let mut values = HashMap::new();
        for (key, value) in data.set {
            if let Some(value) = value {
                values.insert(key.path, value);
            }
        }
        if !values.is_empty() {
            self.server.push(Outgoing::ValuesChanged { values });
        }
        Ok(())
    }

    /// The command stream is single-use: the channel's receiving half is
    /// handed out on the first call, and later calls get an empty stream —
    /// with a loud warning, because a runtime silently losing its command
    /// plane is the worst failure mode.
    async fn commands(&self) -> CommandStream {
        match self.commands.lock().unwrap().take() {
            Some(rx) => Box::pin(rx),
            None => {
                log::warn!(
                    "WsBridge::commands() called more than once; the command \
                     stream was already taken, returning an empty stream"
                );
                Box::pin(futures::stream::empty())
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::server::ServerConfig;

    #[tokio::test]
    async fn send_data_pushes_values_changed() {
        let server = Arc::new(AroraWSServer::new(ServerConfig::default()));
        let bridge = WsBridge::new(server.clone()).await;
        let mut rx = server.subscribe();

        bridge
            .send_data(StateChange::set("face/mouth", Value::F64(0.5)))
            .await
            .expect("send_data");

        match rx.recv().await.expect("a push") {
            Outgoing::ValuesChanged { values } => {
                assert_eq!(values.get("face/mouth"), Some(&Value::F64(0.5)));
            }
            other => panic!("expected ValuesChanged, got {other:?}"),
        }
    }
}