fix(generator): termination for duplex generators (#82)

src/generators/generator.rs

···
 use scru128::Scru128Id;
 use tokio::task::JoinHandle;
 
-use futures::StreamExt;
+use nu_protocol::{ByteStream, ByteStreamType, PipelineData, Signals, Span, Value};
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
 use tokio::io::AsyncReadExt;
-use tokio_stream::wrappers::ReceiverStream;
-
-use nu_protocol::{ByteStream, ByteStreamType, PipelineData, Span, Value};
 
 use crate::nu;
 use crate::nu::ReturnOptions;
···
     };
     let opts: GeneratorScriptOptions = nu_config.deserialize_options().unwrap_or_default();
 
+    // Create and set the interrupt signal on the engine state
+    let interrupt = Arc::new(AtomicBool::new(false));
+    engine.state.set_signals(Signals::new(interrupt.clone()));
+
     let task = GeneratorLoop {
         id: spawn_frame.id,
         context_id: spawn_frame.context_id,
···
     task: &GeneratorLoop,
     rx: tokio::sync::mpsc::Receiver<Frame>,
 ) -> PipelineData {
-    let base_topic = task.topic.clone();
-    let stream = ReceiverStream::new(rx);
-    let stream = stream
-        .filter_map(move |frame: Frame| {
-            let store = store.clone();
-            let topic = format!("{}.send", base_topic);
-            async move {
+    let topic = format!("{}.send", task.topic);
+    let signals = task.engine.state.signals().clone();
+    let mut rx = rx;
+    let iter = std::iter::from_fn(move || loop {
+        if signals.interrupted() {
+            return None;
+        }
+
+        match rx.try_recv() {
+            Ok(frame) => {
                 if frame.topic == topic {
                     if let Some(hash) = frame.hash {
-                        if let Ok(content) = store.cas_read(&hash).await {
-                            return Some(content);
+                        if let Ok(bytes) = store.cas_read_sync(&hash) {
+                            if let Ok(content) = String::from_utf8(bytes) {
+                                return Some(content);
+                            }
                         }
                     }
                 }
-                None
             }
-        })
-        .boxed();
-
-    let handle = tokio::runtime::Handle::current();
-    let mut stream = Some(stream);
-    let iter = std::iter::from_fn(move || {
-        if let Some(ref mut s) = stream {
-            handle.block_on(async { s.next().await })
-        } else {
-            None
+            Err(tokio::sync::mpsc::error::TryRecvError::Empty) => {
+                std::thread::sleep(std::time::Duration::from_millis(10));
+                continue;
+            }
+            Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => {
+                return None;
+            }
         }
     });
 

src/generators/tests.rs

···
 use super::*;
+use std::time::Duration;
 use tempfile::TempDir;
 
 use crate::nu;
···
 }
 
 #[tokio::test]
+async fn test_duplex_terminate_stops() {
+    let (store, engine, ctx) = setup_test_env();
+
+    {
+        let store = store.clone();
+        let engine = engine.clone();
+        tokio::spawn(async move {
+            serve(store, engine).await.unwrap();
+        });
+    }
+
+    let script = r#"{ run: {|| each { |x| $"echo: ($x)" } }, duplex: true }"#;
+    let hash = store.cas_insert(script).await.unwrap();
+
+    store
+        .append(Frame::builder("echo.spawn", ctx.id).hash(hash).build())
+        .unwrap();
+
+    let options = ReadOptions::builder()
+        .context_id(ctx.id)
+        .follow(FollowOption::On)
+        .tail(true)
+        .build();
+    let mut recver = store.read(options).await;
+
+    // expect start
+    let frame = recver.recv().await.unwrap();
+    assert_eq!(frame.topic, "echo.start");
+
+    // terminate while generator waits for input
+    store
+        .append(Frame::builder("echo.terminate", ctx.id).build())
+        .unwrap();
+    let frame = recver.recv().await.unwrap();
+    assert_eq!(frame.topic, "echo.terminate");
+
+    // expect stop frame with reason terminate
+    let frame = recver.recv().await.unwrap();
+    assert_eq!(frame.topic, "echo.stop");
+    assert_eq!(frame.meta.unwrap()["reason"], "terminate");
+
+    store
+        .append(Frame::builder("echo.send", ctx.id).build())
+        .unwrap();
+    let frame = recver.recv().await.unwrap();
+    assert_eq!(frame.topic, "echo.send");
+
+    // ensure no additional frames after stop
+    assert_no_more_frames(&mut recver).await;
+}
+
+#[tokio::test]
 async fn test_spawn_error_eviction() {
     let (store, engine, ctx) = setup_test_env();
 
···
     assert_eq!(stop_frame.meta.as_ref().unwrap()["reason"], "spawn.error");
 
     // Allow ServeLoop to process the spawn.error and evict the generator
-    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+    tokio::time::sleep(Duration::from_millis(50)).await;
 
     let good_script = r#"{ run: {|| "ok" } }"#;
     store