[not working] add cancellation via token

Author: cg-jl

engine/Cargo.lock

@@ -127,6 +127,7 @@ tokio = { version = "1", default-features = false, features = [
   "macros",
   "time",
 ] }
+tokio-util = "0.7"
 
 [workspace.package]
 version = "0.203.1"

engine/Cargo.toml

@@ -31,7 +31,7 @@ mod anthropic;
 mod aws;
 mod google;
 mod openai;
-pub(super) mod request;
+pub mod request;
 mod stream_request;
 mod vertex;
 
@@ -436,6 +436,7 @@ mod tests {
             allow_proxy: bool,
             stream: bool,
             expose_secrets: bool,
+            cancellation_token: Option<tokio_util::sync::CancellationToken>,
         ) -> Result<reqwest::RequestBuilder> {
             unimplemented!("Not used in tests")
         }

engine/baml-runtime/src/internal/llm_client/primitive/mod.rs

@@ -146,7 +146,7 @@ mod cancellation_tests {
 
         assert!(result.is_ok(), "Stream should complete (with cancellation)");
         let (stream_result, _) = result.unwrap();
-        
+
         // Should be cancelled
         assert!(stream_result.is_err());
         let error_msg = stream_result.unwrap_err().to_string();
@@ -185,7 +185,8 @@ mod cancellation_tests {
             "##,
         );
 
-        let runtime = BamlRuntime::from_file_content(".", &files, HashMap::new()).unwrap();
+        let runtime =
+            BamlRuntime::from_file_content(".", &files, HashMap::<&str, &str>::new()).unwrap();
         let ctx = runtime.create_ctx_manager(BamlValue::String("test".to_string()), None);
 
         // Create a stream
@@ -364,18 +365,21 @@ mod cancellation_tests {
 
         // Should be cancelled
         assert!(result.is_err());
-        
+
         // Callback should not have been called due to immediate cancellation
         let final_count = *callback_count.lock().unwrap();
-        assert_eq!(final_count, 0, "Callbacks should not be invoked after cancellation");
+        assert_eq!(
+            final_count, 0,
+            "Callbacks should not be invoked after cancellation"
+        );
     }
 
     /// Test that cancellation works with sync runtime
     #[cfg(not(target_arch = "wasm32"))]
     #[test]
     fn test_sync_stream_cancellation() {
         let rt = tokio::runtime::Runtime::new().unwrap();
-        
+
         rt.block_on(async {
             let mut files = HashMap::new();
             files.insert(
@@ -452,37 +456,38 @@ mod unit_tests {
     #[tokio::test]
     async fn test_cancellation_token_basic() {
         let token = CancellationToken::new();
-        
+
         // Initially not cancelled
         assert!(!token.is_cancelled());
-        
+
         // Cancel the token
         token.cancel();
-        
+
         // Now should be cancelled
         assert!(token.is_cancelled());
-        
+
         // cancelled() future should complete immediately
-        let result = tokio::time::timeout(
-            std::time::Duration::from_millis(100),
-            token.cancelled()
-        ).await;
-        
-        assert!(result.is_ok(), "cancelled() future should complete immediately");
+        let result =
+            tokio::time::timeout(std::time::Duration::from_millis(100), token.cancelled()).await;
+
+        assert!(
+            result.is_ok(),
+            "cancelled() future should complete immediately"
+        );
     }
 
     /// Test CancellationToken with tokio::select!
     #[tokio::test]
     async fn test_cancellation_with_select() {
         let token = CancellationToken::new();
         let token_clone = token.clone();
-        
+
         // Spawn a task that cancels after delay
         tokio::spawn(async move {
             tokio::time::sleep(std::time::Duration::from_millis(100)).await;
             token_clone.cancel();
         });
-        
+
         // Use select to race between work and cancellation
         let result = tokio::select! {
             _ = tokio::time::sleep(std::time::Duration::from_secs(1)) => {
@@ -492,7 +497,7 @@ mod unit_tests {
                 "cancelled"
             }
         };
-        
+
         assert_eq!(result, "cancelled");
     }
 
@@ -501,14 +506,14 @@ mod unit_tests {
     async fn test_child_token_cancellation() {
         let parent_token = CancellationToken::new();
         let child_token = parent_token.child_token();
-        
+
         // Initially neither is cancelled
         assert!(!parent_token.is_cancelled());
         assert!(!child_token.is_cancelled());
-        
+
         // Cancel parent
         parent_token.cancel();
-        
+
         // Both should be cancelled
         assert!(parent_token.is_cancelled());
         assert!(child_token.is_cancelled());

engine/baml-runtime/tests/test_cancellation.rs

@@ -44,7 +44,8 @@ mod integration_cancellation_tests {
             "##,
         );
 
-        let runtime = BamlRuntime::from_file_content(".", &files, HashMap::new()).unwrap();
+        let runtime =
+            BamlRuntime::from_file_content(".", &files, HashMap::<&str, &str>::new()).unwrap();
         let ctx = runtime.create_ctx_manager(BamlValue::String("test".to_string()), None);
 
         // Create a stream (simulating TypeScript BamlStream creation)
@@ -93,12 +94,12 @@ mod integration_cancellation_tests {
 
         // Should complete within timeout due to cancellation
         assert!(result.is_ok(), "Stream should complete due to cancellation");
-        
+
         let (stream_result, _) = result.unwrap();
-        
+
         // Should be cancelled (not a successful completion)
         assert!(stream_result.is_err(), "Stream should be cancelled");
-        
+
         // Should be much faster than the 10-second delay endpoint
         assert!(
             elapsed < Duration::from_secs(2),
@@ -143,7 +144,8 @@ mod integration_cancellation_tests {
             "##,
         );
 
-        let runtime = BamlRuntime::from_file_content(".", &files, HashMap::new()).unwrap();
+        let runtime =
+            BamlRuntime::from_file_content(".", &files, HashMap::<&str, &str>::new()).unwrap();
         let ctx = runtime.create_ctx_manager(BamlValue::String("test".to_string()), None);
 
         // Create multiple streams
@@ -167,7 +169,7 @@ mod integration_cancellation_tests {
 
             let token = CancellationToken::new();
             stream.set_cancellation_token(token.clone());
-            
+
             streams.push(stream);
             tokens.push(token);
         }
@@ -190,7 +192,8 @@ mod integration_cancellation_tests {
                         None,
                         HashMap::new(),
                     ),
-                ).await;
+                )
+                .await;
                 (i, result)
             });
             handles.push(handle);
@@ -201,13 +204,21 @@ mod integration_cancellation_tests {
 
         // Check results
         for (handle_result, (stream_index, timeout_result)) in results.into_iter().enumerate() {
-            assert!(handle_result.is_ok(), "Task {} should complete", stream_index);
-            
+            assert!(
+                handle_result.is_ok(),
+                "Task {} should complete",
+                stream_index
+            );
+
             let (stream_result, _) = timeout_result.unwrap().unwrap();
-            
+
             if stream_index == 1 {
                 // Middle stream should be cancelled
-                assert!(stream_result.is_err(), "Stream {} should be cancelled", stream_index);
+                assert!(
+                    stream_result.is_err(),
+                    "Stream {} should be cancelled",
+                    stream_index
+                );
                 let error_msg = stream_result.unwrap_err().to_string();
                 assert!(
                     error_msg.contains("cancelled") || error_msg.contains("canceled"),
@@ -249,7 +260,8 @@ mod integration_cancellation_tests {
             "##,
         );
 
-        let runtime = BamlRuntime::from_file_content(".", &files, HashMap::new()).unwrap();
+        let runtime =
+            BamlRuntime::from_file_content(".", &files, HashMap::<&str, &str>::new()).unwrap();
         let ctx = runtime.create_ctx_manager(BamlValue::String("test".to_string()), None);
 
         let mut stream = runtime
@@ -299,7 +311,7 @@ mod integration_cancellation_tests {
 
         // Should be cancelled
         assert!(result.is_err());
-        
+
         // Events should be minimal due to quick cancellation
         let final_count = *event_count.lock().unwrap();
         assert!(
@@ -356,7 +368,7 @@ mod integration_cancellation_tests {
 
             let token = CancellationToken::new();
             stream.set_cancellation_token(token.clone());
-            
+
             streams.push(stream);
             tokens.push(token);
         }
@@ -368,19 +380,21 @@ mod integration_cancellation_tests {
 
         // Run all streams - they should all be cancelled quickly
         let start_time = std::time::Instant::now();
-        
+
         let mut handles = Vec::new();
         for mut stream in streams {
             let ctx_clone = ctx.clone();
             let handle = tokio::spawn(async move {
-                stream.run(
-                    None::<fn()>,
-                    None::<fn(baml_runtime::FunctionResult)>,
-                    &ctx_clone,
-                    None,
-                    None,
-                    HashMap::new(),
-                ).await
+                stream
+                    .run(
+                        None::<fn()>,
+                        None::<fn(baml_runtime::FunctionResult)>,
+                        &ctx_clone,
+                        None,
+                        None,
+                        HashMap::new(),
+                    )
+                    .await
             });
             handles.push(handle);
         }

engine/baml-runtime/tests/test_integration_cancellation.rs

@@ -46,6 +46,8 @@ serde_json.workspace = true
 serde-wasm-bindgen = "0.6.5"
 time.workspace = true
 tokio = { workspace = true, features = ["sync"] }
+tokio-util = { workspace = true }
+thiserror = { workspace = true }
 url.workspace = true
 uuid = { version = "1.8", features = ["v4", "js"] }
 wasm-bindgen = "=0.2.92"
@@ -55,6 +57,8 @@ web-time.workspace = true
 either = "1.8.1"
 itertools = "0.13.0"
 once_cell.workspace = true
+
+
 [dependencies.web-sys]
 version = "0.3.69"
 features = [