Fix wasm integ tests, fix playground not streaming (#2504)

Author: aaronvg

engine/Cargo.lock

@@ -157,7 +157,7 @@ uuid = { workspace = true, features = ["v4", "serde", "js"] }
 wasm-bindgen = "^0.2.74"
 wasm-bindgen-futures = "0.4"
 web-sys.workspace = true
-wasmtimer = "0.4.1"
+wasmtimer = "0.4.3"
 tokio = { version = "1", features = ["rt", "sync"] }
 
 

engine/baml-runtime/Cargo.toml

@@ -9,7 +9,10 @@ use internal_baml_core::ir::repr::IntermediateRepr;
 use jsonish::BamlValueWithFlags;
 use serde_json::json;
 use stream_cancel::Tripwire;
-use tokio::time::MissedTickBehavior;
+#[cfg(not(target_family = "wasm"))]
+use tokio::time::*;
+#[cfg(target_family = "wasm")]
+use wasmtimer::tokio::*;
 use web_time::Duration;
 
 use super::{call::CtxWithHttpRequestId, OrchestrationScope, OrchestratorNodeIterator};
@@ -51,7 +54,7 @@ where
     G: Fn(),
 {
     let mut results = Vec::new();
-    let mut total_sleep_duration = std::time::Duration::from_secs(0);
+    let mut total_sleep_duration = web_time::Duration::from_secs(0);
 
     // Create a future that either waits for cancellation or never completes
     let cancel_future = match cancel_tripwire {
@@ -104,7 +107,7 @@ where
                         let mut latest_content_for_parse: Option<String> = None;
                         // Track last parsed payload surfaced to downstream listeners so we can dedupe events
                         let mut last_sent_partial_serialized: Option<String> = None;
-                        let mut parse_interval = tokio::time::interval(std::time::Duration::from_millis(20));
+                        let mut parse_interval = interval(web_time::Duration::from_millis(20));
                         // If parsing falls behind, skip missed ticks so we only parse latest.
                         parse_interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
 

engine/baml-runtime/src/internal/llm_client/orchestrator/stream.rs

@@ -63,7 +63,7 @@ dissimilar = "1.0.4"
 expect-test = "1.1.0"
 indoc.workspace = true
 either = "1.8.1"
-wasm-bindgen-test = "0.3.42"
+wasm-bindgen-test = "0.3.43"
 walkdir = "2.5.0"
 wasm-logger = "0.2.0"
 wasmtimer = "0.4.1"

engine/baml-schema-wasm/Cargo.toml

@@ -2,8 +2,7 @@
 // wasm-pack test --node
 // and make sure to set rust-analyzer target in vscode settings to:   "rust-analyzer.cargo.target": "wasm32-unknown-unknown",
 #[cfg(target_arch = "wasm32")]
-#[cfg(test)]
-mod tests {
+pub mod tests {
     use std::collections::HashMap;
 
     use baml_schema_build::runtime_wasm::WasmProject;
@@ -102,7 +101,7 @@ mod tests {
         // Update BAML file
         let updated_content = "// A COMMENT".to_string();
         project.update_file("main.baml", Some(updated_content.clone()));
-        let project_files = project.files();
+        let _project_files = project.files();
         assert!(project
             .files()
             .contains(&"main.bamlBAML_PATH_SPLTTER// A COMMENT".to_string()));

engine/baml-schema-wasm/tests/test_file_manager.rs

@@ -2,74 +2,27 @@
 // RUST_LOG=info wasm-pack test --node --test test_log_collector
 // and make sure to set rust-analyzer target in vscode settings to:   "rust-analyzer.cargo.target": "wasm32-unknown-unknown",
 
+// Install wasm bindgen cli: cargo install wasm-bindgen-cli --version 0.2.101
 // Browser test command is:
 // RUST_BACKTRACE=1 RUST_LOG=info wasm-pack test --chrome --headless --test test_log_collector -- --nocapture
 #[cfg(target_arch = "wasm32")]
 #[cfg(test)]
-mod tests {
+pub(crate) mod tests {
 
-    use std::collections::HashMap;
+    use std::{cell::RefCell, collections::HashMap, rc::Rc};
 
     use baml_runtime::tracingv2::publisher::publisher::flush;
     // pub static GLOBAL_TRACE_STORAGE: Lazy<Mutex<u32>> = Lazy::new(|| Mutex::new(0));
-    use baml_runtime::{tracingv2::storage::storage::BAML_TRACER, InternalRuntimeInterface};
+    use baml_runtime::tracingv2::storage::storage::BAML_TRACER;
     use baml_schema_build::runtime_wasm::WasmProject;
     use serde_wasm_bindgen::to_value;
-    use wasm_bindgen::JsValue;
+    use wasm_bindgen::{closure::Closure, JsCast, JsValue};
     use wasm_bindgen_test::*;
 
     // instantiate logger
 
     wasm_bindgen_test_configure!(run_in_browser);
 
-    fn sample_baml_content() -> String {
-        r##"
-
-
-        class Email {
-            subject string
-            body string
-            from_address string
-        }
-
-        enum OrderStatus {
-            ORDERED
-            SHIPPED
-            DELIVERED
-            CANCELLED
-        }
-
-        class OrderInfo {
-            order_status OrderStatus
-            tracking_number string?
-            estimated_arrival_date string?
-        }
-
-        client<llm> GPT4Turbo {
-            provider baml-openai-chat
-            options {
-                model gpt-4-1106-preview
-                api_key env.OPENAI_API_KEY
-            }
-        }
-
-        function GetOrderInfo(input: string) -> OrderInfo {
-            client GPT4Turbo
-            prompt #"
-            Extract this info from the email in JSON format:
-            Before you output the JSON, please explain your
-            reasoning step-by-step. Here is an example on how to do this:
-            'If we think step by step we can see that ...
-             therefore the output JSON is:
-            {
-              ... the json schema ...
-            }'
-          "#
-        }
-        "##
-        .to_string()
-    }
-
     // TODO: add a flag to run_test that's debug, to attach a collector on every test. The collector can be created inside run_test. Make the collector call track_function(..)
     // Then we can check the baml_tracer events since they'll be kept.
     #[wasm_bindgen_test]
@@ -90,7 +43,6 @@ mod tests {
         }
     }
 
-
             "##;
         let mut files = HashMap::new();
         files.insert("error.baml".to_string(), sample_baml_content.to_string());
@@ -122,14 +74,120 @@ mod tests {
                 env_vars_js.clone().into(),
                 None,
             )
-            .await;
+            .await
+            .expect("run_test should succeed");
         }
 
         let events = BAML_TRACER.lock().unwrap().events();
         log::info!("Events {events:#?}");
         // TODO: this makes the test hang on Node, but not in browser.
-        flush().await;
+        let _ = flush().await;
 
         log::info!("done!!")
     }
+
+    #[wasm_bindgen_test]
+    async fn test_streaming_run_test_with_expr_events() {
+        let openai_key = match std::env::var("OPENAI_API_KEY") {
+            Ok(key) if !key.trim().is_empty() => key,
+            _ => {
+                log::warn!("Skipping streaming test due to missing OPENAI_API_KEY");
+                return;
+            }
+        };
+
+        BAML_TRACER.lock().unwrap().clear();
+
+        let sample_baml_content = r##"
+    function Func(name: string ) -> string {
+            client "openai/gpt-4o"
+            prompt #"
+            Return the name of {{name}}
+            "#
+    }
+
+    test One {
+        functions [Func]
+        args {
+            name "stream"
+        }
+    }
+
+            "##;
+
+        let mut files = HashMap::new();
+        files.insert(
+            "streaming.baml".to_string(),
+            sample_baml_content.to_string(),
+        );
+        let files_js = to_value(&files).unwrap();
+        let project = WasmProject::new("baml_src", files_js)
+            .map_err(JsValue::from)
+            .unwrap();
+
+        let env_vars = HashMap::from([("OPENAI_API_KEY".to_string(), openai_key)]);
+        let env_vars_js = to_value(&env_vars).unwrap();
+        let undefined_js = JsValue::UNDEFINED;
+
+        let mut current_runtime = project.runtime(env_vars_js.clone(), undefined_js).unwrap();
+
+        let diagnostics = project.diagnostics(&current_runtime);
+        assert!(diagnostics.errors().is_empty());
+
+        let functions = current_runtime.list_functions();
+        let partial_results: Rc<RefCell<Vec<JsValue>>> = Rc::new(RefCell::new(Vec::new()));
+        let expr_events: Rc<RefCell<Vec<JsValue>>> = Rc::new(RefCell::new(Vec::new()));
+
+        for f in functions.iter() {
+            let partial_results_cb = {
+                let responses = partial_results.clone();
+                Closure::wrap(Box::new(move |value: JsValue| {
+                    responses.borrow_mut().push(value);
+                }) as Box<dyn FnMut(JsValue)>)
+            };
+
+            let expr_events_cb = {
+                let events = expr_events.clone();
+                Closure::wrap(Box::new(move |value: JsValue| {
+                    events.borrow_mut().push(value);
+                }) as Box<dyn FnMut(JsValue)>)
+            };
+
+            let on_partial_fn: js_sys::Function = partial_results_cb
+                .as_ref()
+                .unchecked_ref::<js_sys::Function>()
+                .clone();
+            let on_expr_fn: js_sys::Function = expr_events_cb
+                .as_ref()
+                .unchecked_ref::<js_sys::Function>()
+                .clone();
+
+            f.run_test_with_expr_events(
+                &mut current_runtime,
+                "One".to_string(),
+                on_partial_fn,
+                js_sys::Function::new_no_args(""),
+                on_expr_fn,
+                env_vars_js.clone().into(),
+                None,
+            )
+            .await
+            .expect("streaming run_test_with_expr_events should succeed");
+
+            drop(partial_results_cb);
+            drop(expr_events_cb);
+        }
+
+        let tracer = BAML_TRACER.lock().unwrap();
+        assert!(tracer.function_call_count() > 0);
+    }
 }
+
+// use wasm_bindgen_test::*;
+// wasm_bindgen_test_configure!(run_in_browser);
+
+// #[wasm_bindgen_test]
+// fn test_hello_world() {
+//     wasm_logger::init(wasm_logger::Config::new(log::Level::Info));
+//     println!("Hello, world!");
+// }

engine/baml-schema-wasm/tests/test_log_collector.rs

@@ -17,7 +17,7 @@ export interface TabularViewConfig {
   showDuration: boolean
 }
 
-export const testPanelViewTypeAtom = atomWithStorage<TestPanelViewType>('testPanelViewType', TestPanelViewType.TABULAR, sessionStore)
+export const testPanelViewTypeAtom = atomWithStorage<TestPanelViewType>('testPanelViewType', TestPanelViewType.CARD_EXPANDED, sessionStore)
 export const tabularViewConfigAtom = atomWithStorage<TabularViewConfig>('tabularViewConfig', {
   showInputs: true,
   showModel: false,