Add subcommands and args to CLI

Author: modulrmike

Cargo.toml

@@ -9,7 +9,7 @@ roslibrust = { version = "0.16.0", features = ["ros1", "codegen"] }
 [dependencies]
 webrtc = "0.10"
 tokio = { version = "1", features = ["full"] }
-tokio-tungstenite = { version = "0.20", features = ["native-tls"] }
+tokio-tungstenite = { version = "0.20", features = ["rustls-tls-webpki-roots"] }
 tungstenite = "0.20"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
@@ -25,3 +25,9 @@ roslibrust = { version = "0.16.0", features = ["ros1", "rosbridge", "codegen"] }
 async-trait = "0.1.89"
 futures = "0.3.31"
 base64 = "0.22.1"
+clap = { version = "4.5.52", features = ["derive"] }
+dirs = "6.0.0"
+rustls = "*"
+
+[dev-dependencies]
+tempfile = "3.23.0"

src/commands/config.rs

@@ -0,0 +1,90 @@
+use std::path::PathBuf;
+
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Serialize, Deserialize, PartialEq)]
+pub struct AgentConfig {
+    pub robot_id: String,
+    pub signaling_url: String,
+}
+
+pub fn get_default_path() -> Option<PathBuf> {
+    dirs::config_dir().map(|mut path| {
+        path.push("modulr_agent");
+        path.push("config.json");
+        path
+    })
+}
+
+pub fn read_config(override_path: Option<PathBuf>) -> Result<AgentConfig> {
+    let config_path = override_path.or(get_default_path()).ok_or(anyhow::anyhow!(
+        "No configuration file provided and default file cannot be found!"
+    ))?;
+    let config = std::fs::read_to_string(&config_path).map_err(|_| {
+        anyhow::anyhow!(
+            "Unable to read config from file path: {}",
+            &config_path.display()
+        )
+    })?;
+    serde_json::from_str::<AgentConfig>(&config)
+        .map_err(|_| anyhow::anyhow!("Unable to deserialize configuration from file!"))
+}
+
+pub fn write_config(config: &AgentConfig, override_path: Option<PathBuf>) -> Result<()> {
+    let config_path = override_path.or(get_default_path()).ok_or(anyhow::anyhow!(
+        "No configuration file provided and default file path cannot be built!"
+    ))?;
+    std::fs::write(
+        &config_path,
+        serde_json::to_string_pretty(config)
+            .map_err(|_| anyhow::anyhow!("Unable to serialize configuration for writing!"))?,
+    )
+    .map_err(|_| {
+        anyhow::anyhow!(
+            "Unable to write config to file path: {}",
+            &config_path.display()
+        )
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tempfile::NamedTempFile;
+
+    #[test]
+    fn test_path_produced() {
+        let path = get_default_path().unwrap();
+        println!("Default path: {}", path.display());
+        assert!(path.ends_with("modulr_agent/config.json"));
+    }
+
+    #[test]
+    fn test_read_config_roundtrip() {
+        let config = AgentConfig {
+            robot_id: "my_robot_id".to_string(),
+            signaling_url: "my_signaling_url".to_string(),
+        };
+        let serialized = serde_json::to_string_pretty(&config).unwrap();
+        println!("Serialized config: {}", serialized);
+        let temp_file = NamedTempFile::new().unwrap();
+        std::fs::write(temp_file.path(), &serialized).unwrap();
+
+        let deserialized = read_config(Some(temp_file.into_temp_path().to_path_buf())).unwrap();
+
+        assert_eq!(deserialized, config);
+    }
+
+    #[test]
+    fn test_write_config_roundtrip() {
+        let config = AgentConfig {
+            robot_id: "my_robot_id".to_string(),
+            signaling_url: "my_signaling_url".to_string(),
+        };
+        let temp_file = NamedTempFile::new().unwrap();
+        write_config(&config, Some(temp_file.path().to_path_buf())).unwrap();
+        let deserialized = read_config(Some(temp_file.path().to_path_buf())).unwrap();
+        assert_eq!(deserialized, config);
+    }
+}

src/commands/initial_setup.rs

@@ -0,0 +1,26 @@
+use anyhow::Result;
+use clap::Parser;
+use std::path::PathBuf;
+
+use crate::commands::config::{AgentConfig, write_config};
+
+#[derive(Parser, Debug)]
+pub struct InitialSetupArgs {
+    /// Robot ID to use for connection with Modulr services
+    #[arg(short, long)]
+    robot_id: String,
+    /// Signaling URL for establishing WebRTC link
+    #[arg(short, long)]
+    signaling_url: String,
+    /// Override default config path
+    #[arg(short, long, value_name = "FILE")]
+    config_override: Option<PathBuf>,
+}
+
+pub async fn initial_setup(args: InitialSetupArgs) -> Result<()> {
+    let config = AgentConfig {
+        robot_id: args.robot_id,
+        signaling_url: args.signaling_url,
+    };
+    write_config(&config, args.config_override)
+}

src/commands/mod.rs

@@ -0,0 +1,6 @@
+pub mod config;
+mod initial_setup;
+mod start;
+
+pub use initial_setup::*;
+pub use start::*;

src/commands/start.rs

@@ -0,0 +1,143 @@
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use anyhow::Result;
+use bytes::Bytes;
+use clap::Parser;
+use log::{debug, info, warn};
+use tokio::sync::Mutex;
+
+use crate::commands::config::read_config;
+use crate::ros_bridge::Ros1Bridge;
+use crate::ros_bridge::Ros2Bridge;
+use crate::ros_bridge::RosBridge;
+use crate::video_pipeline::VideoPipeline;
+use crate::video_pipeline::VideoPipelineError;
+use crate::webrtc_link::WebRtcLink;
+use crate::webrtc_link::WebRtcLinkError;
+use crate::webrtc_message::WebRtcMessage;
+
+const ROS1: bool = false;
+
+#[derive(Parser, Debug)]
+pub struct StartArgs {
+    /// Override default config path
+    #[arg(short, long, value_name = "FILE")]
+    pub config_override: Option<PathBuf>,
+    /// Allow not checking certificates on WebRTC link
+    #[arg(long, default_value_t = false)]
+    allow_skip_cert_check: bool,
+}
+
+pub async fn start(args: StartArgs) -> Result<()> {
+    let config = read_config(args.config_override)?;
+
+    let robot_id = config.robot_id;
+    let signaling_url = config.signaling_url;
+    let webrtc_link = Arc::new(Mutex::new(WebRtcLink::new(
+        &robot_id,
+        &signaling_url,
+        args.allow_skip_cert_check,
+    )));
+    let pipeline = Arc::new(Mutex::new(VideoPipeline::new()));
+
+    let bridge: Arc<Mutex<dyn RosBridge>> = if ROS1 {
+        Arc::new(Mutex::new(Ros1Bridge::new()))
+    } else {
+        Arc::new(Mutex::new(Ros2Bridge::new()))
+    };
+
+    info!("Creating system components and callbacks");
+
+    // Browser -> WebRTC -> ROS
+    let bridge_clone = Arc::clone(&bridge);
+    webrtc_link
+        .lock()
+        .await
+        .on_webrtc_message(Box::new(move |msg: &WebRtcMessage| {
+            let bridge_clone = Arc::clone(&bridge_clone);
+            let msg_clone = msg.clone();
+            Box::pin(async move {
+                match msg_clone {
+                    WebRtcMessage::MovementCommand(cmd) => {
+                        bridge_clone
+                            .lock()
+                            .await
+                            .post_movement_command(&cmd)
+                            .await
+                            .expect("Failed to post movement command!");
+                    }
+                }
+            })
+        }))
+        .await;
+
+    // TODO: only queue frames in pipeline when WebRTC session is established
+
+    // Pipeline -> WebRTC -> Browser
+    let webrtc_link_clone = Arc::clone(&webrtc_link);
+    pipeline
+        .lock()
+        .await
+        .on_frame_ready(Box::new(move |frame: &Bytes| {
+            let webrtc_link_clone = Arc::clone(&webrtc_link_clone);
+            let frame_clone = Bytes::copy_from_slice(frame);
+            Box::pin(async move {
+                let err = webrtc_link_clone
+                    .lock()
+                    .await
+                    .write_frame(frame_clone)
+                    .await;
+                match err {
+                    Err(WebRtcLinkError::IncorrectStateForOperation) => {
+                        warn!("Still waiting for WebRTC pipeline to connect.")
+                    }
+                    Err(_) => {
+                        panic!("Failed to write frame!")
+                    }
+                    _ => (),
+                }
+            })
+        }))
+        .await;
+
+    // Robot frame -> Pipeline
+    let pipeline_clone = Arc::clone(&pipeline);
+    bridge
+        .lock()
+        .await
+        .on_image_frame_received(Box::new(move |data: &Bytes| {
+            let data_clone = Bytes::copy_from_slice(data);
+            let pipeline_clone = Arc::clone(&pipeline_clone);
+            Box::pin(async move {
+                pipeline_clone
+                    .lock()
+                    .await
+                    .queue_frame(&data_clone)
+                    .await
+                    .expect("Failed to write camera frame to pipeline!");
+            })
+        }))
+        .await;
+
+    info!("Starting all tasks running");
+
+    tokio::spawn(async move {
+        let mut guard = webrtc_link.lock().await;
+        guard.try_connect().await?;
+        guard.try_register().await?;
+        Ok::<(), WebRtcLinkError>(())
+    });
+    let pipeline_clone = Arc::clone(&pipeline);
+    tokio::spawn(async move {
+        pipeline_clone.lock().await.launch().await?;
+        debug!("Finished launching video pipeline");
+        Ok::<(), VideoPipelineError>(())
+    });
+    bridge.lock().await.launch().await?;
+
+    let _ = tokio::signal::ctrl_c().await;
+    info!("Exit requested, cleaning up...");
+    pipeline.lock().await.stop_pipeline().await?;
+    Ok(())
+}