mod.rs

pub mod generator;
pub mod runtime_prompt;
use std::{collections::HashMap, path::PathBuf, str::FromStr};

use anyhow::Context;
use baml_runtime::{
    internal::{
        llm_client::{
            orchestrator::{ExecutionScope, OrchestrationScope, OrchestratorNode},
            LLMResponse,
        },
        prompt_renderer::PromptRenderer,
    },
    internal_baml_diagnostics::SerializedSpan,
    BamlRuntime, BamlSrcReader, DiagnosticsError, FunctionResult, IRHelper,
    InternalRuntimeInterface, RenderCurlSettings, RenderedPrompt,
};
use baml_types::{BamlValue, GeneratorOutputType, ResponseCheck};
use futures::{channel::mpsc, StreamExt};
use indexmap::IndexMap;
use internal_baml_codegen::version_check::{check_version, GeneratorType, VersionCheckMode};
use internal_baml_core::{feature_flags::FeatureFlags, ir::repr::Walker};
use internal_llm_client::AllowedRoleMetadata;
use itertools::join;
use js_sys::{Promise, Uint8Array};
use jsonish::ResponseBamlValue;
use serde::{Deserialize, Serialize};
use wasm_bindgen::{prelude::*, JsError, JsValue};
use wasm_bindgen_futures::JsFuture;

use self::runtime_prompt::WasmScope;
use crate::{
    abort_controller::js_abort_signal_to_tripwire, runtime_wasm::runtime_prompt::WasmPrompt,
};

type JsResult<T> = core::result::Result<T, JsError>;

// trait IntoJs<T> {
//     fn into_js(self) -> JsResult<T>;
// }

// impl<T, E: Into<anyhow::Error> + Send> IntoJs<T> for core::result::Result<T, E> {
//     fn into_js(self) -> JsResult<T> {
//         self.map_err(|e| JsError::new(format!("{:#}", anyhow::Error::from(e)).as_str()))
//     }
// }

//Run: wasm-pack test --firefox --headless  --features internal,wasm
// but for browser we likely need to do
//         wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
// Node is run using: wasm-pack test --node --features internal,wasm

#[wasm_bindgen(start)]
pub fn on_wasm_init() {
    // TODO: set LOG_LEVEL to ::Debug if you wish to see logs.
    // this is disabled by default because its slows down release mode builds.
    cfg_if::cfg_if! {
        if #[cfg(debug_assertions)] {
            const LOG_LEVEL: log::Level = log::Level::Info;
        } else {
            const LOG_LEVEL: log::Level = log::Level::Info;
        }
    };
    // I dont think we need this line anymore -- seems to break logging if you add it.
    //wasm_logger::init(wasm_logger::Config::new(LOG_LEVEL));
    match console_log::init_with_level(LOG_LEVEL) {
        Ok(_) => web_sys::console::log_1(
            &format!("Initialized BAML runtime logging as log::{LOG_LEVEL}").into(),
        ),
        Err(e) => web_sys::console::log_1(
            &format!("Failed to initialize BAML runtime logging: {e:?}").into(),
        ),
    }

    console_error_panic_hook::set_once();
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct WasmProject {
    #[wasm_bindgen(readonly)]
    pub root_dir_name: String,
    // This is the version of the file on disk
    files: HashMap<String, String>,
    // This is the version of the file that is currently being edited
    // (unsaved changes)
    unsaved_files: HashMap<String, String>,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Debug)]
pub struct WasmDiagnosticError {
    errors: DiagnosticsError,
    pub all_files: Vec<String>,
}

// use serde::Serialize;

#[wasm_bindgen(getter_with_clone)]
#[derive(Debug)]
pub struct SymbolLocation {
    pub uri: String,
    pub start_line: usize,
    pub start_character: usize,
    pub end_line: usize,
    pub end_character: usize,
}

#[wasm_bindgen]
impl WasmDiagnosticError {
    #[wasm_bindgen]
    pub fn errors(&self) -> Vec<WasmError> {
        self.errors
            .errors()
            .iter()
            .map(|e| {
                let (start, end) = e.span().line_and_column();

                WasmError {
                    file_path: e.span().file.path(),
                    start_ch: e.span().start,
                    end_ch: e.span().end,
                    start_line: start.0,
                    start_column: start.1,
                    end_line: end.0,
                    end_column: end.1,
                    r#type: "error".to_string(),
                    message: e.message().to_string(),
                }
            })
            .chain(self.errors.warnings().iter().map(|e| {
                let (start, end) = e.span().line_and_column();

                WasmError {
                    file_path: e.span().file.path(),
                    start_ch: e.span().start,
                    end_ch: e.span().end,
                    start_line: start.0,
                    start_column: start.1,
                    end_line: end.0,
                    end_column: end.1,
                    r#type: "warning".to_string(),
                    message: e.message().to_string(),
                }
            }))
            .collect()
    }
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Debug)]
pub struct WasmError {
    #[wasm_bindgen(readonly)]
    pub r#type: String,
    #[wasm_bindgen(readonly)]
    pub file_path: String,
    #[wasm_bindgen(readonly)]
    pub start_ch: usize,
    #[wasm_bindgen(readonly)]
    pub end_ch: usize,
    #[wasm_bindgen(readonly)]
    pub start_line: usize,
    #[wasm_bindgen(readonly)]
    pub start_column: usize,
    #[wasm_bindgen(readonly)]
    pub end_line: usize,
    #[wasm_bindgen(readonly)]
    pub end_column: usize,
    #[wasm_bindgen(readonly)]
    pub message: String,
}

#[wasm_bindgen]
impl WasmProject {
    #[wasm_bindgen]
    pub fn new(root_dir_name: &str, files: JsValue) -> Result<WasmProject, JsError> {
        let files: HashMap<String, String> = serde_wasm_bindgen::from_value(files)?;

        Ok(WasmProject {
            root_dir_name: root_dir_name.to_string(),
            files,
            unsaved_files: HashMap::new(),
        })
    }

    #[wasm_bindgen]
    pub fn files(&self) -> Vec<String> {
        let mut saved_files = self.files.clone();
        self.unsaved_files.iter().for_each(|(k, v)| {
            saved_files.insert(k.clone(), v.clone());
        });
        let formatted_files = saved_files
            .iter()
            .map(|(k, v)| format!("{k}BAML_PATH_SPLTTER{v}"))
            .collect::<Vec<String>>();
        formatted_files
    }

    #[wasm_bindgen]
    pub fn update_file(&mut self, name: &str, content: Option<String>) {
        if let Some(content) = content {
            self.files.insert(name.to_string(), content);
        } else {
            self.files.remove(name);
        }
    }

    #[wasm_bindgen]
    pub fn save_file(&mut self, name: &str, content: &str) {
        self.files.insert(name.to_string(), content.to_string());
        self.unsaved_files.remove(name);
    }

    #[wasm_bindgen]
    pub fn set_unsaved_file(&mut self, name: &str, content: Option<String>) {
        if let Some(content) = content {
            self.unsaved_files.insert(name.to_string(), content);
        } else {
            self.unsaved_files.remove(name);
        }
    }

    #[wasm_bindgen]
    pub fn diagnostics(&self, rt: &WasmRuntime) -> WasmDiagnosticError {
        let mut hm = self.files.iter().collect::<HashMap<_, _>>();
        hm.extend(self.unsaved_files.iter());

        WasmDiagnosticError {
            errors: rt.runtime.internal().diagnostics().clone(),
            all_files: hm.keys().map(|s| s.to_string()).collect(),
        }
    }

    #[wasm_bindgen]
    pub fn runtime(
        &self,
        env_vars: JsValue,
        feature_flags: JsValue,
    ) -> Result<WasmRuntime, JsValue> {
        let mut hm = self.files.iter().collect::<HashMap<_, _>>();
        hm.extend(self.unsaved_files.iter());

        let env_vars: HashMap<String, String> =
            serde_wasm_bindgen::from_value(env_vars).map_err(|e| {
                JsValue::from_str(&format!(
                    "Expected env_vars to be HashMap<string, string>. {e}"
                ))
            })?;

        let feature_flags = if feature_flags.is_undefined() || feature_flags.is_null() {
            FeatureFlags::new()
        } else {
            let flags: Vec<String> =
                serde_wasm_bindgen::from_value(feature_flags).map_err(|e| {
                    JsValue::from_str(&format!("Expected feature_flags to be Array<string>. {e}"))
                })?;
            FeatureFlags::from_vec(flags)
                .map_err(|e| JsValue::from_str(&format!("Invalid feature flags: {e:?}")))?
        };

        BamlRuntime::from_file_content(&self.root_dir_name, &hm, env_vars, feature_flags)
            .map(|r| WasmRuntime { runtime: r })
            .map_err(|e| match e.downcast::<DiagnosticsError>() {
                Ok(e) => {
                    let wasm_error = WasmDiagnosticError {
                        errors: e,
                        all_files: hm.keys().map(|s| s.to_string()).collect(),
                    }
                    .into();
                    wasm_error
                }
                Err(e) => {
                    log::debug!("Error: {e:#?}");
                    JsValue::from_str(&e.to_string())
                }
            })
    }

    #[wasm_bindgen]
    pub fn run_generators(
        &self,
        no_version_check: Option<bool>,
    ) -> Result<Vec<generator::WasmGeneratorOutput>, wasm_bindgen::JsError> {
        let fake_map: HashMap<String, String> = HashMap::new();
        let no_version_check = no_version_check.unwrap_or(false);

        let js_value = serde_wasm_bindgen::to_value(&fake_map).unwrap();
        let empty_flags = JsValue::undefined();
        let runtime = self.runtime(js_value, empty_flags);
        log::info!("Files are: {:#?}", self.files);
        let res = match runtime {
            Ok(runtime) => runtime.run_generators(&self.files, no_version_check),
            Err(e) => Err(wasm_bindgen::JsError::new(
                format!("Failed to create runtime: {e:#?}").as_str(),
            )),
        };

        res
    }
}

#[wasm_bindgen(inspectable, getter_with_clone)]
#[derive(Clone)]
pub struct WasmRuntime {
    runtime: BamlRuntime,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Clone, Debug)]
pub struct WasmFunction {
    #[wasm_bindgen(readonly)]
    pub name: String,
    #[wasm_bindgen(readonly)]
    pub span: WasmSpan,
    #[wasm_bindgen(readonly)]
    pub test_cases: Vec<WasmTestCase>,
    #[wasm_bindgen(readonly)]
    pub test_snippet: String,
    #[wasm_bindgen(readonly)]
    pub signature: String,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Clone, Debug)]
pub struct WasmSpan {
    #[wasm_bindgen(readonly)]
    pub file_path: String,
    #[wasm_bindgen(readonly)]
    pub start: usize,
    #[wasm_bindgen(readonly)]
    pub end: usize,
    #[wasm_bindgen(readonly)]
    pub start_line: usize,
    #[wasm_bindgen(readonly)]
    pub end_line: usize,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Clone, Debug)]
pub struct WasmGeneratorConfig {
    #[wasm_bindgen(readonly)]
    pub output_type: String,
    #[wasm_bindgen(readonly)]
    pub version: String,
    #[wasm_bindgen(readonly)]
    pub span: WasmSpan,
}

impl From<&baml_runtime::internal_baml_diagnostics::Span> for WasmSpan {
    fn from(span: &baml_runtime::internal_baml_diagnostics::Span) -> Self {
        let (start, end) = span.line_and_column();
        WasmSpan {
            file_path: span.file.path().to_string(),
            start: span.start,
            end: span.end,
            start_line: start.0,
            end_line: end.0,
        }
    }
}

impl Default for WasmSpan {
    fn default() -> Self {
        WasmSpan {
            file_path: "".to_string(),
            start: 0,
            end: 0,
            start_line: 0,
            end_line: 0,
        }
    }
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Clone, Debug)]
pub struct WasmParentFunction {
    #[wasm_bindgen(readonly)]
    pub start: usize,
    #[wasm_bindgen(readonly)]
    pub end: usize,
    #[wasm_bindgen(readonly)]
    pub name: String,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Clone, Debug)]
pub struct WasmTestCase {
    #[wasm_bindgen(readonly)]
    pub name: String,
    #[wasm_bindgen(readonly)]
    pub inputs: Vec<WasmParam>,
    #[wasm_bindgen(readonly)]
    pub error: Option<String>,
    #[wasm_bindgen(readonly)]
    pub span: WasmSpan,
    #[wasm_bindgen(readonly)]
    pub parent_functions: Vec<WasmParentFunction>,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Clone, Debug)]
pub struct WasmParam {
    #[wasm_bindgen(readonly)]
    pub name: String,
    #[wasm_bindgen(readonly)]
    pub value: Option<String>,
    #[wasm_bindgen(readonly)]
    pub error: Option<String>,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Debug, Clone)]
pub struct WasmFunctionTestPair {
    #[wasm_bindgen(readonly)]
    pub function_name: String,
    #[wasm_bindgen(readonly)]
    pub test_name: String,
}

#[wasm_bindgen]
pub struct WasmFunctionResponse {
    function_response: baml_runtime::FunctionResult,
    func_test_pair: WasmFunctionTestPair,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Debug)]
pub struct WasmTestResponses {
    responses: Vec<WasmTestResponse>,
}

#[wasm_bindgen]
impl WasmTestResponses {
    // #[wasm_bindgen(typescript_type = "WasmTestResponse | null")]
    #[wasm_bindgen]
    pub fn yield_next(&mut self) -> Option<WasmTestResponse> {
        self.responses.pop()
    }
}

#[wasm_bindgen]
#[derive(Debug)]
#[allow(dead_code)]
pub struct WasmTestResponse {
    test_response: anyhow::Result<baml_runtime::TestResponse>,
    span: Option<String>,
    tracing_project_id: Option<String>,
    func_test_pair: WasmFunctionTestPair,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
pub struct WasmParsedTestResponse {
    #[wasm_bindgen(readonly)]
    pub value: String,
    #[wasm_bindgen(readonly)]
    pub check_count: usize,
    #[wasm_bindgen(readonly)]
    /// JSON-string of the explanation, if there were any ParsingErrors
    pub explanation: Option<String>,
}

#[wasm_bindgen]
#[derive(Clone, Debug)]
pub enum TestStatus {
    Passed,
    LLMFailure,
    ParseFailure,
    FinishReasonFailed,
    ConstraintsFailed,
    AssertFailed,
    UnableToRun,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
pub struct WasmLLMResponse {
    scope: OrchestrationScope,
    pub model: String,
    prompt: RenderedPrompt,
    pub content: String,
    pub start_time_unix_ms: u64,
    pub latency_ms: u64,
    pub input_tokens: Option<u64>,
    pub output_tokens: Option<u64>,
    pub total_tokens: Option<u64>,
    pub stop_reason: Option<String>,
}

#[wasm_bindgen(getter_with_clone, inspectable)]
pub struct WasmLLMFailure {
    scope: OrchestrationScope,
    pub model: Option<String>,
    prompt: RenderedPrompt,
    pub start_time_unix_ms: u64,
    pub latency_ms: u64,
    pub message: String,
    pub code: String,
}

#[wasm_bindgen]
impl WasmLLMFailure {
    #[wasm_bindgen]
    pub fn client_name(&self) -> String {
        self.scope.name()
    }
    pub fn prompt(&self) -> WasmPrompt {
        // TODO: This is a hack. We shouldn't hardcode AllowedRoleMetadata::All
        // here, but instead plumb it through the LLMErrors
        (&self.prompt, &self.scope, &AllowedRoleMetadata::All).into()
    }
}

#[wasm_bindgen]
impl WasmLLMResponse {
    #[wasm_bindgen]
    pub fn client_name(&self) -> String {
        self.scope.name()
    }

    pub fn prompt(&self) -> WasmPrompt {
        // TODO: This is a hack. We shouldn't hardcode AllowedRoleMetadata::All
        // here, but instead plumb it through the LLMErrors
        (&self.prompt, &self.scope, &AllowedRoleMetadata::All).into()
    }
}

#[wasm_bindgen]
impl WasmFunctionResponse {
    pub fn parsed_response(&self) -> Option<String> {
        self.function_response
            .result_with_constraints_content()
            .map(|p| serde_json::to_string(&p.serialize_partial()))
            .map_or_else(|_| None, |s| s.ok())
    }

    #[wasm_bindgen]
    pub fn llm_failure(&self) -> Option<WasmLLMFailure> {
        llm_response_to_wasm_error(
            self.function_response.llm_response(),
            self.function_response.scope(),
        )
    }

    #[wasm_bindgen]
    pub fn llm_response(&self) -> Option<WasmLLMResponse> {
        (
            self.function_response.llm_response(),
            self.function_response.scope(),
        )
            .to_wasm()
    }

    #[wasm_bindgen]
    pub fn func_test_pair(&self) -> WasmFunctionTestPair {
        self.func_test_pair.clone()
    }
}

fn serialize_value_counting_checks(value: &ResponseBamlValue) -> (serde_json::Value, usize) {
    let checks = value
        .0
        .meta()
        .1
        .iter()
        .map(|ResponseCheck { name, status, .. }| (name.clone(), status.clone()))
        .collect::<IndexMap<String, String>>();

    let sub_check_count: usize = value.0.iter().map(|node| node.meta().1.len()).sum();
    let json_value: serde_json::Value = serde_json::to_value(value.serialize_final())
        .unwrap_or("Error converting value to JSON".into());

    let check_count = checks.len() + sub_check_count;

    (json_value, check_count)
}

#[wasm_bindgen]
impl WasmTestResponse {
    #[wasm_bindgen]
    pub fn status(&self) -> TestStatus {
        match &self.test_response {
            Ok(t) => match t.status() {
                baml_runtime::TestStatus::Pass => TestStatus::Passed,
                baml_runtime::TestStatus::NeedsHumanEval(_) => TestStatus::ConstraintsFailed,
                baml_runtime::TestStatus::Fail(r) => match r {
                    baml_runtime::TestFailReason::TestUnspecified(_) => TestStatus::UnableToRun,
                    baml_runtime::TestFailReason::TestLLMFailure(_) => TestStatus::LLMFailure,
                    baml_runtime::TestFailReason::TestParseFailure(_) => TestStatus::ParseFailure,
                    baml_runtime::TestFailReason::TestFinishReasonFailed(_) => {
                        TestStatus::FinishReasonFailed
                    }
                    baml_runtime::TestFailReason::TestConstraintsFailure {
                        failed_assert, ..
                    } => {
                        if failed_assert.is_some() {
                            TestStatus::AssertFailed
                        } else {
                            TestStatus::ConstraintsFailed
                        }
                    }
                },
            },
            Err(_) => TestStatus::UnableToRun,
        }
    }

    fn parsed_response_impl(&self) -> anyhow::Result<WasmParsedTestResponse> {
        let maybe_parsed_response = &self
            .test_response
            .as_ref()
            .ok()
            .context("No test response")?
            .function_response
            .parsed()
            .as_ref();
        let parsed_response = match maybe_parsed_response {
            Some(Ok(value)) => Ok(value),
            _ => Err(anyhow::anyhow!("No parsed value")),
        }
        .context("No parsed value")?;
        let (flattened_checks, check_count) = serialize_value_counting_checks(parsed_response);
        Ok(WasmParsedTestResponse {
            value: serde_json::to_string(&flattened_checks)?,
            check_count,
            explanation: {
                let j = parsed_response.explanation_json();
                if j.is_empty() {
                    None
                } else {
                    Some(serde_json::to_string(&j)?)
                }
            },
        })
    }

    #[wasm_bindgen]
    pub fn parsed_response(&self) -> Option<WasmParsedTestResponse> {
        self.parsed_response_impl().ok()
    }

    #[wasm_bindgen]
    pub fn llm_failure(&self) -> Option<WasmLLMFailure> {
        self.test_response.as_ref().ok().and_then(|r| {
            llm_response_to_wasm_error(
                r.function_response.llm_response(),
                r.function_response.scope(),
            )
        })
    }

    #[wasm_bindgen]
    pub fn llm_response(&self) -> Option<WasmLLMResponse> {
        self.test_response.as_ref().ok().and_then(|r| {
            (
                r.function_response.llm_response(),
                r.function_response.scope(),
            )
                .to_wasm()
        })
    }

    #[wasm_bindgen]
    pub fn failure_message(&self) -> Option<String> {
        match self.test_response.as_ref() {
            Ok(r) => match r.status() {
                baml_runtime::TestStatus::Pass => None,
                baml_runtime::TestStatus::Fail(r) => r.render_error(),
                baml_runtime::TestStatus::NeedsHumanEval(checks) => Some(format!(
                    "Checks require human validation: {}",
                    join(checks, ", ")
                )),
            },
            Err(e) => Some(format!("{e:#}")),
        }
    }

    fn _trace_url(&self) -> anyhow::Result<String> {
        let test_response = match self.test_response.as_ref() {
            Ok(t) => t,
            Err(e) => anyhow::bail!("Failed to get test response: {:?}", e),
        };
        let start_time = match test_response.function_response.llm_response() {
            LLMResponse::Success(s) => s.start_time,
            LLMResponse::LLMFailure(f) => f.start_time,
            _ => anyhow::bail!("Test has no start time"),
        };
        let _start_time = time::OffsetDateTime::from_unix_timestamp(
            start_time
                .duration_since(web_time::UNIX_EPOCH)?
                .as_secs()
                .try_into()?,
        )?
        .format(&time::format_description::well_known::Rfc3339)?;

        // TODO: update this.
        // let event_span_id = self
        //     .span
        //     .as_ref()
        //     .ok_or(anyhow::anyhow!("Test has no span ID"))?
        //     .to_string();
        // let subevent_span_id = test_response
        //     .function_call
        //     .as_ref()
        //     .ok_or(anyhow::anyhow!("Function call has no span ID"))?
        //     .to_string();

        // Ok(format!(
        //     "https://app.boundaryml.com/dashboard/projects/{}/drilldown?start_time={start_time}&eid={event_span_id}&s_eid={subevent_span_id}&test=false&onlyRootEvents=true",
        //     self.tracing_project_id.as_ref().ok_or(anyhow::anyhow!("No project ID specified"))?
        // ))
        Ok("https://app.boundaryml.com/dashboard/projects/".to_string())
    }

    #[wasm_bindgen]
    pub fn trace_url(&self) -> Option<String> {
        self._trace_url().ok()
    }

    #[wasm_bindgen]
    pub fn func_test_pair(&self) -> WasmFunctionTestPair {
        self.func_test_pair.clone()
    }
}

fn llm_response_to_wasm_error(
    r: &baml_runtime::internal::llm_client::LLMResponse,
    scope: &OrchestrationScope,
) -> Option<WasmLLMFailure> {
    match &r {
        LLMResponse::LLMFailure(f) => Some(WasmLLMFailure {
            scope: scope.clone(),
            model: f.model.clone(),
            prompt: f.prompt.clone(),
            start_time_unix_ms: f
                .start_time
                .duration_since(web_time::UNIX_EPOCH)
                .unwrap_or(web_time::Duration::ZERO)
                .as_millis() as u64,
            latency_ms: f.latency.as_millis() as u64,
            message: f.message.clone(),
            code: f.code.to_string(),
        }),
        _ => None,
    }
}

trait ToWasm {
    type Output;
    fn to_wasm(&self) -> Self::Output;
}

impl ToWasm
    for (
        &baml_runtime::internal::llm_client::LLMResponse,
        &OrchestrationScope,
    )
{
    type Output = Option<WasmLLMResponse>;

    fn to_wasm(&self) -> Self::Output {
        match &self.0 {
            baml_runtime::internal::llm_client::LLMResponse::Success(s) => Some(WasmLLMResponse {
                scope: self.1.clone(),
                model: s.model.clone(),
                prompt: s.prompt.clone(),
                content: s.content.clone(),
                start_time_unix_ms: s
                    .start_time
                    .duration_since(web_time::UNIX_EPOCH)
                    .unwrap_or(web_time::Duration::ZERO)
                    .as_millis() as u64,
                latency_ms: s.latency.as_millis() as u64,
                input_tokens: s.metadata.prompt_tokens,
                output_tokens: s.metadata.output_tokens,
                total_tokens: s.metadata.total_tokens,
                stop_reason: s.metadata.finish_reason.clone(),
            }),
            _ => None,
        }
    }
}

trait WithRenderError {
    fn render_error(&self) -> Option<String>;
}

impl WithRenderError for baml_runtime::TestFailReason<'_> {
    fn render_error(&self) -> Option<String> {
        match &self {
            baml_runtime::TestFailReason::TestUnspecified(e) => Some(format!("{e:#}")),
            baml_runtime::TestFailReason::TestLLMFailure(f) => f.render_error(),
            baml_runtime::TestFailReason::TestParseFailure(e)
            | baml_runtime::TestFailReason::TestFinishReasonFailed(e) => {
                match e.downcast_ref::<baml_runtime::errors::ExposedError>() {
                    Some(exposed_error) => match exposed_error {
                        baml_runtime::errors::ExposedError::ValidationError { message, .. } => {
                            Some(message.clone())
                        }
                        baml_runtime::errors::ExposedError::FinishReasonError {
                            message, ..
                        } => Some(message.clone()),
                        baml_runtime::errors::ExposedError::ClientHttpError { message, .. } => {
                            Some(message.clone())
                        }
                        baml_runtime::errors::ExposedError::AbortError => {
                            Some("AbortError".to_string())
                        }
                    },
                    None => Some(format!("{e:#}")),
                }
            }
            baml_runtime::TestFailReason::TestConstraintsFailure {
                checks,
                failed_assert,
            } => {
                let checks_msg = if !checks.is_empty() {
                    let check_msgs = checks.iter().map(|(name, pass)| {
                        format!("{name}: {}", if *pass { "Passed" } else { "Failed" })
                    });
                    format!("Check results:\n{}", join(check_msgs, "\n"))
                } else {
                    String::new()
                };
                let assert_msg = failed_assert
                    .as_ref()
                    .map_or("".to_string(), |name| format!("\nFailed assert: {name}"));
                Some(format!("{checks_msg}{assert_msg}"))
            }
        }
    }
}

impl WithRenderError for baml_runtime::internal::llm_client::LLMResponse {
    fn render_error(&self) -> Option<String> {
        match self {
            baml_runtime::internal::llm_client::LLMResponse::Success(_) => None,
            baml_runtime::internal::llm_client::LLMResponse::LLMFailure(f) => {
                format!("{} {}", f.message, f.code).into()
            }
            baml_runtime::internal::llm_client::LLMResponse::UserFailure(e) => {
                format!("user error: {e}").into()
            }
            baml_runtime::internal::llm_client::LLMResponse::InternalFailure(e) => {
                e.to_string().into()
            }
            baml_runtime::internal::llm_client::LLMResponse::Cancelled(msg) => {
                format!("cancelled: {msg}").into()
            }
        }
    }
}

// Rust-only methods
impl WasmRuntime {
    pub fn run_generators(
        &self,
        input_files: &HashMap<String, String>,
        no_version_check: bool,
    ) -> Result<Vec<generator::WasmGeneratorOutput>, wasm_bindgen::JsError> {
        Ok(self
            .runtime
            // convert the input_files into HashMap(PathBuf, string)
            .run_codegen(
                &input_files
                    .iter()
                    .map(|(k, v)| (PathBuf::from(k), v.clone()))
                    .collect(),
                no_version_check,
            )
            .map_err(|e| JsError::new(format!("{e:#}").as_str()))?
            .into_iter()
            .map(|g| g.into())
            .collect())
    }
}

#[wasm_bindgen]
impl WasmRuntime {
    #[wasm_bindgen]
    pub fn check_if_in_prompt(&self, cursor_idx: usize) -> bool {
        self.runtime.internal().ir().walk_functions().any(|f| {
            f.elem().configs().expect("configs").iter().any(|config| {
                let span = &config.prompt_span;
                cursor_idx >= span.start && cursor_idx <= span.end
            })
        })
    }

    #[wasm_bindgen]
    pub fn list_functions(&self) -> Vec<WasmFunction> {
        let ctx = &self
            .runtime
            .create_ctx_manager(BamlValue::String("wasm".to_string()), None);
        let ctx = ctx.create_ctx_with_default();
        let ctx = ctx.eval_ctx(false);

        self.runtime
            .internal()
            .ir()
            .walk_functions()
            .chain(
                self.runtime
                    .internal()
                    .ir()
                    .expr_fns_as_functions()
                    .iter()
                    .map(|f| Walker {
                        ir: self.runtime.internal().ir(),
                        item: f,
                    }),
            )
            .map(|f| {
                let snippet = format!(
                    r#"test TestName {{
  functions [{name}]
  args {{
{args}
  }}
}}
"#,
                    name = f.name(),
                    args = {
                        // Convert baml_runtime::TypeIR inputs to baml_types::TypeIR and use our improved dummy generator
                        let params = f
                            .inputs()
                            .iter()
                            .map(|(k, runtime_type)| (k.clone(), runtime_type.clone()))
                            .collect::<indexmap::IndexMap<String, _>>();

                        // Use the IR's get_dummy_args method
                        self.runtime
                            .internal()
                            .ir()
                            .get_dummy_args(2, true, &params)
                    }
                );

                let wasm_span = match f.span() {
                    Some(span) => span.into(),
                    None => WasmSpan::default(),
                };

                WasmFunction {
                    name: f.name().to_string(),
                    span: wasm_span,
                    signature: {
                        let inputs = {
                            let params = f
                                .inputs()
                                .iter()
                                .map(|(k, runtime_type)| (k.clone(), runtime_type.clone()))
                                .collect::<indexmap::IndexMap<String, _>>();

                            self.runtime
                                .internal()
                                .ir()
                                .get_dummy_args(2, false, &params)
                                .split('\n')
                                .map(|line| line.trim().to_string())
                                .collect::<Vec<_>>()
                                .join(", ")
                        };

                        format!("({}) -> {}", inputs, f.output())
                    },
                    test_snippet: snippet,
                    test_cases: f
                        .walk_tests()
                        .map(|tc| {
                            let params = match tc.test_case_params(&ctx) {
                                Ok(params) => Ok(params
                                    .iter()
                                    .map(|(k, v)| {
                                        let as_str = match v {
                                            Ok(v) => match serde_json::to_string(v) {
                                                Ok(s) => Ok(s),
                                                Err(e) => Err(e.to_string()),
                                            },
                                            Err(e) => Err(e.to_string()),
                                        };

                                        let (value, error) = match as_str {
                                            Ok(s) => (Some(s), None),
                                            Err(e) => (None, Some(e)),
                                        };

                                        WasmParam {
                                            name: k.to_string(),
                                            value,
                                            error,
                                        }
                                    })
                                    .collect()),
                                Err(e) => Err(e.to_string()),
                            };

                            let (mut params, error) = match params {
                                Ok(p) => (p, None),
                                Err(e) => (Vec::new(), Some(e)),
                            };

                            // Any missing params should be set to an error
                            f.inputs().iter().for_each(|(param_name, t)| {
                                if !params.iter().any(|p| p.name == *param_name) && !t.is_optional()
                                {
                                    params.insert(
                                        0,
                                        WasmParam {
                                            name: param_name.to_string(),
                                            value: None,
                                            error: Some("Missing parameter".to_string()),
                                        },
                                    );
                                }
                            });

                            let wasm_span = match tc.span() {
                                Some(span) => span.into(),
                                None => WasmSpan::default(),
                            };

                            WasmTestCase {
                                name: tc.test_case().name.clone(),
                                inputs: params,
                                error,
                                span: wasm_span,
                                parent_functions: tc
                                    .test_case()
                                    .functions
                                    .iter()
                                    .map(|f| {
                                        let (start, end) = f
                                            .attributes
                                            .span
                                            .as_ref()
                                            .map_or((0, 0), |f| (f.start, f.end));
                                        WasmParentFunction {
                                            start,
                                            end,
                                            name: f.elem.name().to_string(),
                                        }
                                    })
                                    .collect(),
                            }
                        })
                        .collect(),
                }
            })
            .collect()
    }

    #[wasm_bindgen]
    pub fn list_generators(&self) -> Vec<WasmGeneratorConfig> {
        self.runtime
            .codegen_generators()
            .map(|generator| WasmGeneratorConfig {
                output_type: generator.output_type.clone().to_string(),
                version: generator.version.clone(),
                span: WasmSpan {
                    file_path: generator.span.file.path().to_string(),
                    start: generator.span.start,
                    end: generator.span.end,
                    start_line: generator.span.line_and_column().0 .0,
                    end_line: generator.span.line_and_column().1 .0,
                },
            })
            .collect()
    }

    #[wasm_bindgen]
    pub fn check_version(
        generator_version: &str,
        current_version: &str,
        generator_type: &str,
        version_check_mode: &str,
        generator_language: &str,
        is_diagnostic: bool,
    ) -> Option<String> {
        // Convert string parameters to enums
        let generator_type = match generator_type {
            "VSCodeCLI" => GeneratorType::VSCodeCLI,
            "VSCode" => GeneratorType::VSCode,
            "CLI" => GeneratorType::CLI,
            other => return Some(format!("Invalid generator type: {other:?}")),
        };

        let version_check_mode = match version_check_mode {
            "Strict" => VersionCheckMode::Strict,
            "None" => VersionCheckMode::None,
            other => return Some(format!("Invalid version check mode: {other:?}")),
        };

        let Ok(generator_language) = GeneratorOutputType::from_str(generator_language) else {
            return Some(format!(
                "Invalid generator language: {generator_language:?}"
            ));
        };

        check_version(
            generator_version,
            current_version,
            generator_type,
            version_check_mode,
            generator_language,
            is_diagnostic,
        )
        .map(|error| error.msg())
    }

    #[wasm_bindgen]
    pub fn required_env_vars(&self) -> Vec<String> {
        self.runtime
            .internal()
            .ir()
            .required_env_vars()
            .into_iter()
            .map(|s| s.to_string())
            .collect()
    }

    #[wasm_bindgen]
    pub fn search_for_symbol(&self, symbol: &str) -> Option<SymbolLocation> {
        let runtime = self.runtime.internal().ir();

        if let Ok(walker) = runtime.find_enum(symbol) {
            let elem = walker.span().unwrap();

            let ((s_line, s_character), (e_line, e_character)) = elem.line_and_column();
            return Some(SymbolLocation {
                uri: elem.file.path().to_string(), // Use the variable here
                start_line: s_line,
                start_character: s_character,
                end_line: e_line,
                end_character: e_character,
            });
        }
        if let Ok(walker) = runtime.find_class(symbol) {
            let elem = walker.span().unwrap();

            let _uri_str = elem.file.path().to_string(); // Store the String in a variable
            let ((s_line, s_character), (e_line, e_character)) = elem.line_and_column();
            return Some(SymbolLocation {
                uri: elem.file.path().to_string(), // Use the variable here
                start_line: s_line,
                start_character: s_character,
                end_line: e_line,
                end_character: e_character,
            });
        }
        if let Ok(walker) = runtime.find_type_alias(symbol) {
            let elem = walker.span().unwrap();

            let _uri_str = elem.file.path().to_string(); // Store the String in a variable
            let ((s_line, s_character), (e_line, e_character)) = elem.line_and_column();
            return Some(SymbolLocation {
                uri: elem.file.path().to_string(), // Use the variable here
                start_line: s_line,
                start_character: s_character,
                end_line: e_line,
                end_character: e_character,
            });
        }

        if let Ok(walker) = runtime.find_function(symbol) {
            let elem = walker.span().unwrap();

            let _uri_str = elem.file.path().to_string(); // Store the String in a variable
            let ((s_line, s_character), (e_line, e_character)) = elem.line_and_column();
            return Some(SymbolLocation {
                uri: elem.file.path().to_string(), // Use the variable here
                start_line: s_line,
                start_character: s_character,
                end_line: e_line,
                end_character: e_character,
            });
        }

        if let Ok(walker) = runtime.find_client(symbol) {
            let elem = walker.span().unwrap();

            let _uri_str = elem.file.path().to_string(); // Store the String in a variable
            let ((s_line, s_character), (e_line, e_character)) = elem.line_and_column();

            return Some(SymbolLocation {
                uri: elem.file.path().to_string(), // Use the variable here
                start_line: s_line,
                start_character: s_character,
                end_line: e_line,
                end_character: e_character,
            });
        }

        if let Ok(walker) = runtime.find_retry_policy(symbol) {
            let elem = walker.span().unwrap();

            let _uri_str = elem.file.path().to_string(); // Store the String in a variable
            let ((s_line, s_character), (e_line, e_character)) = elem.line_and_column();
            return Some(SymbolLocation {
                uri: elem.file.path().to_string(), // Use the variable here
                start_line: s_line,
                start_character: s_character,
                end_line: e_line,
                end_character: e_character,
            });
        }

        if let Ok(walker) = runtime.find_template_string(symbol) {
            let elem = walker.span().unwrap();
            let _uri_str = elem.file.path().to_string(); // Store the String in a variable
            let ((s_line, s_character), (e_line, e_character)) = elem.line_and_column();
            return Some(SymbolLocation {
                uri: elem.file.path().to_string(), // Use the variable here
                start_line: s_line,
                start_character: s_character,
                end_line: e_line,
                end_character: e_character,
            });
        }

        None
    }

    #[wasm_bindgen]
    pub fn is_valid_class(&self, symbol: &str) -> bool {
        self.runtime.internal().ir().find_class(symbol).is_ok()
    }

    #[wasm_bindgen]
    pub fn is_valid_enum(&self, symbol: &str) -> bool {
        self.runtime.internal().ir().find_enum(symbol).is_ok()
    }

    #[wasm_bindgen]
    pub fn is_valid_type_alias(&self, symbol: &str) -> bool {
        self.runtime.internal().ir().find_type_alias(symbol).is_ok()
    }

    #[wasm_bindgen]
    pub fn is_valid_function(&self, symbol: &str) -> bool {
        self.runtime.internal().ir().find_function(symbol).is_ok()
    }

    #[wasm_bindgen]
    pub fn search_for_class_locations(&self, symbol: &str) -> Vec<SymbolLocation> {
        self.runtime
            .internal()
            .ir()
            .find_class_locations(symbol)
            .into_iter()
            .map(|span| {
                let ((start_line, start_character), (end_line, end_character)) =
                    span.line_and_column();
                SymbolLocation {
                    uri: span.file.path().to_string(),
                    start_line,
                    start_character,
                    end_line,
                    end_character,
                }
            })
            .collect()
    }

    #[wasm_bindgen]
    pub fn search_for_enum_locations(&self, symbol: &str) -> Vec<SymbolLocation> {
        self.runtime
            .internal()
            .ir()
            .find_enum_locations(symbol)
            .into_iter()
            .map(|span| {
                let ((start_line, start_character), (end_line, end_character)) =
                    span.line_and_column();
                SymbolLocation {
                    uri: span.file.path().to_string(),
                    start_line,
                    start_character,
                    end_line,
                    end_character,
                }
            })
            .collect()
    }

    #[wasm_bindgen]
    pub fn search_for_type_alias_locations(&self, symbol: &str) -> Vec<SymbolLocation> {
        self.runtime
            .internal()
            .ir()
            .find_type_alias_locations(symbol)
            .into_iter()
            .map(|span| {
                let ((start_line, start_character), (end_line, end_character)) =
                    span.line_and_column();
                SymbolLocation {
                    uri: span.file.path().to_string(),
                    start_line,
                    start_character,
                    end_line,
                    end_character,
                }
            })
            .collect()
    }

    #[wasm_bindgen]
    pub fn get_function_at_position(
        &self,
        file_name: &str,
        selected_func: &str,
        cursor_idx: usize,
    ) -> Option<WasmFunction> {
        let functions = self.list_functions();

        for function in functions.clone() {
            let span = function.span.clone(); // Clone the span

            if span.file_path.as_str().ends_with(file_name)
                && ((span.start + 1)..=(span.end + 1)).contains(&cursor_idx)
            {
                return Some(function);
            }
        }

        let testcases = self.list_testcases();

        for tc in testcases {
            let span = tc.span;
            if span.file_path.as_str().ends_with(file_name)
                && ((span.start + 1)..=(span.end + 1)).contains(&cursor_idx)
            {
                if let Some(_parent_function) =
                    tc.parent_functions.iter().find(|f| f.name == selected_func)
                {
                    return functions.into_iter().find(|f| f.name == selected_func);
                } else if let Some(first_function) = tc.parent_functions.first() {
                    return functions
                        .into_iter()
                        .find(|f| f.name == first_function.name);
                }
            }
        }

        let testcases = self.list_testcases();

        for tc in testcases {
            let span = tc.span;
            if span.file_path.as_str().ends_with(file_name)
                && ((span.start + 1)..=(span.end + 1)).contains(&cursor_idx)
            {
                if let Some(_parent_function) =
                    tc.parent_functions.iter().find(|f| f.name == selected_func)
                {
                    return functions.into_iter().find(|f| f.name == selected_func);
                } else if let Some(first_function) = tc.parent_functions.first() {
                    return functions
                        .into_iter()
                        .find(|f| f.name == first_function.name);
                }
            }
        }

        None
    }

    #[wasm_bindgen]
    pub fn list_testcases(&self) -> Vec<WasmTestCase> {
        let ctx = self
            .runtime
            .create_ctx_manager(BamlValue::String("wasm".to_string()), None);

        let ctx = ctx.create_ctx_with_default();
        let ctx = ctx.eval_ctx(true);

        self.runtime
            .internal()
            .ir()
            .walk_function_test_pairs()
            .map(|tc| {
                let params = match tc.test_case_params(&ctx) {
                    Ok(params) => Ok(params
                        .iter()
                        .map(|(k, v)| {
                            let as_str = match v {
                                Ok(v) => match serde_json::to_string(v) {
                                    Ok(s) => Ok(s),
                                    Err(e) => Err(e.to_string()),
                                },
                                Err(e) => Err(e.to_string()),
                            };

                            let (value, error) = match as_str {
                                Ok(s) => (Some(s), None),
                                Err(e) => (None, Some(e)),
                            };

                            WasmParam {
                                name: k.to_string(),
                                value,
                                error,
                            }
                        })
                        .collect()),
                    Err(e) => Err(e.to_string()),
                };

                let (mut params, error) = match params {
                    Ok(p) => (p, None),
                    Err(e) => (Vec::new(), Some(e)),
                };
                // Any missing params should be set to an error
                // Any missing params should be set to an error
                tc.function().inputs().iter().for_each(|func_params| {
                    let (param_name, t) = func_params;
                    if !params.iter().any(|p| p.name == *param_name) && !t.is_optional() {
                        params.push(WasmParam {
                            name: param_name.to_string(),
                            value: None,
                            error: Some("Missing parameter".to_string()),
                        });
                    }
                });
                let wasm_span = match tc.span() {
                    Some(span) => span.into(),
                    None => WasmSpan::default(),
                };

                WasmTestCase {
                    name: tc.test_case().name.clone(),
                    inputs: params,
                    error,
                    span: wasm_span,
                    parent_functions: tc
                        .test_case()
                        .functions
                        .iter()
                        .map(|f| {
                            let (start, end) = f
                                .attributes
                                .span
                                .as_ref()
                                .map_or((0, 0), |f| (f.start, f.end));
                            WasmParentFunction {
                                start,
                                end,
                                name: f.elem.name().to_string(),
                            }
                        })
                        .collect(),
                }
            })
            .collect()
    }

    #[wasm_bindgen]
    pub fn get_testcase_from_position(
        &self,
        parent_function: WasmFunction,
        cursor_idx: usize,
    ) -> Option<WasmTestCase> {
        let testcases = parent_function.test_cases;
        for testcase in testcases {
            let span = testcase.clone().span;

            if span.file_path.as_str() == (parent_function.span.file_path)
                && ((span.start + 1)..=(span.end + 1)).contains(&cursor_idx)
            {
                return Some(testcase);
            }
        }
        None
    }

    #[wasm_bindgen]
    pub fn get_function_of_testcase(
        &self,
        file_name: &str,
        cursor_idx: usize,
    ) -> Option<WasmParentFunction> {
        let testcases = self.list_testcases();

        for tc in testcases {
            let span = tc.span;
            if span.file_path.as_str().ends_with(file_name)
                && ((span.start + 1)..=(span.end + 1)).contains(&cursor_idx)
            {
                let first_function = tc
                    .parent_functions
                    .iter()
                    .find(|f| f.start <= cursor_idx && cursor_idx <= f.end)
                    .cloned();

                return first_function;
            }
        }
        None
    }

    #[wasm_bindgen]
    pub async fn run_tests(
        // NOTE: This needs to be `&self` so that the runtime can be read
        // by the UI, e.g to re-enumerate functions. In case you *really* need `&mut` access,
        // consider `RwLock`, ideally only for the data you're going to be mutating and not the
        // entire runtime.
        &self,
        function_test_pairs: js_sys::Array,
        on_partial_response: js_sys::Function,
        get_baml_src_cb: js_sys::Function,
        env: js_sys::Object,
        abort_signal: Option<js_sys::Object>,
    ) -> Result<WasmTestResponses, JsValue> {
        // Convert abort signal to tripwire
        let tripwire = match crate::abort_controller::js_abort_signal_to_tripwire(abort_signal) {
            Ok(tripwire) => tripwire,
            Err(_e) => {
                log::error!("WASM Parallel: Failed to setup abort handler");
                baml_runtime::TripWire::new(None)
            }
        };

        // Create a vector to store all test futures
        let mut test_futures = Vec::new();

        for i in 0..function_test_pairs.length() {
            if let Ok(pair) = js_sys::Reflect::get(&function_test_pairs, &i.into()) {
                if let (Ok(function_name), Ok(test_name)) = (
                    js_sys::Reflect::get(&pair, &JsValue::from_str("functionName")),
                    js_sys::Reflect::get(&pair, &JsValue::from_str("testName")),
                ) {
                    let function_name = function_name.as_string().unwrap_or_default();
                    let test_name = test_name.as_string().unwrap_or_default();

                    let fn_name_copy = function_name.clone();
                    let test_name_copy = test_name.clone();

                    // Create a closure to handle partial responses for this test
                    let on_partial_response_clone = on_partial_response.clone();
                    let cb = Box::new(move |r| {
                        let this = JsValue::NULL;
                        let res = WasmFunctionResponse {
                            function_response: r,
                            func_test_pair: WasmFunctionTestPair {
                                function_name: fn_name_copy.clone(),
                                test_name: test_name_copy.clone(),
                            },
                        }
                        .into();
                        on_partial_response_clone.call1(&this, &res).unwrap();
                    });

                    // Create evaluation context for the test
                    let ctx = self
                        .runtime
                        .create_ctx_manager_for_wasm(js_fn_to_baml_src_reader(
                            get_baml_src_cb.clone(),
                        ));

                    // Reference to the runtime
                    let rt = &self.runtime;
                    let entries = js_sys::Object::entries(&env);
                    let mut env_vars = HashMap::new();
                    for entry in entries.iter() {
                        let arr = entry.dyn_into::<js_sys::Array>().unwrap();
                        let key = arr.get(0).as_string().unwrap();
                        let value = arr.get(1).as_string().unwrap_or_default();
                        env_vars.insert(key, value);
                    }

                    // Clone tripwire for this test
                    let test_tripwire = tripwire.clone();
                    let on_tick = if false { Some(|| {}) } else { None };

                    // Create a future for this test
                    let future = async move {
                        let (test_response, span) = rt
                            .run_test(
                                &function_name,
                                &test_name,
                                &ctx,
                                Some(cb),
                                None,
                                env_vars.clone(),
                                test_tripwire, // Pass tripwire to each test
                                on_tick,
                            )
                            .await;

                        // Return WasmTestResponse for this test
                        WasmTestResponse {
                            test_response,
                            span: Some(span.to_string()),
                            tracing_project_id: rt
                                .tracer_wrapper
                                .get_or_create_tracer(&env_vars)
                                .tracing_project_id(),
                            // tracing_project_id: rt.env_vars().get("BOUNDARY_PROJECT_ID").cloned(),
                            func_test_pair: WasmFunctionTestPair {
                                function_name: function_name.clone(),
                                test_name: test_name.clone(),
                            },
                        }
                    };

                    test_futures.push(future);
                }
            }
        }

        // Run all tests in parallel
        let results = futures::future::join_all(test_futures).await;

        Ok(WasmTestResponses { responses: results })
    }
}

// Define a new struct to store the important information
#[wasm_bindgen(getter_with_clone, inspectable)]
#[derive(Serialize, Deserialize, Debug)]
pub struct SerializableOrchestratorNode {
    pub provider: String,
}

impl From<&OrchestratorNode> for SerializableOrchestratorNode {
    fn from(node: &OrchestratorNode) -> Self {
        SerializableOrchestratorNode {
            provider: node.provider.to_string(),
        }
    }
}

#[cfg(target_arch = "wasm32")]
fn js_fn_to_baml_src_reader(get_baml_src_cb: js_sys::Function) -> BamlSrcReader {
    Some(Box::new(move |path| {
        Box::pin({
            let path = path.to_string();
            let get_baml_src_cb = get_baml_src_cb.clone();
            async move {
                // Windows-specific hotfix: VSCode resolves relative paths relative to workspace root
                // instead of BAML file location. For BAML files directly in baml_src/, prepend "baml_src/".
                // Since WASM can't use cfg!(windows), we detect Windows by checking for backslashes in paths
                // or by checking the user agent, but for simplicity, we'll check if the path contains backslashes.
                let is_windows = web_sys::window()
                    .and_then(|w| w.navigator().user_agent().ok())
                    .map(|ua| ua.contains("Windows"))
                    .unwrap_or(false);

                let adjusted_path =
                    if is_windows && (path.starts_with("../") || path.starts_with("./")) {
                        let result = format!("baml_src/{}", path);
                        web_sys::console::log_1(&wasm_bindgen::JsValue::from_str(&format!(
                            "WASM Windows path fix applied: '{}' → '{}'",
                            path, result
                        )));
                        result
                    } else {
                        web_sys::console::log_1(&wasm_bindgen::JsValue::from_str(&format!(
                            "WASM path unchanged: '{}' (windows={}, relative={})",
                            path,
                            is_windows,
                            path.starts_with("../") || path.starts_with("./")
                        )));
                        path.clone()
                    };

                let null = JsValue::NULL;
                let Ok(read) = get_baml_src_cb.call1(&null, &JsValue::from(adjusted_path)) else {
                    anyhow::bail!("readFileRef did not return a promise");
                };

                let read = JsFuture::from(Promise::unchecked_from_js(read)).await;

                let read = match read {
                    Ok(read) => read,
                    Err(err) => {
                        if let Some(e) = err.dyn_ref::<js_sys::Error>() {
                            if let Some(e_str) = e.message().as_string() {
                                anyhow::bail!("readFileRef failure: {}", e_str);
                            }
                        }

                        anyhow::bail!("readFileRef rejected: {:?}", err);
                    }
                };

                // TODO: how does JsValue -> Uint8Array work without try_from?
                Ok(Uint8Array::from(read).to_vec())
            }
        })
    }))
}

#[cfg(not(target_arch = "wasm32"))]
fn js_fn_to_baml_src_reader(get_baml_src_cb: js_sys::Function) -> BamlSrcReader {
    None
}

#[wasm_bindgen]
pub struct WasmCallContext {
    /// Index of the orchestration graph node to use for the call
    /// Defaults to 0 when unset
    node_index: Option<usize>,
}

#[wasm_bindgen]
impl WasmCallContext {
    #[wasm_bindgen(constructor)]
    pub fn new() -> Self {
        Self { node_index: None }
    }

    #[wasm_bindgen(setter)]
    pub fn set_node_index(&mut self, node_index: Option<usize>) {
        self.node_index = node_index;
    }
}

#[wasm_bindgen]
impl WasmFunction {
    #[wasm_bindgen]
    pub async fn render_prompt_for_test(
        &self,
        rt: &WasmRuntime,
        test_name: String,
        wasm_call_context: &WasmCallContext,
        get_baml_src_cb: js_sys::Function,
        env: js_sys::Object,
    ) -> JsResult<WasmPrompt> {
        let context_manager = rt.runtime.create_ctx_manager(
            BamlValue::String("wasm".to_string()),
            js_fn_to_baml_src_reader(get_baml_src_cb),
        );

        let test_type_builder = rt
            .runtime
            .internal()
            .get_test_type_builder(&self.name, &test_name)
            .map_err(|e| JsError::new(format!("{e:?}").as_str()))?;

        let entries = js_sys::Object::entries(&env);
        let mut env_vars = HashMap::new();
        for entry in entries.iter() {
            let arr = entry.dyn_into::<js_sys::Array>().unwrap();
            let key = arr.get(0).as_string().unwrap();
            let value = arr.get(1).as_string().unwrap_or_default();
            env_vars.insert(key, value);
        }

        let ctx = context_manager
            .create_ctx(
                test_type_builder.as_ref(),
                None,
                env_vars,
                vec![baml_ids::FunctionCallId::new()],
            )
            .map_err(|e| JsError::new(format!("{e:?}").as_str()))?;

        let params = rt
            .runtime
            .get_test_params(&self.name, &test_name, &ctx, false)
            .map_err(|e| JsError::new(format!("{e:?}").as_str()))?;

        rt.runtime
            .internal()
            .render_prompt(&self.name, &ctx, &params, wasm_call_context.node_index)
            .await
            .as_ref()
            .map(|(p, scope, allowed)| (p, scope, allowed).into())
            .map_err(|e| JsError::new(format!("{e:?}").as_str()))
    }

    #[wasm_bindgen]
    pub fn client_name(&self, rt: &WasmRuntime) -> Result<String, JsValue> {
        let rt: &BamlRuntime = &rt.runtime;
        let ctx_manager = rt.create_ctx_manager(BamlValue::String("wasm".to_string()), None);
        let ctx = ctx_manager.create_ctx_with_default();
        let ir = rt.internal().ir();
        let walker = ir
            .find_function(&self.name)
            .map_err(|e| JsValue::from_str(&format!("{e:?}")))?;
        let renderer = PromptRenderer::from_function(&walker, ir, &ctx)
            .map_err(|e| JsValue::from_str(&format!("{e:?}")))?;
        Ok(renderer.client_spec().to_string())
    }

    #[wasm_bindgen]
    pub async fn render_raw_curl_for_test(
        &self,
        rt: &WasmRuntime,
        test_name: String,
        wasm_call_context: &WasmCallContext,
        stream: bool,
        expand_images: bool,
        get_baml_src_cb: js_sys::Function,
        env: js_sys::Object,
        expose_secrets: bool,
    ) -> Result<String, wasm_bindgen::JsError> {
        let context_manager = rt.runtime.create_ctx_manager(
            BamlValue::String("wasm".to_string()),
            js_fn_to_baml_src_reader(get_baml_src_cb),
        );

        let test_type_builder = rt
            .runtime
            .internal()
            .get_test_type_builder(&self.name, &test_name)
            .map_err(|e| JsError::new(format!("{e:?}").as_str()))?;

        let entries = js_sys::Object::entries(&env);
        let mut env_vars = HashMap::new();
        for entry in entries.iter() {
            let arr = entry.dyn_into::<js_sys::Array>().unwrap();
            let key = arr.get(0).as_string().unwrap();
            let value = arr.get(1).as_string().unwrap_or_default();
            env_vars.insert(key, value);
        }

        let ctx = context_manager
            .create_ctx(
                test_type_builder.as_ref(),
                None,
                env_vars,
                vec![baml_ids::FunctionCallId::new()],
            )
            .map_err(|e| JsError::new(format!("{e:?}").as_str()))?;

        let params = rt
            .runtime
            .get_test_params(&self.name, &test_name, &ctx, false)
            .map_err(|e| JsError::new(format!("{e:?}").as_str()))?;

        let result = rt
            .runtime
            .internal()
            .render_prompt(&self.name, &ctx, &params, wasm_call_context.node_index)
            .await;

        let final_prompt = match result {
            Ok((prompt, _, _)) => match prompt {
                RenderedPrompt::Chat(chat_messages) => chat_messages,
                RenderedPrompt::Completion(_) => vec![], // or handle this case differently
            },
            Err(e) => return Err(wasm_bindgen::JsError::new(format!("{e:?}").as_str())),
        };

        rt.runtime
            .internal()
            .render_raw_curl(
                &self.name,
                &ctx,
                &final_prompt,
                RenderCurlSettings {
                    stream,
                    as_shell_commands: !expand_images,
                    expose_secrets,
                },
                wasm_call_context.node_index,
            )
            .await
            .map_err(|e| wasm_bindgen::JsError::new(format!("{e:?}").as_str()))
    }

    #[wasm_bindgen]
    pub async fn run_test_with_expr_events(
        &self,
        rt: &mut WasmRuntime,
        test_name: String,
        on_partial_response: js_sys::Function,
        get_baml_src_cb: js_sys::Function,
        on_expr_event: js_sys::Function,
        env: js_sys::Object,
        abort_signal: Option<js_sys::Object>,
    ) -> Result<WasmTestResponse, JsValue> {
        // Convert abort signal to tripwire
        let tripwire = js_abort_signal_to_tripwire(abort_signal).map_err(JsValue::from)?;

        let rt = &rt.runtime;
        let function_name = self.name.clone();

        let function_name_for_test_pair = function_name.clone();
        let test_name_for_test_pair = test_name.clone();

        // Create the closure to handle partial responses:
        let cb = Box::new(move |r: FunctionResult| {
            let this = JsValue::NULL;
            let res = WasmFunctionResponse {
                function_response: r,
                func_test_pair: WasmFunctionTestPair {
                    function_name: function_name_for_test_pair.clone(),
                    test_name: test_name_for_test_pair.clone(),
                },
            }
            .into();
            on_partial_response.call1(&this, &res).unwrap();
        });

        // Create the channel for expression events
        let (tx, mut rx) = mpsc::unbounded::<Vec<SerializedSpan>>();

        // Spawn a task to handle expression events
        let on_expr_event_clone = on_expr_event.clone();
        wasm_bindgen_futures::spawn_local(async move {
            while let Some(spans) = rx.next().await {
                let this = JsValue::NULL;
                match serde_wasm_bindgen::to_value(&spans) {
                    Ok(res) => {
                        on_expr_event_clone.call1(&this, &res).expect("TODO");
                    }
                    Err(e) => {
                        log::error!("Error serializing spans: {e}");
                    }
                }
            }
        });

        // Create your evaluation context, etc.
        let ctx = rt.create_ctx_manager_for_wasm(js_fn_to_baml_src_reader(get_baml_src_cb));

        let entries = js_sys::Object::entries(&env);
        let mut env_vars = HashMap::new();
        for entry in entries.iter() {
            let arr = entry.dyn_into::<js_sys::Array>().unwrap();
            let key = arr.get(0).as_string().unwrap();
            let value = arr.get(1).as_string().unwrap_or_default();
            env_vars.insert(key, value);
        }

        // Pass the sender to run_test_with_expr_events with tripwire support
        let on_tick = if false { Some(|| {}) } else { None };
        let result = rt
            .run_test_with_expr_events(
                &function_name,
                &test_name,
                &ctx,
                Some(cb),
                Some(tx),
                None,
                env_vars.clone(),
                tripwire,
                on_tick,
            )
            .await;

        let (test_response, span) = result;

        Ok(WasmTestResponse {
            test_response,
            span: Some(span.to_string()),
            tracing_project_id: rt
                .tracer_wrapper
                .get_or_create_tracer(&env_vars)
                .tracing_project_id(),
            func_test_pair: WasmFunctionTestPair {
                function_name,
                test_name,
            },
        })
    }

    #[wasm_bindgen]
    pub async fn run_test(
        &self,
        rt: &mut WasmRuntime,
        test_name: String,
        on_partial_response: js_sys::Function,
        get_baml_src_cb: js_sys::Function,
        env: js_sys::Object,
        abort_signal: Option<js_sys::Object>,
    ) -> Result<WasmTestResponse, JsValue> {
        // Convert abort signal to tripwire
        let tripwire = js_abort_signal_to_tripwire(abort_signal).map_err(JsValue::from)?;

        let rt = &rt.runtime;
        let function_name = self.name.clone();

        let function_name_for_test_pair = function_name.clone();
        let test_name_for_test_pair = test_name.clone();

        // Create the closure to handle partial responses:
        let cb = Box::new(move |r| {
            let this = JsValue::NULL;
            let res = WasmFunctionResponse {
                function_response: r,
                func_test_pair: WasmFunctionTestPair {
                    function_name: function_name_for_test_pair.clone(),
                    test_name: test_name_for_test_pair.clone(),
                },
            }
            .into();
            on_partial_response.call1(&this, &res).unwrap();
        });

        // Create your evaluation context, etc.
        let ctx = rt.create_ctx_manager_for_wasm(js_fn_to_baml_src_reader(get_baml_src_cb));

        let entries = js_sys::Object::entries(&env);
        let mut env_vars = HashMap::new();
        for entry in entries.iter() {
            let arr = entry.dyn_into::<js_sys::Array>().unwrap();
            let key = arr.get(0).as_string().unwrap();
            let value = arr.get(1).as_string().unwrap_or_default();
            env_vars.insert(key, value);
        }
        // Now pass collector_arc to your runtime's run_test with tripwire support
        let on_tick = if false { Some(|| {}) } else { None };
        let result = rt
            .run_test(
                &function_name,
                &test_name,
                &ctx,
                Some(cb),
                None,
                env_vars.clone(),
                tripwire,
                on_tick,
            )
            .await;

        let (test_response, span) = result;

        Ok(WasmTestResponse {
            test_response,
            span: Some(span.to_string()),
            tracing_project_id: rt
                .tracer_wrapper
                .get_or_create_tracer(&env_vars)
                .tracing_project_id(),
            func_test_pair: WasmFunctionTestPair {
                function_name,
                test_name,
            },
        })
    }

    pub fn function_graph(&self, rt: &WasmRuntime) -> Result<String, JsValue> {
        let rt: &BamlRuntime = &rt.runtime;
        let ctx = rt
            .create_ctx_manager(BamlValue::String("wasm".to_string()), None)
            .create_ctx_with_default();
        let graph = rt
            .internal()
            .function_graph(&self.name, &ctx)
            .map_err(|e| JsValue::from_str(&format!("{e:?}")))?;
        Ok(graph)
    }

    pub fn orchestration_graph(&self, rt: &WasmRuntime) -> Result<Vec<WasmScope>, JsValue> {
        let rt: &BamlRuntime = &rt.runtime;

        let ctx = rt
            .create_ctx_manager(BamlValue::String("wasm".to_string()), None)
            .create_ctx_with_default();

        let ir = rt.internal().ir();
        let walker = ir
            .find_function(&self.name)
            .map_err(|e| JsValue::from_str(&format!("{e:?}")))?;
        let renderer = PromptRenderer::from_function(&walker, ir, &ctx)
            .map_err(|e| JsValue::from_str(&format!("{e:?}")))?;
        let client_spec = renderer.client_spec();

        let graph = rt
            .internal()
            .orchestration_graph(client_spec, &ctx)
            .map_err(|e| JsValue::from_str(&format!("{e:?}")))?;

        // Serialize the scopes to JsValue
        let mut scopes = Vec::new();
        for scope in graph {
            scopes.push(WasmScope::from(scope.scope));
        }
        Ok(scopes)
    }
}
trait ToJsValue {
    fn to_js_value(&self) -> JsValue;
}

impl ToJsValue for ExecutionScope {
    fn to_js_value(&self) -> JsValue {
        let obj = js_sys::Object::new();
        let set_property = |obj: &js_sys::Object, key: &str, value: JsValue| {
            js_sys::Reflect::set(obj, &JsValue::from_str(key), &value).is_ok()
        };

        match self {
            ExecutionScope::Direct(name) => {
                set_property(&obj, "type", JsValue::from_str("Direct"));
                set_property(&obj, "name", JsValue::from_str(name));
            }
            ExecutionScope::Retry(name, count, delay) => {
                set_property(&obj, "type", JsValue::from_str("Retry"));
                set_property(&obj, "name", JsValue::from_str(name));
                set_property(&obj, "count", JsValue::from_f64(*count as f64));
                set_property(&obj, "delay", JsValue::from_f64(delay.as_millis() as f64));
            }
            ExecutionScope::RoundRobin(strategy, index) => {
                set_property(&obj, "type", JsValue::from_str("RoundRobin"));
                set_property(
                    &obj,
                    "strategy_name",
                    JsValue::from_str(&format!("{:?}", strategy.name)),
                );
                set_property(&obj, "index", JsValue::from_f64(*index as f64));
            }
            ExecutionScope::Fallback(name, index) => {
                set_property(&obj, "type", JsValue::from_str("Fallback"));
                set_property(&obj, "name", JsValue::from_str(name));
                set_property(&obj, "index", JsValue::from_f64(*index as f64));
            }
        }
        obj.into()
    }
}

impl ToJsValue for OrchestrationScope {
    fn to_js_value(&self) -> JsValue {
        let array = js_sys::Array::new();
        for scope in &self.scope {
            array.push(&scope.to_js_value());
        }
        array.into()
    }
}