feat(zds+kcl): scaffolding to add feature flag for KCL around WASM instance lifecycle (#12453)

Author: therapon

rust/kcl-lib/src/lib.rs

@@ -71,6 +71,7 @@ mod lsp;
 mod modules;
 mod parsing;
 mod project;
+mod runtime_flags;
 mod settings;
 #[cfg(test)]
 mod simulation_tests;
@@ -141,6 +142,10 @@ pub use parsing::ast::types::NodePath;
 pub use parsing::ast::types::NodePathExt;
 pub use parsing::ast::types::Step as NodePathStep;
 pub use project::ProjectManager;
+pub use runtime_flags::KclRuntimeFlags;
+pub use runtime_flags::RuntimeFlag;
+pub use runtime_flags::kcl_runtime_flags;
+pub use runtime_flags::set_kcl_runtime_flags;
 pub use settings::types::Configuration;
 pub use settings::types::project::ProjectConfiguration;
 #[cfg(not(target_arch = "wasm32"))]

rust/kcl-lib/src/parsing/token/mod.rs

@@ -19,8 +19,10 @@ use winnow::{self};
 
 use crate::CompilationIssue;
 use crate::ModuleId;
+use crate::RuntimeFlag;
 use crate::SourceRange;
 use crate::errors::KclError;
+use crate::kcl_runtime_flags;
 use crate::parsing::ast::types::ItemVisibility;
 use crate::parsing::ast::types::VariableKind;
 
@@ -601,8 +603,8 @@ pub(crate) const KCL_LEXER_ENV_VAR: &str = "KCL_LEXER";
 /// `KCL_LEXER` environment variable, so a process can pick either lexer without a
 /// rebuild.
 ///
-/// Precedence: test override > `KCL_LEXER` > [`LexerMode::DEFAULT`]. This mirrors
-/// the existing `KCL_MEMORY_IMPL` selector in `execution::memory`.
+/// Precedence: runtime flags > test override > `KCL_LEXER` >
+/// [`LexerMode::DEFAULT`].
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) enum LexerMode {
     Old,
@@ -615,14 +617,9 @@ impl LexerMode {
 
     /// Resolve the active lexer mode (see precedence on [`LexerMode`]).
     pub(crate) fn resolve() -> Self {
-        #[cfg(test)]
-        if let Some(mode) = Self::test_override() {
-            return mode;
-        }
-
-        match env::var(KCL_LEXER_ENV_VAR) {
-            Ok(value) => Self::parse(&value),
-            Err(env::VarError::NotPresent) => Self::DEFAULT,
+        let env_value = match env::var(KCL_LEXER_ENV_VAR) {
+            Ok(value) => Some(value),
+            Err(env::VarError::NotPresent) => None,
             Err(env::VarError::NotUnicode(value)) => {
                 // Invalid-unicode env var: warn and fall back rather than crash.
                 Self::warn_once(|| {
@@ -631,9 +628,39 @@ impl LexerMode {
                         value.to_string_lossy()
                     )
                 });
-                Self::Old
+                None
             }
+        };
+
+        Self::resolve_from_sources(
+            kcl_runtime_flags().use_new_lexer_parser,
+            Self::test_override_for_resolve(),
+            env_value.as_deref(),
+        )
+    }
+
+    fn resolve_from_sources(runtime_flag: RuntimeFlag, test_override: Option<Self>, env_value: Option<&str>) -> Self {
+        match runtime_flag {
+            RuntimeFlag::On => return Self::New,
+            RuntimeFlag::Off => return Self::Old,
+            RuntimeFlag::Unset => {}
+        }
+
+        if let Some(mode) = test_override {
+            return mode;
         }
+
+        env_value.map(Self::parse).unwrap_or(Self::DEFAULT)
+    }
+
+    #[cfg(test)]
+    fn test_override_for_resolve() -> Option<Self> {
+        Self::test_override()
+    }
+
+    #[cfg(not(test))]
+    fn test_override_for_resolve() -> Option<Self> {
+        None
     }
 
     fn parse(value: &str) -> Self {
@@ -759,10 +786,23 @@ fn lex_legacy(s: &str, module_id: ModuleId) -> Result<TokenStream, KclError> {
 mod lexer_mode_tests {
     use super::LexerMode;
     use super::lex;
+    use crate::KclRuntimeFlags;
     use crate::ModuleId;
+    use crate::RuntimeFlag;
+
+    fn set_runtime_lexer_flag(flag: RuntimeFlag) {
+        crate::set_kcl_runtime_flags(KclRuntimeFlags {
+            use_new_lexer_parser: flag,
+        });
+    }
+
+    fn reset_runtime_lexer_flags() {
+        crate::set_kcl_runtime_flags(KclRuntimeFlags::DEFAULT);
+    }
 
     #[test]
     fn default_mode_is_old() {
+        reset_runtime_lexer_flags();
         assert_eq!(LexerMode::DEFAULT, LexerMode::Old);
     }
 
@@ -780,6 +820,7 @@ mod lexer_mode_tests {
 
     #[test]
     fn override_guard_sets_and_restores_mode() {
+        reset_runtime_lexer_flags();
         // Reserved for dispatch/integration tests; relies on the process-global
         // atomic, which is race-free under nextest's process isolation.
         {
@@ -790,6 +831,75 @@ mod lexer_mode_tests {
         assert_eq!(LexerMode::resolve(), LexerMode::Old);
     }
 
+    #[test]
+    fn runtime_flags_default_to_unset() {
+        reset_runtime_lexer_flags();
+        assert_eq!(
+            crate::kcl_runtime_flags(),
+            KclRuntimeFlags {
+                use_new_lexer_parser: RuntimeFlag::Unset,
+            }
+        );
+    }
+
+    #[test]
+    fn runtime_flag_on_selects_new_lexer() {
+        reset_runtime_lexer_flags();
+        set_runtime_lexer_flag(RuntimeFlag::On);
+        assert_eq!(LexerMode::resolve(), LexerMode::New);
+    }
+
+    #[test]
+    fn runtime_flag_off_selects_old_lexer() {
+        reset_runtime_lexer_flags();
+        set_runtime_lexer_flag(RuntimeFlag::Off);
+        assert_eq!(LexerMode::resolve(), LexerMode::Old);
+    }
+
+    #[test]
+    fn runtime_flag_takes_priority_over_test_override_and_env() {
+        assert_eq!(
+            LexerMode::resolve_from_sources(RuntimeFlag::Off, Some(LexerMode::New), Some("new")),
+            LexerMode::Old
+        );
+        assert_eq!(
+            LexerMode::resolve_from_sources(RuntimeFlag::On, Some(LexerMode::Old), Some("old")),
+            LexerMode::New
+        );
+    }
+
+    #[test]
+    fn unset_runtime_flag_allows_env_to_select_lexer() {
+        assert_eq!(
+            LexerMode::resolve_from_sources(RuntimeFlag::Unset, None, Some("new")),
+            LexerMode::New
+        );
+        assert_eq!(
+            LexerMode::resolve_from_sources(RuntimeFlag::Unset, None, Some("old")),
+            LexerMode::Old
+        );
+    }
+
+    #[test]
+    fn unset_runtime_flag_and_missing_env_selects_default_lexer() {
+        assert_eq!(
+            LexerMode::resolve_from_sources(RuntimeFlag::Unset, None, None),
+            LexerMode::DEFAULT
+        );
+    }
+
+    #[test]
+    fn test_override_takes_priority_over_env() {
+        assert_eq!(
+            LexerMode::resolve_from_sources(RuntimeFlag::Unset, Some(LexerMode::Old), Some("new")),
+            LexerMode::Old
+        );
+        assert_eq!(
+            LexerMode::resolve_from_sources(RuntimeFlag::Unset, Some(LexerMode::New), Some("old")),
+            LexerMode::New
+        );
+    }
+
     /// Exercises the `New` arm of `lex` in default CI: no `KCL_LEXER` env var is
     /// set; the new lexer is selected via the process-global test override (which
     /// is race-free under nextest's process-per-test isolation).
@@ -801,6 +911,7 @@ mod lexer_mode_tests {
     /// not merely that some lexer ran.
     #[test]
     fn lex_dispatches_to_new_lexer() {
+        reset_runtime_lexer_flags();
         let _guard = LexerMode::override_for_test(LexerMode::New);
         assert_eq!(LexerMode::resolve(), LexerMode::New);
 

rust/kcl-lib/src/runtime_flags.rs

@@ -0,0 +1,59 @@
+use std::sync::RwLock;
+
+use serde::Deserialize;
+use serde::Serialize;
+
+/// Runtime representation for feature flags that can be set by the TS app.
+///
+/// TS currently provides a two-state feature answer: `true` means the feature is
+/// on, while `false` covers both explicit off/default behavior and a missing
+/// feature entry. Rust keeps a third state so code that was not initialized
+/// through the TS/wasm path can still fall back to Rust-side defaults, such as
+/// env-based configuration.
+#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq)]
+pub enum RuntimeFlag {
+    /// No TS/wasm runtime flag has been installed; fall back to Rust defaults.
+    #[default]
+    Unset,
+    /// TS observed the feature as on; use the new feature behavior.
+    On,
+    /// TS observed the feature as false or missing; use default behavior.
+    Off,
+}
+
+/// Maps 1-1 to the KCL related flags added to the Admin portal and TS.
+#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
+pub struct KclRuntimeFlags {
+    pub use_new_lexer_parser: RuntimeFlag,
+}
+
+impl KclRuntimeFlags {
+    pub const DEFAULT: Self = Self {
+        use_new_lexer_parser: RuntimeFlag::Unset,
+    };
+}
+
+impl Default for KclRuntimeFlags {
+    fn default() -> Self {
+        Self::DEFAULT
+    }
+}
+
+static KCL_RUNTIME_FLAGS: RwLock<KclRuntimeFlags> = RwLock::new(KclRuntimeFlags::DEFAULT);
+
+pub fn set_kcl_runtime_flags(flags: KclRuntimeFlags) {
+    match KCL_RUNTIME_FLAGS.write() {
+        Ok(mut guard) => *guard = flags,
+        Err(poisoned) => {
+            let mut guard = poisoned.into_inner();
+            *guard = flags;
+        }
+    }
+}
+
+pub fn kcl_runtime_flags() -> KclRuntimeFlags {
+    match KCL_RUNTIME_FLAGS.read() {
+        Ok(guard) => *guard,
+        Err(poisoned) => *poisoned.into_inner(),
+    }
+}

rust/kcl-wasm-lib/src/wasm.rs

@@ -1,6 +1,7 @@
 //! Wasm bindings for `kcl`.
 
 use gloo_utils::format::JsValueSerdeExt;
+use kcl_lib::KclRuntimeFlags;
 use kcl_lib::Program;
 use kcl_lib::SourceRange;
 use kcl_lib::exec::NumericType;
@@ -48,6 +49,15 @@ pub fn parse_wasm(kcl_program_source: &str) -> Result<JsValue, String> {
     JsValue::from_serde(&(program, errs)).map_err(|e| e.to_string())
 }
 
+#[wasm_bindgen]
+pub fn set_kcl_runtime_flags(flags_json: &str) -> Result<(), String> {
+    console_error_panic_hook::set_once();
+
+    let flags: KclRuntimeFlags = serde_json::from_str(flags_json).map_err(|e| e.to_string())?;
+    kcl_lib::set_kcl_runtime_flags(flags);
+    Ok(())
+}
+
 // wasm_bindgen wrapper for recast
 // test for this function and by extension the recaster are done in javascript land src/lang/recast.test.ts
 #[wasm_bindgen]

src/lang/wasmUtils.test.ts

@@ -0,0 +1,68 @@
+import type { ModuleType } from '@src/lib/wasm_lib_wrapper'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => ({
+  getModule: vi.fn(),
+  init: vi.fn(),
+  notifyActiveWasmInstance: vi.fn(),
+  reloadModule: vi.fn(),
+}))
+
+vi.mock('@src/lib/wasm_lib_wrapper', () => ({
+  getModule: mocks.getModule,
+  init: mocks.init,
+  reloadModule: mocks.reloadModule,
+}))
+
+vi.mock('@src/lib/wasmLifecycle', () => ({
+  notifyActiveWasmInstance: mocks.notifyActiveWasmInstance,
+}))
+
+const wasmModule = {
+  default: vi.fn(),
+  import_file_extensions: vi.fn(),
+} as unknown as ModuleType & {
+  default: ReturnType<typeof vi.fn>
+}
+
+describe('wasm utils', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mocks.getModule.mockReturnValue(wasmModule)
+    wasmModule.default.mockResolvedValue(undefined)
+    globalThis.fetch = vi.fn().mockResolvedValue({
+      arrayBuffer: vi.fn().mockResolvedValue(new ArrayBuffer(8)),
+    })
+  })
+
+  it('notifies listeners after browser wasm initialization', async () => {
+    const { initialiseWasm } = await import('@src/lang/wasmUtils')
+
+    await expect(initialiseWasm()).resolves.toBe(wasmModule)
+
+    expect(mocks.reloadModule).toHaveBeenCalled()
+    expect(mocks.init).toHaveBeenCalled()
+    expect(mocks.notifyActiveWasmInstance).toHaveBeenCalledWith(wasmModule)
+  })
+
+  it('starts wasm from a buffer and notifies listeners', async () => {
+    const { startWasmFromBuffer } = await import('@src/lang/wasmUtils')
+    const buffer = new ArrayBuffer(8)
+
+    await expect(startWasmFromBuffer(buffer)).resolves.toBe(wasmModule)
+
+    expect(mocks.reloadModule).toHaveBeenCalled()
+    expect(mocks.init).toHaveBeenCalledWith({ module_or_path: buffer })
+    expect(mocks.notifyActiveWasmInstance).toHaveBeenCalledWith(wasmModule)
+  })
+
+  it('restarts wasm and notifies listeners', async () => {
+    const { restartWasmModule } = await import('@src/lang/wasmUtils')
+
+    await expect(restartWasmModule()).resolves.toBe(wasmModule)
+
+    expect(mocks.reloadModule).toHaveBeenCalled()
+    expect(wasmModule.default).toHaveBeenCalled()
+    expect(mocks.notifyActiveWasmInstance).toHaveBeenCalledWith(wasmModule)
+  })
+})