add core tests, fix polish empty/timeout fallback

Author: boxerab

src-tauri/src/polish.rs

@@ -111,6 +111,10 @@ pub async fn polish(
         return raw.to_string();
     }
     match polish_inner(cfg, raw, app, ctx).await {
+        // An empty completion (no choices, or blank content) would silently
+        // drop the utterance. Keep the raw transcript instead; don't record
+        // the empty result in the rolling context.
+        Ok(out) if out.is_empty() => raw.to_string(),
         Ok(out) => {
             ctx.append(&out);
             out
@@ -191,7 +195,14 @@ async fn polish_inner(
     };
 
     let url = format!("{}/chat/completions", cfg.base_url.trim_end_matches('/'));
-    let mut req = reqwest::Client::new().post(&url).json(&body);
+    // Bound the round-trip so a hung endpoint can't wedge the collector on a
+    // single segment. On timeout reqwest errors, falling through to the raw
+    // transcript like any other polish failure.
+    let client = reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(30))
+        .build()
+        .context("build polish http client")?;
+    let mut req = client.post(&url).json(&body);
 
     // Bearer auth only if the configured env var is set. Lets local servers
     // (Ollama, llama.cpp) work with no key configured.
@@ -223,6 +234,122 @@ async fn polish_inner(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+    use tokio::net::TcpListener;
+
+    /// Spin up a one-shot HTTP server that reads a full request then replies
+    /// with `response`. Returns a `base_url` pointing at it. Draining the
+    /// request body before replying avoids an RST truncating reqwest's read.
+    async fn stub_openai(response: String) -> String {
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap();
+        tokio::spawn(async move {
+            let Ok((mut stream, _)) = listener.accept().await else {
+                return;
+            };
+            let mut buf = Vec::new();
+            let mut tmp = [0u8; 1024];
+            let header_end = loop {
+                let n = match stream.read(&mut tmp).await {
+                    Ok(0) | Err(_) => return,
+                    Ok(n) => n,
+                };
+                buf.extend_from_slice(&tmp[..n]);
+                if let Some(pos) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
+                    break pos + 4;
+                }
+            };
+            let headers = String::from_utf8_lossy(&buf[..header_end]).to_ascii_lowercase();
+            let content_len = headers
+                .lines()
+                .find_map(|l| l.strip_prefix("content-length:"))
+                .and_then(|v| v.trim().parse::<usize>().ok())
+                .unwrap_or(0);
+            while buf.len() < header_end + content_len {
+                match stream.read(&mut tmp).await {
+                    Ok(0) | Err(_) => break,
+                    Ok(n) => buf.extend_from_slice(&tmp[..n]),
+                }
+            }
+            let _ = stream.write_all(response.as_bytes()).await;
+            let _ = stream.flush().await;
+        });
+        format!("http://{addr}/v1")
+    }
+
+    fn http_response(status: &str, body: &str) -> String {
+        format!(
+            "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
+            body.len()
+        )
+    }
+
+    fn enabled_cfg(base_url: String) -> PolishConfig {
+        PolishConfig {
+            enabled: true,
+            base_url,
+            model: "test-model".into(),
+            api_key_env: String::new(),
+            per_app_tone: true,
+        }
+    }
+
+    #[tokio::test]
+    async fn success_returns_model_output_and_records_context() {
+        let body = r#"{"choices":[{"message":{"content":"  Cleaned up text.  "}}]}"#;
+        let cfg = enabled_cfg(stub_openai(http_response("200 OK", body)).await);
+        let ctx = PolishContext::new();
+        let out = polish(&cfg, "um cleaned up text", None, &ctx).await;
+        assert_eq!(out, "Cleaned up text.");
+        assert_eq!(ctx.snapshot(), "Cleaned up text.");
+        assert!(!ctx.is_disabled());
+    }
+
+    #[tokio::test]
+    async fn http_error_falls_back_to_raw_and_disables_session() {
+        let cfg = enabled_cfg(stub_openai(http_response("401 Unauthorized", "{}")).await);
+        let ctx = PolishContext::new();
+        let out = polish(&cfg, "raw words", None, &ctx).await;
+        assert_eq!(
+            out, "raw words",
+            "an API failure must not drop the transcript"
+        );
+        assert!(
+            ctx.is_disabled(),
+            "one failure disables polish for the session"
+        );
+        assert!(
+            ctx.snapshot().is_empty(),
+            "failed polish must not pollute context"
+        );
+    }
+
+    #[tokio::test]
+    async fn empty_completion_falls_back_to_raw() {
+        // Both an empty choices list and blank content must keep the utterance.
+        for body in [
+            r#"{"choices":[]}"#,
+            r#"{"choices":[{"message":{"content":"   "}}]}"#,
+        ] {
+            let cfg = enabled_cfg(stub_openai(http_response("200 OK", body)).await);
+            let ctx = PolishContext::new();
+            let out = polish(&cfg, "keep this", None, &ctx).await;
+            assert_eq!(out, "keep this", "empty completion must not drop text");
+            assert!(
+                ctx.snapshot().is_empty(),
+                "empty output must not enter context"
+            );
+        }
+    }
+
+    #[tokio::test]
+    async fn malformed_body_falls_back_to_raw() {
+        let cfg = enabled_cfg(stub_openai(http_response("200 OK", "this is not json")).await);
+        let ctx = PolishContext::new();
+        let out = polish(&cfg, "keep me", None, &ctx).await;
+        assert_eq!(out, "keep me");
+        assert!(ctx.is_disabled());
+    }
 
     #[test]
     fn context_starts_empty() {

src-tauri/src/transport.rs

@@ -87,7 +87,7 @@ pub async fn run_session(
                     break;
                 }
                 chunks_sent += 1;
-                if chunks_sent.is_multiple_of(25) {
+                if chunks_sent % 25 == 0 {
                     tracing::debug!(chunks_sent, "transport: streaming");
                 }
             }
@@ -149,3 +149,101 @@ fn uuid_v4() -> String {
         .unwrap_or(0);
     format!("qol-{nanos:x}")
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::config::{Config, InjectMethod, PolishConfig};
+    use futures_util::{SinkExt, StreamExt};
+    use std::time::Duration;
+    use tokio::net::TcpListener;
+    use tokio_tungstenite::accept_async;
+
+    fn test_config(url: String) -> Config {
+        Config {
+            aavaaz_url: url,
+            model: "test-model".into(),
+            language: Some("en".into()),
+            hotkey: "Super+Space".into(),
+            polish: PolishConfig {
+                enabled: false,
+                base_url: String::new(),
+                model: String::new(),
+                api_key_env: String::new(),
+                per_app_tone: false,
+            },
+            hotwords: vec!["alpha".into(), "bravo".into()],
+            inject_method: InjectMethod::Type,
+        }
+    }
+
+    #[tokio::test]
+    async fn handshake_then_streams_only_unique_completed_segments() {
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap();
+
+        // Fake Aavaaz: accept one client, capture its handshake, then feed a
+        // mix of messages the client must filter correctly.
+        let server = tokio::spawn(async move {
+            let (tcp, _) = listener.accept().await.unwrap();
+            let mut ws = accept_async(tcp).await.unwrap();
+
+            let handshake = match ws.next().await.unwrap().unwrap() {
+                Message::Text(t) => serde_json::from_str::<serde_json::Value>(&t).unwrap(),
+                other => panic!("expected handshake text frame, got {other:?}"),
+            };
+
+            let msgs = [
+                // garbage that isn't a valid envelope -> ignored
+                "not even json",
+                // partial (completed=false) -> skipped
+                r#"{"segments":[{"start":"0.0","end":"1.0","text":"partial","completed":false}]}"#,
+                r#"{"segments":[{"start":"1.0","end":"2.0","text":"hello","completed":true}]}"#,
+                // exact duplicate (same start/end) -> deduped
+                r#"{"segments":[{"start":"1.0","end":"2.0","text":"hello","completed":true}]}"#,
+                r#"{"segments":[{"start":"2.0","end":"3.0","text":"world","completed":true}]}"#,
+            ];
+            for m in msgs {
+                ws.send(Message::Text(m.to_string())).await.unwrap();
+            }
+            ws.close(None).await.unwrap();
+            handshake
+        });
+
+        let (_audio_tx, audio_rx) = mpsc::unbounded_channel::<Vec<f32>>();
+        let (seg_tx, mut seg_rx) = mpsc::unbounded_channel();
+        let cfg = test_config(format!("ws://{addr}"));
+
+        // _audio_tx stays alive so the audio branch of the select stays
+        // pending; the server's Close frame is what ends the session.
+        tokio::time::timeout(Duration::from_secs(5), run_session(cfg, audio_rx, seg_tx))
+            .await
+            .expect("run_session timed out")
+            .expect("run_session returned an error");
+
+        let mut got = Vec::new();
+        while let Some(seg) = seg_rx.recv().await {
+            got.push(seg.text);
+        }
+        assert_eq!(got, vec!["hello", "world"]);
+
+        let handshake = server.await.unwrap();
+        assert_eq!(handshake["task"], "transcribe");
+        assert_eq!(handshake["model"], "test-model");
+        assert_eq!(handshake["language"], "en");
+        assert_eq!(handshake["use_vad"], true);
+        assert_eq!(handshake["hotwords"], "alpha,bravo");
+        assert_eq!(handshake["send_last_n_segments"], 1);
+        assert!(handshake["uid"].as_str().unwrap().starts_with("qol-"));
+    }
+
+    #[tokio::test]
+    async fn connect_failure_returns_error() {
+        let (_audio_tx, audio_rx) = mpsc::unbounded_channel::<Vec<f32>>();
+        let (seg_tx, _seg_rx) = mpsc::unbounded_channel();
+        // Nothing listens on port 1; connect is refused.
+        let cfg = test_config("ws://127.0.0.1:1".into());
+        let res = run_session(cfg, audio_rx, seg_tx).await;
+        assert!(res.is_err());
+    }
+}