feat: new record api

Author: L-jasmine

src/config.rs

@@ -143,10 +143,19 @@ pub struct Config {
 
     pub hello_wav: Option<String>,
 
+    #[serde(default)]
+    pub record: RecordConfig,
+
     #[serde(flatten)]
     pub config: AIConfig,
 }
 
+#[derive(Debug, Default, Clone, serde::Serialize, serde::Deserialize)]
+pub struct RecordConfig {
+    #[serde(default)]
+    pub callback_url: Option<String>,
+}
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 #[serde(untagged)]
 pub enum AIConfig {

src/main.rs

@@ -90,12 +90,19 @@ async fn routes(
     let mut router = Router::new()
         // .route("/", get(handler))
         .route("/ws/{id}", any(services::ws::ws_handler))
-        .nest("/record", services::file::new_file_service("./record"))
+        .route("/v1/chat/{id}", any(services::ws::ws_handler))
+        .route("/v1/record/{id}", any(services::ws_record::ws_handler))
+        .nest("/downloads", services::file::new_file_service("./record"))
         .layer(axum::Extension(Arc::new(services::ws::WsSetting::new(
             hello_wav,
             config.config,
             tool_set,
-        ))));
+        ))))
+        .layer(axum::Extension(Arc::new(
+            services::ws_record::WsRecordSetting {
+                record_callback_url: config.record.callback_url,
+            },
+        )));
 
     if let Some(real_config) = real_config {
         log::info!(

src/services/ws.rs

@@ -892,7 +892,7 @@ pub enum ClientMsg {
 // return: wav data
 async fn process_socket_io(
     rx: &mut WsRx,
-    audio_tx: tokio::sync::mpsc::Sender<ClientMsg>,
+    audio_tx: ClientTx,
     socket: &mut WebSocket,
 ) -> anyhow::Result<()> {
     loop {

src/services/ws_record.rs

@@ -10,54 +10,66 @@ use axum::{
 };
 use bytes::Bytes;
 
-use crate::services::ws::WsSetting;
+pub struct WsRecordSetting {
+    pub record_callback_url: Option<String>,
+}
+
+async fn post_to_callback_url(callback_url: &str, id: &str, file_path: &str) -> anyhow::Result<()> {
+    let client = reqwest::Client::new();
+
+    let resp = client
+        .post(callback_url)
+        .json(&serde_json::json!({
+            "id": id,
+            "download_uri": format!("/downloads/{}", file_path),
+        }))
+        .send()
+        .await?;
+
+    log::info!(
+        "[Record] {} callback to {} success: {}",
+        id,
+        callback_url,
+        resp.status()
+    );
+
+    Ok(())
+}
 
 pub async fn ws_handler(
-    Extension(pool): Extension<Arc<WsSetting>>,
+    Extension(setting): Extension<Arc<WsRecordSetting>>,
     ws: WebSocketUpgrade,
     Path(id): Path<String>,
 ) -> impl IntoResponse {
     let request_id = uuid::Uuid::new_v4().as_u128();
     log::info!("[Record] {id}:{request_id:x} connected.");
 
     ws.on_upgrade(move |socket| async move {
-        let id = id.clone();
-        let pool = pool.clone();
-        if let Err(e) = handle_socket(socket, &id, pool.clone()).await {
-            log::error!("{id}:{request_id:x} error: {e}");
-        };
+        match handle_socket(socket, &id).await {
+            Ok(file_path) => {
+                if let Some(callback_url) = &setting.record_callback_url {
+                    if let Err(e) = post_to_callback_url(callback_url, &id, &file_path).await {
+                        log::error!("[Record] {} callback to {} failed: {}", id, callback_url, e);
+                    }
+                }
+            }
+            Err(e) => {
+                log::error!("{id}:{request_id:x} error: {e}");
+            }
+        }
         log::info!("{id}:{request_id:x} disconnected.");
     })
 }
 
 enum ProcessMessageResult {
     Audio(Bytes),
-    Submit,
-    Text(String),
-    StartRecord,
-    StartChat,
     Close,
     Skip,
 }
 
 fn process_message(msg: Message) -> ProcessMessageResult {
     match msg {
-        Message::Text(t) => {
-            if let Ok(cmd) = serde_json::from_str::<crate::protocol::ClientCommand>(&t) {
-                match cmd {
-                    crate::protocol::ClientCommand::StartRecord => {
-                        ProcessMessageResult::StartRecord
-                    }
-                    crate::protocol::ClientCommand::StartChat => ProcessMessageResult::StartChat,
-                    crate::protocol::ClientCommand::Submit => ProcessMessageResult::Submit,
-                    crate::protocol::ClientCommand::Text { input } => {
-                        ProcessMessageResult::Text(input)
-                    }
-                }
-            } else {
-                ProcessMessageResult::Skip
-            }
-        }
+        Message::Text(_) => ProcessMessageResult::Skip,
         Message::Binary(d) => ProcessMessageResult::Audio(d),
         Message::Close(c) => {
             if let Some(cf) = c {
@@ -76,29 +88,35 @@ fn process_message(msg: Message) -> ProcessMessageResult {
     }
 }
 
-// TODO: implement recording logic
-async fn handle_socket(
-    mut socket: WebSocket,
-    id: &str,
-    pool: Arc<WsSetting>,
-) -> anyhow::Result<()> {
-    std::fs::create_dir_all(format!("./record/{id}"))?;
+async fn handle_socket(mut socket: WebSocket, id: &str) -> anyhow::Result<String> {
+    let now = chrono::Local::now();
+    let date_str = now.format("%Y%m%d_%H%M%S%z").to_string();
+    let file_path = format!("{id}/record_{date_str}.wav");
+    let path = format!("./record/{file_path}");
 
-    while let Some(message) = socket.recv().await {
-        let message = message.map_err(|e| anyhow::anyhow!("recv ws error: {e}"))?;
+    let mut wav_file = crate::util::UnlimitedWavFileWriter::new(
+        &path,
+        crate::util::WavConfig {
+            sample_rate: 16000,
+            channels: 1,
+            bits_per_sample: 16,
+        },
+    )
+    .await?;
 
+    while let Ok(Some(Ok(message))) =
+        tokio::time::timeout(std::time::Duration::from_secs(60), socket.recv()).await
+    {
         match process_message(message) {
-            ProcessMessageResult::Audio(_) => {}
-            ProcessMessageResult::Submit => {}
-            ProcessMessageResult::Text(_) => {}
-            ProcessMessageResult::Skip => {}
-            ProcessMessageResult::StartRecord => {}
-            ProcessMessageResult::StartChat => {}
+            ProcessMessageResult::Audio(chunk) => {
+                wav_file.write_pcm_data(&chunk).await?;
+            }
             ProcessMessageResult::Close => {
                 return Err(anyhow::anyhow!("ws closed"));
             }
+            ProcessMessageResult::Skip => {}
         }
     }
 
-    Ok(())
+    Ok(file_path)
 }

src/util.rs

@@ -34,7 +34,7 @@ pub fn pcm_to_wav(pcm_data: &[u8], config: WavConfig) -> Vec<u8> {
     cursor.write_all(&file_size.to_le_bytes()).unwrap(); // ChunkSize (little-endian)
     cursor.write_all(b"WAVE").unwrap(); // Format
 
-    // fmt 子块
+    // fmt subchunk
     cursor.write_all(b"fmt ").unwrap(); // Subchunk1ID
     cursor.write_all(&16u32.to_le_bytes()).unwrap(); // Subchunk1Size (PCM = 16)
     cursor.write_all(&1u16.to_le_bytes()).unwrap(); // AudioFormat (PCM = 1)
@@ -46,7 +46,7 @@ pub fn pcm_to_wav(pcm_data: &[u8], config: WavConfig) -> Vec<u8> {
         .write_all(&config.bits_per_sample.to_le_bytes())
         .unwrap(); // BitsPerSample
 
-    // data 子块
+    // data subchunk
     cursor.write_all(b"data").unwrap(); // Subchunk2ID
     cursor.write_all(&data_size.to_le_bytes()).unwrap(); // Subchunk2Size
 
@@ -279,3 +279,65 @@ pub fn get_samples_i16(reader: &mut wav_io::reader::Reader) -> Result<Vec<i16>,
     }
     Ok(result)
 }
+
+pub struct UnlimitedWavFileWriter {
+    pub config: WavConfig,
+    pub file: tokio::fs::File,
+}
+
+impl UnlimitedWavFileWriter {
+    pub async fn new(path: &str, config: WavConfig) -> anyhow::Result<Self> {
+        let file = tokio::fs::File::create_new(path).await.map_err(|e| {
+            anyhow::anyhow!(
+                "Failed to create wav file at path {}: {}",
+                path,
+                e.to_string()
+            )
+        })?;
+        Ok(Self { config, file })
+    }
+
+    pub async fn write_wav_header(&mut self) -> anyhow::Result<()> {
+        use tokio::io::AsyncWriteExt;
+
+        let bytes_per_sample = self.config.bits_per_sample / 8;
+        let byte_rate =
+            self.config.sample_rate * self.config.channels as u32 * bytes_per_sample as u32;
+        let block_align = self.config.channels * bytes_per_sample;
+        let data_size = 0xFFFFFFFFu32; // unknown data size
+        let file_size = 0x7FFFFFFFu32;
+
+        self.file.write_all(b"RIFF").await?;
+        self.file.write_all(&file_size.to_le_bytes()).await?; // ChunkSize (little-endian)
+        self.file.write_all(b"WAVE").await?; // Format
+
+        // fmt subchunk
+        self.file.write_all(b"fmt ").await?; // Subchunk1ID
+        self.file.write_all(&16u32.to_le_bytes()).await?; // Subchunk1Size (PCM = 16)
+        self.file.write_all(&1u16.to_le_bytes()).await?; // AudioFormat (PCM = 1)
+        self.file
+            .write_all(&self.config.channels.to_le_bytes())
+            .await?; // NumChannels
+        self.file
+            .write_all(&self.config.sample_rate.to_le_bytes())
+            .await?; // SampleRate
+        self.file.write_all(&byte_rate.to_le_bytes()).await?; // ByteRate
+        self.file.write_all(&block_align.to_le_bytes()).await?; // BlockAlign
+        self.file
+            .write_all(&self.config.bits_per_sample.to_le_bytes())
+            .await?; // BitsPerSample
+
+        // data subchunk
+        self.file.write_all(b"data").await?; // Subchunk2ID
+        self.file.write_all(&data_size.to_le_bytes()).await?; // Subchunk2Size
+
+        Ok(())
+    }
+
+    pub async fn write_pcm_data(&mut self, pcm_data: &[u8]) -> anyhow::Result<()> {
+        use tokio::io::AsyncWriteExt;
+
+        self.file.write_all(pcm_data).await?;
+        Ok(())
+    }
+}